diff --git a/.ci/run_test.sh b/.ci/run_test.sh
deleted file mode 100644
index afeac3a3..00000000
--- a/.ci/run_test.sh
+++ /dev/null
@@ -1,51 +0,0 @@
-set -ex
-TEST_DIR=$1
-OUTPUT_DIR=$2
-
-
-run_test(){
-    echo "$OUTPUT_DIR"
-    ./"$1" 
-    LOCAL_SUCCESS=$?
-    if [[ $LOCAL_SUCCESS == 0 ]]; then
-        echo "PASSED : $1 " >> "$OUTPUT_DIR"/results.txt
-    else
-        echo "FAILED : $1 " >> "$OUTPUT_DIR"/results.txt
-    fi
-}
-
-
-pushd $TEST_DIR
-etcd --listen-client-urls http://0.0.0.0:2379 --advertise-client-urls http://10.0.0.1:2379 &
-export MC_GID_INDEX=1
-run_test transport_uint_test
-run_test tcp_transport_test
-export MC_GID_INDEX=1
-
-# sudo pkill etcd
-# sleep 1
-# etcd --listen-client-urls http://0.0.0.0:2379 --advertise-client-urls http://10.0.0.1:2379 &
-# sudo fuser -k 14345/tcp
-export MC_GID_INDEX=1 && ./rdma_transport_test --mode=target  --metadata_server=127.0.0.1:2379 --local_server_name=127.0.0.2:14345 --device_name=erdma_0 &
-export MC_GID_INDEX=1 && ./rdma_transport_test --metadata_server=127.0.0.1:2379 --segment_id=127.0.0.2:14345 --local_server_name=127.0.0.3:14346 --device_name=erdma_1
-LOCAL_SUCCESS=$?
-if [[ $LOCAL_SUCCESS == 0 ]]; then
-    echo "PASSED : rdma_transport_test " >> "$OUTPUT_DIR"/results.txt
-else
-    echo "FAILED : rdma_transport_test " >> "$OUTPUT_DIR"/results.txt
-fi
-
-# sudo pkill etcd
-# sleep 1
-# etcd --listen-client-urls http://0.0.0.0:2379 --advertise-client-urls http://10.0.0.1:2379 &
-# export MC_GID_INDEX=1
-# sudo fuser -k 15345/tcp
-# export MC_GID_INDEX=1 && ./rdma_transport_test --mode=target  --metadata_server=127.0.0.1:2379 --local_server_name=127.0.0.2:14345 --device_name=erdma_0 &
-export MC_GID_INDEX=1 && ./rdma_transport_test2 --metadata_server=127.0.0.1:2379 --segment_id=127.0.0.2:14345 --local_server_name=127.0.0.3:14346 --device_name=erdma_1
-LOCAL_SUCCESS=$?
-if [[ $LOCAL_SUCCESS == 0 ]]; then
-    echo "PASSED : rdma_transport_test2 " >> "$OUTPUT_DIR"/results.txt
-else
-    echo "FAILED : rdma_transport_test2 " >> "$OUTPUT_DIR"/results.txt
-fi
-popd
\ No newline at end of file
diff --git a/.claude-plugin/marketplace.json b/.claude-plugin/marketplace.json
new file mode 100644
index 00000000..e15c8550
--- /dev/null
+++ b/.claude-plugin/marketplace.json
@@ -0,0 +1,42 @@
+{
+  "$schema": "https://json.schemastore.org/claude-code-marketplace.json",
+  "name": "mooncake",
+  "owner": {
+    "name": "Mooncake Team"
+  },
+  "metadata": {
+    "description": "Built-in Claude Code skills for working with Mooncake: deployment troubleshooting, local CI validation, and the Python API."
+  },
+  "plugins": [
+    {
+      "name": "mooncake-troubleshoot",
+      "source": {
+        "source": "git-subdir",
+        "url": "https://github.com/kvcache-ai/Mooncake.git",
+        "path": ".claude/skills/mooncake-troubleshoot"
+      },
+      "description": "Systematically diagnose Mooncake deployment and runtime issues (services, RDMA, env vars, connectivity, logs) and propose fixes.",
+      "homepage": "https://kvcache-ai.github.io/Mooncake/skills/mooncake-troubleshoot.html"
+    },
+    {
+      "name": "mooncake-ci-local",
+      "source": {
+        "source": "git-subdir",
+        "url": "https://github.com/kvcache-ai/Mooncake.git",
+        "path": ".claude/skills/mooncake-ci-local"
+      },
+      "description": "Run Mooncake pre-PR local validation via scripts/run_ci_test.sh, reproducing the reproducible parts of GitHub Actions.",
+      "homepage": "https://kvcache-ai.github.io/Mooncake/skills/mooncake-ci-local.html"
+    },
+    {
+      "name": "mooncake-api",
+      "source": {
+        "source": "git-subdir",
+        "url": "https://github.com/kvcache-ai/Mooncake.git",
+        "path": ".claude/skills/mooncake-api"
+      },
+      "description": "Work with the Mooncake Store, Transfer Engine, and EP/Backend Python APIs.",
+      "homepage": "https://kvcache-ai.github.io/Mooncake/skills/mooncake-api.html"
+    }
+  ]
+}
diff --git a/.claude/skills/mooncake-api/.claude-plugin/plugin.json b/.claude/skills/mooncake-api/.claude-plugin/plugin.json
new file mode 100644
index 00000000..81bad701
--- /dev/null
+++ b/.claude/skills/mooncake-api/.claude-plugin/plugin.json
@@ -0,0 +1,12 @@
+{
+  "$schema": "https://json.schemastore.org/claude-code-plugin-manifest.json",
+  "name": "mooncake-api",
+  "description": "Work with the Mooncake Store, Transfer Engine, and EP/Backend Python APIs.",
+  "author": {
+    "name": "Mooncake Team"
+  },
+  "homepage": "https://kvcache-ai.github.io/Mooncake/skills/mooncake-api.html",
+  "repository": "https://github.com/kvcache-ai/Mooncake",
+  "license": "Apache-2.0",
+  "keywords": ["mooncake", "python", "api", "transfer-engine", "store"]
+}
diff --git a/.claude/skills/mooncake-api/SKILL.md b/.claude/skills/mooncake-api/SKILL.md
index daa1b6c5..d0edb73a 100644
--- a/.claude/skills/mooncake-api/SKILL.md
+++ b/.claude/skills/mooncake-api/SKILL.md
@@ -1,4 +1,7 @@
-<!-- Let's perfect this skill together. -->
+---
+name: mooncake-api
+description: Help users work with the Mooncake Python APIs for distributed storage and high-performance data transfer. Use when working with Mooncake Store (distributed KV cache), Transfer Engine (RDMA/TCP transfers), service setup (master, metadata server), PyTorch tensors in the Store, zero-copy/buffer management, batch operations and replication, or Mooncake EP / Backend (Expert Parallelism). Trigger on questions about MooncakeDistributedStore, TransferEngine, put/get, put_tensor, register_buffer, ReplicateConfig, or the mooncake.store / mooncake.engine / mooncake.pg Python modules.
+---
 
 # Mooncake Python API Skill
 
diff --git a/.claude/skills/mooncake-ci-local/.claude-plugin/plugin.json b/.claude/skills/mooncake-ci-local/.claude-plugin/plugin.json
new file mode 100644
index 00000000..afb3699b
--- /dev/null
+++ b/.claude/skills/mooncake-ci-local/.claude-plugin/plugin.json
@@ -0,0 +1,12 @@
+{
+  "$schema": "https://json.schemastore.org/claude-code-plugin-manifest.json",
+  "name": "mooncake-ci-local",
+  "description": "Run Mooncake pre-PR local validation via scripts/run_ci_test.sh, reproducing the reproducible parts of GitHub Actions.",
+  "author": {
+    "name": "Mooncake Team"
+  },
+  "homepage": "https://kvcache-ai.github.io/Mooncake/skills/mooncake-ci-local.html",
+  "repository": "https://github.com/kvcache-ai/Mooncake",
+  "license": "Apache-2.0",
+  "keywords": ["mooncake", "ci", "validation", "pre-pr"]
+}
diff --git a/.claude/skills/mooncake-ci-local/SKILL.md b/.claude/skills/mooncake-ci-local/SKILL.md
index 0729957a..546bcc24 100644
--- a/.claude/skills/mooncake-ci-local/SKILL.md
+++ b/.claude/skills/mooncake-ci-local/SKILL.md
@@ -1,266 +1,159 @@
 ---
 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".
+description: Run Mooncake pre-PR local validation through scripts/run_ci_test.sh. Use this skill whenever the user wants to validate a branch before opening or submitting a PR, run local CI, run ci test, check changes before PR, reproduce GitHub Actions locally, or force a full pre-submit verification. Trigger on phrases like "提交 PR 前验证", "run ci test", "run local CI", "check my branch", "test before PR", "pre-submit validation", and "reproduce CI locally".
 ---
 
-# Mooncake CI Local Test Runner
+# Mooncake Pre-PR Local Validation
 
-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.
+Use `bash scripts/run_ci_test.sh` as the default entry point. This is the single local lane for PR-before-submit validation, and it already coordinates the reproducible parts of GitHub Actions.
 
-## CI Test Layers
+## Default Entry Point
 
-### Layer 1 — C++ Unit Tests (ctest)
-**CI equivalent:** `build` job in `ci.yml` — "Test (in build env) with coverage"
+When the user asks for any of the following, run the repo script first instead of reconstructing the workflow by hand:
 
-**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
+- 提交 PR 前本地验证
+- run ci test
+- run local CI
+- check my branch before PR
+- reproduce CI locally
 
-# 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:
+Default command:
 
 ```bash
-bash .claude/skills/mooncake-ci-local/scripts/check-prerequisites.sh
+bash scripts/run_ci_test.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
+What this script already covers:
 
-**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 .
-```
+- GitHub-like `paths-filter` against `origin/main`
+- `typos`
+- `scripts/code_format.sh --check`
+- default CMake configure/build/install in `build-ci-local`
+- `ctest`
+- wheel build in `build-wheel-local`
+- wheel installation validation
+- `scripts/run_tests.sh`
+- selected Python API and integration tests
+- per-stage summary and logs under `local_test/run-ci-logs/<timestamp>/`
 
-**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
+## Standard Agent Workflow
 
-### Step 3 — Run and report
+1. Run `bash scripts/run_ci_test.sh` from the repo root unless the user explicitly asks for a narrower subset.
+2. Read the stage summary instead of dumping raw terminal output.
+3. Report these items back to the user:
+   - passed stages
+   - failed stages
+   - blocked stages
+   - unsupported stages
+   - whether `paths-filter` skipped downstream stages
+   - the log directory under `local_test/run-ci-logs/...`
+4. If there is a failure, inspect the corresponding stage log and summarize the root cause.
 
-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 Options
 
-## 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
+Force a full lane even if `paths-filter` would skip downstream stages:
 
 ```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
+bash scripts/run_ci_test.sh --skip-path-filter
 ```
+
+Use another base ref:
+
+```bash
+bash scripts/run_ci_test.sh --base origin/main
+```
+
+Auto-install missing dependencies:
+
+```bash
+bash scripts/run_ci_test.sh --install-deps
+```
+
+Keep services running for follow-up debugging:
+
+```bash
+bash scripts/run_ci_test.sh --keep-services
+```
+
+## Minimal Example
+
+User prompt:
+
+- 提交 PR 前，帮我跑一遍本地 CI 验证当前分支。
+
+Expected action:
+
+```bash
+bash scripts/run_ci_test.sh
+```
+
+If the user wants to ignore changed-path optimization and force the full lane:
+
+```bash
+bash scripts/run_ci_test.sh --skip-path-filter
+```
+
+See also `.claude/skills/mooncake-ci-local/examples/minimal.md`.
+
+## How To Interpret Results
+
+- `passed`: the stage succeeded locally.
+- `failed`: the stage reproduced a real local failure and needs investigation.
+- `blocked`: local environment or dependency issue prevented execution.
+- `unsupported`: intentionally not run in the local lane because it needs external platforms, special hardware, or a non-default build.
+
+If `paths-filter` skips downstream stages, explain that the current branch changed only non-source paths relative to the selected base.
+
+## Current Local Coverage
+
+Included by default:
+
+- spell check
+- code format check
+- default ASan CMake lane in `build-ci-local`
+- `ctest`
+- wheel build and installation test
+- `scripts/run_tests.sh`
+- selected Python API tests
+
+Unsupported by design in the default local lane:
+
+- Ascend jobs
+- T-one integration jobs
+- MUSA jobs
+- Docker image build jobs
+- CUDA 13 wheel jobs
+- PG-backend tests absent from the default wheel build
+- Python drain-http API stage in the local ASan lane
+
+## Targeted Reruns For Debugging
+
+Use targeted reruns only after the full script identifies a failing area, or when the user explicitly asks for a smaller scope.
+
+Rerun a specific C++ test pattern:
+
+```bash
+cd build-ci-local
+MC_METADATA_SERVER=http://127.0.0.1:8080/metadata DEFAULT_KV_LEASE_TTL=500 ctest -R <pattern> --output-on-failure
+```
+
+Rerun the Python wheel integration lane:
+
+```bash
+source test_env/bin/activate
+MC_STORE_MEMCPY=false TEST_SSD_OFFLOAD_IN_EVICT=true ./scripts/run_tests.sh
+```
+
+Rerun the safetensor unittest:
+
+```bash
+source test_env/bin/activate
+python -m unittest mooncake-wheel.tests.test_safetensor_functions
+```
+
+## Notes For The Agent
+
+- Prefer the repo script over rebuilding the CI workflow step by step.
+- Preserve the separation between `build-ci-local` and `build-wheel-local`.
+- Summarize failing stages from their logs instead of pasting raw output.
+- If the user only asks whether the branch is safe before opening a PR, the default answer path is `bash scripts/run_ci_test.sh`.
diff --git a/.claude/skills/mooncake-ci-local/examples/minimal.md b/.claude/skills/mooncake-ci-local/examples/minimal.md
new file mode 100644
index 00000000..3d081c73
--- /dev/null
+++ b/.claude/skills/mooncake-ci-local/examples/minimal.md
@@ -0,0 +1,30 @@
+# Minimal Example
+
+Goal: validate the current branch before opening or submitting a PR.
+
+User prompt:
+
+- 提交 PR 前，帮我跑一遍本地 CI 验证当前分支。
+
+Expected action:
+
+```bash
+bash scripts/run_ci_test.sh
+```
+
+Force a full rerun even if `paths-filter` would skip downstream stages:
+
+```bash
+bash scripts/run_ci_test.sh --skip-path-filter
+```
+
+Result location:
+
+- `local_test/run-ci-logs/<timestamp>/`
+
+Typical report format:
+
+- passed stages
+- failed or blocked stages
+- unsupported stages
+- first actionable failure and its log path
diff --git a/.claude/skills/mooncake-troubleshoot/.claude-plugin/plugin.json b/.claude/skills/mooncake-troubleshoot/.claude-plugin/plugin.json
new file mode 100644
index 00000000..239659ad
--- /dev/null
+++ b/.claude/skills/mooncake-troubleshoot/.claude-plugin/plugin.json
@@ -0,0 +1,12 @@
+{
+  "$schema": "https://json.schemastore.org/claude-code-plugin-manifest.json",
+  "name": "mooncake-troubleshoot",
+  "description": "Systematically diagnose Mooncake deployment and runtime issues (services, RDMA, env vars, connectivity, logs) and propose fixes.",
+  "author": {
+    "name": "Mooncake Team"
+  },
+  "homepage": "https://kvcache-ai.github.io/Mooncake/skills/mooncake-troubleshoot.html",
+  "repository": "https://github.com/kvcache-ai/Mooncake",
+  "license": "Apache-2.0",
+  "keywords": ["mooncake", "troubleshooting", "rdma", "deployment"]
+}
diff --git a/.devcontainer/Dockerfile b/.devcontainer/Dockerfile
index 7835f4c0..a5729ef9 100644
--- a/.devcontainer/Dockerfile
+++ b/.devcontainer/Dockerfile
@@ -12,7 +12,7 @@ RUN apt-get install -y libibverbs-dev \
     libjsoncpp-dev \
     libnuma-dev \
     libpython3-dev \
-    libboost-all-dev \
+    libboost-dev \
     libssl-dev \
     libgrpc-dev \
     libgrpc++-dev \
@@ -26,6 +26,7 @@ RUN apt-get install -y libibverbs-dev \
     libzstd-dev \
     libmsgpack-dev \
     libgflags-dev \
+    libxxhash-dev \
     pkg-config \
     patchelf
 
diff --git a/.devcontainer/devcontainer.json b/.devcontainer/devcontainer.json
index 7df63069..dea0cbb2 100644
--- a/.devcontainer/devcontainer.json
+++ b/.devcontainer/devcontainer.json
@@ -1,6 +1,8 @@
 {
   "name": "Mooncake Dev",
-  "dockerFile": "Dockerfile",
+  "build": {
+    "dockerfile": "Dockerfile"
+  },
   "runArgs": [
     "--cap-add=SYS_PTRACE",
     "--cap-add=NET_RAW",
@@ -18,7 +20,7 @@
     "vscode": {
       "settings": {
         "terminal.integrated.shell.linux": "/bin/bash",
-        "bazel.buildifierFixOnFormat": true,
+        "bazel.buildifierFixOnFormat": true
       }
     }
   }
diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS
index 6641f3f1..eabf4b29 100644
--- a/.github/CODEOWNERS
+++ b/.github/CODEOWNERS
@@ -15,6 +15,7 @@
 /mooncake-store @ykwd @stmatengss @XucSh @YiXR
 /mooncake-store/*/ha/ @Libotry @YiXR @00fish0
 /mooncake-transfer-engine @alogfans @doujiang24 @chestnut-Q
+/mooncake-transfer-engine/tent @alogfans @doujiang24 @chestnut-Q @staryxchen @00fish0 @dtcccc 
 /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
diff --git a/.github/labeler.yml b/.github/labeler.yml
index dad607da..aa76d0be 100644
--- a/.github/labeler.yml
+++ b/.github/labeler.yml
@@ -22,3 +22,43 @@ Mooncake EP:
 Installation:
   - changed-files:
     - any-glob-to-any-file: 'mooncake-wheel/**/*'
+
+P2P Store:
+  - changed-files:
+    - any-glob-to-any-file: 'mooncake-p2p-store/**/*'
+
+Integration:
+  - changed-files:
+    - any-glob-to-any-file: 'mooncake-integration/**/*'
+
+Common:
+  - changed-files:
+    - any-glob-to-any-file: 'mooncake-common/**/*'
+
+CI/Build:
+  - changed-files:
+    - any-glob-to-any-file:
+      - '.github/**/*'
+      - 'CMakeLists.txt'
+      - '.pre-commit-config.yaml'
+      - 'dependencies.sh'
+
+Documentation:
+  - changed-files:
+    - any-glob-to-any-file:
+      - 'docs/**/*'
+      - '*.md'
+
+Tests:
+  - changed-files:
+    - any-glob-to-any-file:
+      - 'scripts/test_*'
+      - 'mooncake-wheel/tests/**/*'
+      - 'scripts/tone_tests/**/*'
+
+Ascend/NPU:
+  - changed-files:
+    - any-glob-to-any-file:
+      - '**/ascend*/**'
+      - 'scripts/ascend/**'
+      - '.github/workflows/*ascend*'
diff --git a/.github/pull_request_template.md b/.github/pull_request_template.md
index 541ebfd2..a926d407 100644
--- a/.github/pull_request_template.md
+++ b/.github/pull_request_template.md
@@ -7,10 +7,11 @@
 - [ ] Transfer Engine (`mooncake-transfer-engine`)
 - [ ] Mooncake Store (`mooncake-store`)
 - [ ] Mooncake EP (`mooncake-ep`)
+- [ ] Mooncake PG (`mooncake-pg`)
 - [ ] Integration (`mooncake-integration`)
 - [ ] P2P Store (`mooncake-p2p-store`)
 - [ ] Python Wheel (`mooncake-wheel`)
-- [ ] PyTorch Backend (`mooncake-pg`)
+- [ ] Common (`mooncake-common`)
 - [ ] Mooncake RL (`mooncake-rl`)
 - [ ] CI/CD
 - [ ] Docs
@@ -23,15 +24,38 @@
 - [ ] Refactor
 - [ ] Breaking change
 - [ ] Documentation update
+- [ ] Performance improvement
 - [ ] Other
 
 ## How Has This Been Tested?
 
-<!-- Please describe the tests you've run to verify your changes. -->
+<!-- Describe the tests you've run. Include commands and output if applicable. -->
+
+**Test commands:**
+```bash
+# Example: bash scripts/run_ci_test.sh
+```
+
+**Test results:**
+- [ ] Unit tests pass
+- [ ] Integration tests pass (if applicable)
+- [ ] Manual testing done (describe below)
 
 ## Checklist
 
-- [ ] I have performed a self-review of my own code.
-- [ ] I have formatted my own code using `./scripts/code_format.sh` before submitting.
-- [ ] I have updated the documentation.
-- [ ] I have added tests to prove my changes are effective.
+- [ ] I have performed a self-review of my own code
+- [ ] I have formatted my code using `./scripts/code_format.sh`
+- [ ] I have run `pre-commit run --all-files` and all hooks pass
+- [ ] I have updated the documentation (if applicable)
+- [ ] I have added tests to prove my changes are effective
+- [ ] For changes >500 LOC: I have filed an RFC issue
+
+## AI Assistance Disclosure
+
+<!-- If AI tools were used, please disclose. This helps reviewers focus attention. -->
+
+- [ ] No AI tools were used
+- [ ] AI tools were used (specify below)
+
+<!-- If AI tools were used, briefly describe which tool and what it helped with.
+     The human submitter is responsible for understanding and defending all changes. -->
diff --git a/.github/workflows/assistant.yml b/.github/workflows/assistant.yml
new file mode 100644
index 00000000..59e76321
--- /dev/null
+++ b/.github/workflows/assistant.yml
@@ -0,0 +1,60 @@
+name: Qoder Assistant
+
+on:
+  issue_comment:
+    types: [created]
+  pull_request_review_comment:
+    types: [created]
+
+jobs:
+  qoder-assistant:
+    if: |
+      contains(github.event.comment.body, '@qoder') && 
+      !endsWith(github.event.comment.user.login, '[bot]')
+    runs-on: ubuntu-latest
+    permissions:
+      contents: read
+      issues: write
+      pull-requests: write
+      id-token: write
+
+    steps:
+      - name: Checkout repository
+        uses: actions/checkout@v4
+        with:
+          fetch-depth: 0
+
+      - name: Build Arguments
+        id: build_args
+        run: |
+          ARGS="REPO: ${{ github.repository }}
+          REQUEST_SOURCE: ${{ github.event_name }}
+          THREAD_ID: ${{ github.event.comment.node_id }}
+          COMMENT_ID: ${{ github.event.comment.id }}
+          AUTHOR: ${{ github.event.comment.user.login }}
+          BODY: ${{ github.event.comment.body }}
+          URL: ${{ github.event.comment.html_url }}
+          IS_PR: ${{ github.event.issue.pull_request != null || github.event_name == 'pull_request_review_comment' }}
+          ISSUE_OR_PR_NUMBER: ${{ github.event.issue.number || github.event.pull_request.number }}"
+          
+          if [ -n "${{ github.event.comment.pull_request_review_id }}" ]; then
+            ARGS="$ARGS
+          REVIEW_ID: ${{ github.event.comment.pull_request_review_id }}"
+          fi
+          
+          if [ -n "${{ github.event.comment.in_reply_to_id }}" ]; then
+            ARGS="$ARGS
+          REPLY_TO_COMMENT_ID: ${{ github.event.comment.in_reply_to_id }}"
+          fi
+          
+          echo "args<<EOF" >> $GITHUB_OUTPUT
+          echo "$ARGS" >> $GITHUB_OUTPUT
+          echo "EOF" >> $GITHUB_OUTPUT
+
+      - name: Run Qoder Assistant
+        uses: QoderAI/qoder-action@v0
+        with:
+          qoder_personal_access_token: ${{ secrets.QODER_PERSONAL_ACCESS_TOKEN }}
+          prompt: |
+            /assistant
+            ${{ steps.build_args.outputs.args }}
diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
index e72c7e46..5b6eb46e 100644
--- a/.github/workflows/ci.yml
+++ b/.github/workflows/ci.yml
@@ -43,6 +43,9 @@ jobs:
       with:
         python-version: ${{ matrix.python-version }}
 
+    - name: Install Rust toolchain
+      uses: dtolnay/rust-toolchain@stable
+
     - name: Install and start etcd
       run: |
         wget https://github.com/etcd-io/etcd/releases/download/v3.6.1/etcd-v3.6.1-linux-amd64.tar.gz
@@ -72,6 +75,11 @@ jobs:
         sudo apt-get update
         sudo apt-get install -y lcov gcovr ninja-build
 
+    - name: Test HugeTLB sizing helper
+      run: |
+        python3 scripts/test_hicache_hugepage_requirements.py
+      shell: bash
+
     - name: Set up coverage compilation flags
       run: |
         echo "Setting up coverage compilation flags..."
@@ -90,17 +98,13 @@ jobs:
           core.exportVariable('ACTIONS_RESULTS_URL', process.env.ACTIONS_RESULTS_URL || '');
           core.exportVariable('ACTIONS_RUNTIME_TOKEN', process.env.ACTIONS_RUNTIME_TOKEN || '');
 
-    - name: Run sccache stat for check
-      shell: bash
-      run: ${SCCACHE_PATH} --show-stats
-
     - name: Configure project with coverage support
       run: |
         sudo apt update -y
         sudo bash -x dependencies.sh -y
         mkdir build
         cd build
-        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
+        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 -DENABLE_DEBUG_SYMBOLS=OFF
       shell: bash
 
     - name: Build project
@@ -118,6 +122,11 @@ jobs:
         bash build.sh ../../build/mooncake-transfer-engine/nvlink-allocator/
       shell: bash
 
+    - name: Run sccache stat for check
+      if: ${{ env.SCCACHE_PATH != '' }}
+      shell: bash
+      run: ${SCCACHE_PATH} --show-stats
+
     - name: Start Metadata Server
       run: |
         cd mooncake-transfer-engine/example/http-metadata-server-python
@@ -125,6 +134,39 @@ jobs:
         python ./bootstrap_server.py &
       shell: bash
 
+    - name: Run Mooncake Store Rust smoke test and benchmark
+      run: |
+        $GITHUB_WORKSPACE/build/mooncake-store/src/mooncake_master \
+          --eviction_high_watermark_ratio=0.95 \
+          --cluster_id=ci_rust_test_cluster \
+          --port 50051 &
+        MASTER_PID=$!
+        sleep 3
+        cd mooncake-store/rust
+        export LD_LIBRARY_PATH=$GITHUB_WORKSPACE/build/mooncake-asio:$GITHUB_WORKSPACE/build/mooncake-store/src:$GITHUB_WORKSPACE/build/mooncake-store/src/cachelib_memory_allocator:$GITHUB_WORKSPACE/build/mooncake-transfer-engine/src:$GITHUB_WORKSPACE/build/mooncake-transfer-engine/src/common/base:$GITHUB_WORKSPACE/build/mooncake-common/etcd:$LD_LIBRARY_PATH
+        export MOONCAKE_BUILD_DIR=$GITHUB_WORKSPACE/build
+        export MOONCAKE_STORE_LIB_DIR=$GITHUB_WORKSPACE/build/mooncake-store/src
+        export MOONCAKE_STORE_INCLUDE_DIR=$GITHUB_WORKSPACE/mooncake-store/include
+        # This job builds Mooncake with -DENABLE_ASAN=ON, so the C++ libraries
+        # the Rust package links against carry undefined __asan_* references. Opt
+        # in to linking the ASan runtime; build.rs emits -lasan first, which
+        # keeps libasan first in the initial library list as ASan requires.
+        # Non-sanitized builds leave this unset and link without ASan.
+        export MOONCAKE_LINK_ASAN=1
+        export MC_METADATA_SERVER=http://127.0.0.1:8080/metadata
+        export MC_RUST_STORE_RUN_INTEGRATION=true
+        export MC_RUST_STORE_MASTER_ADDR=127.0.0.1:50051
+        export MC_RUST_STORE_LOCAL_HOSTNAME=127.0.0.1
+        export MC_RUST_STORE_PROTOCOL=tcp
+        export MC_RUST_STORE_DEVICE_NAME=
+        cargo test --test minimal_smoke -- --nocapture
+        MC_RUST_BENCH_ITERATIONS=4 \
+        MC_RUST_BENCH_VALUE_SIZE=4096 \
+        MC_RUST_BENCH_WARMUP=1 \
+          cargo run --release --example store_benchmark
+        kill $MASTER_PID 2>/dev/null || true
+      shell: bash
+
     - name: Run Go store binding integration tests
       run: |
         $GITHUB_WORKSPACE/build/mooncake-store/src/mooncake_master \
@@ -137,7 +179,7 @@ jobs:
         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"
+        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 -lmlx5 -ljsoncpp -lzstd -lcurl -luring -lasan -lm -lgcov -lxxhash -lyaml-cpp"
         # 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/...
@@ -262,7 +304,7 @@ jobs:
         bash -x dependencies.sh -y
         mkdir build
         cd build
-        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
+        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 -DENABLE_DEBUG_SYMBOLS=OFF
       shell: bash
 
     - name: Build project
@@ -273,6 +315,106 @@ jobs:
         cmake --install .
       shell: bash
 
+  build-arm64:
+    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'))
+    runs-on: ubuntu-22.04-arm
+    strategy:
+      matrix:
+        python-version: ['3.10', '3.12']
+    env:
+      CI: "true"
+      CU13_BUILD: "1"
+      TORCH_CUDA_ARCH_LIST: "9.0"
+      SCCACHE_GHA_ENABLED: "true"
+      CUDA_HOME: "/usr/local/cuda"
+
+    steps:
+    - uses: actions/checkout@v4
+      with:
+        persist-credentials: false
+
+    - name: Set up Python ${{ matrix.python-version }}
+      uses: actions/setup-python@v5
+      with:
+        python-version: ${{ matrix.python-version }}
+
+    - name: Free up disk space
+      run: |
+        sudo rm -rf /usr/share/dotnet
+        sudo rm -rf /opt/ghc
+        sudo rm -rf /opt/hostedtoolcache/CodeQL
+
+    - name: Install CUDA Toolkit 13.0 (arm64 SBSA)
+      run: |
+        wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/sbsa/cuda-keyring_1.1-1_all.deb
+        sudo dpkg -i cuda-keyring_1.1-1_all.deb
+        sudo apt-get update
+        sudo apt-get install -y cuda-toolkit-13-0
+        echo "/usr/local/cuda/bin" >> $GITHUB_PATH
+        /usr/local/cuda/bin/nvcc --version
+      shell: bash
+
+    - name: Run sccache-cache
+      uses: mozilla-actions/sccache-action@v0.0.9
+
+    - name: Configure sccache
+      uses: actions/github-script@v7
+      with:
+        script: |
+          core.exportVariable('ACTIONS_RESULTS_URL', process.env.ACTIONS_RESULTS_URL || '');
+          core.exportVariable('ACTIONS_RUNTIME_TOKEN', process.env.ACTIONS_RUNTIME_TOKEN || '');
+
+    - name: Configure project
+      run: |
+        sudo apt update -y
+        sudo bash -x dependencies.sh -y
+        mkdir build
+        cd build
+        cmake -G Ninja .. -DUSE_HTTP=ON -DUSE_CUDA=ON -DUSE_MNNVL=ON -DWITH_EP=OFF -DWITH_STORE_RUST=OFF -DENABLE_SCCACHE=ON -DCMAKE_BUILD_TYPE=Release
+      shell: bash
+
+    - name: Build project
+      run: |
+        export CUDA_HOME=/usr/local/cuda
+        for dir in /usr/local/cuda/lib64/stubs /usr/local/cuda/targets/*/lib/stubs; do
+          if [ -d "$dir" ]; then
+            export LIBRARY_PATH="$dir:${LIBRARY_PATH:-}"
+          fi
+        done
+        cd build
+        cmake --build .
+        sudo cmake --install .
+      shell: bash
+
+    - name: Run sccache stat for check
+      if: ${{ env.SCCACHE_PATH != '' }}
+      shell: bash
+      run: ${SCCACHE_PATH} --show-stats
+
+    - name: Generate Python version tag
+      id: generate_tag_arm64
+      run: |
+        echo "python_version_tag=$(echo ${{ matrix.python-version }} | tr -d '.')" >> $GITHUB_OUTPUT
+      shell: bash
+
+    - name: Build Python wheel
+      run: |
+        PYTHON_VERSION=${{ matrix.python-version }} OUTPUT_DIR=dist-py${{ steps.generate_tag_arm64.outputs.python_version_tag }} ./scripts/build_wheel.sh
+      shell: bash
+
+    - name: Upload wheel artifact
+      uses: actions/upload-artifact@v4
+      with:
+        name: wheel-arm64-cuda13-py${{ steps.generate_tag_arm64.outputs.python_version_tag }}
+        path: mooncake-wheel/dist-py${{ steps.generate_tag_arm64.outputs.python_version_tag }}/*.whl
+
   test-wheel-ubuntu:
     needs: [spell-check, clang-format, build-flags]
     if: >-
@@ -430,7 +572,7 @@ jobs:
         MC_FORCE_TCP: "true"
       run: |
         source test_env/bin/activate
-        python -m unittest mooncake-wheel.tests.test_mooncake_backend_cpu
+        python mooncake-pg/tests/test_pg_collectives.py
       shell: bash
 
     - name: Test Safetensor Functions
@@ -458,6 +600,7 @@ jobs:
       BUILD_WITH_EP: "1"
       TORCH_CUDA_ARCH_LIST: "8.0;9.0"
       SCCACHE_GHA_ENABLED: "true"
+      PIP_NO_CACHE_DIR: "1"
 
     steps:
     - uses: actions/checkout@v4
@@ -496,10 +639,6 @@ jobs:
           core.exportVariable('ACTIONS_RESULTS_URL', process.env.ACTIONS_RESULTS_URL || '');
           core.exportVariable('ACTIONS_RUNTIME_TOKEN', process.env.ACTIONS_RUNTIME_TOKEN || '');
 
-    - name: Run sccache stat for check
-      shell: bash
-      run: ${SCCACHE_PATH} --show-stats
-
     - name: Install dependencies
       run: |
         sudo apt update -y
@@ -518,7 +657,7 @@ jobs:
         cd build
         export PATH=/usr/local/nvidia/bin:/usr/local/nvidia/lib64:$PATH
         export LD_LIBRARY_PATH=/usr/local/cuda/lib64/stubs:$LD_LIBRARY_PATH
-        cmake -G Ninja .. -DUSE_ETCD=OFF -DUSE_CXL=ON -DUSE_REDIS=ON -DUSE_HTTP=ON -DWITH_METRICS=ON -DBUILD_UNIT_TESTS=ON -DBUILD_EXAMPLES=ON -DENABLE_SCCACHE=ON -DUSE_CUDA=OFF -DUSE_MNNVL=OFF -DUSE_UB=OFF -DCMAKE_EXE_LINKER_FLAGS="-L/usr/local/cuda/lib64/stubs"
+        cmake -G Ninja .. -DUSE_ETCD=OFF -DUSE_CXL=ON -DUSE_REDIS=ON -DUSE_HTTP=ON -DWITH_METRICS=ON -DBUILD_UNIT_TESTS=ON -DBUILD_EXAMPLES=ON -DENABLE_SCCACHE=ON -DUSE_CUDA=OFF -DUSE_MNNVL=OFF -DUSE_UB=OFF -DCMAKE_EXE_LINKER_FLAGS="-L/usr/local/cuda/lib64/stubs" -DENABLE_DEBUG_SYMBOLS=OFF
         cmake --build .
         sudo cmake --install .
         df -h
@@ -528,7 +667,7 @@ jobs:
       run: |
         mkdir build
         cd build
-        cmake -G Ninja .. -DUSE_ETCD=ON -DUSE_CXL=ON -DUSE_REDIS=ON -DUSE_HTTP=ON -DWITH_STORE=ON -DWITH_P2P_STORE=ON -DWITH_METRICS=ON -DBUILD_UNIT_TESTS=ON -DBUILD_EXAMPLES=ON -DENABLE_SCCACHE=ON -DUSE_CUDA=ON -DUSE_MNNVL=OFF -DUSE_UB=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" -DENABLE_DEBUG_SYMBOLS=OFF
       shell: bash
       # TODO: lack USE_NVMEOF,USE_MNNVL
 
@@ -545,7 +684,7 @@ jobs:
     - name: Configure project with unit tests and examples
       run: |
         cd build
-        cmake -G Ninja .. -DBUILD_UNIT_TESTS=ON -DBUILD_EXAMPLES=ON -DWITH_STORE_RUST=ON -DENABLE_SCCACHE=ON
+        cmake -G Ninja .. -DBUILD_UNIT_TESTS=ON -DBUILD_EXAMPLES=ON -DWITH_STORE_RUST=ON -DENABLE_SCCACHE=ON -DENABLE_DEBUG_SYMBOLS=OFF
       shell: bash
 
     - name: Build project with unit tests and examples
@@ -557,19 +696,28 @@ jobs:
         sudo cmake --install .
       shell: bash
 
-    - name: Check Mooncake Store Rust bindings and example
+    - name: Check Mooncake Store Rust bindings, examples, and tests
       run: |
+        # libcuda.so.1 (SONAME of the CUDA stub) must be findable at runtime.
+        # The toolkit stubs dir only ships libcuda.so; create the versioned symlink.
+        if [ -f /usr/local/cuda/lib64/stubs/libcuda.so ] && \
+           [ ! -e /usr/local/cuda/lib64/stubs/libcuda.so.1 ]; then
+          sudo ln -s libcuda.so /usr/local/cuda/lib64/stubs/libcuda.so.1
+        fi
+        export LD_LIBRARY_PATH=/usr/local/cuda/lib64/stubs:$LD_LIBRARY_PATH
         cd mooncake-store/rust
+        export MOONCAKE_BUILD_DIR=$GITHUB_WORKSPACE/build
+        cargo test --lib
         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
+          cargo test --examples --tests --no-run
       shell: bash
 
     - name: Configure project
       run: |
         cd build
         rm -r */tests
-        cmake -G Ninja .. -DBUILD_UNIT_TESTS=OFF -DBUILD_EXAMPLES=OFF -DUSE_HTTP=ON -DENABLE_SCCACHE=ON -DUSE_CXL=ON -DWITH_EP=ON -DEP_TORCH_VERSIONS="2.9.0;2.9.1;2.10.0;2.11.0"
+        cmake -G Ninja .. -DBUILD_UNIT_TESTS=OFF -DBUILD_EXAMPLES=OFF -DUSE_HTTP=ON -DENABLE_SCCACHE=ON -DUSE_CXL=ON -DWITH_EP=ON -DEP_TORCH_VERSIONS="2.9.1;2.10.0;2.11.0;2.12.0;2.12.1" -DENABLE_DEBUG_SYMBOLS=OFF
       shell: bash
 
     - name: Build project
@@ -581,6 +729,26 @@ jobs:
         sudo cmake --install .
       shell: bash
 
+    - name: Configure project with TENT
+      run: |
+        mkdir build-tent
+        cd build-tent
+        cmake -G Ninja .. -DUSE_TENT=ON -DUSE_HTTP=ON -DENABLE_SCCACHE=ON -DBUILD_UNIT_TESTS=ON -DBUILD_EXAMPLES=ON -DENABLE_DEBUG_SYMBOLS=OFF
+      shell: bash
+
+    - name: Build project with TENT
+      run: |
+        cd build-tent
+        cmake --build .
+        sudo cmake --install .
+      shell: bash
+
+    - name: Test (TENT)
+      run: |
+        cd build-tent
+        ctest --test-dir mooncake-transfer-engine/tent/tests -j --output-on-failure
+      shell: bash
+
     - name: Build nvlink_allocator.so
       run: |
         mkdir -p build/mooncake-transfer-engine/nvlink-allocator
@@ -591,6 +759,11 @@ jobs:
         bash build.sh ../../build/mooncake-transfer-engine/nvlink-allocator/
       shell: bash
 
+    - name: Run sccache stat for check
+      if: ${{ env.SCCACHE_PATH != '' }}
+      shell: bash
+      run: ${SCCACHE_PATH} --show-stats
+
     - name: Generate Python version tag
       id: generate_tag_flags
       run: |
@@ -638,10 +811,10 @@ jobs:
   spell-check:
     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')
+      (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
@@ -654,10 +827,10 @@ jobs:
   clang-format:
     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')
+      (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
@@ -680,7 +853,7 @@ jobs:
           echo "Error: code_format.sh not found or not executable"
           exit 1
         fi
-        
+
         # Determine base ref for comparison
         if [ "${{ github.event_name }}" == "pull_request" ]; then
           # For PRs: compare against the target branch
@@ -701,6 +874,48 @@ jobs:
       shell: bash
 
 
+  docs-check:
+    name: Check Sphinx docs build
+    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:
+    - uses: actions/checkout@v4
+      with:
+        fetch-depth: 2
+        persist-credentials: false
+
+    - uses: dorny/paths-filter@v3
+      id: filter
+      with:
+        filters: |
+          docs:
+            - 'docs/**'
+            - 'requirements_docs.txt'
+
+    - name: Set up Python
+      if: steps.filter.outputs.docs == 'true'
+      uses: actions/setup-python@v5
+      with:
+        python-version: '3.12'
+
+    - name: Install documentation dependencies
+      if: steps.filter.outputs.docs == 'true'
+      run: |
+        python -m pip install --upgrade pip
+        pip install -r requirements_docs.txt
+
+    - name: Build docs with strict mode
+      if: steps.filter.outputs.docs == 'true'
+      run: |
+        cd docs
+        make html SPHINXOPTS=-W
+      shell: bash
+
+
   check-paths:
     if: >-
       github.event_name == 'push' ||
@@ -746,6 +961,18 @@ jobs:
     uses: ./.github/workflows/ci_cu13.yml
     secrets: inherit
 
+  build-wheel-efa:
+    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_efa.yml
+    secrets: inherit
+
   ascend-test:
     needs: [build, check-paths]
     if: needs.check-paths.outputs.should-run-downstream == 'true'
@@ -764,12 +991,15 @@ jobs:
     needs:
       - spell-check
       - clang-format
+      - docs-check
       - build
       - build-musa
+      - build-arm64
       - build-flags
       - build-docker
       - test-wheel-ubuntu
       - build-wheel-cu13
+      - build-wheel-efa
       - ascend-test
       - integration-test
     runs-on: ubuntu-latest
diff --git a/.github/workflows/ci_ascend.yml b/.github/workflows/ci_ascend.yml
index c3f39c98..13f79a24 100644
--- a/.github/workflows/ci_ascend.yml
+++ b/.github/workflows/ci_ascend.yml
@@ -4,7 +4,7 @@ on:
   workflow_call:
     inputs:
       checkout_ref:
-        description: 'Git ref to checkout (PR head SHA for pull_request_target)'
+        description: 'Git ref to checkout (PR head SHA)'
         required: false
         type: string
 
@@ -49,7 +49,7 @@ jobs:
       if: steps.checkout_code.outcome == 'failure'
       shell: bash
       env:
-        ASCEND_GITHUB_MIRROR_URLS: ${{ vars.ASCEND_GITHUB_MIRROR_URLS }}
+        ASCEND_GITHUB_MIRROR_URLS: 'https://ghfast.top/'
         CHECKOUT_REF: ${{ inputs.checkout_ref || github.sha }}
       run: |
         set -euo pipefail
@@ -108,7 +108,7 @@ jobs:
     - name: Configure CMake
       shell: bash
       env:
-        ASCEND_GITHUB_MIRROR_URLS: ${{ vars.ASCEND_GITHUB_MIRROR_URLS }}
+        ASCEND_GITHUB_MIRROR_URLS: 'https://ghfast.top/'
       run: |
         source /usr/local/Ascend/cann-9.0.0/set_env.sh
         pwd
@@ -163,7 +163,9 @@ jobs:
         cmake .. \
           -DUSE_ASCEND_DIRECT=ON \
           -DBUILD_EXAMPLES=OFF \
-          -DBUILD_UNIT_TESTS=OFF
+          -DBUILD_UNIT_TESTS=OFF \
+          -DENABLE_DEBUG_SYMBOLS=OFF \
+          -DBUILD_BENCHMARK=OFF
 
     - name: Build
       shell: bash
@@ -171,7 +173,7 @@ jobs:
         source /usr/local/Ascend/cann-9.0.0/set_env.sh
         echo "Building..."
         cd build
-        cmake --build . -j$(nproc)
+        cmake --build . -j48
         cmake --install .
         echo "Mooncake installed successfully."
 
diff --git a/.github/workflows/ci_cu13.yml b/.github/workflows/ci_cu13.yml
index 0fae01f7..828491f3 100644
--- a/.github/workflows/ci_cu13.yml
+++ b/.github/workflows/ci_cu13.yml
@@ -52,10 +52,6 @@ jobs:
           core.exportVariable('ACTIONS_RESULTS_URL', process.env.ACTIONS_RESULTS_URL || '');
           core.exportVariable('ACTIONS_RUNTIME_TOKEN', process.env.ACTIONS_RUNTIME_TOKEN || '');
 
-    - name: Run sccache stat for check
-      shell: bash
-      run: ${SCCACHE_PATH} --show-stats
-
     - name: Install dependencies
       run: |
         sudo apt update -y
@@ -75,7 +71,7 @@ jobs:
           -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" \
+          -DEP_TORCH_VERSIONS="2.9.1;2.10.0;2.11.0;2.12.0;2.12.1" \
           -DWITH_METRICS=ON \
           -DBUILD_UNIT_TESTS=OFF \
           -DBUILD_EXAMPLES=ON \
@@ -83,7 +79,8 @@ jobs:
           -DBUILD_BENCHMARK=ON \
           -DUSE_CUDA=ON \
           -DCMAKE_BUILD_TYPE=Release \
-          -DCMAKE_EXE_LINKER_FLAGS="-L/usr/local/cuda/lib64/stubs"
+          -DCMAKE_EXE_LINKER_FLAGS="-L/usr/local/cuda/lib64/stubs" \
+          -DENABLE_DEBUG_SYMBOLS=OFF
       shell: bash
 
     - name: Build project
@@ -106,6 +103,11 @@ jobs:
         bash build.sh ../../build/mooncake-transfer-engine/nvlink-allocator/
       shell: bash
 
+    - name: Run sccache stat for check
+      if: ${{ env.SCCACHE_PATH != '' }}
+      shell: bash
+      run: ${SCCACHE_PATH} --show-stats
+
     - name: Generate Python version tag
       id: generate_tag
       run: |
diff --git a/.github/workflows/ci_efa.yml b/.github/workflows/ci_efa.yml
new file mode 100644
index 00000000..64daa8c8
--- /dev/null
+++ b/.github/workflows/ci_efa.yml
@@ -0,0 +1,154 @@
+name: 'Build Wheel (AWS EFA)'
+
+on:
+  workflow_call: {}
+
+# Builds the AWS EFA (libfabric) wheel variants on a stock ubuntu runner.
+# No EFA hardware is required to *build*: USE_EFA only needs the libfabric
+# headers/lib to compile and link. auditwheel later excludes libfabric/libefa
+# from the wheel so they resolve to the user's system EFA install
+# (/opt/amazon/efa/lib) at runtime. The distro libfabric (1.x) is ABI-forward-
+# compatible with the AWS EFA libfabric (2.x) that loads at runtime; the EFA
+# transport only uses long-stable fi_* core APIs.
+#
+# Two variants, since the EFA transport's memory path is CUDA-aware
+# (FI_HMEM_CUDA / GPUDirect under USE_CUDA=ON, FI_HMEM=system otherwise):
+#   efa            USE_CUDA=ON   (GPU)
+#   efa-non-cuda   USE_CUDA=OFF  (CPU/DRAM)
+# PR validation builds one python version per variant to keep CI cheap;
+# the release workflow builds the full python matrix.
+jobs:
+  build-wheel-efa:
+    runs-on: ubuntu-22.04
+    strategy:
+      matrix:
+        include:
+          - variant: cuda
+            use_cuda: "ON"
+            build_env: "EFA_BUILD"
+            python-version: "3.12"
+          - variant: non-cuda
+            use_cuda: "OFF"
+            build_env: "EFA_NON_CUDA_BUILD"
+            python-version: "3.10"
+    env:
+      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
+      with:
+        python-version: ${{ matrix.python-version }}
+
+    - name: Free up disk space
+      run: |
+        sudo rm -rf /usr/share/dotnet
+        sudo rm -rf /opt/ghc
+        sudo rm -rf /opt/hostedtoolcache/CodeQL
+        sudo rm -rf /usr/local/lib/android
+        df -h
+
+    - name: Install CUDA Toolkit
+      if: matrix.use_cuda == 'ON'
+      uses: Jimver/cuda-toolkit@v0.2.24
+      with:
+        cuda: '12.8.1'
+        method: 'network'
+        sub-packages: '["nvcc", "nvrtc-dev"]'
+        non-cuda-sub-packages: '["libcusparse-dev", "libcublas-dev", "libcusolver-dev"]'
+
+    - name: Run sccache-cache
+      uses: mozilla-actions/sccache-action@v0.0.9
+
+    - name: Configure sccache
+      uses: actions/github-script@v7
+      with:
+        script: |
+          core.exportVariable('ACTIONS_RESULTS_URL', process.env.ACTIONS_RESULTS_URL || '');
+          core.exportVariable('ACTIONS_RUNTIME_TOKEN', process.env.ACTIONS_RUNTIME_TOKEN || '');
+
+    - name: Install dependencies
+      run: |
+        sudo apt update -y
+        sudo apt install -y ninja-build libfabric-dev libfabric1
+        sudo bash -x dependencies.sh -y
+        df -h
+      shell: bash
+
+    - name: Configure project
+      run: |
+        mkdir build
+        cd build
+        EXTRA_FLAGS=""
+        if [ "${{ matrix.use_cuda }}" = "ON" ]; then
+          EXTRA_FLAGS="-DCMAKE_EXE_LINKER_FLAGS=-L/usr/local/cuda/lib64/stubs"
+        fi
+        cmake -G Ninja .. \
+          -DUSE_ETCD=ON \
+          -DUSE_HTTP=ON \
+          -DWITH_STORE=ON \
+          -DWITH_METRICS=ON \
+          -DBUILD_UNIT_TESTS=OFF \
+          -DBUILD_EXAMPLES=ON \
+          -DENABLE_SCCACHE=ON \
+          -DBUILD_BENCHMARK=ON \
+          -DUSE_EFA=ON \
+          -DUSE_CUDA=${{ matrix.use_cuda }} \
+          -DLIBFABRIC_INCLUDE_DIR=/usr/include \
+          -DLIBFABRIC_LIBRARY=/usr/lib/x86_64-linux-gnu/libfabric.so \
+          -DCMAKE_BUILD_TYPE=Release \
+          -DENABLE_DEBUG_SYMBOLS=OFF \
+          $EXTRA_FLAGS
+      shell: bash
+
+    - name: Build project
+      run: |
+        if [ "${{ matrix.use_cuda }}" = "ON" ]; then
+          export LD_LIBRARY_PATH=/usr/local/cuda/lib64/stubs:$LD_LIBRARY_PATH
+          export LIBRARY_PATH=/usr/local/cuda/lib64/stubs:$LIBRARY_PATH
+        fi
+        cd build
+        cmake --build .
+        sudo cmake --install .
+        df -h
+      shell: bash
+
+    - name: Run sccache stat for check
+      if: ${{ env.SCCACHE_PATH != '' }}
+      shell: bash
+      run: ${SCCACHE_PATH} --show-stats
+
+    - name: Generate Python version tag
+      id: generate_tag
+      run: |
+        echo "python_version_tag=$(echo ${{ matrix.python-version }} | tr -d '.')" >> $GITHUB_OUTPUT
+      shell: bash
+
+    - name: Build Python wheel
+      run: |
+        export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/lib
+        export ${{ matrix.build_env }}=1
+        PYTHON_VERSION=${{ matrix.python-version }} OUTPUT_DIR=dist-py${{ steps.generate_tag.outputs.python_version_tag }} ./scripts/build_wheel.sh
+      shell: bash
+
+    - name: Verify libfabric is excluded from the wheel
+      run: |
+        WHL=$(ls mooncake-wheel/dist-py${{ steps.generate_tag.outputs.python_version_tag }}/*.whl | head -1)
+        echo "Inspecting $WHL"
+        if unzip -l "$WHL" | grep -iE 'libfabric|libefa'; then
+          echo "::error::libfabric/libefa must NOT be bundled in the EFA wheel"
+          exit 1
+        fi
+        echo "OK: libfabric/libefa correctly excluded (resolve to system EFA at runtime)"
+      shell: bash
+
+    - name: Upload Python wheel artifact
+      uses: actions/upload-artifact@v4
+      with:
+        name: mooncake-wheel-efa-${{ matrix.variant }}-ubuntu-py${{ steps.generate_tag.outputs.python_version_tag }}
+        path: mooncake-wheel/dist-py${{ steps.generate_tag.outputs.python_version_tag }}/*.whl
diff --git a/.github/workflows/code-review.yml b/.github/workflows/code-review.yml
new file mode 100644
index 00000000..94622c08
--- /dev/null
+++ b/.github/workflows/code-review.yml
@@ -0,0 +1,29 @@
+name: Qoder Auto Code Review
+
+on:
+  pull_request:
+    types: [opened, synchronize, reopened]
+
+jobs:
+  qoder-review:
+    # Skip fork and cross-repo PRs (head repo must match this repository)
+    if: github.event.pull_request.head.repo.full_name == github.repository
+    runs-on: ubuntu-latest
+    permissions:
+      contents: read
+      pull-requests: write
+      id-token: write
+    
+    steps:
+      - name: Checkout repository
+        uses: actions/checkout@v4
+        with:
+          fetch-depth: 0
+      
+      - name: Run Qoder Code Review
+        uses: QoderAI/qoder-action@v0
+        with:
+          qoder_personal_access_token: ${{ secrets.QODER_PERSONAL_ACCESS_TOKEN }}
+          prompt: |
+            /review-pr
+            REPO:${{ github.repository }} PR_NUMBER:${{ github.event.pull_request.number }}
diff --git a/.github/workflows/issue-bot.yml b/.github/workflows/issue-bot.yml
new file mode 100644
index 00000000..c199c049
--- /dev/null
+++ b/.github/workflows/issue-bot.yml
@@ -0,0 +1,71 @@
+name: Issue Bot
+
+on:
+  issues:
+    types: [opened]
+  schedule:
+    # Daily at 03:00 UTC
+    - cron: "0 3 * * *"
+  workflow_dispatch:
+
+permissions:
+  issues: write
+
+jobs:
+  welcome-and-assign:
+    if: github.event_name == 'issues'
+    runs-on: ubuntu-latest
+    steps:
+      - name: Post welcome comment
+        uses: actions/github-script@v7
+        with:
+          script: |
+            const issue = context.payload.issue;
+            const body = [
+              `Thanks for opening this issue, @${issue.user.login}!`,
+              '',
+              '| Field | Value |',
+              '|-------|-------|',
+              `| **Issue** | #${issue.number} |`,
+              `| **GitHub user ID** | \`${issue.user.id}\` |`,
+              `| **Reporter** | @${issue.user.login} |`,
+              '',
+              'A maintainer will triage this when possible. To help us respond faster, please include:',
+              '',
+              '- Mooncake version or commit SHA',
+              '- Environment (OS, CUDA/driver, RDMA stack if relevant)',
+              '- Steps to reproduce and expected vs. actual behavior',
+              '',
+              `Useful links: [Documentation](https://kvcache-ai.github.io/Mooncake/) · [Contributing guide](https://github.com/${context.repo.owner}/${context.repo.repo}/blob/main/CONTRIBUTING.md)`,
+              '',
+              '> This message was posted automatically by the issue bot.',
+            ].join('\n');
+
+            await github.rest.issues.createComment({
+              owner: context.repo.owner,
+              repo: context.repo.repo,
+              issue_number: issue.number,
+              body,
+            });
+
+  close-stale-issues:
+    if: github.event_name != 'issues'
+    runs-on: ubuntu-latest
+    steps:
+      - name: Close inactive issues (3+ months)
+        uses: actions/stale@v9
+        with:
+          repo-token: ${{ secrets.GITHUB_TOKEN }}
+          days-before-stale: 90
+          days-before-close: 7
+          stale-issue-label: stale
+          close-issue-label: auto-closed
+          stale-issue-message: >
+            This issue has had no activity for 90 days and will be closed in 7 days
+            if there is no further activity. Please comment or react if it should
+            stay open.
+          close-issue-message: >
+            Closing due to 3 months of inactivity. If this is still relevant,
+            please comment and we can reopen.
+          exempt-issue-labels: pinned,keep-open,security
+          operations-per-run: 100
diff --git a/.github/workflows/pr-tidy.yml b/.github/workflows/pr-tidy.yml
new file mode 100644
index 00000000..82342963
--- /dev/null
+++ b/.github/workflows/pr-tidy.yml
@@ -0,0 +1,31 @@
+name: PR Tidy
+
+on:
+  pull_request_target:
+    types: [opened]
+
+permissions:
+  pull-requests: write
+
+jobs:
+  tidy:
+    runs-on: ubuntu-latest
+    steps:
+      - name: Clean PR description
+        uses: actions/github-script@v7
+        with:
+          script: |
+            const pr = context.payload.pull_request;
+            if (!pr.body) return;
+
+            // Remove HTML comment blocks (template instructions)
+            let body = pr.body.replace(/<!--[\s\S]*?-->/g, '').trim();
+
+            if (body !== pr.body) {
+              await github.rest.pulls.update({
+                owner: context.repo.owner,
+                repo: context.repo.repo,
+                pull_number: pr.number,
+                body: body
+              });
+            }
diff --git a/.github/workflows/pre-release.yaml b/.github/workflows/pre-release.yaml
new file mode 100644
index 00000000..4e74b164
--- /dev/null
+++ b/.github/workflows/pre-release.yaml
@@ -0,0 +1,549 @@
+name: Pre-Release
+
+# Dry-run of the release pipelines: build wheels like Release / Release Non-CUDA /
+# Release CUDA 13, validate artifacts, but do not create a GitHub Release or publish to PyPI.
+#
+# Trigger by pushing a pre-release tag, for example:
+#   git tag v1.0.0-rc1 && git push origin v1.0.0-rc1
+on:
+  push:
+    tags:
+      - 'v*-rc*'
+      - 'v*-alpha*'
+      - 'v*-beta*'
+      - 'v*-pre*'
+
+env:
+  SCCACHE_GHA_ENABLED: "true"
+
+jobs:
+  build-cuda:
+    name: Build (CUDA 12)
+    runs-on: ubuntu-22.04
+    permissions:
+      contents: read
+    strategy:
+      matrix:
+        python-version: ['3.10', '3.11', '3.12', '3.13']
+    env:
+      BUILD_WITH_EP: "1"
+      TORCH_CUDA_ARCH_LIST: "8.0;9.0"
+    steps:
+      - name: Checkout source
+        uses: actions/checkout@v4
+
+      - name: Set version from tag
+        run: echo "VERSION=${GITHUB_REF_NAME#v}" >> $GITHUB_ENV
+
+      - name: Set up Python ${{ matrix.python-version }}
+        uses: actions/setup-python@v5
+        with:
+          python-version: ${{ matrix.python-version }}
+
+      - name: Free up disk space
+        run: |
+          sudo rm -rf /usr/share/dotnet
+          sudo rm -rf /opt/ghc
+          sudo rm -rf /opt/hostedtoolcache/CodeQL
+          sudo rm -rf /usr/local/lib/android
+          df -h
+
+      - name: Install CUDA Toolkit
+        uses: Jimver/cuda-toolkit@v0.2.24
+        with:
+          cuda: '12.8.1'
+          linux-local-args: '["--toolkit"]'
+          method: 'network'
+          sub-packages: '["nvcc", "nvrtc-dev"]'
+          non-cuda-sub-packages: '["libcusparse-dev", "libcublas-dev", "libcusolver-dev"]'
+
+      - name: Run sccache-cache
+        uses: mozilla-actions/sccache-action@v0.0.9
+
+      - name: Configure sccache
+        uses: actions/github-script@v7
+        with:
+          script: |
+            core.exportVariable('ACTIONS_RESULTS_URL', process.env.ACTIONS_RESULTS_URL || '');
+            core.exportVariable('ACTIONS_RUNTIME_TOKEN', process.env.ACTIONS_RUNTIME_TOKEN || '');
+
+      - name: Configure project
+        run: |
+          sudo apt update -y
+          sudo bash -x dependencies.sh -y
+          mkdir build
+          cd build
+          cmake .. -DBUILD_UNIT_TESTS=OFF -DUSE_HTTP=ON -DUSE_ETCD=ON -DUSE_CUDA=ON -DWITH_EP=ON -DEP_TORCH_VERSIONS="2.9.1;2.10.0;2.11.0;2.12.0;2.12.1" -DSTORE_USE_ETCD=ON -DENABLE_SCCACHE=ON -DCMAKE_BUILD_TYPE=Release
+        shell: bash
+
+      - name: Build project
+        run: |
+          export LIBRARY_PATH=/usr/local/cuda/lib64/stubs:$LIBRARY_PATH
+          cd build
+          make -j
+          sudo -E make install
+        shell: bash
+
+      - name: Build nvlink_allocator.so
+        run: |
+          export PATH=/usr/local/nvidia/bin:/usr/local/nvidia/lib64:$PATH
+          export LD_LIBRARY_PATH=/usr/local/cuda/lib64/stubs:$LD_LIBRARY_PATH
+          export LIBRARY_PATH=/usr/local/cuda/lib64/stubs:$LIBRARY_PATH
+          mkdir -p build/mooncake-transfer-engine/nvlink-allocator
+          cd mooncake-transfer-engine/nvlink-allocator
+          bash build.sh ../../build/mooncake-transfer-engine/nvlink-allocator/
+        shell: bash
+
+      - name: Run sccache stat for check
+        if: ${{ env.SCCACHE_PATH != '' }}
+        shell: bash
+        run: ${SCCACHE_PATH} --show-stats
+
+      - name: Generate Python version tag
+        id: generate_tag_release
+        run: |
+          echo "python_version_tag=$(echo ${{ matrix.python-version }} | tr -d '.')" >> $GITHUB_OUTPUT
+        shell: bash
+
+      - name: Build Python wheel
+        run: |
+          export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/lib
+          PYTHON_VERSION=${{ matrix.python-version }} OUTPUT_DIR=dist-py${{ steps.generate_tag_release.outputs.python_version_tag }} ./scripts/build_wheel.sh
+        env:
+          VERSION: ${{ env.VERSION }}
+
+      - name: Upload Python wheel artifact
+        uses: actions/upload-artifact@v4
+        with:
+          name: pre-release-cuda-py${{ steps.generate_tag_release.outputs.python_version_tag }}
+          path: mooncake-wheel/dist-py${{ steps.generate_tag_release.outputs.python_version_tag }}/*.whl
+
+  build-non-cuda:
+    name: Build (Non-CUDA)
+    runs-on: ubuntu-22.04
+    permissions:
+      contents: read
+    strategy:
+      matrix:
+        python-version: ['3.10', '3.11', '3.12', '3.13']
+    env:
+      BUILD_WITH_EP: "0"
+      NON_CUDA_BUILD: "1"
+    steps:
+      - name: Checkout source
+        uses: actions/checkout@v4
+
+      - name: Set version from tag
+        run: echo "VERSION=${GITHUB_REF_NAME#v}" >> $GITHUB_ENV
+
+      - name: Set up Python ${{ matrix.python-version }}
+        uses: actions/setup-python@v5
+        with:
+          python-version: ${{ matrix.python-version }}
+
+      - name: Free up disk space
+        run: |
+          sudo rm -rf /usr/share/dotnet
+          sudo rm -rf /opt/ghc
+          sudo rm -rf /opt/hostedtoolcache/CodeQL
+
+      - name: Run sccache-cache
+        uses: mozilla-actions/sccache-action@v0.0.9
+
+      - name: Configure sccache
+        uses: actions/github-script@v7
+        with:
+          script: |
+            core.exportVariable('ACTIONS_RESULTS_URL', process.env.ACTIONS_RESULTS_URL || '');
+            core.exportVariable('ACTIONS_RUNTIME_TOKEN', process.env.ACTIONS_RUNTIME_TOKEN || '');
+
+      - name: Configure project
+        run: |
+          sudo apt update -y
+          sudo bash -x dependencies.sh -y
+          mkdir build
+          cd build
+          cmake .. -DUSE_HTTP=ON -DUSE_ETCD=ON -DUSE_CUDA=OFF -DWITH_EP=OFF -DSTORE_USE_ETCD=ON -DENABLE_SCCACHE=ON -DCMAKE_BUILD_TYPE=Release
+        shell: bash
+
+      - name: Build project
+        run: |
+          cd build
+          make -j
+          sudo make install
+        shell: bash
+
+      - name: Run sccache stat for check
+        if: ${{ env.SCCACHE_PATH != '' }}
+        shell: bash
+        run: ${SCCACHE_PATH} --show-stats
+
+      - name: Generate Python version tag
+        id: generate_tag_release
+        run: |
+          echo "python_version_tag=$(echo ${{ matrix.python-version }} | tr -d '.')" >> $GITHUB_OUTPUT
+        shell: bash
+
+      - name: Build Python wheel
+        run: |
+          export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/lib
+          PYTHON_VERSION=${{ matrix.python-version }} OUTPUT_DIR=dist-py${{ steps.generate_tag_release.outputs.python_version_tag }} ./scripts/build_wheel.sh
+        env:
+          VERSION: ${{ env.VERSION }}
+
+      - name: Upload Python wheel artifact
+        uses: actions/upload-artifact@v4
+        with:
+          name: pre-release-non-cuda-py${{ steps.generate_tag_release.outputs.python_version_tag }}
+          path: mooncake-wheel/dist-py${{ steps.generate_tag_release.outputs.python_version_tag }}/*.whl
+
+  build-cuda13:
+    name: Build (CUDA 13)
+    runs-on: ubuntu-22.04
+    permissions:
+      contents: read
+    strategy:
+      matrix:
+        python-version: ['3.10', '3.11', '3.12', '3.13']
+    env:
+      BUILD_WITH_EP: "1"
+      CU13_BUILD: "1"
+      TORCH_CUDA_ARCH_LIST: "8.0;9.0"
+    steps:
+      - name: Checkout source
+        uses: actions/checkout@v4
+
+      - name: Set version from tag
+        run: echo "VERSION=${GITHUB_REF_NAME#v}" >> $GITHUB_ENV
+
+      - name: Set up Python ${{ matrix.python-version }}
+        uses: actions/setup-python@v5
+        with:
+          python-version: ${{ matrix.python-version }}
+
+      - name: Free up disk space
+        run: |
+          sudo rm -rf /usr/share/dotnet
+          sudo rm -rf /opt/ghc
+          sudo rm -rf /opt/hostedtoolcache/CodeQL
+          sudo rm -rf /usr/local/lib/android
+          df -h
+
+      - name: Install CUDA Toolkit 13
+        uses: Jimver/cuda-toolkit@v0.2.29
+        with:
+          cuda: '13.0.2'
+          linux-local-args: '["--toolkit"]'
+          method: 'network'
+          sub-packages: '["nvcc", "nvrtc-dev"]'
+          non-cuda-sub-packages: '["libcusparse-dev", "libcublas-dev", "libcusolver-dev"]'
+
+      - name: Run sccache-cache
+        uses: mozilla-actions/sccache-action@v0.0.9
+
+      - name: Configure sccache
+        uses: actions/github-script@v7
+        with:
+          script: |
+            core.exportVariable('ACTIONS_RESULTS_URL', process.env.ACTIONS_RESULTS_URL || '');
+            core.exportVariable('ACTIONS_RUNTIME_TOKEN', process.env.ACTIONS_RUNTIME_TOKEN || '');
+
+      - name: Configure project
+        run: |
+          sudo apt update -y
+          sudo bash -x dependencies.sh -y
+          mkdir build
+          cd build
+          cmake .. -DBUILD_UNIT_TESTS=OFF -DUSE_HTTP=ON -DUSE_ETCD=ON -DUSE_CUDA=ON -DWITH_EP=ON -DEP_TORCH_VERSIONS="2.9.1;2.10.0;2.11.0;2.12.0;2.12.1" -DSTORE_USE_ETCD=ON -DENABLE_SCCACHE=ON -DCMAKE_BUILD_TYPE=Release
+        shell: bash
+
+      - name: Build project
+        run: |
+          export LIBRARY_PATH=/usr/local/cuda/lib64/stubs:$LIBRARY_PATH
+          cd build
+          make -j
+          sudo make install
+        shell: bash
+
+      - name: Build nvlink_allocator.so
+        run: |
+          export PATH=/usr/local/nvidia/bin:/usr/local/nvidia/lib64:$PATH
+          export LD_LIBRARY_PATH=/usr/local/cuda/lib64/stubs:$LD_LIBRARY_PATH
+          export LIBRARY_PATH=/usr/local/cuda/lib64/stubs:$LIBRARY_PATH
+          mkdir -p build/mooncake-transfer-engine/nvlink-allocator
+          cd mooncake-transfer-engine/nvlink-allocator
+          bash build.sh ../../build/mooncake-transfer-engine/nvlink-allocator/
+        shell: bash
+
+      - name: Run sccache stat for check
+        if: ${{ env.SCCACHE_PATH != '' }}
+        shell: bash
+        run: ${SCCACHE_PATH} --show-stats
+
+      - name: Generate Python version tag
+        id: generate_tag_release
+        run: |
+          echo "python_version_tag=$(echo ${{ matrix.python-version }} | tr -d '.')" >> $GITHUB_OUTPUT
+        shell: bash
+
+      - name: Build Python wheel
+        run: |
+          export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/lib
+          PYTHON_VERSION=${{ matrix.python-version }} OUTPUT_DIR=dist-py${{ steps.generate_tag_release.outputs.python_version_tag }} ./scripts/build_wheel.sh
+        env:
+          VERSION: ${{ env.VERSION }}
+
+      - name: Upload Python wheel artifact
+        uses: actions/upload-artifact@v4
+        with:
+          name: pre-release-cuda13-py${{ steps.generate_tag_release.outputs.python_version_tag }}
+          path: mooncake-wheel/dist-py${{ steps.generate_tag_release.outputs.python_version_tag }}/*.whl
+
+  build-cuda-arm64:
+    name: Build (CUDA 12, arm64)
+    runs-on: ubuntu-22.04-arm
+    permissions:
+      contents: read
+    strategy:
+      matrix:
+        python-version: ['3.10', '3.11', '3.12', '3.13']
+    env:
+      TORCH_CUDA_ARCH_LIST: "9.0"
+      CUDA_HOME: "/usr/local/cuda"
+    steps:
+      - name: Checkout source
+        uses: actions/checkout@v4
+
+      - name: Set version from tag
+        run: echo "VERSION=${GITHUB_REF_NAME#v}" >> $GITHUB_ENV
+
+      - name: Set up Python ${{ matrix.python-version }}
+        uses: actions/setup-python@v5
+        with:
+          python-version: ${{ matrix.python-version }}
+
+      - name: Free up disk space
+        run: |
+          sudo rm -rf /usr/share/dotnet
+          sudo rm -rf /opt/ghc
+          sudo rm -rf /opt/hostedtoolcache/CodeQL
+
+      - name: Install CUDA Toolkit 12.8 (arm64 SBSA)
+        run: |
+          wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/sbsa/cuda-keyring_1.1-1_all.deb
+          sudo dpkg -i cuda-keyring_1.1-1_all.deb
+          sudo apt-get update
+          sudo apt-get install -y cuda-toolkit-12-8
+          echo "/usr/local/cuda/bin" >> $GITHUB_PATH
+          /usr/local/cuda/bin/nvcc --version
+        shell: bash
+
+      - name: Run sccache-cache
+        uses: mozilla-actions/sccache-action@v0.0.9
+
+      - name: Configure sccache
+        uses: actions/github-script@v7
+        with:
+          script: |
+            core.exportVariable('ACTIONS_RESULTS_URL', process.env.ACTIONS_RESULTS_URL || '');
+            core.exportVariable('ACTIONS_RUNTIME_TOKEN', process.env.ACTIONS_RUNTIME_TOKEN || '');
+
+      - name: Configure project
+        run: |
+          sudo apt update -y
+          sudo bash -x dependencies.sh -y
+          mkdir build
+          cd build
+          cmake -G Ninja .. -DUSE_HTTP=ON -DUSE_CUDA=ON -DUSE_MNNVL=ON -DWITH_EP=OFF -DWITH_STORE_RUST=OFF -DENABLE_SCCACHE=ON -DCMAKE_BUILD_TYPE=Release
+        shell: bash
+
+      - name: Build project
+        run: |
+          export CUDA_HOME=/usr/local/cuda
+          for dir in /usr/local/cuda/lib64/stubs /usr/local/cuda/targets/*/lib/stubs; do
+            if [ -d "$dir" ]; then
+              export LIBRARY_PATH="$dir:${LIBRARY_PATH:-}"
+            fi
+          done
+          cd build
+          cmake --build .
+          sudo cmake --install .
+        shell: bash
+
+      - name: Run sccache stat for check
+        if: ${{ env.SCCACHE_PATH != '' }}
+        shell: bash
+        run: ${SCCACHE_PATH} --show-stats
+
+      - name: Generate Python version tag
+        id: generate_tag_arm64
+        run: |
+          echo "python_version_tag=$(echo ${{ matrix.python-version }} | tr -d '.')" >> $GITHUB_OUTPUT
+        shell: bash
+
+      - name: Build Python wheel
+        run: |
+          export CUDA_HOME=/usr/local/cuda
+          PYTHON_VERSION=${{ matrix.python-version }} OUTPUT_DIR=dist-py${{ steps.generate_tag_arm64.outputs.python_version_tag }} ./scripts/build_wheel.sh
+        env:
+          VERSION: ${{ env.VERSION }}
+        shell: bash
+
+      - name: Upload Python wheel artifact
+        uses: actions/upload-artifact@v4
+        with:
+          name: pre-release-cuda-arm64-py${{ steps.generate_tag_arm64.outputs.python_version_tag }}
+          path: mooncake-wheel/dist-py${{ steps.generate_tag_arm64.outputs.python_version_tag }}/*.whl
+
+  build-cuda13-arm64:
+    name: Build (CUDA 13, arm64)
+    runs-on: ubuntu-22.04-arm
+    permissions:
+      contents: read
+    strategy:
+      matrix:
+        python-version: ['3.10', '3.11', '3.12', '3.13']
+    env:
+      CU13_BUILD: "1"
+      TORCH_CUDA_ARCH_LIST: "9.0"
+      CUDA_HOME: "/usr/local/cuda"
+    steps:
+      - name: Checkout source
+        uses: actions/checkout@v4
+
+      - name: Set version from tag
+        run: echo "VERSION=${GITHUB_REF_NAME#v}" >> $GITHUB_ENV
+
+      - name: Set up Python ${{ matrix.python-version }}
+        uses: actions/setup-python@v5
+        with:
+          python-version: ${{ matrix.python-version }}
+
+      - name: Free up disk space
+        run: |
+          sudo rm -rf /usr/share/dotnet
+          sudo rm -rf /opt/ghc
+          sudo rm -rf /opt/hostedtoolcache/CodeQL
+
+      - name: Install CUDA Toolkit 13.0 (arm64 SBSA)
+        run: |
+          wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/sbsa/cuda-keyring_1.1-1_all.deb
+          sudo dpkg -i cuda-keyring_1.1-1_all.deb
+          sudo apt-get update
+          sudo apt-get install -y cuda-toolkit-13-0
+          echo "/usr/local/cuda/bin" >> $GITHUB_PATH
+          /usr/local/cuda/bin/nvcc --version
+        shell: bash
+
+      - name: Run sccache-cache
+        uses: mozilla-actions/sccache-action@v0.0.9
+
+      - name: Configure sccache
+        uses: actions/github-script@v7
+        with:
+          script: |
+            core.exportVariable('ACTIONS_RESULTS_URL', process.env.ACTIONS_RESULTS_URL || '');
+            core.exportVariable('ACTIONS_RUNTIME_TOKEN', process.env.ACTIONS_RUNTIME_TOKEN || '');
+
+      - name: Configure project
+        run: |
+          sudo apt update -y
+          sudo bash -x dependencies.sh -y
+          mkdir build
+          cd build
+          cmake -G Ninja .. -DUSE_HTTP=ON -DUSE_CUDA=ON -DUSE_MNNVL=ON -DWITH_EP=OFF -DWITH_STORE_RUST=OFF -DENABLE_SCCACHE=ON -DCMAKE_BUILD_TYPE=Release
+        shell: bash
+
+      - name: Build project
+        run: |
+          export CUDA_HOME=/usr/local/cuda
+          for dir in /usr/local/cuda/lib64/stubs /usr/local/cuda/targets/*/lib/stubs; do
+            if [ -d "$dir" ]; then
+              export LIBRARY_PATH="$dir:${LIBRARY_PATH:-}"
+            fi
+          done
+          cd build
+          cmake --build .
+          sudo cmake --install .
+        shell: bash
+
+      - name: Run sccache stat for check
+        if: ${{ env.SCCACHE_PATH != '' }}
+        shell: bash
+        run: ${SCCACHE_PATH} --show-stats
+
+      - name: Generate Python version tag
+        id: generate_tag_arm64
+        run: |
+          echo "python_version_tag=$(echo ${{ matrix.python-version }} | tr -d '.')" >> $GITHUB_OUTPUT
+        shell: bash
+
+      - name: Build Python wheel
+        run: |
+          export CUDA_HOME=/usr/local/cuda
+          PYTHON_VERSION=${{ matrix.python-version }} OUTPUT_DIR=dist-py${{ steps.generate_tag_arm64.outputs.python_version_tag }} ./scripts/build_wheel.sh
+        env:
+          VERSION: ${{ env.VERSION }}
+        shell: bash
+
+      - name: Upload Python wheel artifact
+        uses: actions/upload-artifact@v4
+        with:
+          name: pre-release-cuda13-arm64-py${{ steps.generate_tag_arm64.outputs.python_version_tag }}
+          path: mooncake-wheel/dist-py${{ steps.generate_tag_arm64.outputs.python_version_tag }}/*.whl
+
+  validate-release:
+    name: Validate release artifacts
+    needs: [build-cuda, build-non-cuda, build-cuda13, build-cuda-arm64, build-cuda13-arm64]
+    runs-on: ubuntu-22.04
+    permissions:
+      contents: read
+    steps:
+      - name: Checkout source
+        uses: actions/checkout@v4
+
+      - name: Download all pre-release wheel artifacts
+        uses: actions/download-artifact@v4
+        with:
+          path: mooncake-wheel/dist-all
+          pattern: pre-release-*
+
+      - name: Prepare wheels for validation
+        run: |
+          mkdir -p mooncake-wheel/dist-release
+          find mooncake-wheel/dist-all -name "*.whl" -exec cp {} mooncake-wheel/dist-release/ \;
+          echo "Pre-release tag: ${GITHUB_REF_NAME}"
+          echo "Collected wheels:"
+          ls -la mooncake-wheel/dist-release/
+          wheel_count=$(find mooncake-wheel/dist-release -name "*.whl" | wc -l)
+          echo "wheel_count=${wheel_count}" >> "$GITHUB_ENV"
+          if [ "${wheel_count}" -lt 20 ]; then
+            echo "Expected at least 20 wheels (x86: 4 Python versions x 3 variants; arm64: 4 Python versions x 2 CUDA variants), found ${wheel_count}"
+            exit 1
+          fi
+
+      - name: Set up Python
+        uses: actions/setup-python@v5
+        with:
+          python-version: '3.12'
+
+      - name: Validate wheels with twine
+        run: |
+          pip install twine
+          twine check mooncake-wheel/dist-release/*.whl
+
+      - name: Upload validated wheels as workflow artifacts
+        uses: actions/upload-artifact@v4
+        with:
+          name: pre-release-wheels-${{ github.ref_name }}
+          path: mooncake-wheel/dist-release/*.whl
+          retention-days: 14
+
+      - name: Pre-release validation summary
+        run: |
+          echo "## Pre-Release validation passed" >> "$GITHUB_STEP_SUMMARY"
+          echo "" >> "$GITHUB_STEP_SUMMARY"
+          echo "- Tag: \`${GITHUB_REF_NAME}\`" >> "$GITHUB_STEP_SUMMARY"
+          echo "- Wheels built: ${wheel_count}" >> "$GITHUB_STEP_SUMMARY"
+          echo "- PyPI publish: skipped (pre-release dry run)" >> "$GITHUB_STEP_SUMMARY"
+          echo "- GitHub Release upload: skipped (pre-release dry run)" >> "$GITHUB_STEP_SUMMARY"
diff --git a/.github/workflows/release-cuda13.yaml b/.github/workflows/release-cuda13.yaml
index 42d8fc0b..e12bb16b 100644
--- a/.github/workflows/release-cuda13.yaml
+++ b/.github/workflows/release-cuda13.yaml
@@ -10,6 +10,7 @@ env:
 jobs:
   build:
     runs-on: ubuntu-22.04
+    container: pytorch/manylinux2_28-builder:cuda13.0
     permissions:
       contents: write
     strategy:
@@ -23,27 +24,15 @@ jobs:
       - name: Checkout source
         uses: actions/checkout@v4
 
-      - name: Set up Python ${{ matrix.python-version }}
-        uses: actions/setup-python@v5
-        with:
-          python-version: ${{ matrix.python-version }}
+      - name: Mark workspace safe for git (container runs as root)
+        run: git config --global --add safe.directory '*'
 
-      - name: Free up disk space
+      - name: Select Python ${{ matrix.python-version }} from manylinux image
         run: |
-          sudo rm -rf /usr/share/dotnet
-          sudo rm -rf /opt/ghc
-          sudo rm -rf /opt/hostedtoolcache/CodeQL
-          sudo rm -rf /usr/local/lib/android
-          df -h
-
-      - name: Install CUDA Toolkit 13
-        uses: Jimver/cuda-toolkit@v0.2.29
-        with:
-          cuda: '13.0.2'
-          linux-local-args: '["--toolkit"]'
-          method: 'network'
-          sub-packages: '["nvcc", "nvrtc-dev"]'
-          non-cuda-sub-packages: '["libcusparse-dev", "libcublas-dev", "libcusolver-dev"]'
+          PYV_NODOT=$(echo "${{ matrix.python-version }}" | tr -d '.')
+          PYBIN="/opt/python/cp${PYV_NODOT}-cp${PYV_NODOT}/bin"
+          echo "$PYBIN" >> "$GITHUB_PATH"
+          "$PYBIN/pip" install --quiet "cmake<4" setuptools wheel
 
       - name: Run sccache-cache
         uses: mozilla-actions/sccache-action@v0.0.9
@@ -55,17 +44,13 @@ jobs:
             core.exportVariable('ACTIONS_RESULTS_URL', process.env.ACTIONS_RESULTS_URL || '');
             core.exportVariable('ACTIONS_RUNTIME_TOKEN', process.env.ACTIONS_RUNTIME_TOKEN || '');
 
-      - name: Run sccache stat for check
-        shell: bash
-        run: ${SCCACHE_PATH} --show-stats
-
       - name: Configure project
         run: |
-          sudo apt update -y
-          sudo bash -x dependencies.sh -y
+          bash -x dependencies.sh -y
+          echo "/usr/local/go/bin" >> "$GITHUB_PATH"
           mkdir build
           cd build
-          cmake .. -DBUILD_UNIT_TESTS=OFF -DUSE_HTTP=ON -DUSE_ETCD=ON -DUSE_CUDA=ON -DWITH_EP=ON -DEP_TORCH_VERSIONS="2.9.0;2.9.1;2.10.0;2.11.0" -DSTORE_USE_ETCD=ON -DENABLE_SCCACHE=ON -DCMAKE_BUILD_TYPE=Release
+          cmake .. -DBUILD_UNIT_TESTS=OFF -DUSE_HTTP=ON -DUSE_ETCD=ON -DUSE_CUDA=ON -DWITH_EP=ON -DEP_TORCH_VERSIONS="2.9.1;2.10.0;2.11.0;2.12.0;2.12.1" -DSTORE_USE_ETCD=ON -DENABLE_SCCACHE=ON -DCMAKE_BUILD_TYPE=Release -DPython3_EXECUTABLE="$(which python3)"
         shell: bash
 
       - name: Build project
@@ -73,7 +58,7 @@ jobs:
           export LIBRARY_PATH=/usr/local/cuda/lib64/stubs:$LIBRARY_PATH
           cd build
           make -j
-          sudo make install
+          make install
         shell: bash
 
       - name: Build nvlink_allocator.so
@@ -86,6 +71,11 @@ jobs:
           bash build.sh ../../build/mooncake-transfer-engine/nvlink-allocator/
         shell: bash
 
+      - name: Run sccache stat for check
+        if: ${{ env.SCCACHE_PATH != '' }}
+        shell: bash
+        run: ${SCCACHE_PATH} --show-stats
+
       - name: Generate Python version tag
         id: generate_tag_release
         run: |
@@ -106,8 +96,102 @@ jobs:
           name: mooncake-wheel-cuda13-py${{ steps.generate_tag_release.outputs.python_version_tag }}
           path: mooncake-wheel/dist-py${{ steps.generate_tag_release.outputs.python_version_tag }}/*.whl
 
+  build-arm64:
+    if: ${{ !contains(github.ref_name, '-') }}
+    runs-on: ubuntu-22.04-arm
+    permissions:
+      contents: write
+    strategy:
+      matrix:
+        python-version: ['3.10', '3.11', '3.12', '3.13']
+    env:
+      CU13_BUILD: "1"
+      TORCH_CUDA_ARCH_LIST: "9.0"
+      CUDA_HOME: "/usr/local/cuda"
+    steps:
+      - name: Checkout source
+        uses: actions/checkout@v4
+
+      - name: Set up Python ${{ matrix.python-version }}
+        uses: actions/setup-python@v5
+        with:
+          python-version: ${{ matrix.python-version }}
+
+      - name: Free up disk space
+        run: |
+          sudo rm -rf /usr/share/dotnet
+          sudo rm -rf /opt/ghc
+          sudo rm -rf /opt/hostedtoolcache/CodeQL
+
+      - name: Install CUDA Toolkit 13.0 (arm64 SBSA)
+        run: |
+          wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/sbsa/cuda-keyring_1.1-1_all.deb
+          sudo dpkg -i cuda-keyring_1.1-1_all.deb
+          sudo apt-get update
+          sudo apt-get install -y cuda-toolkit-13-0
+          echo "/usr/local/cuda/bin" >> $GITHUB_PATH
+          /usr/local/cuda/bin/nvcc --version
+        shell: bash
+
+      - name: Run sccache-cache
+        uses: mozilla-actions/sccache-action@v0.0.9
+
+      - name: Configure sccache
+        uses: actions/github-script@v7
+        with:
+          script: |
+            core.exportVariable('ACTIONS_RESULTS_URL', process.env.ACTIONS_RESULTS_URL || '');
+            core.exportVariable('ACTIONS_RUNTIME_TOKEN', process.env.ACTIONS_RUNTIME_TOKEN || '');
+
+      - name: Configure project
+        run: |
+          sudo apt update -y
+          sudo bash -x dependencies.sh -y
+          mkdir build
+          cd build
+          cmake -G Ninja .. -DUSE_HTTP=ON -DUSE_CUDA=ON -DUSE_MNNVL=ON -DWITH_EP=OFF -DWITH_STORE_RUST=OFF -DENABLE_SCCACHE=ON -DCMAKE_BUILD_TYPE=Release
+        shell: bash
+
+      - name: Build project
+        run: |
+          export CUDA_HOME=/usr/local/cuda
+          for dir in /usr/local/cuda/lib64/stubs /usr/local/cuda/targets/*/lib/stubs; do
+            if [ -d "$dir" ]; then
+              export LIBRARY_PATH="$dir:${LIBRARY_PATH:-}"
+            fi
+          done
+          cd build
+          cmake --build .
+          sudo cmake --install .
+        shell: bash
+
+      - name: Run sccache stat for check
+        if: ${{ env.SCCACHE_PATH != '' }}
+        shell: bash
+        run: ${SCCACHE_PATH} --show-stats
+
+      - name: Generate Python version tag
+        id: generate_tag_release
+        run: |
+          echo "python_version_tag=$(echo ${{ matrix.python-version }} | tr -d '.')" >> $GITHUB_OUTPUT
+        shell: bash
+
+      - name: Build Python wheel
+        run: |
+          export CUDA_HOME=/usr/local/cuda
+          export LD_LIBRARY_PATH="${LD_LIBRARY_PATH:-}:/usr/local/lib"
+          PYTHON_VERSION=${{ matrix.python-version }} OUTPUT_DIR=dist-py${{ steps.generate_tag_release.outputs.python_version_tag }} ./scripts/build_wheel.sh
+        shell: bash
+
+      - name: Upload Python wheel artifact
+        uses: actions/upload-artifact@v4
+        with:
+          name: mooncake-wheel-cuda13-arm64-py${{ steps.generate_tag_release.outputs.python_version_tag }}
+          path: mooncake-wheel/dist-py${{ steps.generate_tag_release.outputs.python_version_tag }}/*.whl
+
   publish-release:
-    needs: build
+    if: ${{ !contains(github.ref_name, '-') }}
+    needs: [build, build-arm64]
     runs-on: ubuntu-22.04
     permissions:
       contents: write
@@ -120,7 +204,7 @@ jobs:
         uses: actions/download-artifact@v4
         with:
           path: mooncake-wheel/dist-all
-          pattern: mooncake-wheel-cuda13-py*
+          pattern: mooncake-wheel-cuda13*
 
       - name: Prepare wheels for release
         run: |
diff --git a/.github/workflows/release-efa-non-cuda.yaml b/.github/workflows/release-efa-non-cuda.yaml
new file mode 100644
index 00000000..9916ba4b
--- /dev/null
+++ b/.github/workflows/release-efa-non-cuda.yaml
@@ -0,0 +1,155 @@
+name: Release EFA Non-CUDA
+
+on:
+  push:
+    tags:
+      - 'v*'
+
+# Publishes the AWS EFA (libfabric) non-CUDA wheel variant:
+#   mooncake-transfer-engine-efa-non-cuda   USE_EFA=ON USE_CUDA=OFF (CPU/DRAM only)
+# The CUDA variant (mooncake-transfer-engine-efa) is built by release-efa.yaml
+# — split into its own workflow because each PyPI package publishes from a
+# dedicated release workflow (trusted publisher / artifact pattern is
+# per-package), mirroring release.yaml vs release-non-cuda.yaml.
+#
+# No EFA hardware is needed to build: USE_EFA only needs libfabric headers/lib
+# to compile/link. auditwheel excludes libfabric/libefa from the wheel so they
+# resolve to the user's system AWS EFA install (/opt/amazon/efa/lib) at
+# runtime. The distro libfabric (1.x) used to build is ABI-forward-compatible
+# with the AWS EFA libfabric (2.x) loaded at runtime; the EFA transport uses
+# only long-stable fi_* core APIs.
+env:
+  SCCACHE_GHA_ENABLED: "true"
+jobs:
+  build:
+    runs-on: ubuntu-22.04
+    permissions:
+      contents: write
+    strategy:
+      matrix:
+        python-version: ['3.10', '3.11', '3.12', '3.13']
+    env:
+      EFA_NON_CUDA_BUILD: "1"
+    steps:
+      - name: Checkout source
+        uses: actions/checkout@v4
+
+      - name: Set up Python ${{ matrix.python-version }}
+        uses: actions/setup-python@v5
+        with:
+          python-version: ${{ matrix.python-version }}
+
+      - name: Free up disk space
+        run: |
+          sudo rm -rf /usr/share/dotnet
+          sudo rm -rf /opt/ghc
+          sudo rm -rf /opt/hostedtoolcache/CodeQL
+          sudo rm -rf /usr/local/lib/android
+
+      - name: Run sccache-cache
+        uses: mozilla-actions/sccache-action@v0.0.9
+
+      - name: Configure sccache
+        uses: actions/github-script@v7
+        with:
+          script: |
+            core.exportVariable('ACTIONS_RESULTS_URL', process.env.ACTIONS_RESULTS_URL || '');
+            core.exportVariable('ACTIONS_RUNTIME_TOKEN', process.env.ACTIONS_RUNTIME_TOKEN || '');
+
+      - name: Configure project
+        run: |
+          sudo apt update -y
+          sudo apt install -y libfabric-dev libfabric1
+          sudo bash -x dependencies.sh -y
+          mkdir build
+          cd build
+          cmake .. \
+            -DUSE_HTTP=ON \
+            -DUSE_ETCD=ON \
+            -DUSE_CUDA=OFF \
+            -DWITH_EP=OFF \
+            -DSTORE_USE_ETCD=ON \
+            -DUSE_EFA=ON \
+            -DLIBFABRIC_INCLUDE_DIR=/usr/include \
+            -DLIBFABRIC_LIBRARY=/usr/lib/x86_64-linux-gnu/libfabric.so \
+            -DENABLE_SCCACHE=ON \
+            -DCMAKE_BUILD_TYPE=Release
+        shell: bash
+
+      - name: Build project
+        run: |
+          cd build
+          make -j
+          sudo make install
+        shell: bash
+
+      - name: Run sccache stat for check
+        if: ${{ env.SCCACHE_PATH != '' }}
+        shell: bash
+        run: ${SCCACHE_PATH} --show-stats
+
+      - name: Generate Python version tag
+        id: generate_tag_release
+        run: |
+          echo "python_version_tag=$(echo ${{ matrix.python-version }} | tr -d '.')" >> $GITHUB_OUTPUT
+        shell: bash
+
+      - name: Build Python wheel
+        run: |
+          export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/lib
+          PYTHON_VERSION=${{ matrix.python-version }} OUTPUT_DIR=dist-py${{ steps.generate_tag_release.outputs.python_version_tag }} ./scripts/build_wheel.sh
+        env:
+          VERSION: ${{ env.VERSION }}
+
+      - name: Verify libfabric is excluded from the wheel
+        run: |
+          WHL=$(ls mooncake-wheel/dist-py${{ steps.generate_tag_release.outputs.python_version_tag }}/*.whl | head -1)
+          echo "Inspecting $WHL"
+          if unzip -l "$WHL" | grep -iE 'libfabric|libefa'; then
+            echo "::error::libfabric/libefa must NOT be bundled in the EFA wheel"
+            exit 1
+          fi
+          echo "OK: libfabric/libefa correctly excluded"
+        shell: bash
+
+      - name: Upload Python wheel artifact
+        uses: actions/upload-artifact@v4
+        with:
+          name: mooncake-wheel-efa-non-cuda-py${{ steps.generate_tag_release.outputs.python_version_tag }}
+          path: mooncake-wheel/dist-py${{ steps.generate_tag_release.outputs.python_version_tag }}/*.whl
+
+  publish-release:
+    if: ${{ !contains(github.ref_name, '-') }}
+    needs: build
+    runs-on: ubuntu-22.04
+    permissions:
+      contents: write
+      id-token: write
+    steps:
+      - name: Checkout source
+        uses: actions/checkout@v4
+
+      - name: Download all wheel artifacts
+        uses: actions/download-artifact@v4
+        with:
+          path: mooncake-wheel/dist-all
+          pattern: mooncake-wheel-efa-non-cuda-py*
+
+      - name: Prepare wheels for release
+        run: |
+          mkdir -p mooncake-wheel/dist-release
+          find mooncake-wheel/dist-all -name "*.whl" -exec cp {} mooncake-wheel/dist-release/ \;
+          echo "Collected wheels for release:"
+          ls -la mooncake-wheel/dist-release/
+
+      - name: Upload wheels to GitHub Release
+        uses: softprops/action-gh-release@v1
+        with:
+          files: mooncake-wheel/dist-release/*.whl
+
+      - name: Publish package to PyPI
+        if: github.repository == 'kvcache-ai/Mooncake'
+        uses: pypa/gh-action-pypi-publish@release/v1
+        with:
+          packages-dir: mooncake-wheel/dist-release/
+          password: ${{ secrets.PYPI_API_TOKEN }}
diff --git a/.github/workflows/release-efa.yaml b/.github/workflows/release-efa.yaml
new file mode 100644
index 00000000..cc327358
--- /dev/null
+++ b/.github/workflows/release-efa.yaml
@@ -0,0 +1,166 @@
+name: Release EFA
+
+on:
+  push:
+    tags:
+      - 'v*'
+
+# Publishes the AWS EFA (libfabric) CUDA wheel variant:
+#   mooncake-transfer-engine-efa            USE_EFA=ON USE_CUDA=ON  (GPU, GPUDirect/FI_HMEM_CUDA)
+# The non-CUDA variant (mooncake-transfer-engine-efa-non-cuda) is built by
+# release-efa-non-cuda.yaml — split into its own workflow because each PyPI
+# package publishes from a dedicated release workflow (trusted publisher /
+# artifact pattern is per-package), mirroring release.yaml vs release-non-cuda.yaml.
+#
+# No EFA hardware is needed to build: USE_EFA only needs libfabric headers/lib
+# to compile/link. auditwheel excludes libfabric/libefa from the wheel so they
+# resolve to the user's system AWS EFA install (/opt/amazon/efa/lib) at
+# runtime. The distro libfabric (1.x) used to build is ABI-forward-compatible
+# with the AWS EFA libfabric (2.x) loaded at runtime; the EFA transport uses
+# only long-stable fi_* core APIs.
+env:
+  SCCACHE_GHA_ENABLED: "true"
+jobs:
+  build:
+    runs-on: ubuntu-22.04
+    permissions:
+      contents: write
+    strategy:
+      matrix:
+        python-version: ['3.10', '3.11', '3.12', '3.13']
+    env:
+      EFA_BUILD: "1"
+    steps:
+      - name: Checkout source
+        uses: actions/checkout@v4
+
+      - name: Set up Python ${{ matrix.python-version }}
+        uses: actions/setup-python@v5
+        with:
+          python-version: ${{ matrix.python-version }}
+
+      - name: Free up disk space
+        run: |
+          sudo rm -rf /usr/share/dotnet
+          sudo rm -rf /opt/ghc
+          sudo rm -rf /opt/hostedtoolcache/CodeQL
+          sudo rm -rf /usr/local/lib/android
+
+      - name: Install CUDA Toolkit
+        uses: Jimver/cuda-toolkit@v0.2.24
+        with:
+          cuda: '12.8.1'
+          method: 'network'
+          sub-packages: '["nvcc", "nvrtc-dev"]'
+          non-cuda-sub-packages: '["libcusparse-dev", "libcublas-dev", "libcusolver-dev"]'
+
+      - name: Run sccache-cache
+        uses: mozilla-actions/sccache-action@v0.0.9
+
+      - name: Configure sccache
+        uses: actions/github-script@v7
+        with:
+          script: |
+            core.exportVariable('ACTIONS_RESULTS_URL', process.env.ACTIONS_RESULTS_URL || '');
+            core.exportVariable('ACTIONS_RUNTIME_TOKEN', process.env.ACTIONS_RUNTIME_TOKEN || '');
+
+      - name: Configure project
+        run: |
+          sudo apt update -y
+          sudo apt install -y libfabric-dev libfabric1
+          sudo bash -x dependencies.sh -y
+          mkdir build
+          cd build
+          cmake .. \
+            -DUSE_HTTP=ON \
+            -DUSE_ETCD=ON \
+            -DUSE_CUDA=ON \
+            -DWITH_EP=OFF \
+            -DSTORE_USE_ETCD=ON \
+            -DUSE_EFA=ON \
+            -DLIBFABRIC_INCLUDE_DIR=/usr/include \
+            -DLIBFABRIC_LIBRARY=/usr/lib/x86_64-linux-gnu/libfabric.so \
+            -DENABLE_SCCACHE=ON \
+            -DCMAKE_BUILD_TYPE=Release \
+            -DCMAKE_EXE_LINKER_FLAGS=-L/usr/local/cuda/lib64/stubs
+        shell: bash
+
+      - name: Build project
+        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 -j
+          sudo make install
+        shell: bash
+
+      - name: Run sccache stat for check
+        if: ${{ env.SCCACHE_PATH != '' }}
+        shell: bash
+        run: ${SCCACHE_PATH} --show-stats
+
+      - name: Generate Python version tag
+        id: generate_tag_release
+        run: |
+          echo "python_version_tag=$(echo ${{ matrix.python-version }} | tr -d '.')" >> $GITHUB_OUTPUT
+        shell: bash
+
+      - name: Build Python wheel
+        run: |
+          export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/lib
+          PYTHON_VERSION=${{ matrix.python-version }} OUTPUT_DIR=dist-py${{ steps.generate_tag_release.outputs.python_version_tag }} ./scripts/build_wheel.sh
+        env:
+          VERSION: ${{ env.VERSION }}
+
+      - name: Verify libfabric is excluded from the wheel
+        run: |
+          WHL=$(ls mooncake-wheel/dist-py${{ steps.generate_tag_release.outputs.python_version_tag }}/*.whl | head -1)
+          echo "Inspecting $WHL"
+          if unzip -l "$WHL" | grep -iE 'libfabric|libefa'; then
+            echo "::error::libfabric/libefa must NOT be bundled in the EFA wheel"
+            exit 1
+          fi
+          echo "OK: libfabric/libefa correctly excluded"
+        shell: bash
+
+      - name: Upload Python wheel artifact
+        uses: actions/upload-artifact@v4
+        with:
+          name: mooncake-wheel-efa-py${{ steps.generate_tag_release.outputs.python_version_tag }}
+          path: mooncake-wheel/dist-py${{ steps.generate_tag_release.outputs.python_version_tag }}/*.whl
+
+  publish-release:
+    if: ${{ !contains(github.ref_name, '-') }}
+    needs: build
+    runs-on: ubuntu-22.04
+    permissions:
+      contents: write
+      id-token: write
+    steps:
+      - name: Checkout source
+        uses: actions/checkout@v4
+
+      - name: Download all wheel artifacts
+        uses: actions/download-artifact@v4
+        with:
+          path: mooncake-wheel/dist-all
+          pattern: mooncake-wheel-efa-py*
+
+      - name: Prepare wheels for release
+        run: |
+          mkdir -p mooncake-wheel/dist-release
+          find mooncake-wheel/dist-all -name "*.whl" -exec cp {} mooncake-wheel/dist-release/ \;
+          echo "Collected wheels for release:"
+          ls -la mooncake-wheel/dist-release/
+
+      - name: Upload wheels to GitHub Release
+        uses: softprops/action-gh-release@v1
+        with:
+          files: mooncake-wheel/dist-release/*.whl
+
+      - name: Publish package to PyPI
+        if: github.repository == 'kvcache-ai/Mooncake'
+        uses: pypa/gh-action-pypi-publish@release/v1
+        with:
+          packages-dir: mooncake-wheel/dist-release/
+          password: ${{ secrets.PYPI_API_TOKEN }}
diff --git a/.github/workflows/release-musa.yaml b/.github/workflows/release-musa.yaml
new file mode 100644
index 00000000..577c2344
--- /dev/null
+++ b/.github/workflows/release-musa.yaml
@@ -0,0 +1,161 @@
+name: Release MUSA
+
+on:
+  push:
+    tags:
+      - 'v*'
+
+jobs:
+  build:
+    if: ${{ !contains(github.ref_name, '-') }}
+    runs-on: ubuntu-22.04
+    container: mthreads/musa:rc4.3.0-devel-ubuntu22.04-amd64
+
+    permissions:
+      contents: write
+
+    strategy:
+      max-parallel: 2
+      matrix:
+        python-version: ['3.9', '3.10', '3.11', '3.12', '3.13']
+
+    env:
+      BUILD_WITH_EP: "0"
+      MUSA_BUILD: "1"
+
+    steps:
+      - name: Checkout code
+        uses: actions/checkout@v4
+        with:
+          fetch-depth: 0
+
+      - name: Mark repository as safe
+        shell: bash
+        run: git config --global --add safe.directory "$GITHUB_WORKSPACE"
+
+      - name: Setup Python ${{ matrix.python-version }}
+        uses: actions/setup-python@v5
+        with:
+          python-version: ${{ matrix.python-version }}
+
+      - name: Set Python bin
+        shell: bash
+        run: |
+          set -euo pipefail
+          PYTHON_BIN="$(command -v python${{ matrix.python-version }})"
+          "$PYTHON_BIN" --version
+          "$PYTHON_BIN" -m pip --version
+          echo "PYTHON_BIN=${PYTHON_BIN}" >> "$GITHUB_ENV"
+
+      - name: Install dependencies
+        shell: bash
+        run: |
+          set -eo pipefail
+          export PYTHONPATH="${PYTHONPATH:-}"
+          export CMAKE_PREFIX_PATH="${CMAKE_PREFIX_PATH:-}"
+          bash -x dependencies.sh -y
+          echo "PATH=/usr/local/go/bin:/usr/local/musa/bin:${PATH}" >> "$GITHUB_ENV"
+          echo "LD_LIBRARY_PATH=/usr/local/musa/lib:${LD_LIBRARY_PATH:-}" >> "$GITHUB_ENV"
+          echo "LIBRARY_PATH=/usr/local/musa/lib:${LIBRARY_PATH:-}" >> "$GITHUB_ENV"
+
+      - name: Configure project
+        shell: bash
+        run: |
+          set -eo pipefail
+          export PYTHONPATH="${PYTHONPATH:-}"
+          export CMAKE_PREFIX_PATH="${CMAKE_PREFIX_PATH:-}"
+          rm -rf build
+          mkdir build
+          cd build
+          cmake_args=(
+            -DUSE_MUSA=ON
+            -DUSE_HTTP=ON
+            -DUSE_ETCD=ON
+            -DSTORE_USE_ETCD=ON
+            -DBUILD_UNIT_TESTS=OFF
+            -DENABLE_DEBUG_SYMBOLS=OFF
+            -DCMAKE_BUILD_TYPE=Release
+            -DPython3_EXECUTABLE="${PYTHON_BIN}"
+          )
+          cmake -G Ninja .. "${cmake_args[@]}"
+
+      - name: Build project
+        shell: bash
+        run: |
+          set -eo pipefail
+          export PYTHONPATH="${PYTHONPATH:-}"
+          export CMAKE_PREFIX_PATH="${CMAKE_PREFIX_PATH:-}"
+          cd build
+          cmake --build . -j"$(nproc)"
+
+      - name: Build MUSA allocator
+        shell: bash
+        run: |
+          set -eo pipefail
+          mkdir -p build/mooncake-transfer-engine/nvlink-allocator
+          cd mooncake-transfer-engine/nvlink-allocator
+          bash build.sh --use-mcc ../../build/mooncake-transfer-engine/nvlink-allocator/
+
+      - name: Install project
+        shell: bash
+        run: |
+          set -eo pipefail
+          cd build
+          cmake --install .
+
+      - name: Generate Python version tag
+        id: generate_tag_release
+        shell: bash
+        run: |
+          echo "python_version_tag=$(echo ${{ matrix.python-version }} | tr -d '.')" >> $GITHUB_OUTPUT
+
+      - name: Build Python wheel
+        shell: bash
+        run: |
+          set -eo pipefail
+          export LD_LIBRARY_PATH="${LD_LIBRARY_PATH:-}:/usr/local/lib:/usr/local/musa/lib"
+          MUSA_BUILD=1 PYTHON_VERSION=${{ matrix.python-version }} OUTPUT_DIR=dist-musa-py${{ steps.generate_tag_release.outputs.python_version_tag }} ./scripts/build_wheel.sh
+
+      - name: Upload Python wheel artifact
+        uses: actions/upload-artifact@v4
+        with:
+          name: mooncake-wheel-musa-x86_64-py${{ steps.generate_tag_release.outputs.python_version_tag }}
+          path: mooncake-wheel/dist-musa-py${{ steps.generate_tag_release.outputs.python_version_tag }}/*.whl
+
+  publish-release:
+    if: ${{ !contains(github.ref_name, '-') }}
+    needs: build
+    runs-on: ubuntu-22.04
+    environment: pypi
+
+    permissions:
+      contents: write
+      id-token: write
+
+    steps:
+      - name: Checkout source
+        uses: actions/checkout@v4
+
+      - name: Download all wheel artifacts
+        uses: actions/download-artifact@v4
+        with:
+          path: mooncake-wheel/dist-all
+          pattern: mooncake-wheel-musa-*
+
+      - name: Prepare wheels for release
+        run: |
+          mkdir -p mooncake-wheel/dist-release
+          find mooncake-wheel/dist-all -name "*.whl" -exec cp {} mooncake-wheel/dist-release/ \;
+          echo "Collected wheels for release:"
+          ls -la mooncake-wheel/dist-release/
+
+      - name: Upload wheels to GitHub Release
+        uses: softprops/action-gh-release@v1
+        with:
+          files: mooncake-wheel/dist-release/*.whl
+
+      - name: Publish package to PyPI
+        if: github.repository == 'kvcache-ai/Mooncake'
+        uses: pypa/gh-action-pypi-publish@release/v1
+        with:
+          packages-dir: mooncake-wheel/dist-release/
diff --git a/.github/workflows/release-non-cuda.yaml b/.github/workflows/release-non-cuda.yaml
index ab9b8424..18527173 100644
--- a/.github/workflows/release-non-cuda.yaml
+++ b/.github/workflows/release-non-cuda.yaml
@@ -43,10 +43,6 @@ jobs:
             core.exportVariable('ACTIONS_RESULTS_URL', process.env.ACTIONS_RESULTS_URL || '');
             core.exportVariable('ACTIONS_RUNTIME_TOKEN', process.env.ACTIONS_RUNTIME_TOKEN || '');
 
-      - name: Run sccache stat for check
-        shell: bash
-        run: ${SCCACHE_PATH} --show-stats
-
       - name: Configure project
         run: |
           sudo apt update -y
@@ -63,6 +59,11 @@ jobs:
           sudo make install
         shell: bash
 
+      - name: Run sccache stat for check
+        if: ${{ env.SCCACHE_PATH != '' }}
+        shell: bash
+        run: ${SCCACHE_PATH} --show-stats
+
       - name: Generate Python version tag
         id: generate_tag_release
         run: |
@@ -84,6 +85,7 @@ jobs:
           path: mooncake-wheel/dist-py${{ steps.generate_tag_release.outputs.python_version_tag }}/*.whl
 
   publish-release:
+    if: ${{ !contains(github.ref_name, '-') }}
     needs: build
     runs-on: ubuntu-22.04
     permissions:
@@ -120,4 +122,3 @@ jobs:
         with:
           packages-dir: mooncake-wheel/dist-release/
           password: ${{ secrets.PYPI_API_TOKEN }}
-
diff --git a/.github/workflows/release-npu.yaml b/.github/workflows/release-npu.yaml
new file mode 100644
index 00000000..dae2fd9b
--- /dev/null
+++ b/.github/workflows/release-npu.yaml
@@ -0,0 +1,196 @@
+name: Release Ascend NPU
+
+on:
+  push:
+    tags:
+      - 'v*'
+  workflow_dispatch:
+    inputs:
+      tag:
+        description: 'Release tag (e.g. v0.1.0)'
+        required: true
+        type: string
+
+jobs:
+  build:
+    if: ${{ (github.event_name == 'push' && !contains(github.ref_name, '-')) || (github.event_name == 'workflow_dispatch' && startsWith(inputs.tag, 'v') && !contains(inputs.tag, '-')) }}
+
+    strategy:
+      max-parallel: 2
+      matrix:
+        arch: [aarch64, x86_64]
+        python-version: ['3.9', '3.10', '3.11', '3.12', '3.13']
+        include:
+          - arch: aarch64
+            runner: ubuntu-22.04-arm
+            cann_arch: aarch64
+          - arch: x86_64
+            runner: ubuntu-22.04
+            cann_arch: x86_64
+
+    runs-on: ${{ matrix.runner }}
+
+    permissions:
+      contents: write
+
+    env:
+      BUILD_WITH_EP: "0"
+      NPU_BUILD: "1"
+
+    steps:
+      - name: Checkout code
+        uses: actions/checkout@v4
+        with:
+          ref: ${{ github.event_name == 'workflow_dispatch' && format('refs/tags/{0}', inputs.tag) || github.ref }}
+          fetch-depth: 0
+
+      - name: Setup Python ${{ matrix.python-version }}
+        uses: actions/setup-python@v5
+        with:
+          python-version: ${{ matrix.python-version }}
+
+      - name: Set Python bin
+        shell: bash
+        run: |
+          set -euo pipefail
+          PYTHON_BIN="$(command -v python${{ matrix.python-version }})"
+          "$PYTHON_BIN" --version
+          "$PYTHON_BIN" -m pip --version
+          echo "PYTHON_BIN=${PYTHON_BIN}" >> "$GITHUB_ENV"
+
+      - name: Install CANN Toolkit
+        shell: bash
+        run: |
+          set -euo pipefail
+          CANN_BASE_URL="https://ascend.devcloud.huaweicloud.com/cann/run/software"
+          CANN_VERSION=$(curl -s "${CANN_BASE_URL}/" | grep -oP '[0-9]+\.[0-9]+\.[0-9]+(?=/)' | sort -V | tail -1)
+          echo "Latest CANN version: ${CANN_VERSION}"
+          CANN_URL="${CANN_BASE_URL}/${CANN_VERSION}/${{ matrix.cann_arch }}/Ascend-cann-toolkit_${CANN_VERSION}_linux-${{ matrix.cann_arch }}.run"
+          echo "Downloading CANN ${CANN_VERSION} from ${CANN_URL}"
+          wget -q --show-progress -O /tmp/cann_toolkit.run "${CANN_URL}"
+          chmod +x /tmp/cann_toolkit.run
+          sudo /tmp/cann_toolkit.run --install --install-for-all --install-path=/usr/local/Ascend --quiet
+          rm -f /tmp/cann_toolkit.run
+
+      - name: Build and install HIXL
+        shell: bash
+        run: |
+          set -eo pipefail
+          export PYTHONPATH="${PYTHONPATH:-}"
+          export CMAKE_PREFIX_PATH="${CMAKE_PREFIX_PATH:-}"
+          source /usr/local/Ascend/cann/set_env.sh
+          sudo ${PYTHON_BIN} -m pip install setuptools
+          cd /tmp
+          git clone https://gitcode.com/cann/hixl.git
+          cd hixl
+          bash build.sh -j$(nproc)
+          sudo ./build_out/cann-hixl_*.run --full --quiet --pylocal --install-path=/usr/local/Ascend
+
+      - name: Install dependencies
+        shell: bash
+        run: |
+          set -eo pipefail
+          export PYTHONPATH="${PYTHONPATH:-}"
+          export CMAKE_PREFIX_PATH="${CMAKE_PREFIX_PATH:-}"
+          source /usr/local/Ascend/cann/set_env.sh
+          sudo bash -x dependencies.sh -y
+          sudo bash scripts/ascend/dependencies_ascend_installation.sh
+          echo "PATH=/usr/local/go/bin:${PATH}" >> "$GITHUB_ENV"
+
+      - name: Configure project
+        shell: bash
+        run: |
+          set -eo pipefail
+          export PYTHONPATH="${PYTHONPATH:-}"
+          export CMAKE_PREFIX_PATH="${CMAKE_PREFIX_PATH:-}"
+          source /usr/local/Ascend/cann/set_env.sh
+          rm -rf build
+          mkdir build
+          cd build
+          cmake_args=(
+            -DUSE_ASCEND_DIRECT=ON
+            -DUSE_ETCD=ON
+            -DSTORE_USE_ETCD=ON
+            -DBUILD_UNIT_TESTS=OFF
+            -DCMAKE_BUILD_TYPE=Release
+            -DPython3_EXECUTABLE="${PYTHON_BIN}"
+            -DCMAKE_INSTALL_RPATH_USE_LINK_PATH=OFF
+            -DCMAKE_BUILD_WITH_INSTALL_RPATH=OFF
+            -DCMAKE_SKIP_INSTALL_RPATH=ON
+          )
+          cmake .. "${cmake_args[@]}"
+
+      - name: Build project
+        shell: bash
+        run: |
+          set -eo pipefail
+          export PYTHONPATH="${PYTHONPATH:-}"
+          export CMAKE_PREFIX_PATH="${CMAKE_PREFIX_PATH:-}"
+          source /usr/local/Ascend/cann/set_env.sh
+          export PATH="/usr/local/go/bin:${PATH}"
+          cd build
+          cmake --build . -j"$(nproc)"
+          sudo cmake --install .
+
+      - name: Generate Python version tag
+        id: generate_tag_release
+        shell: bash
+        run: |
+          echo "python_version_tag=$(echo ${{ matrix.python-version }} | tr -d '.')" >> $GITHUB_OUTPUT
+
+      - name: Build Python wheel
+        shell: bash
+        run: |
+          set -eo pipefail
+          export PYTHONPATH="${PYTHONPATH:-}"
+          export CMAKE_PREFIX_PATH="${CMAKE_PREFIX_PATH:-}"
+          source /usr/local/Ascend/cann/set_env.sh
+          export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/lib
+          NPU_BUILD=1 PYTHON_VERSION=${{ matrix.python-version }} OUTPUT_DIR=dist-npu-py${{ steps.generate_tag_release.outputs.python_version_tag }} ./scripts/build_wheel.sh
+
+      - name: Upload Python wheel artifact
+        uses: actions/upload-artifact@v4
+        with:
+          name: mooncake-wheel-npu-${{ matrix.arch }}-py${{ steps.generate_tag_release.outputs.python_version_tag }}
+          path: mooncake-wheel/dist-npu-py${{ steps.generate_tag_release.outputs.python_version_tag }}/*.whl
+
+  publish-release:
+    if: ${{ (github.event_name == 'push' && !contains(github.ref_name, '-')) || (github.event_name == 'workflow_dispatch' && startsWith(inputs.tag, 'v') && !contains(inputs.tag, '-')) }}
+    needs: build
+    runs-on: ubuntu-22.04
+    environment: pypi
+
+    permissions:
+      contents: write
+      id-token: write
+
+    steps:
+      - name: Checkout source
+        uses: actions/checkout@v4
+        with:
+          ref: ${{ github.event_name == 'workflow_dispatch' && format('refs/tags/{0}', inputs.tag) || github.ref }}
+
+      - name: Download all wheel artifacts
+        uses: actions/download-artifact@v4
+        with:
+          path: mooncake-wheel/dist-all
+          pattern: mooncake-wheel-npu-*
+
+      - name: Prepare wheels for release
+        run: |
+          mkdir -p mooncake-wheel/dist-release
+          find mooncake-wheel/dist-all -name "*.whl" -exec cp {} mooncake-wheel/dist-release/ \;
+          echo "Collected wheels for release:"
+          ls -la mooncake-wheel/dist-release/
+
+      - name: Upload wheels to GitHub Release
+        uses: softprops/action-gh-release@v1
+        with:
+          tag_name: ${{ github.event_name == 'workflow_dispatch' && inputs.tag || github.ref_name }}
+          files: mooncake-wheel/dist-release/*.whl
+
+      - name: Publish package to PyPI
+        if: github.repository == 'kvcache-ai/Mooncake'
+        uses: pypa/gh-action-pypi-publish@release/v1
+        with:
+          packages-dir: mooncake-wheel/dist-release/
diff --git a/.github/workflows/release.yaml b/.github/workflows/release.yaml
index 26d5398d..51ea8e60 100644
--- a/.github/workflows/release.yaml
+++ b/.github/workflows/release.yaml
@@ -9,7 +9,10 @@ env:
   SCCACHE_GHA_ENABLED: "true"
 jobs:
   build:
+    # Skip semver pre-release tags (e.g. v1.0.0-rc1); those are handled by pre-release.yaml.
+    if: ${{ !contains(github.ref_name, '-') }}
     runs-on: ubuntu-22.04
+    container: pytorch/manylinux2_28-builder:cuda12.8
     permissions:
       contents: write
     strategy:
@@ -22,27 +25,15 @@ jobs:
       - name: Checkout source
         uses: actions/checkout@v4
 
-      - name: Set up Python ${{ matrix.python-version }}
-        uses: actions/setup-python@v5
-        with:
-          python-version: ${{ matrix.python-version }}
+      - name: Mark workspace safe for git (container runs as root)
+        run: git config --global --add safe.directory '*'
 
-      - name: Free up disk space
+      - name: Select Python ${{ matrix.python-version }} from manylinux image
         run: |
-          sudo rm -rf /usr/share/dotnet
-          sudo rm -rf /opt/ghc
-          sudo rm -rf /opt/hostedtoolcache/CodeQL
-          sudo rm -rf /usr/local/lib/android
-          df -h
-
-      - name: Install CUDA Toolkit
-        uses: Jimver/cuda-toolkit@v0.2.24
-        with:
-          cuda: '12.8.1'
-          linux-local-args: '["--toolkit"]'
-          method: 'network'
-          sub-packages: '["nvcc", "nvrtc-dev"]'
-          non-cuda-sub-packages: '["libcusparse-dev", "libcublas-dev", "libcusolver-dev"]'
+          PYV_NODOT=$(echo "${{ matrix.python-version }}" | tr -d '.')
+          PYBIN="/opt/python/cp${PYV_NODOT}-cp${PYV_NODOT}/bin"
+          echo "$PYBIN" >> "$GITHUB_PATH"
+          "$PYBIN/pip" install --quiet "cmake<4" setuptools wheel
 
       - name: Run sccache-cache
         uses: mozilla-actions/sccache-action@v0.0.9
@@ -54,25 +45,21 @@ jobs:
             core.exportVariable('ACTIONS_RESULTS_URL', process.env.ACTIONS_RESULTS_URL || '');
             core.exportVariable('ACTIONS_RUNTIME_TOKEN', process.env.ACTIONS_RUNTIME_TOKEN || '');
 
-      - name: Run sccache stat for check
-        shell: bash
-        run: ${SCCACHE_PATH} --show-stats
-
       - name: Configure project
         run: |
-          sudo apt update -y
-          sudo bash -x dependencies.sh -y
+          bash -x dependencies.sh -y
+          echo "/usr/local/go/bin" >> "$GITHUB_PATH"
           mkdir build
           cd build
-          cmake .. -DBUILD_UNIT_TESTS=OFF -DUSE_HTTP=ON -DUSE_ETCD=ON -DUSE_CUDA=ON -DWITH_EP=ON -DEP_TORCH_VERSIONS="2.9.0;2.9.1;2.10.0;2.11.0" -DSTORE_USE_ETCD=ON -DENABLE_SCCACHE=ON -DCMAKE_BUILD_TYPE=Release
+          cmake .. -DBUILD_UNIT_TESTS=OFF -DUSE_HTTP=ON -DUSE_ETCD=ON -DUSE_CUDA=ON -DWITH_EP=ON -DEP_TORCH_VERSIONS="2.9.1;2.10.0;2.11.0;2.12.0;2.12.1" -DSTORE_USE_ETCD=ON -DENABLE_SCCACHE=ON -DCMAKE_BUILD_TYPE=Release -DPython3_EXECUTABLE="$(which python3)"
         shell: bash
 
       - name: Build project
         run: |
-          export LIBRARY_PATH=/usr/local/cuda/lib64/stubs:LIBRARY_PATH
+          export LIBRARY_PATH=/usr/local/cuda/lib64/stubs:$LIBRARY_PATH
           cd build
           make -j
-          sudo make install
+          make install
         shell: bash
 
       - name: Build nvlink_allocator.so
@@ -85,6 +72,11 @@ jobs:
           bash build.sh ../../build/mooncake-transfer-engine/nvlink-allocator/
         shell: bash
 
+      - name: Run sccache stat for check
+        if: ${{ env.SCCACHE_PATH != '' }}
+        shell: bash
+        run: ${SCCACHE_PATH} --show-stats
+
       - name: Generate Python version tag
         id: generate_tag_release
         run: |
@@ -105,8 +97,100 @@ jobs:
           name: mooncake-wheel-py${{ steps.generate_tag_release.outputs.python_version_tag }}
           path: mooncake-wheel/dist-py${{ steps.generate_tag_release.outputs.python_version_tag }}/*.whl
 
+  build-arm64:
+    if: ${{ !contains(github.ref_name, '-') }}
+    runs-on: ubuntu-22.04-arm
+    permissions:
+      contents: write
+    strategy:
+      matrix:
+        python-version: ['3.10', '3.11', '3.12', '3.13']
+    env:
+      TORCH_CUDA_ARCH_LIST: "9.0"
+      CUDA_HOME: "/usr/local/cuda"
+    steps:
+      - name: Checkout source
+        uses: actions/checkout@v4
+
+      - name: Set up Python ${{ matrix.python-version }}
+        uses: actions/setup-python@v5
+        with:
+          python-version: ${{ matrix.python-version }}
+
+      - name: Free up disk space
+        run: |
+          sudo rm -rf /usr/share/dotnet
+          sudo rm -rf /opt/ghc
+          sudo rm -rf /opt/hostedtoolcache/CodeQL
+
+      - name: Install CUDA Toolkit 12.8 (arm64 SBSA)
+        run: |
+          wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/sbsa/cuda-keyring_1.1-1_all.deb
+          sudo dpkg -i cuda-keyring_1.1-1_all.deb
+          sudo apt-get update
+          sudo apt-get install -y cuda-toolkit-12-8
+          echo "/usr/local/cuda/bin" >> $GITHUB_PATH
+          /usr/local/cuda/bin/nvcc --version
+        shell: bash
+
+      - name: Run sccache-cache
+        uses: mozilla-actions/sccache-action@v0.0.9
+
+      - name: Configure sccache
+        uses: actions/github-script@v7
+        with:
+          script: |
+            core.exportVariable('ACTIONS_RESULTS_URL', process.env.ACTIONS_RESULTS_URL || '');
+            core.exportVariable('ACTIONS_RUNTIME_TOKEN', process.env.ACTIONS_RUNTIME_TOKEN || '');
+
+      - name: Configure project
+        run: |
+          sudo apt update -y
+          sudo bash -x dependencies.sh -y
+          mkdir build
+          cd build
+          cmake -G Ninja .. -DUSE_HTTP=ON -DUSE_CUDA=ON -DUSE_MNNVL=ON -DWITH_EP=OFF -DWITH_STORE_RUST=OFF -DENABLE_SCCACHE=ON -DCMAKE_BUILD_TYPE=Release
+        shell: bash
+
+      - name: Build project
+        run: |
+          export CUDA_HOME=/usr/local/cuda
+          for dir in /usr/local/cuda/lib64/stubs /usr/local/cuda/targets/*/lib/stubs; do
+            if [ -d "$dir" ]; then
+              export LIBRARY_PATH="$dir:${LIBRARY_PATH:-}"
+            fi
+          done
+          cd build
+          cmake --build .
+          sudo cmake --install .
+        shell: bash
+
+      - name: Run sccache stat for check
+        if: ${{ env.SCCACHE_PATH != '' }}
+        shell: bash
+        run: ${SCCACHE_PATH} --show-stats
+
+      - name: Generate Python version tag
+        id: generate_tag_release
+        run: |
+          echo "python_version_tag=$(echo ${{ matrix.python-version }} | tr -d '.')" >> $GITHUB_OUTPUT
+        shell: bash
+
+      - name: Build Python wheel
+        run: |
+          export CUDA_HOME=/usr/local/cuda
+          export LD_LIBRARY_PATH="${LD_LIBRARY_PATH:-}:/usr/local/lib"
+          PYTHON_VERSION=${{ matrix.python-version }} OUTPUT_DIR=dist-py${{ steps.generate_tag_release.outputs.python_version_tag }} ./scripts/build_wheel.sh
+        shell: bash
+
+      - name: Upload Python wheel artifact
+        uses: actions/upload-artifact@v4
+        with:
+          name: mooncake-wheel-arm64-py${{ steps.generate_tag_release.outputs.python_version_tag }}
+          path: mooncake-wheel/dist-py${{ steps.generate_tag_release.outputs.python_version_tag }}/*.whl
+
   publish-release:
-    needs: build
+    needs: [build, build-arm64]
     runs-on: ubuntu-22.04
     permissions:
       contents: write
diff --git a/.github/workflows/stale.yml b/.github/workflows/stale.yml
new file mode 100644
index 00000000..d4b92c8f
--- /dev/null
+++ b/.github/workflows/stale.yml
@@ -0,0 +1,41 @@
+name: Close Stale Issues and PRs
+
+on:
+  schedule:
+    - cron: '37 2 * * *'
+  workflow_dispatch: {}
+
+permissions:
+  issues: write
+  pull-requests: write
+
+jobs:
+  stale:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/stale@v9
+        with:
+          repo-token: ${{ secrets.GITHUB_TOKEN }}
+          stale-issue-message: >
+            This issue has been automatically marked as stale because it has not had
+            recent activity. It will be closed in 30 days if no further activity occurs.
+            If this issue is still relevant, please comment to keep it open.
+          stale-pr-message: >
+            This PR has been automatically marked as stale because it has not had
+            recent activity. It will be closed in 14 days if no further activity occurs.
+            Please update the PR or comment to keep it open.
+          close-issue-message: >
+            This issue was closed because it has been inactive for 120 days.
+            Feel free to reopen if this is still relevant.
+          close-pr-message: >
+            This PR was closed because it has been inactive for 104 days.
+            Feel free to reopen if you'd like to continue the work.
+          days-before-stale: 120
+          days-before-close: 30
+          days-before-pr-stale: 120
+          days-before-pr-close: 30
+          stale-issue-label: 'stale'
+          stale-pr-label: 'stale'
+          exempt-issue-labels: 'pinned,security,RFC'
+          exempt-pr-labels: 'pinned,work-in-progress'
+          operations-per-run: 100
diff --git a/.gitignore b/.gitignore
index dafcb87c..4a3ec467 100644
--- a/.gitignore
+++ b/.gitignore
@@ -206,3 +206,9 @@ _codeql_detected_source_root
 
 # CodeBuddy Memory
 .codebuddy/
+
+# core dumps
+core_*
+# MacOS
+.DS_Store
+.envrc
diff --git a/.typos.toml b/.typos.toml
index b3fe85df..320972c8 100644
--- a/.typos.toml
+++ b/.typos.toml
@@ -1,13 +1,17 @@
 [default]
-extend-ignore-words = ["CANN", "ASO", "fre", "wqs"]
+extend-ignore-words = ["CANN", "ASO", "fre", "wqs", "hsa"]
 
 [default.extend-words]
 CANN = "CANN"
 ASO = "ASO"
 fre = "fre"
 wqs = "wqs"
+# AMD HSA runtime symbol prefix (hsa_*, hsaRes, hsaErr, etc.) — used by the
+# ROCm dmabuf MR registration path.
+hsa = "hsa"
+HPE = "HPE"
 
 [files]
 extend-exclude = [
   "mooncake-transfer-engine/tent/include/tent/thirdparty/nlohmann/json.h",
-]
\ No newline at end of file
+]
diff --git a/CMakeLists.txt b/CMakeLists.txt
index c6dd0699..79c79383 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -19,6 +19,7 @@ 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)
+option(USE_NOF "build mooncake store with NoF SSD pool support" OFF)
 
 include(${CMAKE_CURRENT_SOURCE_DIR}/mooncake-common/SetupPython.cmake)
 add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/extern/pybind11)
@@ -56,6 +57,12 @@ if (STORE_USE_K8S_LEASE)
   add_compile_definitions(STORE_USE_K8S_LEASE)
 endif()
 
+if (USE_NOF)
+  add_compile_definitions(USE_NOF)
+else()
+  message(STATUS "USE_NOF=OFF, NoF SSD pool support is disabled")
+endif()
+
 option(STORE_USE_JEMALLOC "Use jemalloc in mooncake store master" OFF)
 
 # Define ASIO macros before building targets that include ASIO headers.
@@ -63,6 +70,7 @@ add_compile_definitions(ASIO_SEPARATE_COMPILATION ASIO_DYN_LINK)
 
 add_subdirectory(mooncake-common)
 include_directories(mooncake-common/etcd)
+include_directories(mooncake-common/k8s-lease)
 include_directories(mooncake-common/include)
 
 if (WITH_TE)
@@ -94,8 +102,10 @@ if (WITH_EP)
     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}")
+    if(USE_CUDA)
+      find_package(CUDAToolkit REQUIRED)
+      message(STATUS "Detected CUDA version: ${CUDAToolkit_VERSION}")
+    endif()
 
     # 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.
@@ -136,6 +146,7 @@ if (WITH_EP)
         "-DTORCH_CUDA_ARCH_LIST=${_torch_cuda_arch_list_pipe}"
         "-DSTAGING_DIR=${EP_PG_STAGING_DIR}"
         "-DENGINE_SO_PATH=$<TARGET_FILE:engine>"
+        "-DEP_USE_MUSA=$<IF:$<BOOL:${USE_MUSA}>,1,0>"
         -P "${CMAKE_CURRENT_SOURCE_DIR}/mooncake-ep/BuildEpExt.cmake"
       COMMENT "Building Mooncake EP Python extension(s)"
       DEPENDS engine
@@ -152,6 +163,7 @@ if (WITH_EP)
         "-DTORCH_CUDA_ARCH_LIST=${_torch_cuda_arch_list_pipe}"
         "-DSTAGING_DIR=${EP_PG_STAGING_DIR}"
         "-DENGINE_SO_PATH=$<TARGET_FILE:engine>"
+        "-DEP_USE_MUSA=$<IF:$<BOOL:${USE_MUSA}>,1,0>"
         -P "${CMAKE_CURRENT_SOURCE_DIR}/mooncake-pg/BuildPgExt.cmake"
       COMMENT "Building Mooncake PG Python extension(s)"
       DEPENDS engine mooncake_ep_ext
diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md
index fdab919a..b18fa536 100644
--- a/CONTRIBUTING.md
+++ b/CONTRIBUTING.md
@@ -70,7 +70,7 @@ git commit -m "chore: pre-commit autoupdate"
 
 If clang-format is missing, install it (Ubuntu example):
 ```bash
-sudo apt-get update && sudo apt-get install -y clang-format
+sudo apt-get update && sudo apt-get install -y clang-format-20
 ```
 
 You can temporarily skip hooks:
diff --git a/FAST25-release/README.md b/FAST25-release/README.md
new file mode 100644
index 00000000..395186de
--- /dev/null
+++ b/FAST25-release/README.md
@@ -0,0 +1,17 @@
+```json
+{
+    "timestamp": 27482,
+    "input_length": 6955,
+    "output_length": 52,
+    "hash_ids": [46, 47, 48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 2353, 2354]
+}
+{
+    "timestamp": 30535,
+    "input_length": 6472,
+    "output_length": 26,
+    "hash_ids": [46, 47, 48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 2366]
+}
+```
+The above presents two samples from our trace dataset. The trace includes the timing of request arrivals, the number of input tokens, the number of output tokens, and the remapped block hash. To protect our customers' privacy, we applied several mechanisms to remove user-related information while preserving the dataset's utility for simulated evaluation. More descriptions of the trace (e.g., up to 50% cache hit ratio) can be found in Section 4 of the technical report.
+
+**_Update [Feb 21, 2025]: The updated [traces](./traces) used in our FAST'25 paper have been released! Please refer to the paper's appendix (found [here](Mooncake-FAST25.pdf)) for more details._**
\ No newline at end of file
diff --git a/MAINTAINERS.md b/MAINTAINERS.md
index 3e978d9c..297963d9 100644
--- a/MAINTAINERS.md
+++ b/MAINTAINERS.md
@@ -12,6 +12,6 @@ Current list of codeowners on this project:
 | ---------------------------- | ---------------------------- | ---------------------------- | ---------------------------- |
 | <img src="image/partners/approaching_ai_logo.png" width="120"/> | <img src="image/partners/ant_group_logo.png" width="120"/> | <img src="image/partners/huawei_logo.png" width="120"/> | <img src="image/partners/nvidia_logo.png" width="120"/> |
 | <img src="image/partners/moore_thread_logo.jpg" width="120"/> | <img src="image/partners/tencent_logo.png" width="120"/> | <img src="image/partners/volcengine_logo.png" width="120"/> | <img src="image/partners/amd_logo.png" width="120"/> |
-| <img src="image/partners/ieitsystems_logo.png" width="120"/> | | | |
+| <img src="image/partners/ieitsystems_logo.png" width="120"/> | <img src="image/partners/sunrise_logo.png" width="120"/> | <img src="image/partners/hygon_logo.png" width="120"/> | |
 
 Want to include your company logo? Just open a Pull Request!
\ No newline at end of file
diff --git a/README.md b/README.md
index 4a1e2ca9..b08dbfee 100644
--- a/README.md
+++ b/README.md
@@ -6,34 +6,41 @@
   <a href="https://www.usenix.org/system/files/fast25-qin.pdf" target="_blank"><strong>Paper</strong></a>
   | <a href="https://www.usenix.org/system/files/fast25_slides-qin.pdf" target="_blank"><strong>Slides</strong></a>
   | <a href="FAST25-release/traces" target="_blank"><strong>Traces</strong></a>
-  | <a href="https://arxiv.org/abs/2407.00079" target="_blank"><strong>Technical Report</strong></a>
-  | <a href="https://kvcache-ai.github.io/Mooncake/" target="_blank"><strong>Blog</strong></a>
+  | <a href="https://kvcache-ai.github.io/Mooncake/" target="_blank"><strong>Documentation</strong></a>
+  | <a href="https://kvcache.ai/" target="_blank"><strong>Blog</strong></a>
   | <a href="https://join.slack.com/t/mooncake-project/shared_invite/zt-3qx4x35ea-zSSTqTHItHJs9SCoXLOSPA" target="_blank"><strong>Slack</strong></a>
   <br />
   <br />
 
-  [![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)
+  [![PyPI - Downloads](https://static.pepy.tech/badge/mooncake-transfer-engine?period=month)](https://pypi.org/project/mooncake-transfer-engine)
   [![GitHub commit activity](https://img.shields.io/github/commit-activity/w/kvcache-ai/Mooncake)](https://github.com/kvcache-ai/Mooncake/graphs/commit-activity)
   [![license](https://img.shields.io/github/license/kvcache-ai/mooncake.svg)](https://github.com/kvcache-ai/Mooncake/blob/main/LICENSE-APACHE)
+  [![Docker](https://img.shields.io/docker/v/kvcacheai/mooncake?label=docker&logo=docker&logoColor=white&color=2496ED)](https://hub.docker.com/r/kvcacheai/mooncake)
+  <br />
 
+  [![PyPI](https://img.shields.io/pypi/v/mooncake-transfer-engine)](https://pypi.org/project/mooncake-transfer-engine)
+  [![PyPI CUDA <=12.9](https://img.shields.io/static/v1?label=pypi&message=CUDA%20%3C%3D12.9&color=76B900)](https://pypi.org/project/mooncake-transfer-engine)
+  [![PyPI CUDA 13.0/13.1](https://img.shields.io/static/v1?label=pypi&message=CUDA%2013.0%2F13.1&color=76B900)](https://pypi.org/project/mooncake-transfer-engine-cuda13)
+  [![PyPI Non-CUDA](https://img.shields.io/static/v1?label=pypi&message=non-CUDA&color=00BFFF)](https://pypi.org/project/mooncake-transfer-engine-non-cuda/)
+  [![PyPI NPU](https://img.shields.io/static/v1?label=pypi&message=NPU&color=F87171)](https://pypi.org/project/mooncake-transfer-engine-npu/)
 </div>
 <br/>
 
 Mooncake is the serving platform for  <a href="https://kimi.ai/"><img src="image/kimi.png" alt="icon" style="height: 16px; vertical-align: middle;"> Kimi</a>, a leading LLM service provided by <a href="https://www.moonshot.cn/"><img src="image/moonshot.jpg" alt="icon" style="height: 16px; vertical-align: middle;"> Moonshot AI</a>.
-Now both the Transfer Engine and Mooncake Store are open-sourced!
-This repository also hosts its technical report and the open-sourced traces.
+Under real workloads, Mooncake’s innovative architecture enables Kimi to handle 75% more requests while adhering to SLOs.
 
 <h2 id="updates">🔄 Updates</h2>
 
+- **May 7, 2026**: 🚀 [vLLM officially features Mooncake Store](https://vllm.ai/blog/mooncake-store) — a deep dive into how Mooncake's distributed KVCache engine supercharges vLLM inference with high-throughput, memory-efficient, cross-instance KV cache sharing!
+- **Apr 29, 2026**: SGLang introduces [RDMA-based P2P weight transfer for large-scale distributed RL](https://lmsys.org/blog/2026-04-29-p2p-update/) using Mooncake TransferEngine, achieving 7x faster weight updates for the 1T-parameter Kimi-K2 model (53s → 7.2s) with zero-copy RDMA transfer across thousands of GPUs.
 - **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.
+
+<details>
+<summary>More</summary>
+
 - **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.
@@ -63,127 +70,136 @@ This repository also hosts its technical report and the open-sourced traces.
  - **June 27, 2024**: We present a series of Chinese blogs with more discussions on <a href="https://zhuanlan.zhihu.com/p/705754254">zhihu 1</a>, <a href="https://zhuanlan.zhihu.com/p/705910725">2</a>, <a href="https://zhuanlan.zhihu.com/p/706204757">3</a>, <a href="https://zhuanlan.zhihu.com/p/707997501">4</a>, <a href="https://zhuanlan.zhihu.com/p/9461861451">5</a>, <a href="https://zhuanlan.zhihu.com/p/1939988652114580803">6</a>, <a href="https://zhuanlan.zhihu.com/p/1959366095443064318">7</a>.
  - **June 26, 2024**: Initial technical report release.
 
+</details>
 
 <h2 id="overview">🎉 Overview</h2>
 
-Mooncake features a KVCache-centric disaggregated architecture that separates the prefill and decoding clusters. It also leverages the underutilized CPU, DRAM, and SSD resources of the GPU cluster to implement a disaggregated KVCache pool.
-
-![architecture](image/architecture.png)
-
-The core of Mooncake is its KVCache-centric scheduler, which balances maximizing overall effective throughput while meeting latency-related Service Level Objectives (SLOs). Unlike traditional studies that assume all requests will be processed, Mooncake faces challenges in highly overloaded scenarios. To mitigate these, we developed a prediction-based early rejection policy. Experiments show that Mooncake excels in long-context scenarios. Compared to the baseline method, Mooncake can achieve up to a 525% increase in throughput in certain simulated scenarios while adhering to SLOs. Under real workloads, Mooncake’s innovative architecture enables <a href="https://kimi.ai/">Kimi</a> to handle 75% more requests.
-
-<h2 id="components">🧩 Components</h2>
-
 <!-- ![components](image/components.png) -->
-<img src=image/components.png width=74% />
+<div align="center">
+  <img src=image/components.png width=74% />
+</div>
 
-**Mooncake Core Component: Transfer Engine (TE)**
-The core of Mooncake is the Transfer Engine (TE), which provides a unified interface for batched data transfer across various storage devices and network links. Supporting multiple protocols including TCP, RDMA, CXL/shared-memory, and NVMe over Fabric (NVMe-of), TE is designed to enable fast and reliable data transfer for AI workloads. Compared to Gloo (used by Distributed PyTorch) and traditional TCP, TE achieves significantly lower I/O latency, making it a superior solution for efficient data transmission.
+Mooncake is an infrastructure project for large-scale LLM inference and training. It features a KV cache-centric disaggregated architecture that separates prefill and decode clusters, while leveraging otherwise underutilized CPU, DRAM, and SSD resources in GPU clusters to build a disaggregated KV cache pool.
 
-**P2P Store and Mooncake Store**
-Both P2P Store and Mooncake Store are built on the Transfer Engine and provide key/value caching for different scenarios. P2P Store focuses on sharing temporary objects (e.g., checkpoint files) across nodes in a cluster, preventing bandwidth saturation on a single machine. Mooncake Store, on the other hand, supports distributed pooled KVCache, specifically designed for XpYd disaggregation to enhance resource utilization and system performance.
-
-**Mooncake Integration with Leading LLM Inference Systems**
-Mooncake has been seamlessly integrated with several popular large language model (LLM) inference systems. Through collaboration with the vLLM and SGLang teams, Mooncake now officially supports prefill-decode disaggregation. By leveraging the high-efficiency communication capabilities of RDMA devices, Mooncake significantly improves inference efficiency in prefill-decode disaggregation scenarios, providing robust technical support for large-scale distributed inference tasks.
-In addition, Mooncake has been successfully integrated with SGLang's Hierarchical KV Caching, vLLM's prefill serving, and LMCache, augmenting KV cache management capabilities across large-scale inference scenarios.
-
-**Elastic Expert Parallelism Support**
-Mooncake adds elasticity and fault tolerance support for MoE model inference, enabling inference systems to remain responsive and recoverable in the event of GPU failures or changes in resource configuration. This functionality includes automatic faulty rank detection and can work with the EPLB module to dynamically route tokens to healthy ranks during inference.
-
-**Tensor-Centric Ecosystem**
-Mooncake establishes a full-stack, Tensor-oriented AI infrastructure where Tensors serve as the fundamental data carrier. The ecosystem spans from the Transfer Engine, which accelerates Tensor data movement across heterogeneous storage (DRAM/VRAM/NVMe), to the P2P Store and Mooncake Store for distributed management of Tensor objects (e.g., Checkpoints and KVCache), up to the Mooncake Backend enabling Tensor-based elastic distributed computing. This architecture is designed to maximize Tensor processing efficiency for large-scale model inference and training.
+Mooncake includes a high-performance Transfer Engine for low-latency data movement across heterogeneous networks and accelerators; Mooncake Store for distributed KV cache and model-weight management; and Mooncake EP & PG for elastic MoE serving. Deeply integrated with ecosystems such as SGLang and vLLM, Mooncake helps LLM systems improve cache reuse, reduce serving latency, and scale efficiently across multi-node clusters.
 
 <h2 id="show-cases">🔥 Show Cases</h2>
 
-### Use Transfer Engine Standalone ([Guide](https://kvcache-ai.github.io/Mooncake/design/transfer-engine/index.html))
+### Transfer Engine (TE)
 
-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).
+The core of Mooncake is the Transfer Engine (TE), a high-performance data transfer framework. TE offers a unified interface for batched data movement across diverse storage, network, and accelerator environments. By supporting multiple transport protocols, topology-aware routing, multi-NIC bandwidth aggregation, and automatic failover, TE delivers low-latency, scalable, and robust data transmission for distributed AI workloads. See the [Transfer Engine guide](https://kvcache-ai.github.io/Mooncake/design/transfer-engine/index.html) for details.
+
+<details>
+<summary>Highlights</summary>
 
-#### 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*.
 
-- **Topology aware path selection.** Transfer Engine can *select optimal devices* based on the location (NUMA affinity, etc.) of both source and destination.
+- **Topology-aware path selection.** Transfer Engine can *select optimal devices* based on the location (NUMA affinity, etc.) of both source and destination.
 
-- **More robust against temporary network errors.** Once transmission fails, Transfer Engine will try to use alternative paths for data delivery automatically.
+- **Robust against temporary network errors.** Once transmission fails, Transfer Engine will try to use alternative paths for data delivery automatically.
 
-#### Performance
-With 40 GB of data (equivalent to the size of the KVCache generated by 128k tokens in the LLaMA3-70B model), Mooncake Transfer Engine delivers up to **87 GB/s** and **190 GB/s** of bandwidth in 4×200 Gbps and 8×400 Gbps RoCE networks respectively, which are about **2.4x and 4.6x faster** than the TCP protocol.
+- **Superior performance at scale.** With 40 GB of data (equivalent to the size of the KVCache generated by 128k tokens in the LLaMA3-70B model), Mooncake Transfer Engine delivers up to **87 GB/s** and **190 GB/s** of bandwidth in 4×200 Gbps and 8×400 Gbps RoCE networks respectively, which are about **2.4x and 4.6x faster** than the TCP protocol.
 
 <!-- ![transfer-engine-performance.png](image/transfer-engine-performance.png) -->
 <img src=image/transfer-engine-performance.png width=75% />
 
-### P2P Store  ([Guide](https://kvcache-ai.github.io/Mooncake/design/p2p-store.html))
-P2P Store is built on the Transfer Engine and supports sharing temporary objects between peer nodes in a cluster. P2P Store is ideal for scenarios like checkpoint transfer, where data needs to be rapidly and efficiently shared across a cluster.
-**P2P Store has been used in the checkpoint transfer service of Moonshot AI.**
+- **Broad support for heterogeneous transports and accelerators.** Transfer Engine provides unified data transfer across diverse protocols, including TCP, RDMA, AWS EFA, NVMe-oF, NVLink, HIP, Barex, CXL, and Ascend-family transports. When built with the corresponding runtime, Transfer Engine can detect accelerator memory and select suitable transport paths for efficient data movement across CUDA, MUSA, HIP, MACA, Cambricon MLU, and Ascend-enabled environments. 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
-- **Decentralized architecture.** P2P Store leverages a pure client-side architecture with global metadata managed by the etcd service.
+- **Widely adopted across the LLM ecosystem.** TE is used in production inference stacks such as [SGLang](https://github.com/sgl-project/sglang), [vLLM](https://github.com/vllm-project/vllm), [TensorRT-LLM](https://github.com/NVIDIA/TensorRT-LLM), [vLLM-Ascend](https://github.com/vllm-project/vllm-ascend), [checkpoint-engine](https://github.com/MoonshotAI/checkpoint-engine), and [NIXL](https://github.com/ai-dynamo/nixl), among others, to efficiently transfer KV cache, embeddings, model weights, and other data.
 
-- **Efficient data distribution.** Designed to enhance the efficiency of large-scale data distribution, P2P Store *avoids bandwidth saturation* issues by allowing replicated nodes to share data directly. This reduces the CPU/RDMA NIC pressures of data providers (e.g., trainers).
+</details>
 
-<!-- #### Performance
-Thanks to the high performance of Transfer Engine, P2P Stores can also distribute objects with full utilization of *hardware incoming bandwidth* (e.g., A 25Gbps NIC was used in the following figure, and the throughput of get replica is about 3.1 GB/s). -->
+### Mooncake Store
 
-<!-- ![p2p-store.gif](image/p2p-store.gif) -->
+Mooncake Store is a high-performance distributed key-value cache storage engine designed for LLM inference. Built on the Transfer Engine, it stores and manages reusable KV caches and model weights across inference clusters, with support for efficient object storage, replication, eviction, and high-bandwidth data transfer. See the [Mooncake Store guide](https://kvcache-ai.github.io/Mooncake/design/mooncake-store.html) for details.
 
-### Mooncake Store ([Guide](https://kvcache-ai.github.io/Mooncake/design/mooncake-store.html))
-Mooncake Store is a distributed KVCache storage engine specialized for LLM inference based on Transfer Engine. It is the central component of the KVCache-centric disaggregated architecture. The goal of Mooncake Store is to store the reusable KV caches across various locations in an inference cluster. Mooncake Store has been supported in  [SGLang's Hierarchical KV Caching](https://lmsys.org/blog/2025-09-10-sglang-hicache/), [vLLM's prefill serving](https://docs.vllm.ai/en/latest/features/disagg_prefill.html) and is now integrated with [LMCache](https://kvcache-ai.github.io/Mooncake/getting_started/examples/lmcache-integration.html) to provide enhanced KVCache management capabilities.
+<details>
+<summary>Highlights</summary>
 
-#### Highlights
-- **Multi-replica support**: Mooncake Store supports storing multiple data replicas for the same object, effectively alleviating hotspots in access pressure.
+- **High bandwidth utilization.** Mooncake Store supports large-object striping, parallel I/O, and end-to-end zero-copy data transfer, fully utilizing aggregated bandwidth across multiple NICs.
 
-- **High bandwidth utilization**: Mooncake Store supports striping and parallel I/O transfer of large objects, fully utilizing multi-NIC aggregated bandwidth for high-speed data reads and writes.
+- **Multi-tier cache hierarchy**. Mooncake Store supports a multi-level cache design across DRAM and SSD/NVMe, enabling larger cache capacity.
 
-### SGLang Integration ([Guide](https://kvcache-ai.github.io/Mooncake/getting_started/examples/sglang-integration/hicache-integration-v1.html))
+- **Elastic and disaggregated storage.** Mooncake Store decouples KVCache storage from inference engines, allowing storage nodes to be dynamically added or removed while keeping cached data independent from engine restarts, upgrades, and scheduling decisions.
 
-SGLang officially supports Mooncake Store as a [HiCache storage backend](https://lmsys.org/blog/2025-09-10-sglang-hicache/). This integration enables scalable KV cache retention and high-performance access for large-scale LLM serving scenarios.
+- **Programmatic object management.** Mooncake Store allows applications to control object placement and lifecycle through per-object policies, including replica counts, preferred segments, soft pin, and hard pin. These controls help inference systems protect important KV caches and model weights while guiding replication, placement, and eviction behavior.
+
+- **Broad ecosystem adoption.** Mooncake Store is used across the LLM systems ecosystem as a high-performance distributed storage backend for KV caches, hidden states, and model weights. It supports integrations with [SGLang's Hierarchical KV Caching](https://lmsys.org/blog/2025-09-10-sglang-hicache/), [vLLM's prefill serving](https://docs.vllm.ai/en/latest/features/disagg_prefill.html), and [LMCache](https://kvcache-ai.github.io/Mooncake/getting_started/examples/lmcache-integration.html), and has been adopted by systems such as [TorchSpec](https://pytorch.org/blog/torchspec-speculative-decoding-training-at-scale/) and [TransferQueue](https://github.com/Ascend/TransferQueue) to decouple inference, training, and reinforcement-learning workloads through efficient state management and asynchronous data movement.
+
+</details>
+
+### Mooncake EP and Process Group (PG)
+
+Mooncake EP and Mooncake PG extend Mooncake from high-performance data movement to fault-tolerant distributed execution for large-scale MoE inference. Mooncake EP adapts DeepEP-style expert-parallel dispatch and combine operations with rank activeness awareness, while Mooncake PG provides a PyTorch distributed process-group backend with collective communication primitives that can detect failed ranks, report failures to upper layers, and recover ranks without restarting the entire inference service. See the [Mooncake EP & Backend guide](https://kvcache-ai.github.io/Mooncake/python-api-reference/ep-backend.html) for details.
+
+<details>
+<summary>Highlights</summary>
+
+- **Fault-tolerant expert parallelism.** Mooncake EP adds `active_ranks` awareness to expert-parallel dispatch and combine APIs, allowing MoE inference systems to route around failed ranks and continue serving with healthy experts.
+
+- **DeepEP-compatible programming model.** Mooncake EP keeps the API largely consistent with DeepEP's low-latency mode, making it easier for inference engines to adopt fault-tolerant expert parallelism without rewriting their MoE communication stack.
+
+- **PyTorch ProcessGroup integration.** Mooncake PG can be registered as a `torch.distributed` backend, enabling standard collective APIs such as `all_gather` while using Mooncake's communication and failure-reporting mechanisms underneath.
+
+- **Elastic rank recovery.** Mooncake PG exposes recovery-oriented primitives such as peer-state polling and rank recovery, allowing replacement processes to rejoin existing process groups and helping inference services recover from partial failures.
+
+- **SGLang integration for production MoE serving.** Mooncake's collective backend and expert-parallel kernels are integrated into SGLang to support fault-tolerant expert-parallel inference for large MoE models, including Elastic Expert Parallel serving scenarios.
+
+</details>
+
+### Tensor-Centric Ecosystem
+
+Mooncake establishes a full-stack, Tensor-oriented AI infrastructure where Tensors serve as the fundamental data carrier. The ecosystem spans from the Transfer Engine, which accelerates Tensor data movement across heterogeneous storage (DRAM/VRAM/NVMe), to Mooncake Store for distributed management of Tensor objects (e.g., KVCache and model weight), up to the Mooncake Backend enabling Tensor-based elastic distributed computing. This architecture is designed to maximize Tensor processing efficiency for large-scale model inference and training.
+
+### SGLang Integration ([Guide](https://kvcache-ai.github.io/Mooncake/getting_started/examples/sglang-integration/index.html))
+
+Mooncake is deeply integrated into [SGLang](https://github.com/sgl-project/sglang/) as a high-performance communication and storage backend. These integrations enable efficient KV cache transfer in PD-disaggregated serving, scalable multi-level KV caching through HiCache, fault-tolerant expert-parallel inference, high-performance multimodal pipeline data movement, and fast RDMA-based weight synchronization for large-scale RL training. Together, Mooncake and SGLang provide a production-oriented foundation for building elastic, high-throughput, and resource-efficient LLM and multimodal serving systems.
+
+<details>
+<summary>Details</summary>
+
+- **PD Disaggregated Serving:** SGLang officially supports Mooncake Transfer Engine as a backend for disaggregated serving and KV cache transfer, enabling prefill and decode workers to exchange KV cache data efficiently across devices and machines.
 
-#### Highlights
 - **Hierarchical KV Caching**: Mooncake Store serves as an external storage backend in SGLang's HiCache system, extending RadixAttention with multi-level KV cache storage across device, host, and remote storage layers.
-- **Flexible Cache Management**: Supports multiple cache policies including write-through, write-through-selective, and write-back modes, with intelligent prefetching strategies for optimal performance.
-- **Comprehensive Optimizations**: Features advanced data plane optimizations including page-first memory layout for improved I/O efficiency, zero-copy mechanisms for reduced memory overhead, GPU-assisted I/O kernels delivering fast CPU-GPU transfers, and layer-wise overlapping for concurrent KV cache loading while computation executes.
-- **Elastic Expert Parallel**: Mooncake's collective communication backend and expert parallel kernels are integrated into SGLang to enable fault-tolerant expert parallel inference ([sglang#11657](https://github.com/sgl-project/sglang/pull/11657)).
-- **Significant Performance Gains**: The multi-turn benchmark demonstrates substantial performance improvements over the non-HiCache setting. See our [benchmark report](https://kvcache-ai.github.io/Mooncake/performance/sglang-hicache-benchmark-results-v1.html) for more details.
-- **Community Feedback**: Effective KV caching significantly reduces TTFT by eliminating redundant and costly re-computation. Integrating SGLang HiCache with the Mooncake service enables scalable KV cache retention and high-performance access. In our evaluation, we tested the DeepSeek-R1-671B model under PD-disaggregated deployment using in-house online requests sampled from a general QA scenario. On average, cache hits achieved an 84% reduction in TTFT compared to full re-computation. – Ant Group
 
-### vLLM Integration ([Guide v0.2](https://kvcache-ai.github.io/Mooncake/getting_started/examples/vllm-integration/vllm-integration-v0.2.html))
-To optimize LLM inference, the vLLM community is working on supporting [disaggregated prefilling (PR 10502)](https://github.com/vllm-project/vllm/pull/10502). This feature allows separating the **prefill** phase from the **decode** phase in different processes. The vLLM uses `nccl` and `gloo` as the transport layer by default, but currently it cannot efficiently decouple both phases in different machines.
+- **Elastic Expert Parallel**: Mooncake's collective communication backend and expert parallel kernels are integrated into SGLang to enable fault-tolerant expert parallel inference ([Elastic EP](https://www.lmsys.org/blog/2026-03-25-eep-partial-failure-tolerance/)).
 
-We have implemented vLLM integration, which uses Transfer Engine as the network layer instead of `nccl` and `gloo`, to support **inter-node KVCache transfer** [(PR 10884)](https://github.com/vllm-project/vllm/pull/10884). Transfer Engine provides simpler interfaces and more efficient use of RDMA devices.
+- **Cloud-Native SGLang HiCache Deployment with RBG**: The [RBG](https://github.com/sgl-project/rbg) + SGLang HiCache + Mooncake integration provides a role-based, out-of-the-box cloud-native deployment solution that is elastic, scalable, and optimized for high-performance inference workloads.
 
-We will soon release the new vLLM integration based on Mooncake Store, which supports xPyD prefill/decode disaggregation.
+- **Encode-Prefill-Decode Disaggregation for Multimodal Serving**: SGLang introduces Encode-Prefill-Decode disaggregation with Mooncake as a transfer backend. This enables compute-intensive multimodal encoders, such as Vision Transformers, to be decoupled from language model workers while transferring large embeddings efficiently through Mooncake’s RDMA-based engine.
 
-**_Update[Dec 16, 2024]: Here is the latest vLLM Integration ([Guide v0.2](https://kvcache-ai.github.io/Mooncake/getting_started/examples/vllm-integration/vllm-integration-v0.2.html)) that is based on vLLM's main branch._**
+- **SGLang-Omni Multi-Stage Pipeline Data Transfer**: [SGLang-Omni](https://github.com/sgl-project/sglang-omni) integrates Mooncake as a relay backend for efficient cross-stage tensor and blob transfer in multimodal serving pipelines. This enables high-performance data movement between heterogeneous components such as thinker, talker, codec, and vocoder stages.
 
-#### Performance
-By supporting Topology Aware Path Selection and multi-card bandwidth aggregation, Mean TTFT of vLLM with Transfer Engine is up to 25% lower than traditional TCP-based transports.
-In the future, we will further improve TTFT through GPUDirect RDMA and zero-copy.
+- **RDMA-Based P2P Weight Transfer for Distributed RL**: SGLang adopts Mooncake TransferEngine for RDMA-based peer-to-peer weight transfer in large-scale distributed reinforcement learning. This enables zero-copy weight updates across thousands of GPUs and significantly accelerates synchronization for trillion-parameter models.
 
-| Backend/Setting                                         | Output Token Throughput (tok/s) | Total Token Throughput (tok/s) | Mean TTFT (ms) | Median TTFT (ms) | P99 TTFT (ms)|
-|---------------------------------------------------------|---------------------------------|--------------------------------|----------------|------------------|---------------|
-| Transfer Engine (RDMA) | 12.06                           | 2042.74                        | 1056.76        | 635.00           | 4006.59       |
-| TCP  | 12.05                           | 2041.13                        | 1414.05        | 766.23          | 6035.36       |
+</details>
 
-- Click [here](https://kvcache-ai.github.io/Mooncake/performance/vllm-benchmark-results-v0.2.html) to access detailed benchmark results.
+### vLLM Integration ([Guide](https://kvcache-ai.github.io/Mooncake/getting_started/examples/vllm-integration/index.html))
 
-**More advanced features are coming soon, so stay tuned!**
+Mooncake integrates with [vLLM](https://github.com/vllm-project/vllm) to accelerate large language model serving through high-performance KV cache transfer and distributed KV cache storage. The integration supports both disaggregated prefill-decode serving and cross-instance KV cache sharing, helping vLLM deployments reduce TTFT, improve cache reuse, and scale more efficiently across multi-node inference clusters.
 
-<h2 id="quick-start">🚀 Quick Start</h2>
+<details>
+<summary>Details</summary>
 
-### Before using Mooncake
+- **Disaggregated prefill-decode serving**: Mooncake enables vLLM to split prefill and decode workloads across different nodes. Through MooncakeConnector, vLLM transfers KV cache blocks from prefill workers to decode workers using Mooncake’s high-performance transfer engine, allowing prefill and decode resources to scale independently while keeping cross-node KV transfer overhead low.
 
-Mooncake is designed and optimized for high-speed RDMA networks. Though Mooncake supports TCP-only data transfer, we **strongly** recommend users to evaluate the functionality and performance of Mooncake with RDMA network support.
+- **Distributed KV cache pooling and sharing**: [Mooncake Store extends vLLM](https://vllm.ai/blog/2026-05-06-mooncake-store) from isolated per-instance KV caches to a shared, cluster-level KV cache pool. Through MooncakeStoreConnector, multiple vLLM instances can store, retrieve, and reuse KV cache blocks based on hash-based prefix caching, reducing redundant prefill computation and improving cache efficiency for workloads with repeated prefixes, especially agentic and multi-turn serving scenarios.
 
-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`.
+- **vLLM-Omni stage communication**: Mooncake also integrates with [vLLM-Omni](https://github.com/vllm-project/vllm-omni) through `MooncakeTransferEngineConnector` and `MooncakeStoreConnector`, enabling efficient cross-node data exchange between vLLM-Omni stages.
 
-### Use Python package
-The simplest way to use Mooncake Transfer Engine is using `pip`:
+</details>
 
-**For CUDA-enabled systems:**
+<h2 id="supported-hardware">🖥️ Supported Hardware</h2>
+
+Mooncake supports hardware backends across accelerator vendors, cloud fabrics, and standard datacenter interconnects, as listed below. See the [supported protocols](https://kvcache-ai.github.io/Mooncake/getting_started/supported-protocols.html) and [Transfer Engine design docs](https://kvcache-ai.github.io/Mooncake/design/transfer-engine/index.html) for details.
+
+| <img src="image/partners/nvidia_logo.png" width="120" alt="NVIDIA"/> | <img src="image/partners/huawei_logo.png" width="120" alt="Huawei"/> | <img src="image/partners/amd_logo.png" width="120" alt="AMD"/> | <img src="image/hardwares/cambricon_logo.png" width="120" alt="Cambricon"/> | <img src="image/partners/moore_thread_logo.jpg" width="120" alt="Moore Threads"/> | <img src="image/partners/aws-logo.png" width="120" alt="AWS"/> |
+| --- | --- | --- | --- | --- | --- |
+| <img src="image/hardwares/MetaX_logo.png" width="120" alt="MetaX"/> | <img src="image/hardwares/T-Head_logo.png" width="120" alt="T-Head"/> | <img src="image/partners/aliyun_logo.png" width="120" alt="Alibaba Cloud"/> | <img src="image/partners/sunrise_logo.png" width="120" alt="Sunrise"/> | <img src="image/partners/hygon_logo.png" width="120" alt="Hygon"/> | |
+
+<h2 id="quick-start">🚀 Getting Started</h2>
+
+Install Mooncake using `pip`. The `mooncake-transfer-engine` package includes Mooncake Transfer Engine, Mooncake Store, Mooncake EP and PG:
 
 - CUDA < 13.0
 ```bash
@@ -194,108 +210,84 @@ pip install mooncake-transfer-engine
 pip install mooncake-transfer-engine-cuda13
 ```
 
-**For non-CUDA systems:**
-```bash
-pip install mooncake-transfer-engine-non-cuda
+In addition to CUDA, Mooncake also supports other accelerator backends, along with flexible installation and deployment options. See the guides below for details:
+
+- [Quick Start](https://kvcache-ai.github.io/Mooncake/getting_started/quick-start.html)
+- [Build from Source](https://kvcache-ai.github.io/Mooncake/getting_started/build.html)
+- [Deployment Guide](https://kvcache-ai.github.io/Mooncake/deployment/mooncake-store-deployment-guide.html)
+
+
+### Skills for AI Assistants
+
+Mooncake ships a set of **built-in skills** under [`.claude/skills`](.claude/skills) — reusable, task-focused playbooks that an AI coding assistant (such as Claude Code) invokes automatically when your request matches, or that you can run as a slash command.
+
+<details>
+<summary>Details</summary>
+
+| Skill | Description |
+|-------|-------------|
+| `/mooncake-troubleshoot` | Diagnose Mooncake deployment and runtime issues (services, RDMA, env vars, logs). |
+| `/mooncake-ci-local` | Run pre-PR local validation via `scripts/run_ci_test.sh`. |
+| `/mooncake-api` | Work with the Mooncake Store, Transfer Engine, and EP/Backend Python APIs. |
+
+Install them without cloning the repository via the [Claude Code plugin marketplace](https://code.claude.com/docs/en/plugin-marketplaces):
+
+```text
+/plugin marketplace add kvcache-ai/Mooncake --sparse .claude-plugin
+/plugin install mooncake-troubleshoot@mooncake
+/plugin install mooncake-ci-local@mooncake
+/plugin install mooncake-api@mooncake
 ```
 
-> [!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.
+The `--sparse .claude-plugin` flag fetches only the marketplace catalog, and each plugin is published as a `git-subdir` source, so installing one fetches only that single skill directory — never the whole repo. If you are already working inside a Mooncake checkout, the skills under `.claude/skills/` load automatically with no setup.
 
-### Use Docker image
-Mooncake supports Docker-based deployment, see [Build Guide](https://kvcache-ai.github.io/Mooncake/getting_started/build.html) in detail.
+</details>
 
-To produce an image that compiles Mooncake from source, builds the wheel via `scripts/build_wheel.sh`, and installs that wheel inside the container, use `build-wheel.dockerfile`:
+<h2 id="trace">📦 Open Source Traces and Tools </h2>
 
-```bash
-docker build -f docker/mooncake.Dockerfile \
-  --build-arg PYTHON_VERSION=3.10 \
-  --build-arg EP_TORCH_VERSIONS="2.9.1" \
-  -t mooncake:from-source .
-```
+We open-source anonymized request traces containing request arrival times, input and output token counts, and remapped block hashes. These traces are designed to support reproducible simulation and evaluation of caching behavior while preserving user privacy. The released traces and related details are available in [FAST25-release](FAST25-release).
 
-The resulting image already has a virtual environment at `/opt/venv` with the freshly built wheel installed. Launch it with GPU/RDMA access as needed, for example:
-
-```bash
-docker run --gpus all --network host -it mooncake:from-source /bin/bash
-```
-
-> [!NOTE]
-> Make sure you build the image from the repository root so that Git metadata and submodules are available inside the build context.
-
-### Build and use binaries
-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.
-
-The build and installation steps are as follows:
-1. Retrieve source code from GitHub repo
-   ```bash
-   git clone https://github.com/kvcache-ai/Mooncake.git
-   cd Mooncake
-   ```
-
-2. Install dependencies
-   ```bash
-   bash dependencies.sh
-   ```
-
-3. Compile Mooncake and examples
-   ```bash
-   mkdir build
-   cd build
-   cmake ..
-   make -j
-   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>
-
-- [x] First release of Mooncake and integrate with latest vLLM
-- [ ] Share KV caches across multiple serving engines
-- [ ] User and developer documentation
-
-<h2 id="trace">📦 Open Source Trace</h2>
-
-```json
-{
-    "timestamp": 27482,
-    "input_length": 6955,
-    "output_length": 52,
-    "hash_ids": [46, 47, 48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 2353, 2354]
-}
-{
-    "timestamp": 30535,
-    "input_length": 6472,
-    "output_length": 26,
-    "hash_ids": [46, 47, 48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 2366]
-}
-```
-The above presents two samples from our trace dataset. The trace includes the timing of request arrivals, the number of input tokens, the number of output tokens, and the remapped block hash. To protect our customers' privacy, we applied several mechanisms to remove user-related information while preserving the dataset's utility for simulated evaluation. More descriptions of the trace (e.g., up to 50% cache hit ratio) can be found in Section 4 of the technical report.
-
-**_Update[Feb 21, 2025]: The updated [traces](FAST25-release/traces) used in our FAST'25 paper have been released! Please refer to the paper's appendix (found [here](FAST25-release/Mooncake-FAST25.pdf)) for more details._**
+Together with the released traces, we also provide two KV cache analysis tools: a [KV Cache Size Calculator](https://kvcache.ai/tools/kv-cache-size-calculator/) for calculating cache capacity across popular LLM model families, and a [KV Cache Hit Rate Simulator](https://kvcache.ai/tools/kv-cache-hit-rate-simulator/) for analyzing KV cache hit rates and planning cache capacity under different workloads and models. These tools help users better understand KV cache storage costs and caching effectiveness when analyzing or reproducing serving workloads. The tools are open-sourced [here](https://github.com/kvcache-ai/kvcache-blog).
 
 <h2 id="citation">📑 Citation</h2>
-Please kindly cite our paper if you find the paper or the traces are useful:
+Please kindly cite our papers if you find the papers or the traces are useful:
 
 ```bibtex
+@inproceedings{qin2025mooncake,
+  author    = {Ruoyu Qin and Zheming Li and Weiran He and Jialei Cui and Feng Ren and Mingxing Zhang and Yongwei Wu and Weimin Zheng and Xinran Xu},
+  title     = {Mooncake: Trading More Storage for Less Computation {\textemdash} A {KVCache-centric} Architecture for Serving {LLM} Chatbot},
+  booktitle = {23rd USENIX Conference on File and Storage Technologies (FAST 25)},
+  year      = {2025},
+  isbn      = {978-1-939133-45-8},
+  address   = {Santa Clara, CA},
+  pages     = {155--170},
+  url       = {https://www.usenix.org/conference/fast25/presentation/qin},
+  publisher = {USENIX Association},
+  month     = {feb},
+}
+```
+
+<details>
+<summary>More</summary>
+
+```bibtex
+@misc{ren2026tentdeclarativeslicespraying,
+  title     = {TENT: A Declarative Slice Spraying Engine for Performant and Resilient Data Movement in Disaggregated LLM Serving},
+  author    = {Feng Ren and Ruoyu Qin and Teng Ma and Shangming Cai and Zheng Liu and Chao Lei and Dejiang Zhu and Ke Yang and Zheming Li and Jialei Cui and Weixiao Huang and Yikai Zhao and Yineng Zhang and Hao Wu and Xiang Gao and Yuhao Fu and Jinlei Jiang and Yongwei Wu and Mingxing Zhang},
+  year      = {2026},
+  eprint    = {2604.00368},
+  archivePrefix = {arXiv},
+  primaryClass  = {cs.DC},
+  url       = {https://arxiv.org/abs/2604.00368},
+}
+
+@article{sun2026survivingpartialrankfailures,
+  title     = {Surviving Partial Rank Failures in Wide Expert-Parallel MoE Inference},
+  author    = {Xun Sun and Shaoyuan Chen and Pingchuan Ma and Yue Chen and Ziwei Yuan and Zhanhao Cao and Han Han and Shangming Cai and Teng Ma and Xuchun Shang and Xinpeng Zhao and Ke Yang and Junlin Wei and Lianzhi Lin and Yuji Liu and Feng Ren and Haoran Hu and Cheng Wan and Yingdi Shan and Yongwei Wu and Mingxing Zhang},
+  year      = {2026},
+  url       = {https://arxiv.org/abs/2605.10670},
+}
+
 @article{qin2025mooncake_tos,
   author    = {Qin Ruoyu and Li Zheming and He Weiran and Cui Jialei and Tang Heyi and Ren Feng and Ma Teng and Cai Shangming and Zhang Yineng and Zhang Mingxing and Wu Yongwei and Zheng Weimin and Xu Xinran},
   title     = {Mooncake: A KVCache-centric Disaggregated Architecture for LLM Serving},
@@ -310,19 +302,6 @@ Please kindly cite our paper if you find the paper or the traces are useful:
   keywords  = {Machine learning system, LLM serving, KVCache},
 }
 
-@inproceedings{qin2025mooncake,
-  author    = {Ruoyu Qin and Zheming Li and Weiran He and Jialei Cui and Feng Ren and Mingxing Zhang and Yongwei Wu and Weimin Zheng and Xinran Xu},
-  title     = {Mooncake: Trading More Storage for Less Computation {\textemdash} A {KVCache-centric} Architecture for Serving {LLM} Chatbot},
-  booktitle = {23rd USENIX Conference on File and Storage Technologies (FAST 25)},
-  year      = {2025},
-  isbn      = {978-1-939133-45-8},
-  address   = {Santa Clara, CA},
-  pages     = {155--170},
-  url       = {https://www.usenix.org/conference/fast25/presentation/qin},
-  publisher = {USENIX Association},
-  month     = {feb},
-}
-
 @article{qin2024mooncake_arxiv,
   title  = {Mooncake: A KVCache-centric Disaggregated Architecture for LLM Serving},
   author = {Ruoyu Qin and Zheming Li and Weiran He and Mingxing Zhang and Yongwei Wu and Weimin Zheng and Xinran Xu},
@@ -330,3 +309,5 @@ Please kindly cite our paper if you find the paper or the traces are useful:
   url    = {https://arxiv.org/abs/2407.00079},
 }
 ```
+
+</details>
diff --git a/benchmarks/storage_benchmark_v1/__init__.py b/benchmarks/storage_benchmark_v1/__init__.py
new file mode 100644
index 00000000..2973545a
--- /dev/null
+++ b/benchmarks/storage_benchmark_v1/__init__.py
@@ -0,0 +1,41 @@
+"""
+Mooncake KVCache Storage Benchmark Library
+"""
+
+__version__ = "2.0.0"
+
+from benchmark import (
+    main,
+    run_benchmark,
+    print_results,
+    StorageBenchmark,
+    TraceReplay,
+    KVCacheRequest,
+)
+from storage import Storage, KVKey, KVValue, DiskHashTable, SSDStorage
+from layout import (
+    KVLayout, MLALayout, KVEntry, MLA_MODEL_CONFIG, get_model_config, create_layout,
+)
+
+__all__ = [
+    # Main
+    'main',
+    'run_benchmark',
+    'print_results',
+    'StorageBenchmark',
+    'TraceReplay',
+    'KVCacheRequest',
+    # Storage
+    'Storage',
+    'KVKey',
+    'KVValue',
+    'DiskHashTable',
+    'SSDStorage',
+    # Layout
+    'KVLayout',
+    'MLALayout',
+    'KVEntry',
+    'MLA_MODEL_CONFIG',
+    'get_model_config',
+    'create_layout',
+]
diff --git a/benchmarks/storage_benchmark_v1/__main__.py b/benchmarks/storage_benchmark_v1/__main__.py
new file mode 100644
index 00000000..2ae8bc89
--- /dev/null
+++ b/benchmarks/storage_benchmark_v1/__main__.py
@@ -0,0 +1,11 @@
+"""
+Entry point for running the benchmark as a module
+
+Usage:
+    python -m storage_benchmark --model=glm5 --scenario=toolagent --max-requests=100
+"""
+
+from benchmark import main
+
+if __name__ == '__main__':
+    main()
diff --git a/benchmarks/storage_benchmark_v1/benchmark.py b/benchmarks/storage_benchmark_v1/benchmark.py
new file mode 100644
index 00000000..3f5ba5ed
--- /dev/null
+++ b/benchmarks/storage_benchmark_v1/benchmark.py
@@ -0,0 +1,486 @@
+#!/usr/bin/env python3
+"""
+Mooncake KVCache Storage Benchmark
+
+Complete benchmark tool with CLI interface.
+"""
+
+import argparse
+import json
+import sys
+import time
+import statistics
+import signal
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Iterator, List, Dict, Any
+
+from storage import DiskHashTable
+from layout import get_model_config, create_layout
+
+
+# ============================================================================
+# Data Structures
+# ============================================================================
+
+@dataclass
+class KVCacheRequest:
+    """KVCache request from trace"""
+    timestamp: float
+    hash_ids: List[int]
+    input_length: int
+    output_length: int
+
+
+# ============================================================================
+# Trace Replay
+# ============================================================================
+
+class TraceReplay:
+    """Trace replay handler"""
+
+    def __init__(self, trace_path: str):
+        self.trace_path = trace_path
+
+    def load_all(self) -> List[KVCacheRequest]:
+        """Load all requests from trace file"""
+        requests = []
+        with open(self.trace_path, 'r', encoding='utf-8') as f:
+            for line in f:
+                line = line.strip()
+                if line:
+                    req = json.loads(line)
+                    requests.append(KVCacheRequest(
+                        timestamp=req.get('timestamp', 0),
+                        hash_ids=req.get('hash_ids', []),
+                        input_length=req.get('input_length', 0),
+                        output_length=req.get('output_length', 0),
+                    ))
+        return requests
+
+
+# ============================================================================
+# Storage Benchmark
+# ============================================================================
+
+class StorageBenchmark:
+    """KVCache storage benchmark
+
+    Processes KVCache requests using layout-generated access patterns.
+    """
+
+    def __init__(self, storage_dir: str, model_config: dict,
+                 page_size_tokens: int = 512,
+                 max_pages: int = 100000,
+                 fsync_mode: str = 'none', fsync_batch_size: int = 100):
+        """Initialize benchmark
+
+        Args:
+            storage_dir: Directory for storage files
+            model_config: Model configuration dict
+            page_size_tokens: Tokens per page (default: 512)
+            max_pages: Maximum number of pages
+            fsync_mode: When to fsync ('none', 'batch', 'always', 'end')
+            fsync_batch_size: Number of writes between fsync in batch mode
+        """
+        self.model_config = model_config
+        self.layout = create_layout(model_config, page_size_tokens)
+        self.page_size_bytes = self.layout.value_size_bytes
+
+        # Initialize storage
+        self.storage = DiskHashTable(
+            storage_dir=storage_dir,
+            page_size=self.page_size_bytes,
+            max_pages=max_pages,
+            fsync_mode=fsync_mode,
+            fsync_batch_size=fsync_batch_size
+        )
+
+        # Statistics
+        self.stats = {
+            'total_requests': 0,
+            'total_tokens': 0,
+            'read_pages': 0,
+            'write_pages': 0,
+            'page_hits': 0,
+            'request_latencies_ms': [],
+        }
+
+    def process_request(self, req: KVCacheRequest) -> float:
+        """Process a KVCache request
+
+        Args:
+            req: KVCache request
+
+        Returns:
+            Total latency in milliseconds
+        """
+        self.stats['total_requests'] += 1
+        self.stats['total_tokens'] += req.input_length + req.output_length
+
+        total_latency = 0.0
+
+        # Process each access requirement from layout
+        for access in self.layout.get_operations(req):
+            if self.storage.exists(access.page_id):
+                # Page exists, perform READ
+                total_latency += self.storage.read(
+                    access.page_id,
+                    offset_in_page=access.offset_in_page,
+                    length=access.length
+                )
+                self.stats['read_pages'] += 1
+                self.stats['page_hits'] += 1
+            else:
+                # Page doesn't exist, perform WRITE
+                total_latency += self.storage.write(
+                    access.page_id,
+                    offset_in_page=access.offset_in_page,
+                    length=access.length
+                )
+                self.stats['write_pages'] += 1
+
+        latency_ms = total_latency if total_latency > 0 else 0.0
+        if latency_ms > 0:
+            self.stats['request_latencies_ms'].append(latency_ms)
+        return latency_ms
+
+    def get_stats(self) -> Dict:
+        """Get statistics"""
+        storage_stats = self.storage.get_stats()
+        request_latencies = self.stats['request_latencies_ms']
+
+        if request_latencies:
+            sorted_latencies = sorted(request_latencies)
+            n = len(sorted_latencies)
+
+            def get_percentile(p: float) -> float:
+                idx = int(n * p)
+                return sorted_latencies[idx] if idx < n else sorted_latencies[-1]
+
+            latency_stats = {
+                'avg_ms': statistics.mean(request_latencies),
+                'p50_ms': sorted_latencies[n // 2],
+                'p95_ms': get_percentile(0.95),
+                'p99_ms': get_percentile(0.99),
+            }
+        else:
+            latency_stats = {'avg_ms': 0, 'p50_ms': 0, 'p95_ms': 0, 'p99_ms': 0}
+
+        total_pages = self.stats['read_pages'] + self.stats['write_pages']
+
+        return {
+            'total_requests': self.stats['total_requests'],
+            'total_tokens': self.stats['total_tokens'],
+            'total_pages': total_pages,
+            'read_pages': self.stats['read_pages'],
+            'write_pages': self.stats['write_pages'],
+            'page_hits': self.stats['page_hits'],
+            'page_hit_rate': self.stats['read_pages'] / total_pages if total_pages > 0 else 0,
+            'write_ratio': self.stats['write_pages'] / total_pages if total_pages > 0 else 0,
+            'latency': latency_stats,
+            'storage': storage_stats,
+        }
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self.close()
+        return False
+
+    def close(self, force_sync: bool = True):
+        self.storage.close(force_sync=force_sync)
+
+
+# ============================================================================
+# Benchmark Runner
+# ============================================================================
+
+def get_max_page_id(requests: List[KVCacheRequest]) -> int:
+    max_id = 0
+    for req in requests:
+        if req.hash_ids:
+            max_id = max(max_id, max(req.hash_ids))
+    return max_id
+
+
+def run_benchmark(trace_path: str, storage_dir: str, model_config: dict,
+                  max_requests: int = None, max_pages: int = None,
+                  page_size_tokens: int = 512,
+                  fsync_mode: str = 'none', fsync_batch_size: int = 100) -> Dict:
+    """Run benchmark
+
+    Args:
+        trace_path: Trace file path
+        storage_dir: Storage directory
+        model_config: Model configuration
+        max_requests: Maximum number of requests (None = all)
+        max_pages: Maximum number of pages (None = auto-calculate)
+        page_size_tokens: Tokens per page
+        fsync_mode: When to fsync
+        fsync_batch_size: Number of writes between fsync
+
+    Returns:
+        Benchmark results dictionary
+    """
+    print(f"\n{'='*80}")
+    print(f"Running: {Path(trace_path).name}")
+    print(f"Model: {model_config['name']}")
+    print(f"Layers: {model_config['num_layers']}")
+    print(f"Page size: {page_size_tokens} tokens")
+    print(f"{'='*80}")
+
+    # Load trace
+    replay = TraceReplay(trace_path)
+    requests = replay.load_all()
+
+    if max_requests:
+        requests = requests[:max_requests]
+
+    print(f"Loaded {len(requests)} requests")
+
+    # Find max page_id from trace
+    max_page_id = get_max_page_id(requests)
+    max_pages_needed = max_page_id + 1  # page_id is 0-based
+
+    # Create layout to get page size
+    layout = create_layout(model_config, page_size_tokens)
+    page_size_bytes = layout.value_size_bytes
+
+    # Determine max_pages
+    if max_pages is None:
+        # Use max page_id from trace
+        max_pages = max_pages_needed
+        # Round up to next thousand for cleaner numbers
+        max_pages = ((max_pages + 999) // 1000) * 1000
+    else:
+        # User specified max_pages
+        pass
+
+    max_size_gb = max_pages * page_size_bytes / (1024**3)
+    trace_size_gb = max_pages_needed * page_size_bytes / (1024**3)
+
+    print(f"\n[Storage Configuration]")
+    print(f"  Max page_id in trace:             {max_page_id:,}")
+    print(f"  Pages needed (trace):             {max_pages_needed:,}")
+    print(f"  Trace storage size:               {trace_size_gb:.2f} GB")
+    print(f"  Max pages configured:             {max_pages:,}")
+    print(f"  Max storage available:            {max_size_gb:.2f} GB")
+
+    if max_pages_needed > max_pages:
+        shortfall = max_pages_needed - max_pages
+        shortfall_gb = shortfall * page_size_bytes / (1024**3)
+        compression_ratio = max_pages / max_pages_needed
+        print(f"\n  ⚠️  Storage insufficient: {shortfall:,} pages shortfall ({shortfall_gb:.2f} GB)")
+        print(f"  ⚠️  Consider increasing --max-pages to at least {max_pages_needed:,} for full simulation")
+    else:
+        surplus = max_pages - max_pages_needed
+        surplus_pct = (surplus / max_pages) * 100 if max_pages > 0 else 0
+        print(f"  ✓ Direct mapping: all {max_pages_needed:,} logical pages uniquely mapped")
+
+    # Run benchmark
+    with StorageBenchmark(
+        storage_dir=storage_dir,
+        model_config=model_config,
+        page_size_tokens=page_size_tokens,
+        max_pages=max_pages,
+        fsync_mode=fsync_mode,
+        fsync_batch_size=fsync_batch_size
+    ) as benchmark:
+        start_time = time.perf_counter()
+        try:
+            for i, req in enumerate(requests):
+                benchmark.process_request(req)
+                # Print progress for each request
+                elapsed = time.perf_counter() - start_time
+                qps = (i + 1) / elapsed if elapsed > 0 else 0
+                stats = benchmark.get_stats()
+                storage = stats.get('storage', {})
+                read_latency = storage.get('read', {}).get('avg_ms', 0)
+                write_latency = storage.get('write', {}).get('avg_ms', 0)
+                read_mb = storage.get('read', {}).get('mb', 0)
+                write_mb = storage.get('write', {}).get('mb', 0)
+                read_time = storage.get('read', {}).get('time_s', 0)
+                write_time = storage.get('write', {}).get('time_s', 0)
+                read_mbps = read_mb / read_time if read_time > 0 else 0
+                write_mbps = write_mb / write_time if write_time > 0 else 0
+                print(f"  [{i+1:5d}/{len(requests)}] ids={len(req.hash_ids):3d} "
+                      f"tokens={req.input_length+req.output_length:6d} | "
+                      f"QPS={qps:7.2f} | "
+                      f"R={stats['read_pages']:6d} ({read_latency:6.2f}ms, {read_mbps:6.1f}MB/s) | "
+                      f"W={stats['write_pages']:6d} ({write_latency:6.2f}ms, {write_mbps:6.1f}MB/s)")
+        except KeyboardInterrupt:
+            print(f"\n\n{'='*80}")
+            print(f"Interrupted! Showing partial results:")
+            print(f"{'='*80}")
+            elapsed = time.perf_counter() - start_time
+            stats = benchmark.get_stats()
+            print_results([{
+                'trace_file': Path(trace_path).name,
+                'total_requests': i + 1,
+                'io_time_s': elapsed,
+                'requests_per_second': (i + 1) / elapsed if elapsed > 0 else 0,
+                'model': model_config['name'],
+                'fsync_mode': fsync_mode,
+                **stats,
+            }])
+            sys.exit(0)
+
+        elapsed = time.perf_counter() - start_time
+        stats = benchmark.get_stats()
+
+    return {
+        'trace_file': Path(trace_path).name,
+        'total_requests': len(requests),
+        'io_time_s': elapsed,
+        'requests_per_second': len(requests) / elapsed if elapsed > 0 else 0,
+        'model': model_config['name'],
+        'fsync_mode': fsync_mode,
+        **stats,
+    }
+
+
+# ============================================================================
+# Output Formatting
+# ============================================================================
+
+def format_storage_stats(stats: Dict, title: str = "Storage"):
+    """Format storage statistics with clear read/write separation"""
+    storage = stats.get('storage', {})
+    read_stats = storage.get('read', {})
+    write_stats = storage.get('write', {})
+
+    output = []
+    output.append(f"\n[{title}]")
+
+    # General info
+    output.append(f"\n[General]")
+    output.append(f"  Model:            {stats.get('model', 'N/A')}")
+    output.append(f"  Requests:         {stats.get('total_requests', 0):,}")
+    output.append(f"  Tokens:           {stats.get('total_tokens', 0):,}")
+    output.append(f"  Total I/O Time:    {stats.get('io_time_s', 0):.3f} s")
+    output.append(f"  QPS:              {stats.get('requests_per_second', 0):.2f}")
+    output.append(f"  Hit Rate:         {stats.get('page_hit_rate', 0):.2%}")
+
+    # Read Stats
+    output.append(f"\n[Read Operations]")
+    output.append(f"  Count:            {read_stats.get('count', 0):,}")
+    output.append(f"  Data Volume:      {read_stats.get('mb', 0):.2f} MB")
+    read_time = read_stats.get('time_s', 0)
+    read_mbps = read_stats.get('mb', 0) / read_time if read_time > 0 else 0
+    output.append(f"  Total Time:       {read_time:.3f} s")
+    output.append(f"  Bandwidth:        {read_mbps:.2f} MB/s")
+    output.append(f"  Latency:")
+    output.append(f"    Avg:            {read_stats.get('avg_ms', 0):.3f} ms")
+    output.append(f"    P50:            {read_stats.get('p50_ms', 0):.3f} ms")
+    output.append(f"    P95:            {read_stats.get('p95_ms', 0):.3f} ms")
+    output.append(f"    P99:            {read_stats.get('p99_ms', 0):.3f} ms")
+
+    # Write Stats
+    output.append(f"\n[Write Operations]")
+    output.append(f"  Count:            {write_stats.get('count', 0):,}")
+    output.append(f"  Data Volume:      {write_stats.get('mb', 0):.2f} MB")
+    write_time = write_stats.get('time_s', 0)
+    write_mbps = write_stats.get('mb', 0) / write_time if write_time > 0 else 0
+    output.append(f"  Total Time:       {write_time:.3f} s")
+    output.append(f"  Bandwidth:        {write_mbps:.2f} MB/s")
+    output.append(f"  Latency:")
+    output.append(f"    Avg:            {write_stats.get('avg_ms', 0):.3f} ms")
+    output.append(f"    P50:            {write_stats.get('p50_ms', 0):.3f} ms")
+    output.append(f"    P95:            {write_stats.get('p95_ms', 0):.3f} ms")
+    output.append(f"    P99:            {write_stats.get('p99_ms', 0):.3f} ms")
+
+    # Storage Info
+    output.append(f"\n[Storage Info]")
+    output.append(f"  Max Pages:        {storage.get('max_pages', 0):,}")
+    output.append(f"  Written Pages:    {storage.get('written_pages', 0):,}")
+    output.append(f"  Sync Count:       {storage.get('sync_count', 0):,}")
+
+    return "\n".join(output)
+
+
+def print_results(results: List[Dict]):
+    """Print 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(format_storage_stats(r))
+
+
+# ============================================================================
+# CLI Entry Point
+# ============================================================================
+
+def main():
+    """Main entry point"""
+    parser = argparse.ArgumentParser(
+        description='Mooncake KVCache Storage Benchmark',
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+    )
+
+    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, default='glm5', choices=['glm5', 'kimi-k2.6'],
+                       help='Model preset')
+    parser.add_argument('--page-size-tokens', type=int, default=512,
+                       help='Page size in tokens (default: 512)')
+    parser.add_argument('--max-requests', type=int, default=None,
+                       help='Maximum number of requests')
+    parser.add_argument('--max-pages', type=int, default=2000,
+                       help='Maximum number of pages')
+    parser.add_argument('--fsync-mode', type=str, choices=['batch', 'always', 'end', 'none'],
+                       default='none', help='When to fsync')
+    parser.add_argument('--fsync-batch-size', type=int, default=100,
+                       help='Number of writes between fsync')
+
+    args = parser.parse_args()
+
+    print(f"\n{'='*80}")
+    print(f"{'Mooncake KVCache Storage Benchmark':^80}")
+    print(f"{'='*80}")
+
+    model_config = get_model_config(args.model)
+    print(f"Model: {args.model} ({model_config['num_layers']} layers)")
+
+    # Determine 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),
+                model_config,
+                args.max_requests,
+                args.max_pages,
+                args.page_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)
+    else:
+        print("Error: No trace files were successfully processed.", file=sys.stderr)
+        sys.exit(1)
+
+if __name__ == '__main__':
+    main()
diff --git a/benchmarks/storage_benchmark_v1/doc/README.md b/benchmarks/storage_benchmark_v1/doc/README.md
new file mode 100644
index 00000000..fa490aad
--- /dev/null
+++ b/benchmarks/storage_benchmark_v1/doc/README.md
@@ -0,0 +1,121 @@
+# KVCache Storage Benchmark v1
+
+## Overview
+
+The KVCache Storage Benchmark is a tool for evaluating storage performance of KVCache workloads. It simulates real-world cache access patterns using trace replay and measures storage I/O performance with detailed statistics.
+
+## Usage
+
+### Basic Usage
+
+```bash
+cd benchmarks/storage_benchmark_v1
+python benchmark.py --scenario conversation \
+                    --trace-dir /path/to/Mooncake/FAST25-release/traces \
+                    --storage-dir /path/to/test/drive
+```
+
+### Command Line Options
+
+| Option | Default | Description |
+|--------|---------|-------------|
+| `--trace-dir` | `../../FAST25-release/traces` | Directory containing trace files |
+| `--scenario` | `toolagent` | Test scenario: `conversation`, `synthetic`, `toolagent`, or `all` |
+| `--storage-dir` | `/tmp/mooncake_bench` | Directory for storage files |
+| `--model` | `glm5` | Model preset: `glm5` or `kimi-k2.6` |
+| `--page-size-tokens` | `512` | Page size in tokens |
+| `--max-requests` | `None` | Maximum number of requests to process |
+| `--max-pages` | `2000` | Maximum number of pages (creates modulo mapping if trace is larger) |
+| `--fsync-mode` | `none` | When to fsync: `none`, `batch`, `always`, or `end` |
+| `--fsync-batch-size` | `100` | Number of writes between fsync in batch mode |
+
+## Output Format
+
+### Progress Output
+
+During execution, each request displays real-time statistics:
+
+```
+[    10/12031] ids= 35 tokens= 18060 | QPS=   2.45 | R=    36 ( 22.01ms, 2435.2MB/s) | W=   963 ( 19.35ms, 2770.1MB/s)
+```
+
+Fields:
+- `[N/Total]`: Current request progress
+- `ids=N`: Number of hash_ids in this request
+- `tokens=N`: Total tokens (input + output)
+- `QPS=X`: Queries per second (overall)
+- `R=N (latency, bandwidth)`: Read count, average latency, bandwidth
+- `W=N (latency, bandwidth)`: Write count, average latency, bandwidth
+
+### Final Results
+
+```
+================================================================================
+  [1/1] toolagent_trace.jsonl
+================================================================================
+
+[General]
+  Model:            glm5
+  Requests:         12031
+  Tokens:           123456789
+  Total I/O Time:   245.123 s
+  QPS:              49.07
+  Hit Rate:         3.25%
+
+[Read Operations]
+  Count:            390
+  Data Volume:      20919.62 MB
+  Total Time:       8.590 s
+  Bandwidth:         2435.67 MB/s
+  Latency:
+    Avg:             22.032 ms
+    P50:             21.456 ms
+    P95:             28.912 ms
+    P99:             35.234 ms
+
+[Write Operations]
+  Count:            11641
+  Data Volume:      654321.45 MB
+  Total Time:       236.533 s
+  Bandwidth:         2765.89 MB/s
+  Latency:
+    Avg:             20.312 ms
+    P50:             19.876 ms
+    P95:             25.123 ms
+    P99:             31.456 ms
+
+[Storage Info]
+  Max Pages:        2000
+  Written Pages:    2000
+  Sync Count:       0
+```
+
+## Modulo Mapping
+
+When the trace requires more pages than `--max-pages`, modulo mapping is enabled:
+
+```
+physical_page_id = logical_page_id % max_pages
+```
+
+This allows simulating large traces (millions of pages) with limited storage (thousands of pages). The first-hit/read-then-write logic is preserved by tracking written logical pages in memory.
+
+**Example**: With `--max-pages 2000`, logical page IDs 0-1999 map directly to physical pages 0-1999. Logical page 2000 maps to physical page 0, logical page 2001 maps to physical page 1, etc.
+
+## Graceful Interruption
+
+Press `Ctrl-C` at any time to stop the benchmark and view partial results. The output will display all statistics collected up to the interruption point, using the same format as final results.
+
+## Troubleshooting
+
+### Insufficient Storage Warning
+```
+⚠️  Modulo mapping ENABLED (limited storage)
+⚠️  Storage insufficient: 14,257,620 pages shortfall (762.34 GB)
+```
+Increase `--max-pages` to reduce modulo mapping effects.
+
+### Low Bandwidth or High Latency
+- Check fsync mode (`--fsync-mode none` for best performance)
+- Verify disk performance with `fio` or `dd`
+- Verify storage device health
diff --git a/benchmarks/storage_benchmark_v1/layout/__init__.py b/benchmarks/storage_benchmark_v1/layout/__init__.py
new file mode 100644
index 00000000..5d00a789
--- /dev/null
+++ b/benchmarks/storage_benchmark_v1/layout/__init__.py
@@ -0,0 +1,17 @@
+"""
+KVCache Layout Module
+
+Provides layout interface and implementations for different model architectures.
+"""
+
+from .interface import KVLayout, StorageAccess
+from .mla import MLALayout, MLA_MODEL_CONFIG, get_model_config, create_layout
+
+__all__ = [
+    'KVLayout',
+    'StorageAccess',
+    'MLALayout',
+    'MLA_MODEL_CONFIG',
+    'get_model_config',
+    'create_layout',
+]
diff --git a/benchmarks/storage_benchmark_v1/layout/interface.py b/benchmarks/storage_benchmark_v1/layout/interface.py
new file mode 100644
index 00000000..569bc104
--- /dev/null
+++ b/benchmarks/storage_benchmark_v1/layout/interface.py
@@ -0,0 +1,55 @@
+"""
+KVCache Layout Interface
+
+Defines the abstract interface for KVCache storage layouts.
+Layout layer converts normalized requests into storage access requirements.
+"""
+
+from abc import ABC, abstractmethod
+from typing import Iterator, Any
+from dataclasses import dataclass
+
+
+@dataclass
+class StorageAccess:
+    """Storage access requirement
+
+    Represents a need to access a page. Whether to READ or WRITE is determined
+    by the upper layer based on whether the page already exists.
+    """
+    page_id: int             # Page ID (hash_id from trace)
+    offset_in_page: int = 0  # Offset within the page (default: 0)
+    length: int = None       # Number of bytes (default: entire page)
+
+    def __repr__(self):
+        if self.offset_in_page == 0 and self.length is None:
+            return f"Access(page_id={self.page_id})"
+        else:
+            return f"Access(page_id={self.page_id}, offset={self.offset_in_page}, length={self.length})"
+
+
+class KVLayout(ABC):
+    """Abstract interface for KVCache storage layout
+
+    Converts normalized KVCache requests into storage access requirements.
+
+    Request format:
+        - hash_ids: List[int] - chunk/page identifiers
+        - input_length: int - input token count
+        - output_length: int - output token count
+
+    Output:
+        - Iterator of StorageAccess (page access requirements)
+    """
+
+    @abstractmethod
+    def get_operations(self, request: Any) -> Iterator[StorageAccess]:
+        """Generate storage access requirements for a request
+
+        Args:
+            request: KVCache request with hash_ids, input_length, output_length
+
+        Yields:
+            StorageAccess: Page access requirements (READ vs WRITE decided by upper layer)
+        """
+        pass
diff --git a/benchmarks/storage_benchmark_v1/layout/mla.py b/benchmarks/storage_benchmark_v1/layout/mla.py
new file mode 100644
index 00000000..6d7a3864
--- /dev/null
+++ b/benchmarks/storage_benchmark_v1/layout/mla.py
@@ -0,0 +1,174 @@
+"""
+MLA (Multi-head Latent Attention) KVCache Layout
+
+Implements the KVLayout interface for MLA architecture.
+"""
+
+from typing import Iterator, Any
+
+from .interface import KVLayout, StorageAccess
+
+
+# ============================================================================
+# MLA Model Configurations
+# ============================================================================
+
+# MLA Model configurations
+# Source: https://kvcache.ai/tools/kv-cache-calculator/
+# MLA architecture: hash_id -> {layer_0: [pages], layer_1: [pages], ...}
+MLA_MODEL_CONFIG = {
+    # GLM-5: 78 layers, 64 tokens/page
+    # KV: 78 layers × 64 tokens × (512+64+128) × 2 = 90,112 bytes/page
+    # Per token: 1,408 bytes
+    "glm5": {
+        "name": "glm5",
+        "num_layers": 78,
+        "kv_lora_rank": 512,
+        "qk_rope_head_dim": 64,
+        "index_head_dim": 128,
+        "kv_precision_bytes": 2,  # BF16
+        "indexer_precision_bytes": 2,  # BF16
+    },
+
+    # Kimi-K2.6: 61 layers, 64 tokens/page
+    # KV: 61 layers × 64 tokens × (512+64) × 2 = 73,728 bytes/page
+    # Per token: 1,152 bytes
+    "kimi-k2.6": {
+        "name": "kimi-k2.6",
+        "num_layers": 61,
+        "kv_lora_rank": 512,
+        "qk_rope_head_dim": 64,
+        "index_head_dim": 0,  # Kimi doesn't use separate indexer
+        "kv_precision_bytes": 2,  # BF16
+        "indexer_precision_bytes": 0,
+    },
+}
+
+
+class MLALayout(KVLayout):
+    """MLA (Multi-head Latent Attention) KVCache layout
+
+    MLA Architecture:
+    - Each hash_id corresponds to a 512-token chunk
+    - Each hash_id maps to ONE complete entry containing all layers
+    - Entry contains KV + Indexer data for all layers for 512 tokens
+    - Value size is fixed per hash_id (includes all layers)
+
+    Value Size Calculation (per hash_id entry):
+        per_layer_size = 512 × (kv_lora_rank + qk_rope_head_dim + index_head_dim) × precision_bytes
+        value_size = per_layer_size × num_layers
+
+    For GLM-5 with 512 tokens per entry:
+        per_layer_size = 512 × (512 + 64 + 128) × 2 = 720,896 bytes
+        value_size = 720,896 × 78 = 56,229,888 bytes = 53.6 MiB
+
+    Key Pattern: hash_id → single entry (all layers included)
+    Total Keys = len(hash_ids)
+
+    Used in: GLM-5, Kimi-K2.6
+    """
+
+    def __init__(self, num_layers: int, kv_lora_rank: int, qk_rope_head_dim: int,
+                 index_head_dim: int, precision_bytes: int, page_size_tokens: int = 512):
+        """Initialize MLA layout
+
+        Args:
+            num_layers: Number of transformer layers
+            kv_lora_rank: KV LoRA rank dimension
+            qk_rope_head_dim: QK rope head dimension
+            index_head_dim: Indexer head dimension
+            precision_bytes: Precision in bytes (BF16=2, INT8=1, INT4=0.5)
+            page_size_tokens: Tokens per page (default: 512)
+        """
+        self.num_layers = num_layers
+        self.kv_lora_rank = kv_lora_rank
+        self.qk_rope_head_dim = qk_rope_head_dim
+        self.index_head_dim = index_head_dim
+        self.precision_bytes = precision_bytes
+        self.page_size_tokens = page_size_tokens
+
+        # Calculate fixed value size per entry (per hash_id)
+        # Each entry contains KV + Indexer for all layers for page_size_tokens
+        # Per layer: page_size_tokens × (kv_lora_rank + qk_rope_head_dim + index_head_dim) × precision_bytes
+        # Total: per_layer_size × num_layers
+        per_layer_size = page_size_tokens * (kv_lora_rank + qk_rope_head_dim + index_head_dim) * precision_bytes
+        self.value_size_bytes = per_layer_size * num_layers
+
+        # Store page_size for backward compatibility
+        self.page_size = self.value_size_bytes
+
+    def get_operations(self, request: Any) -> Iterator[StorageAccess]:
+        """Generate storage access requirements for a request
+
+        For MLA architecture:
+        - Each hash_id corresponds to one complete page (512 tokens, all layers)
+        - Generate one access requirement per hash_id
+
+        Args:
+            request: KVCache request with hash_ids, input_length, output_length
+
+        Yields:
+            StorageAccess: Page access requirements
+        """
+        for hash_id in request.hash_ids:
+            yield StorageAccess(
+                page_id=hash_id,
+                offset_in_page=0,
+                length=self.value_size_bytes
+            )
+
+
+# ============================================================================
+# Utility Functions
+# ============================================================================
+
+def get_model_config(model_name: str) -> dict:
+    """Get MLA model configuration by name
+
+    Args:
+        model_name: Model identifier (e.g., 'glm5', 'kimi-k2.6')
+
+    Returns:
+        dict: Model configuration
+
+    Raises:
+        KeyError: If model name is not found
+    """
+    if model_name not in MLA_MODEL_CONFIG:
+        available = ", ".join(MLA_MODEL_CONFIG.keys())
+        raise KeyError(f"Unknown model: {model_name}. Available: {available}")
+    return MLA_MODEL_CONFIG[model_name].copy()
+
+
+def create_layout(model_config: dict, page_size_tokens: int = 64) -> MLALayout:
+    """Create an MLALayout instance from model configuration
+
+    Args:
+        model_config: Model configuration dictionary with fields:
+            - num_layers: Number of transformer layers
+            - kv_lora_rank: KV LoRA rank dimension
+            - qk_rope_head_dim: QK rope head dimension
+            - index_head_dim: Indexer head dimension
+            - kv_precision_bytes: Precision in bytes (BF16=2, INT8=1)
+        page_size_tokens: Tokens per page (default 64)
+
+    Returns:
+        MLALayout: Layout instance for MLA architecture
+
+    Raises:
+        ValueError: If model configuration is invalid
+    """
+    required_fields = ['num_layers', 'kv_lora_rank', 'qk_rope_head_dim',
+                       'index_head_dim', 'kv_precision_bytes']
+    for field in required_fields:
+        if field not in model_config:
+            raise ValueError(f"Missing required field: {field}")
+
+    return MLALayout(
+        num_layers=model_config['num_layers'],
+        kv_lora_rank=model_config['kv_lora_rank'],
+        qk_rope_head_dim=model_config['qk_rope_head_dim'],
+        index_head_dim=model_config['index_head_dim'],
+        precision_bytes=model_config['kv_precision_bytes'],
+        page_size_tokens=page_size_tokens,
+    )
diff --git a/benchmarks/storage_benchmark_v1/storage/__init__.py b/benchmarks/storage_benchmark_v1/storage/__init__.py
new file mode 100644
index 00000000..40fe3856
--- /dev/null
+++ b/benchmarks/storage_benchmark_v1/storage/__init__.py
@@ -0,0 +1,13 @@
+"""
+KVCache Storage Module
+
+Provides storage backend implementations for KVCache systems.
+"""
+
+from .interface import Storage
+from .disk import DiskHashTable
+
+__all__ = [
+    'Storage',
+    'DiskHashTable',
+]
diff --git a/benchmarks/storage_benchmark_v1/storage/disk.py b/benchmarks/storage_benchmark_v1/storage/disk.py
new file mode 100644
index 00000000..286ff751
--- /dev/null
+++ b/benchmarks/storage_benchmark_v1/storage/disk.py
@@ -0,0 +1,328 @@
+"""
+Simple SSD Hash Table Storage
+
+Each key maps to a complete page entry.
+"""
+
+import os
+import time
+from pathlib import Path
+from typing import Dict, List, Optional, Any
+
+from .interface import Storage
+
+
+def calc_percentiles(data):
+    """Calculate latency percentiles"""
+    if not data:
+        return {'avg_ms': 0, 'p50_ms': 0, 'p95_ms': 0, 'p99_ms': 0}
+    import statistics
+    sorted_data = sorted(data)
+    n = len(sorted_data)
+    def get_percentile(p):
+        idx = int(n * p / 100)
+        if idx >= n: idx = n - 1
+        return sorted_data[idx]
+    return {
+        'avg_ms': statistics.mean(data),
+        'p50_ms': get_percentile(50),
+        'p95_ms': get_percentile(95),
+        'p99_ms': get_percentile(99),
+    }
+
+
+class DiskHashTable(Storage):
+    """Simple disk-based hash table
+
+    File Layout:
+    Page 0: offset = 0 * page_size
+    Page 1: offset = 1 * page_size
+    Page 2: offset = 2 * page_size
+    ...
+    Page N: offset = N * page_size
+
+    When max_pages < actual page_id range, uses modulo mapping:
+    actual_page_id = page_id % max_pages
+    This allows simulating large traces with limited storage.
+    """
+
+    def __init__(self, storage_dir: str, page_size: int,
+                 max_pages: int = 100000,
+                 fsync_mode: str = 'batch', fsync_batch_size: int = 100):
+        """Initialize disk hash table
+
+        Args:
+            storage_dir: Storage directory
+            page_size: Size of each entry (page) in bytes
+            max_pages: Maximum number of entries (creates circular mapping if trace is larger)
+            fsync_mode: When to fsync ('batch', 'always', 'end', 'none')
+            fsync_batch_size: Writes between fsync
+        """
+        self.storage_dir = Path(storage_dir)
+        self.page_size = page_size
+        self.max_pages = max_pages
+        self.max_page_id = max_pages - 1
+        self.fsync_mode = fsync_mode
+        self.fsync_batch_size = fsync_batch_size
+        self.storage_dir.mkdir(parents=True, exist_ok=True)
+        self.storage_file = self.storage_dir / "data.bin"
+        self._allocate_file()
+        self.fd = None
+        self._buffer = os.urandom(page_size)
+        self.stats = {
+            'read_count': 0,
+            'write_count': 0,
+            'read_bytes': 0,
+            'write_bytes': 0,
+            'read_latencies_ms': [],
+            'write_latencies_ms': [],
+            'read_time_s': 0.0,
+            'write_time_s': 0.0,
+            'sync_count': 0,
+            'hit': 0,
+            'miss': 0,
+        }
+        self._pending_syncs = 0
+        self._written_pages: set = set()
+
+    def _allocate_file(self):
+        """Pre-allocate disk space efficiently"""
+        file_size = self.max_pages * self.page_size
+        if not self.storage_file.exists():
+            print(f"  [Storage] Creating file: {self.storage_file}")
+            print(f"  [Storage] Requested size: {file_size / (1024**3):.2f} GB ({self.max_pages:,} pages × {self.page_size} bytes = {file_size:,} bytes)")
+            # Use fallocate for efficient preallocation (Linux)
+            fd = os.open(self.storage_file, os.O_RDWR | os.O_CREAT | os.O_TRUNC, 0o644)
+            try:
+                # Try fallocate first (Linux specific, much faster)
+                try:
+                    import fcntl
+                    fcntl.fallocate(fd, 0, file_size)
+                    method = "fallocate"
+                except (ImportError, AttributeError, OSError):
+                    # Fallback to seek+write method
+                    os.lseek(fd, file_size - 1, os.SEEK_SET)
+                    os.write(fd, b'\0')
+                    os.fsync(fd)
+                    method = "seek+write"
+            finally:
+                os.close(fd)
+            actual_size = self.storage_file.stat().st_size if self.storage_file.exists() else 0
+            print(f"  [Storage] Pre-allocated {actual_size / (1024**3):.2f} GB using {method}")
+        else:
+            actual_size = self.storage_file.stat().st_size
+            actual_pages = actual_size // self.page_size
+            print(f"  [Storage] Reusing existing file: {self.storage_file}")
+            print(f"  [Storage] Current file size: {actual_size / (1024**3):.2f} GB ({actual_pages:,} pages × {self.page_size} bytes = {actual_size:,} bytes)")
+
+    def _get_fd(self):
+        if self.fd is None:
+            print(f"  [Storage] Opening file: {self.storage_file}")
+            self.fd = os.open(self.storage_file, os.O_RDWR | os.O_CREAT, 0o644)
+            if self.storage_file.exists():
+                actual_size = self.storage_file.stat().st_size
+                actual_pages = actual_size // self.page_size
+                print(f"  [Storage] File size: {actual_size / (1024**3):.2f} GB ({actual_pages:,} pages × {self.page_size} bytes = {actual_size:,} bytes)")
+        return self.fd
+
+    # ========================================================================
+    # Core operations
+    # ========================================================================
+
+    def _map_page_id(self, page_id: int) -> int:
+        """Map logical page_id to physical page_id using modulo
+
+        This allows simulating large traces with limited storage space.
+
+        Args:
+            page_id: Logical page_id (from trace)
+
+        Returns:
+            Physical page_id in storage (0 to max_pages-1)
+        """
+        return page_id % self.max_pages
+
+    def read(self, page_id: int, offset_in_page: int = 0, length: int = None) -> float:
+        """Read entry from disk
+
+        Args:
+            page_id: Logical page ID (hash_id from trace)
+            offset_in_page: Offset within the page (default: 0)
+            length: Number of bytes to read (default: entire page)
+
+        Returns:
+            Read latency in ms
+        """
+        if length is None:
+            length = self.page_size - offset_in_page
+
+        # Validate parameters
+        if offset_in_page < 0 or offset_in_page >= self.page_size:
+            raise ValueError(f"offset_in_page {offset_in_page} out of range [0, {self.page_size})")
+        if length <= 0 or offset_in_page + length > self.page_size:
+            raise ValueError(f"length {length} invalid with offset_in_page {offset_in_page} (page_size={self.page_size})")
+
+        # Map to physical page_id and calculate offset
+        physical_page_id = self._map_page_id(page_id)
+        offset = physical_page_id * self.page_size + offset_in_page
+        start = time.perf_counter()
+
+        try:
+            fd = self._get_fd()
+            os.pread(fd, length, offset)
+
+            latency = (time.perf_counter() - start) * 1000.0
+            self.stats['read_count'] += 1
+            self.stats['read_bytes'] += length
+            self.stats['read_latencies_ms'].append(latency)
+            self.stats['read_time_s'] += latency / 1000.0
+            self.stats['hit'] += 1
+            return latency
+        except OSError as e:
+            print(f"Read error (page_id={page_id}, physical_page_id={physical_page_id}): {e}")
+            return 0.0
+
+    def write(self, page_id: int, offset_in_page: int = 0, length: int = None) -> float:
+        """Write entry to disk
+
+        Args:
+            page_id: Logical page ID (hash_id from trace)
+            offset_in_page: Offset within the page (default: 0)
+            length: Number of bytes to write (default: entire page)
+
+        Returns:
+            Write latency in ms
+        """
+        if length is None:
+            length = self.page_size - offset_in_page
+
+        # Validate parameters
+        if offset_in_page < 0 or offset_in_page >= self.page_size:
+            raise ValueError(f"offset_in_page {offset_in_page} out of range [0, {self.page_size})")
+        if length <= 0 or offset_in_page + length > self.page_size:
+            raise ValueError(f"length {length} invalid with offset_in_page {offset_in_page} (page_size={self.page_size})")
+
+        # Map to physical page_id and calculate offset
+        physical_page_id = self._map_page_id(page_id)
+        offset = physical_page_id * self.page_size + offset_in_page
+        start = time.perf_counter()
+
+        try:
+            fd = self._get_fd()
+            # Use corresponding portion of buffer
+            os.pwrite(fd, self._buffer[:length], offset)
+            write_done = time.perf_counter()
+
+            # Fsync
+            if self.fsync_mode == 'always':
+                os.fsync(fd)
+                self.stats['sync_count'] += 1
+                self._pending_syncs = 0
+                latency = (time.perf_counter() - start) * 1000.0
+            elif self.fsync_mode == 'batch':
+                self._pending_syncs += 1
+                if self._pending_syncs >= self.fsync_batch_size:
+                    os.fsync(fd)
+                    self.stats['sync_count'] += 1
+                    self._pending_syncs = 0
+                latency = (write_done - start) * 1000.0
+            else:
+                latency = (write_done - start) * 1000.0
+
+            self._written_pages.add(page_id)
+            self.stats['write_count'] += 1
+            self.stats['write_bytes'] += length
+            self.stats['write_latencies_ms'].append(latency)
+            self.stats['write_time_s'] += latency / 1000.0
+            self.stats['miss'] += 1
+            return latency
+        except OSError as e:
+            print(f"Write error (page_id={page_id}, offset_in_page={offset_in_page}, length={length}): {e}")
+            return 0.0
+
+    # ========================================================================
+    # Storage interface methods
+    # ========================================================================
+
+    def exists(self, page_id: int) -> bool:
+        """Check if logical page has been written
+
+        Args:
+            page_id: Logical page ID (hash_id from trace)
+
+        Returns:
+            True if this logical page has been written before
+        """
+        return page_id in self._written_pages
+
+    def delete(self, page_id: int) -> bool:
+        """Delete page (no-op in direct mapping)
+
+        Args:
+            page_id: Page ID to delete
+
+        Returns:
+            True (always succeeds in direct mapping)
+        """
+        # No-op since we don't track which pages have been written
+        return True
+
+    def get_stats(self) -> Dict[str, Any]:
+        """Get statistics"""
+        import statistics
+
+        def calc_stats(latencies):
+            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,
+                'time_s': self.stats['read_time_s'],
+                **calc_stats(self.stats['read_latencies_ms'])
+            },
+            'write': {
+                'count': self.stats['write_count'],
+                'mb': self.stats['write_bytes'] / 1024 / 1024,
+                'time_s': self.stats['write_time_s'],
+                **calc_stats(self.stats['write_latencies_ms'])
+            },
+            'sync_count': self.stats['sync_count'],
+            'max_pages': self.max_pages,
+            'written_pages': len(self._written_pages),
+            'page_hits': self.stats['hit'],
+            'page_misses': self.stats['miss'],
+        }
+
+    # ========================================================================
+    # Resource management
+    # ========================================================================
+
+    def close(self, force_sync: bool = True):
+        """Close file"""
+        if force_sync and self.fsync_mode in ['end', 'batch']:
+            if self.fd is not None:
+                try:
+                    os.fsync(self.fd)
+                    self.stats['sync_count'] += 1
+                except OSError:
+                    pass
+
+        if self.fd is not None:
+            try:
+                os.close(self.fd)
+            except OSError:
+                pass
+            self.fd = None
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self.close()
+        return False
diff --git a/benchmarks/storage_benchmark_v1/storage/interface.py b/benchmarks/storage_benchmark_v1/storage/interface.py
new file mode 100644
index 00000000..e7c2aa02
--- /dev/null
+++ b/benchmarks/storage_benchmark_v1/storage/interface.py
@@ -0,0 +1,93 @@
+"""
+KVCache Storage Interface
+
+Simplified storage interface for KVCache benchmark.
+Key = page_id (int), Value = fixed-size bytes.
+"""
+
+from abc import ABC, abstractmethod
+from typing import Dict, Any
+
+
+class Storage(ABC):
+    """Abstract base class for KVCache storage
+
+    Simplified design:
+    - Key: page_id (int, hash_id from trace)
+    - Value: fixed-size bytes (page_size)
+    - Direct mapping: page_id -> offset -> payload
+
+    Core operations: read, write, exists, delete.
+    """
+
+    @abstractmethod
+    def read(self, page_id: int, offset_in_page: int = 0, length: int = None) -> float:
+        """Read page from disk
+
+        Args:
+            page_id: Page ID (hash_id from trace)
+            offset_in_page: Offset within the page (default: 0)
+            length: Number of bytes to read (default: entire page)
+
+        Returns:
+            Read latency in milliseconds
+        """
+        pass
+
+    @abstractmethod
+    def write(self, page_id: int, offset_in_page: int = 0, length: int = None) -> float:
+        """Write page to disk
+
+        Args:
+            page_id: Page ID (hash_id from trace)
+            offset_in_page: Offset within the page (default: 0)
+            length: Number of bytes to write (default: entire page)
+
+        Returns:
+            Write latency in milliseconds
+        """
+        pass
+
+    @abstractmethod
+    def exists(self, page_id: int) -> bool:
+        """Check if page exists
+
+        Args:
+            page_id: Page ID (hash_id from trace)
+
+        Returns:
+            True if page_id is within valid range
+        """
+        pass
+
+    @abstractmethod
+    def delete(self, page_id: int) -> bool:
+        """Delete page (no-op in direct mapping)
+
+        Args:
+            page_id: Page ID to delete
+
+        Returns:
+            True (always succeeds in direct mapping)
+        """
+        pass
+
+    @abstractmethod
+    def get_stats(self) -> Dict[str, Any]:
+        """Get storage statistics
+
+        Returns:
+            Dictionary containing storage statistics
+        """
+        pass
+
+    def close(self):
+        """Close the storage and release resources"""
+        pass
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self.close()
+        return False
diff --git a/dependencies.sh b/dependencies.sh
index 607e58af..bbbd82aa 100755
--- a/dependencies.sh
+++ b/dependencies.sh
@@ -24,6 +24,7 @@ NC="\033[0m" # No Color
 REPO_ROOT=`pwd`
 GITHUB_PROXY=${GITHUB_PROXY:-"https://github.com"}
 GOVER=1.25.9
+OS_RELEASE_FILE=${OS_RELEASE_FILE:-/etc/os-release}
 
 # Function to print section headers
 print_section() {
@@ -48,23 +49,57 @@ check_success() {
     fi
 }
 
+read_os_release_value() {
+    local key="$1"
+    awk -F= -v key="$key" '
+        $1 == key {
+            value = $0
+            sub(/^[^=]*=/, "", value)
+            gsub(/^"|"$/, "", value)
+            print value
+            exit
+        }
+    ' "$OS_RELEASE_FILE"
+}
+
+# Function to detect OS
+detect_os() {
+    if [ -f "$OS_RELEASE_FILE" ]; then
+        ID=$(read_os_release_value ID)
+        VERSION_ID=$(read_os_release_value VERSION_ID)
+        OS=$(echo "$ID" | tr '[:upper:]' '[:lower:]')
+        OS_VERSION=$VERSION_ID
+    elif [ -f /etc/redhat-release ]; then
+        OS="centos"
+    else
+        print_error "Cannot detect OS. Supported OS: Ubuntu, Debian, CentOS, RHEL, Rocky, AlmaLinux, EulerOS, and openEuler."
+    fi
+
+    echo -e "${GREEN}Detected OS: $OS ${OS_VERSION:-unknown}${NC}"
+}
+
 if [ $(id -u) -ne 0 ]; then
 	print_error "Require root permission, try sudo ./dependencies.sh"
 fi
 
 # Parse command line arguments
 SKIP_CONFIRM=false
+INSTALL_SPDK=false
 for arg in "$@"; do
     case $arg in
         -y|--yes)
             SKIP_CONFIRM=true
             ;;
+        --with-spdk)
+            INSTALL_SPDK=true
+            ;;
         -h|--help)
             echo -e "${YELLOW}Mooncake Dependencies Installer${NC}"
             echo -e "Usage: ./dependencies.sh [OPTIONS]"
             echo -e "\nOptions:"
-            echo -e "  -y, --yes    Skip confirmation and install all dependencies"
-            echo -e "  -h, --help   Show this help message and exit"
+            echo -e "  -y, --yes       Skip confirmation and install all dependencies"
+            echo -e "  --with-spdk     Install SPDK for NVMe-oF support"
+            echo -e "  -h, --help      Show this help message and exit"
             exit 0
             ;;
     esac
@@ -77,6 +112,9 @@ echo -e "The following components will be installed:"
 echo -e "  - System packages (build tools, libraries)"
 echo -e "  - Git submodules (including pybind11 and yalantinglibs)"
 echo -e "  - Go $GOVER"
+if [ "$INSTALL_SPDK" = true ]; then
+    echo -e "  - SPDK (for NVMe-oF support)"
+fi
 echo
 
 # Ask for confirmation unless -y flag is used
@@ -89,51 +127,100 @@ if [ "$SKIP_CONFIRM" = false ]; then
     fi
 fi
 
+# Detect OS
+detect_os
 
 # Update package lists
 print_section "Updating package lists"
-apt-get update
-check_success "Failed to update package lists"
+if [ "$OS" = "ubuntu" ] || [ "$OS" = "debian" ]; then
+    apt-get update
+    check_success "Failed to update package lists"
+elif [ "$OS" = "centos" ] || [ "$OS" = "rhel" ] || [ "$OS" = "rocky" ] || [ "$OS" = "almalinux" ] || [ "$OS" = "euleros" ] || [ "$OS" = "openeuler" ]; then
+    yum install -y dnf-plugins-core epel-release || true
+    yum config-manager --set-enabled powertools || yum config-manager --set-enabled crb || true
+    yum clean all
+    yum makecache
+    check_success "Failed to update package lists"
+else
+    print_error "Unsupported OS: $OS"
+fi
 
 # Install system packages
 print_section "Installing system packages"
 echo -e "${YELLOW}This may take a few minutes...${NC}"
 
-SYSTEM_PACKAGES="build-essential \
-                  cmake \
-                  ninja-build \
-                  git \
-                  wget \
-                  unzip \
-                  libibverbs-dev \
-                  libgoogle-glog-dev \
-                  libgtest-dev \
-                  libjsoncpp-dev \
-                  libunwind-dev \
-                  libnuma-dev \
-                  libpython3-dev \
-                  libboost-all-dev \
-                  libssl-dev \
-                  libgrpc-dev \
-                  libgrpc++-dev \
-                  libprotobuf-dev \
-                  libyaml-cpp-dev \
-                  protobuf-compiler-grpc \
-                  libcurl4-openssl-dev \
-                  libhiredis-dev \
-                  liburing-dev \
-                  libjemalloc-dev \
-                  libmsgpack-dev \
-                  libzstd-dev \
-                  libasio-dev \
-                  libxxhash-dev \
-                  pkg-config \
-                  patchelf \
-                  libc6-dev \
-                  libc-bin"
+if [ "$OS" = "ubuntu" ] || [ "$OS" = "debian" ]; then
+    SYSTEM_PACKAGES="build-essential \
+                     cmake \
+                     ninja-build \
+                     git \
+                     wget \
+                     unzip \
+                     libibverbs-dev \
+                     libgoogle-glog-dev \
+                     libgtest-dev \
+                     libjsoncpp-dev \
+                     libunwind-dev \
+                     libnuma-dev \
+                     libpython3-dev \
+                     libboost-all-dev \
+                     libssl-dev \
+                     libgrpc-dev \
+                     libgrpc++-dev \
+                     libprotobuf-dev \
+                     libyaml-cpp-dev \
+                     protobuf-compiler-grpc \
+                     libcurl4-openssl-dev \
+                     libhiredis-dev \
+                     liburing-dev \
+                     libjemalloc-dev \
+                     libmsgpack-dev \
+                     libzstd-dev \
+                     libasio-dev \
+                     libxxhash-dev \
+                     pkg-config \
+                     patchelf \
+                     libc6-dev \
+                     libc-bin"
+
+    apt-get install -y $SYSTEM_PACKAGES
+    check_success "Failed to install system packages"
+
+elif [ "$OS" = "centos" ] || [ "$OS" = "rhel" ] || [ "$OS" = "rocky" ] || [ "$OS" = "almalinux" ] || [ "$OS" = "euleros" ] || [ "$OS" = "openeuler" ]; then
+    SYSTEM_PACKAGES="@development \
+                     cmake \
+                     git \
+                     wget \
+                     rdma-core-devel \
+                     glog-devel \
+                     gflags-devel \
+                     gtest-devel \
+                     jsoncpp-devel \
+                     libunwind-devel \
+                     numactl-devel \
+                     python3-devel \
+                     boost1.78-devel \
+                     openssl-devel \
+                     protobuf-devel \
+                     yaml-cpp-devel \
+                     libcurl-devel \
+                     hiredis-devel \
+                     liburing-devel \
+                     jemalloc-devel \
+                     msgpack-devel \
+                     libzstd-devel \
+                     pkgconf-pkg-config \
+                     elfutils-libelf-devel \
+                     patchelf  \
+                     xxhash-devel \
+                     libbsd-devel"
+
+    yum install -y $SYSTEM_PACKAGES
+    check_success "Failed to install system packages"
+else
+    print_error "Unsupported OS: $OS"
+fi
 
-apt-get install -y $SYSTEM_PACKAGES
-check_success "Failed to install system packages"
 print_success "System packages installed successfully"
 
 # Initialize and update git submodules
@@ -185,15 +272,16 @@ cd "${REPO_ROOT}"
 print_section "Verifying essential build tools"
 
 # Verify getconf and ldd (required for glibc version detection in build_wheel.sh)
-# Both are provided by libc-bin, which is included in SYSTEM_PACKAGES
-if ! command -v getconf >/dev/null 2>&1; then
-    print_error "getconf not found after installing system packages. This should not happen."
+if [ "$OS" = "ubuntu" ] || [ "$OS" = "debian" ]; then
+    if ! command -v getconf >/dev/null 2>&1; then
+        print_error "getconf not found after installing system packages. This should not happen."
+    fi
+    if ! command -v ldd >/dev/null 2>&1; then
+        print_error "ldd not found after installing system packages. This should not happen."
+    fi
+    print_success "getconf found: $(getconf --version 2>&1 | head -1)"
+    print_success "ldd found: $(ldd --version 2>&1 | head -1)"
 fi
-if ! command -v ldd >/dev/null 2>&1; then
-    print_error "ldd not found after installing system packages. This should not happen."
-fi
-print_success "getconf found: $(getconf --version 2>&1 | head -1)"
-print_success "ldd found: $(ldd --version 2>&1 | head -1)"
 
 print_section "Installing Go $GOVER"
 
@@ -224,8 +312,6 @@ install_go() {
         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
@@ -241,19 +327,16 @@ install_go() {
         print_error "Failed to download Go $GOVER from all mirrors"
     fi
 
-    # Install Go
     echo "Installing Go $GOVER..."
     tar -C /usr/local -xzf "${GO_TARBALL}"
     check_success "Failed to install Go $GOVER"
 
-    # Clean up downloaded file
     rm -f "${GO_TARBALL}"
     check_success "Failed to clean up Go installation file"
 
     print_success "Go $GOVER installed successfully"
 }
 
-# Check if Go is already installed
 if command -v go &> /dev/null; then
     GO_VERSION=$(go version | awk '{print $3}')
     if [[ "$GO_VERSION" == "go$GOVER" ]]; then
@@ -273,8 +356,7 @@ 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.
+# Set GOPROXY only if Go download fell back to a CN mirror
 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}"
@@ -287,6 +369,69 @@ elif [ -n "$GOPROXY" ]; then
     echo -e "${GREEN}GOPROXY already set to: ${GOPROXY}${NC}"
 fi
 
+# Install SPDK if requested
+if [ "$INSTALL_SPDK" = true ]; then
+    print_section "Installing SPDK"
+
+    cd "${REPO_ROOT}/extern"
+    check_success "Failed to change to extern directory"
+
+    # Remove existing SPDK if present
+    if [ -d "spdk" ]; then
+        echo -e "${YELLOW}SPDK directory already exists. Removing for fresh install...${NC}"
+        rm -rf spdk
+        check_success "Failed to remove existing SPDK directory"
+    fi
+
+    # Clone SPDK
+    echo "Cloning SPDK from ${GITHUB_PROXY}/spdk/spdk.git..."
+    git clone ${GITHUB_PROXY}/spdk/spdk.git
+    check_success "Failed to clone SPDK"
+
+    cd spdk
+    check_success "Failed to change to SPDK directory"
+
+    # Checkout specific version
+    echo "Checking out SPDK version v23.01.1..."
+    git checkout v23.01.1
+    check_success "Failed to checkout SPDK version v23.01.1"
+
+    # Initialize submodules
+    echo "Initializing SPDK submodules..."
+    git submodule update --init
+    check_success "Failed to initialize SPDK submodules"
+
+    # Install SPDK dependencies
+    echo "Installing SPDK dependencies..."
+    ./scripts/pkgdep.sh
+    check_success "Failed to install SPDK dependencies"
+
+    # Configure SPDK with RDMA support
+    echo "Configuring SPDK with RDMA support..."
+    ./configure --with-rdma
+    check_success "Failed to configure SPDK"
+
+    # Build SPDK
+    echo "Building SPDK (using $(nproc) cores)..."
+    make -j$(nproc)
+    check_success "Failed to build SPDK"
+
+    # Install SPDK
+    echo "Installing SPDK..."
+    make install
+    check_success "Failed to install SPDK"
+
+    # Copy DPDK libraries to system library path
+    if ls dpdk/build/lib/*.a >/dev/null 2>&1; then
+        echo "Copying DPDK libraries to /usr/local/lib..."
+        cp dpdk/build/lib/*.a /usr/local/lib/
+        check_success "Failed to copy DPDK libraries"
+    fi
+
+    print_success "SPDK installed successfully"
+    cd "${REPO_ROOT}"
+fi
+
 # Return to the repository root
 cd "${REPO_ROOT}"
 
@@ -298,6 +443,13 @@ echo -e "  ${GREEN}✓${NC} System packages"
 echo -e "  ${GREEN}✓${NC} yalantinglibs"
 echo -e "  ${GREEN}✓${NC} Git submodules"
 echo -e "  ${GREEN}✓${NC} Go $GOVER"
+if [ "$INSTALL_SPDK" = true ]; then
+    echo -e "  ${GREEN}✓${NC} SPDK (v23.01.1)"
+fi
 echo
 echo -e "You can now build and run Mooncake."
 echo -e "${YELLOW}Note: You may need to restart your terminal or run 'source ~/.bashrc' to use Go.${NC}"
+
+if [ "$INSTALL_SPDK" = true ]; then
+    echo -e "${YELLOW}Note: SPDK requires hugepages and RDMA configuration. Please refer to SPDK documentation for setup.${NC}"
+fi
diff --git a/docker/mooncake.Dockerfile b/docker/mooncake.Dockerfile
index 60d7314d..3dda2e6c 100644
--- a/docker/mooncake.Dockerfile
+++ b/docker/mooncake.Dockerfile
@@ -116,6 +116,7 @@ RUN apt-get update && \
 
 # Copy wheels produced in builder stage and install them via pip
 COPY --from=builder /workspace/mooncake-wheel/dist /tmp/mooncake-wheel
+COPY --chmod=755 scripts/check_hicache_hugepage_requirements.py /usr/local/bin/mooncake-hicache-sizing
 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"]
diff --git a/docs/source/api-reference/cpp/index.md b/docs/source/api-reference/cpp/index.md
new file mode 100644
index 00000000..35e77191
--- /dev/null
+++ b/docs/source/api-reference/cpp/index.md
@@ -0,0 +1,16 @@
+# C/C++ API Reference
+
+| Module | Description |
+|--------|-------------|
+| [Transfer Engine C++ API](../../design/transfer-engine/cpp-api) | `TransferEngine` class — memory registration, batch transfer, segment management, RDMA transport |
+| [TENT C++ API](../../design/tent/cpp-api) | `mooncake::tent::TransferEngine` — next-gen transfer engine with automatic transport selection and fault tolerance |
+| [Mooncake Store Client C++ API](mooncake-store) | `Client` class — `Put`/`Get`/`Remove`/`Replicate` operations, `BufferAllocatorBase` interface |
+
+:::{toctree}
+:maxdepth: 1
+:hidden:
+
+../../design/transfer-engine/cpp-api
+../../design/tent/cpp-api
+mooncake-store
+:::
diff --git a/docs/source/api-reference/cpp/mooncake-store.md b/docs/source/api-reference/cpp/mooncake-store.md
new file mode 100644
index 00000000..41c3028a
--- /dev/null
+++ b/docs/source/api-reference/cpp/mooncake-store.md
@@ -0,0 +1,178 @@
+# Mooncake Store C++ API Reference
+
+## Client C++ API
+
+### Constructor and Initialization `Init`
+
+```C++
+ErrorCode Init(const std::string& local_hostname,
+               const std::string& metadata_connstring,
+               const std::string& protocol,
+               void** protocol_args,
+               const std::string& master_server_entry);
+```
+
+Initializes the Mooncake Store client. The parameters are as follows:
+- `local_hostname`: The `IP:Port` of the local machine or an accessible domain name (default value used if port is not included)
+- `metadata_connstring`: The address of the metadata service (e.g., etcd/Redis) required for Transfer Engine initialization
+- `protocol`: The protocol supported by the Transfer Engine, including RDMA and TCP
+- `protocol_args`: Protocol parameters required by the Transfer Engine
+- `master_server_entry`: The address information of the Master (`IP:Port` for default mode and `etcd://IP:Port;IP:Port;...;IP:Port` for high availability mode)
+
+### Get
+
+```C++
+tl::expected<void, ErrorCode> Get(const std::string& object_key,
+                                  std::vector<Slice>& slices);
+```
+
+`Get` retrieves the value of `object_key` into the provided `slices`. The returned data is guaranteed to be complete and correct. Each slice must reference local DRAM/VRAM memory that has been pre-registered with `registerLocalMemory(addr, len)` (not the global segments that contribute to the distributed memory pool). When persistence is enabled and the requested data is not found in the distributed memory pool, `Get` will fall back to loading the data from SSD.
+
+### Put
+
+```C++
+tl::expected<void, ErrorCode> Put(const ObjectKey& key,
+                                  std::vector<Slice>& slices,
+                                  const ReplicateConfig& config);
+```
+
+`Put` stores the value associated with `key` in the distributed memory pool. The `config` parameter allows specifying the required number of replicas as well as the preferred segment for storing the value. When persistence is enabled, `Put` also asynchronously triggers a persistence operation to SSD.
+
+**Replication Guarantees and Best Effort Behavior:**
+- Each slice of an object is guaranteed to be replicated to different segments, ensuring distribution across separate storage nodes
+- Different slices from different objects may be placed in the same segment
+- Replication operates on a best-effort basis: if insufficient space is available for all requested replicas, the object will still be written with as many replicas as possible
+
+The data structure details of `ReplicateConfig` are as follows:
+
+```C++
+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++
+tl::expected<void, ErrorCode> Remove(const ObjectKey& key);
+```
+
+Used to delete the object corresponding to the specified key. This interface marks all data replicas associated with the key in the storage engine as deleted, without needing to communicate with the corresponding storage node (Client).
+
+### CreateCopyTask
+
+```C++
+tl::expected<UUID, ErrorCode> CreateCopyTask(
+    const std::string& key,
+    const std::vector<std::string>& targets);
+```
+
+`CreateCopyTask` creates an asynchronous copy task that will be executed by the client's task execution system. This is useful when you want to submit multiple copy operations without waiting for each one to complete. The task is submitted to the master service, assigned a unique task ID, and executed asynchronously by an available client. The task status can be queried using `QueryTask`.
+
+**Task Execution and Result Reporting:**
+1. **Task Assignment**: The master service assigns the task to an available client during the client's periodic ping operation
+2. **Task Execution**: The assigned client executes the copy operation asynchronously in a background thread pool
+3. **Result Reporting**: Upon completion (success or failure), the client automatically reports the result to the master service via `MarkTaskToComplete`:
+   - On success: `status = SUCCESS`, `message = "Task completed successfully"`
+   - On failure: `status = FAILED`, `message = <error description>`
+4. **Status Query**: You can query the task status at any time using `QueryTask` to monitor progress
+
+### CreateMoveTask
+
+```C++
+tl::expected<UUID, ErrorCode> CreateMoveTask(
+    const std::string& key,
+    const std::string& source,
+    const std::string& target);
+```
+
+`CreateMoveTask` creates an asynchronous move task that will be executed by the client's task execution system. This is useful when you want to submit multiple move operations without waiting for each one to complete. The task is submitted to the master service, assigned a unique task ID, and executed asynchronously by an available client. The task status can be queried using `QueryTask`.
+
+**Task Execution and Result Reporting:**
+1. **Task Assignment**: The master service assigns the task to an available client during the client's periodic ping operation
+2. **Task Execution**: The assigned client executes the move operation asynchronously in a background thread pool
+3. **Result Reporting**: Upon completion (success or failure), the client automatically reports the result to the master service via `MarkTaskToComplete`:
+   - On success: `status = SUCCESS`, `message = "Task completed successfully"`
+   - On failure: `status = FAILED`, `message = <error description>`
+4. **Status Query**: You can query the task status at any time using `QueryTask` to monitor progress
+
+### QueryTask
+
+```C++
+tl::expected<QueryTaskResponse, ErrorCode> QueryTask(const UUID& task_id);
+```
+
+`QueryTask` queries the status of an asynchronous task (copy or move). This allows you to monitor the progress of task-based operations. The response includes task status, type, creation time, last update time, assigned client, and status message.
+
+The data structure details of `QueryTaskResponse` are as follows:
+
+```C++
+struct QueryTaskResponse {
+    UUID id;                                    // Task UUID
+    TaskType type;                              // Task type (REPLICA_COPY or REPLICA_MOVE)
+    TaskStatus status;                          // Task status (PENDING, PROCESSING, SUCCESS, or FAILED)
+    int64_t created_at_ms_epoch;                // Task creation timestamp in milliseconds
+    int64_t last_updated_at_ms_epoch;           // Last update timestamp in milliseconds
+    UUID assigned_client;                       // UUID of the client assigned to execute the task
+    std::string message;                        // Status message or error description
+};
+```
+
+### BatchQueryIp
+
+```C++
+tl::expected<std::unordered_map<UUID, std::vector<std::string>, boost::hash<UUID>>, ErrorCode>
+BatchQueryIp(const std::vector<UUID>& client_ids);
+```
+
+Used to batch query the IP addresses for multiple client IDs. For each client ID in the input list, this interface retrieves the unique IP addresses from all segments mounted by that client. The operation is performed on the Master Service and returns a map from client ID to their IP address lists. Only client IDs that have successfully mounted segments are included in the result map. This is useful for discovering the network locations of storage nodes in the cluster.
+
+### BatchReplicaClear
+
+```C++
+tl::expected<std::vector<std::string>, ErrorCode>
+BatchReplicaClear(const std::vector<std::string>& object_keys,
+                  const UUID& client_id,
+                  const std::string& segment_name);
+```
+
+Used to batch clear replicas for multiple object keys belonging to a specific client ID. This interface allows clearing replicas either on a specific segment or across all segments. If segment_name is empty, all replicas of the specified objects are cleared (the objects are deleted entirely). If segment_name is provided, only replicas located on that specific segment are cleared. The operation is performed on the Master Service and returns a list of object keys that were successfully cleared. Only objects that belong to the specified `client_id`, have expired leases, and meet the clearing criteria are processed. This is useful for managing storage resources and cleaning up data on specific storage nodes.
+
+### QueryByRegex
+
+```C++
+tl::expected<std::unordered_map<std::string, std::vector<Replica::Descriptor>>, ErrorCode>
+QueryByRegex(const std::string& str);
+```
+
+Used to query the replica information for all objects whose keys match the given regular expression. This is useful for batch operations or for retrieving a group of related objects. The operation is performed on the Master and returns a map of keys to their replica lists.
+
+### RemoveByRegex
+
+```C++
+tl::expected<long, ErrorCode> RemoveByRegex(const ObjectKey& str);
+```
+
+Used to delete all objects from the store whose keys match the specified regular expression. This provides a powerful way to perform bulk deletions. The command returns the number of objects that were successfully removed.
diff --git a/docs/source/api-reference/http/index.md b/docs/source/api-reference/http/index.md
new file mode 100644
index 00000000..e946409f
--- /dev/null
+++ b/docs/source/api-reference/http/index.md
@@ -0,0 +1,14 @@
+# HTTP
+
+| Module | Description |
+|--------|-------------|
+| [HTTP Service](../../http-api-reference/http-service) | RESTful endpoints for cluster management, metrics, and data inspection |
+| [Conductor Indexer API](../../design/conductor/indexer-api-design) | Cache-aware routing: service registration, KV event subscription, prefix cache-hit query APIs |
+
+:::{toctree}
+:maxdepth: 1
+:hidden:
+
+../../http-api-reference/http-service
+../../design/conductor/indexer-api-design
+:::
diff --git a/docs/source/api-reference/python/index.md b/docs/source/api-reference/python/index.md
new file mode 100644
index 00000000..3bb15f03
--- /dev/null
+++ b/docs/source/api-reference/python/index.md
@@ -0,0 +1,18 @@
+# Python
+
+| Module | Description |
+|--------|-------------|
+| [Mooncake Store](../../python-api-reference/mooncake-store) | Distributed KV cache storage client — `put`/`get`/`remove`/`replicate` operations |
+| [DataProto Structured Object Transfer](../../python-api-reference/dataproto-structured-object-transfer) | Structured-object helpers for storing and retrieving DataProto-like payloads |
+| [Transfer Engine](../../python-api-reference/transfer-engine) | High-performance RDMA/TCP data transfer between nodes |
+| [EP Backend](../../python-api-reference/ep-backend) | Expert-parallel backend for large MoE model deployment |
+
+:::{toctree}
+:maxdepth: 1
+:hidden:
+
+../../python-api-reference/mooncake-store
+../../python-api-reference/dataproto-structured-object-transfer
+../../python-api-reference/transfer-engine
+../../python-api-reference/ep-backend
+:::
diff --git a/docs/source/api-reference/rust/index.md b/docs/source/api-reference/rust/index.md
new file mode 100644
index 00000000..30ab8e83
--- /dev/null
+++ b/docs/source/api-reference/rust/index.md
@@ -0,0 +1,15 @@
+# Rust
+
+| Module | Description |
+|--------|-------------|
+| [Mooncake Store](./mooncake-store) | Safe Rust bindings for the distributed KV-cache store (`mooncake_store`) |
+| [Transfer Engine](./transfer-engine) | Rust bindings for Transfer Engine (FFI wrapper used by the Rust example crate) |
+
+```{toctree}
+:maxdepth: 1
+:hidden:
+
+mooncake-store
+transfer-engine
+```
+
diff --git a/docs/source/api-reference/rust/mooncake-store.md b/docs/source/api-reference/rust/mooncake-store.md
new file mode 100644
index 00000000..f589246d
--- /dev/null
+++ b/docs/source/api-reference/rust/mooncake-store.md
@@ -0,0 +1,223 @@
+# Mooncake Store Rust API
+
+This page documents the Rust crate `mooncake_store` (located at `mooncake-store/rust`).
+It is a **safe wrapper** around the Mooncake Store C API (`store_c.h`).
+
+For deployment and service prerequisites, also see:
+
+- Mooncake Store deployment guide: `deployment/mooncake-store-deployment-guide`
+- Error code reference: `troubleshooting/error-code`
+
+## Build & runtime prerequisites
+
+The Rust crate links against the C++ Mooncake build outputs.
+
+- **Build**:
+  - Build Mooncake with Store + Rust enabled via CMake: `-DWITH_STORE=ON -DWITH_STORE_RUST=ON`
+  - Or build with Cargo after exporting the CMake build directory / include paths (see `mooncake-store/rust/README.md`).
+- **Runtime**:
+  - Dynamic linker must find Mooncake shared libraries (typically via `LD_LIBRARY_PATH` pointing at the CMake build outputs).
+  - The store client requires:
+    - a **metadata server** (HTTP metadata or etcd, depending on your deployment)
+    - `mooncake_master`
+
+## Quick start (copy-paste)
+
+```rust
+use mooncake_store::MooncakeStore;
+
+fn main() -> Result<(), mooncake_store::StoreError> {
+    let store = MooncakeStore::new()?;
+    store.setup(
+        "127.0.0.1",
+        "http://127.0.0.1:8080/metadata",
+        512 << 20, // global_segment_size
+        128 << 20, // local_buffer_size
+        "tcp",
+        "",
+        "127.0.0.1:50051",
+    )?;
+
+    store.put("hello", b"world", None)?;
+    let value = store.get("hello")?;
+    assert_eq!(value, b"world");
+
+    store.remove("hello", false)?;
+    Ok(())
+}
+```
+
+## API reference
+
+### `MooncakeStore`
+
+#### `new() -> Result<MooncakeStore, StoreError>`
+
+Allocate a new store handle (uninitialised). You must call `setup()` before any data operations.
+
+#### `setup(...) -> Result<(), StoreError>`
+
+Initialise the store client and establish connections.
+
+Parameters:
+
+- `local_hostname`: IP/hostname for this node.
+- `metadata_server`: metadata URI, for example:
+  - HTTP: `"http://127.0.0.1:8080/metadata"`
+  - etcd: `"etcd://127.0.0.1:2379"`
+- `global_segment_size`: per-segment size in bytes.
+- `local_buffer_size`: local staging buffer size in bytes.
+- `protocol`: transport protocol string (for example `"tcp"` / `"rdma"`).
+- `device_name`: device selector; empty string means auto-select (when supported by the backend).
+- `master_server_addr`: `mooncake_master` address, e.g. `"127.0.0.1:50051"`.
+
+Returns `Ok(())` on success, otherwise `StoreError::OperationFailed(code)`.
+
+#### `health_check() -> Result<(), StoreError>`
+
+Connectivity health check. Returns `Ok(())` when the backend is reachable.
+
+#### `put(key, value, config) -> Result<(), StoreError>`
+
+Store `value` under `key`. This is a **copying** API: `value` is copied into store-managed buffers.
+
+- `config`: optional replication settings (`ReplicateConfig`).
+
+#### `get(key) -> Result<Vec<u8>, StoreError>`
+
+Retrieve the full value for `key` into a newly allocated `Vec<u8>`.
+
+Notes:
+
+- Internally calls `get_size()` to allocate an exact-sized buffer, then `get_into()` to fill it.
+- A missing key or backend failure can surface as `OperationFailed(...)` because the C API does not provide a distinct NotFound code in all paths.
+
+#### `unsafe get_into(key, buffer, size) -> Result<i64, StoreError>`
+
+Retrieve the value for `key` into a caller-provided buffer.
+
+- **Returns**: number of bytes written on success.
+- **Safety**: `buffer` must point to at least `size` bytes of writable valid memory.
+
+#### `is_exist(key) -> Result<bool, StoreError>`
+
+Existence check.
+
+- `Ok(true)` if exists, `Ok(false)` if missing.
+- Any other return code becomes `StoreError::OperationFailed(code)`.
+
+#### `get_size(key) -> Result<i64, StoreError>`
+
+Get the stored value size in bytes.
+
+Important limitation:
+
+- The underlying C API uses a single negative return code for multiple error conditions, so Rust surfaces errors as `OperationFailed(raw_code)` without distinguishing NotFound.
+
+#### `get_hostname() -> Result<String, StoreError>`
+
+Returns the hostname (and potentially port) that the store client is registered under.
+
+#### `remove(key, force) -> Result<(), StoreError>`
+
+Remove a key.
+
+- If `force = true`, the key is removed even if another client is reading it.
+
+#### `remove_by_regex(pattern, force) -> Result<i64, StoreError>`
+
+Remove keys matching a regex pattern. Returns number of removed keys.
+
+#### `remove_all(force) -> Result<i64, StoreError>`
+
+Remove **all** keys. Returns number of removed keys.
+
+### Zero-copy APIs (advanced)
+
+The Rust wrapper exposes zero-copy APIs that map directly to the underlying RDMA-capable C++ store.
+
+#### `unsafe register_buffer(buffer, size) -> Result<(), StoreError>`
+
+Register a memory region for zero-copy operations.
+
+- **Safety**: `buffer` must remain valid and pinned until `unregister_buffer()` is called.
+- This is required before calling `put_from()` or other registered-memory operations.
+
+#### `unsafe unregister_buffer(buffer) -> Result<(), StoreError>`
+
+Deregister a previously registered buffer.
+
+#### `unsafe put_from(key, buffer, size, config) -> Result<(), StoreError>`
+
+Store from a registered buffer.
+
+- **Safety**: `buffer` must have been registered via `register_buffer()` and be at least `size` bytes.
+
+### Batch APIs
+
+Batch forms are useful when you want to amortize RPC overhead.
+
+#### `unsafe batch_put_from(keys, buffers, sizes, config) -> Result<Vec<i32>, StoreError>`
+
+Batch version of `put_from()`.
+
+- **Returns**: per-key result codes (0 = success, non-zero = error code for that key).
+- **Safety**: each `buffers[i]` must be registered and valid for `sizes[i]` bytes.
+
+#### `unsafe batch_get_into(keys, buffers, sizes) -> Result<Vec<i64>, StoreError>`
+
+Batch version of `get_into()`.
+
+- **Returns**: per-key bytes written (≥ 0) or error code (< 0).
+- **Safety**: each destination buffer must be writable and at least `sizes[i]` bytes.
+
+#### `batch_is_exist(keys) -> Result<Vec<bool>, StoreError>`
+
+Batch existence check. Errors are returned as `OperationFailed(code)`.
+
+### `ReplicateConfig`
+
+Replication settings for write operations (`put`, `put_from`, `batch_put_from`).
+
+Fields:
+
+- `replica_num`: number of replicas (0 means “use server default”).
+- `with_soft_pin`: prefer retaining the object in memory (soft pin).
+- `with_hard_pin`: never evict (hard pin).
+- `preferred_segments`: whitelist of segment names that should host a replica.
+
+Example:
+
+```rust
+use mooncake_store::{MooncakeStore, ReplicateConfig};
+
+fn write_with_replication(store: &MooncakeStore) -> Result<(), mooncake_store::StoreError> {
+    let cfg = ReplicateConfig {
+        replica_num: 2,
+        with_soft_pin: true,
+        with_hard_pin: false,
+        preferred_segments: vec!["seg-a".to_string(), "seg-b".to_string()],
+    };
+
+    store.put("k", b"v", Some(&cfg))?;
+    Ok(())
+}
+```
+
+### `StoreError`
+
+Errors returned by the Rust wrapper.
+
+- `NullHandle`: store handle allocation failed.
+- `InvalidString`: input string contained an interior `\0` (cannot form C string).
+- `OperationFailed(i32)`: underlying C layer returned a non-zero / negative code.
+- `NotFound`: convenience for consumers that implement a NotFound check externally.
+- `InvalidArgument(String)`: wrapper-level argument validation failure (e.g. mismatched array lengths).
+
+## Safety & thread-safety
+
+- `MooncakeStore` is `Send + Sync` (the underlying C object is internally synchronised).
+- Methods that accept raw pointers are marked `unsafe`:
+  - You must uphold Rust aliasing and lifetime rules for buffers passed to FFI.
+  - For zero-copy operations, buffers must be registered and remain valid until unregistered.
+
diff --git a/docs/source/api-reference/rust/transfer-engine.md b/docs/source/api-reference/rust/transfer-engine.md
new file mode 100644
index 00000000..a7279115
--- /dev/null
+++ b/docs/source/api-reference/rust/transfer-engine.md
@@ -0,0 +1,207 @@
+# Transfer Engine Rust API
+
+This page documents the Rust bindings living under `mooncake-transfer-engine/rust`.
+
+At the time of writing, the crate (`transfer_engine_rust`) is primarily used as a **Rust-side binding + example binary**. The public Rust types are implemented in `src/transfer_engine.rs` and wrap the Transfer Engine C API (`transfer_engine_c.h`).
+
+For Transfer Engine design docs and non-Rust APIs, see:
+
+- Transfer Engine design docs: `design/transfer-engine/index`
+- Transfer Engine C++ API: `design/transfer-engine/cpp-api`
+
+## Build & runtime prerequisites
+
+The Rust crate uses bindgen + CMake to link against the Transfer Engine C/C++ build outputs.
+
+- **Build**:
+  - Requires Rust toolchain and libclang (bindgen).
+  - The crate has `build.rs` that expects to find / build the native library via CMake.
+- **Runtime**:
+  - Dynamic linker must find Transfer Engine shared libraries.
+  - You need a metadata server backend (commonly etcd) and a reachable peer segment registry.
+
+## Mental model
+
+The Transfer Engine operates on **segments** and **transfer batches**:
+
+- You create a `TransferEngine` bound to:
+  - `metadata_uri` (for example, etcd endpoint)
+  - `local_server_name` (this node's address/name)
+  - `rpc_port` (RPC listener port)
+- You register local memory regions as RDMA-capable buffers.
+- You open a remote segment to obtain a `segment_id` (an integer handle).
+- You allocate a batch id for a fixed number of transfer requests.
+- You submit a batch of `TransferRequest`.
+- You poll status per task id inside the batch, then free the batch id.
+
+## API reference
+
+### Enums
+
+#### `OpcodeEnum`
+
+- `OpcodeEnum::Read`
+- `OpcodeEnum::Write`
+
+Used by `TransferRequest.opcode`.
+
+#### `TransferStatusEnum`
+
+Status values returned by the C layer. Common values you will check for:
+
+- `Completed`
+- `Failed`
+- `Timeout`
+
+### Structs
+
+#### `TransferRequest`
+
+One transfer operation inside a batch.
+
+Fields:
+
+- `opcode: OpcodeEnum`
+- `source: *mut c_void`: local source/destination pointer (depends on opcode).
+- `target_id: i32`: segment id returned by `open_segment()`.
+- `target_offset: u64`: byte offset inside the target segment.
+- `length: u64`: transfer length in bytes.
+
+#### `BufferEntry`
+
+Used for batch memory registration:
+
+- `addr: *mut c_void`
+- `length: u64`
+
+### `TransferEngine`
+
+#### `new(metadata_uri, local_server_name, rpc_port) -> anyhow::Result<TransferEngine>`
+
+Create a new engine handle.
+
+Notes:
+
+- `metadata_uri` and `local_server_name` are passed through `CString`; interior `\0` bytes will error.
+- The wrapper currently disables `auto_discover` in the underlying C call.
+
+#### `discover_topology() -> anyhow::Result<()>`
+
+Trigger topology discovery.
+
+#### `install_transport(proto) -> anyhow::Result<()>`
+
+Install a transport by name (e.g. `"tcp"`, `"rdma"`, `"efa"` depending on build/runtime support).
+
+#### `register_local_memory(addr, length, location) -> anyhow::Result<()>`
+
+Register a local memory region for zero-copy transfers.
+
+- `addr`: pointer to the memory region.
+- `length`: size in bytes.
+- `location`: location string such as `"cpu:0"`.
+
+#### `unregister_local_memory(addr) -> anyhow::Result<()>`
+
+Unregister a previously registered memory region.
+
+#### `register_local_memory_batch(buffer_list, location) -> anyhow::Result<()>`
+
+Batch register multiple local buffers.
+
+- No-op when `buffer_list` is empty.
+
+#### `unregister_local_memory_batch(buffer_list) -> anyhow::Result<()>`
+
+Batch unregister multiple local buffers.
+
+#### `open_segment(name: String) -> anyhow::Result<i32>`
+
+Open a remote segment by name and get a segment id.
+
+#### `close_segment(segment_id: i32) -> anyhow::Result<()>`
+
+Close a previously opened segment.
+
+#### `warmup_efa_segment(name: &str) -> anyhow::Result<()>`
+
+Eagerly establish EFA endpoints so the first `submit_transfer()` does not pay the serial connection setup cost.
+
+- No-op on non-EFA transports.
+- Call after `open_segment()` and after the metadata server has published the peer's NIC list.
+
+#### `sync_segment_cache() -> anyhow::Result<()>`
+
+Synchronize segment cache from metadata.
+
+#### `allocate_batch_id(batch_size) -> anyhow::Result<u64>`
+
+Allocate a batch id for `batch_size` transfer requests.
+
+You must call `free_batch_id(batch_id)` after all tasks are done.
+
+#### `submit_transfer(batch_id, requests) -> anyhow::Result<()>`
+
+Submit a batch transfer request list.
+
+- No-op when `requests` is empty.
+- The wrapper converts each `TransferRequest` into the C representation (`transfer_request_t`).
+
+#### `get_transfer_status(batch_id, task_id) -> anyhow::Result<(i32, u64)>`
+
+Get status for one task in a batch.
+
+- `task_id` is an index inside the batch, typically `0..batch_size`.
+- Returns `(status_code, transferred_bytes)`.
+
+The `status_code` maps to values in `TransferStatusEnum` (represented as `i32`).
+
+#### `free_batch_id(batch_id) -> anyhow::Result<()>`
+
+Free a previously allocated batch id.
+
+## Minimal usage example (pseudo-code)
+
+The crate's `src/main.rs` contains a full benchmark-style example. The following sketch shows the typical control flow:
+
+```rust
+use std::ffi::c_void;
+use transfer_engine_rust::transfer_engine::{OpcodeEnum, TransferEngine, TransferRequest};
+
+fn main() -> anyhow::Result<()> {
+    let engine = TransferEngine::new("127.0.0.1:2379", "127.0.0.1", 12345)?;
+    engine.discover_topology()?;
+    engine.install_transport("tcp")?;
+
+    // Register local memory (example only; you must allocate and pin memory appropriately).
+    let mut buffer = vec![0u8; 4096];
+    engine.register_local_memory(buffer.as_mut_ptr() as *mut c_void, buffer.len(), "cpu:0")?;
+
+    let seg_id = engine.open_segment("target-seg".to_string())?;
+    let batch_id = engine.allocate_batch_id(1)?;
+
+    let mut reqs = [TransferRequest {
+        opcode: OpcodeEnum::Write,
+        source: buffer.as_mut_ptr() as *mut c_void,
+        target_id: seg_id,
+        target_offset: 0,
+        length: buffer.len() as u64,
+    }];
+
+    engine.submit_transfer(batch_id, &mut reqs)?;
+    let (status, bytes) = engine.get_transfer_status(batch_id, 0)?;
+    println!("status={status}, bytes={bytes}");
+
+    engine.free_batch_id(batch_id)?;
+    engine.close_segment(seg_id)?;
+    engine.unregister_local_memory(buffer.as_mut_ptr() as *mut c_void)?;
+    Ok(())
+}
+```
+
+## Safety & thread-safety
+
+- The wrapper marks `TransferEngine` as `Send + Sync`, but it owns an FFI handle (`transfer_engine_t`).
+- All pointer-based arguments must satisfy Rust’s aliasing and lifetime rules.
+- You must ensure registered memory remains valid until it is unregistered.
+
diff --git a/docs/source/conf.py b/docs/source/conf.py
index 817b116d..a505aa4f 100644
--- a/docs/source/conf.py
+++ b/docs/source/conf.py
@@ -41,35 +41,28 @@ extensions = [
     "sphinx.ext.linkcode",
     "sphinx.ext.intersphinx",
     "sphinx_copybutton",
-    "autodoc2",
     "myst_parser",
     "sphinxarg.ext",
     "sphinx_design",
     "sphinx_togglebutton",
+    "sphinxcontrib.mermaid",
 ]
 myst_enable_extensions = [
     "colon_fence",
     "fieldlist",
 ]
-autodoc2_packages = [
-    {
-        "path": "../../Mooncake",
-        "exclude_dirs": ["__pycache__", "third_party"],
-    },
-]
-autodoc2_output_dir = "api"
-autodoc2_render_plugin = "myst"
-autodoc2_hidden_objects = ["dunder", "private", "inherited"]
-autodoc2_sort_names = True
-autodoc2_index_template = None
-
+myst_fence_as_directive = ["mermaid"]
 # Add any paths that contain templates here, relative to this directory.
 templates_path = ['_templates']
 
 # List of patterns, relative to source directory, that match files and
 # directories to ignore when looking for source files.
 # This pattern also affects html_static_path and html_extra_path.
-exclude_patterns: list[str] = ["**/*.template.md", "**/*.inc.md"]
+exclude_patterns: list[str] = [
+    "**/*.template.md",
+    "**/*.inc.md",
+    "zh_archive/**",
+]
 
 # Exclude the prompt "$" when copying code
 copybutton_prompt_text = r"\$ "
diff --git a/docs/source/deployment/mooncake-store-deployment-guide.md b/docs/source/deployment/mooncake-store-deployment-guide.md
index 46f44f32..93fe6492 100644
--- a/docs/source/deployment/mooncake-store-deployment-guide.md
+++ b/docs/source/deployment/mooncake-store-deployment-guide.md
@@ -1,81 +1,228 @@
-# Mooncake Store Deployment & Operations Guide
+# Mooncake Store Deployment & Tuning Guide
 
-This page summarizes useful flags, environment variables, and HTTP endpoints to help advanced users tune Mooncake Master and observe metrics.
+This guide covers minimal deployment, and operational tuning of Mooncake Store.
 
-## Master Startup Flags (with defaults)
+## Architecture Overview
 
-- RPC Related
-  - `--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.
 
-- Metrics
-  - `--enable_metric_reporting` (bool, default `true`): Periodically log master metrics to INFO.
-  - `--metrics_port` (int, default `9003`): HTTP port for `/metrics` endpoints.
+![architecture](../image/mooncake-store-preview.png)
 
-- HTTP Metadata Server For Mooncake Transfer Engine
-  - `--enable_http_metadata_server` (bool, default `false`): Enable embedded HTTP metadata server.
-  - `--http_metadata_server_host` (str, default `0.0.0.0`): Metadata bind host.
-  - `--http_metadata_server_port` (int, default `8080`): Metadata TCP port.
+**Master Service** (`mooncake_master`): The central coordinator. It manages cluster membership, allocates object storage across client nodes, and enforces eviction/placement policies. Runs as a standalone process.
 
-- Allocation Strategy
-  - `--allocation_strategy` (str, default `random`): Memory allocation strategy for replica placement. Available options:
-    - `random`: Pure random selection across segments (baseline, fastest).
-    - `free_ratio_first`: Free-ratio-first strategy. Samples multiple candidates and selects those with highest free space ratio for better load balancing.
+**Client Node**: Each node contributes DRAM (and optionally VRAM/SSD) to form the distributed cache pool. Clients communicate with the master over RPC for control operations (`Put`/`Get`/`Remove`), but transfer actual data directly between each other via the Transfer Engine — the master is never in the data path.
 
-- Eviction and TTLs
-  - `--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.
+**Metadata Service**: A separate service (etcd, Redis, or HTTP) used by the Transfer Engine for peer discovery and configuration. The master's embedded HTTP metadata server can replace an external etcd/Redis for simple deployments. We also provide a P2P handshake mechanism (`P2PHANDSHAKE`) that enables decentralized metadata management by storing metadata locally on each node, eliminating the need for a centralized service — this is the simplest metadata handshake method and the recommended starting point (see [Quick Start](#quick-start)).
 
-- High Availability (optional)
-  - `--enable_ha` (bool, default `false`): Enable HA (requires etcd).
-  - `--etcd_endpoints` (str, default empty unless HA config): etcd endpoints, semicolon separated.
-  - `--client_ttl` (int64, default `10` s): Client alive TTL after last ping (HA mode).
-  - `--cluster_id` (str, default `mooncake_cluster`): Cluster ID for persistence in HA mode.
+For a detailed design discussion, see the [Mooncake Store Design](../design/mooncake-store.md).
 
-- Task Manager (optional)
-  - `--max_total_finished_tasks` (uint32, default `10000`): Maximum number of finished tasks to keep in memory. When this limit is reached, the oldest finished tasks will be pruned from memory.
-  - `--max_total_pending_tasks` (uint32, default `10000`): Maximum number of pending tasks that can be queued in memory. When this limit is reached, new task submissions will fail with `TASK_PENDING_LIMIT_EXCEEDED` error.
-  - `--max_total_processing_tasks` (uint32, default `10000`): Maximum number of tasks that can be processing simultaneously. When this limit is reached, no new tasks will be popped from the pending queue until some processing tasks complete.
-  - `--max_retry_attempts` (uint32, default `10`): Maximum number of retry attempts for failed tasks. Tasks that fail with `NO_AVAILABLE_HANDLE` error will be retried up to this many times before being marked as failed.
+---
 
-- DFS Storage (optional)
-  - `--root_fs_dir` (str, default empty): DFS mount directory for storage backend, used in Multi-layer Storage Support.
-  - `--global_file_segment_size` (int64, default `int64_max`): Maximum available space for DFS segments.
+## Quick Start
 
-- Snapshot / Restore (optional)
-  - `--enable_snapshot` (bool, default `false`): Enable periodic snapshot of master metadata data (effective when using the `offset` memory allocator).
-  - `--snapshot_interval_seconds` (uint64, default `600`): Interval in seconds between periodic snapshots of master data.
-  - `--snapshot_child_timeout_seconds` (uint64, default `300`): Timeout in seconds for each snapshot child process.
-  - `--snapshot_retention_count` (uint32, default `2`): Number of recent snapshots to keep. Older snapshots beyond this limit will be automatically deleted.
-  - `--snapshot_backend_type` (str, required when snapshot enabled): Snapshot storage backend type: `local` for local filesystem, `s3` for S3 storage.
-  - `--snapshot_backup_dir` (str, default empty): Optional local directory for snapshot backup. If empty (default), local backup is disabled. When set, it serves two purposes: (1) during snapshot persistence, data will be saved locally as a fallback if uploading to the backend fails; (2) during restore, downloaded metadata will also be saved to this directory as a local backup.
-  - `--enable_snapshot_restore` (bool, default `false`): Enable restore from the latest snapshot at master startup.
-  - **Environment variable** `MOONCAKE_SNAPSHOT_LOCAL_PATH` (**required** when `--snapshot_backend_type=local`): Persistent directory path for local snapshot storage. This variable **must** be set before starting the master; there is no default value. Example: `export MOONCAKE_SNAPSHOT_LOCAL_PATH=/data/mooncake_snapshots`.
+Deploy a minimal single-node Mooncake Store in three steps.
 
-  > **Warning: Managed Directory**
-  >
-  > The snapshot storage path (`MOONCAKE_SNAPSHOT_LOCAL_PATH` for local backend, or S3 bucket for S3 backend) is a **managed directory** exclusively controlled by the Mooncake snapshot system. **DO NOT store other files or data in this directory.** Old snapshots exceeding `--snapshot_retention_count` will be automatically and permanently deleted during cleanup. Use a dedicated, isolated directory for snapshot storage to avoid accidental data loss.
+### 1. Start the Metadata Service
 
-Example (enable embedded HTTP metadata and metrics):
+This quick start uses **P2P handshake** — the simplest option, with **nothing to start**: each node exchanges and stores Transfer Engine metadata locally during connection setup. You just pass the literal string `P2PHANDSHAKE` as the client's `metadata_server` (step 3).
+
+For large or long-lived clusters, use the master's embedded HTTP metadata server or an external etcd/Redis instead — see [Deployment Scenarios](#deployment-scenarios).
+
+### 2. Start the Master Service
+
+With P2P handshake the master needs no metadata-server flags:
+
+```bash
+mooncake_master
+```
+
+On success the master logs a single line like:
+
+```
+Master service started on port 50051, max_threads=4, ...
+```
+
+The master's default RPC port is `50051`. (To embed an HTTP metadata server instead of using P2P, add `--enable_http_metadata_server=true --http_metadata_server_port=8080`.)
+
+(start-a-store-client)=
+### 3. Start a Store Client
+
+A client contributes DRAM (and optionally SSD) to the cluster. The simplest way is to embed Mooncake in a Python process and call `store.setup(...)` with `metadata_server="P2PHANDSHAKE"`:
+
+```python
+from mooncake.store import MooncakeDistributedStore
+
+store = MooncakeDistributedStore()
+store.setup(
+    local_hostname="localhost",
+    metadata_server="P2PHANDSHAKE",           # decentralized; no metadata service
+    global_segment_size=3200 * 1024 * 1024,   # DRAM contributed to the cluster
+    local_buffer_size=512 * 1024 * 1024,      # Transfer Engine buffer
+    protocol="tcp",
+    rdma_devices="",                          # keyword is rdma_devices (not device_name)
+    master_server_addr="127.0.0.1:50051",     # keyword is master_server_addr
+)
+```
+
+There are **three ways** to run a client — programmatic (above), a standalone `mooncake_store_service` process (configured via `MOONCAKE_*`), and the `mooncake_client` real-client RPC process. See [Reference: Client Configuration & Tuning](#reference-client-configuration-tuning) for all three, with full parameter/env tables.
+
+**What just happened:**
+
+1. The client registered itself with the master via RPC.
+2. The master allocated a 3.2 GB segment on this node and added it to the cluster's memory pool.
+3. The client is now ready to serve `Put`/`Get`/`Remove` requests.
+
+### Run the Stress Benchmark
+
+Mooncake Store includes sample programs for validating C++ and Python integrations. The [stress benchmark script](gh-file:mooncake-store/tests/stress_cluster_benchmark.py) can be used to verify a two-role prefill/decode setup.
+
+Configure the script with command-line flags (run with `--help` for the full list):
+
+- `--local-hostname`: the local machine's reachable IP address or hostname.
+- `--metadata-server`: the Transfer Engine metadata service, e.g. `P2PHANDSHAKE`, `http://127.0.0.1:8080/metadata`, or an etcd address.
+- `--master-server`: the Mooncake Store master address. Use `IP:Port` in default mode, or `etcd://IP:Port;IP:Port;...;IP:Port` in etcd-backed HA mode.
+- `--protocol`: transport, `tcp` / `rdma` / `cxl` / `ascend` (defaults to `rdma`).
+
+Then start the roles:
+
+```bash
+python3 mooncake-store/tests/stress_cluster_benchmark.py --role prefill
+python3 mooncake-store/tests/stress_cluster_benchmark.py --role decode
+```
+
+For RDMA, topology auto-discovery and NIC filters can be passed through environment variables:
+
+```bash
+MC_MS_AUTO_DISC=1 MC_MS_FILTERS="mlx5_1,mlx5_2" python3 mooncake-store/tests/stress_cluster_benchmark.py --role prefill
+MC_MS_AUTO_DISC=1 MC_MS_FILTERS="mlx5_1,mlx5_2" python3 mooncake-store/tests/stress_cluster_benchmark.py --role decode
+```
+
+The absence of errors indicates successful data transfer.
+
+### Verify Installed Examples
+
+For a Python integration check, run `mooncake-store/tests/distributed_object_store_provider.py` after starting the metadata service and `mooncake_master`.
+
+For a C++ integration check, run `build/mooncake-store/tests/client_integration_test` after building tests and starting the required services.
+
+### Verify
+
+```bash
+# Health check — master metrics endpoint
+curl -s http://localhost:9003/metrics/summary
+
+# List registered clients
+# (exposed through the store's Python API or RPC)
+```
+
+---
+
+## Deployment Scenarios
+
+### Single-Node (TCP) — Development / Quick Evaluation
+
+The simplest deployment, as shown in [Quick Start](#quick-start). A single `mooncake_master` orchestrates clients over TCP. Suitable for development, testing, and single-host evaluation.
 
 ```bash
 mooncake_master \
   --enable_http_metadata_server=true \
   --http_metadata_server_host=0.0.0.0 \
-  --http_metadata_server_port=8080 \
-  --rpc_thread_num=64 \
-  --metrics_port=9003 \
-  --enable_metric_reporting=true
+  --http_metadata_server_port=8080
 ```
 
-Example (resolve the master RPC address from a stable interface name in a container):
+Limitation: the master is a single point of failure. If it crashes, cluster operations pause until it is restored.
+
+---
+
+### High-Availability (etcd) — Production HA
+
+Runs a cluster of master instances coordinated through etcd. If the leader fails, the remaining instances elect a new leader automatically.
+
+```bash
+# Start each master instance with:
+mooncake_master \
+  --enable_ha=true \
+  --etcd_endpoints="10.0.0.1:2379;10.0.0.2:2379;10.0.0.3:2379" \
+  --rpc_address=10.0.0.1
+```
+
+Each instance must specify its own reachable `--rpc_address`. The etcd cluster used for HA can be shared with or separate from the Transfer Engine's metadata etcd.
+
+**Client addressing:** to reach an HA cluster, clients must use the `etcd://` master-address form (so they can discover the current leader) instead of a single `IP:Port` — set `master_server_addr` (Method A) / `MOONCAKE_MASTER` (Method B) / `--master_server_address` (Method C) to `etcd://10.0.0.1:2379;10.0.0.2:2379;...`.
+
+---
+
+### High-Availability (Redis) — Alternative HA Backend
+
+Same HA semantics but using Redis instead of etcd for leader election:
+
+```bash
+mooncake_master \
+  --enable_ha=true \
+  --ha_backend_type=redis \
+  --ha_backend_connstring="redis://127.0.0.1:6379" \
+  --rpc_address=10.0.0.1
+```
+
+**Client addressing:** clients reach a Redis-backed HA cluster with the `redis://connstring` master-address form (e.g. `redis://127.0.0.1:6379`) for `master_server_addr` / `MOONCAKE_MASTER` / `--master_server_address`, instead of a single `IP:Port`.
+
+
+---
+
+### Snapshot & Restore — Backup / Disaster Recovery
+
+```{caution}
+Metadata Snapshot And Restore is experimental feature.
+```
+
+Periodically persist master metadata to local disk or S3, enabling recovery from a recent snapshot after a crash.
+
+```bash
+export MOONCAKE_SNAPSHOT_LOCAL_PATH=/data/mooncake_snapshots
+
+mooncake_master \
+  --enable_snapshot=true \
+  --snapshot_interval_seconds=300 \
+  --snapshot_retention_count=5 \
+  --snapshot_object_store_type=local \
+  --enable_snapshot_restore=true
+```
+
+---
+
+### Tiered Storage with SSD Offload — Cost-Effective Capacity
+
+Extends the cache pool from DRAM to SSD while keeping normal reads and writes on the distributed memory path. With `--enable_offload=true`, completed memory writes are queued for asynchronous SSD persistence through the master control plane. Set `--offload_on_evict=true` to defer that SSD write until the memory eviction path selects an object for reclamation. When `--promotion_on_hit=true`, SSD-only objects can be promoted back to DRAM after repeated reads; admission is gated by `--promotion_admission_threshold`.
+
+```bash
+mooncake_master \
+  --enable_offload=true \
+  --offload_on_evict=true \
+  --promotion_on_hit=true \
+  --promotion_admission_threshold=2 \
+  --root_fs_dir=/mnt/ssd_cache \
+  --enable_http_metadata_server=true \
+  --http_metadata_server_port=8080
+```
+
+---
+
+### CXL-Aware Allocation — Memory Tiering
+
+When the host has CXL-attached memory, the master can preferentially allocate new objects on the CXL tier, reserving local DRAM for latency-sensitive operations.
+
+```bash
+mooncake_master \
+  --enable_cxl=true \
+  --cxl_path=/dev/dax0.0 \
+  --cxl_size=17179869184 \
+  --allocation_strategy=cxl
+```
+
+---
+
+### Container / Dynamic Network Interface
+
+When the master runs in a container with a dynamic IP, use `--rpc_interface` to resolve the RPC address from a stable interface name:
 
 ```bash
 mooncake_master \
@@ -85,86 +232,558 @@ mooncake_master \
   --http_metadata_server_port=8080
 ```
 
-This resolves the current IPv4 address of `eth0` at startup and uses it as the final `rpc_address`.
+The master resolves the current IPv4 address of `eth0` at startup and uses it as the advertised RPC address.
 
-Example (use free-ratio-first allocation strategy for better load balancing):
 
-```bash
-mooncake_master \
-  --allocation_strategy=free_ratio_first \
-  --enable_http_metadata_server=true \
-  --http_metadata_server_port=8080
-```
-
-**Tips:**
-
-In addition to command-line flags, the Master also supports configuration via JSON and YAML files. For example:
-
-```bash
-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`:
-
-- `GET /metrics` — Prometheus format (`text/plain; version=0.0.4`).
-- `GET /metrics/summary` — Human-readable summary.
-
-Examples:
+The master exposes Prometheus-style metrics on `--metrics_port`:
 
 ```bash
+# Prometheus format
 curl -s http://<master_host>:9003/metrics
+
+# Human-readable summary
 curl -s http://<master_host>:9003/metrics/summary
 ```
 
-## Client/Engine Tuning (Env Vars, with defaults)
+When tenant quota is enabled, `/metrics` also includes per-tenant quota gauges and quota counters:
 
-- Topology discovery (Store Client → Transfer Engine)
-  - `MC_MS_AUTO_DISC` (default `1`): Auto-discover NIC/GPU topology. Set `0` to disable and provide `rdma_devices` manually.
-  - `MC_MS_FILTERS` (default empty): Optional comma-separated NIC whitelist when auto-discovery is enabled (e.g., `mlx5_0,mlx5_2`).
-  - If `MC_MS_AUTO_DISC=0`, pass `rdma_devices` (comma-separated) to the Python `setup(...)` call.
+- `mooncake_tenant_quota_requested_bytes{tenant_id}`
+- `mooncake_tenant_quota_effective_bytes{tenant_id}`
+- `mooncake_tenant_quota_used_bytes{tenant_id}`
+- `mooncake_tenant_quota_reserved_bytes{tenant_id}`
+- `mooncake_tenant_quota_committed_count{tenant_id}`
+- `mooncake_tenant_quota_metadata_object_count{tenant_id}`
+- `mooncake_tenant_quota_over_quota{tenant_id}`
+- `mooncake_tenant_quota_explicit_policy{tenant_id}`
+- `mooncake_tenant_quota_reject_total{tenant_id,reason}`
+- `mooncake_tenant_evict_bytes_total{tenant_id}`
+- `mooncake_tenant_quota_allocatable_capacity_bytes`
+- `mooncake_tenant_quota_requested_bytes_sum`
+- `mooncake_tenant_quota_effective_bytes_sum`
 
-- Transfer Engine metrics (disabled by default)
-  - `MC_TE_METRIC` (default `0`/unset): Set to `1` to enable periodic engine metrics logging. **Note:** Not supported when using Transfer Engine TENT.
-  - `MC_TE_METRIC_INTERVAL_SECONDS` (default `5`): Positive integer seconds between reports (effective only if metrics enabled).
+---
 
-- Client metrics (enabled by default)
-  - `MC_STORE_CLIENT_METRIC` (default `1`): Client-side metrics on by default; set `0` to disable entirely.
-  - `MC_STORE_CLIENT_METRIC_INTERVAL` (default `0`): Reporting interval in seconds; `0` collects but does not periodically report.
+## Tenant Quota Management
 
-- Local memcpy optimization (Store transfer path)
-  - `MC_STORE_MEMCPY` (default `0`/false): Set to `1` to prefer local memcpy when source/destination are on the same client.
+Tenant quota admission is disabled by default. Enable strict multi-tenant mode on the master when you want memory writes admitted against connector-managed per-tenant quota:
 
-## Set the Log Level for yalantinglibs coro_rpc and coro_http
-By default, the log level is set to warning. You can customize it using the following environment variable:
+```bash
+mooncake_master \
+  --enable_multi_tenants=true \
+  --tenant_quota_connector_type=file \
+  --tenant_quota_connector_uri=/etc/mooncake/tenant_quotas.yaml
+```
 
-`export MC_YLT_LOG_LEVEL=info`
+The v1 connector is a writable YAML file. The file must use schema version `1`; tenant names must be non-empty, unique, must not start with `_`, and must not contain NUL or control characters; quotas must be positive integers with optional `B`, `KB`, `MB`, `GB`, or `TB` units:
 
-This sets the log level for yalantinglibs (including coro_rpc and coro_http) to info.
+```yaml
+version: 1
 
-Available log levels: trace, debug, info, warn (or warning), error, and critical.
+tenants:
+  - name: tenant-a
+    quota: 200GB
+
+  - name: tenant-b
+    quota: 500GB
+```
+
+When strict multi-tenant mode is enabled, write requests must include a registered tenant. The `default` tenant is not special unless it is explicitly registered in the connector policy.
+
+The same HTTP port used for metrics exposes the tenant quota admin API:
+
+```bash
+# List tenant quota snapshots
+curl -s http://<master_host>:9003/api/v1/tenant_quotas
+
+# Query one tenant
+curl -s "http://<master_host>:9003/api/v1/tenant_quotas?tenant_id=tenant-a"
+
+# Upsert an explicit policy. Explicit tenant policies must be positive.
+curl -s -X PUT "http://<master_host>:9003/api/v1/tenant_quotas?tenant_id=tenant-a" \
+  -H 'Content-Type: application/json' \
+  -d '{"requested_quota_bytes":2147483648}'
+
+# Delete an explicit policy. The tenant must not own objects or quota usage.
+curl -s -X DELETE "http://<master_host>:9003/api/v1/tenant_quotas?tenant_id=tenant-a"
+```
+
+Each tenant quota snapshot returns:
+
+```json
+{
+  "success": true,
+  "data": {
+    "tenant_id": "tenant-a",
+    "requested_quota_bytes": 2147483648,
+    "effective_quota_bytes": 2147483648,
+    "used_bytes": 0,
+    "reserved_bytes": 0,
+    "committed_count": 0,
+    "metadata_object_count": 0,
+    "over_quota": false,
+    "has_explicit_policy": true
+  }
+}
+```
+
+In HA mode, quota admin requests are served only by the active master service. Standby, candidate, or inactive services return HTTP 503. If strict multi-tenant mode is disabled, the quota admin API returns HTTP 409 with `UNAVAILABLE_IN_CURRENT_MODE`. Deleting a non-empty tenant returns HTTP 409 with `TENANT_NOT_EMPTY`.
+
+---
 
 ## Quick Tips
 
 - 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.
+- For detailed SSD offload configuration (storage backends, eviction policies, io_uring), see the [SSD Offload guide](ssd-offload).
+- For NVMe-oF SSD pool configuration see the [NVMe-oF SSD Pool Deployment Guide](nvmf-ssd-deployment-guide)
+- For experimental 3FS (USRBIO) integration as a persistent storage backend, see the [3FS USRBIO Plugin guide](../getting_started/plugin-usage/3FS-USRBIO-Plugin).
+- For detailed monitoring and observation see [Observability](../getting_started/observability)
 
+:::{toctree}
+:maxdepth: 1
+:hidden:
+
+ssd-offload
+NvMe-Of SSD Pool<nvmf-ssd-deployment-guide>
+HF3FS Plugin (Experimental)<../getting_started/plugin-usage/3FS-USRBIO-Plugin>
+../getting_started/observability
+:::
 
 ---
 
-:::{toctree}
-:caption: Advanced Topics
-:maxdepth: 1
+## Reference: Master Startup Flags
 
-ssd-offload
-:::
+### RPC
+
+| Flag | Default | Description |
+|------|---------|-------------|
+| `--rpc_port` | `0` → effective `50051` | RPC listen port. The literal default is `0`, which falls back to the deprecated `--port` (default `50051`) |
+| `--rpc_thread_num` | `0` → effective `min(4, CPU cores)` | RPC worker threads. The literal default is `0`, which falls back to the deprecated `--max_threads` → `min(4, CPU cores)` |
+| `--rpc_address` | `0.0.0.0` | RPC bind address |
+| `--rpc_interface` | empty | Network interface to resolve RPC address at startup (overrides `--rpc_address`) |
+| `--rpc_conn_timeout_seconds` | `0` | Idle connection timeout; `0` disables |
+| `--rpc_enable_tcp_no_delay` | `true` | Enable TCP_NODELAY |
+
+### Logging
+
+The master uses glog. When `--log_dir` is set, all severities are merged into a single journal file in that directory (`mooncake_master.INFO.<date>-<time>.<pid>`), reachable through the stable `mooncake_master.INFO` symlink.
+
+glog's standard flags (`--log_dir`, `--max_log_size`, `--logtostderr`, ...) control the rest.
+
+### Metrics
+
+| Flag | Default | Description |
+|------|---------|-------------|
+| `--enable_metric_reporting` | `true` | Periodically log master metrics |
+| `--metrics_port` | `9003` | HTTP port for `/metrics` endpoints |
+
+### HTTP Metadata Server (Embedded)
+
+| Flag | Default | Description |
+|------|---------|-------------|
+| `--enable_http_metadata_server` | `false` | Enable embedded HTTP metadata server |
+| `--http_metadata_server_host` | `0.0.0.0` | Metadata bind host |
+| `--http_metadata_server_port` | `8080` | Metadata TCP port |
+| `--enable_metadata_cleanup_on_timeout` | `false` | Delete a client's stale HTTP metadata (`mooncake/[<cluster>/]ram/<segment>` and `mooncake/[<cluster>/]rpc_meta/<segment>`) when its heartbeat times out (see below) |
+
+### Stale Metadata Cleanup on Client Timeout
+
+When a client crashes or is force-killed (`kill -9`, OOM, node failure), it cannot
+run its normal cleanup, leaving stale entries on the HTTP metadata server
+(`mooncake/[<cluster>/]ram/<segment>` and `mooncake/[<cluster>/]rpc_meta/<segment>`).
+The HTTP metadata server has no heartbeat of its own, so these entries linger and
+can mislead nodes that later connect or restart with different RDMA parameters.
+
+With `--enable_metadata_cleanup_on_timeout=true`, the Master Service reuses its
+existing client-heartbeat monitor: when a client's `--client_ttl` expires, in
+addition to unmounting the segment it also removes that client's `ram/` and
+`rpc_meta/` keys from the HTTP metadata server. It supports both deployment
+topologies:
+
+- **Co-located** (`--enable_http_metadata_server=true`): the master removes the
+  keys via a direct in-process call (no network overhead).
+- **Separately deployed** HTTP metadata server: the master derives the metadata
+  server address from the cluster's existing configuration and removes the keys
+  via HTTP `DELETE`. The address is read, in priority order, from:
+  1. the `MOONCAKE_TE_META_DATA_SERVER` environment variable (the same Transfer
+     Engine metadata connection string the clients use, e.g.
+     `http://host:8080/metadata`), then
+  2. the `metadata_server` field of the JSON file pointed to by
+     `MOONCAKE_CONFIG_PATH`.
+
+Notes:
+- Only `http(s)` metadata servers are supported; `etcd`/`redis`/`P2PHANDSHAKE`
+  backends are not cleaned up (a warning is logged and cleanup stays disabled).
+- The feature is opt-in and best-effort: if no co-located server is enabled and
+  no HTTP metadata address can be derived, the master logs a warning and
+  disables cleanup. Remote `DELETE` failures are logged but never block the
+  client-monitor thread or the main process.
+- Respects `MC_METADATA_CLUSTER_ID` for custom key prefixes (matching the
+  Transfer Engine).
+
+```bash
+# Co-located metadata server
+mooncake_master \
+  --enable_http_metadata_server=true \
+  --enable_metadata_cleanup_on_timeout=true \
+  --client_ttl=10
+
+# Separately-deployed HTTP metadata server (address derived from the env var)
+export MOONCAKE_TE_META_DATA_SERVER=http://metadata-host:8080/metadata
+mooncake_master \
+  --enable_metadata_cleanup_on_timeout=true \
+  --client_ttl=10
+```
+
+### Memory Allocator
+
+| Flag | Default | Description |
+|------|---------|-------------|
+| `--memory_allocator` | `offset` | Memory allocator: `offset` (default) or `cachelib` |
+
+### Allocation Strategy
+
+| Flag | Default | Description |
+|------|---------|-------------|
+| `--allocation_strategy` | `random` | `random` (pure random, fastest), `free_ratio_first` (best load balance), or `cxl` (prefer CXL memory) |
+
+### PutStart Timeouts
+
+| Flag | Default | Description |
+|------|---------|-------------|
+| `--put_start_discard_timeout_sec` | `30` | Seconds before an uncompleted `PutStart` is discarded |
+| `--put_start_release_timeout_sec` | `600` (10 min) | Seconds before `PutStart`-allocated space is released
+
+### Eviction & TTLs
+
+| Flag | Default | Description |
+|------|---------|-------------|
+| `--default_kv_lease_ttl` | `5000` ms | Lease TTL for KV objects. Supports `5000ms`, `5s`, `30m`, `1h` |
+| `--default_kv_soft_pin_ttl` | `1800000` ms | Soft pin TTL (30 min) |
+| `--allow_evict_soft_pinned_objects` | `true` | Allow evicting soft-pinned objects |
+| `--eviction_ratio` | `0.05` | Fraction evicted at high watermark |
+| `--eviction_high_watermark_ratio` | `0.95` | Usage ratio triggering eviction |
+| `--client_ttl` | `10` s | Seconds before a silent client is considered disconnected |
+
+### Tenant Quota
+
+| Flag | Default | Description |
+|------|---------|-------------|
+| `--enable_multi_tenants` | `false` | Enable strict tenant registration and per-tenant memory quota admission |
+| `--tenant_quota_connector_type` | `file` | Tenant quota policy connector type |
+| `--tenant_quota_connector_uri` | empty | Connector URI; for `file`, the writable YAML policy path |
+
+### High Availability
+
+**Master Node High Availability**
+| Flag | Default | Description |
+|------|---------|-------------|
+| `--enable_ha` | `false` | Enable HA mode |
+| `--ha_backend_type` | `etcd` | HA backend: `etcd`, `redis`, or `k8s` |
+| `--ha_backend_connstring` | empty | HA backend connection string |
+| `--etcd_endpoints` | empty | etcd endpoints, semicolon separated (when `--ha_backend_type=etcd`) |
+| `--cluster_id` | `mooncake_cluster` | Cluster ID for HA persistence |
+
+```{caution}
+Metadata Snapshot And Restore is experimental feature.
+```
+
+**Metadata Snapshot And Restore**
+
+| Flag | Default | Description |
+|------|---------|-------------|
+| `--enable_snapshot` | `false` | Enable periodic metadata snapshot |
+| `--snapshot_interval_seconds` | `600` (10 min) | Interval between snapshots |
+| `--snapshot_child_timeout_seconds` | `300` (5 min) | Timeout per snapshot child process |
+| `--snapshot_retention_count` | `2` | Number of recent snapshots retained |
+| `--snapshot_object_store_type` | required | Object store: `local` or `s3` |
+| `--snapshot_catalog_store_type` | empty | Catalog store: `embedded` or `redis` |
+| `--snapshot_catalog_store_connstring` | empty | Catalog store connection string (required for `redis`) |
+| `--snapshot_backup_dir` | empty | Optional local backup directory |
+| `--enable_snapshot_restore` | `false` | Restore from latest snapshot at startup |
+
+**Environment variable:** `MOONCAKE_SNAPSHOT_LOCAL_PATH` (required when `--snapshot_object_store_type=local`) — persistent directory for local snapshots.
+
+```{warning}
+The snapshot storage path is a **managed directory** exclusively controlled by Mooncake. Old snapshots exceeding `--snapshot_retention_count` are automatically deleted. Use a dedicated directory to avoid data loss.
+```
+
+### Task Manager
+
+| Flag | Default | Description |
+|------|---------|-------------|
+| `--max_total_finished_tasks` | `10000` | Max finished tasks kept in memory |
+| `--max_total_pending_tasks` | `10000` | Max queued pending tasks |
+| `--max_total_processing_tasks` | `10000` | Max simultaneously processing tasks |
+| `--pending_task_timeout_sec` | `300` (5 min) | Timeout for pending tasks (`0` = no timeout) |
+| `--processing_task_timeout_sec` | `300` (5 min) | Timeout for processing tasks (`0` = no timeout) |
+| `--max_retry_attempts` | `10` | Max retries for failed tasks (`NO_AVAILABLE_HANDLE`) |
+
+### Offload / Tiered Storage
+
+Flags for controlling data movement between DRAM and SSD.
+
+| Flag | Default | Description |
+|------|---------|-------------|
+| `--enable_offload` | `false` | Enable offload from DRAM to SSD |
+| `--offload_on_evict` | `false` | Defer offload to eviction time rather than at `Put` |
+| `--offload_force_evict` | `false` | Force-evict objects exceeding capacity without offload |
+| `--offloading_queue_limit` | `50000` | Max number of objects allowed in the offloading queue per local disk segment. Increase to allow more objects to be offloaded to SSD before force-eviction kicks in |
+| `--offload_cap_ratio` | `0.5` | Per-cycle offload cap as a fraction of `offloading_queue_limit` (range `[0.0, 1.0]`). Controls how many objects can be queued for offload in a single eviction cycle before falling back to force-evict |
+| `--promotion_on_hit` | `false` | Promote SSD-resident keys to DRAM on read hit |
+| `--promotion_admission_threshold` | `2` | Min CountMinSketch count to allow promotion (`1` = disable gating) |
+| `--promotion_max_per_heartbeat` | `1` | Max promotion tasks handed to a single client per heartbeat. Each task is a synchronous SSD-read + RDMA-write on the client; serializing them avoids blocking past the client-liveness window |
+| `--promotion_queue_limit` | `50000` | Max in-flight promotion tasks |
+| `--quota_bytes` | `0` (90% of capacity) | Storage quota in bytes |
+| `--enable_disk_eviction` | `true` | Enable disk eviction |
+
+Start with `--enable_offload=true` for eager asynchronous SSD persistence after `Put` completion. Add `--offload_on_evict=true` when you want SSD writes to happen only when memory pressure selects an object for eviction. Add `--promotion_on_hit=true` to allow hot SSD-only data to be promoted back to DRAM, and tune `--promotion_admission_threshold` to control how many observed reads are required before promotion is queued.
+
+When `--offload_on_evict=true` is active, each `BatchEvict` cycle can queue at most `offloading_queue_limit * offload_cap_ratio` objects for SSD offload (default: `50000 * 0.5 = 25000`); objects exceeding this cap fall back to force-evict (discard) if `--offload_force_evict=true`, otherwise they remain in memory. For SSD-heavy workloads where NVMe bandwidth is underutilized while the KV-cache hit rate suffers, raise both `--offloading_queue_limit` and `--offload_cap_ratio` so more objects per cycle are actually persisted to SSD instead of discarded. Example: `--offloading_queue_limit=500000 --offload_cap_ratio=0.8` yields a per-cycle cap of `400000` (vs the default `25000`).
+
+### CXL Memory
+
+| Flag | Default | Description |
+|------|---------|-------------|
+| `--enable_cxl` | `false` | Enable CXL memory support |
+| `--cxl_path` | `/dev/dax0.0` | DAX device path for CXL memory |
+| `--cxl_size` | `8GB` (`8589934592`) | CXL memory size in bytes |
+
+When `--allocation_strategy=cxl` is set alongside `--enable_cxl=true`, the master preferentially allocates new objects on CXL memory.
+
+### DFS Storage
+
+| Flag | Default | Description |
+|------|---------|-------------|
+| `--root_fs_dir` | empty | DFS mount directory for multi-layer storage backend |
+| `--global_file_segment_size` | `INT64_MAX` (unlimited) | Max available space for DFS segments; default does not cap DFS usage |
+
+### NoF (NVMe-oF SSD Pool)
+
+```{caution}
+NVMe-oF SSD Pool (NoF) is an experimental feature.
+```
+
+Master-side flags for the NVMe-oF SSD pool. They control eviction within the NoF SSD tier and the heartbeat used to detect and unmount unresponsive NoF segments. For the client-side NoF I/O tuning (`MC_NOF_*`), see the [NVMe-oF SSD Pool Deployment Guide](nvmf-ssd-deployment-guide.md).
+
+| Flag | Default | Description |
+|------|---------|-------------|
+| `--nof_eviction_ratio` | `0.05` | Fraction of objects evicted when NoF SSD space is full |
+| `--nof_eviction_high_watermark_ratio` | `0.95` | Usage ratio that triggers eviction in the NoF SSD tier |
+| `--nof_heartbeat_interval_sec` | `10` | How often the master probes each mounted NoF segment |
+| `--nof_heartbeat_probe_timeout_ms` | `1000` | Timeout for a single NoF heartbeat probe |
+| `--nof_heartbeat_failures_threshold` | `3` | Consecutive NoF heartbeat failures before a segment is unmounted |
+
+### Master Configuration File
+
+In addition to CLI flags, the master accepts JSON/YAML config files:
+
+```bash
+mooncake_master --config_path=mooncake-store/conf/master.yaml
+```
+
+```yaml
+rpc_interface: "eth0"
+rpc_port: 50051
+```
+
+---
+
+(reference-client-configuration-tuning)=
+## Reference: Client Configuration & Tuning
+
+A client is configured through one of the **methods** introduced in [Start a Store Client](#start-a-store-client), plus a shared family of engine-tuning variables:
+
+- **Method A — Programmatic (`setup()` arguments)**: you pass configuration as explicit Python arguments. `MOONCAKE_*` variables are **not** read in this method.
+- **Method B — Service / Integration (`MOONCAKE_*` + CLI)**: `mooncake.mooncake_store_service` and the vLLM/SGLang connectors read `MOONCAKE_*` environment variables (via `MooncakeConfig`).
+- **Method C — Resource-owning real client (`mooncake_client`)**: configured through `mooncake_client` CLI flags (see the **Method C** subsection below).
+- **Engine runtime tuning (`MC_*`)**: low-level variables read by the C++ Transfer Engine / store client at runtime. They are orthogonal to the above and **apply to all methods**.
+
+The Method A arguments and the `MOONCAKE_*` variables are the **same logical fields in two forms** (Method B maps onto Method A); note that the `mooncake_client` CLI (Method C) uses yet another spelling for some of them (e.g. `--device_names`, `--master_server_address`).
+
+### Method A — Programmatic (`setup()` arguments)
+
+Arguments of `MooncakeDistributedStore.setup(...)`:
+
+| Argument | Type | Default | Description |
+|----------|------|---------|-------------|
+| `local_hostname` | str | required | This node's hostname / IP |
+| `metadata_server` | str | required | `P2PHANDSHAKE` / `http://…:8080/metadata` / etcd address |
+| `global_segment_size` | int (bytes) | required | DRAM contributed to the cluster (the sample uses 3.2 GB) |
+| `local_buffer_size` | int (bytes) | required | Transfer Engine buffer |
+| `protocol` | str | required | `tcp` / `rdma` / `efa` / `cxl` / `ascend` |
+| `rdma_devices` | str | required | RDMA NIC(s), comma-separated (pass `""` for non-RDMA). **Keyword is `rdma_devices`, not `device_name`** |
+| `master_server_addr` | str | required | Master `host:port`. **Keyword is `master_server_addr`, not `master_server_address`** |
+| `engine` | TransferEngine | `None` | *(advanced)* Reuse an existing Transfer Engine instance instead of creating one |
+| `enable_ssd_offload` | bool | `false` | *(advanced)* Enable client-side SSD offload |
+| `ssd_offload_path` | str | empty | *(advanced)* SSD offload directory |
+| `tenant_id` | str | `default` | *(advanced)* Tenant identifier |
+
+```{note}
+The first seven arguments have **no Python default** — the C++ defaults are not exposed by the pybind binding, so they must all be supplied (a bare `setup(local_hostname, metadata_server)` raises `TypeError`). Only `engine` / `enable_ssd_offload` / `ssd_offload_path` / `tenant_id` are optional. Also, in Method A only the `MC_*` engine variables below have any effect — `MOONCAKE_*` are ignored.
+```
+
+### Method B — Service / Integration (`MOONCAKE_*` + CLI)
+
+`python -m mooncake.mooncake_store_service` (and the vLLM/SGLang connectors) build their configuration through `MooncakeConfig`, resolved in this order:
+
+1. `--config <path>` CLI argument → load from that JSON file.
+2. Otherwise `MOONCAKE_CONFIG_PATH` (if set) → load from that file; else read the `MOONCAKE_*` variables below.
+3. `-D key=value` CLI overrides individual fields (keys must match the `MooncakeConfig` field names, e.g. `-Dmaster_server_address=...`).
+
+```{note}
+The store service CLI only accepts `--config`, `-D/--define`, `--port`, and `--max-wait-time`. There are **no** `--local_hostname` / `--metadata_server` / `--master_server` flags — use the `MOONCAKE_*` variables (or `-D`) instead.
+```
+
+| Variable | Maps to (`setup()` arg) | Default | Description |
+|----------|-------------------------|---------|-------------|
+| `MOONCAKE_MASTER` | `master_server_addr` | — (required unless `MOONCAKE_CONFIG_PATH`) | Master `host:port` |
+| `MOONCAKE_TE_META_DATA_SERVER` | `metadata_server` | `P2PHANDSHAKE` | `P2PHANDSHAKE` / `http://…:8080/metadata` / etcd address |
+| `MOONCAKE_PROTOCOL` | `protocol` | `tcp` | `tcp` / `rdma` / `efa` / `cxl` / `ascend` |
+| `MOONCAKE_DEVICE` | `rdma_devices` | empty | RDMA/EFA device(s), comma-separated; `auto-discovery` supported |
+| `MOONCAKE_GLOBAL_SEGMENT_SIZE` | `global_segment_size` | `3355443200` (3.125 GiB) | DRAM contributed; accepts byte integer **or** suffixed form like `500gb` |
+| `MOONCAKE_LOCAL_BUFFER_SIZE` | `local_buffer_size` | `1073741824` (1 GiB) | Transfer Engine buffer; same parsing as above |
+| `MOONCAKE_LOCAL_HOSTNAME` | `local_hostname` | `localhost` | |
+| `MOONCAKE_OFFLOAD_ENABLED` | `enable_ssd_offload` | `false` | Client-side SSD offload |
+| `MOONCAKE_OFFLOAD_FILE_STORAGE_PATH` | `ssd_offload_path` | empty | Offload directory |
+| `MOONCAKE_CONFIG_PATH` | — | unset | Path to a JSON config file (takes precedence over the variables above) |
+
+```{note}
+`MooncakeConfig` (Method B) defaults `global_segment_size`/`local_buffer_size` to 3.125 GiB / 1 GiB. A direct `setup()` (Method A) has **no** default for these — they are required arguments. Unlike `MC_STORE_LOCAL_HOT_CACHE_SIZE` (raw bytes only), `MOONCAKE_GLOBAL_SEGMENT_SIZE` / `MOONCAKE_LOCAL_BUFFER_SIZE` accept human-readable suffixes (`kb`/`mb`/`gb`/…) because they are parsed by `MooncakeConfig`.
+```
+
+**Launch examples:**
+
+```bash
+# P2P handshake
+MOONCAKE_MASTER=127.0.0.1:50051 \
+MOONCAKE_TE_META_DATA_SERVER=P2PHANDSHAKE \
+python -m mooncake.mooncake_store_service
+
+# HTTP metadata server
+MOONCAKE_MASTER=127.0.0.1:50051 \
+MOONCAKE_TE_META_DATA_SERVER=http://127.0.0.1:8080/metadata \
+python -m mooncake.mooncake_store_service
+```
+
+Or via a JSON config file. The service also exposes a lightweight HTTP API (on `--port`, default `8080`) for manual `Get`/`Put` debugging:
+
+```json
+{
+  "local_hostname": "localhost",
+  "metadata_server": "http://127.0.0.1:8080/metadata",
+  "global_segment_size": 268435456,
+  "local_buffer_size": 268435456,
+  "protocol": "tcp",
+  "device_name": "",
+  "master_server_address": "127.0.0.1:50051"
+}
+```
+
+```bash
+python -m mooncake.mooncake_store_service --config=<config_path> --port=8081
+```
+
+### Method C — Resource-owning Real Client (`mooncake_client`)
+
+Run the `mooncake_client` binary as a standalone RPC process that owns storage resources; application processes (vLLM / SGLang) use lightweight **dummy clients** to forward requests to it. It connects to the master and listens on port `50052` by default.
+
+```bash
+mooncake_client \
+  --global_segment_size="4GB" \
+  --master_server_address="127.0.0.1:50051" \
+  --metadata_server="http://127.0.0.1:8080/metadata"
+```
+
+| Flag | Default | Description |
+|------|---------|-------------|
+| `--host` | `0.0.0.0` | Client service bind host |
+| `--port` | `50052` | Client service listen port |
+| `--global_segment_size` | `4 GB` | Global segment size contributed by the client |
+| `--master_server_address` | `127.0.0.1:50051` | Master service address |
+| `--metadata_server` | `http://127.0.0.1:8080/metadata` | Transfer Engine metadata service |
+| `--protocol` | `tcp` | Transfer protocol |
+| `--device_names` | empty | Transfer device name(s), comma-separated |
+| `--threads` | `1` | Client worker thread count |
+| `--enable_offload` | `false` | Enable client-side SSD offload |
+| `--start_offload_rpc_server` | `true` | Start the offload RPC server for dummy clients |
+
+### Engine Runtime Tuning (`MC_*`)
+
+The following `MC_*` variables are read directly by the engine/client at runtime and **apply to all methods (A, B, and C)**.
+
+#### Runtime Protocol
+
+| Variable | Default | Description |
+|----------|---------|-------------|
+| `MC_RPC_PROTOCOL` | `tcp` | RPC transport protocol between master and clients: `tcp` or `rdma` |
+| `MC_RPC_TIMEOUT_MS` | `30000` | Per-request deadline (ms) for all client→master RPCs. Applies uniformly to every RPC method. A negative value disables the timeout. On expiry the call returns `RPC_TIMEOUT` |
+| `MC_RPC_CONNECT_TIMEOUT_MS` | `30000` | Connection-establishment timeout (ms) for the master RPC client |
+| `MC_USE_TENT` / `MC_USE_TEV1` | unset | Set to any value to enable the TENT (next-gen) transfer engine |
+| `MC_STORE_CLUSTER_ID` | unset | Cluster ID label attached to client metrics |
+
+#### Topology Discovery
+
+| Variable | Default | Description |
+|----------|---------|-------------|
+| `MC_MS_AUTO_DISC` | unset | Auto-discover NIC/GPU topology. Set `1` to force on, `0` to provide `rdma_devices` manually. When unset, auto-discovery is **off** except for `rdma`/`efa` protocols when no `rdma_devices` are given, where it defaults **on**. Ignored when TENT is enabled |
+| `MC_MS_FILTERS` | empty | Comma-separated NIC whitelist (e.g., `mlx5_0,mlx5_2`) |
+
+When `MC_MS_AUTO_DISC=0`, pass `rdma_devices` (comma-separated) to the Python `setup()` call.
+
+#### Transfer Engine Metrics (disabled by default)
+
+| Variable | Default | Description |
+|----------|---------|-------------|
+| `MC_TE_METRIC` | `0` | Set to `1` to enable engine metrics. Not supported with TENT |
+| `MC_TE_METRIC_INTERVAL_SECONDS` | `5` | Seconds between reports |
+
+#### Client Metrics (enabled by default)
+
+| Variable | Default | Description |
+|----------|---------|-------------|
+| `MC_STORE_CLIENT_METRIC` | `1` | Set `0` to disable |
+| `MC_STORE_CLIENT_METRIC_INTERVAL` | `0` | Reporting interval; `0` collects but does not periodically report |
+| `MC_STORE_CLIENT_MIN_PORT` | `12300` | Min local port for client connections |
+| `MC_STORE_CLIENT_MAX_PORT` | `14300` | Max local port for client connections |
+
+#### Local Hot Cache
+
+Local hot cache provides a DRAM read cache on top of SSD-resident objects for faster access.
+
+| Variable | Default | Description |
+|----------|---------|-------------|
+| `MC_STORE_LOCAL_HOT_CACHE_SIZE` | unset | Size of the local hot cache **in raw bytes** (decimal integer, e.g., `8589934592` for 8 GB). Suffixed forms like `"8gb"` are **not** parsed. Set to a positive value to enable the hot cache |
+| `MC_STORE_LOCAL_HOT_BLOCK_SIZE` | `16777216` (16 MB) | Block size for hot cache **in raw bytes** (decimal integer, e.g., `2097152` for 2 MB). Suffixed forms like `"2mb"` are **not** parsed. Only read when the hot cache is enabled |
+| `MC_STORE_LOCAL_HOT_CACHE_USE_SHM` | unset | Set `1` to use memfd-backed shared memory |
+| `MC_STORE_LOCAL_HOT_ADMISSION_THRESHOLD` | unset | Minimum CountMinSketch count before a key is admitted to hot cache |
+
+#### Local Memory Optimization
+
+| Variable | Default | Description |
+|----------|---------|-------------|
+| `MC_STORE_MEMCPY` | auto | Prefer local memcpy when source/destination are on the same client. When unset, auto-detected by transport: **enabled** in a TCP-only environment, **disabled** when an RDMA/other transport is available. Accepts `1`/`true`/`yes`/`on` or `0`/`false`/`no`/`off` to override |
+| `MC_STORE_CLIENT_SETUP_RETRIES` | `20` | Number of times to retry client registration on failure |
+| `MC_CXL_DEV_SIZE` | unset | CXL device size in raw bytes for client-side allocation. **Required when `protocol="cxl"`** — the client aborts at startup if it is missing |
+
+#### MMap Buffer & HugePages
+
+| Variable | Default | Description |
+|----------|---------|-------------|
+| `MC_STORE_USE_HUGEPAGE` | unset | Set `1` to request HugeTLB-backed `mmap()` |
+| `MC_STORE_HUGEPAGE_SIZE` | `2MB` | Supported: `2MB`, `1GB` |
+| `MC_MMAP_ARENA_POOL_SIZE` | unset | Pre-allocated arena pool size (e.g., `8gb`). Explicitly set to enable the arena |
+| `MC_DISABLE_MMAP_ARENA` | unset | Disable arena, fall back to per-call `mmap()`. Accepts `1`/`true`/`yes`/`on` (or `0`/`false`/`no`/`off`) |
+
+#### yalantinglibs Log Level
+
+```bash
+export MC_YLT_LOG_LEVEL=info
+```
+
+Available: `trace`, `debug`, `info`, `warn` (or `warning`), `error`, `critical`. When unset (or set to an unrecognized value), the level defaults to `warn`.
diff --git a/docs/source/deployment/nvmf-ssd-deployment-guide.md b/docs/source/deployment/nvmf-ssd-deployment-guide.md
new file mode 100644
index 00000000..eb855838
--- /dev/null
+++ b/docs/source/deployment/nvmf-ssd-deployment-guide.md
@@ -0,0 +1,286 @@
+# Mooncake NVMe-oF SSD Pool Deployment Guide
+
+## Overview
+
+This guide shows how to attach an NVMe-oF SSD pool to Mooncake Store. The
+deployment has two main phases:
+
+- Start Mooncake services built with NoF support enabled.
+- Create SPDK NVMe-oF targets on SSD pool nodes and register their namespaces
+  with the Mooncake master.
+
+After registration, the master reports the registered NVMe-oF namespaces as a
+remote SSD pool in its metrics, and clients can place NoF replicas through
+Mooncake Store.
+
+## 1. Build Mooncake with NoF Support
+
+Follow the "Build with NVMe-oF SSD Pool" section in the
+[Build Guide](../getting_started/build.md) to install SPDK dependencies and
+build Mooncake with `-DUSE_NOF=ON`.
+
+## 2. Deploy Mooncake Services
+
+### 2.1 Node Topology
+
+- **Mooncake service node**: 192.168.65.81. This node runs the master, metadata, and store services.
+- **SSD pool nodes**: 192.168.65.56 and 192.168.65.57. These nodes provide SSD storage resources.
+
+### 2.2 Deploy the Master Service
+
+```bash
+mooncake_master --rpc_address=192.168.65.81
+```
+
+### 2.3 Deploy the Metadata Service
+
+```bash
+python3 -m mooncake.http_metadata_server --host=192.168.65.81 --port=8080
+```
+
+If an aiohttp-related error occurs during startup, install aiohttp:
+
+```bash
+pip3 install aiohttp
+```
+
+### 2.4 Deploy the Store Service
+
+#### Configure `store_service.json`
+
+Create `store_service.json` under `/home`:
+
+```json
+{
+  "local_hostname": "localhost",
+  "metadata_server": "http://192.168.65.81:8080/metadata",
+  "master_server_address": "192.168.65.81:50051",
+  "protocol": "rdma",
+  "device_name": "mlx5_0",
+  "global_segment_size": "50gb",
+  "local_buffer_size": 0
+}
+```
+
+**Notes**:
+
+- `device_name`: Run `ibv_devices` on node 192.168.65.81 to check the RDMA device name.
+
+#### Start the Service
+
+The store service initializes the SPDK environment during startup. Configure hugepages on the store service node, 192.168.65.81:
+
+```bash
+echo 512 > /proc/sys/vm/nr_hugepages
+```
+
+Only a small number of hugepages is required during startup. In most cases, 512 hugepages are sufficient.
+
+Start the store service:
+
+```bash
+python3 -m mooncake.mooncake_store_service --config=/home/store_service.json --port=8081
+```
+
+If a timeout error occurs during startup, check whether a proxy is configured on node 192.168.65.81. If a proxy is configured, unset the proxy configuration and try again.
+
+## 3. Deploy the NVMe-oF SSD Pool
+
+### 3.1 Prerequisites
+
+1. Configure passwordless SSH login from the Mooncake node, 192.168.65.81, to the SSD pool nodes, 192.168.65.56 and 192.168.65.57.
+   See [OpenSSH key-based authentication](https://help.ubuntu.com/community/SSH/OpenSSH/Keys).
+2. Build SPDK on each SSD pool node in advance.
+   See [SPDK build instructions](https://github.com/spdk/spdk/blob/master/README.md#build).
+
+### 3.2 Install SSH Dependencies
+
+```bash
+python3 -m pip install "paramiko>=3.4.0"
+```
+
+### 3.3 Deploy the SSD Pool
+
+#### Deployment Command
+
+```bash
+python3 -m mooncake.spdk_tgt_create \
+    --spdk_target_info="ip:192.168.65.56 path:/home/spdk pci:0000:01:00.0,0000:02:00.0" \
+    --spdk_target_info="ip:192.168.65.57 path:/home/spdk" \
+    --core-mask=0xff \
+    --transport-type=RDMA \
+    --max-queue-depth=128 \
+    --max-io-qpairs-per-ctrlr=127 \
+    --max-io-size=4096 \
+    --in-capsule-data-size=131072 \
+    --io-unit-size=131072 \
+    --max-aq-depth=128 \
+    --num-shared-buffers=4096 \
+    --buf-cache-size=32
+```
+
+#### Parameters
+
+| Parameter | Description |
+|-----------|-------------|
+| `ip` | IP address of the target node. |
+| `path` | SPDK installation path on the target node. |
+| `pci` | PCI addresses of SSDs to register with the target. Use commas to separate multiple PCI addresses. If this field is omitted, SPDK-ready or unmounted NVMe devices on the target node are registered. |
+| `--core-mask` | CPU core mask used to start `nvmf_tgt` with `-m`. The default value is `0xff`. |
+
+**Tip**: Run `/path/scripts/setup.sh status` on a target node to list available PCI addresses.
+
+#### Transport Options
+
+The transport options are passed to the SPDK `nvmf_create_transport` RPC. If an option is not specified, the tool uses the default value listed below.
+
+| Option | Default | Description |
+|--------|---------|-------------|
+| `--transport-type` | `RDMA` | NVMe-oF transport type. |
+| `--max-queue-depth` | `128` | Maximum number of outstanding I/O operations per queue. |
+| `--max-io-qpairs-per-ctrlr` | `127` | Maximum number of I/O queue pairs per controller. |
+| `--max-io-size` | `4096` | Maximum I/O size, in bytes. |
+| `--in-capsule-data-size` | `131072` | Maximum in-capsule data size, in bytes. |
+| `--io-unit-size` | `131072` | I/O unit size, in bytes. |
+| `--max-aq-depth` | `128` | Maximum number of admin commands per admin queue. |
+| `--num-shared-buffers` | `4096` | Number of pooled data buffers available to the transport. |
+| `--buf-cache-size` | `32` | Number of shared buffers reserved for each poll group. |
+
+## 4. Register the NVMe-oF SSD Pool
+
+### 4.1 Register All SSDs
+
+```bash
+python3 -m mooncake.mooncake_ssd_register \
+    --master_server_address=192.168.65.81:50051 \
+    --spdk_target_info="ip:192.168.65.56 path:/home/spdk" \
+    --spdk_target_info="ip:192.168.65.57 path:/root/spdk"
+```
+
+#### Parameters
+
+| Parameter | Description |
+|-----------|-------------|
+| `--master_server_address` | IP address and port of the master service node. The default port is 50051. |
+| `--spdk_target_info` | Target node information, including `ip`, the node IP address, and `path`, the SPDK installation path. |
+| `--username` | SSH username used to connect to target nodes. The default value is `root`. |
+| `--port` | SSH port used to connect to target nodes. The default value is `22`. |
+| `--password` | SSH password used to connect to target nodes. |
+| `--key-file` | SSH private key file used to connect to target nodes. |
+
+## 5. Unregister the NVMe-oF SSD Pool
+
+### 5.1 Unregister a Specific SSD
+
+```bash
+python3 -m mooncake.mooncake_ssd_unregister \
+    --master_server_address=192.168.65.81:50051 \
+    --spdk_target_info="ip:192.168.65.56 ns:1 nqn:nqn.2016-06.io.spdk:cnode1"
+```
+
+#### Parameters
+
+| Parameter | Description |
+|-----------|-------------|
+| `--master_server_address` | IP address and port of the master service node. The default port is 50051. |
+| `--spdk_target_info` | Disk information to unregister, including `ip`, the node IP address, `ns`, the namespace ID, and `nqn`, the subsystem NQN. |
+| `--username` | SSH username used to connect to target nodes. The default value is `root`. |
+| `--port` | SSH port used to connect to target nodes. The default value is `22`. |
+| `--password` | SSH password used to connect to target nodes. |
+| `--key-file` | SSH private key file used to connect to target nodes. |
+
+### 5.2 Get Target Disk Information
+
+Enter the SPDK directory on the target node and run the following commands.
+
+1. Show subsystem information, including NQN and namespace IDs:
+
+```bash
+./scripts/rpc.py nvmf_get_subsystems
+```
+
+2. Show disk details, including block size and PCI address:
+
+```bash
+./scripts/rpc.py bdev_get_bdevs
+```
+
+## 6. Performance Tests
+
+### 6.1 Use the Built-in Benchmark Tool
+
+```bash
+./build/mooncake-store/benchmarks/nof_worker_pool_bench \
+    --endpoints='traddr:192.168.65.56 trsvcid:4420 subnqn:nqn.2016-06.io.spdk:cnode1 trtype:RDMA adrfam:IPv4 ns:1, traddr:192.168.65.56 trsvcid:4420 subnqn:nqn.2016-06.io.spdk:cnode1 trtype:RDMA adrfam:IPv4 ns:2' \
+    --op=read \
+    --io_size=1048576 \
+    --iodepth=8 \
+    --warmup_sec=3 \
+    --duration_sec=30
+```
+
+#### Parameters
+
+| Parameter | Description |
+|-----------|-------------|
+| `--endpoints` | Disk information used for the test. |
+| `--op` | I/O operation type, either `read` or `write`. |
+| `--io_size` | Block size in bytes. |
+| `--iodepth` | Read/write queue depth. |
+| `--warmup_sec` | Warmup duration in seconds. |
+| `--duration_sec` | Test duration in seconds. |
+
+### 6.2 Use NoF with vLLM + LMCache
+
+For the general VLLM + LMCache + Mooncake deployment flow, see
+[vLLM V1 Disaggregated Serving with Mooncake Store and LMCache](../getting_started/examples/vllm-integration/vllmv1-lmcache-integration.md).
+After the NVMe-oF SSD pool is registered with Mooncake, add the NoF-specific
+Mooncake configuration below.
+
+#### NoF Environment Variables
+
+```bash
+export LMCACHE_CONFIG_FILE="/path/vllm-lmcache-mooncake-config.yaml"
+export MC_NOF_WORKERS=4
+export MC_NOF_SUBMIT_CHUNK_BYTES=$((1 << 17))  # 128KB
+export MC_NOF_INFLIGHT_BYTES_LIMIT=$((1 << 25))  # 32MB
+```
+
+#### NoF LMCache Configuration
+
+```yaml
+chunk_size: 256
+remote_url: "mooncakestore://192.168.65.81:50051/"
+remote_serde: "naive"
+local_cpu: True
+max_local_cpu_size: 8
+enable_mooncake_nof_pool: True
+
+extra_config:
+  local_hostname: "localhost"
+  metadata_server: "http://192.168.65.81:8080/metadata"
+  master_server_address: "192.168.65.81:50051"
+  global_segment_size: 0
+  local_buffer_size: 1073741824
+  protocol: "rdma"
+  device_name: "mlx5_0"
+```
+
+**Notes**:
+
+- `enable_mooncake_nof_pool=True` enables writing KV cache objects to the
+  registered NoF pool.
+- `global_segment_size: 0` means the inference process does not contribute a
+  memory segment to the Mooncake cluster.
+- Keep `local_buffer_size` non-zero because the client still needs local
+  staging buffers for Mooncake transfers.
+- The parameters in `extra_config` should use the same Mooncake master,
+  metadata server, protocol, and RDMA device as the store service.
+
+| Environment Variable | Description | Default |
+|----------------------|-------------|---------|
+| `MC_NOF_WORKERS` | Number of worker threads used to process SPDK NoF I/O operations. | 4 |
+| `MC_NOF_SUBMIT_CHUNK_BYTES` | Size of each I/O operation submitted to SPDK. | 128KB |
+| `MC_NOF_INFLIGHT_BYTES_LIMIT` | Maximum number of in-flight I/O bytes allowed in the system. | 32MB |
+
+These three parameters together provide QoS control for SPDK NoF I/O.
diff --git a/docs/source/deployment/ssd-offload.md b/docs/source/deployment/ssd-offload.md
index 46830a9e..439f5f29 100644
--- a/docs/source/deployment/ssd-offload.md
+++ b/docs/source/deployment/ssd-offload.md
@@ -4,7 +4,14 @@
 
 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.
+For measured TTFT and throughput impact in multi-turn workloads, see [Mooncake SSD Offload Benchmark](../performance/ssd-offload-benchmark-results.md).
+
+SSD offload requires the **Real Client** and supports two deployment modes:
+
+- **Mode A: Embedded Real Client** — the Python process embeds the Real Client, and SSD offload runs inside the Python process.
+- **Mode B: Standalone Real Client + DummyClient** — a standalone `mooncake_client` process runs SSD offload, and the Python process connects via a DummyClient.
+
+In both modes, all SSD reads and writes happen within the Real Client (embedded or standalone).
 
 ## Startup Steps
 
@@ -19,12 +26,33 @@ mkdir -p /nvme/mooncake_offload
 ```bash
 mooncake_master \
     --rpc_port=50051 \
-    --enable-offload true
+    --enable_offload=true
 ```
 
-### Step 3: Start the real client with SSD offload enabled
+### Step 3A (Mode A): Start the application with embedded Real Client
 
-Use the `--enable_offload` flag to enable SSD offload, and set environment variables to specify the storage path and backend:
+Use the `--enable_ssd_offload` flag to enable SSD offload, and set environment variables to specify the storage path and backend:
+
+```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,  # 512 MB
+    protocol="rdma",
+    rdma_devices="eth0",
+    master_server_addr="127.0.0.1:50051",
+    enable_ssd_offload=True,
+    ssd_offload_path="/nvme/mooncake_offload"
+)
+```
+
+### Step 3B (Mode B): Start the standalone real client with SSD offload enabled and connect it with the dummy client
+
+Set the same SSD offload environment variables as in Mode A.
 
 ```bash
 export MOONCAKE_OFFLOAD_FILE_STORAGE_PATH=/nvme/mooncake_offload
@@ -41,33 +69,28 @@ mooncake_client \
     --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.
+The application (e.g., SGLang) connects to the standalone real client via the `MooncakeDistributedStore` Python SDK. The Python process acts as a DummyClient; SSD offload and fallback loading are handled by the standalone real client process.
 
 ```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",
+store.setup_dummy(
+    mem_pool_size=4 * 1024 * 1024 * 1024,  # 4 GB
+    local_buffer_size=512 * 1024 * 1024,  # 512 MB
+    server_address="<machine IP>:50052"   # mooncake_client RPC address from Step 3
 )
 ```
 
+> **Note:** When using `bucket_storage_backend` or `file_per_key_storage_backend`, the real client scans existing SSD metadata on startup and reports it to the master automatically. `offset_allocator_storage_backend` is the exception: it truncates its data file during initialization and does not recover previously offloaded objects after a restart.
+
 ---
 
 ## Real Client Parameters
 
 | Flag | Default | Description |
 |------|---------|-------------|
+| `--metadata_server` | `http://127.0.0.1:8080/metadata` | Metadata server connection string |
 | `--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 |
@@ -92,6 +115,8 @@ store.setup(
 | `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_CLIENT_BUFFER_GC_INTERVAL_SECONDS` | `1` | Interval for reclaiming expired offload buffers; defaults to the heartbeat interval in the current implementation |
+| `MOONCAKE_OFFLOAD_CLIENT_BUFFER_GC_TTL_MS` | `5000` | Lease time for buffers returned by `batch_get_offload_object` before GC reclaims them |
 | `MOONCAKE_OFFLOAD_USE_URING` | `false` | Enable io_uring for async file I/O |
 
 ### Bucket backend settings
@@ -103,7 +128,16 @@ Applies when `MOONCAKE_OFFLOAD_STORAGE_BACKEND_DESCRIPTOR=bucket_storage_backend
 | `MOONCAKE_OFFLOAD_BUCKET_SIZE_LIMIT_BYTES` | `268435456` (256 MB) | Max size per bucket |
 | `MOONCAKE_OFFLOAD_BUCKET_KEYS_LIMIT` | `500` | Max keys per bucket |
 | `MOONCAKE_OFFLOAD_BUCKET_MAX_TOTAL_SIZE` | `0` | Eviction threshold in bytes. When set to `0`, the backend uses **90% of the physical disk capacity** as the quota — it does not mean unlimited. Set an explicit value to control disk usage precisely. |
-| `MOONCAKE_OFFLOAD_BUCKET_EVICTION_POLICY` | `none` | Eviction policy: `none` / `fifo` / `lru` |
+| `MOONCAKE_OFFLOAD_BUCKET_EVICTION_POLICY` | `fifo` | Eviction policy: `none` / `fifo` / `lru` |
+
+### File-per-key backend settings
+
+Applies when `MOONCAKE_OFFLOAD_STORAGE_BACKEND_DESCRIPTOR=file_per_key_storage_backend`.
+
+| Environment Variable | Default | Description |
+|---|---|---|
+| `MOONCAKE_OFFLOAD_FSDIR` | `file_per_key_dir` | Subdirectory name created under `MOONCAKE_OFFLOAD_FILE_STORAGE_PATH` |
+| `ENABLE_EVICTION` | `true` | Enables local-storage eviction logic for this backend |
 
 ---
 
@@ -176,7 +210,8 @@ The following example starts a master and a real client on a single machine.
 
 ```bash
 mooncake_master \
-    --rpc_port=50051
+    --rpc_port=50051 \
+    --enable_offload=true
 ```
 
 ### Start the real client (new terminal)
@@ -202,7 +237,7 @@ mooncake_client \
 ## 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.
+- On real client restart, `bucket_storage_backend` and `file_per_key_storage_backend` scan existing SSD metadata and report it to the master, so previously offloaded objects remain accessible. `offset_allocator_storage_backend` does not support restart recovery.
 - 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.
 
@@ -210,7 +245,7 @@ mooncake_client \
 
 ```bash
 # Node A — runs the master and its own real client
-mooncake_master --rpc_port=50051 --enable-offload true &
+mooncake_master --rpc_port=50051 --enable_offload=true &
 
 export MOONCAKE_OFFLOAD_FILE_STORAGE_PATH=/nvme/mooncake_offload
 mooncake_client \
@@ -244,8 +279,8 @@ mooncake_client \
 
 ### 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.
+- Confirm `--enable_offload=true` is passed to both `mooncake_client` and `mooncake_master`. Both binaries define the same gflags boolean; hyphenated aliases may work, but the underscored spelling matches the source.
+- Check that `MOONCAKE_OFFLOAD_FILE_STORAGE_PATH` points to an existing, writable directory. If the path is invalid, real client setup fails during `FileStorageConfig::Validate()`.
 - 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
@@ -253,6 +288,7 @@ mooncake_client \
 - 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.
+- The path must already exist and must be a directory. Passing a regular file path also causes setup to fail.
 
 ### "Failed to register buffer with UringFile" warning in logs
 
diff --git a/docs/source/design/architecture.md b/docs/source/design/architecture.md
index a3edfe8d..434def0f 100644
--- a/docs/source/design/architecture.md
+++ b/docs/source/design/architecture.md
@@ -12,7 +12,7 @@ Mooncake:
 
 ## Architectural Overview
 ![architecture](../image/mooncake-store.png)
-- Mooncake provides object-level operations, i.e. `Get/Put/List/Del`, and also supports dynamically configurating replication strategies (`Replicate` operations);
+- Mooncake provides object-level operations, i.e. `Get/Put/List/Del`, and also supports dynamically configuring replication strategies (`Replicate` operations);
 - Mooncake supports zero-copy and multi-NIC data transfer over VRAM/DRAM/NVMe SSD. This feature is supported by Transfer Engine, which has been open-sourced;
 - **The master node** centrally manages the mappings of objects to VRAM/DRAM/NVM buffers. The master node also drives **managed pool buffer nodes** to achieve data transfer by calling Transfer Engine's APIs;
 - **Managed pool buffer nodes** mainly provide DRAM space for storing objects.
diff --git a/docs/source/design/conductor/conductor-architecture-design.md b/docs/source/design/conductor/conductor-architecture-design.md
new file mode 100644
index 00000000..e21cb74e
--- /dev/null
+++ b/docs/source/design/conductor/conductor-architecture-design.md
@@ -0,0 +1,182 @@
+# Mooncake Conductor Architecture
+
+## Overview
+
+Mooncake Conductor is the kv-cache indexer used by cache-aware routers. It
+subscribes to KV cache events from inference engines or storage backends,
+normalizes those events, maintains a global prefix cache table, and exposes
+HTTP APIs for dynamic service registration and cache-hit queries.
+
+The design goal is to let routers answer a simple scheduling question:
+for this request prefix, which registered instance has the best reusable KV
+cache locality, and on which cache tiers is that prefix available?
+
+## Architecture
+
+```mermaid
+flowchart TD
+    subgraph Router["Gateway or router"]
+        HTTPClient["HTTP client"]
+    end
+
+    subgraph Conductor["Mooncake Conductor"]
+        EventManager["EventManager"]
+        ZMQClient["ZMQClient"]
+        KVEventHandler["KVEventHandler"]
+        PrefixCacheTable["PrefixCacheTable"]
+    end
+
+    subgraph Publishers["vLLM, SGLang, Mooncake, or cache daemons"]
+        InferenceService["Inference or cache service"]
+        KVEvent["KV event stream"]
+    end
+
+    EventManager -->|creates and manages| ZMQClient
+    ZMQClient -->|registers handler| KVEventHandler
+    InferenceService -->|publishes events| KVEvent
+    KVEvent -->|ZMQ SUB or DEALER| ZMQClient
+    ZMQClient -->|decoded batches| KVEventHandler
+    KVEventHandler -->|normalizes events| PrefixCacheTable
+    HTTPClient -->|/query, /register, /unregister| EventManager
+    EventManager -->|cache-hit computation| PrefixCacheTable
+    PrefixCacheTable -->|hit result by instance| EventManager
+    EventManager -->|HTTP response| HTTPClient
+```
+
+## Components
+
+| Component | Responsibility |
+|---|---|
+| `EventManager` | Owns the Conductor lifecycle, HTTP server, dynamic registration, active service map, and tenant-to-instance map. |
+| `ZMQClient` | Connects to publisher endpoints, consumes event frames, decodes event batches, tracks sequence numbers, and requests replay after reconnects. |
+| `KVEventHandler` | Adapts decoded engine events into Conductor store/remove events enriched with registration metadata. |
+| `PrefixCacheTable` | Maintains model-context-specific prefix maps, engine-hash to conductor-hash mappings, medium metadata, DP-rank metadata, and query-time hit computation. |
+
+
+## Data model
+
+The prefix index is scoped by `ModelContext`:
+
+```text
+(tenant_id, model_name, lora_name, block_size, additional_salt, instance_id)
+```
+
+Within each context, Conductor stores:
+
+- a mapping from engine-provided block hash to Conductor prefix hash;
+- a prefix hash map that records replica count, medium set, DP-rank set, and
+  per-instance access metadata;
+- a DP-rank set used to report rank-level hit information.
+
+The current implementation computes complete-block prefix hashes from token IDs
+and ignores trailing partial blocks during `/query`.
+
+## Event flow
+
+1. A service is registered statically from `conductor_config.json` or
+   dynamically through `POST /register`.
+2. `EventManager` creates one `ZMQClient` per `(instance_id, tenant_id,
+   dp_rank)` service key.
+3. `ZMQClient` subscribes to the publisher endpoint and consumes frames in the
+   form `[topic, sequence, payload]`.
+4. The payload is decoded into a batch of engine events. Today, the implemented
+   parser supports vLLM `BlockStored` and `BlockRemoved` msgpack events.
+5. `KVEventHandler` enriches events with registration metadata such as model,
+   LoRA, tenant, instance, block size, and additional salt.
+6. `PrefixCacheTable` updates the prefix map for stored or removed blocks.
+7. If a reconnect detects missed sequence numbers, the `replay endpoint` can be
+   used to request missed events.
+
+## Query flow
+
+1. A router obtains prompt token IDs, usually from an engine tokenizer endpoint.
+2. The router calls `POST /query` with `model`, `token_ids`, `block_size`, and
+   optional `tenant_id`, `instance_id`, `lora_name`, and `cache_salt`.
+3. Conductor computes complete-block prefix hashes for the request.
+4. The prefix table is scanned in order. The first miss terminates the scan so
+   prefix continuity is preserved.
+5. Conductor returns per-instance `longest_matched`, medium hit counts, and
+   DP-rank hit counts.
+6. The router selects the best target instance and forwards the request.
+
+## Dynamic registration
+
+Conductor supports runtime registration so routers or control planes can add
+and remove KV event publishers without restarting the process.
+
+```json
+{
+  "endpoint": "tcp://127.0.0.1:5557",
+  "replay_endpoint": "tcp://127.0.0.1:5558",
+  "type": "vLLM",
+  "modelname": "qwen2.5",
+  "lora_name": "",
+  "tenant_id": "default",
+  "instance_id": "vllm-prefill-node1",
+  "block_size": 128,
+  "dp_rank": 0,
+  "additionalsalt": ""
+}
+```
+
+Static configuration uses the same fields under `kvevent_instance`:
+
+```json
+{
+  "http_server_port": 13333,
+  "kvevent_instance": {
+    "vllm-prefill-node1": {
+      "endpoint": "tcp://127.0.0.1:5557",
+      "replay_endpoint": "tcp://127.0.0.1:5558",
+      "type": "vLLM",
+      "modelname": "qwen2.5",
+      "lora_name": "",
+      "tenant_id": "default",
+      "instance_id": "vllm-prefill-node1",
+      "block_size": 128,
+      "dp_rank": 0,
+      "additionalsalt": ""
+    }
+  }
+}
+```
+
+## Environment variables
+
+| Variable | Default | Description |
+|---|---|---|
+| `CONDUCTOR_LOG_LEVEL` | `INFO` | Log level: `DEBUG`, `INFO`, `WARN`, or `ERROR`. |
+| `CONDUCTOR_CONFIG_PATH` | `~/.mooncake/conductor_config.json` | Path to the static configuration file. |
+| `CONDUCTOR_SEED` | random | Legacy seed option for hash computation experiments. |
+
+## Build and run
+
+```bash
+cd mooncake-conductor/conductor-ctrl
+go mod tidy
+go build -o mooncake_conductor .
+```
+
+```bash
+export CONDUCTOR_CONFIG_PATH=../example/conductor_config.json
+export CONDUCTOR_LOG_LEVEL=INFO
+./mooncake_conductor
+```
+
+## Project structure
+
+```text
+mooncake-conductor/
++-- conductor-ctrl/
+|   +-- common/        # shared types, helpers, and sync map
+|   +-- kvevent/       # EventManager and KVEventHandler
+|   +-- prefixindex/   # prefix cache table and hit computation
+|   +-- zmq/           # ZMQ client, event decoding, event types
+|   +-- main.go        # process entry point
++-- example/           # demo config and cache-aware proxy
++-- build.sh
++-- CMakeLists.txt
+```
+
+See [Indexer API](./indexer-api-design.md) for the HTTP API and KV Events wire
+format.
diff --git a/docs/source/design/conductor/indexer-api-design.md b/docs/source/design/conductor/indexer-api-design.md
index 876cfaf6..e79367e6 100644
--- a/docs/source/design/conductor/indexer-api-design.md
+++ b/docs/source/design/conductor/indexer-api-design.md
@@ -1,224 +1,362 @@
-# Mooncake Conductor Indexer
+# Mooncake Conductor Indexer API
 
-## Introduction
-The Mooncake Conductor Indexer is a specialized service designed to efficiently track and report token hit counts across various caching levels for different model instances. It provides a list of APIs that allow users to query token hit statistics based on token ID or chunked token hash, thereby facilitating optimized the performance of LLM inference.The figure below illustrates the architecture of Mooncake KVindexer: ![Mooncake KVindexer](../../image/conductor/architecture.png)
+## Overview
 
-## tiered storage & pools
-We drew inspiration from the definition of [KVBM components](https://github.com/ai-dynamo/dynamo/blob/main/docs/kvbm/kvbm_components.md) and divided the KV cache into three levels: G1, G2, and G3. The detailed introduction is as follows:
+Mooncake Conductor is a KV cache indexer used by routers and gateways to make
+cache-aware scheduling decisions. It consumes KV cache events from inference
+engines or storage backends, maintains prefix-hit metadata across cache tiers,
+and exposes HTTP APIs for service registration and cache-hit queries.
 
-- **Device Pool(G1)**: Device-resident KV block pool. Allocates mutable device blocks, registers completed blocks (immutable), serves lookups by sequence hash, and is the target for onboarding (Host→Device, Disk→Device).
-- **Host Pool(G2)**: Mooncake registered memory KV pool. Receives Device offloads (Device→Host), can onboard to Device (Host→Device), and offloads to Disk. For high-performance, zero-copy data transfers, it utilizes the Mooncake Transfer-Engine.
-- **Disk Pool(G3)**: SSD NVMe-backed KV pool. Receives Host offloads (Host→Disk), and provides large space for storing KV. 
+This document incorporates the latest API direction from:
 
+- [RFC #1403: Mooncake KV-Store Indexer API Standardization](https://github.com/kvcache-ai/Mooncake/issues/1403)
+- [RFC #1527: KV Events API Standardization](https://github.com/kvcache-ai/Mooncake/issues/1527)
 
-## Indexer API
+The Conductor can serve multiple model groups in one process. Each query is
+scoped by model identity, block size, LoRA identity, tenant isolation, and the
+registered instance that can receive traffic.
+
+## Concepts
+
+### Storage tiers
+
+The indexer tracks KV cache availability across three logical tiers:
+
+- **G1, Device Pool**: Device-resident KV blocks, such as GPU, NPU, HBM, or
+  other accelerator memory owned by inference engines.
+- **G2, Host Pool**: CPU or host DRAM KV blocks, including Mooncake registered
+  memory pools.
+- **G3, Disk Pool**: SSD, 3FS, DFS, NFS, or other disk-backed KV storage.
+
+The `medium` field identifies the concrete tier or device type. Common values
+are `gpu`, `cpu`, and `disk`. Engines may add other values as new media are
+supported.
+
+### Identity dimensions
+
+KV cache hits are interpreted under the following dimensions:
+
+| Dimension | Description |
+|---|---|
+| `model_name` or `model` | Model identifier. KV blocks from different models are incompatible. |
+| `block_size` | Number of tokens per KV block. Different block sizes produce different token-to-block mappings. |
+| `additional_salt` | Opaque salt used to separate hash namespaces for quantization, model revision, tenant isolation, or other deployment-specific dimensions. |
+|`cache_salt`| Ensure cached data blocks are kept separate for different customers|
+| `lora_name` | LoRA adapter name. Empty or `null` means the base model. |
+| `tenant_id` | Upstream tenant or customer identity. Used for isolation and to keep query output bounded. |
+| `instance_id` | Routable API server or engine instance returned by the Indexer API. Routers use this value as the scheduling target. |
+| `backend_id` | KV Events identity for the entity that owns the KV blocks. It may be an inference worker, a Mooncake storage daemon, or another cache backend. |
+| `medium` | Cache medium where the blocks are present. |
+| `dp_rank` | Data-parallel rank that owns or can serve the blocks. |
+
+`instance_id` and `backend_id` intentionally have different meanings.
+`instance_id` is the router-facing target in the Indexer API. `backend_id` is
+the event-facing cache owner in the KV Events API. In deployments where cache
+storage is decoupled from inference workers, `backend_id` can identify a cache
+daemon while `instance_id` still identifies the engine endpoint that receives
+requests.
+
+## Hashing standard
+
+The standardized event contract recommends **XXH3-64 with seed `S`**.
+
+- **Local block hash**:
+  `XXH3(token_bytes_le, S)`, where tokens are little-endian `u32` values
+  concatenated for one block.
+- **Rolling sequence hash**:
+  The first block uses `seq_hash[0] = local_block_hash[0]`. Each subsequent
+  block uses:
+
+  ```text
+  seq_hash[i] = XXH3(seq_hash[i-1]_le || local_block_hash[i]_le, S)
+  ```
+
+  Here `||` means byte concatenation, not a logical OR.
+
+All hashes used by the standardized KV Events API are rolling sequence hashes.
+A `seq_hash` identifies the whole prefix up to that block depth, so equal
+prefixes produce equal hashes until the first differing block.
+
+If an engine does not follow the standardized hashing scheme, it must provide
+`token_ids` in `stored` events so the consumer can recompute the indexer's hash
+representation.
+
+## HTTP APIs
+
+### `POST /register`
+
+Registers a KV event publisher and starts consuming events from it.
+
+```json
+{
+  "endpoint": "tcp://1.1.1.1:5557",
+  "replay_endpoint": "tcp://1.1.1.1:5558",
+  "type": "vLLM",
+  "modelname": "deepseek",
+  "lora_name": "sql-adapter",
+  "tenant_id": "default",
+  "instance_id": "vllm-prefill-node1",
+  "block_size": 128,
+  "dp_rank": 0,
+  "additionalsalt": "w8a8"
+}
+```
+
+| Field | Required | Description |
+|---|---|---|
+| `endpoint` | Yes | ZMQ KV event publisher endpoint. |
+| `replay_endpoint` | No | ZMQ replay endpoint used to recover missed events. |
+| `type` | Yes | Publisher type, such as `vLLM`, `SGLang`, or `Mooncake`. |
+| `modelname` | Yes | Model name for this publisher. This is the HTTP API wire name for `model_name`. |
+| `lora_name` | No | LoRA adapter name. Empty or omitted means base model. |
+| `tenant_id` | No | Tenant identity. Defaults to `default`. |
+| `instance_id` | Yes | Router-facing engine or API server instance identity. |
+| `block_size` | Yes | KV block size in tokens. |
+| `dp_rank` | Yes | Data-parallel rank for this publisher. |
+| `additionalsalt` | No | HTTP API wire name for `additional_salt`. Defaults to an empty string. |
+
+Successful response:
+
+```json
+{
+  "status": "registered successfully",
+  "instance_id": "vllm-prefill-node1"
+}
+```
+
+### `POST /unregister`
+
+Stops consuming events for a registered publisher.
+
+```json
+{
+  "type": "vLLM",
+  "modelname": "deepseek",
+  "lora_name": "sql-adapter",
+  "tenant_id": "default",
+  "instance_id": "vllm-prefill-node1",
+  "block_size": 128,
+  "dp_rank": 0
+}
+```
+
+`tenant_id` defaults to `default`. The current implementation removes the
+subscription identified by `(instance_id, tenant_id, dp_rank)`.
+
+Successful response:
+
+```json
+{
+  "status": "unregistered successfully",
+  "removed_instances": ["vllm-prefill-node1|default|0"]
+}
+```
 
 ### `POST /query`
- query token hit count.
-- **Input**:
-    - **Body** (JSON):
-        ```json
-        {
-            "model": "deepseek",
-            "lora_name": "xx-adapter",
-            "lora_id": 12, // defined for backward compatibility and should not be used together with `lora_name`
-            "token_ids": [1, 15, 100],
-            "tenant_id": None,
-            "cache_salt": None,
-        }
-        ```
-    - **Parameter Description**:
-        - `model`: (required, string) model name
-        - `lora_name`: (optional, string) The name of the LoRA adapter, default is `None`(indicating no LoRA adapter is used)
-        - `lora_id`: (optional, int) The ID of the LoRA adapter. This parameter is defined for backward compatibility and should not be used together with `lora_name`(Only one of them can be specified). Default is `-1`(indicating no LoRA adapter is used)
-        - `token_ids`: (required, [int]) prompt token id list
-        - `tenant_id`: (optional, int) In a multi-tenant architecture, tenant_id is the key identifier for distinguishing and isolating data from different tenants (such as different companies or user groups). All data operations are logically isolated based on this ID. If you provide it, the indexer will only return the token hit information for this tenant. Default is None(meaning there is only one tenant)
-        - `cache_salt`: (optional, int) An optional salt value to ensure cached data blocks are kept separate for different customers. This prevents one customer's kv-index data from being served to another. Default is None, meaning no salt is used.
-    - **example**:
-        ```json
-        {
-            "model": "deepseek-v3",
-            "lora_name": "sql_adapter",
-            "token_ids": [101, 15, 100, 55, 89],
-        }
-        ```
-- **Output**:
-    ```json
-    {
-        "data": {
-            "tenant_id": {
-                "api_server_unique_name": {
-                    "longest_matched": 100, // the number of longest prefix matched token among multiple DPs(if there are)
-                    "GPU": 20,
-                    "DP": {
-                        0: 10,
-                        1: 20
-                    },
-                    "CPU": 60,
-                    "DISK": 10
-                },
-                ... // other engine instance
-            },
-            ... // other tenant
-        }
-    }
-    ```
-    - **Parameter Description**
-        - `tenant_id`: tenant id, only used in multi-tenant scenario.
-        - `api_server_unique_name`: it is a unique name for a LLM API server endpoint in the engine side. For example, two service instances are currently started separately by running the `vllm server` command, and they are registered in the indexer with different names(such as vllm-1,vllm-2)
-        - `longest_matched`: the number of longest prefix matched token among G1/G2/G3. Indexer will sequentially query the hit status of each token-block according to the prefix order. If it hits, count the situation of this token-block at each level; If it missed, terminate the query (ensuring prefix continuity).
-        - `GPU`, `CPU`, `DISK`: token ids hit count for each tiered storage medium. The Indexer will track the storage status of KV-cache across various media. This requires different KV publishers to inform the Indexer of the actual storage medium type via kv-events. The following examples list several common names, such as using GPU or NPU to represent the Device Pool, using CPU to represent the Host Pool, and using DISK to represent the Disk Pool.
-        - `DP`: token ids hit count for each DP rank.
-    - **example**:
 
-        Assume the input token_ids are [101, 15, 100, 55, 89, 63], the block_size is 2, and the dp2 strategy is enabled. There are three block hashes to match [H1, H2, H3], where H1 hits in GPU (dp0, dp1), CPU, and DISK; H2 hits in GPU (dp0) and CPU; and H3 hits in DISK.
-        ```json
-        {
-            "vllm-1": {
-                "longest_matched": 6,
-                "GPU": 4,
-                "DP": {
-                    0: 4,
-                    1: 2
-                },
-                "CPU": 4,
-                "DISK": 4
-            }
-        }
-        ```
+Query cache hits by token IDs.
+
+```json
+{
+  "model": "deepseek",
+  "lora_name": "sql-adapter",
+  "token_ids": [101, 15, 100, 55, 89],
+  "tenant_id": "default",
+  "instance_id": "vllm-prefill-node1",
+  "block_size": 64,
+  "cache_salt": "w8a8"
+}
+```
+
+| Field | Required | Description |
+|---|---|---|
+| `model` | Yes | Model name. |
+| `lora_name` | No | LoRA adapter name. Empty or omitted means base model. |
+| `lora_id` | No | Deprecated compatibility field. Do not use together with `lora_name`. |
+| `token_ids` | Yes | Prompt token IDs. Only complete blocks are considered. |
+| `tenant_id` | No | Tenant identity. Defaults to `default`. |
+| `instance_id` | No | If set, query one instance. If omitted, query all instances registered under the tenant. |
+| `block_size` | Yes | KV block size in tokens. |
+| `cache_salt` | No | Query-side hash namespace salt. Corresponds to the event `additional_salt` concept. |
+
+Response:
+
+```json
+{
+  "default": {
+    "vllm-prefill-node1": {
+      "longest_matched": 256,
+      "GPU": 128,
+      "DP": {
+        "0": 128,
+        "1": 256
+      },
+      "CPU": 256,
+      "DISK": 0
+    }
+  }
+}
+```
+
+| Field | Description |
+|---|---|
+| `longest_matched` | Longest continuous prefix hit in tokens across all tracked media and DP ranks for this instance. |
+| `GPU`, `CPU`, `DISK` | Matched prefix tokens available on each medium. Names are examples; future media can be added. |
+| `DP` | Matched prefix tokens grouped by data-parallel rank. |
 
 ### `POST /query_by_hash`
- query token hit count by chunked_token hash key. Each model service uses its own independent page_size, `longest_matched = page_size * matched hash_key`
-- **Input**:
-    - **Body** (JSON):
-        ```json
-        {
-            "model": "deepseek",
-            "lora_name": "xx-adapter",
-            "lora_id": 12, // defined for backward compatibility and should not be used together with `lora_name`
-            "block_hash": ["hash_key_by_chunked_tokens"],
-            "tenant_id": None,
-            "cache_salt": None,
-        }
-        ```
-    - **Parameter Description**:
-        - `model`: (required, string) model name
-        - `lora_name`: (optional, string) The name of the LoRA adapter, default is `None`(indicating no LoRA adapter is used)
-        - `lora_id`: (optional, int) The ID of the LoRA adapter. This parameter is defined for backward compatibility and should not be used together with `lora_name`(Only one of them can be specified). Default is `-1`(indicating no LoRA adapter is used)
-        - `block_hash`: (required, [int]) chunk_token hash list
-        - `tenant_id`: (optional, int) In a multi-tenant architecture, tenant_id is the key identifier for distinguishing and isolating data from different tenants (such as different companies or user groups). All data operations are logically isolated based on this ID. If you provide it, the indexer will only return the token hit information for this tenant. Default is None(meaning there is only one tenant)
-        - `cache_salt`: (optional, int) An optional salt value to ensure cached data blocks are kept separate for different customers. This prevents one customer's kv-index data from being served to another. Default is None, meaning no salt is used.
-- **Output**:
-    ```json
-    {
-        "data": {
-            "tenant_id": {
-                "api_server_unique_name": {
-                    "longest_matched": 100, // the number of longest prefix matched token among multiple DPs(if there are)
-                    "GPU": 20,
-                    "DP": {
-                        0: 10,
-                        1: 20
-                    },
-                    "CPU": 60,
-                    "DISK": 10
-                },
-                ... // other engine instance
-            },
-            ... // other tenant
-        }
-    }
-    ```
-    - **Parameter Description**: The output result is same as `/query` api.
 
+Queries cache hits by precomputed rolling sequence hashes. This API avoids
+sending long token lists over the network.
 
-## Indexer KVEvents Structure
-Typically, the device pool is used for loading model weights, with the remaining space registered for KV blocks by the model inference service runtime, the host pool and disk pool are managed uniformly by the Mooncake Store. There is a difference in the management unit for KV data between the two: the device pool uses blocks as the smallest unit for KV data, while the host pool and disk pool use Mooncake Store Objects as the smallest unit for KV storage. In practice, users may split a complete KV block into multiple Mooncake Store Objects for maintenance according to parallel strategies such as tensor parallelism (tp) and context parallelism (cp).
-
-### G1 KVEvents
-[vLLM](https://github.com/vllm-project/vllm/blob/main/vllm/distributed/kv_events.py) and [SGLang](https://github.com/sgl-project/sglang/blob/main/python/sglang/srt/disaggregation/kv_events.py) publishes events using the `EventBatch` structure, with each batch containing three types of events:
-
-- `BlockStored`：Adds a single KV block.
-- `BlockRemoved`：Removes a single KV block.
-- `AllBlocksCleared`：Clears all KV blocks.
-
-```py
-EventBatch：
+```json
 {
-    ts: float， # timestamp
-    events：list[BlockStored | BlockRemoved | AllBlocksCleared]，
-    data_parallel_rank: int | None = None,  # vLLM use this to indicate dp rank
-    attn_dp_rank: int | None = None,        # SGLang use this to indicate dp rank
+  "model": "deepseek",
+  "lora_name": "sql-adapter",
+  "seq_hashes": [1234567890, 9876543210],
+  "tenant_id": "default",
+  "instance_id": "vllm-prefill-node1",
+  "block_size": 64,
+  "cache_salt": "w8a8"
 }
-
-
-BlockStored:
-{
-    block_hashes: list[int]
-    parent_block_hash: int | None
-    token_ids: list[int]
-    block_size: int
-
-    lora_id: int | None
-    """Deprecated: use `lora_name` for KV block key hash.
-    Retained for backward compatibility.
-    """
-    
-    medium: str | None
-    """KV cache is categorized by tier. Currently, the following types are supported:
-    set "GPU", "NPU" for device pool(G1), 
-    set "CPU" for host pool(G2), 
-    set "DISK" for disk pool(G3).
-    In the future, more medium types can be supported for each tier. For example, "TPU" and "AMD" could be added for device pool.
-    """
-    lora_name: str | None
-}
-
-BlockRemoved:
-{
-    block_hashes: list[int]
-    lora_name: str | None
-}
-
 ```
 
-### G2/G3 KVEvents
-Mooncake Store is a distributed key-value (KV) store. To ensure system consistency, a timestamp is assigned to each KVEvent for maintenance.
+For compatibility with earlier drafts, clients may call the hash list
+`block_hash`, but new clients should use `seq_hashes` to make it explicit that
+the values are rolling sequence hashes rather than local block hashes.
 
-Mooncake publishes events using the `EventBatch` structure, with each batch containing three types of events:
-- `BlockStoreEvent`：Adds a single Mooncake Store Object.
-- `BlockUpdateEvent`：Updates a single Mooncake Store Object.
-- `RemoveAllEvent`：Removes all Mooncake Store Objects.
+The response shape is the same as `/query`. For this API,
+`longest_matched = block_size * matched_hash_count`.
 
-```cpp
-EventBatch {
-    std::vector<std::variant<BlockStoreEvent, BlockUpdateEvent, RemoveAllEvent>> events,
-}
 
-BlockStoreEvent {
-    "BlockStoreEvent", // Event type identifier, string type
-    float ts,          // timestamp
-    std::string mooncake_key,
-    std::vector<std::string> addr_list, // Storage location of each replica
+## KV Events API
 
-    uint64_t block_hash,
-    uint64_t parent_block_hash,
-    std::vector<uint32_t> token_id,
-    uint32_t block_size,
+The KV Events API is the wire contract between cache owners and indexers.
+Conductor normalizes engine-specific events into this model.
 
-    std::string model_name,
-    std::string lora_name,
-    uint32_t lora_id, // Retained for backward compatibility
-}
+### Event envelope
 
-BlockUpdateEvent {
-    "BlockUpdateEvent",
-    float ts,          // timestamp
-    mooncake_key,
-    std::vector<std::string> addr_list,
-}
+Every standardized event carries the same envelope:
 
-RemoveAllEvent {
-    "RemoveAllEvent"
+```json
+{
+  "event_id": 42,
+  "timestamp": 1739145600000,
+  "event_type": "stored",
+  "model_name": "llama-3.1-8b",
+  "block_size": 64,
+  "additional_salt": null,
+  "lora_name": null,
+  "tenant_id": "default",
+  "backend_id": "worker-0",
+  "medium": "gpu",
+  "dp_rank": 0
 }
 ```
+
+| Field | Type | Description |
+|---|---|---|
+| `event_id` | `u64` | Monotonically increasing sequence number scoped by the event stream dimensions. Authoritative for ordering. |
+| `timestamp` | `u64 or null` | Unix epoch milliseconds. Informational only, not used for ordering. |
+| `event_type` | `string` | One of `stored`, `removed`, or `cleared`. |
+| `model_name` | `string or null` | Model identifier. |
+| `block_size` | `u32 or null` | Tokens per block. |
+| `additional_salt` | `string or null` | Opaque deployment salt or namespace. |
+| `lora_name` | `string or null` | LoRA adapter name, or `null` for the base model. |
+| `tenant_id` | `string` | Tenant or customer identity. |
+| `backend_id` | `string` | Entity that owns the KV blocks. This can be an engine worker or a decoupled cache daemon. |
+| `medium` | `string or null` | Cache medium such as `gpu`, `cpu`, or `disk`. |
+| `dp_rank` | `u32 or null` | Data-parallel rank. |
+
+Events must be processed in consecutive `event_id` order within each stream
+identified by `(model_name, block_size, additional_salt, lora_name, tenant_id,
+backend_id, medium, dp_rank)`.
+
+### `stored`
+
+Published when one or more consecutive blocks are committed to a KV cache.
+
+```json
+{
+  "event_id": 42,
+  "timestamp": 1739145600000,
+  "event_type": "stored",
+  "model_name": "llama-3.1-8b",
+  "block_size": 64,
+  "additional_salt": null,
+  "lora_name": null,
+  "tenant_id": "default",
+  "backend_id": "worker-0",
+  "medium": "gpu",
+  "dp_rank": 0,
+  "seq_hashes": [1234567890, 9876543210, 1122334455],
+  "base_block_idx": 5,
+  "parent_hash": 9999999999,
+  "token_ids": null
+}
+```
+
+| Field | Type | Description |
+|---|---|---|
+| `seq_hashes` | `u64[]` | Rolling sequence hashes of consecutive stored blocks. |
+| `base_block_idx` | `u32 or null` | Zero-based depth of the first block in this event. |
+| `parent_hash` | `u64 or null` | Rolling sequence hash at depth `base_block_idx - 1`; `null` at the root. |
+| `token_ids` | `u32[] or null` | Tokens across all blocks in this event. Required when the publisher does not use the standardized hash. |
+
+At least one of `base_block_idx` or `parent_hash` must be present so the
+consumer can locate the blocks in the sequence.
+
+### `removed`
+
+Published when one or more blocks are evicted.
+
+```json
+{
+  "event_id": 43,
+  "timestamp": 1739145601000,
+  "event_type": "removed",
+  "model_name": "llama-3.1-8b",
+  "block_size": 64,
+  "additional_salt": null,
+  "lora_name": null,
+  "tenant_id": "default",
+  "backend_id": "worker-0",
+  "medium": "gpu",
+  "dp_rank": 0,
+  "seq_hashes": [1122334455],
+  "base_block_idx": 7
+}
+```
+
+`seq_hashes` is required. `base_block_idx` is optional but recommended for
+collision detection and observability.
+
+### `cleared`
+
+Published when all blocks for the event stream dimensions are purged.
+
+```json
+{
+  "event_id": 44,
+  "timestamp": 1739145602000,
+  "event_type": "cleared",
+  "model_name": "llama-3.1-8b",
+  "block_size": 64,
+  "additional_salt": null,
+  "lora_name": null,
+  "tenant_id": "default",
+  "backend_id": "worker-0",
+  "medium": "gpu",
+  "dp_rank": 0
+}
+```
+
+No additional payload fields are required.
+
+## Compatibility notes
+
+The current Conductor implementation consumes vLLM ZMQ msgpack batches and
+normalizes `BlockStored` and `BlockRemoved` into the internal prefix index.
+Registration metadata supplies fields such as `modelname`, `tenant_id`,
+`instance_id`, `block_size`, and `additionalsalt` when the engine event does
+not carry the full standardized envelope.
diff --git a/docs/source/design/engram.md b/docs/source/design/engram.md
new file mode 100644
index 00000000..0372ee44
--- /dev/null
+++ b/docs/source/design/engram.md
@@ -0,0 +1,144 @@
+# EngramStore Backend
+
+Mooncake provides EngramStore as the storage backend for Engram embedding tables.
+
+The scope is intentionally narrow:
+
+- the caller defines the physical table layout
+- the caller uploads one table per head
+- the caller provides precomputed row ids with shape `[B, L, H]`
+- Mooncake returns the selected rows as `[B, L, H, D]`
+
+Mooncake does not implement tokenizer compression, N-gram hashing, query logic,
+or any other model-side Engram algorithm.
+
+## Current Backend Boundary
+
+The current implementation is intentionally conservative. It keeps EngramStore
+on top of the existing Store interfaces and does not depend on:
+
+- transfer scatter read
+- grouped transfer task
+- `get_into_range`
+- `batch_query`
+- local direct mapping
+- query cache
+- remote gather control-plane changes
+
+Those optimizations are deferred to follow-up PRs so that the EngramStore backend can
+land first as a small, reviewable unit.
+
+## Configuration
+
+`EngramStoreConfig` contains the physical layout for one EngramStore layer:
+
+- `table_vocab_sizes`: per-head table sizes `[N_0, N_1, ..., N_{H-1}]`
+- `embedding_dim`: row width `D`
+
+For `layer_id`, Mooncake generates one store key per head:
+
+```text
+engram:l{layer_id}:h{head_idx}
+```
+
+Each key stores a `float32` table with shape `[N_h, D]`.
+
+## Public Interface
+
+Python:
+
+- `EngramStore(layer_id, config, store=None)`
+- `populate(embedding_buffers)`
+- `lookup(row_ids)`
+- `remove_from_store(force=False)`
+- `get_table_vocab_sizes()`
+- `get_store_keys()`
+- `get_num_heads()`
+- `get_embedding_dim()`
+
+The Python `store` argument accepts the existing `MooncakeDistributedStore`
+wrapper, or `None` for metadata-only construction.
+
+C++:
+
+- constructor `EngramStore(int layer_id, const EngramStoreConfig&, std::shared_ptr<PyClient>)`
+- `populate(...)`
+- `lookup_rows(...)`
+- `lookup_rows_contiguous(...)`
+- `remove_from_store(...)`
+- metadata getters matching the Python surface
+
+## Data Contract
+
+Populate expects one NumPy `float32` array per head:
+
+```text
+embedding_buffers[h].shape == [N_h, D]
+```
+
+Lookup accepts either:
+
+- nested Python lists with logical shape `[B, L, H]`, or
+- a contiguous NumPy `int64` array with shape `[B, L, H]`
+
+Lookup returns:
+
+```text
+output.shape == [B, L, H, D]
+```
+
+## Populate Flow
+
+Populate follows the existing Store write path:
+
+1. validate that exactly one table is provided for each head
+2. validate that every table matches `[N_h, D]`
+3. verify that the target head-table keys do not already exist
+4. register each embedding table buffer
+5. upload all head tables with `batch_put_from(...)`
+6. unregister the staging buffers
+
+`populate(...)` is defined as a create-only operation for one EngramStore layer. To
+reuse a `layer_id`, first remove the old tables with `remove_from_store(...)`.
+
+If upload fails after some head tables have already been written, or if publish
+finishes but post-write buffer cleanup fails, the backend best-effort removes
+the keys written by the failed populate attempt before returning an error.
+
+## Lookup Flow
+
+Each lookup follows the same simplified backend flow:
+
+1. validate the `row_ids` shape and bounds
+2. build per-head byte ranges for the requested rows
+3. issue one `get_into_ranges(...)` call to materialize those rows into the output buffer
+
+For NumPy `row_ids`, the binding uses a contiguous fast path and builds ranges
+directly from the input tensor without first converting the entire input into a
+nested C++ container.
+
+## Validation
+
+The backend enforces these invariants:
+
+- `table_vocab_sizes` is non-empty and every entry is positive
+- `embedding_dim` is positive
+- `populate(...)` receives exactly one table per head
+- every populated table matches `[N_h, D]`
+- `lookup(...)` receives a non-empty `[B, L, H]` input
+- every row id satisfies `0 <= row_ids[..., h] < N_h`
+
+## Validation Status
+
+This backend is covered by:
+
+- correctness tests in `scripts/test_engram_store.py`
+- benchmark coverage in `scripts/bench_engram_store_27b.py`
+
+`scripts/test_engram_store.py` can run against an existing Mooncake deployment through
+`MOONCAKE_CONFIG_PATH` / `MOONCAKE_MASTER`, or it can start a local
+`mooncake_master` instance automatically for a self-contained TCP test run.
+
+By default, the benchmark exercises `engram_store.populate(...)` directly. Its
+fallback populate paths are gated behind `ENGRAM_ALLOW_POPULATE_FALLBACK=1` so
+they do not silently mask regressions in the current implementation.
diff --git a/docs/source/design/mooncake-backend-pg.md b/docs/source/design/mooncake-backend-pg.md
new file mode 100644
index 00000000..39835605
--- /dev/null
+++ b/docs/source/design/mooncake-backend-pg.md
@@ -0,0 +1,224 @@
+# Mooncake Backend (PG) Design
+
+Mooncake Backend is a `torch.distributed` ProcessGroup backend for Mooncake. It
+provides collective and point-to-point communication primitives, rank-health
+tracking, and elastic recovery hooks for inference systems that need to keep
+serving after partial rank failures.
+
+This document is intended for developers who maintain Mooncake PG itself or
+integrate it into higher-level serving systems.
+
+## Goals
+
+Mooncake Backend is designed to:
+
+- integrate with PyTorch through the standard ProcessGroup extension mechanism;
+- expose `mooncake` for accelerator tensors and `mooncake-cpu` for CPU tensors;
+- support common collective APIs used by inference engines;
+- track active and inactive ranks so collectives can continue after failures;
+- allow replacement ranks to publish metadata, join an existing group, and be
+  activated by healthy ranks;
+- reuse Mooncake Transfer Engine and topology information for data movement.
+
+Non-goals:
+
+- It is not a drop-in replacement for every NCCL/Gloo behavior. Validate each
+  collective, dtype, and topology required by the application.
+- Elastic recovery is an explicit protocol. The backend does not silently add a
+  new process to all collectives without application coordination.
+
+## Relationship with `torch.distributed`
+
+Mooncake registers two PyTorch backends when the PG extension module is imported:
+
+- `mooncake-cpu`, registered for CPU devices;
+- `mooncake`, registered for accelerator devices such as CUDA or MUSA depending
+  on the build.
+
+The backend class itself derives from `c10d::ProcessGroup`. Applications use
+regular PyTorch APIs such as `dist.init_process_group()`, `dist.all_reduce()`,
+`dist.new_group()`, and `dist.batch_isend_irecv()`.
+
+Point-to-point dispatch in PyTorch expects a `c10d::Backend` object. Mooncake PG
+therefore includes a lightweight P2P shim that delegates `send` and `recv` calls
+back to the owning `MooncakeBackend` instance.
+
+## Main runtime objects
+
+### `MooncakeBackendOptions`
+
+`MooncakeBackendOptions` carries Mooncake-specific process-group configuration:
+
+| Field | Meaning |
+| --- | --- |
+| `activeRanks_` | Rank-health tensor exposed to collectives and user code. |
+| `isExtension_` | Whether this process is a joining/replacement rank. |
+| `maxWorldSize_` | Optional reserved capacity for future ranks. |
+
+The `activeRanks_` tensor must be `torch.int32`. It must be on CPU for
+`mooncake-cpu` and on the accelerator device for `mooncake`. When
+`maxWorldSize_` is set, `activeRanks_` must be sized to `maxWorldSize_` so the
+backend can reserve inactive rank slots.
+
+### Transfer group metadata
+
+Each backend owns shared metadata for:
+
+- the current rank and backend index;
+- current capacity (`size`) and visible active size (`activeSize`);
+- host and device active-rank masks;
+- peer connection state;
+- rank-local and rank-global mapping;
+- store handles and extension state used by recovery;
+- P2P proxy and connection poller state.
+
+`size` is the reserved capacity. `activeSize` is the visible group size returned
+by `dist.get_world_size()`. With `max_world_size`, `size` may be larger than
+`activeSize`; inactive slots are masked out by the active-rank state.
+
+### Transfer Engine ownership
+
+By default, Mooncake Backend initializes its own Transfer Engine. Advanced
+integrations may call `pg.set_transfer_engine(engine)` before
+`init_process_group()` to inject an external Transfer Engine. In that mode, the
+caller owns the engine and must keep it alive until all Mooncake process groups
+using it are destroyed.
+
+## Initialization lifecycle
+
+The initialization flow is:
+
+1. Python imports `mooncake.pg`, which loads a PyTorch-version-specific native
+   extension.
+2. The extension registers `mooncake` / `mooncake-cpu` with PyTorch.
+3. `dist.init_process_group()` invokes the backend factory with PyTorch
+   distributed options and optional `MooncakeBackendOptions`.
+4. The backend initializes active-rank masks and reserved rank slots.
+5. Non-extension ranks publish local peer metadata and wait until current peers
+   are connected.
+6. Extension ranks enter local-only mode and wait for the explicit join protocol.
+
+The important distinction is that reserving capacity does not automatically make
+future ranks active. New ranks are masked until `recover_ranks()` activates them.
+
+## Active ranks and dynamic world size
+
+Mooncake PG tracks two related concepts:
+
+- **Reserved capacity** (`size`): how many rank slots the backend knows about.
+- **Visible active size** (`activeSize`): the current group size visible through
+  PyTorch APIs.
+
+When `max_world_size` is larger than the initial `world_size`, Mooncake reserves
+extra slots but marks them inactive. This lets healthy ranks poll for joiner
+metadata and activate the joiners later without reconstructing the process group.
+
+`extend_group_size_to(size)` can also increase capacity. Newly extended ranks
+start inactive; the application must call `get_peer_state()` and
+`recover_ranks()` before they participate in collectives.
+
+## Elastic recovery protocol
+
+Mooncake PG uses a two-phase protocol for recovery and scale-up:
+
+```mermaid
+sequenceDiagram
+    participant H as Healthy ranks
+    participant J as Joining rank
+    participant S as Store / metadata
+
+    H->>S: init_process_group(world_size=M, max_world_size=N)
+    J->>S: init_process_group(world_size=N, is_extension=True)
+    J->>S: publish local peer metadata
+    H->>S: get_peer_state(join_ranks)
+    H->>S: recover_ranks(join_ranks)
+    S-->>J: extension state
+    J->>J: join_group() returns
+    H->>J: collectives include recovered ranks
+```
+
+Healthy rank responsibilities:
+
+1. Reserve capacity with `max_world_size` or `extend_group_size_to()`.
+2. Poll `get_peer_state(backend, ranks)` from all healthy ranks in a consistent
+   order.
+3. Call `recover_ranks(backend, ranks)` once candidate ranks are connected.
+4. Refresh higher-level components, such as Mooncake EP buffers, if they cache
+   transport metadata.
+
+Joining rank responsibilities:
+
+1. Initialize the process group with `is_extension=True`.
+2. Publish local peer metadata through the backend initialization path.
+3. Call `join_group(backend)` and block until healthy ranks publish extension
+   state.
+4. Re-enter normal collectives after `join_group()` returns.
+
+## Subgroup semantics
+
+Mooncake PG follows PyTorch process-group ordering requirements. All processes
+that participate in a parent group should call `dist.new_group()` in a consistent
+order. This is especially important for elastic subgroups because healthy ranks
+and joining ranks must agree on store prefixes and backend indices.
+
+For split-rank elastic patterns, create subgroups using the current membership
+on healthy ranks and the eventual membership on joining ranks, while preserving
+the same creation order. The PG elastic tests contain executable examples of this
+pattern.
+
+## Collective behavior
+
+Collectives use the backend active-rank state to skip inactive ranks. The exact
+implementation varies by operation and device type, but the high-level contract
+is:
+
+- active ranks participate in the collective;
+- inactive ranks are not waited on;
+- if communication detects a rank failure, active-rank state can be updated;
+- user code can read the current mask with `pg.get_active_ranks(backend)`.
+
+The backend currently implements common collective APIs including all-reduce,
+broadcast, all-gather, reduce-scatter, all-to-all, barrier, reduce, gather,
+scatter, and single-tensor P2P send/recv.
+
+## Failure and recovery boundaries
+
+Mooncake PG exposes low-level recovery primitives; higher-level systems are
+responsible for policy decisions such as:
+
+- which ranks are safe to replace;
+- when to stop routing traffic to a failed rank;
+- how to recreate model state on a replacement process;
+- when to refresh EP, scheduler, or application-level metadata;
+- how to coordinate subgroup recovery.
+
+Avoid assuming that `recover_ranks()` alone reconstructs all higher-level state.
+It activates the process-group communication path; the application still owns
+model weights, KV-cache state, routing policy, and request scheduling.
+
+## Testing checklist for PG changes
+
+When modifying PG internals, run at least:
+
+```bash
+# CPU functional tests
+python -m unittest discover -s mooncake-pg/tests -k CPU -v
+
+# CUDA functional tests, when GPUs are available
+python -m unittest discover -s mooncake-pg/tests -k CUDA -v
+
+# Collective benchmark smoke test
+PYTHONPATH=mooncake-pg \
+python mooncake-pg/benchmark/pgbench.py \
+  --collective all_reduce --backend mooncake --device cuda -g 2 -b 8 -e 1M -f 2
+```
+
+Also run elastic tests for changes that touch active ranks, metadata polling,
+subgroups, `extend_group_size_to()`, `get_peer_state()`, `recover_ranks()`, or
+`join_group()`.
+
+## Related documentation
+
+- [Mooncake EP design](mooncake-ep.md)
+- [Python API reference](../python-api-reference/ep-backend.md)
+- [PG/EP troubleshooting](../troubleshooting/pg-ep-troubleshooting.md)
diff --git a/docs/source/design/mooncake-ep.md b/docs/source/design/mooncake-ep.md
new file mode 100644
index 00000000..e0734e0b
--- /dev/null
+++ b/docs/source/design/mooncake-ep.md
@@ -0,0 +1,247 @@
+# Mooncake EP Design
+
+Mooncake EP is the expert-parallel communication runtime used for MoE token
+dispatch and combine. It follows the DeepEP low-latency programming model while
+adding Mooncake transport integration and rank activeness awareness.
+
+This document explains the runtime at a level useful for developers maintaining
+Mooncake EP or integrating it into inference engines.
+
+## Goals
+
+Mooncake EP is designed to:
+
+- provide low-latency dispatch/combine operations for expert-parallel MoE
+  inference;
+- keep the Python programming model close to DeepEP low-latency mode;
+- use Mooncake device transports for fast intra-node and inter-node movement;
+- detect failed source ranks through timeout-aware kernels;
+- interoperate with Mooncake Backend (PG) for bootstrap metadata exchange and
+  rank-health state.
+
+## High-level data flow
+
+MoE inference with Mooncake EP has three phases:
+
+1. **Dispatch**: each rank sends token hidden states to the ranks that own the
+   selected experts. The receiver packs tokens by local expert.
+2. **Expert compute**: each rank runs its local experts over the packed inputs.
+3. **Combine**: expert outputs are routed back to the original token owners and
+   reduced with routing weights.
+
+```mermaid
+flowchart LR
+    A[Local tokens x + topk_idx] --> B[dispatch]
+    B --> C[Packed local expert inputs]
+    C --> D[Local expert kernels]
+    D --> E[Packed expert outputs]
+    E --> F[combine]
+    F --> G[Combined local token outputs]
+```
+
+## Relationship with Mooncake Backend (PG)
+
+Mooncake EP is constructed from a `torch.distributed` process group:
+
+```python
+from mooncake.mooncake_ep_buffer import Buffer
+
+buffer = Buffer(group, num_ep_buffer_bytes)
+```
+
+The group should normally be a Mooncake Backend process group. EP uses it to
+exchange RDMA memory-region metadata, QP information, GID/LID information, and
+IPC handles. The backend active-rank state is also used by the fallback path and
+by peer synchronization.
+
+After PG recovery changes membership, call `Buffer.update_ep_member()` to refresh
+EP peer metadata and QPs before relying on the recovered ranks for dispatch and
+combine.
+
+## Runtime objects
+
+### Python wrapper
+
+`mooncake.mooncake_ep_buffer.Buffer` is the user-facing wrapper. It owns:
+
+- the process group;
+- the native EP buffer runtime;
+- the fallback flag;
+- Python fallback buffers used when fast path is unavailable.
+
+`Buffer.connect()` exchanges peer metadata. It first tries the RDMA/IBGDA path,
+then exchanges IPC handles for intra-node P2P. If neither fast path is usable, it
+falls back to a Python implementation based on PyTorch collectives.
+
+### Native buffer
+
+The native `MooncakeEpBuffer` owns or references:
+
+- rank and world-size metadata;
+- a GDR workspace buffer;
+- P2P and RDMA device transports;
+- a communication stream;
+- temporary workspace used by dispatch/combine kernels.
+
+If an external Transfer Engine is supplied by a higher layer, EP can reference
+device transports owned by that engine; otherwise EP creates and owns the
+transports itself.
+
+### Buffer layout
+
+EP allocates paired send/receive buffers for double-buffered dispatch/combine.
+Each pair contains:
+
+- RDMA send signal buffer;
+- RDMA receive signal buffer;
+- RDMA send data buffer;
+- RDMA receive data buffer.
+
+The workspace size returned by `Buffer.get_ep_buffer_size_hint()` is derived from
+`num_max_dispatch_tokens_per_rank`, `hidden`, `num_ranks`, and `num_experts`.
+Size this for peak dispatch demand, not the average request.
+
+## Fast paths and fallback
+
+Mooncake EP has three broad execution modes:
+
+| Mode | Purpose | Notes |
+| --- | --- | --- |
+| IBGDA / RDMA fast path | Inter-node GPU memory movement. | Requires RDMA-capable environment and successful metadata/QP setup. |
+| P2P / IPC fast path | Intra-node peer access such as NVLink. | Requires peer accessibility and IPC handle exchange. |
+| Python fallback | Functional fallback for unsupported environments. | Slower; useful for correctness and limited testing. |
+
+The runtime reports whether IBGDA is disabled and whether the fast path is usable
+through native helper methods surfaced in Python. The Python wrapper updates its
+fallback flag after metadata exchange.
+
+## Metadata exchange
+
+During `Buffer.connect()`, ranks exchange:
+
+- RDMA memory-region address and key;
+- local QP numbers;
+- LID and GID information;
+- subnet prefix and interface ID;
+- CUDA IPC handles for local peer access;
+- current active-rank mask from the backend.
+
+The exchange uses `dist.all_gather()` and `dist.all_to_all()` on the process
+group. For this reason, the process group must already be initialized and healthy
+before the EP buffer is constructed or refreshed.
+
+## Dispatch internals
+
+`dispatch()` takes local token hidden states and selected expert IDs. It sends
+tokens to expert-owner ranks and packs received tokens into local-expert-major
+layout.
+
+Important inputs:
+
+- `x`: `[num_tokens, hidden]` token hidden states;
+- `topk_idx`: `[num_tokens, top_k]` global expert IDs;
+- `active_ranks`: `[num_ranks]` int32 rank-health tensor;
+- `num_max_dispatch_tokens_per_rank`: receive capacity per source rank;
+- `num_experts`: global expert count, divisible by `num_ranks`;
+- `timeout_us`: failure-detection timeout.
+
+Important outputs:
+
+- packed local-expert input tensor;
+- per-local-expert receive counts;
+- source/layout metadata handle used by `combine()`;
+- event/hook synchronization helpers.
+
+When `use_fp8=True`, dispatch returns packed FP8 data and FP32 scales. The local
+expert path must either consume that format directly or dequantize before expert
+compute.
+
+## Combine internals
+
+`combine()` sends local expert outputs back to token-owner ranks and applies
+routing weights. It consumes the handle produced by the matching `dispatch()`
+call.
+
+For zero-copy combine, call `get_next_combine_buffer(handle)`, write expert
+outputs into that buffer, then call `combine(..., zero_copy=True)` with the same
+handle. Do not reuse a handle across unrelated dispatch/combine pairs.
+
+## Rank activeness and timeout behavior
+
+Mooncake EP receives an `active_ranks` tensor in both dispatch and combine. The
+native kernels poll receive signals from source ranks. If `timeout_us` is not
+`-1` and a source rank does not make progress before the timeout, the kernel can
+mark `active_ranks[src_rank] = 0` and skip that source.
+
+This EP-level tensor is rank-level and should have shape `[num_ranks]`. It is
+related to, but not automatically identical to, the backend-level active-rank
+mask passed through `MooncakeBackendOptions`. Integrations should propagate
+health updates consistently between scheduling logic, PG state, and EP buffers.
+
+## Stream synchronization model
+
+Mooncake EP operations return an `EventOverlap` object and, optionally, a hook:
+
+- If `return_recv_hook=False`, call `event.current_stream_wait()` before using
+  the output tensors on the current stream.
+- If `return_recv_hook=True`, call the returned `hook()` at the chosen overlap
+  point.
+- If `async_finish=True`, the wrapper records extra tensors in the event helper
+  to keep lifetimes safe for asynchronous use and CUDA graph scenarios.
+
+Always make synchronization explicit when composing EP with custom expert
+kernels, CUDA graphs, or application-level streams.
+
+## Recovery integration
+
+When a rank fails, multiple layers must be updated:
+
+1. PG active-rank state must stop collectives from waiting for the failed rank.
+2. Scheduler / MoE routing should stop assigning tokens to unavailable experts.
+3. Replacement ranks should join through the PG elastic protocol.
+4. EP buffers should refresh peer metadata with `update_ep_member()` after the
+   process group activates recovered ranks.
+
+EP timeout detection can mark a rank inactive in the EP-level tensor, but higher
+layers still need to coordinate recovery and routing decisions.
+
+## Performance considerations
+
+- Prefer fast path operation with working RDMA/IBGDA or P2P peer access.
+- Size `num_max_dispatch_tokens_per_rank` for the worst expected per-rank token
+  count to avoid overflow.
+- Use `return_recv_hook=True` or `async_finish=True` only when the application
+  deliberately overlaps communication and compute.
+- `use_fp8=True` reduces dispatch bandwidth but requires FP8-aware expert code
+  or explicit dequantization.
+- Repeatedly reconstructing EP buffers is expensive; refresh membership only
+  when PG membership changes.
+
+## Developer test checklist
+
+For EP changes, run correctness across:
+
+- BF16 and FP8 dispatch;
+- zero-copy and non-zero-copy combine;
+- synchronous event wait and hook-based synchronization;
+- fallback and fast path when available;
+- failure simulation with finite `timeout_us`;
+- multi-rank topologies where `num_experts % num_ranks == 0`.
+
+Useful entry points:
+
+```bash
+# EP grid correctness test
+python mooncake-ep/tests/test_ep_grid.py
+
+# Wheel-level EP smoke test
+python mooncake-wheel/tests/test_mooncake_ep.py
+```
+
+Adapt launch commands to the target environment and number of GPUs.
+
+## Related documentation
+
+- [Mooncake Backend (PG) design](mooncake-backend-pg.md)
+- [Python API reference](../python-api-reference/ep-backend.md)
+- [PG/EP troubleshooting](../troubleshooting/pg-ep-troubleshooting.md)
diff --git a/docs/source/design/mooncake-store.md b/docs/source/design/mooncake-store.md
index a171e6e2..2ebb67bd 100644
--- a/docs/source/design/mooncake-store.md
+++ b/docs/source/design/mooncake-store.md
@@ -6,10 +6,11 @@ Mooncake Store is a high-performance **distributed key-value (KV) cache storage
 
 Unlike traditional caching systems such as Redis or Memcached, Mooncake Store is positioned as **a distributed KV cache rather than a generic caching system**. The key difference is that in the latter, the key is derived from the value through hashing, so value is immutable after inserting (although the key/value pair may be evicted).
 
-Mooncake Store provides low-level object storage and management capabilities, including configurable caching and eviction strategies that offers high memory efficiency and is specifically designed to accelerate LLM inference performance.
+Mooncake Store provides low-level object storage and management capabilities, including configurable caching and eviction strategies that offer high memory efficiency and is specifically designed to accelerate LLM inference performance.
 
 Key features of Mooncake Store include:
 - **Object-level storage operations**: Mooncake Store provides simple and easy-to-use object-level APIs, including `Put`, `Get`, and `Remove` operations.
+- **Optional object grouping**: Related objects can carry an optional group ID so that the Master can route their metadata to the same shard and apply best-effort shared lifecycle behavior.
 - **Multi-replica support**: Mooncake Store supports storing multiple data replicas for the same object, effectively alleviating hotspots in access pressure. Each slice within an object is guaranteed to be placed in different segments, while different objects' slices may share segments. Replication operates on a best-effort basis.
 - **Strong consistency**: Mooncake Store guarantees that `Get` operations always return correct and complete data. Once an object has been successfully `Put`, it remains immutable until removal, ensuring that all subsequent `Get` requests retrieve the most recent value.
 - **Zero-copy, bandwidth-saturating transfers**: Powered by the Transfer Engine, Mooncake Store eliminates redundant memory copies and exploits multi-NIC GPUDirect RDMA pooling to drive data across the network at full line rate while keeping CPU overhead negligible.
@@ -48,198 +49,33 @@ If the current leader fails or becomes partitioned from the network, the remaini
 
 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.
 
+(client-c-api)=
 ## Client C++ API
 
-### Constructor and Initialization `Init`
+The `Client` class provides the primary interface for Mooncake Store operations:
 
-```C++
-ErrorCode Init(const std::string& local_hostname,
-               const std::string& metadata_connstring,
-               const std::string& protocol,
-               void** protocol_args,
-               const std::string& master_server_entry);
-```
+| API | Description |
+|-----|-------------|
+| `Init` | Initialize the client with metadata server, protocol, and master address |
+| `Get` | Retrieve object data into pre-registered local memory slices |
+| `Put` | Store object data with configurable replication and persistence |
+| `Upsert` / `BatchUpsert` | Insert or update with existing placement reuse |
+| `Remove` | Delete an object and all its replicas |
+| `CreateCopyTask` / `CreateMoveTask` | Asynchronous cross-node data transfer |
+| `QueryTask` | Monitor the status of async copy/move tasks |
+| `BatchQueryIp` | Discover network locations of storage nodes |
+| `BatchReplicaClear` | Batch clear replicas on specific segments |
+| `QueryByRegex` / `RemoveByRegex` | Query or delete objects matching a regex |
 
-Initializes the Mooncake Store client. The parameters are as follows:
-- `local_hostname`: The `IP:Port` of the local machine or an accessible domain name (default value used if port is not included)
-- `metadata_connstring`: The address of the metadata service (e.g., etcd/Redis) required for Transfer Engine initialization
-- `protocol`: The protocol supported by the Transfer Engine, including RDMA and TCP
-- `protocol_args`: Protocol parameters required by the Transfer Engine
-- `master_server_entry`: The address information of the Master (`IP:Port` for default mode and `etcd://IP:Port;IP:Port;...;IP:Port` for high availability mode)
+For full API signatures, parameter details, and usage examples, see the [Mooncake Store C++ API Reference](../api-reference/cpp/mooncake-store.md).
 
-### Get
-
-```C++
-tl::expected<void, ErrorCode> Get(const std::string& object_key,
-                                  std::vector<Slice>& slices);
-```
-
-![mooncake-store-simple-get](../image/mooncake-store-simple-get.png)
-
-`Get` retrieves the value of `object_key` into the provided `slices`. The returned data is guaranteed to be complete and correct. Each slice must reference local DRAM/VRAM memory that has been pre-registered with `registerLocalMemory(addr, len)` (not the global segments that contribute to the distributed memory pool). When persistence is enabled and the requested data is not found in the distributed memory pool, `Get` will fall back to loading the data from SSD.
-
-### Put
-
-```C++
-tl::expected<void, ErrorCode> Put(const ObjectKey& key,
-                                  std::vector<Slice>& slices,
-                                  const ReplicateConfig& config);
-```
-
-![mooncake-store-simple-put](../image/mooncake-store-simple-put.png)
-
-`Put` stores the value associated with `key` in the distributed memory pool. The `config` parameter allows specifying the required number of replicas as well as the preferred segment for storing the value. When persistence is enabled, `Put` also asynchronously triggers a persistence operation to SSD.
-
-**Replication Guarantees and Best Effort Behavior:**
-- Each slice of an object is guaranteed to be replicated to different segments, ensuring distribution across separate storage nodes
-- Different slices from different objects may be placed in the same segment
-- Replication operates on a best-effort basis: if insufficient space is available for all requested replicas, the object will still be written with as many replicas as possible
-
-The data structure details of `ReplicateConfig` are as follows:
-
-```C++
-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++
-tl::expected<void, ErrorCode> Remove(const ObjectKey& key);
-```
-
-Used to delete the object corresponding to the specified key. This interface marks all data replicas associated with the key in the storage engine as deleted, without needing to communicate with the corresponding storage node (Client).
-
-### CreateCopyTask
-
-```C++
-tl::expected<UUID, ErrorCode> CreateCopyTask(
-    const std::string& key,
-    const std::vector<std::string>& targets);
-```
-
-![mooncake-store-create-copy-task](../image/mooncake-store-client-create-copy-task.png)
-
-`CreateCopyTask` creates an asynchronous copy task that will be executed by the client's task execution system. This is useful when you want to submit multiple copy operations without waiting for each one to complete. The task is submitted to the master service, assigned a unique task ID, and executed asynchronously by an available client. The task status can be queried using `QueryTask`.
-
-**Task Execution and Result Reporting:**
-1. **Task Assignment**: The master service assigns the task to an available client during the client's periodic ping operation
-2. **Task Execution**: The assigned client executes the copy operation asynchronously in a background thread pool
-3. **Result Reporting**: Upon completion (success or failure), the client automatically reports the result to the master service via `MarkTaskToComplete`:
-   - On success: `status = SUCCESS`, `message = "Task completed successfully"`
-   - On failure: `status = FAILED`, `message = <error description>`
-4. **Status Query**: You can query the task status at any time using `QueryTask` to monitor progress
-
-### CreateMoveTask
-
-```C++
-tl::expected<UUID, ErrorCode> CreateMoveTask(
-    const std::string& key,
-    const std::string& source,
-    const std::string& target);
-```
-
-![mooncake-store-create-move-task](../image/mooncake-store-client-create-move-task.png)
-
-`CreateMoveTask` creates an asynchronous move task that will be executed by the client's task execution system. This is useful when you want to submit multiple move operations without waiting for each one to complete. The task is submitted to the master service, assigned a unique task ID, and executed asynchronously by an available client. The task status can be queried using `QueryTask`.
-
-**Task Execution and Result Reporting:**
-1. **Task Assignment**: The master service assigns the task to an available client during the client's periodic ping operation
-2. **Task Execution**: The assigned client executes the move operation asynchronously in a background thread pool
-3. **Result Reporting**: Upon completion (success or failure), the client automatically reports the result to the master service via `MarkTaskToComplete`:
-   - On success: `status = SUCCESS`, `message = "Task completed successfully"`
-   - On failure: `status = FAILED`, `message = <error description>`
-4. **Status Query**: You can query the task status at any time using `QueryTask` to monitor progress
-
-### QueryTask
-
-```C++
-tl::expected<QueryTaskResponse, ErrorCode> QueryTask(const UUID& task_id);
-```
-
-`QueryTask` queries the status of an asynchronous task (copy or move). This allows you to monitor the progress of task-based operations. The response includes task status, type, creation time, last update time, assigned client, and status message.
-
-The data structure details of `QueryTaskResponse` are as follows:
-
-```C++
-struct QueryTaskResponse {
-    UUID id;                                    // Task UUID
-    TaskType type;                              // Task type (REPLICA_COPY or REPLICA_MOVE)
-    TaskStatus status;                          // Task status (PENDING, PROCESSING, SUCCESS, or FAILED)
-    int64_t created_at_ms_epoch;                // Task creation timestamp in milliseconds
-    int64_t last_updated_at_ms_epoch;           // Last update timestamp in milliseconds
-    UUID assigned_client;                       // UUID of the client assigned to execute the task
-    std::string message;                        // Status message or error description
-};
-```
-
-### BatchQueryIp
-
-```C++
-tl::expected<std::unordered_map<UUID, std::vector<std::string>, boost::hash<UUID>>, ErrorCode>
-BatchQueryIp(const std::vector<UUID>& client_ids);
-```
-
-Used to batch query the IP addresses for multiple client IDs. For each client ID in the input list, this interface retrieves the unique IP addresses from all segments mounted by that client. The operation is performed on the Master Service and returns a map from client ID to their IP address lists. Only client IDs that have successfully mounted segments are included in the result map. This is useful for discovering the network locations of storage nodes in the cluster.
-
-### BatchReplicaClear
-
-```C++
-tl::expected<std::vector<std::string>, ErrorCode>
-BatchReplicaClear(const std::vector<std::string>& object_keys,
-                  const UUID& client_id,
-                  const std::string& segment_name);
-```
-
-Used to batch clear replicas for multiple object keys belonging to a specific client ID. This interface allows clearing replicas either on a specific segment or across all segments. If segment_name is empty, all replicas of the specified objects are cleared (the objects are deleted entirely). If segment_name is provided, only replicas located on that specific segment are cleared. The operation is performed on the Master Service and returns a list of object keys that were successfully cleared. Only objects that belong to the specified `client_id`, have expired leases, and meet the clearing criteria are processed. This is useful for managing storage resources and cleaning up data on specific storage nodes.
-
-### QueryByRegex
-
-```C++
-tl::expected<std::unordered_map<std::string, std::vector<Replica::Descriptor>>, ErrorCode>
-QueryByRegex(const std::string& str);
-```
-
-Used to query the replica information for all objects whose keys match the given regular expression. This is useful for batch operations or for retrieving a group of related objects. The operation is performed on the Master and returns a map of keys to their replica lists.
-
-### RemoveByRegex
-
-```C++
-tl::expected<long, ErrorCode> RemoveByRegex(const ObjectKey& str);
-```
-
-Used to delete all objects from the store whose keys match the specified regular expression. This provides a powerful way to perform bulk deletions. The command returns the number of objects that were successfully removed.
-
-### Master Service
+## Master Service
 
 The cluster's available resources are viewed as a large resource pool, managed centrally by a Master process for space allocation and guiding data replication
 
 **Note: The Master Service does not take over any data flow, only providing corresponding metadata information.**
 
-#### Snapshot & Restore
+### Snapshot & Restore
 
 To reduce cache warm-up time after a master restart, the Master Service supports periodic snapshots of its in-memory metadata and recovery from these snapshots.
 
@@ -255,7 +91,47 @@ To reduce cache warm-up time after a master restart, the Master Service supports
 >
 > The snapshot storage location is **exclusively managed** by the Mooncake snapshot system. Old snapshots are automatically deleted during cleanup. **DO NOT store other files in this location.** Use a dedicated, isolated storage for snapshots.
 
-#### Master Service APIs
+### Tenant Quota
+
+The Master Service can optionally enforce strict multi-tenant memory quota admission. This feature is disabled by default. When `enable_multi_tenants=false`, request tenant IDs are ignored for object placement, all objects use the `default` namespace, and tenant quota management requests return `UNAVAILABLE_IN_CURRENT_MODE`.
+
+When strict multi-tenant mode is enabled, the tenant quota policy is loaded from the configured connector. The v1 connector is a writable YAML file configured by `tenant_quota_connector_type=file` and `tenant_quota_connector_uri=<path>`. Tenants must be explicitly present in that connector policy before they can write. Missing tenants, empty tenants, and an unregistered `default` tenant are rejected with `TENANT_NOT_REGISTERED`.
+
+The YAML policy uses schema version `1`:
+
+```yaml
+version: 1
+
+tenants:
+  - name: tenant-a
+    quota: 200GB
+```
+
+Tenant names must be non-empty, unique, must not start with `_`, and must not contain NUL or control characters. Quotas must be positive integers and may use `B`, `KB`, `MB`, `GB`, or `TB` units.
+
+Effective quota is recomputed from the current registered memory capacity:
+
+- If explicit tenant requests fit within the registered memory capacity, tenants receive their requested quotas and remaining capacity stays unallocated.
+- If explicit tenant requests exceed registered memory capacity, explicit tenants receive quota scaled proportionally by request size.
+- Remainders are assigned deterministically by tenant ID, so repeated recomputes produce stable results.
+- Tenants present in restored metadata but missing from the connector policy become in-memory orphans with requested quota `0`, effective quota `0`, and `over_quota=true` while they still own metadata. Reads and removals are allowed so operators can clean them up; writes remain blocked until the tenant is re-registered or emptied.
+
+`PutStart` and size-changing `UpsertStart` charge quota before memory is allocated. If the first reservation fails, the master performs tenant-scoped memory eviction for the target tenant and retries the reservation. The retry is bounded to two eviction attempts. Tenant quota eviction scans only the target tenant, skips hard-pinned objects, honors soft-pin eviction configuration, and preserves grouped-object lease safety checks.
+
+Admin policy changes are persisted before the final in-memory policy is applied. `PUT` writes the connector first and then applies the policy in memory. `DELETE` first marks the tenant unregistered in memory to block concurrent writes, verifies the tenant is empty, writes the connector, and rolls back the in-memory mark if the connector write fails. The admin HTTP API exposes:
+
+| Method | Path | Description |
+|--------|------|-------------|
+| `GET` | `/api/v1/tenant_quotas` | List quota snapshots for active or explicit tenants |
+| `GET` | `/api/v1/tenant_quotas?tenant_id=<tenant>` | Query one tenant quota snapshot |
+| `PUT` | `/api/v1/tenant_quotas?tenant_id=<tenant>` | Create or update a tenant quota policy |
+| `DELETE` | `/api/v1/tenant_quotas?tenant_id=<tenant>` | Delete an empty tenant quota policy |
+
+Tenant quota snapshots include `tenant_id`, `requested_quota_bytes`, `effective_quota_bytes`, `used_bytes`, `reserved_bytes`, `committed_count`, `metadata_object_count`, `over_quota`, and `has_explicit_policy`.
+
+Snapshots restore object runtime state only. Tenant quota policy is always loaded from the connector after metadata restore, then usage and effective quota are rebuilt from restored metadata and current registered capacity. If the connector cannot be loaded in strict multi-tenant mode, startup fails.
+
+### Master Service APIs
 
 The protobuf definition between Master and Client is as follows:
 
@@ -498,7 +374,7 @@ message UnMountSegmentResponse {
 
 When the space needs to be released, this interface is used to remove the previously mounted resources from the Master Service.
 
-#### Object Information Maintenance
+### Object Information Maintenance
 
 The Master Service needs to maintain mappings related to buffer allocators and object metadata to efficiently manage memory resources and precisely control replica states in multi-replica scenarios. Additionally, the Master Service uses read-write locks to protect critical data structures, ensuring data consistency and security in multi-threaded environments. The following are the interfaces maintained by the Master Service for storage space information:
 
@@ -590,7 +466,25 @@ tl::expected<long, ErrorCode> RemoveByRegex(const std::string& str);
 
 The Client requests the Master Service to delete all replicas corresponding to the specified key or for all object keys that match the specified regular expression.
 
-### Buffer Allocator
+### Optional Object Groups
+
+Object groups are intended for workloads where one logical cache entry is split into multiple Mooncake Store objects, such as separate K/V tensors, parallel shards, or auxiliary index objects. Without grouping, these objects are managed independently, so lease refresh and memory eviction may affect different parts of the same logical entry at different times.
+
+Mooncake Store remains an object-oriented KV cache: objects are still put, queried, and removed by key. For workloads where one logical cache unit is represented by multiple physical objects, callers may attach optional group metadata through `ReplicateConfig::group_ids` during `Put`, `BatchPut`, `Upsert`, or `BatchUpsert`.
+
+For single-object writes, `group_ids` contains one entry. For batch writes, it must have the same length as the key list, and entry `i` is the group ID for key `i`. An empty string stores that key as ungrouped, and leaving the field unset preserves the legacy ungrouped behavior. For an existing object, group membership is immutable: `Upsert` may preserve the existing group, but it cannot move the object to another group or clear its group while the object exists.
+
+On the Master side, group state is tenant-scoped. Objects with a non-empty group ID are routed to the metadata shard selected by `hash(group_id)`, and the Master keeps a tenant-scoped object-to-group routing index so existing key-based APIs can still locate grouped objects. The Master tracks only the current member set of each group; it does not require an expected member count, a member index, or a commit protocol for group completeness.
+
+Group metadata affects lifecycle behavior on a best-effort basis:
+
+- `ExistKey` and `GetReplicaList` refresh the lease, and the soft-pin timeout if present, for the current members of the group.
+- Memory eviction expands a grouped candidate to the group's current members and then applies the existing per-object safety checks. Members with active leases, hard pins, soft pins when soft-pin eviction is disabled, incomplete writes, busy replicas, or unavailable replica states are skipped.
+- Object removal APIs, copy/move tasks, and NoF eviction keep their existing object-level semantics. Group routing and membership metadata are cleaned up when objects are removed.
+
+This design is intentionally lightweight and backward compatible. Grouping should be treated as a lifecycle hint for related objects, not as a transactional guarantee that all members are created, made visible, or evicted atomically.
+
+## Buffer Allocator
 
 The buffer allocator serves as a low-level memory management component within the Mooncake Store system, primarily responsible for efficient memory allocation and deallocation. It builds upon underlying memory allocators to perform its functions.
 
@@ -600,6 +494,8 @@ Mooncake Store provides two concrete implementations of `BufferAllocatorBase`:
 
 **OffsetBufferAllocator (default and recommended)**: This allocator is derived from [OffsetAllocator](https://github.com/sebbbi/OffsetAllocator), which uses a custom bin-based allocation strategy that supports fast hard realtime `O(1)` offset allocation with minimal fragmentation. Mooncake Store optimizes this allocator based on the specific memory usage characteristics of LLM inference workloads, thereby enhancing memory utilization in LLM scenarios.
 
+For measured utilization and allocation latency across LLM-style workloads, see [Allocator Performance](../performance/allocator-benchmark-result.md).
+
 **CachelibBufferAllocator (deprecated)**: This allocator leverages Facebook's [CacheLib](https://github.com/facebook/CacheLib) to manage memory using a slab-based allocation strategy. It provides efficient memory allocation with good fragmentation resistance and is well-suited for high-performance scenarios. However, in our modified version, it does not handle workloads with highly variable object sizes effectively, so it is currently marked as deprecated.
 
 Users can choose the allocator that best matches their performance and memory usage requirements through the `--memory-allocator` startup parameter of `master_service`.
@@ -620,7 +516,15 @@ class BufferAllocatorBase {
 
 3. **`deallocate` Function**: This function is automatically triggered by the `BufHandle` destructor. It calls the internal allocator to release the associated memory and updates the handle’s status to `BufStatus::UNREGISTERED`.
 
-### AllocationStrategy
+### Client Local Buffer and Python BufferPool
+
+Each Store client can also create a setup-time local buffer through `local_buffer_size`. This memory is registered once with the Transfer Engine and managed by `ClientBufferAllocator` for short-lived client-side staging work.
+
+The Python `BufferPool` reuses this existing local buffer instead of allocating a second registered arena. A pool lease is a sub-allocation from `client_buffer_allocator_`, so the common path avoids per-lease `register_buffer()` and `unregister_buffer()` calls. The pool still keeps the Python-facing lease API, memoryview lifetime checks, blocking acquire semantics, and optional `max_regions` concurrency limiting.
+
+This is a soft-isolation policy: internal Store paths and external Python leases share the local registered buffer, allowing bursty external usage when memory is available rather than reserving a hard partition. If the local buffer is temporarily exhausted, `BufferPool` can allocate and register a short-lived overflow buffer so bursts do not immediately surface as upper-layer errors; that overflow region is unregistered as soon as the lease is released. If callers need to cap long-lived external pressure, they should use pool-level controls such as `max_regions`, `max_bytes`, or acquire timeouts.
+
+## AllocationStrategy
 AllocationStrategy is a strategy class for efficiently managing memory resource allocation and replica storage location selection in a distributed environment. It is mainly used in the following scenarios:
 - Determining the allocation locations for object storage replicas.
 - Selecting suitable read/write paths among multiple replicas.
@@ -630,7 +534,7 @@ AllocationStrategy is used in conjunction with the Master Service and the underl
 - Master Service: Determines the target locations for replica allocation via `AllocationStrategy`.
 - Buffer Allocator: Executes the actual memory allocation and release tasks.
 
-#### APIs
+### APIs
 
 `Allocate`: Finds suitable storage segments from available storage resources to allocate space of a specified size for multiple replicas. Uses best-effort semantics, meaning it allocates as many replicas as possible even if the full requested count cannot be satisfied.
 
@@ -652,7 +556,7 @@ virtual tl::expected<std::vector<Replica>, ErrorCode> Allocate(
   - On success: vector of allocated replicas (may be fewer than requested due to resource constraints, but at least 1)
   - On failure: ErrorCode::NO_AVAILABLE_HANDLE if no replicas can be allocated, ErrorCode::INVALID_PARAMS for invalid configuration
 
-#### Allocation Strategies
+### Allocation Strategies
 
 Mooncake Store provides multiple built-in allocation strategies to control how storage space is distributed across segments. Users can select a strategy via the `--allocation_strategy` flag when starting the master service:
 
@@ -662,7 +566,7 @@ Mooncake Store provides multiple built-in allocation strategies to control how s
 
 Valid values are: `random` (default), `free_ratio_first`, `cxl` (case-sensitive).
 
-##### How to Choose
+#### How to Choose
 
 | Strategy | Best For | Trade-off |
 |---|---|---|
@@ -678,7 +582,9 @@ Valid values are: `random` (default), `free_ratio_first`, `cxl` (case-sensitive)
 
 **Use `cxl`** only when your hardware includes CXL (Compute Express Link) memory devices and you want to allocate data exclusively on CXL segments.
 
-##### Strategy Details
+For benchmark data comparing `random` and `free_ratio_first` across segment counts, replica counts, and skewed capacities, see [AllocationStrategy Performance](../performance/allocation-strategy-benchmark-result.md).
+
+#### Strategy Details
 
 **`random` — RandomAllocationStrategy**
 
@@ -715,21 +621,25 @@ Specialized for CXL (Compute Express Link) memory hardware. Unlike the other str
 
 Limitations: This strategy only supports single-replica allocation (does not distribute across multiple segments) and does not support the `AllocateFrom()` interface.
 
-### Eviction Policy
+## Eviction Policy
 
 When a `PutStart` request fails due to insufficient memory, or when the eviction thread detects that space usage has reached the configured high watermark (95% by default, configurable via `-eviction_high_watermark_ratio`), an eviction task is triggered to free up space by evicting a portion of objects (5% by default, configurable via `-eviction_ratio`). Similar to `Remove`, evicted objects are simply marked as deleted, with no data transfer required.
 
 Currently, an approximate LRU policy is adopted, where the least recently used objects are preferred for eviction. To avoid data races and corruption, objects currently being read or written by clients should not be evicted. For this reason, objects that have leases or have not been marked as complete by `PutEnd` requests will be ignored by the eviction task.
 
-### Lease
+For grouped objects, memory eviction resolves the current group membership and attempts to reclaim eligible members together.
+
+## Lease
 
 To avoid data conflicts, a per-object lease is granted whenever an `ExistKey` request or a `GetReplicaListRequest` request succeeds. While the lease is active, the object is protected from `Remove`, `RemoveAll`, and `Eviction` operations. Specifically, a `Remove` request targeting a leased object will fail, and a `RemoveAll` request will only delete objects without an active lease. This ensures that the object’s data can be safely read as long as the lease has not expired.
 
+For grouped objects, a successful `ExistKey` or `GetReplicaList` refreshes the lease for the current members of the group, so recently accessed members are less likely to be separated by memory eviction.
+
 However, if the lease expires before a `Get` operation finishes reading the data, the operation will be considered failed, and no data will be returned, in order to prevent potential data corruption.
 
 The default lease TTL is 5 seconds and is configurable via a startup parameter of `master_service`.
 
-### Soft Pin
+## Soft Pin
 
 For important and frequently used objects, such as system prompts, Mooncake Store provides a soft pin mechanism. When putting an object, it can be configured to enable soft pin. During eviction, objects that are not soft pinned are prioritized for eviction. Soft pinned objects are only evicted when memory is insufficient and no other objects are eligible for eviction.
 
@@ -743,7 +653,7 @@ 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
+## 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.
 
@@ -755,23 +665,23 @@ Key differences from soft pin:
 - 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
+## 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.
 
-#### `PutStart` Preemption
+### `PutStart` Preemption
 
 If an object receives neither a `PutEnd` nor a `PutRevoke` request within `put_start_discard_timeout` (default: 30 seconds) after its `PutStart`, any subsequent `PutStart` request for the same object will be allowed to "preempt" the previous `PutStart`. This enables the new request to proceed with writing the object, thereby preventing a single faulty Client from permanently blocking access to that object. Note that during such preemption, the storage space allocated by the old `PutStart` is not reused; instead, new space is allocated for the preempting `PutStart`. The space previously allocated by the old `PutStart` will be reclaimed via the mechanism described below.
 
-#### Space Reclaim
+### Space Reclaim
 
 Replica space allocated during a `PutStart` is considered releasable by the Master if the write operation is neither completed (via `PutEnd`) nor canceled (via `PutRevoke`) within `put_start_release_timeout` (default: 10 minutes) after the `PutStart`. When object eviction is triggered—either due to allocation failures or because storage utilization exceeds the configured threshold—these releasable replica spaces are prioritized for release to reclaim storage capacity.
 
-### Preferred Segment Allocation
+## Preferred Segment Allocation
 
 Mooncake Store provides a **preferred segment allocation** feature that allows users to specify a preferred storage segment (node) for object allocation. This feature is particularly useful for optimizing data locality and reducing network overhead in distributed scenarios.
 
-#### How It Works
+### How It Works
 
 The preferred segment allocation feature is implemented through the `AllocationStrategy` system and is controlled via the `preferred_segment` field in the `ReplicateConfig` structure:
 
@@ -795,23 +705,23 @@ When a `Put` operation is initiated with a non-empty `preferred_segment` value,
 - **Data Locality**: By preferring local segments, applications can reduce network traffic and improve access performance for frequently used data.
 - **Load Balancing**: Applications can distribute data across specific nodes to achieve better load distribution.
 
-### Multi-layer Storage Support
+## Multi-layer Storage Support
 
 This system provides support for a hierarchical cache architecture, enabling efficient data access through a combination of in-memory caching and persistent storage. Data is initially stored in memory cache and asynchronously backed up to a Distributed File System (DFS), forming a two-tier "memory-SSD persistent storage" cache structure.
 
-#### Enabling Persistence Functionality
+### Enabling Persistence Functionality
 
 When the user specifies `--root_fs_dir=/path/to/dir` when starting the master, and this path is a valid DFS-mounted directory on all machines where the clients reside, Mooncake Store's tiered caching functionality will work properly. Additionally, during master initialization, a `cluster_id` is loaded. This ID can be specified during master initialization (`--cluster_id=xxxx`). If not specified, the default value `mooncake_cluster` will be used. Subsequently, the root directory for client persistence will be `<root_fs_dir>/<cluster_id>`.
 
 ​Note​​: When enabling this feature, the user must ensure that the DFS-mounted directory (`root_fs_dir=/path/to/dir`) is valid and consistent across all client hosts. If some clients have invalid or incorrect mount paths, it may cause abnormal behavior in Mooncake Store.
 
-#### Persistent Storage Space Configuration​
+### Persistent Storage Space Configuration​
 Mooncake provides configurable DFS available space. Users can specify `--global_file_segment_size=1048576` when starting the master, indicating a maximum usable space of 1MB on DFS.
 The current default setting is the maximum value of int64 (as we generally do not restrict DFS storage usage), which is displayed as `infinite` in `mooncake_maseter`'s console logs.
 **Notice**  The DFS cache space configuration must be used together with the `--root_fs_dir` parameter. Otherwise, you will observe that the `SSD Storage` usage consistently shows: `0 B / 0 B`
 **Notice** The capability for file eviction on DFS has not been provided yet
 
-#### Data Access Mechanism
+### Data Access Mechanism
 
 The persistence feature also follows Mooncake Store's design principle of separating control flow from data flow. The read/write operations of kvcache objects are completed on the client side, while the query and management functions of kvcache objects are handled on the master side. In the file system, the key -> kvcache object index information is maintained by a fixed indexing mechanism, with each file corresponding to one kvcache object (the filename serves as the associated key name).
 
@@ -820,24 +730,33 @@ After enabling the persistence feature:
 - For each `Put` or `BatchPut` operation, both a synchronous memory pool write operation and an asynchronous DFS persistence operation will be initiated.
 - For each `Get` or `BatchGet` operation, if the corresponding kvcache is not found in the memory pool, the system will attempt to read the file data from DFS and return it to the user.
 
-#### 3FS USRBIO Plugin
+### 3FS USRBIO Plugin (Experimental)
+
+```{note}
+This integration is **experimental** and incomplete; see the plugin page for details before relying on it.
+```
+
 If you need to use 3FS's native API (USRBIO) to achieve high-performance persistent file reads and writes, you can refer to the configuration instructions in this document [3FS USRBIO Plugin](../getting_started/plugin-usage/3FS-USRBIO-Plugin.md).
 
-### Builtin Metadata Server
+## Builtin Metadata Server
 Mooncake Store provides a built-in HTTP metadata server as an alternative to etcd for storing cluster metadata. This feature is particularly useful for development environments or scenarios where etcd is not available.
-#### Configuration Parameters
+### Configuration Parameters
 The HTTP metadata server can be configured using the following parameters:
 - **`enable_http_metadata_server`** (boolean, default: `false`): Enables the built-in HTTP metadata server instead of using etcd. When set to `true`, the master service will start an embedded HTTP server that handles metadata operations.
 - **`http_metadata_server_port`** (integer, default: `8080`): Specifies the TCP port on which the HTTP metadata server will listen for incoming connections. This port must be available and not conflict with other services.
 - **`http_metadata_server_host`** (string, default: `"0.0.0.0"`): Specifies the host address for the HTTP metadata server to bind to. Use `"0.0.0.0"` to listen on all available network interfaces, or specify a specific IP address for security purposes.
-#### Environment Variables
+### Environment Variables
 - MC_STORE_CLUSTER_ID: Identify the metadata when multiple cluster share the same master, default 'mooncake'.
 - MC_STORE_MEMCPY: Enables or disables local memcpy optimization, set to 1/true to enable, 0/false to disable.
 - MC_STORE_CLIENT_METRIC: Enables client metric reporting, enabled by default; set to 0/false to disable.
 - MC_STORE_CLIENT_METRIC_INTERVAL: Reporting interval in seconds, default 0 (collects but does not report).
+- MC_STORE_CLIENT_MIN_PORT: Minimum local port for client connections (default 12300). Must be in range 1024–32767 or 61000–65535; falls back to default on invalid input.
+- MC_STORE_CLIENT_MAX_PORT: Maximum local port for client connections (default 14300). Same range constraints; must be ≥ MC_STORE_CLIENT_MIN_PORT.
 - MC_STORE_USE_HUGEPAGE: Enables huge page support, disabled by default.
 - MC_STORE_HUGEPAGE_SIZE: Specifies the page size of the huge page to use, default 2M.
-#### Usage Example
+- MC_MMAP_ARENA_POOL_SIZE: Size of the pre-allocated arena pool for mmap buffer allocations. Accepts human-readable sizes (e.g., `"8gb"`, `"20gb"`). Providing this variable explicitly enables the arena; when enabled via gflag without an env override, the default pool size is `8gb`. The arena is allocated once at first use and serves subsequent allocations via lock-free atomic bump pointer (~50ns per allocation vs ~1000ns for direct mmap).
+- MC_DISABLE_MMAP_ARENA: Set to `1` to disable the arena allocator and fall back to per-call `mmap()`, even if the arena was explicitly requested. Also accepts `true`, `yes`, or `on`. This must be set before the first Mooncake mmap-buffer allocation in the process. Useful for debugging or memory-constrained environments where pre-allocating a pool is not desirable.
+### Usage Example
 To start the master service with the HTTP metadata server enabled:
 ```bash
 ./build/mooncake-store/src/mooncake_master \
@@ -848,195 +767,15 @@ To start the master service with the HTTP metadata server enabled:
 When enabled, the HTTP metadata server will start automatically and provide metadata services for the Mooncake Store cluster. This eliminates the need for an external etcd deployment, simplifying the setup process for development and testing environments.
 Note that the HTTP metadata server is designed for single-node deployments and does not provide the high availability features that etcd offers. For production environments requiring high availability, etcd is still the recommended choice.
 
+For detailed guidance on monitoring master metrics, Prometheus endpoints, and health checks, see the [Observability guide](../getting_started/observability.md).
+
 ## Mooncake Store Python API
 
 **Complete Python API Documentation**: [https://kvcache-ai.github.io/Mooncake/python-api-reference/mooncake-store.html](https://kvcache-ai.github.io/Mooncake/python-api-reference/mooncake-store.html)
 
-## Compilation and Usage
-Mooncake Store is compiled together with other related components (such as the Transfer Engine).
-
-For default mode:
-```
-mkdir build && cd build
-cmake .. # default mode
-make
-sudo make install # Install Python interface support package
-```
-
-High availability mode:
-```
-mkdir build && cd build
-cmake .. -DSTORE_USE_ETCD # compile etcd wrapper that depends on go
-make
-sudo make install # Install Python interface support package
-```
-
-**Note:** To use high availability mode, only `-DSTORE_USE_ETCD` is required. `-DUSE_ETCD` is a compilation option for the **Transfer Engine** and is **not related** to the high availability mode.
-
-### Starting the Transfer Engine's Metadata Service
-Mooncake Store uses the Transfer Engine as its core transfer engine, so it is necessary to start the metadata service (etcd/redis/http). The startup and configuration of the `metadata` service can be referred to in the relevant sections of [Transfer Engine](./transfer-engine/index.md). **Special Note**: For the etcd service, by default, it only provides services for local processes. You need to modify the listening options (IP to 0.0.0.0 instead of the default 127.0.0.1). You can use commands like curl to verify correctness.
-
-### Starting the Master Service
-The Master Service runs as an independent process, provides gRPC interfaces externally, and is responsible for the metadata management of Mooncake Store (note that the Master Service does not reuse the metadata service of the Transfer Engine). The default listening port is `50051`. After compilation, you can directly run `mooncake_master` located in the `build/mooncake-store/src/` directory. After starting, the Master Service will output the following content in the log:
-```
-Starting Mooncake Master Service
-Port: 50051
-Max threads: 4
-Master service listening on 0.0.0.0:50051
-```
-
-**High availability mode**:
-
-HA mode relies on an etcd service for coordination. If Transfer Engine also uses etcd as its metadata service, the etcd cluster used by Mooncake Store can either be shared with or separate from the one used by Transfer Engine.
-
-HA mode allows deployment of multiple master instances to eliminate the single point of failure. Each master instance must be started with the following parameters:
-```
---enable-ha: enables high availability mode
---etcd-endpoints: specifies endpoints for etcd service, separated by ';'
---rpc-address: the RPC address of this instance. Note that the address specified here should be accessible to the client.
-```
-
-For example:
-```
-./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-address=10.0.0.1
-```
-
-### Starting the Sample Program
-Mooncake Store provides various sample programs, including interface forms based on C++ and Python. Below is an example of how to run using `stress_cluster_benchmark`.
-
-1. Open `stress_cluster_benchmark.py` and update the initialization settings based on your network environment. Pay particular attention to the following fields:
-`local_hostname`: the IP address of the local machine
-`metadata_server`: the address of the Transfer Engine metadata service
-`master_server_address`: the address of the Master Service
-**Note**: The format of `master_server_address` depends on the deployment mode. In default mode, use the format `IP:Port`, specifying the address of a single master node. In HA mode, use the format `etcd://IP:Port;IP:Port;...;IP:Port`, specifying the addresses of the etcd cluster endpoints.
-For example:
-```python
-import os
-import time
-
-from distributed_object_store import DistributedObjectStore
-
-store = DistributedObjectStore()
-# Protocol used by the transfer engine, optional values are "rdma" or "tcp"
-protocol = os.getenv("PROTOCOL", "tcp")
-# Device name used by the transfer engine
-device_name = os.getenv("DEVICE_NAME", "ibp6s0")
-# Hostname of this node in the cluster, port number is randomly selected from (12300-14300)
-local_hostname = os.getenv("LOCAL_HOSTNAME", "localhost")
-# Metadata service address of the Transfer Engine, here etcd is used as the metadata service
-metadata_server = os.getenv("METADATA_ADDR", "127.0.0.1:2379")
-# The size of the Segment mounted by each node to the cluster, allocated by the Master Service after mounting, in bytes
-global_segment_size = 3200 * 1024 * 1024
-# Local buffer size registered with the Transfer Engine, in bytes
-local_buffer_size = 512 * 1024 * 1024
-# Address of the Master Service of Mooncake Store
-master_server_address = os.getenv("MASTER_SERVER", "127.0.0.1:50051")
-# Data length for each put()
-value_length = 1 * 1024 * 1024
-# Total number of requests sent
-max_requests = 1000
-# Initialize Mooncake Store Client
-retcode = store.setup(
-    local_hostname,
-    metadata_server,
-    global_segment_size,
-    local_buffer_size,
-    protocol,
-    device_name,
-    master_server_address,
-)
-```
-
-2. Run `ROLE=prefill python3 ./stress_cluster_benchmark.py` on one machine to start the Prefill node.
-   For "rdma" protocol, you can also enable topology auto discovery and filters, e.g., `ROLE=prefill MC_MS_AUTO_DISC=1 MC_MS_FILTERS="mlx5_1,mlx5_2" python3 ./stress_cluster_benchmark.py`.
-
-3. Run `ROLE=decode python3 ./stress_cluster_benchmark.py` on another machine to start the Decode node.
-   For "rdma" protocol, you can also enable topology auto discovery and filters, e.g., `ROLE=decode MC_MS_AUTO_DISC=1 MC_MS_FILTERS="mlx5_1,mlx5_2" python3 ./stress_cluster_benchmark.py`.
-
-The absence of error messages indicates successful data transfer.
-
-### Starting the Client as Standalone Process and accessing via RPC
-To start a RPC type **real** `Client` as a standalone process, you can use the following command:
-
-```bash
-./build/mooncake-store/src/mooncake_client \
-    --global_segment_size="4GB" \
-    --master_server_address="localhost:50051" \
-    --metadata_server="http://localhost:8080/metadata"
-```
-
-Next, a **real** `Client` instance is created and connected to the Master Service. The **real** `Client` instance is listening on port 50052 as default.
-If you want to send requests to it, a **dummy** `Client` should be used in the application process (e.g., vLLM, SGLang). You can start a **dummy** `Client`
-with specific parameters defined in the application.
-
-The **real** `Client` can be configured using the following parameters:
-
-- **`host`**: (string, default: "0.0.0.0"): The hostname of the client.
-
-- **`port`**: (int, default: 50052): The port number the client service listens on.
-
-- **`global_segment_size`**: (string, default: "4GB"): The size of the global segment to be allocated by the client.
-
-- **`master_server_address`**: (string, default: "localhost:50051"): The address of the Master Service.
-
-- **`metadata_server`**: (string, default: "http://localhost:8080/metadata"): The address of the metadata service.
-
-- **`protocol`**: (string, default: "tcp"): The protocol used by the Transfer Engine.
-
-- **`device_name`**: (string, default: ""): The device name used by the Transfer Engine.
-
-- **`threads`**: (int, default: 1): The number of threads used by the client.
-
-### Starting the Client as Standalone Process and accessing via HTTP
-
-Use `mooncake-wheel/mooncake/mooncake_store_service.py` to start a **real** `Client` as a standalone process and accessing via HTTP.
-
-First, create and save a configuration file in JSON format. For example:
-
-```
-{
-    "local_hostname": "localhost",
-    "metadata_server": "http://localhost:8080/metadata",
-    "global_segment_size": 268435456,
-    "local_buffer_size": 268435456,
-    "protocol": "tcp",
-    "device_name": "",
-    "master_server_address": "localhost:50051"
-}
-```
-
-Then run `mooncake_store_service.py`. This program starts an HTTP server alongside the **real** `Client`. Through this server, users can manually perform operations such as `Get` and `Put`, which is useful for debugging.
-
-The main startup parameters include:
-
-* `config`: Path to the configuration file.
-* `port`: Port number for the HTTP server.
-
-Suppose the `mooncake_transfer_engine` wheel package is already installed, the following command starts the program:
-```bash
-python -m mooncake.mooncake_store_service --config=[config_path] --port=8081
-```
-
-### Set the Log Level for yalantinglibs coro_rpc and coro_http
-By default, the log level is set to warning. You can customize it using the following environment variable:
-`export MC_YLT_LOG_LEVEL=info`
-This sets the log level for yalantinglibs (including coro_rpc and coro_http) to info.
-Available log levels: trace, debug, info, warn (or warning), error, and critical.
-
-## Example Code
-
-#### Python Usage Example
-We provide a reference example `distributed_object_store_provider.py`, located in the `mooncake-store/tests` directory. To check if the related components are properly installed, you can run etcd and Master Service (`mooncake_master`) in the background on the same server, and then execute this Python program in the foreground. It should output a successful test result.
-
-#### C++ Usage Example
-The C++ API of Mooncake Store provides more low-level control capabilities. We provide a reference example `client_integration_test`, located in the `mooncake-store/tests` directory. To check if the related components are properly installed, you can run etcd and Master Service (`mooncake_master`) on the same server, and then execute this C++ program (located in the `build/mooncake-store/tests` directory). It should output a successful test result.
-
 ## Version Management Policy
 
-The current version of Mooncake Store is defined in [`CMakeLists.txt`](../../mooncake-store/CMakeLists.txt) as `project(MooncakeStore VERSION 2.0.0)`.
+The current version of Mooncake Store is defined in [`CMakeLists.txt`](gh-file:mooncake-store/CMakeLists.txt) as `project(MooncakeStore VERSION 2.0.0)`.
 
 When to bump the version:
 
diff --git a/docs/source/design/ssd-free-ratio-first-allocation.md b/docs/source/design/ssd-free-ratio-first-allocation.md
new file mode 100644
index 00000000..337f171c
--- /dev/null
+++ b/docs/source/design/ssd-free-ratio-first-allocation.md
@@ -0,0 +1,152 @@
+# SSD Free-Ratio-First Allocation Design
+
+## Overview
+
+Mooncake Store distributes KV cache objects across multiple memory segments hosted on different nodes. The master's allocation strategy decides which segment receives each new object replica. When using DDR-only strategies such as `random` or `free_ratio_first`, the allocator ignores SSD state entirely. In deployments where some segments have SSD offload enabled and others do not, this blind allocation can concentrate traffic on a small subset of segments whose SSD capacity is quickly exhausted, while segments with ample SSD headroom remain underutilized.
+
+This document describes `SsdFreeRatioFirstAllocationStrategy`, an allocation strategy that ranks candidate segments by their SSD free ratio and preferentially allocates to segments with the most available SSD space. It integrates with the existing allocation framework and adds SSD usage tracking to `LocalDiskSegment` so that the master can make informed placement decisions.
+
+---
+
+## Architecture
+
+```
+                         Allocation Request
+                                │
+                                ▼
+                    ┌───────────────────────┐
+                    │  Sample candidate      │
+                    │  segments (up to       │
+                    │  6 * replica_num)      │
+                    └───────────┬───────────┘
+                                │
+                                ▼
+                    ┌───────────────────────┐
+                    │  Compute SSD free     │
+                    │  ratio per candidate  │
+                    └───────────┬───────────┘
+                                │
+                                ▼
+                    ┌───────────────────────┐
+                    │  Sort by SSD free     │
+                    │  ratio descending     │
+                    └───────────┬───────────┘
+                                │
+                                ▼
+                    ┌───────────────────────┐
+                    │  Allocate from top    │
+                    │  candidates           │
+                    └───────────┬───────────┘
+                                │
+                     ┌──────────┴──────────┐
+                     │  Remaining replicas? │
+                     └──────┬──────┬───────┘
+                            │      │
+                       Yes ◄┘      └► No → Done
+                            │
+                            ▼
+                    ┌───────────────────────┐
+                    │  Fallback: random     │
+                    │  allocation for       │
+                    │  remaining replicas   │
+                    └───────────────────────┘
+```
+
+The flow follows the same high-level structure as the existing `FreeRatioFirstAllocationStrategy`: sample a subset of candidates, compute a ranking metric, sort, and allocate from the top. The key difference is that the ranking metric is SSD free ratio rather than DRAM free ratio.
+
+`MasterService` only passes an `SsdMetricsProvider` when the effective allocation strategy is `SSD_FREE_RATIO_FIRST`. Non-SSD strategies receive `nullptr`, so they avoid unnecessary local-disk segment access.
+
+---
+
+## Core Algorithm
+
+### Candidate sampling
+
+For each allocation request requesting `replica_num` replicas, the strategy samples `min(6 * replica_num, total_segments)` candidate segments. This bounded sampling keeps the sorting cost predictable regardless of cluster size while still providing a statistically diverse candidate set.
+
+### SSD free ratio
+
+For each sampled segment, the SSD free ratio is computed as:
+
+```
+ssd_free_ratio = (ssd_total_capacity - ssd_used_bytes) / ssd_total_capacity
+```
+
+A segment with 1 TB total SSD capacity and 200 GB used has an SSD free ratio of 0.80. A segment whose SSD is full has a ratio of 0.0.
+
+Before calculating the ratio, `ssd_used_bytes` is clamped to `[0, ssd_total_capacity]`. This keeps transient concurrent accounting drift from producing a negative free ratio or a value greater than 1.0. If no SSD metrics provider is available, or if the reported total capacity is not positive, the strategy treats the segment as fully free.
+
+### Sorting
+
+Candidates are sorted by SSD free ratio in descending order. Segments with more available SSD space appear first and are preferred for allocation.
+
+### Preferred segments
+
+As with other allocation strategies, segments marked as preferred by the caller are handled first. Preferred segments bypass the SSD free ratio ranking and are allocated immediately if they have sufficient capacity.
+
+### Fallback to random allocation
+
+After allocating from the SSD-ranked candidates, any remaining replicas that could not be satisfied are allocated using the standard random strategy as a fallback. This ensures that allocation succeeds even when SSD metrics are unavailable (for example, on segments without SSD offload configured).
+
+---
+
+## SSD Usage Tracking
+
+### `ssd_used_bytes` counter
+
+`LocalDiskSegment` maintains `ssd_total_capacity_bytes`, updated by `ReportSsdCapacity`, and an atomic counter `ssd_used_bytes` that tracks the total number of bytes currently occupied by offloaded replicas on the segment's SSD. `ssd_used_bytes` is updated alongside metadata changes:
+
+- **Increment**: `NotifyOffloadSuccess` increments `ssd_used_bytes` by the object size only after the master successfully adds a `LOCAL_DISK` replica to the object entry. If the object has already disappeared from metadata, the notification is ignored and the counter is not changed.
+- **Decrement**: The master decrements `ssd_used_bytes` via `ReleaseLocalDiskUsage` whenever a `LOCAL_DISK` replica is removed from metadata — on full object deletion (`EraseMetadata`), and on partial replica removal through the shared erase helper (`EraseReplicasWithCacheTotalAccounting`, used by batch clear, revoke, and stale-handle cleanup) and the local-disk eviction path. `ReleaseLocalDiskUsage` iterates only `LOCAL_DISK` replicas and is a no-op for other replica types, so it is safe to call on any replica set.
+
+The counter is atomic to allow concurrent updates from multiple RPC handler threads without requiring a separate lock. The allocation strategy treats it as an eventually consistent placement signal and clamps it before computing the free ratio.
+
+### `SsdMetricsProvider` interface
+
+`ScopedLocalDiskSegmentAccess` implements the `SsdMetricsProvider` interface, which exposes two methods:
+
+| Method | Return type | Description |
+|--------|-------------|-------------|
+| `getSsdTotalCapacity` | `int64_t` | Total SSD capacity configured for the segment, in bytes |
+| `getSsdUsedBytes` | `int64_t` | Current SSD usage, read from `ssd_used_bytes` |
+
+The allocation strategy queries these methods through the `SsdMetricsProvider` interface, keeping the strategy decoupled from the concrete segment implementation.
+
+---
+
+## Configuration Parameters
+
+| Parameter | Type | Default | Description |
+|-----------|------|---------|-------------|
+| `allocation_strategy` | string | `"random"` | Set to `"ssd_free_ratio_first"` to enable SSD-ratio-based load balancing |
+
+The parameter is passed as a gflag to the master process at startup.
+
+---
+
+## Code Structure
+
+| File | Change |
+|------|--------|
+| `mooncake-store/include/types.h` | Add `SSD_FREE_RATIO_FIRST` enum value to the allocation strategy enum |
+| `mooncake-store/include/allocation_strategy.h` | Add `SsdMetricsProvider` interface and `SsdFreeRatioFirstAllocationStrategy` class |
+| `mooncake-store/include/segment.h` | Add `ssd_total_capacity_bytes` and `ssd_used_bytes` fields to `LocalDiskSegment`; inherit `SsdMetricsProvider` |
+| `mooncake-store/src/segment.cpp` | Implement `getSsdTotalCapacity` and `getSsdUsedBytes` |
+| `mooncake-store/src/master_service.cpp` | Pass SSD metrics only to `SSD_FREE_RATIO_FIRST`; update `ssd_used_bytes` after successful `NotifyOffloadSuccess` metadata insertion; release usage when `LOCAL_DISK` replicas are erased |
+
+---
+
+## Usage Example
+
+Start the master with the SSD free-ratio-first strategy:
+
+```bash
+./mooncake_master --allocation_strategy=ssd_free_ratio_first
+```
+
+With this configuration:
+
+1. The master samples up to `6 * replica_num` candidate segments for each allocation request.
+2. Candidates are ranked by SSD free ratio (descending).
+3. Allocation proceeds from the top-ranked candidates.
+4. Any remaining replicas fall back to random allocation.
diff --git a/docs/source/design/ssd-offload.md b/docs/source/design/ssd-offload.md
index a9e42448..d9430f5d 100644
--- a/docs/source/design/ssd-offload.md
+++ b/docs/source/design/ssd-offload.md
@@ -6,6 +6,8 @@ Mooncake Store supports offloading KV cache objects from distributed memory to l
 
 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.
 
+For multi-turn conversation benchmark results, see [Mooncake SSD Offload Benchmark](../performance/ssd-offload-benchmark-results.md).
+
 ---
 
 ## Architecture
@@ -242,4 +244,6 @@ To prevent `io_uring`'s `FOLL_LONGTERM` page pinning from failing on systems wit
 
 ## 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.
+On startup, `FileStorage::Init` calls `StorageBackend::ScanMeta`, which reads 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 for the backends that preserve restart metadata, namely `BucketStorageBackend` and the file-per-key backend.
+
+`OffsetAllocatorStorageBackend` is the exception. It truncates its pre-allocated data file during initialization and clears its in-memory metadata, so previously offloaded objects are not recoverable after a real client restart.
diff --git a/docs/source/design/tent/cpp-api.md b/docs/source/design/tent/cpp-api.md
index 319ae92a..fc68a22c 100644
--- a/docs/source/design/tent/cpp-api.md
+++ b/docs/source/design/tent/cpp-api.md
@@ -189,6 +189,9 @@ struct Request {
     SegmentID target_id;
     uint64_t target_offset;
     size_t length;
+    int priority = PRIO_HIGH;
+    std::optional<std::string> policy_name;
+    TransportType transport_hint = UNSPEC;
 };
 ```
 
@@ -197,6 +200,9 @@ struct Request {
 - `target_id`: Segment ID obtained from `openSegment`.
 - `target_offset`: Offset within the target segment.
 - `length`: Number of bytes to transfer.
+- `priority`: Scheduling priority. Used by the QoS layer; see [qos.md](qos.md).
+- `policy_name`: Optional. When set, the request matches the named entry instead of the first-matching policy.
+- `transport_hint`: Optional. `UNSPEC` (default) defers to `TransportSelector`. Any other `TransportType` pins this request onto that transport for its first try.
 
 #### TransferStatus
 
@@ -522,7 +528,8 @@ Location strings identify device affinity: `"cpu:0"`, `"cuda:0"`, `"cuda:1"`, et
 
 ```cpp
 enum TransportType {
-    RDMA = 0,
+    UNSPEC = 0,
+    RDMA,
     MNNVL,
     SHM,
     NVLINK,
@@ -530,7 +537,7 @@ enum TransportType {
     IOURING,
     TCP,
     AscendDirect,
-    UNSPEC
+    SUNRISE_LINK,
 };
 ```
 
diff --git a/docs/source/design/tent/failover.md b/docs/source/design/tent/failover.md
new file mode 100644
index 00000000..676b986a
--- /dev/null
+++ b/docs/source/design/tent/failover.md
@@ -0,0 +1,197 @@
+# TENT Failover
+
+TENT hides transfer failures from the application by recovering inside the data path.
+This document describes how the recovery works, which knobs control it, and how it is tested.
+
+The design has two layers:
+
+1. **Cross-transport failover** in `TransferEngineImpl`. When a transport fails a task at the completion stage, the engine moves that task to the next available transport (for example RDMA → TCP). Submit-stage failures are not retried today; see Known Gaps.
+2. **Intra-RDMA rail recovery** in `RailMonitor`. When a specific (local NIC, remote NIC) rail keeps failing, the monitor pauses it with exponential cooldown; a successful transfer or the cooldown expiry brings it back.
+
+Application code submits a batch and polls `getTransferStatus`. It never sees a `FAILED` task as long as any healthy path remains and the failover budget is not exhausted.
+
+## Fault Model
+
+TENT focuses on three kinds of transient faults:
+
+| Fault | Surface | Recovery action |
+|-------|---------|-----------------|
+| Work request completion error (WC error) | RDMA worker sees a bad completion | Rail-level `markFailed` + task-level resubmit |
+| QP / endpoint failure | `submitTransferTasks` returns non-OK | *Not retried today*: task surfaces as `FAILED`. See Known Gaps. |
+| Peer disconnect mid-transfer | `getTransferStatus` returns `FAILED` | Cross-transport failover |
+
+Permanent or application-visible errors (invalid arguments, out-of-memory, segment not found) are *not* retried; they are returned to the caller as-is.
+
+## Architecture
+
+```
+                +-------------------------+
+ submitTransfer |  TransferEngineImpl     |
+ ---------------> classify by TransportType
+                |  submitTransferTasks    |----failure----+
+                +-------------------------+               |
+                      |                                   v
+                      |                       resubmitTransferTask
+                      |                       (bump priority, pick next
+                      |                        transport, resubmit)
+                      v
+         +---------------------------+
+         | RdmaTransport / workers   |
+         |   +---------------------+ |
+         |   | RailMonitor         | |
+         |   |  per-rail state     | |
+         |   |  cooldown / recover | |
+         |   +---------------------+ |
+         +---------------------------+
+```
+
+* Each request is owned by one `TaskInfo`. `type` names the transport currently executing the task; `xport_priority` is the index into the ranked fallback list; `failover_count` caps how many times we may re-resolve the transport.
+* The ranked fallback list comes from `getTransportType(req, priority)`. Priority 0 yields the best available transport; increasing priority walks down the list; `UNSPEC` means no transport left.
+* RDMA rail state lives in `RailMonitor`. Its lifecycle is independent of the task-level state machine: a rail can be paused while tasks keep flowing on other rails.
+
+## State Machine
+
+### Cross-transport failover
+
+`resubmitTransferTask` is the single entry point that promotes a failing task to the next transport:
+
+```
+++task.failover_count
+if failover_count > max_failover_attempts  -> return error (exhausted)
+
+task.xport_priority++
+type = resolveTransport(task.request, task.xport_priority)
+if type == UNSPEC                          -> return error (no transport)
+
+transport_list_[type]->submitTransferTasks(...)
+```
+
+It has two callers, one per recoverable failure surface:
+
+1. **Completion-stage failure.** `getTransferStatus(batch_id, task_id, status)` and the batch-form overload call `resubmitTransferTask` once per `FAILED` completion. On success the task is re-marked `PENDING` so the aggregated batch status does not latch to `FAILED` because of a task that is actually retrying.
+
+2. **Exhaustion.** When the budget is hit, `resubmitTransferTask` sets the returned status to `InvalidEntry("Failover limit exceeded, all transports exhausted")`. Callers leave `task.type` unchanged; the task then reports `FAILED` through the normal status flow.
+
+Submit-stage failures (`submitTransferTasks` returning non-OK) are **not** retried today. They mark the task as `UNSPEC`, and `getTransferStatus` short-circuits to `FAILED`. See Known Gaps for why.
+
+### RDMA rail recovery
+
+Inside `RdmaTransport`, each completion drives the rail monitor:
+
+* Bad completion → `rail.markFailed(local_nic, remote_nic)`
+* Good completion → `rail.markRecovered(local_nic, remote_nic)`
+
+`markFailed` bumps `error_count` inside `error_window_`. Once the count hits `error_threshold_` the rail is paused until `now + cooldown_`; the cooldown doubles on every repeat failure up to `kMaxCooldown` (300 s).
+
+`markRecovered` clears the error count, un-pauses the rail, and resets the exponential-backoff memory so the next failure cycle starts from the initial cooldown. A fast path returns without work when the rail is already healthy, which is the common case on the completion hot path.
+
+`available(local, remote)` is the gate every work request passes through before posting. If the cooldown has expired, `available` itself resets all backoff state (error count, resume time, cooldown) and logs `Rail recovered: ... (cooldown expired)`. Otherwise it returns false and the scheduler picks another rail via `findBestRemoteDevice`.
+
+This produces two independent recovery signals — cooldown expiry and live success — so a flaky rail does not stall forever if no other rail is posted to, and a recovered rail returns to service at the first good completion instead of waiting for the full cooldown.
+
+## Configuration
+
+All knobs live in the top-level `transfer-engine.json`. Defaults are safe for production; tune only if you have evidence.
+
+| Key | Default | Meaning |
+|-----|---------|---------|
+| `enable_auto_failover_on_poll` | `true` | Controls whether `getTransferStatus` automatically resubmits tasks that report a recoverable `FAILED` completion. Set to `false` to make status polling observational only; internal completion paths can still trigger failover/resubmit. |
+| `max_failover_attempts` | `3` | Upper bound on `resubmitTransferTask` calls per task. `0` disables cross-transport failover entirely. `1` allows exactly one switch. |
+| `transports/rdma/rail_error_threshold` | `3` | Number of failures inside `rail_error_window_secs` that trips a rail into the paused state. |
+| `transports/rdma/rail_error_window_secs` | `10` | Sliding window for counting rail errors. A failure older than the window resets `error_count` to 1. |
+| `transports/rdma/rail_cooldown_secs` | `30` | Initial cooldown after tripping. Doubles on each repeat failure, capped at 300 s. |
+
+The RDMA keys are read by `RailMonitor::load`. Example:
+
+```json
+{
+  "enable_auto_failover_on_poll": true,
+  "max_failover_attempts": 3,
+  "transports": {
+    "rdma": {
+      "rail_error_threshold": 3,
+      "rail_error_window_secs": 10,
+      "rail_cooldown_secs": 30
+    }
+  }
+}
+```
+
+## Observability
+
+### Metric
+
+`tent_transport_failover_total` is a counter incremented once per successful transport switch inside `resubmitTransferTask`. A non-zero rate means the engine is actively recovering; a sudden jump usually points at a single bad link or flaky peer.
+
+The counter is only built when TENT is compiled with `-DTENT_METRICS_ENABLED=ON` (see `metrics.md`). Without that flag the macro is a no-op.
+
+### Log keywords
+
+| Keyword | Interpretation |
+|---------|----------------|
+| `Transport failover: X -> Y (attempt N/M)` | A task has successfully switched transports. |
+| `Task failover limit reached (M), last transport=X` | Task exhausted its budget and will surface `FAILED`. |
+| `No more transports available after X failed` | `resolveTransport` returned `UNSPEC`; no further fallback exists for this request. |
+| `Rail recovered: local_nic=... remote_nic=... (cooldown expired)` | Cooldown elapsed and the rail is back in service. |
+| `Rail recovered: ... (un-paused by successful transfer)` | Live success on a previously paused rail brought it back early. |
+
+## Testing
+
+Real hardware faults are hard to stage, so TENT tests the failover machinery with decorator-style fault injection.
+
+### FaultProxyTransport
+
+`FaultProxyTransport` wraps any `Transport` and injects four policy-driven faults:
+
+* `submit_fail_rate` — probability that `submitTransferTasks` returns an error.
+* `status_corrupt_rate` — probability that `getTransferStatus` flips `COMPLETED → FAILED`.
+* `fail_after_n_submits` — deterministic variant: succeed the first N submits, then always fail.
+* `fail_install` — make `install()` fail, simulating a transport that cannot come up.
+
+Because it implements the `Transport` interface, the engine sees an ordinary transport. All failover paths (`submitTransfer`, `getTransferStatus`, `resubmitTransferTask`) run unmodified.
+
+### Test-only injection hook
+
+`TransferEngineImpl::swapTransportForTest` replaces the transport in one slot after `construct()`. This is the only way the end-to-end test can wrap the real transport with `FaultProxyTransport` without bypassing `resolveTransport` or `resubmitTransferTask`. Production code never calls it.
+
+### End-to-end suite
+
+The end-to-end failover test suite drives the real `TransferEngineImpl` with fake transports (`FakeTransport`) wrapped in `FaultProxyTransport`. It uses a `p2p` metadata backend on `127.0.0.1` so no external services are required — the whole suite is self-contained.
+
+Current cases:
+
+| Test | What it exercises |
+|------|-------------------|
+| `StatusCorruptionTriggersFailoverToSecondary` | Primary reports `FAILED` in `getTransferStatus`; engine must resubmit on the secondary. |
+| `BothTransportsFailExhaustsFailoverBudget` | Both transports fail at the completion stage; task must surface `FAILED` once the budget is drained. |
+| `MixedFaultsAcrossManySubmissions` | 10 one-request batches with 30% completion corruption on RDMA; every task must end `COMPLETED`, and submit-counter math must hold. |
+| `MaxFailoverAttemptsZeroDisablesFailover` | `max_failover_attempts = 0` → the first completion fault is permanent, TCP is never touched. |
+| `MaxFailoverAttemptsOneAllowsSingleFailover` | `max_failover_attempts = 1` → one switch allowed; RDMA fault → TCP success. |
+| `PerTaskFailoverCountsAreIndependent` | A failing task must not consume another task's budget; `failover_count` is strictly per-task. |
+
+A test-local `PerRequestFaultProxy` (in the same file) subclasses `FaultProxyTransport` to take a `std::function` predicate, remembers which sub-task ids it marked as "poisoned" at submit time, and flips only those completions from `COMPLETED` to `FAILED` at status-query time.
+
+### Running manually
+
+The TENT tests are **not** in CI today (the upstream workflow builds with `USE_TENT=OFF`). Run them locally:
+
+```bash
+cmake -S . -B build-tent -DUSE_TENT=ON -DUSE_CUDA=OFF
+cmake --build build-tent --target tent_failover_test tent_engine_failover_e2e_test -j
+./build-tent/mooncake-transfer-engine/tent/tests/tent_failover_test
+./build-tent/mooncake-transfer-engine/tent/tests/tent_engine_failover_e2e_test
+```
+
+Setting `USE_CUDA=OFF` forces `CpuPlatform`, which always reports `MTYPE_CPU`. With `USE_CUDA=ON` on a host without a GPU, `cudaPointerGetAttributes` fails, `getMemoryType` returns `MTYPE_UNKNOWN`, every transport reports unavailable, and `resolveTransport` returns `UNSPEC` before the fault injection ever runs.
+
+Companion unit tests cover the rail monitor and related building blocks: `tent_rail_monitor_test`, `tent_failover_test`, `tent_fault_proxy_test`.
+
+## Known Gaps
+
+* **Submit-stage failures do not trigger failover.** When `submitTransferTasks` returns non-OK, every task in that call is marked `UNSPEC` and surfaces as `FAILED`. A naive retry loop here is unsafe for two reasons:
+  1. **Merged requests.** When `merge_requests` is enabled (default), `task_id_list[type]` contains both the real merged task and its derived aliases. Resubmitting per task-id re-posts one logical transfer multiple times on the fallback transport, breaking the deduplication the merge pass established.
+  2. **Partial enqueue.** Some transports (for example `ShmTransport::submitTransferTasks`, `NVLinkTransport::submitTransferTasks`) enqueue or start work for earlier requests in `request_list` before returning an error on a later one. The return status alone does not tell us which tasks partially succeeded, so a blanket resubmit would duplicate already-started transfers.
+  A safe submit-stage recovery needs either (a) a transport-level "atomic submit" capability flag plus per-task skip of derived ids, or (b) per-request status returned from `submitTransferTasks`. Neither exists today.
+* `markRecovered` (and cooldown expiry in `available`) clears the exponential-backoff memory entirely. A rail that flaps repeatedly therefore does not accumulate a growing cooldown across recovery cycles. If this becomes a problem the fix is to decay rather than reset.
+* Cross-transport failover is driven purely by return status; there is no latency-based "this transport is healthy but too slow, try another" signal. That belongs to the scheduler, not this document.
+* TENT tests are not exercised by CI. A follow-up can add a CI job that builds with `-DUSE_TENT=ON -DUSE_CUDA=OFF` and runs the `tent_*` test targets; none of the code in this document changes in that case.
diff --git a/docs/source/design/tent/overview.md b/docs/source/design/tent/overview.md
index b56e69b7..32887435 100644
--- a/docs/source/design/tent/overview.md
+++ b/docs/source/design/tent/overview.md
@@ -75,11 +75,7 @@ The design favors predictable behavior and operational simplicity over manual tu
 
 ## TENT C++ API Reference 
 
-:::{toctree}
-:maxdepth: 1
-
-cpp-api
-:::
+[TENT C++ API Reference](cpp-api.md)
 
 ## TENT Metrics System
 
@@ -87,4 +83,36 @@ cpp-api
 :maxdepth: 1
 
 metrics
-:::
\ No newline at end of file
+:::
+
+## TENT Transport Selection
+
+:::{toctree}
+:maxdepth: 1
+
+transport-selector
+:::
+
+## TENT Quality of Service
+
+:::{toctree}
+:maxdepth: 1
+
+qos
+:::
+
+## TENT Slice Spraying
+
+:::{toctree}
+:maxdepth: 1
+
+slice-spraying
+:::
+
+## TENT Failover
+
+:::{toctree}
+:maxdepth: 1
+
+failover
+:::
diff --git a/docs/source/design/tent/qos.md b/docs/source/design/tent/qos.md
new file mode 100644
index 00000000..06ecfc27
--- /dev/null
+++ b/docs/source/design/tent/qos.md
@@ -0,0 +1,396 @@
+# TENT Quality of Service (QoS)
+
+## Overview
+
+TENT provides Quality of Service (QoS) support to ensure that high-priority requests receive preferential treatment in multi-tenant and multi-workload environments. This document describes TENT's QoS architecture and configuration.
+
+## Background
+
+In shared RDMA clusters, different types of transfers have different priority requirements:
+
+1. **Metadata and Control Messages**: Require low latency, small size
+2. **Interactive Queries**: Require low to medium latency, medium size
+3. **Bulk Data Transfer**: Can tolerate higher latency, large size
+
+Without QoS, low-priority bulk transfers can monopolize bandwidth and cause high tail latency for critical requests.
+
+TENT addresses this through:
+- **Per-worker priority queues** for intra-process isolation
+- **Global time-sliced coordination** for inter-process isolation
+- **Priority-aware device filtering** for NUMA-aware scheduling
+
+## Architecture
+
+### Priority Levels
+
+TENT supports three priority levels:
+
+| Priority | Value | Description | Use Cases |
+|----------|-------|-------------|-----------|
+| `PRIO_HIGH` | 0 | High-priority requests | Metadata, control messages, latency-sensitive operations |
+| `PRIO_MEDIUM` | 1 | Medium-priority requests | Interactive queries, serving workloads |
+| `PRIO_LOW` | 2 | Low-priority requests | Bulk data transfer, background jobs |
+
+### Per-Worker Priority Queues
+
+Each worker thread maintains separate queues for each priority level:
+
+```
+┌─────────────────────────────────────┐
+│         Worker Thread               │
+├─────────────────────────────────────┤
+│  PRIO_HIGH Queue     │              │
+│  PRIO_MEDIUM Queue   │              │
+│  PRIO_LOW Queue      │              │
+├─────────────────────────────────────┤
+│  Dequeue Priority: HIGH→MEDIUM→LOW   │
+└─────────────────────────────────────┘
+```
+
+**Scheduling Logic**:
+1. Always drain HIGH priority queue first
+2. Only process MEDIUM when HIGH is empty
+3. Only process LOW when both HIGH and MEDIUM are empty
+
+**Priority Promotion (Anti-Starvation)**:
+To prevent low-priority requests from starving indefinitely, TENT implements timeout-based priority promotion:
+- MEDIUM priority requests are promoted to HIGH after waiting too long
+- LOW priority requests are promoted to MEDIUM after waiting too long
+- Promotion checks run periodically (every 1ms by default)
+
+This ensures that:
+- High-priority requests normally never wait behind lower-priority work
+- Low-priority requests eventually get serviced even under continuous high-priority load
+
+### Global Slot Coordination
+
+For multi-process environments, TENT implements global time-sliced coordination using shared memory:
+
+```
+Time slices rotate every N milliseconds:
+
+Slot 0 (0-Nms):     Only HIGH priority requests allowed
+Slot 1 (N-2Nms):    MEDIUM + HIGH priority requests allowed  
+Slot 2 (2N-3Nms):   All priorities allowed
+...repeats...
+```
+
+**Default Configuration**: 2ms per slot (6ms full cycle)
+
+This mechanism ensures that:
+- High-priority requests get dedicated service windows
+- No process can monopolize bandwidth indefinitely
+- Fair access across process boundaries
+
+### Shared Memory Structure
+
+The global slot state is maintained in shared memory:
+
+```cpp
+struct SharedHeader {
+    uint64_t magic;              // Magic number for validation
+    int32_t version;             // Format version
+    std::atomic<int> current_slot;  // Current global slot (0, 1, or 2)
+    pthread_mutex_t global_mutex;    // For synchronization (robust)
+};
+```
+
+**Operations**:
+- Background thread rotates slot every N milliseconds
+- Workers check `canSend()` before processing requests
+- Only requests with priority ≤ slot level are processed
+
+## Configuration
+
+### Priority Filtering
+
+```json
+{
+  "transports": {
+    "rdma": {
+      "enable_priority_filtering": true,
+      "local_rotation_interval_us": 200,
+      "priority_promotion_timeout_us": 10000
+    }
+  }
+}
+```
+
+| Parameter | Type | Default | Description |
+|-----------|------|---------|-------------|
+| `enable_priority_filtering` | bool | `true` | Enable priority-based device filtering |
+| `local_rotation_interval_us` | int | `200` | Local device priority rotation interval (microseconds) |
+| `priority_promotion_timeout_us` | int | `10000` | Timeout for priority promotion (microseconds) |
+
+### Global Coordination
+
+```json
+{
+  "transports": {
+    "rdma": {
+      "slot_rotation_interval_ms": 2,
+      "shared_quota_shm_path": "/mooncake_rdma_slots"
+    }
+  }
+}
+```
+
+| Parameter | Type | Default | Description |
+|-----------|------|---------|-------------|
+| `slot_rotation_interval_ms` | int | `2` | Global slot rotation interval (milliseconds) |
+| `shared_quota_shm_path` | string | `""` | Shared memory path for multi-process coordination |
+
+**Note**: Leave `shared_quota_shm_path` empty to disable global coordination (single-process mode).
+
+## Usage Examples
+
+### Example 1: Latency-Critical Workload
+
+For workloads where high-priority requests must have minimal latency:
+
+```json
+{
+  "transports": {
+    "rdma": {
+      "enable_priority_filtering": true,
+      "slot_rotation_interval_ms": 1,
+      "shared_quota_shm_path": "/mooncake_rdma_slots"
+    }
+  }
+}
+```
+
+**Effect**: High-priority requests get dedicated windows every 1ms.
+
+### Example 2: Single-Process Mode
+
+For single-process deployments where global coordination is not needed:
+
+```json
+{
+  "transports": {
+    "rdma": {
+      "enable_priority_filtering": true,
+      "shared_quota_shm_path": ""
+    }
+  }
+}
+```
+
+**Effect**: Per-worker priority queues only, no cross-process coordination.
+
+### Example 3: Bulk-Friendly Configuration
+
+For workloads where low-priority bulk transfers should not be starved:
+
+```json
+{
+  "transports": {
+    "rdma": {
+      "slot_rotation_interval_ms": 10,
+      "shared_quota_shm_path": "/mooncake_rdma_slots"
+    }
+  }
+}
+```
+
+**Effect**: Longer slots allow more low-priority work to complete.
+
+## Priority Assignment
+
+### Setting Request Priority
+
+Priority is assigned when creating `Request` objects. The default priority is `PRIO_HIGH`.
+
+### C++ API
+
+```cpp
+// In C++
+#include "tent/transfer_engine.h"
+
+using namespace mooncake::tent;
+
+// Create request with default priority (HIGH)
+Request req;
+req.opcode = Request::OpCode::READ;
+req.source = buffer;
+req.target_id = segment_id;
+req.target_offset = 0;
+req.length = size;
+// req.priority is PRIO_HIGH by default
+
+// Or specify priority explicitly
+req.priority = PRIO_MEDIUM;  // or PRIO_HIGH, PRIO_LOW
+
+// Submit the request
+engine.submitTransfer(batch_id, {req});
+```
+
+### Python API
+
+```python
+# In Python
+import tent
+
+# Create request with default priority (HIGH)
+req = tent.Request(
+    opcode=tent.OpCode.READ,
+    source=buffer_addr,
+    target_id=segment_id,
+    target_offset=0,
+    length=size
+)
+# req.priority is tent.PRIO_HIGH by default
+
+# Or specify priority in constructor
+req = tent.Request(
+    opcode=tent.OpCode.READ,
+    source=buffer_addr,
+    target_id=segment_id,
+    target_offset=0,
+    length=size,
+    priority=tent.PRIO_LOW  # or PRIO_HIGH, PRIO_MEDIUM
+)
+
+# Or set after creation
+req.priority = tent.PRIO_MEDIUM
+
+# Submit the request
+engine.submit_transfer(batch_id, [req])
+```
+
+### C API
+
+```c
+// In C
+#include "tent/transfer_engine.h"
+
+// Create request with priority
+tent_request_t req = {
+    .opcode = OPCODE_READ,
+    .source = buffer,
+    .target_id = segment_id,
+    .target_offset = 0,
+    .length = size,
+    .priority = 0  // 0=HIGH, 1=MEDIUM, 2=LOW
+};
+
+// Submit the request
+tent_submit(engine, batch_id, &req, 1);
+```
+
+## Performance Considerations
+
+### Trade-offs
+
+| Configuration | High-Priority Latency | Low-Priority Throughput | Fairness |
+|---------------|----------------------|------------------------|----------|
+| Short slot interval (1ms) | Excellent | Poor | High |
+| Default slot interval (2ms) | Good | Fair | High |
+| Long slot interval (10ms) | Fair | Good | Medium |
+| No global coordination | Variable | Excellent | Low (per-process only) |
+
+### Starvation Prevention
+
+TENT prevents starvation through two mechanisms:
+
+1. **Global slot mechanism**:
+   - **HIGH priority**: Never starved (always allowed in slot 0)
+   - **MEDIUM priority**: Never starved (allowed in slots 1 and 2)
+   - **LOW priority**: Never starved (always allowed in slot 2)
+
+2. **Priority promotion timeout**:
+   - Low-priority requests waiting longer than `priority_promotion_timeout_us` are promoted
+   - MEDIUM → HIGH promotion ensures medium priority gets service
+   - LOW → MEDIUM promotion ensures low priority eventually gets service
+   - Configurable via `priority_promotion_timeout_us` (default 10ms)
+
+### Tuning Guidelines
+
+1. **Start with default settings** (2ms slot interval)
+2. **Measure tail latency** for each priority level
+3. **Adjust slot interval** based on observations:
+   - If HIGH priority latency is too high: decrease interval
+   - If LOW priority throughput is too low: increase interval
+
+## Troubleshooting
+
+### Problem: High-priority requests have high latency
+
+**Symptoms**: `PRIO_HIGH` requests experiencing unexpected delays
+
+**Possible causes**:
+1. Slot interval too long
+2. Global coordination not enabled
+3. Worker threads blocked on LOW priority work
+
+**Solution**:
+```json
+{
+  "slot_rotation_interval_ms": 1,
+  "enable_priority_filtering": true
+}
+```
+
+### Problem: Low-priority transfers starved
+
+**Symptoms**: `PRIO_LOW` requests making no progress
+
+**Possible causes**:
+1. HIGH priority load is continuous
+2. Slot interval too short
+
+**Solution**: Increase slot interval to give LOW priority more time:
+```json
+{
+  "slot_rotation_interval_ms": 10
+}
+```
+
+### Problem: Shared memory creation fails
+
+**Symptoms**: Error messages about `/mooncake_rdma_slots`
+
+**Possible causes**:
+1. Permission issues (need write access to `/dev/shm`)
+2. Stale shared memory from previous run
+
+**Solution**:
+```bash
+# Remove stale shared memory
+rm -f /dev/shm/mooncake_rdma_slots
+
+# Or use a different path
+{
+  "shared_quota_shm_path": "/mooncake_rdma_slots_v2"
+}
+```
+
+## Monitoring
+
+### Traffic Statistics
+
+Monitor per-device traffic distribution:
+
+```cpp
+device_selector_->printTrafficStats();
+```
+
+Output example:
+```
+=== Device Traffic Statistics ===
+Dev 0: Total=10.5 GB, EWMA BW=45.23 Gbps, Inflight=0 bytes
+Dev 1: Total=8.2 GB, EWMA BW=42.18 Gbps, Inflight=0 bytes
+Dev 2: Total=0.5 GB, EWMA BW=38.91 Gbps, Inflight=0 bytes
+Dev 3: Total=0.3 GB, EWMA BW=39.12 Gbps, Inflight=0 bytes
+```
+
+### Priority Statistics
+
+Monitor queue depths for each priority level (requires instrumentation).
+
+## References
+
+- [TENT Overview](overview.md)
+- [TENT Slice Spraying](slice-spraying.md)
+- [TENT C++ API](cpp-api.md)
diff --git a/docs/source/design/tent/slice-spraying.md b/docs/source/design/tent/slice-spraying.md
new file mode 100644
index 00000000..5cf6d3bc
--- /dev/null
+++ b/docs/source/design/tent/slice-spraying.md
@@ -0,0 +1,398 @@
+# TENT Slice Spraying
+
+## Overview
+
+This document describes TENT's Slice Spraying mechanism, which enables efficient data movement in multi-rail RDMA environments through intelligent device selection and adaptive load balancing.
+
+## Background
+
+In multi-rail RDMA environments, naive round-robin striping leads to suboptimal performance because:
+
+1. **NUMA Effects**: Cross-NUMA access incurs additional latency and reduces effective bandwidth
+2. **Load Imbalance**: Static striping cannot adapt to dynamic load conditions
+3. **Heterogeneous Link Quality**: Different rails may have different effective bandwidth due to congestion or hardware characteristics
+
+TENT addresses these issues through:
+- **NUMA-aware device selection** with configurable penalties
+- **EWMA-based bandwidth estimation** for adaptive load balancing
+- **Dynamic multi-path allocation** for large transfers
+
+## Architecture
+
+### Device Selector
+
+The `DeviceSelector` component is responsible for choosing which RDMA device(s) to use for each transfer request. It operates in two modes:
+
+#### Baseline Mode (Round-Robin)
+
+When `enable_smart_scheduling = false`, the selector uses simple round-robin within the highest-priority device tier (typically local NUMA devices):
+
+```
+For each request:
+  1. Find first non-empty device tier (local NUMA preferred)
+  2. Select devices round-robin within that tier
+  3. Ignore lower-priority tiers
+```
+
+**Characteristics**:
+- Deterministic behavior
+- No runtime overhead for tracking
+- Consistent with original TE behavior
+- Does not adapt to load conditions
+
+#### Smart Mode (EWMA-Based Selection)
+
+When `enable_smart_scheduling = true`, the selector uses an EWMA-based algorithm:
+
+```
+For each request:
+  1. Calculate predicted completion time for each device:
+     predicted_time = (inflight_bytes + slice_bytes) / ewma_bandwidth
+
+  2. Apply NUMA penalty based on tier:
+     score = predicted_time × numa_tier_weights[tier]
+
+  3. Select device(s) with minimum score:
+     - Single slice: best device only
+     - Multiple slices: weighted distribution across devices
+
+  4. Update EWMA bandwidth on completion:
+     ewma_bandwidth = α × ewma_bandwidth + (1 - α) × observed_bandwidth
+     where α = bandwidth_learning_rate
+```
+
+**Characteristics**:
+- Adapts to changing load conditions
+- Prefers local NUMA devices
+- Spreads load across multiple rails
+- Higher runtime overhead
+
+### NUMA-Aware Selection
+
+Devices are organized into tiers based on NUMA distance:
+
+| Tier | Description | Default Penalty |
+|------|-------------|-----------------|
+| Rank 0 | Local NUMA | 1.0 (baseline) |
+| Rank 1 | Remote NUMA (tier 1) | 5.0 |
+| Rank 2 | Remote NUMA (tier 2) | 10.0 |
+
+The penalty is applied as a multiplier to predicted completion time, making remote devices less attractive unless local devices are heavily loaded.
+
+### EWMA Bandwidth Estimation
+
+Each device maintains an EWMA (Exponentially Weighted Moving Average) of its effective bandwidth:
+
+```
+initial_value = theoretical_bandwidth
+
+on_transfer_complete:
+  observed_bandwidth = transfer_size / transfer_time
+  ewma_bandwidth = α × ewma_bandwidth + (1 - α) × observed_bandwidth
+  ewma_bandwidth = clamp(ewma_bandwidth,
+                        0.1 × theoretical,
+                        10.0 × theoretical)
+```
+
+where `α = bandwidth_learning_rate`.
+
+**Note on terminology**: The EWMA formula uses α as the coefficient for the old value. Therefore:
+- **Lower α** (closer to 0) → more weight on new observations → **faster adaptation**
+- **Higher α** (closer to 1) → more weight on old value → **slower adaptation**
+
+Examples:
+- α = 0: `ewma_bandwidth = observed_bandwidth` (full adaptation, always use new value)
+- α = 1: `ewma_bandwidth = ewma_bandwidth` (no learning, never update)
+- α = 0.01: `ewma_bandwidth = 0.01 × old + 0.99 × new` (default, gradual adaptation)
+
+The EWMA provides:
+- **Memory**: Recent observations have more influence than old ones
+- **Stability**: Smooths out transient fluctuations
+- **Adaptability**: Tracks gradual changes in link quality
+
+### Multi-Path Allocation
+
+For large transfers, TENT distributes slices across multiple devices:
+
+**Single Path** (small requests):
+- All slices go to the single best device
+- Minimizes coordination overhead
+
+**Multi Path** (large requests):
+- **Normal mode** (99% of calls): Slices distributed proportionally to device capacity
+  - Each device gets: `(device_weight / total_weight) × num_slices`
+  - Remaining slices assigned to best device
+- **Probe mode** (1% of calls, every 100th call): Slices distributed round-robin
+  - Purpose: Ensure all devices are continuously sampled for EWMA updates
+  - Prevents EWMA starvation for less-used devices
+
+### Request Flow
+
+```
+┌──────────────┐
+│ Application  │
+└──────┬───────┘
+       │ submitTransfer()
+       ▼
+┌──────────────────────────────────────┐
+│  RdmaTransport::submitTransferTasks  │
+│  - Split large requests into slices   │
+│  - Call DeviceSelector for allocation │
+│  - Only if num_slices >= max_slice_count/2 │
+└──────┬───────────────────────────────┘
+       │
+       ▼
+┌──────────────────────────────────────┐
+│     DeviceSelector::allocate         │
+│  ┌────────────────────────────────┐  │
+│  │ smart_selection_enabled?       │  │
+│  └────┬──────────────────────┬────┘  │
+│       │ Yes                  │ No     │
+│       ▼                     ▼        │
+│  ┌─────────┐          ┌─────────┐   │
+│  │  Smart  │          │ Baseline│   │
+│  │  Mode   │          │   Mode  │   │
+│  └────┬────┘          └────┬────┘   │
+│       │                    │         │
+│       └────────┬───────────┘         │
+│                ▼                     │
+│  ┌────────────────────────────────┐  │
+│  │  Return slice_dev_ids          │  │
+│  └────────────────────────────────┘  │
+└──────────────────────────────────────┘
+```
+
+## Configuration
+
+All slice spraying parameters are configurable via the configuration file:
+
+### Core Scheduling
+
+```json
+{
+  "transports": {
+    "rdma": {
+      "enable_smart_scheduling": true
+    }
+  }
+}
+```
+
+| Parameter | Type | Default | Description |
+|-----------|------|---------|-------------|
+| `enable_smart_scheduling` | bool | `true` | Enable EWMA-based selection (false = round-robin) |
+
+### NUMA Penalties
+
+```json
+{
+  "transports": {
+    "rdma": {
+      "numa_penalties": [1.0, 5.0, 10.0]
+    }
+  }
+}
+```
+
+| Parameter | Type | Default | Description |
+|-----------|------|---------|-------------|
+| `numa_penalties` | array[float] | `[1.0, 5.0, 10.0]` | Penalty multipliers for each NUMA tier |
+
+**Guidelines**:
+- Higher values = stronger preference for local devices
+- Set all to `1.0` to disable NUMA awareness
+- Increase remote penalties if cross-NUMA latency is high
+
+### Bandwidth Estimation
+
+```json
+{
+  "transports": {
+    "rdma": {
+      "bandwidth_learning_rate": 0.01,
+      "ewma_min_bandwidth_multiplier": 0.1,
+      "ewma_max_bandwidth_multiplier": 10.0
+    }
+  }
+}
+```
+
+| Parameter | Type | Default | Description |
+|-----------|------|---------|-------------|
+| `bandwidth_learning_rate` | float | `0.01` | EWMA learning rate (0.0 = full adaptation, 1.0 = no learning) |
+| `ewma_min_bandwidth_multiplier` | float | `0.1` | Minimum bandwidth as fraction of theoretical |
+| `ewma_max_bandwidth_multiplier` | float | `10.0` | Maximum bandwidth as fraction of theoretical |
+
+**Guidelines**:
+- Lower α (e.g., 0.001) → faster adaptation, more volatile → responds quickly to changes
+- Higher α (e.g., 0.1) → slower adaptation, more stable → smooths out transient fluctuations
+- Default α = 0.01 provides balanced adaptation
+- Multipliers constrain EWMA to reasonable range [0.1×, 10.0×] of theoretical bandwidth
+
+### Device Selection Scoring
+
+```json
+{
+  "transports": {
+    "rdma": {
+      "score_jitter_range": 1e-9,
+      "score_epsilon": 1e-12
+    }
+  }
+}
+```
+
+| Parameter | Type | Default | Description |
+|-----------|------|---------|-------------|
+| `score_jitter_range` | float | `1e-9` | Random jitter range to avoid deterministic selection |
+| `score_epsilon` | float | `1e-12` | Small value to prevent division by zero |
+
+### Bandwidth Constants
+
+```json
+{
+  "transports": {
+    "rdma": {
+      "default_bandwidth_gbps": 400.0,
+      "min_bandwidth_gbps": 10.0,
+      "max_bandwidth_gbps": 800.0
+    }
+  }
+}
+```
+
+| Parameter | Type | Default | Description |
+|-----------|------|---------|-------------|
+| `default_bandwidth_gbps` | float | `400.0` | Default NIC bandwidth when topology info unavailable |
+| `min_bandwidth_gbps` | float | `10.0` | Minimum valid NIC bandwidth (Gbps) |
+| `max_bandwidth_gbps` | float | `800.0` | Maximum valid NIC bandwidth (Gbps) |
+
+**Notes**:
+- These constants define the valid range and default for device bandwidth
+- Used in EWMA calculations and theoretical bandwidth estimation
+- If a device's reported bandwidth is outside [min, max], default_bandwidth is used
+
+## Usage Examples
+
+### Example 1: Latency-Sensitive Workload
+
+For latency-sensitive queries where local NUMA access is critical:
+
+```json
+{
+  "transports": {
+    "rdma": {
+      "enable_smart_scheduling": true,
+      "numa_penalties": [1.0, 100.0, 1000.0],
+      "bandwidth_learning_rate": 0.001
+    }
+  }
+}
+```
+
+**Effect**: Strongly prefers local devices, slow adaptation for stability.
+
+### Example 2: Bulk Data Transfer
+
+For bulk transfers where throughput is more important than latency:
+
+```json
+{
+  "transports": {
+    "rdma": {
+      "enable_smart_scheduling": true,
+      "numa_penalties": [1.0, 2.0, 3.0],
+      "bandwidth_learning_rate": 0.1
+    }
+  }
+}
+```
+
+**Effect**: Allows cross-NUMA transfers, fast adaptation to load.
+
+### Example 3: Baseline Mode
+
+For deterministic performance matching original TE:
+
+```json
+{
+  "transports": {
+    "rdma": {
+      "enable_smart_scheduling": false
+    }
+  }
+}
+```
+
+**Effect**: Round-robin within local NUMA tier, no adaptation, minimal overhead.
+
+## Performance Considerations
+
+### Overhead Comparison
+
+| Mode | CPU Overhead | Adaptability | NUMA Awareness |
+|------|--------------|--------------|----------------|
+| Baseline | Minimal | None | Tier-based (static) |
+| Smart | Moderate | EWMA-based | Dynamic + penalty |
+
+### When to Use Each Mode
+
+**Use Baseline Mode when**:
+- Workload is uniform and predictable
+- Deterministic performance is required
+- CPU overhead must be minimized
+- All devices are in same NUMA node
+
+**Use Smart Mode when**:
+- Workload is heterogeneous
+- Link quality varies over time
+- NUMA effects are significant
+- Maximum throughput is desired
+
+### Tuning Guidelines
+
+1. **Start with baseline mode** to establish performance baseline
+2. **Enable smart mode** with conservative parameters:
+   - `numa_penalties = [1.0, 2.0, 5.0]`
+   - `bandwidth_learning_rate = 0.01`
+3. **Monitor performance** and adjust based on observations:
+   - If cross-NUMA transfers are too frequent: increase remote penalties
+   - If adaptation is too slow (EWMA not keeping up with load changes): decrease α
+   - If performance is unstable (too much fluctuation): increase α
+
+## Troubleshooting
+
+### Problem: All requests go to cross-NUMA devices
+
+**Symptoms**: Poor performance, high latency
+
+**Diagnosis**:
+```cpp
+device_selector_->printTrafficStats();
+```
+
+**Solution**: Check `numa_penalties` configuration. Ensure local devices have lowest penalty (1.0).
+
+### Problem: Performance worse than baseline
+
+**Symptoms**: Smart mode slower than baseline mode
+
+**Possible causes**:
+1. Learning rate too high (volatile decisions)
+2. NUMA penalties too low (not preferring local)
+3. Score jitter too large (too much randomness)
+
+**Solution**: Use more conservative:
+```json
+{
+  "bandwidth_learning_rate": 0.001,
+  "numa_penalties": [1.0, 10.0, 100.0],
+  "score_jitter_range": 1e-12
+}
+```
+
+## References
+
+- [TENT Overview](overview.md)
+- [TENT QoS](qos.md)
+- [TENT C++ API](cpp-api.md)
diff --git a/docs/source/design/tent/transport-selector.md b/docs/source/design/tent/transport-selector.md
new file mode 100644
index 00000000..d1a90e08
--- /dev/null
+++ b/docs/source/design/tent/transport-selector.md
@@ -0,0 +1,246 @@
+# TENT Transport Selector
+
+## Overview
+
+The Transport Selector is responsible for choosing the optimal transport and devices for each transfer request based on configuration policies.
+
+## Request Priority
+
+TENT uses a unified priority system across all components. See [QoS.md](qos.md) for detailed description of priority levels and their usage throughout the system.
+
+Quick reference:
+- `"high"` / `0` - High-priority requests (metadata, control, latency-sensitive)
+- `"medium"` / `1` - Medium-priority requests (interactive queries, serving)
+- `"low"` / `2` - Low-priority requests (bulk transfer, background jobs)
+
+## Configuration-Based Transport Selection
+
+Transport selection is driven by configuration with pattern-based rules.
+
+### Configuration Example
+
+```json
+{
+  "policy": [
+    {
+      "name": "high_prio_fast",
+      "segment_type": "memory",
+      "priority": "high",
+      "devices": ["mlx5_0", "mlx5_1", "mlx5_2"],
+      "transports": ["nvlink", "rdma", "shm"]
+    },
+    {
+      "name": "low_prio_slow",
+      "segment_type": "memory",
+      "priority": "low",
+      "devices": ["mlx5_0"],
+      "transports": ["rdma", "tcp"]
+    },
+    {
+      "name": "file_storage",
+      "segment_type": "file",
+      "transports": ["gds", "io_uring", "rdma"]
+    }
+  ]
+}
+```
+
+### Policy Fields
+
+| Field | Type | Required | Description |
+|-------|------|----------|-------------|
+| `name` | string | Yes | Policy identifier (for logging) |
+| `segment_type` | string | Yes | `"memory"` or `"file"` |
+| `priority` | string or int | No | Match only requests with this priority: `"high"` (0), `"medium"` (1), `"low"` (2) |
+| `devices` | array[string] | No | List of allowed device names (empty = all devices) |
+| `transports` | array[string] | No | Transport preference list (evaluated in order) |
+
+### Memory Type Filters
+
+For `memory` segments, you can filter by source/destination memory type:
+
+| Pattern | Matches |
+|---------|---------|
+| `"cuda"` | CUDA GPU memory |
+| `"cpu"` | CPU/host memory |
+| `"hip"` | ROCm/HIP GPU memory |
+| `"npu"` | Ascend NPU memory |
+| `"*"` | Any memory type |
+
+### Size Filters
+
+Restrict policies to specific transfer sizes:
+
+```json
+{
+  "name": "small_transfers",
+  "segment_type": "memory",
+  "min_size": 0,
+  "max_size": 1048576,
+  "transports": ["shm", "nvlink"]
+}
+```
+
+## Device Mask
+
+The `devices` field in a policy creates a bitmask that restricts which NICs can be used:
+
+```json
+{
+  "name": "use_nic_0_only",
+  "segment_type": "memory",
+  "devices": ["mlx5_0"],
+  "transports": ["rdma"]
+}
+```
+
+This is translated internally to a 64-bit bitmask where each bit represents one device:
+- `devices: ["mlx5_0"]` → `device_mask = 0x0001` (bit 0 set)
+- `devices: ["mlx5_1", "mlx5_2"]` → `device_mask = 0x0006` (bits 1 and 2 set)
+- `devices: []` (empty) → `device_mask = ~0ULL` (all devices)
+
+## Transport Fallback
+
+When multiple transports are listed in the `transports` array, they act as fallback options:
+
+```json
+{
+  "transports": ["rdma", "tcp"]
+}
+```
+
+This means:
+1. Try RDMA first
+2. If RDMA is unavailable or fails, fall back to TCP
+
+### Fallback with `transport_index`
+
+For programmatic control, the `transport_index` parameter selects which transport to use:
+
+```cpp
+// transport_index = 0 → First transport (rdma)
+// transport_index = 1 → Second transport (tcp)
+// transport_index = 2 → Third transport (if exists)
+auto result = selector.select(context, transports, transport_index);
+```
+
+## Per-request override (`transport_hint`)
+
+`Request::transport_hint` lets a caller bypass the policy lookup for one request at a time. It's the per-request analogue of `transport_index` and sits on top of the configured policies, so callers can keep the global policy config as the default while still pinning specific requests.
+
+```cpp
+Request r{};
+r.opcode = Request::WRITE;
+r.source = local_ptr;
+r.target_id = seg;
+r.target_offset = 0;
+r.length = 4096;
+r.transport_hint = TransportType::TCP;   // UNSPEC (default) defers to policy
+engine.submitTransfer(batch_id, {r});
+```
+
+### Semantics
+
+* **`UNSPEC` (default)**: unchanged — `TransportSelector::select(...)` runs as documented above.
+* **Pinned**: the request's first attempt uses the hinted transport. On failover (when enabled), the engine asks the selector for the next candidate **with the hint excluded**, so failover never loops back onto the transport that already failed.
+
+## Default Behavior
+
+If no `policy` is configured, TENT falls back to original behavior:
+
+| Segment Type | Default Transport Order |
+|--------------|-------------------------|
+| File | GDS → IOURING → RDMA |
+| Memory | Uses `buffer_transports` order from buffer registration |
+
+## Complete Example
+
+```json
+{
+  "policy": [
+    {
+      "name": "high_priority_local",
+      "segment_type": "memory",
+      "same_machine": true,
+      "priority": "high",
+      "transports": ["nvlink", "shm"]
+    },
+    {
+      "name": "high_priority_remote",
+      "segment_type": "memory",
+      "same_machine": false,
+      "priority": "high",
+      "local_memory": "cuda",
+      "remote_memory": "cuda",
+      "devices": ["mlx5_0", "mlx5_1", "mlx5_2"],
+      "transports": ["rdma"]
+    },
+    {
+      "name": "medium_priority",
+      "segment_type": "memory",
+      "priority": "medium",
+      "transports": ["rdma"]
+    },
+    {
+      "name": "bulk_transfer",
+      "segment_type": "memory",
+      "priority": "low",
+      "min_size": 104857600,
+      "transports": ["rdma", "tcp"]
+    },
+    {
+      "name": "file_ops",
+      "segment_type": "file",
+      "transports": ["gds", "io_uring"]
+    }
+  ]
+}
+```
+
+## Unified Priority System
+
+The priority value propagates through multiple layers:
+
+```
+Request.priority (application layer)
+    ↓
+TransportSelector policy matching
+    ↓
+DeviceSelector allocation (QoS filtering)
+    ↓
+Worker thread scheduling queues
+```
+
+**Layer interactions**:
+
+1. **TransportSelector**: Policy's `priority` field filters which requests match
+2. **DeviceSelector**: `priority` parameter controls device eligibility (QoS mode)
+3. **Workers**: Separate queues per priority level with strict draining order
+
+See [QoS.md](qos.md) for details on worker scheduling and global slot coordination.
+
+## Data Flow
+
+```
+Request with priority
+    ↓
+TransportSelector.select(context, transports, transport_index)
+    ↓
+Match policy by:
+  - segment_type (file/memory)
+  - priority (exact match if specified in policy)
+  - location constraints
+  - size constraints
+    ↓
+Build device_mask from policy.devices
+    ↓
+Select transport from policy.transports[transport_index]
+    ↓
+Return SelectionResult { transport, device_mask }
+    ↓
+RdmaTransport.submitTransferTasks(batch, requests)
+    ↓
+DeviceSelector.allocate(..., request.priority, batch.device_mask)
+    ↓
+Worker scheduling (separate queues per priority)
+```
diff --git a/docs/source/design/transfer-engine/ascend_direct_transport.md b/docs/source/design/transfer-engine/ascend_direct_transport.md
index 4f8ee63b..217133e2 100644
--- a/docs/source/design/transfer-engine/ascend_direct_transport.md
+++ b/docs/source/design/transfer-engine/ascend_direct_transport.md
@@ -48,7 +48,7 @@ The following environment variables can be configured to control Ascend Direct T
 
 | Variable | Description | Default Value | Example |
 |----------|-------------|---------------|---------|
-| `ASCEND_AUTO_CONNECT` | Enable automatic connection management | 0 (disabled) | `ASCEND_AUTO_CONNECT=1` |
+| `ASCEND_AUTO_CONNECT` | Enable automatic connection management | 1 (enabled) | `ASCEND_AUTO_CONNECT=0` |
 | `ASCEND_ENABLE_USE_FABRIC_MEM` | Enable fabric memory transfer mode in Mooncake Store (A3 only) | 0 (disabled) | `ASCEND_ENABLE_USE_FABRIC_MEM=1` |
 | `ASCEND_USE_ASYNC_TRANSFER` | Enable asynchronous transfer mode | 0 (disabled) | `ASCEND_USE_ASYNC_TRANSFER=1` |
 | `ASCEND_GLOBAL_RESOURCE_CONFIG` | Global resource configuration | - | `ASCEND_GLOBAL_RESOURCE_CONFIG="{\"fabric_memory.max_capacity\":32}"` |
@@ -64,7 +64,7 @@ The following environment variables can be configured to control Ascend Direct T
 
 **Detailed Descriptions:**
 
-- **ASCEND_AUTO_CONNECT**: Requires CANN 9.0 or later. Default is 0, recommended to enable on supported versions: link can be automatically disconnected when the remote end goes offline abnormally.
+- **ASCEND_AUTO_CONNECT**: Requires CANN 9.0 or later. Default is 1: link can be automatically disconnected when the remote end goes offline abnormally. Set `ASCEND_AUTO_CONNECT=0` to disable.
 - **ASCEND_ENABLE_USE_FABRIC_MEM**: Requires CANN 9.0+ and HDK 26.0+. Recommended when using Mooncake Store on supported A3 platforms: it can significantly improve transmission performance.
 - **ASCEND_USE_ASYNC_TRANSFER**: Requires CANN 8.5+. Enables HIXL asynchronous transfer mode, defaults to synchronous mode.
 - **ASCEND_GLOBAL_RESOURCE_CONFIG**: Configures HIXL global resources. Refer to HIXL documentation for `OPTION_GLOBAL_RESOURCE_CONFIG` settings.
@@ -100,4 +100,4 @@ The following environment variables can be configured to control Ascend Direct T
 
 12. **Fabric Memory mode**: On the A3, with the latest drivers and CANN installed, when using Mooncake store, the ASCEND_ENABLE_USE_FABRIC_MEM environment variable can be set to enable fabric memory transfer mode (which allows direct access remote HOST memory).
 
-13. **Auto Connect**: The auto connect feature can be enabled by configuring the `ASCEND_AUTO_CONNECT` environment variable. The default value is 0 (disabled).
+13. **Auto Connect**: The auto connect feature is enabled by default (`ASCEND_AUTO_CONNECT=1`). Set `ASCEND_AUTO_CONNECT=0` to disable.
diff --git a/docs/source/design/transfer-engine/cpp-api.md b/docs/source/design/transfer-engine/cpp-api.md
index eee3629c..44c9ff9d 100644
--- a/docs/source/design/transfer-engine/cpp-api.md
+++ b/docs/source/design/transfer-engine/cpp-api.md
@@ -337,7 +337,7 @@ The HTTP server should implement three following RESTful APIs, while the metadat
 2. `PUT /metadata?key=$KEY`: Update the metadata corresponding to `$KEY` to the value of the request body.
 3. `DELETE /metadata?key=$KEY`: Delete the metadata corresponding to `$KEY`.
 
-For specific implementation, refer to the demo service implemented in Golang at [mooncake-transfer-engine/example/http-metadata-server](../../../mooncake-transfer-engine/example/http-metadata-server).
+For specific implementation, refer to the demo service implemented in Golang at [mooncake-transfer-engine/example/http-metadata-server](gh-dir:mooncake-transfer-engine/example/http-metadata-server).
 
 ### Initialization
 
diff --git a/docs/source/design/transfer-engine/efa_transport.md b/docs/source/design/transfer-engine/efa_transport.md
index f1fd0df3..c1c10a68 100644
--- a/docs/source/design/transfer-engine/efa_transport.md
+++ b/docs/source/design/transfer-engine/efa_transport.md
@@ -4,6 +4,7 @@ This document describes how to build and use Mooncake with AWS Elastic Fabric Ad
 
 ## Prerequisites
 
+(efa-prerequisites-driver)=
 ### 1. AWS EFA Driver and libfabric
 
 EFA driver and libfabric should be pre-installed on AWS instances with EFA support (e.g., p6-b300.48xlarge, p6-b200.48xlarge, p5en.48xlarge, p5e.48xlarge, p5.48xlarge).
@@ -32,14 +33,23 @@ sudo ./dependencies.sh -y
 
 This installs all system packages, git submodules (including pybind11 and yalantinglibs), and Go.
 
-**Additional EFA-specific dependencies** (not covered by `dependencies.sh`):
+> **Note:** The EFA driver and libfabric are **not** installed by `dependencies.sh`. They must be pre-installed on the instance (see section 1 above).
+
+## Installing from PyPI (recommended)
+
+Pre-built EFA wheels are published to PyPI by the official release pipeline, so most users do not need to build from source. The EFA transport's memory path is CUDA-aware, so two variants are published:
 
 ```bash
-# gflags is needed by transfer_engine_bench and EFA unit tests
-sudo apt-get install -y libgflags-dev
+# GPU memory transfers (e.g., KV cache in vLLM) — built with USE_CUDA=ON
+pip install mooncake-transfer-engine-efa
+
+# CPU/DRAM-only transfers — built with USE_CUDA=OFF
+pip install mooncake-transfer-engine-efa-non-cuda
 ```
 
-> **Note:** The EFA driver and libfabric are **not** installed by `dependencies.sh`. They must be pre-installed on the instance (see section 1 above).
+> **Note:** These wheels deliberately do **not** bundle `libfabric`/`libefa` (see the runtime note in [Building a Distributable Wheel](#efa-distributable-wheel)). They resolve to the system AWS EFA installation at runtime, so the EFA driver and libfabric from the [Prerequisites](#efa-prerequisites-driver) must still be present on the instance. Make sure `/opt/amazon/efa/lib` is on `LD_LIBRARY_PATH`.
+
+To build from source instead (for development, an unreleased revision, or a custom configuration), follow the sections below.
 
 ## Building Mooncake with EFA Support
 
@@ -87,6 +97,31 @@ cp mooncake-common/libasio.so ../mooncake-wheel/mooncake/
 pip install -e ../mooncake-wheel --no-build-isolation
 ```
 
+(efa-distributable-wheel)=
+### 3. Building a Distributable Wheel (optional)
+
+To produce a relocatable wheel for distribution (instead of the editable install above), use `scripts/build_wheel.sh`, which runs `auditwheel repair` to bundle non-system dependencies:
+
+```bash
+# After the cmake/make build above completes:
+PYTHON_VERSION=3.13 BUILD_DIR=build bash scripts/build_wheel.sh 3.13 dist
+pip install dist/mooncake_transfer_engine-*.whl
+```
+
+To produce a wheel whose package name matches the published variants (`mooncake-transfer-engine-efa` / `-efa-non-cuda`), set the corresponding build-variant environment variable — this is exactly what the release pipeline does:
+
+```bash
+# GPU build (cmake was configured with USE_CUDA=ON):
+EFA_BUILD=1 PYTHON_VERSION=3.13 BUILD_DIR=build bash scripts/build_wheel.sh 3.13 dist
+
+# CPU build (cmake was configured with USE_CUDA=OFF):
+EFA_NON_CUDA_BUILD=1 PYTHON_VERSION=3.13 BUILD_DIR=build bash scripts/build_wheel.sh 3.13 dist
+```
+
+> **CI/CD:** EFA wheels are built and published automatically — see `.github/workflows/ci_efa.yml` (per-PR build validation) and `.github/workflows/release-efa.yaml` (tagged release to GitHub Release + PyPI). No EFA hardware is required to *build* the wheel: only the libfabric headers/library are needed to compile and link, which the CI runner obtains from the distro `libfabric-dev` package.
+
+> **Important (EFA builds):** `auditwheel repair` excludes `libfabric` and `libefa` from the wheel so they resolve to the system EFA installation (`/opt/amazon/efa/lib`) at runtime. This is required because the in-process `aws-ofi-nccl` plugin (loaded by NCCL) links the **same** system `libfabric`. If the wheel bundled its own copy, the process would load two independent libfabric instances — Mooncake's bundled one and NCCL's system one — and whichever initializes first claims the EFA device, leaving the other with an empty provider list (`fi_getinfo: provider efa output empty list`). NCCL then silently falls back to the TCP provider and cross-node collectives such as `all_gather_object` hang. Excluding libfabric/libefa (see `scripts/build_wheel.sh`) keeps a single shared libfabric in the process. If you are on an older Mooncake build whose wheel still bundles libfabric, force the system copy with `export LD_PRELOAD=/opt/amazon/efa/lib/libfabric.so.1` as a workaround.
+
 ## Verification
 
 Test EFA transport initialization:
@@ -99,7 +134,7 @@ result = te.initialize('127.0.0.1', 'P2PHANDSHAKE', 'efa', '')
 print(f'Initialize result: {result}')  # Should be 0
 
 # You should see logs like:
-# EFA device (libfabric): rdmap79s0, domain: rdmap79s0-rdm, provider: efa
+# EFA device (libfabric): rdmap79s0, domain: rdmap79s0-rdm, fabric: efa, provider: efa
 ```
 
 ## Unit Tests
@@ -119,34 +154,43 @@ The test suite includes:
 | `WriteAndRead` | Write then read with data integrity check |
 | `MultiWrite` | Batch write (16 requests) |
 | `StressMultipleBatches` | Stress test (20 batches x 8 requests) |
+| `WarmupSegmentLoopback` | `warmupSegment()` handshake path + idempotent re-call |
+| `WarmupSegmentNotFound` | `warmupSegment()` fails cleanly for an unknown segment |
+| `RegisterMemoryBatch` | `registerLocalMemoryBatch` / `unregisterLocalMemoryBatch` round-trip |
+| `LargeTransfer` | 128 MB buffer, 64 x 1 MB slices — exercises WR / CQ pacing |
+| `RepeatedOpenSegment` | `openSegment()` on the same peer repeatedly still transfers correctly |
 
-You can also run all unit tests via CTest:
+You can also run the EFA tests via CTest:
 
 ```bash
-cd build && ctest --output-on-failure
+cd build && ctest --output-on-failure -R 'efa'
 ```
 
-Environment variables for test configuration:
-
-```bash
-export MC_METADATA_SERVER=P2PHANDSHAKE     # default
-export MC_LOCAL_SERVER_NAME=127.0.0.1:12345  # default
-```
+> **Note:** `ctest --output-on-failure` without a filter runs every test in the build, including TCP / metadata / master-service suites that require an etcd server or a running `mooncake_master`. Those will fail or hang on a machine that is only provisioned for EFA testing — the failures are not EFA-specific. Use `-R 'efa'` to restrict the run to the EFA tests.
 
 ## Performance Benchmark
 
 Use `transfer_engine_bench` to measure EFA transport throughput between two nodes.
 
-### Target Node (receiver)
+The following commands are the GPU-to-GPU configuration that produces the headline numbers in the [Benchmark Results](#benchmark-results) tables (≈ 350 GB/s write on a p5en.48xlarge pair, ≈ 302 GB/s on p6-b200.48xlarge). Two things matter the most:
+
+- `--gpu_id=-1` on **both** sides — this fans buffers across every GPU, which in turn lets both NUMA nodes' NICs saturate. Pinning a single GPU (the default `--gpu_id=0`) halves throughput because half the NICs end up cross-NUMA.
+- `--block_size=1048576` (1MB, not the 64 KB default) — each block becomes one `fi_write` / `fi_read`, so larger blocks amortize per-op overhead and are the main knob for hitting line rate.
+
+### 1. Target Node (receiver)
 
 ```bash
 ./build/mooncake-transfer-engine/example/transfer_engine_bench \
     --mode=target \
     --protocol=efa \
-    --metadata_server=P2PHANDSHAKE
+    --metadata_server=P2PHANDSHAKE \
+    --buffer_size=4294967296 \
+    --gpu_id=-1
 ```
 
-### Initiator Node (sender)
+`--buffer_size` must be at least as large as the initiator's `--buffer_size` — the initiator writes into offsets `[0, buffer_size)` on the target, so keep these in sync.
+
+### 2. Initiator Node (sender)
 
 ```bash
 ./build/mooncake-transfer-engine/example/transfer_engine_bench \
@@ -156,45 +200,78 @@ Use `transfer_engine_bench` to measure EFA transport throughput between two node
     --segment_id=<target_hostname>:<target_port> \
     --operation=write \
     --duration=10 \
-    --threads=8 \
-    --block_size=65536 \
+    --threads=16 \
+    --block_size=1048576 \
     --batch_size=128 \
-    --buffer_size=1073741824 \
+    --buffer_size=4294967296 \
+    --gpu_id=-1 \
     --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`).
 
+> **CPU-to-CPU** (no GPUs): build with `-DUSE_CUDA=OFF`, **or** pass `--use_vram=false` to a CUDA-enabled binary. Drop `--gpu_id=-1` in that case — the bench will spread buffers across NUMA nodes instead.
+
+> **Why `threads=16` and not 32:** the SRD shared endpoint caps outstanding WRs per NIC (default 256 — see `MC_MAX_WR`). With `threads × batch ≤ NICs × max_wr` the CQ never saturates; going higher triggers backoff and times out. 32 threads × 128 batch = 4096 slices chasing 16 × 256 = 4096 WRs has no headroom, so the steady-state config settles at 16 threads.
+
 ### Key Parameters
 
 | Parameter | Default | Description |
 |-----------|---------|-------------|
-| `--block_size` | 65536 | Bytes per transfer request |
+| `--mode` | initiator | `initiator` (sender) or `target` (receiver) |
+| `--protocol` | rdma | Transport protocol; use `efa` here |
+| `--metadata_server` | `192.168.3.77:2379` | etcd address or `P2PHANDSHAKE` for standalone use |
+| `--segment_id` | `192.168.3.76` | Initiator only: `<target_host>:<target_port>` from the target's startup log |
+| `--operation` | read | `read` or `write` |
+| `--block_size` | 65536 | Bytes per transfer request; **1 MB (1048576) is the main knob for EFA throughput** |
 | `--batch_size` | 128 | Requests per batch |
-| `--threads` | 12 | Concurrent submission threads |
-| `--buffer_size` | 1 GB | Total buffer size (per GPU when `--gpu_id=-1`) |
+| `--threads` | 12 | Concurrent submission threads (initiator) |
+| `--buffer_size` | 1 GB | Buffer size (per GPU when `--gpu_id=-1`, otherwise total) |
 | `--duration` | 10 | Test duration in seconds |
-| `--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`) |
+| `--report_unit` | GB | `GB\|GiB\|Gb\|MB\|MiB\|Mb\|KB\|KiB\|Kb` |
+| `--use_vram` | true | Allocate from GPU VRAM (requires `-DUSE_CUDA=ON`); pass `--use_vram=false` for CPU-to-CPU on a CUDA build |
+| `--gpu_id` | 0 | GPU device ID when `--use_vram=true`; `-1` fans buffers across every GPU and is what actually saturates all NICs in a GPU-to-GPU run |
+| `--init_mem` | true | Zero-fill the allocated buffer; rarely needs to change |
+| `--auto_discovery` | false | Auto-discover topology on init; off for reproducible runs |
 
-| 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 on EFA slicing:** EFA transport does not split each transfer into fixed-size slices the way RDMA transport does — each transfer is sent as a single `fi_write` / `fi_read` whose size equals `block_size`, round-robin'd across NICs per request. **`block_size` is the key tuning parameter** for EFA throughput.
 
 > **Note:** `buffer_size` must be >= `block_size * batch_size * threads`. The benchmark auto-adjusts if too small.
 
+(benchmark-results)=
 ### Benchmark Results
 
-#### p6-b200.48xlarge (B200, 8 EFA × 400 Gbps)
+#### 1. 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. Numbers below are post-SRD-shared-endpoint (#1944) on a fresh `main` build with the DLAMI pytorch env (CUDA 13).
+
+**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 | 758.88 GB/s | 720.35 GB/s |
+| block=1MB, threads=32, batch=64 | 753.32 GB/s | **755.78 GB/s** |
+| **block=1MB, threads=32, batch=32** | **780.33 GB/s** | - |
+| block=1MB, threads=64, batch=32 | 780.23 GB/s | - |
+
+> **Peak: 780 GB/s write**, reaching ~97.5% 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`, or `--use_vram=false` on a CUDA build, `--buffer_size=4294967296`):
+
+| Configuration | Write | Read |
+|---------------|-------|------|
+| **block=1MB, threads=32, batch=128** | **282.93 GB/s** | **270.47 GB/s** |
+| block=1MB, threads=16, batch=128 | 282.04 GB/s | 249.67 GB/s |
+| block=1MB, threads=32, batch=64 | 282.84 GB/s | 256.26 GB/s |
+
+> CPU-to-CPU is bounded by DRAM bandwidth on the Xeon 8559C — write throughput is essentially flat across thread/batch combinations (~282 GB/s), confirming DRAM controller saturation rather than a NIC or in-flight-WR limit.
+
+#### 2. p6-b200.48xlarge (B200, 8 EFA × 400 Gbps)
 
 Tested on two p6-b200.48xlarge instances in the same AWS placement group.
 
+> **Note:** numbers below predate the SRD shared-endpoint refactor (#1944) and current EFA tuning work. They are a lower bound for the current code; we will re-sweep and update when a B200 pair is available again.
+
 **GPU-to-GPU** (build with `-DUSE_CUDA=ON`, `--gpu_id=-1` for all 8 GPUs):
 
 | Configuration | Write | Read |
@@ -208,121 +285,223 @@ Tested on two p6-b200.48xlarge instances in the same AWS placement group.
 |---------------|-------|------|
 | 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 |
-| 128KB | 32 | 128 | default        | 92.33 GB/s |
-| 128KB | 32 | 128 | 256KB          | 156.18 GB/s |
-| 128KB | 48 | 128 | 256KB          | 160.34 GB/s |
-
-> **Note:** These results predate the EFA striping optimization. With the current code, `MC_SLICE_SIZE` no longer affects EFA performance. Use `--block_size=1048576` (1MB) instead, which achieves 222 GB/s.
-
-</details>
-
-#### p6-b300.48xlarge (B300, 16 EFA × 400 Gbps)
-
-Tested on two p6-b300.48xlarge instances (Intel Xeon Platinum 8559C, 8× B300, 16 EFA devices) in the same AWS placement group.
-
-**GPU-to-GPU** (build with `-DUSE_CUDA=ON`, `--gpu_id=-1` for all 8 GPUs, `--buffer_size=2147483648`):
-
-| Configuration | Write | Read |
-|---------------|-------|------|
-| block=1MB, threads=16, batch=128 | 701 GB/s | **697 GB/s** |
-| **block=1MB, threads=32, batch=64** | **752 GB/s** | 713 GB/s |
-| block=1MB, threads=32, batch=32 | 751 GB/s | - |
-| block=1MB, threads=64, batch=32 | 728 GB/s | - |
-
-> **Peak: 752 GB/s write**, reaching ~94% of the 800 GB/s theoretical line rate (16×400 Gbps). GPUDirect RDMA bypasses DRAM entirely (HBM3e → PCIe switch → NIC), so performance is not bottlenecked by CPU memory bandwidth.
-
-**CPU-to-CPU** (build with `-DUSE_CUDA=OFF`):
-
-| Configuration | Write | Read |
-|---------------|-------|------|
-| **block=1MB, threads=32, batch=128, buf=4GB** | **230 GB/s** | 180 GB/s |
-| block=16MB, threads=32, batch=8, buf=8GB (striping off) | 233 GB/s | - |
-
-> CPU-to-CPU is bounded by DRAM bandwidth (~250 GB/s/socket on Xeon 8559C). Per-NIC sampling shows NUMA-0 NICs at 90 Gbps and NUMA-1 NICs at 53 Gbps, confirming DRAM controller saturation rather than NIC limit.
-
-#### p5en.48xlarge (H200, 16 EFA × 200 Gbps)
+#### 3. 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):
+**GPU-to-GPU** (build with `-DUSE_CUDA=ON`, `--gpu_id=-1` for all 8 GPUs, `--buffer_size=4294967296`):
 
 | 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 |
+| block=1MB, threads=8, batch=128 | 318.71 GB/s | 277.06 GB/s |
+| **block=1MB, threads=16, batch=128** | **365.66 GB/s** | 284.22 GB/s |
+| **block=1MB, threads=16, batch=32** | 297.23 GB/s | **303.78 GB/s** |
+| block=1MB, threads=32, batch=64 | 357.12 GB/s | 279.61 GB/s |
+| block=1MB, threads=32, batch=128 | 364.21 GB/s | 250.90 GB/s |
+| block=1MB, threads=48, batch=64 | 363.47 GB/s | 268.43 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.
+> **Peak write: 365 GB/s** at `threads=16, batch=128` — ~91% of the 400 GB/s theoretical line rate (16×200 Gbps). Write saturates on batch size, so `batch=128` outperforms smaller batches as long as `threads × batch ≤ 16 × 256 = 4096` (the shared-endpoint WR cap). **Peak read: 304 GB/s** at `threads=16, batch=32` — reads tolerate smaller in-flight queues, and throughput drops as batch grows.
 
-> **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`):
+**CPU-to-CPU** (build with `-DUSE_CUDA=OFF`, or `--use_vram=false` on a CUDA build, `--buffer_size=4294967296`):
 
 | 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** |
+| block=1MB, threads=8, batch=128 | 210.76 GB/s | 209.93 GB/s |
+| block=1MB, threads=16, batch=128 | 212.71 GB/s | 211.21 GB/s |
+| **block=1MB, threads=16, batch=32** | 211.67 GB/s | **212.18 GB/s** |
+| block=1MB, threads=32, batch=128 | 212.99 GB/s | 210.33 GB/s |
+| **block=1MB, threads=48, batch=32** | **213.57 GB/s** | 206.92 GB/s |
 
-> CPU-to-CPU throughput is bottlenecked by DRAM bandwidth (~155 GB/s per NUMA node, measured with STREAM Copy).
+> CPU-to-CPU is DRAM-bound — throughput is essentially flat (~205–214 GB/s) across every thread / batch combination that doesn't hit the WR cap. Peak write 213.57 GB/s, peak read 212.18 GB/s.
 
-#### Cross-Transport Comparison
+**block_size sweep** (p5en GPU-to-GPU, `threads=16 batch=128 buf=4GB --gpu_id=-1`):
 
-| 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 |
+| block | Write | Read |
+|-------|-------|------|
+| 64 KB (default) | 97.11 GB/s | 96.87 GB/s |
+| 128 KB | 190.82 GB/s | 203.88 GB/s |
+| 256 KB | 324.90 GB/s | 296.19 GB/s |
+| 512 KB | 357.62 GB/s | 289.47 GB/s |
+| **1 MB (recommended)** | **352.59 GB/s** | **301.14 GB/s** |
+| 2 MB | 366.87 GB/s | 302.25 GB/s |
 
-**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).
+> The 64 KB default only reaches ~26% of peak. Write throughput climbs steeply up to ~512 KB and plateaus between 1 MB and 2 MB; read saturates at ~256 KB. 1 MB is the recommended value — within a few percent of the 2 MB peak with more headroom for `batch_size` under the shared-endpoint WR cap.
+
+**buffer_size sweep** (p5en GPU-to-GPU, `threads=16 batch=128 block=1MB --gpu_id=-1`):
+
+| buffer_size | Write | Read |
+|-------------|-------|------|
+| 2 GB (min: `block × batch × threads`) | 353.51 GB/s | 297.03 GB/s |
+| 4 GB | 364.86 GB/s | 293.79 GB/s |
+
+> `buffer_size` only needs to satisfy `buffer_size ≥ block_size × batch_size × threads` (the bench auto-adjusts if smaller, but silently). Anything larger than that minimum does not change throughput — 2 GB vs 4 GB differs by ~3% on write, read is flat within noise. The example commands use 4 GB because it is safe for any reasonable threads/batch combination without having to recompute the minimum.
+
+#### 4. p5.48xlarge (H100, 32 EFA × 100 Gbps)
+
+Tested on two p5.48xlarge instances (AMD EPYC 7R13, 8× H100 80GB, 32 EFA devices) in the same AWS placement group. Per-NIC line rate is half of p5en's, but with twice the NIC count the aggregate ceiling is the same 400 GB/s. The shared-endpoint WR cap scales with NIC count: `32 NICs × 256 = 8192` in-flight slots, so `threads × batch_size ≤ 8192` (vs 4096 on p5en).
+
+**GPU-to-GPU** (build with `-DUSE_CUDA=ON`, `--gpu_id=-1` for all 8 GPUs, `--buffer_size=4294967296`):
+
+| Configuration | Write | Read |
+|---------------|-------|------|
+| block=1MB, threads=8, batch=128 | 335.11 GB/s | - |
+| block=1MB, threads=16, batch=128 | 388.52 GB/s | 379.10 GB/s |
+| block=1MB, threads=32, batch=64 | 388.83 GB/s | 379.78 GB/s |
+| **block=1MB, threads=32, batch=128** | **388.90 GB/s** | **381.64 GB/s** |
+| block=1MB, threads=16, batch=32 | - | 356.94 GB/s |
+| block=1MB, threads=32, batch=32 | - | 380.66 GB/s |
+
+> **Peak write: 389 GB/s** at `threads=32, batch=128` — ~97% of the 400 GB/s theoretical line rate (32×100 Gbps). The plateau is wide: any `(threads, batch)` between `(16, 128)` and `(32, 128)` lands within 0.1% of peak. **Peak read: 382 GB/s** at `threads=32, batch=128` — unlike p5en, reads on this host scale with batch size up to 128 because the wider 32-NIC fabric absorbs larger in-flight queues without backoff. `(32, 256)` and `(64, 128)` (both at the 8192 WR cap) fail with no headroom for retries.
+
+**CPU-to-CPU** (build with `-DUSE_CUDA=OFF`, or `--use_vram=false` on a CUDA build, `--buffer_size=4294967296`):
+
+| Configuration | Write | Read |
+|---------------|-------|------|
+| block=1MB, threads=16, batch=128 | 39.83 GB/s | 40.43 GB/s |
+| block=1MB, threads=32, batch=64 | 47.31 GB/s | 48.06 GB/s |
+| block=1MB, threads=32, batch=128 | 47.75 GB/s | 48.62 GB/s |
+| block=1MB, threads=32, batch=32 | 55.66 GB/s | 56.68 GB/s |
+| **block=1MB, threads=48, batch=16** | **63.60 GB/s** | 63.00 GB/s |
+| **block=1MB, threads=64, batch=16** | 63.39 GB/s | 63.05 GB/s |
+| block=1MB, threads=32, batch=16 | 63.21 GB/s | 60.82 GB/s |
+| block=1MB, threads=96, batch=32 | 57.61 GB/s | 59.63 GB/s |
+
+> **Peak: ~64 GB/s** on both write and read — far below the GPU-to-GPU number despite identical NIC count. The bottleneck is DDR4-3200 DRAM bandwidth on the EPYC 7R13 (Milan): `batch=16` consistently wins because larger in-flight queues only deepen DRAM contention without unlocking new NIC capacity. p5.48xlarge CPU-to-CPU runs around **3× slower than p5en** (DDR5 Xeon 8488C, ~213 GB/s) at the same NIC aggregate. For PD KV transfer, the GPU-to-GPU path is the relevant one.
+
+### Single-host loopback
+
+EFA NICs have no hardware loopback short-circuit: when a transfer's source and destination resolve to the same host, the data does not go out on the wire as GPUDirect/device RDMA. libfabric handles the same-host case in software, and there are **two distinct provider knobs** that select how:
+
+- **`FI_EFA_ENABLE_SHM_TRANSFER`** (default `1`, on): when on, the EFA provider routes same-host peers through the **`shm` provider** — verifiable at runtime, where libfabric reports `Opened fabric: shm` alongside `Opened fabric: efa` even on a default (device-RDMA-enabled) configuration. This SHM path is the one that supplies the same-host memcpy fast path; it is active **by default**, independent of `FI_EFA_USE_DEVICE_RDMA`.
+- **`FI_EFA_USE_DEVICE_RDMA`** (default `1` after #2041): controls whether the EFA RDM data path uses device RDMA vs libfabric's emulated RDM path. It is a provider-level flag resolved at `fi_getinfo` time; Mooncake does not wrap it.
+
+```{warning}
+**GPU (FI_HMEM_CUDA) buffers — known segfault.** The default same-host **SHM**
+path (`FI_EFA_ENABLE_SHM_TRANSFER=1`) performs a **host `memcpy` into the
+destination buffer** during SHM SAR reassembly, *without* honoring an
+`FI_HMEM_CUDA` destination's iface. On a GPU buffer this writes host memory
+straight into a device pointer and **segfaults** on the first same-host
+transfer — `__memcpy_avx_unaligned` ← `ofi_copy_to_mr_iov` ← `smr_copy_from_sar`
+← `efa_rdm_cq_readfrom` ← `fi_cq_read`. Reported upstream as
+[ofiwg/libfabric#12328](https://github.com/ofiwg/libfabric/issues/12328).
+
+- **Same-process self-loopback** (e.g. a TP-colocated rank reading its own
+  registered GPU weights — checkpoint-engine p2p weight update) is handled
+  inside Mooncake: `EfaContext::tryLoopbackCopy` detects a same-process peer
+  (matched on `local_server_name`, which embeds this process's unique RPC
+  port) and satisfies the transfer with a local `cudaMemcpy` instead of routing
+  it over EFA, so it never reaches the broken SHM path.
+- **Same-host cross-process GPU transfers** are *not* short-circuited (the
+  peer is a different process / address space). Until libfabric#12328 is fixed,
+  set `FI_EFA_ENABLE_SHM_TRANSFER=0` on such processes — same-host transfers
+  then fall back to device RDMA, which is GPU-aware and correct.
+```
+
+For **host (DRAM) buffers** the SHM memcpy path is safe (host→host copy) and is the same-host fast path the measurements below exercise.
+
+Measured on p5.48xlarge (1 NIC, ~1.2 GiB per `put_from` call, host DRAM buffer, same-host producer/consumer in **separate processes**):
+
+| same-host path | per-write latency |
+|---|---:|
+| device RDMA (NIC round-trip, no fast-path for loopback) | ~830 ms |
+| SHM memcpy fast path (default) | ~390 ms |
+
+For reference, a cross-host `put_from` of the same payload (device RDMA, 1 NIC) is ~340 ms — i.e., driving a same-host loopback through the NIC is *slower* than going over the wire to another host, because the NIC has no fast-path for loopback. Cross-host transfers always use device RDMA and are unaffected by `FI_EFA_ENABLE_SHM_TRANSFER`: leave it at its default on any process that also talks to remote peers.
 
 ### Tuning Tips
 
-- **Use `--block_size=1048576` (1MB)** — this is the most important tuning parameter for EFA. Each `block_size`-sized transfer becomes a single `fi_write`/`fi_read` call, so larger blocks amortize per-operation overhead. 1MB gives ~2× throughput over the 64KB default.
-- `MC_SLICE_SIZE` has **no effect** on EFA transport (it only applies to RDMA transport). Use `block_size` instead.
-- Increase `--threads` to 32-48 to saturate multiple EFA devices (2-4 threads per device is a good starting point)
-- For **CPU-to-CPU**: use `--block_size=1048576` (1MB) with NUMA-split (separate instances per NUMA node) for best results
-- For **GPU-to-GPU**: use `--block_size=1048576` (1MB), `--gpu_id=-1` (all GPUs), and `--buffer_size=2147483648` (2GB max per GPU). Write peaks at threads=32, read at threads=16
-- Keep `--batch_size` such that `block_size * batch_size * threads <= buffer_size`
-- Allocate buffers on both NUMA nodes for balanced NIC utilization (the bench tool does this by default for CPU mode)
-- On 16-NIC instances (p5en), writes are NUMA-sensitive: 8 local-NUMA NICs reach 90 Gbps each, while 8 cross-NUMA NICs only reach ~20 Gbps without NUMA-split
+- **Use `--block_size=1048576` (1MB)** — the single most important knob. The 64 KB default reaches only ~26% of peak. 1 MB is within a few percent of the 2 MB plateau while leaving headroom for `batch_size` under the shared-endpoint WR cap.
+- **Keep `threads × batch_size ≤ num_nics × max_wr`** — under the SRD shared endpoint each NIC carries one `fid_ep` with a 256 WR cap (`MC_MAX_WR`), giving `16 NICs × 256 = 4096` in-flight slots on a 16-NIC host (b300, b200, p5en) and `32 NICs × 256 = 8192` on p5. Exceeding this trips "timed out waiting for CQ drain". `threads=16, batch=128` is a solid baseline on 16-NIC hosts; on 32-NIC p5, `threads=32, batch=128` works the same way.
+- **Write vs read:** write benefits from larger batches (peak at `batch=128`); on 16-NIC p5en read prefers smaller queues (peak at `batch=32`), but on 32-NIC p5 reads scale up to `batch=128` because the wider fabric absorbs larger in-flight queues.
+- For **GPU-to-GPU**: pass `--gpu_id=-1` on **both** sides so buffers fan out across every GPU. Pinning a single GPU halves throughput because half the NICs end up cross-NUMA.
+- For **CPU-to-CPU**: DRAM bandwidth is the ceiling. NUMA-split (separate initiator/target instances per NUMA node) can help reduce contention when one instance can't saturate both nodes.
+- `--buffer_size` only needs `≥ block × batch × threads`; larger values do not improve throughput. The example commands use 4 GB because that is safe for any reasonable config.
 
-### Eager endpoint warmup (first-request latency)
+### 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).
+Peer addressing resolves lazily: `fi_av_insert()` and the metadata handshake fire on the first send to each `(local_NIC, peer_NIC)` pair. On 16-NIC hosts, the first few `submitTransfer` calls carry this cost before steady state.
 
-Mooncake exposes an explicit eager-warmup API to eliminate the stall:
+**Measured on p5en (16 × 16 NICs, cross-node, 1 MB write, 3 reps, median):**
 
-- 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 |
+| | cold submit #0 (no warmup) | `warmupSegment()` (all 256 pairs) |
 |---|---:|---:|
-| 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 |
+| SRD shared endpoint (#1944) | **26 ms** | **1.1 s** |
+| Per-peer `fid_ep` (upstream main) | 99 ms | 17 s |
+| Speedup | **~4×** | **~15×** |
 
-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.
+The SRD shared-endpoint refactor (#1944) speeds up first-request latency two different ways:
+
+- **Without any code change from callers** — the cold `submitTransfer` is ~4× faster (26 ms vs 99 ms), because the shared endpoint removes the per-peer `fi_endpoint` / `fi_enable` that used to dominate. This is what existing Mooncake callers (vLLM, SGLang, etc.) will see.
+- **For callers that want sub-10 ms first-request latency**, an explicit eager-warmup API lets you pay the handshake cost up front, outside the critical path:
+  - 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)`
+  - Python: `engine.warmup_efa_segment(segment_name)`
+
+Call once per peer right after `openSegment`. The call is idempotent. Under this refactor `warmupSegment` itself is ~15× faster than the pre-#1944 code (1.1 s vs 17 s), bounded by the peer's single-threaded handshake RPC daemon (`accept` + JSON parse serialized on one thread), so it scales linearly with the number of fresh NIC pairs.
+
+vLLM and SGLang do not currently call `warmupSegment` — they go through the generic `TransferEngine` interface and pick up the 4× cold-submit speedup automatically. The API is there for direct Mooncake callers that want the larger win.
+
+#### Reproducing the numbers — `efa_first_submit_probe`
+
+`mooncake-transfer-engine/example/efa_first_submit_probe.cpp` is a two-host probe that isolates the two latency costs `transfer_engine_bench` hides inside its 10-second throughput average:
+
+1. **Cold first-submit** — the handshake + `fi_av_insert()` that fires on the first send to each `(local_NIC, peer_NIC)` pair.
+2. **Eager warmup** — how much of that is paid up front by an explicit `warmupSegment()` call.
+
+It is EFA-specific (there is no `warmupSegment` on the RDMA / TCP transports — they establish connections at `connect` time, so "pre-warming" has no meaning there) and intentionally not wired into `ctest`: it needs two hosts.
+
+**Run**:
+
+```bash
+# On the target host:
+./build/mooncake-transfer-engine/example/efa_first_submit_probe \
+    --mode=target \
+    --metadata_server=P2PHANDSHAKE \
+    --local_server_name=$(hostname):12345
+
+# Note the "[target] ready, addr=<ip>:<port>" line, then on the initiator:
+./build/mooncake-transfer-engine/example/efa_first_submit_probe \
+    --mode=initiator \
+    --metadata_server=P2PHANDSHAKE \
+    --local_server_name=$(hostname):12346 \
+    --segment_id=<target_ip>:<target_port> \
+    --warmup=1 \
+    --iters=5
+```
+
+**What it prints**:
+
+```
+warmup:   1094.74 ms (rc=0)           # present only with --warmup=1
+submit #0:   6.95 ms                   # first user-visible submit
+submit #1:   5.45 ms                   # steady state
+submit #2:   5.92 ms
+submit #3:   5.67 ms
+submit #4:   0.09 ms
+```
+
+- `warmup:` is the wall time of `EfaTransport::warmupSegment()` — it opens every `(local_NIC, peer_NIC)` pair (16 × 16 = 256 on p5en) up front.
+- `submit #0` through `#N-1` are the timings of single 1 MB `submitTransfer + poll`. With `--warmup=1` they are all in the steady-state regime. With `--warmup=0`, `#0` pays the handshake cost, `#1+` are steady state.
+
+**Flags**:
+
+| Flag | Default | What it controls |
+|---|---|---|
+| `--warmup` | `true` | Call `warmupSegment()` before the first submit |
+| `--iters` | `5` | How many timed submits after warmup |
+| `--xfer_size` | `1<<20` (1 MB) | Bytes per submit |
+| `--buffer_size` | `1<<30` (1 GB) | Registered buffer size |
+| `--local_server_name` | `hostname:12345` | Local metadata advertise name |
+
+**Use cases**:
+
+- **Deciding whether your application needs `warmupSegment()`**: if `submit #0` in the `--warmup=0` run is acceptable for your use case, you don't need to call `warmupSegment` at all.
+- **Comparing PR branches**: run the probe on the same hardware on this PR's branch vs `main` (or whatever upstream you're benchmarking against) to see per-NIC-pair handshake cost directly, without having the number drowned in a 10-second throughput average.
 
 ## Usage with vLLM
 
-### Prefill Instance
+### 1. Prefill Instance
 
 ```bash
 VLLM_MOONCAKE_BOOTSTRAP_PORT=8998 \
@@ -332,7 +511,7 @@ vllm serve <model_path> -tp 8 \
    --kv-transfer-config '{"kv_connector":"MooncakeConnector","kv_role":"kv_producer","kv_connector_extra_config":{"mooncake_protocol":"efa"}}'
 ```
 
-### Decode Instance
+### 2. Decode Instance
 
 ```bash
 vllm serve <model_path> -tp 8 \
@@ -341,105 +520,166 @@ vllm serve <model_path> -tp 8 \
    --kv-transfer-config '{"kv_connector":"MooncakeConnector","kv_role":"kv_consumer","kv_connector_extra_config":{"mooncake_protocol":"efa"}}'
 ```
 
+### 3. Router
+
+Front the prefill / decode pair with `vllm-router`. **`PREFILL_HOST` / `DECODE_HOST` must be each node's reachable private IP, not `127.0.0.1`** — the router forwards these addresses to the peer for the KV handshake; localhost will fail with `Connection refused` and stall traffic.
+
+```bash
+vllm-router --policy round_robin \
+  --vllm-pd-disaggregation \
+  --prefill http://<prefill_ip>:8010 \
+  --decode  http://<decode_ip>:8020 \
+  --kv-connector mooncake \
+  --host 0.0.0.0 --port 30000
+```
+
+> **Do not add `--intra-node-data-parallel-size` here.** The prefill / decode instances above are launched with `-tp 8` (pure tensor parallelism, data parallel size = 1), so there is no intra-node DP to advertise. Only pass `--intra-node-data-parallel-size N` when your instances actually run `N`-way data parallelism per node (e.g. you launched them with `--data-parallel-size N`); setting it to match `-tp 8` is wrong and will misroute requests.
+
 ## Usage with SGLang
 
-SGLang's Mooncake integration currently hardcodes the `"rdma"` protocol. To use EFA transport, apply the provided patch and set environment variables.
+SGLang's PD-disaggregation Mooncake integration reads the transport from `MOONCAKE_PROTOCOL`. Set it to `efa` and select the Mooncake backend with `--disaggregation-transfer-backend mooncake`.
 
-### 1. Apply EFA Patch
+### 1. Apply EFA Patch (only if SGLang version predates PR #25083)
 
-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):
+Older SGLang releases hardcode `"rdma"` in the transfer engine init. [SGLang PR #25083](https://github.com/sgl-project/sglang/pull/25083) has been **merged into SGLang `main`**, so the protocol is now read from `MOONCAKE_PROTOCOL`. If your SGLang build includes that PR (any recent `main` or release built after it), **this step is unnecessary** — skip to step 2.
+
+Only if you are pinned to an older release that predates PR #25083, apply the [patch script](https://github.com/whn09/kimi-k2-sglang/blob/main/patch_sglang_efa.sh):
 
 ```bash
 bash patch_sglang_efa.sh
 ```
 
-This is idempotent and safe to rerun.
+The script is idempotent and safe to rerun.
 
 ### 2. Environment Variables
 
+Only one Mooncake-specific env is required:
+
 ```bash
 export MOONCAKE_PROTOCOL=efa
+```
+
+If your container does not already export libfabric/EFA paths in its `Dockerfile`, also set:
+
+```bash
 export FI_PROVIDER=efa
 export FI_EFA_USE_DEVICE_RDMA=1
-export GLOO_SOCKET_IFNAME=enp71s0  # adjust to your instance's primary interface
+export LD_LIBRARY_PATH=/opt/amazon/efa/lib:$LD_LIBRARY_PATH
 ```
 
-For multi-node expert parallelism (EP) deployments, also set:
+> **Note on additional `MC_*` knobs:** `MC_NUM_CQ_PER_CTX`, `MC_MAX_WR`, `MC_MAX_CQE_PER_CTX`, `MC_SLICE_SIZE`, and `MC_EFA_STRIPING_THRESHOLD` are **not** required at typical PD-disagg loads — the SRD shared-endpoint refactor (#1944) makes them redundant up to high concurrency on 1k/1k traffic. Treat them as emergency switches for CQ-overflow or long-running drift symptoms.
+
+> **`MC_EFA_CQ_THREADS`** — caps the number of CQ polling threads spawned by the EFA transport. Default is `1`, which reaches 99.93% of peak GPU-to-GPU throughput while saving CPU for other workloads. Set to `0` to disable the cap (one poller per EFA context — the legacy behavior). Higher values (e.g., `MC_EFA_CQ_THREADS=4`) are available as an escape hatch for throughput tuning but rarely help in practice.
+>
+> ```bash
+> export MC_EFA_CQ_THREADS=1   # default: single CQ poller (recommended)
+> export MC_EFA_CQ_THREADS=0   # disable cap: one poller per EFA context (legacy)
+> ```
+>
+> If the value exceeds the number of EFA contexts, it is safely ignored (no excess threads are created).
+
+### 3. Prefill Instance
 
 ```bash
-export NVSHMEM_REMOTE_TRANSPORT=libfabric
-export NVSHMEM_LIBFABRIC_PROVIDER=efa
+MOONCAKE_PROTOCOL=efa \
+sglang serve <model_path> \
+  --trust-remote-code \
+  --tp 8 --dp 2 --enable-dp-attention --enable-dp-lm-head \
+  --host 0.0.0.0 --port 8010 \
+  --disaggregation-mode prefill \
+  --disaggregation-transfer-backend mooncake \
+  --disaggregation-bootstrap-port 8998
 ```
 
-> **Warning:** Do **not** set NVSHMEM variables on single-node deployments — doing so causes segmentation faults.
-
-### 3. Docker Launch Example
+### 4. Decode Instance
 
 ```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
+MOONCAKE_PROTOCOL=efa \
+sglang serve <model_path> \
+  --trust-remote-code \
+  --tp 8 --dp 2 --enable-dp-attention --enable-dp-lm-head \
+  --host 0.0.0.0 --port 8020 \
+  --disaggregation-mode decode \
+  --disaggregation-transfer-backend mooncake \
+  --disaggregation-bootstrap-port 8998
 ```
 
-> **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)).
+### 5. Router
+
+Front the pair with `sglang_router` from the prefill host. **`PREFILL_HOST` must be the prefill node's reachable IP, not `127.0.0.1`** — the router forwards this address to the decode node for the bootstrap_room handshake; localhost will fail with `Connection refused` and stall traffic at 0/N.
+
+```bash
+python3 -m sglang_router.launch_router \
+  --pd-disaggregation \
+  --prefill "http://<prefill_ip>:8010" 8998 \
+  --decode  "http://<decode_ip>:8020" \
+  --policy round_robin \
+  --host 0.0.0.0 --port 8000
+```
+
+The trailing `8998` after `--prefill` must match the prefill's `--disaggregation-bootstrap-port`.
 
 ## Technical Details
 
 ### Why libfabric instead of ibverbs?
 
-AWS EFA exposes RDMA-like devices through the ibverbs interface, but does not support the full ibverbs API. Specifically:
-- Queue Pair (QP) creation fails with "Operation not supported" (error 95)
-- EFA requires using libfabric's `FI_EP_RDM` (Reliable Datagram Message) endpoint type
+AWS EFA exposes an RDMA-capable device through the ibverbs interface, but it does **not** implement the full ibverbs API. In particular, EFA only supports **SRD** (Scalable Reliable Datagram) and **UD** (Unreliable Datagram) queue pairs — it does **not** support the **RC** (Reliable Connection) queue pairs that Mooncake's RDMA (`rdma`) transport is built on. Attempting to create an RC QP on an EFA device fails (`EOPNOTSUPP`), and SRD has no one-sided RC-style `ibv_post_send(RDMA_WRITE)` verb in the public ibverbs API.
+
+The portable way to drive EFA's SRD transport is libfabric, whose EFA provider exposes SRD through the `FI_EP_RDM` (Reliable Datagram Message) endpoint type and implements `fi_write` / `fi_read` (one-sided RMA) on top of it. Mooncake's EFA transport therefore targets libfabric directly rather than ibverbs.
 
 ### EFA Transport Architecture
 
+Under the SRD shared-endpoint model every peer is addressed through one `fid_ep` per local NIC — peers are AV-slot entries, not separate endpoints.
+
 ```
-┌─────────────────────────────────────────────────────┐
-│                   EfaTransport                       │
-├─────────────────────────────────────────────────────┤
-│  EfaContext (per device)                            │
-│  ├── fi_info      (fabric info)                     │
-│  ├── fid_fabric   (fabric handle)                   │
-│  ├── fid_domain   (protection domain)               │
-│  ├── fid_av       (address vector for peer lookup)  │
-│  ├── fid_cq       (completion queues)               │
-│  └── fid_mr       (memory regions)                  │
-├─────────────────────────────────────────────────────┤
-│  EfaEndpoint (per connection)                       │
-│  ├── fid_ep       (RDM endpoint)                    │
-│  ├── fi_addr_t    (peer address)                    │
-│  └── local_addr   (local endpoint address)          │
-└─────────────────────────────────────────────────────┘
+┌───────────────────────────────────────────────────────────┐
+│                   EfaTransport                            │
+├───────────────────────────────────────────────────────────┤
+│  EfaContext (per local NIC)                               │
+│  ├── fid_fabric     (fabric handle)                       │
+│  ├── fid_domain     (protection domain)                   │
+│  ├── fid_av         (address vector — one slot per peer)  │
+│  ├── fid_cq         (completion queues)                   │
+│  ├── fid_mr         (memory regions)                      │
+│  ├── shared_ep_     (the single fid_ep that serves every  │
+│  │                   peer via fi_addr_t lookup in the AV) │
+│  └── peer_map_      (full "host:port@nic" path ->         │
+│                      EfaEndPoint; the RPC port is NOT     │
+│                      stripped — see note below)           │
+├───────────────────────────────────────────────────────────┤
+│  EfaEndPoint (per peer)                                   │
+│  └── peer_fi_addr_  (AV slot index for this peer; sends   │
+│                      route through the owning context's    │
+│                      shared_ep_ with this fi_addr_t)      │
+└───────────────────────────────────────────────────────────┘
 ```
 
+> **Peer-map keying.** `peer_map_` is keyed by the **full** `host:port@nic` path, *not* a port-stripped form. Under SGLang DP > 1 each DP worker on a peer host is a separate process with its own Mooncake `TransferEngine` and its own P2PHANDSHAKE RPC port; they share host + NIC but have distinct EFA addresses. Normalizing the port away would collapse every DP worker on that host onto one `EfaEndPoint`, so each arriving handshake would look like a "peer reconnected" to the previous holder and trigger `fi_av_remove` + `fi_av_insert` churn on every KV transfer. Keeping the port in the key costs nothing in steady state (the port is stable for a worker's lifetime).
+
 ### Thread Safety
 
-The EFA transport requests `FI_THREAD_SAFE` from the libfabric provider and adds per-endpoint spinlocks to serialize `fi_write`/`fi_read` calls. This is necessary because:
+The EFA transport requests `FI_THREAD_SAFE` at the domain level and guards the shared endpoint with a single `post_lock_` spinlock (one per `EfaContext`, i.e. one per local NIC) 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`/`fi_read` without serialization corrupts provider internals, causing completions to silently vanish
+- Multiple submission threads may route slices through the same shared endpoint concurrently.
+- libfabric's EFA RDM endpoints are not thread-safe for concurrent `fi_write`/`fi_read` even under `FI_THREAD_SAFE` at the domain level — concurrent posts corrupt provider internals and completions silently vanish.
 
-CQ completion queues are polled by dedicated worker threads (one per EFA device) that run independently of submission threads.
+CQ completion queues are polled by dedicated worker threads that run independently of submission threads. The poller count is `min(MC_EFA_CQ_THREADS, num_EFA_devices)`; `MC_EFA_CQ_THREADS` defaults to `1`, so a single poller round-robins every context's CQ (which already reaches ~99.9% of peak — see the SGLang env-var note above). Set `MC_EFA_CQ_THREADS=0` to lift the cap and spawn one poller per EFA device (the legacy behavior).
 
 ### EFA vs RoCE RDMA
 
 | Feature | EFA (libfabric SRD) | RoCE (ibverbs) |
 |---------|--------------------|--------------------|
-| Protocol | Scalable Reliable Datagram | RDMA over Converged Ethernet |
-| Endpoint type | `FI_EP_RDM` (message-based) | Queue Pairs (true RDMA) |
-| Write operation | Software-emulated via messages + ACKs | Hardware-offloaded one-sided RDMA |
-| CPU overhead | Moderate (provider processes ACKs) | Minimal (NIC handles everything) |
-| Throughput CPU-to-CPU (8×400G) | 222 GB/s (tuned) | ~190 GB/s |
-| Throughput GPU-to-GPU (16×200G) | 347 GB/s (tuned) | N/A |
-| Throughput GPU-to-GPU (8×400G) | 313 GB/s (tuned) | N/A |
+| Protocol | Scalable Reliable Datagram (SRD) | RDMA over Converged Ethernet |
+| QP type | SRD / UD (no RC) | RC (Reliable Connection) |
+| Endpoint type | `FI_EP_RDM` (connectionless) | Queue Pairs (connection-oriented) |
+| Reliability / ordering | Reliable delivery, **unordered** (SRD sprays across paths) | Reliable, in-order |
+| Write operation | One-sided `fi_write` over SRD; device-RDMA-offloaded by default (`FI_EFA_USE_DEVICE_RDMA=1`), falls back to libfabric's emulated RMA only if device RDMA is disabled | Hardware-offloaded one-sided RDMA |
+| Throughput GPU-to-GPU (16×200G, p5en) | 365 GB/s (tuned) | N/A |
+| Throughput CPU-to-CPU (16×200G, p5en) | 213 GB/s (tuned) | — |
 | AWS availability | All EFA-enabled instances | Not available on AWS |
 
+> Mooncake requests libfabric API ≥ 1.18 at `fi_getinfo`, which makes `FI_EFA_USE_DEVICE_RDMA=1` the default on every supported EFA generation (p5/p5e included). On this path `fi_write` / `fi_read` are hardware-offloaded one-sided RMA over SRD — the host CPU is not in the data path. The software-emulated RMA path only applies if you explicitly set `FI_EFA_USE_DEVICE_RDMA=0`.
+
 ### Supported AWS Instance Types
 
 - p6-b300.48xlarge (16 EFA devices × 400 Gbps = 6,400 Gbps, `rdmap*` naming)
diff --git a/docs/source/design/transfer-engine/index.md b/docs/source/design/transfer-engine/index.md
index f5fb3875..1a9ddd84 100644
--- a/docs/source/design/transfer-engine/index.md
+++ b/docs/source/design/transfer-engine/index.md
@@ -13,6 +13,7 @@ As shown in the diagram, each specific client corresponds to a `TransferEngine`,
 
 Mooncake Transfer Engine provides interfaces through the `TransferEngine` class (located in `mooncake-transfer-engine/include/transfer_engine.h`), where the specific data transfer functions for different backends are implemented by the `Transport` class, currently supporting `TcpTransport`, `RdmaTransport`, `EfaTransport`, `NVMeoFTransport`, `NvlinkTransport`, `IntraNodeNvlinkTransport`, and `HipTransport`.
 
+(segment)=
 ### Segment
 Segment represents a collection of source address ranges and target address ranges available during the data transfer process in Transfer Engine. That is, all local and remote addresses involved in `BatchTransfer` requests must be within the valid segment range. Transfer Engine supports the following two types of Segments.
 
@@ -28,6 +29,7 @@ In addition, Transfer Engine also supports registering some **local DRAM areas**
 #### 2. NVMeof Segment
 Transfer Engine also leverages the NVMeof protocol to support direct data transfer from files on NVMe to DRAM/VRAM via PCIe, without going through the CPU and achieving zero-copy. Users need to follow the instructions to mount remote storage nodes locally and use the `openSegment` interface for reference to complete data read/write operations.
 
+(batchtransfer)=
 ### BatchTransfer
 
 With the help of Transfer Engine, Mooncake Store can achieve local DRAM/VRAM reading and writing of specified parts in valid segments  through TCP, (GPUDirect) RDMA, NVMe-of protocols, etc.
@@ -75,6 +77,8 @@ If a connection fails due to link errors, it is removed from the
 endpoint pools on both sides and re-established during the
 next data transfer attempt.
 
+Evicted and deleted endpoints are moved to an internal `waiting_list_` and reclaimed asynchronously once their outstanding slices drain. Reclaim runs on every new endpoint insertion, and additionally on a ~1 Hz heartbeat from the per-context `monitorWorker`, so the waiting list drains even under failure load where new insertions stall while evictions continue.
+
 ### Fault Handling
 In a multi-NIC environment, one common failure scenario is the temporary unavailability of a specific NIC, while other routes may still connect two nodes.
 Mooncake Store is designed to adeptly manage such temporary
@@ -146,7 +150,7 @@ After successfully compiling Transfer Engine, the test program `transfer_engine_
    The initiator node can also configure the following test parameters: `--operation` (can be `"read"` or `"write"`), `batch_size`, `block_size`, `duration`, `threads`, etc.
 
 > [!NOTE]
-> If an exception occurs during execution, it is usually due to incorrect parameter settings. It is recommended to refer to the [troubleshooting document](troubleshooting.md) for preliminary troubleshooting.
+> If an exception occurs during execution, it is usually due to incorrect parameter settings. It is recommended to refer to the [troubleshooting document](../../troubleshooting/troubleshooting.md) for preliminary troubleshooting.
 
 ### Sample Run
 
@@ -165,7 +169,7 @@ Transfer Engine provides batch-based read/write transfers between segments (DRAM
 ### Multi-Transport Management
 
 The `TransferEngine` class internally manages multiple backend `Transport` classes.
-And it will discover the toplogy between CPU/CUDA and RDMA devices automatically
+And it will discover the topology between CPU/CUDA and RDMA devices automatically
 (more device types are working in progress, feedbacks are welcome when the automatic discovery mechanism is not accurate),
 and it will install `Transport` automatically based on the topology.
 
@@ -250,7 +254,7 @@ The HTTP server should implement three following RESTful APIs, while the metadat
 2. `PUT /metadata?key=$KEY`: Update the metadata corresponding to `$KEY` to the value of the request body.
 3. `DELETE /metadata?key=$KEY`: Delete the metadata corresponding to `$KEY`.
 
-For specific implementation, refer to the demo service implemented in Golang at [mooncake-transfer-engine/example/http-metadata-server](../../../mooncake-transfer-engine/example/http-metadata-server).
+For specific implementation, refer to the demo service implemented in Golang at [mooncake-transfer-engine/example/http-metadata-server](gh-dir:mooncake-transfer-engine/example/http-metadata-server).
 
 ## Using Transfer Engine to Your Projects
 
@@ -263,7 +267,7 @@ To support the operational needs of P2P Store, Transfer Engine provides a Golang
 When compiling the project, enable the `-DWITH_P2P_STORE=ON` option to compile the P2P Store example program at the same time.
 
 ### Using Rust Interface
-Under `mooncake-transfer-engine/rust`, the Rust interface implementation of TransferEngine is provided, and a Rust version of the benchmark is implemented based on the interface, similar to [transfer_engine_bench.cpp](../../../mooncake-transfer-engine/example/transfer_engine_bench.cpp). To compile the rust example, you need to install the Rust SDK and add `-DWITH_RUST_EXAMPLE=ON` in the cmake command.
+Under `mooncake-transfer-engine/rust`, the Rust interface implementation of TransferEngine is provided, and a Rust version of the benchmark is implemented based on the interface, similar to [transfer_engine_bench.cpp](gh-file:mooncake-transfer-engine/example/transfer_engine_bench.cpp). To compile the rust example, you need to install the Rust SDK and add `-DWITH_RUST_EXAMPLE=ON` in the cmake command.
 
 ## Advanced Runtime Options
 For advanced users, TransferEngine provides the following advanced runtime options, all of which can be passed in through **environment variables**.
@@ -272,8 +276,12 @@ For advanced users, TransferEngine provides the following advanced runtime optio
 - `MC_NUM_COMP_CHANNELS_PER_CTX` The number of Completion Channel created per device instance, default value 1
 - `MC_IB_PORT` The IB port number used per device instance, default value 1
 - `MC_IB_TC` Adjust RDMA NIC Traffic Class when switch/NIC defaults differ or for traffic planning. Default value -1
+- `MC_IB_SL` Set the InfiniBand Service Level (0-15) of RDMA QPs. The switch maps SL to a Virtual Lane for QoS isolation, e.g. to steer KV-cache traffic into a different VL than Expert-Parallel all-to-all traffic that shares the same NIC. -1 keeps the default (0). Default value -1
 - `MC_IB_PCI_RELAXED_ORDERING` Setting the PCIe ordering to relaxed for the network adapter sometimes results in better performance. Can set 1 to enable RO function. Default value 0
+- `MC_MLX5_QP_UDP_SPORTS` Comma-separated list of UDP source ports (0-65535) used to override the RoCEv2 UDP source port of each QP, for spreading traffic across different ECMP/LAG paths. QP at index *i* uses `list[i % size]`. Default empty (driver chooses). **Requires** an mlx5 NIC + RoCEv2, and the binary built with `-DUSE_MLX5DV=ON`. Recommend ports in the dynamic range 49152-65535. Example: `MC_MLX5_QP_UDP_SPORTS="49152,49153,49154,49155"`
+- `MC_MLX5_QP_LAG_PORT_BALANCE` Set to `1` or `true` to enable automatic LAG port balancing across bonded physical ports. QP at index *i* is pinned to port `(i % num_lag_ports) + 1`; the number of LAG ports is queried from hardware via `mlx5dv_query_device` at startup and printed in the device log. If the device is not in LAG mode the setting is a no-op. Default: disabled. **Requires** the binary built with `-DUSE_MLX5DV=ON`. Example: `MC_MLX5_QP_LAG_PORT_BALANCE=1`
 - `MC_GID_INDEX` The GID index used per device instance, default value 3 (or the maximum value supported by the platform)
+- `MC_PKEY_INDEX` The QP `pkey_index` (partition key table index) used when transitioning the QP to the INIT state. Valid range: 0 to 65535. Default value 0. Set this when the partition key required for your fabric is not at index 0 of the HCA's pkey table
 - `MC_MAX_CQE_PER_CTX` The CQ buffer size per device instance, default value 4096
 - `MC_MAX_EP_PER_CTX` The maximum number of active EndPoint per device instance, default value 65536. **Note:** For versions prior to 0.3.7.post1, the default value is 256, and it cannot be manually set to 65536. The maximum supported value is 65535!
 - `MC_NUM_QP_PER_EP` The number of QPs per EndPoint, the more the number, the better the fine-grained I/O performance, default value 2
@@ -284,9 +292,11 @@ For advanced users, TransferEngine provides the following advanced runtime optio
 - `MC_WORKERS_PER_CTX` The number of asynchronous worker threads corresponding to each device instance
 - `MC_SLICE_SIZE` The segmentation granularity of user requests in Transfer Engine
 - `MC_RETRY_CNT` The maximum number of retries in Transfer Engine
+- `MC_AUTO_GID_MAX_RETRIES` The maximum number of automatic local GID reprobe retries during classic RDMA handshake recovery. Default value 2. Set to 0 to disable automatic GID retry.
 - `MC_LOG_LEVEL` This option can be set as `TRACE`/`INFO`/`WARNING`/`ERROR` (see [glog doc](https://github.com/google/glog/blob/master/docs/logging.md)), and more detailed logs will be output during runtime
 - `MC_DISABLE_METACACHE` Disable local meta cache to prevent transfer failure due to dynamic memory registrations, which may downgrades the performance
 - `MC_HANDSHAKE_LISTEN_BACKLOG` The backlog size of socket listening for handshaking, default value is 128
+- `MC_HANDSHAKE_CONNECT_TIMEOUT` Connect timeout in seconds for outbound handshake-port requests (QP handshake, probe, notify, metadata exchange), default value is 5. Bounds the stall when the peer address is unreachable; without it, a connect to an unroutable address (e.g. a removed node) blocks for the kernel's full TCP SYN retry cycle, which can take minutes
 - `MC_HANDSHAKE_MAX_LENGTH` The maximum handshake message length in bytes for P2P mode. Valid range: 1MB to 128MB. Default value is 1MB (1048576 bytes). Increase this value when using a single RDMA instance with many registered memory buffers (>10,000) to avoid handshake failures. Example: set to 10485760 for 10MB
 - `MC_LOG_DIR` Specify the directory path for log redirection files. If invalid, log to stderr instead.
 - `MC_REDIS_PASSWORD` The password for Redis storage plugin, only takes effect when Redis is specified as the metadata server. If not set, no authentication will be attempted to log in to the Redis.
@@ -298,19 +308,26 @@ For advanced users, TransferEngine provides the following advanced runtime optio
 - `MC_FORCE_MNNVL` Force to use Multi-Node NVLink as the active transport regardless whether RDMA devices are installed.
 - `MC_INTRA_NVLINK` Enable intra-node NVLINK transport, and cannot be used together with MC_FORCE_MNNVL.
 - `MC_FORCE_TCP` Force to use TCP as the active transport regardless whether RDMA devices are installed.
-- `MC_MIN_PRC_PORT` Specifies the minimum port number for RPC service. The default value is 15000.
-- `MC_MAX_PRC_PORT` Specifies the maximum port number for RPC service. The default value is 17000.
+- `MC_MIN_RPC_PORT` Specifies the minimum port number for RPC service. The default value is 15000.
+- `MC_MAX_RPC_PORT` Specifies the maximum port number for RPC service. The default value is 17000.
 - `MC_PATH_ROUNDROBIN` Use round-robin mode in the RDMA path selection. This may be beneficial for transferring large bulks.
+- `MC_TE_FILTERS` Optional comma-separated whitelist of IB device names (e.g. `mlx5_0,mlx5_2`) for legacy Transfer Engine topology discovery. When unset, all available devices are discovered.
+- `WITH_NVIDIA_PEERMEM` When set to `1`, `ON`, or `TRUE`, Mooncake uses `ibv_reg_mr()` directly for GPU memory registration (requires the `nvidia-peermem` kernel module). By default (unset or `0`), Mooncake uses the DMA-BUF path which does not require `nvidia-peermem`.
 - `MC_ENDPOINT_STORE_TYPE` Choose FIFO Endpoint Store (`FIFO`) or Sieve Endpoint Store (`SIEVE`), default is `SIEVE`.
-- `MC_TCP_ENABLE_CONNECTION_POOL` Enable TCP Connection Pool to avoid excessive sockets. 
+- `MC_TCP_ENABLE_CONNECTION_POOL` Enable TCP Connection Pool to avoid excessive sockets.
+- `MC_TCP_SLICE_SIZE` The segmentation granularity (in bytes) of TCP transport for splitting large transfers into socket read/write operations. Corresponds to `MC_SLICE_SIZE` for RDMA. Default value 65536 (64KB).
 
 ## C++ API Reference
 
-::::{toctree}
+For the complete C++ API reference, see [Transfer Engine C++ API](../../api-reference/cpp/index).
+
+## Supported Protocols
+
+:::{toctree}
 :maxdepth: 1
 
-cpp-api
-::::
+../../getting_started/supported-protocols
+:::
 
 ## EFA Transport (AWS)
 
@@ -330,6 +347,15 @@ ascend_transport
 heterogeneous_ascend
 :::
 
+## Sunrise Link Transport Component
+
+:::{toctree}
+:maxdepth: 1
+
+kunpeng_ub_transport
+sunrise_link_transport
+:::
+
 ## Benchmark and Tuning Guide
 
 :::{toctree}
diff --git a/docs/source/design/transfer-engine/kunpeng_ub_transport.md b/docs/source/design/transfer-engine/kunpeng_ub_transport.md
index ead3a6c8..a751b997 100644
--- a/docs/source/design/transfer-engine/kunpeng_ub_transport.md
+++ b/docs/source/design/transfer-engine/kunpeng_ub_transport.md
@@ -50,7 +50,7 @@ sudo apt-get install -y \
     libjsoncpp-dev \
     libnuma-dev \
     libibverbs-dev \
-    libboost-all-dev \
+    libboost-dev \
     libcurl4-openssl-dev \
     libgtest-dev \
     libmsgpack-dev \
diff --git a/docs/source/design/transfer-engine/sunrise_link_transport.md b/docs/source/design/transfer-engine/sunrise_link_transport.md
new file mode 100644
index 00000000..eaf8c446
--- /dev/null
+++ b/docs/source/design/transfer-engine/sunrise_link_transport.md
@@ -0,0 +1,129 @@
+# Sunrise Link Transport
+
+## Overview
+Sunrise Link is a GPU-to-GPU direct communication interconnect developed by Sunrise. In Mooncake, it is integrated as a GPU transport backend in the TENT transfer framework. It relies on Tang Runtime (`tangrt`) for device memory allocation, pointer attribute queries, peer copy, and IPC handle operations, and is registered as `SUNRISE_LINK` in TENT.
+
+At runtime, `TransferEngineImpl` loads Sunrise Link Transport when `USE_SUNRISE` is enabled at build time and `transports/sunrise_link/enable=true` is set in config.
+
+### Transfer Pipeline
+
+1. **Init**: `dlopen` `libptml_shared.so` + `libtangrt_shared.so`, initialize PTML.
+2. **Topology**: `ptmlPtlinkPhytopoDetect` discovers inter-chip ports, builds C2C mapping (`local_chipid, remote_chipid → local_port`).
+3. **Registration**: `tangIpcGetMemHandle` generates IPC handles for GPU buffers.
+4. **Execution**: Same device → `tangMemcpy`; Cross-device → `tangMemcpyPeer` / `tangMemcpyPeer_v2`, fallback to `tangDeviceGetPeerPointer` + C2C; Remote → `tangIpcOpenMemHandle` + peer copy; Host ↔ Device → `tangMemcpy` with direction.
+5. **Async**: When size exceeds `async_memcpy_threshold`, uses `tangMemcpyAsync` with per-device stream.
+
+---
+
+## Additional Dependencies
+In addition to Mooncake's base dependencies, Sunrise Link Transport requires Tang Runtime:
+
+- **Header path**: `/usr/local/tangrt/include`
+- **Library path**: `/usr/local/tangrt/lib/linux-x86_64`
+- **Shared library**: `libtangrt_shared.so` (loaded from the default installation path at runtime)
+
+Make sure runtime libraries related to `tangrt` and `ptml` can be found by the dynamic linker (for example, via `LD_LIBRARY_PATH` or system library search paths).
+
+---
+
+## Build and Compile
+
+**Prerequisites**
+
+- Tang Runtime is installed and available (default path: `/usr/local/tangrt`)
+- Build environment can access Mooncake and its base dependencies
+
+**CMake Configuration**
+
+```bash
+# Clone Mooncake
+git clone https://github.com/kvcache-ai/Mooncake.git
+cd Mooncake
+
+# Enable TENT + Sunrise Link
+mkdir build && cd build
+cmake .. -DUSE_TENT=ON -DUSE_SUNRISE=ON
+
+# Build
+make -j$(nproc)
+```
+
+If Tang Runtime is installed at a non-default location:
+
+```bash
+cmake .. -DUSE_TENT=ON -DUSE_SUNRISE=ON -DMC_TANGRT_ROOT=/opt/tangrt
+```
+
+---
+
+## Run and Test
+
+`transfer_engine_bench` supports the `sunrise_link` protocol for basic connectivity and performance validation.
+
+```bash
+# Terminal 1: target
+./transfer_engine_bench \
+  --mode=target \
+  --protocol=sunrise_link \
+  --local_server_name=10.0.0.2 \
+  --metadata_server=P2PHANDSHAKE \
+  --gpu_id=0
+
+# Terminal 2: initiator
+./transfer_engine_bench \
+  --mode=initiator \
+  --protocol=sunrise_link \
+  --metadata_server=P2PHANDSHAKE \
+  --segment_id=10.0.0.2:$PORT \
+  --gpu_id=0 \
+  --block_size=8388608 \
+  --batch_size=32
+```
+
+> Note: when `metadata_server=P2PHANDSHAKE`, the target node may listen on a dynamically assigned port. Replace `$PORT` in `--segment_id` with the actual port printed in target logs.
+
+### Using `tebench` (TENT Backend)
+
+Role is determined by `--target_seg_name`: empty → target, otherwise → initiator.
+
+```bash
+# Terminal 1: target
+./tebench --backend=tent --metadata_type=p2p --xport_type=sunrise_link \
+  --seg_type=VRAM --local_gpu_id=0
+
+# Terminal 2: initiator
+./tebench --backend=tent --metadata_type=p2p --xport_type=sunrise_link \
+  --seg_type=VRAM --target_seg_name=192.168.172.52:168 \
+  --local_gpu_id=1 --target_gpu_id=0
+```
+
+---
+
+## Environment Variables
+
+| Variable | Description | Default |
+|----------|-------------|---------|
+| `MC_TANGRT_ROOT` | Tang Runtime root directory | `/usr/local/tangrt` |
+| `MC_TANGRT_LIB_DIR` | Direct path to `.so` dir. Overrides `MC_TANGRT_ROOT` + arch lookup. | (unset) |
+| `MC_TANGRT_LIB_ARCH` | Architecture subdirectory under `<root>/lib/` | `linux-x86_64` |
+
+---
+
+## Configuration Options (TENT)
+
+Sunrise Link behavior can be tuned through config file options:
+
+- `transports/sunrise_link/enable`: enable or disable the transport (default: `true`)
+- `transports/sunrise_link/async_memcpy_threshold`: threshold for async memcpy in MiB, 0 to disable (default: `0`)
+- `transports/sunrise_link/enable_port_striping`: enable multi-port striping for large cross-device copies >4 MiB (default: `false`)
+
+Tune these options based on workload characteristics and hardware topology.
+
+---
+
+## Important Notes
+
+1. **Device Context**: Internally manages `tangSetDevice`. When calling Tang API directly, set the correct device before querying pointer attributes.
+2. **IPC Handle Lifecycle**: IPC handles are cached and released in `uninstall()`.
+3. **Concurrent Peer Copy**: Per-device mutexes protect context-bound ops. Cross-device copies acquire source/destination mutexes in fixed order (lower device ID first) to prevent deadlocks.
+4. **Mapped Host Pointer Normalization**: `tangHostAlloc` host pointers with corresponding device pointers are auto-normalized for correct peer copy path selection.
diff --git a/docs/source/design/unified-parallel-tensor-io.md b/docs/source/design/unified-parallel-tensor-io.md
new file mode 100644
index 00000000..ff9f1f3b
--- /dev/null
+++ b/docs/source/design/unified-parallel-tensor-io.md
@@ -0,0 +1,405 @@
+# Unified Parallel Tensor IO
+
+## Goal
+
+This document defines the single source of truth for Mooncake's next-generation tensor IO API. The purpose is to keep implementation aligned around one explicit design and avoid drifting into ad hoc TP/EP/DP/PP-specific interfaces.
+
+The target outcome is:
+
+- one unified write API family
+- one unified read API family
+- one unified upsert API family
+- TP-specific APIs retained only as compatibility wrappers
+- minimal caller-facing inputs, with layout and planning details derived internally
+- one stable abstraction that can cover TP / DP / EP / PP combinations without exploding the public API surface
+
+## Design principles
+
+1. **One API family, not one family per parallelism type.**
+   Public APIs should not branch into separate long-term `*_with_tp`, `*_with_ep`, `*_with_pp`, and `*_with_dp` families.
+2. **Parallelism is modeled as explicit axes.**
+   `mixed` must not be a vague enum value. A shard should be described by the exact axis coordinates that identify it.
+3. **Reads must encode caller intent explicitly.**
+   A read request must say whether the caller wants the stored view, a target shard/view, or the reconstructed full tensor.
+4. **Only require true caller intent.**
+   If Mooncake can derive a field from the input tensor, stored metadata, or planning context, it should not be required in the public API.
+5. **Planner/runtime complexity stays internal.**
+   Byte ranges, payload offsets, reconstruction plans, and lowering to `get_into_ranges(...)` are runtime concerns, not public API concerns.
+6. **Compatibility wrappers stay thin.**
+   Existing TP methods should lower to the same unified implementation path.
+
+## Core model
+
+### Parallelism is an axis list
+
+A tensor object may be identified by one axis or by multiple orthogonal axes.
+
+Examples:
+
+- pure TP shard: `[TP]`
+- DP + TP shard: `[DP, TP]`
+- PP + TP shard: `[PP, TP]`
+- DP + PP + TP shard: `[DP, PP, TP]`
+- DP + PP + EP + TP shard: `[DP, PP, EP, TP]`
+
+So the design should not use:
+
+```python
+kind = "mixed"
+```
+
+Instead it should use:
+
+```python
+class ParallelAxis:
+    kind: Literal["tp", "dp", "ep", "pp"]
+    rank: int
+    size: int
+    split_dim: Optional[int] = None
+    expert_id: Optional[int] = None
+    stage_id: Optional[int] = None
+```
+
+```python
+class TensorParallelism:
+    axes: list[ParallelAxis]
+```
+
+Single-axis cases are just special cases of this model:
+
+- TP only: `axes=[TP(...)]`
+- EP only: `axes=[EP(...)]`
+- PP + TP: `axes=[PP(...), TP(...)]`
+
+## Public API shape
+
+### Write / upsert side
+
+```python
+def put_tensor_with_parallelism(
+    self,
+    key: str,
+    tensor: torch.Tensor,
+    parallelism: TensorParallelism | None = None,
+    replica: ReplicateConfig | None = None,
+) -> int
+```
+
+```python
+def batch_put_tensor_with_parallelism(
+    self,
+    keys: list[str],
+    tensors: list[torch.Tensor],
+    parallelisms: list[TensorParallelism | None] | None = None,
+    replica: ReplicateConfig | None = None,
+) -> list[int]
+```
+
+```python
+def upsert_tensor_with_parallelism(
+    self,
+    key: str,
+    tensor: torch.Tensor,
+    parallelism: TensorParallelism | None = None,
+    replica: ReplicateConfig | None = None,
+) -> int
+```
+
+```python
+def batch_upsert_tensor_with_parallelism(
+    self,
+    keys: list[str],
+    tensors: list[torch.Tensor],
+    parallelisms: list[TensorParallelism | None] | None = None,
+    replica: ReplicateConfig | None = None,
+) -> list[int]
+```
+
+### Read side
+
+Read-side requests need one extra structure because the request must encode not only the target parallel coordinates, but also the materialization mode.
+
+```python
+class ReadTarget:
+    mode: Literal["as_stored", "shard", "full"]
+    parallelism: TensorParallelism | None = None
+```
+
+```python
+def get_tensor_with_parallelism(
+    self,
+    key: str,
+    target: ReadTarget | None = None,
+) -> torch.Tensor
+```
+
+```python
+def batch_get_tensor_with_parallelism(
+    self,
+    keys: list[str],
+    targets: list[ReadTarget | None] | None = None,
+) -> list[torch.Tensor]
+```
+
+Optional zero-copy forms follow the same model:
+
+```python
+def get_tensor_with_parallelism_into(
+    self,
+    key: str,
+    buffer_ptr: int,
+    size: int,
+    target: ReadTarget | None = None,
+) -> torch.Tensor
+```
+
+```python
+def batch_get_tensor_with_parallelism_into(
+    self,
+    keys: list[str],
+    buffer_ptrs: list[int],
+    sizes: list[int],
+    targets: list[ReadTarget | None] | None = None,
+) -> list[torch.Tensor]
+```
+
+## Caller inputs vs internal derived fields
+
+### Caller must provide
+
+These are true caller intent and belong in the public API:
+
+- the tensor itself
+- whether the write target is full or shard-like
+- the shard identity axes when writing a shard object
+- the read target mode: `as_stored`, `shard`, or `full`
+- the target axis coordinates when requesting a target shard/view
+- `ReplicateConfig` when the caller wants replication / publish-like behavior
+
+### Mooncake should derive internally
+
+These should not be mandatory public inputs when they are derivable:
+
+- logical shape
+- local shard shape
+- metadata encoding/version details
+- payload offsets
+- source and destination byte ranges
+- reconstruction plans
+- lowering to `get_into_ranges(...)`
+
+## Why `ReadTarget` is required
+
+`parallelism` alone is not enough to describe a read.
+
+For example, if a tensor is stored as TP shards, a request that carries TP axis coordinates is ambiguous unless it also says whether the caller wants:
+
+- the stored local shard
+- a target shard/view
+- the reconstructed full tensor
+
+So reads must explicitly encode:
+
+```python
+ReadTarget(mode="as_stored")
+ReadTarget(mode="shard", parallelism=...)
+ReadTarget(mode="full")
+```
+
+This distinction must not be guessed from the axis metadata.
+
+## Mixed-parallel scenarios in training and inference
+
+`mixed` is not a mode. It means the shard identity needs more than one axis coordinate to be uniquely described.
+
+### Scenario table
+
+| Scenario | Typical axis list | Meaning |
+|---|---|---|
+| TP training | `[TP]` | one TP slice of a logical tensor |
+| DP + TP training | `[DP, TP]` | one TP slice within one DP replica/group |
+| PP + TP training | `[PP, TP]` | one TP slice owned by one pipeline stage |
+| DP + PP + TP training | `[DP, PP, TP]` | one TP slice in one stage in one DP replica |
+| DP + PP + EP + TP training | `[DP, PP, EP, TP]` | one expert-local TP slice in one stage and one DP replica |
+| TP inference | `[TP]` | one TP slice used by one inference rank |
+| PP + TP inference | `[PP, TP]` | one TP slice owned by one inference pipeline stage |
+| EP inference | `[EP]` or `[EP, TP]` | one expert shard, optionally further TP-sliced |
+| multi-replica serving inference | `[DP, TP]` or `[DP, PP, TP]` | one shard scoped to a serving replica plus model-parallel axes |
+
+### Important distinctions
+
+- TP and EP often affect the tensor's physical layout directly.
+- PP and serving-replica / DP often act more like ownership or scope tags, even when they do not themselves change the byte layout inside the local shard.
+
+## Write-side matrix
+
+The unified write family should be driven by the identity of the object being written, not by method-name proliferation.
+
+| Caller holds | Wants to store | API | Required `parallelism` |
+|---|---|---|---|
+| full tensor | full tensor | `put_tensor_with_parallelism(..., parallelism=None)` | none |
+| full tensor | TP shard | `put_tensor_with_parallelism(...)` | `axes=[TP(rank,size,split_dim)]` |
+| full tensor | DP-scoped shard/replica | `put_tensor_with_parallelism(...)` | `axes=[DP(rank,size)]` plus layout axis if actually sharded |
+| full tensor | EP shard | `put_tensor_with_parallelism(...)` | `axes=[EP(rank,size,expert_id)]` plus `split_dim` if needed |
+| full tensor | PP stage shard | `put_tensor_with_parallelism(...)` | `axes=[PP(rank,size,stage_id)]` |
+| full tensor | combined shard | `put_tensor_with_parallelism(...)` | explicit axis list such as `[PP(...), TP(...)]` |
+| shard tensor | shard object | `put_tensor_with_parallelism(...)` | explicit axis list describing that shard identity |
+
+For TP-containing **multi-axis** layouts, the write semantic is now: the caller may pass the **full source tensor**, and the provided TP rank/layout tells Mooncake which uniform shard to materialize and persist. Callers no longer need to pre-split the tensor themselves for `dp_tp` / `pp_tp` / `ep_tp` style writes.
+
+Single-axis TP compatibility wrappers and the preserved plain-TP `with_parallelism` behavior still accept shard input rather than auto-materializing from a full tensor.
+
+Pure DP still does not invent a split axis on its own. If the stored object is actually sharded, the request must still include the layout axis that defines the shard shape.
+
+The same matrix applies to `upsert_tensor_with_parallelism(...)`.
+
+## Read-side matrix
+
+The unified read family should be driven by `ReadTarget`.
+
+| Stored layout | Caller wants | API | `ReadTarget` |
+|---|---|---|---|
+| full tensor | stored full tensor | `get_tensor_with_parallelism(...)` | `None` or `mode="as_stored"` |
+| TP shard object | stored shard | `get_tensor_with_parallelism(...)` | `mode="as_stored"` |
+| TP shard set | target TP shard | `get_tensor_with_parallelism(...)` | `mode="shard", parallelism=TP(...)` |
+| TP shard set | full tensor | `get_tensor_with_parallelism(...)` | `mode="full"` |
+| EP shard set | target expert shard | `get_tensor_with_parallelism(...)` | `mode="shard", parallelism=EP(...)` |
+| EP shard set | full tensor | `get_tensor_with_parallelism(...)` | `mode="full"` |
+| PP + TP shard set | target PP+TP shard | `get_tensor_with_parallelism(...)` | `mode="shard", parallelism=[PP(...), TP(...)]` |
+| mixed shard set | full tensor | `get_tensor_with_parallelism(...)` | `mode="full"` |
+
+## Source-layout to target-layout matrix
+
+The planner must eventually cover these conversions, but the public API should remain the same across all of them.
+
+| Stored source layout | Requested target layout | Support model |
+|---|---|---|
+| TP | TP | direct shard fetch or shard-local fast path |
+| TP | full | reconstruct full tensor |
+| TP | EP / PP / DP / mixed | planner-driven remap |
+| EP | EP | direct shard fetch or shard-local fast path |
+| EP | full | reconstruct full tensor |
+| EP | TP / PP / DP / mixed | planner-driven remap |
+| PP | PP | direct shard fetch |
+| PP | full | reconstruct full tensor when meaningful |
+| PP | TP / EP / DP / mixed | planner-driven remap |
+| mixed | mixed | direct fetch if exact match, else planner-driven remap |
+| mixed | full | reconstruct full tensor |
+| mixed | TP / EP / PP / DP | planner-driven remap |
+
+The important point is that these combinations must not create public API explosion.
+
+## Compatibility wrappers
+
+Existing TP APIs remain compatibility wrappers.
+
+Conceptually:
+
+```python
+put_tensor_with_tp(key, tensor, tp_rank, tp_size, split_dim)
+```
+
+lowers to:
+
+```python
+put_tensor_with_parallelism(
+    key,
+    tensor,
+    TensorParallelism(axes=[TP(rank=tp_rank, size=tp_size, split_dim=split_dim)]),
+)
+```
+
+and:
+
+```python
+get_tensor_with_tp(key, tp_rank, tp_size, split_dim)
+```
+
+lowers to:
+
+```python
+get_tensor_with_parallelism(
+    key,
+    ReadTarget(
+        mode="shard",
+        parallelism=TensorParallelism(
+            axes=[TP(rank=tp_rank, size=tp_size, split_dim=split_dim)]
+        ),
+    ),
+)
+```
+
+Existing TP behavior stays stable, but implementation should flow through the unified path.
+
+## Runtime lowering direction
+
+Internally, reads should be planner-driven.
+
+When a request can be lowered to explicit:
+
+- key
+- src offset
+- dst offset
+- size
+
+ranges, the runtime should reuse `get_into_ranges(...)` to assemble the result directly into the output buffer.
+
+Otherwise it can fall back to the simpler whole-object path.
+
+This keeps the public interface stable while allowing future planner work to add DP / TP / EP / PP remapping and optimized reconstruction without another API redesign.
+
+## Current implementation status
+
+The current `store_py.cpp` implementation now reflects the main shape of this design:
+
+- unified write APIs are exposed as `put_tensor_with_parallelism(...)` and `batch_put_tensor_with_parallelism(...)`
+- unified read APIs are exposed as `get_tensor_with_parallelism(...)` and `batch_get_tensor_with_parallelism(...)`
+- unified upsert APIs are exposed as `upsert_tensor_with_parallelism(...)` and `batch_upsert_tensor_with_parallelism(...)`
+- zero-copy `_into` and `_from` variants exist for the unified API family
+- TP-specific APIs remain available as compatibility wrappers and should not be treated as the long-term surface area
+
+### Implemented write-side convenience: `writer_partitions`
+
+Batch write and batch upsert paths also support `writer_partitions` as a convenience input for full tensors that should be written as stored shards.
+
+This is intentionally narrower than the full `TensorParallelism` model:
+
+- it is a write-side convenience, not a replacement for `TensorParallelism`
+- it is primarily for batch full-tensor writes where the caller already knows rank / size / split_dim per item
+- it should not change the unified read-side abstraction
+
+`writer_partitions` remains a separate explicit route. The newer TP-containing `parallelism` write semantic now overlaps with it for the common case of “full tensor in, store one requested shard”, but `writer_partitions` is still useful when the caller wants a lighter write-side request shape without constructing `TensorParallelism` objects.
+
+### Implemented read-side behavior
+
+Read-side support includes:
+
+- returning the stored local object (`mode="as_stored"`)
+- returning a target shard (`mode="shard"`)
+- reconstructing the full tensor (`mode="full"`)
+- lowering reconstruction-oriented paths onto existing runtime helpers such as `get_into_ranges(...)` where appropriate
+
+### Compatibility boundary
+
+The project should continue to preserve this boundary:
+
+- keep old TP APIs functional
+- keep them thin
+- do not expand the old TP-specific family as the primary interface
+- document and evolve the unified `*_with_parallelism` family instead
+
+## Scope still intentionally limited
+
+This document describes the stable public API direction, but not every theoretical source-layout to target-layout remap is fully implemented.
+
+In particular, the design should continue to avoid over-promising planner coverage for arbitrary remaps across all DP / TP / EP / PP combinations until those paths are explicitly implemented and tested.
+
+The safe documented contract today is:
+
+1. define the public API structures and signatures clearly
+2. align `store_py.cpp` with `axes + ReadTarget`
+3. keep TP wrappers working by lowering into the unified path
+4. do not introduce a new long-term `kind="mixed"` model
+5. do not push planner internals such as byte ranges and derived shapes into public arguments
+6. do not expand into many parallelism-specific public methods
+7. reuse existing runtime helpers like `get_into_ranges(...)` rather than inventing a parallel reconstruction path from scratch
diff --git a/docs/source/getting_started/build.md b/docs/source/getting_started/build.md
index 168fe1cf..b2d80cee 100644
--- a/docs/source/getting_started/build.md
+++ b/docs/source/getting_started/build.md
@@ -17,10 +17,10 @@ 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**: The CUDA version includes Mooncake-EP and GPU topology detection, requiring CUDA 12.1+. The non-CUDA version is for environments without CUDA dependencies, but it still requires the system runtime libraries used by the transfer stack. On Ubuntu, install them with `sudo apt-get update && sudo apt-get install -y libcurl4 libibverbs1 rdma-core librdmacm1 libnuma1 liburing2`.
 > **Note**: MLU support is currently source-build only. If you need Cambricon MLU memory support, install Neuware and build with `-DUSE_MLU=ON`.
 
-## Automatic
+## Automatic Build
 
 ### Recommended Version
 - OS: Ubuntu 22.04 LTS+
@@ -30,7 +30,7 @@ pip install mooncake-transfer-engine-non-cuda
 ### Steps
 1. Install dependencies, stable Internet connection is required:
    ```bash
-   bash dependencies.sh
+   sudo bash dependencies.sh
    ```
 
 2. In the root directory of this project, run the following commands:
@@ -45,7 +45,25 @@ pip install mooncake-transfer-engine-non-cuda
    sudo make install
    ```
 
-## Manual
+**Build with NVMe-oF SSD Pool**
+
+To enable the NVMe-oF SSD pool, install the SPDK dependencies and build
+Mooncake with `USE_NOF` enabled:
+
+```bash
+sudo bash dependencies.sh --with-spdk
+
+mkdir build
+cd build
+cmake .. -DUSE_NOF=ON
+make -j
+sudo make install
+```
+
+`-DUSE_NOF=ON` builds the NoF registration APIs and deployment tools. Use
+`-DUSE_NOF=OFF` or omit the option when the NVMe-oF SSD pool is not needed.
+
+## Manual Build
 
 ### Recommended Version
 - cmake: 3.22.x
@@ -70,7 +88,7 @@ pip install mooncake-transfer-engine-non-cuda
                        libnuma-dev \
                        libunwind-dev \
                        libpython3-dev \
-                       libboost-all-dev \
+                       libboost-dev \
                        libssl-dev \
                        pybind11-dev \
                        libcurl4-openssl-dev \
@@ -100,12 +118,15 @@ pip install mooncake-transfer-engine-non-cuda
     ```
 
 2. If you want to compile the GPUDirect support module, first follow the instructions in https://docs.nvidia.com/cuda/cuda-installation-guide-linux/ to install CUDA (ensure to enable `nvidia-fs` for proper `cuFile` module compilation). After that:
-    1) Follow Section 3.7 in https://docs.nvidia.com/cuda/gpudirect-rdma/ to install `nvidia-peermem` for enabling GPU-Direct RDMA
-    2) Configure `LIBRARY_PATH` and `LD_LIBRARY_PATH` to ensure linking of `cuFile`, `cudart`, and other libraries during compilation:
+    1) Configure `LIBRARY_PATH` and `LD_LIBRARY_PATH` to ensure linking of `cuFile`, `cudart`, and other libraries during compilation:
     ```bash
     export LIBRARY_PATH=$LIBRARY_PATH:/usr/local/cuda/lib64
     export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/cuda/lib64
     ```
+    ```{admonition} GPU-Direct RDMA
+    :class: note
+    Mooncake could use the DMA-BUF path for GPU-Direct RDMA, which does **not** require the `nvidia-peermem` kernel module. If you prefer the DMA-BUF path, please set the runtime environment variable `WITH_NVIDIA_PEERMEM=0` before starting Mooncake. If you prefer the legacy `ibv_reg_mr` path (which requires `nvidia-peermem`), set the runtime environment variable `WITH_NVIDIA_PEERMEM=1`. See Section 3.7 of https://docs.nvidia.com/cuda/gpudirect-rdma/ for instructions on installing `nvidia-peermem`.
+    ```
 
 3. If you want to compile the Moore Mthreads GPUDirect support module, first follow the instructions in https://docs.mthreads.com/musa-sdk/musa-sdk-doc-online/install_guide to install MUSA. After that:
     1) Install `mthreads-peermem` for enabling GPU-Direct RDMA
@@ -149,7 +170,20 @@ pip install mooncake-transfer-engine-non-cuda
     - `-DMACA_LIB_DIR=/path/to/maca/lib64`
     - `-DMACA_RUNTIME_LIBS="mcruntime;mxc-runtime64;rt"` (semicolon-separated CMake list)
 
-6. Install yalantinglibs
+6. If you want to compile Huawei Ascend NPU support, first install the Ascend CANN Toolkit following the instructions at https://www.hiascend.com/document. After that:
+    1) Source `set_env.sh` in the CANN installation directory to configure the build environment (no need to manually set `ASCEND_HOME_PATH` or other related environment variables).
+    2) Mooncake provides two Ascend NPU transport paths, choose one as needed:
+       - `-DUSE_ASCEND_DIRECT=ON` (**recommended**): Ascend Direct transport based on the ADXL engine. (refer to [Version Compatibility Guide](https://gitcode.com/cann/hixl/wiki/Mooncake%20+%20HIXL%20%E5%BF%AB%E9%80%9F%E4%B8%8A%E6%89%8B%E6%8C%87%E5%8D%97.md) for details).
+       - `-DUSE_UBSHMEM=ON`: Shared memory transport based on CANN VMM APIs (requires CANN >= 9.0.0, driver >= 26.0.0, Lingqu >= 1.5).
+
+    Example for building with Ascend NPU:
+    ```bash
+    source /usr/local/Ascend/cann/set_env.sh
+    cmake .. -DUSE_ASCEND_DIRECT=ON
+    make -j
+    ```
+
+7. Install yalantinglibs
     ```bash
     git clone https://github.com/alibaba/yalantinglibs.git
     cd yalantinglibs
@@ -159,7 +193,7 @@ pip install mooncake-transfer-engine-non-cuda
     make install
     ```
 
-7. In the root directory of this project, run the following commands:
+8. In the root directory of this project, run the following commands:
    ```bash
    mkdir build
    cd build
@@ -167,23 +201,56 @@ pip install mooncake-transfer-engine-non-cuda
    make -j
    ```
 
-8. Install Mooncake python package and mooncake_master executable
+9. Install Mooncake python package and mooncake_master executable
    ```bash
    make install
    ```
 
 ## Use Mooncake in Docker Containers
-Mooncake supports Docker-based deployment. What you need is to get Docker image by `docker pull alogfans/mooncake`.
+Mooncake supports Docker-based deployment. You can either build the image from
+this repository with `docker/mooncake.Dockerfile` or substitute a published
+tag that matches the release you want to run.
 For the container to use the host's network resources, you need to add the `--device` option when starting the container. The following is an example.
 
 ```
 # In host
-sudo docker run --net=host --device=/dev/infiniband/uverbs0 --device=/dev/infiniband/rdma_cm --ulimit memlock=-1 -t -i mooncake:v0.9.0 /bin/bash
+sudo docker build -f docker/mooncake.Dockerfile -t mooncake:from-source .
+sudo docker run --net=host --device=/dev/infiniband/uverbs0 --device=/dev/infiniband/rdma_cm --ulimit memlock=-1 -t -i mooncake:from-source /bin/bash
 # Run transfer engine in container
 cd /Mooncake-main/build/mooncake-transfer-engine/example
 ./transfer_engine_bench --device_name=ibp6s0 --metadata_server=10.1.101.3:2379 --mode=target --local_server_name=10.1.100.3
 ```
 
+For SGLang HiCache deployments inside Docker, reserve HugeTLB pages on the host before starting the container and pass the allocator settings through the container environment:
+
+```bash
+python3 scripts/check_hicache_hugepage_requirements.py \
+  --tp-size 4 \
+  --hicache-size 64gb \
+  --global-segment-size 8gb \
+  --arena-pool-size 56gb \
+  --available-hugetlb 512gb
+
+sudo sysctl -w vm.nr_hugepages=262144
+grep -E 'HugePages_Total|HugePages_Free|Hugepagesize' /proc/meminfo
+
+sudo docker run --gpus all \
+  --net=host \
+  --ipc=host \
+  --ulimit memlock=-1 \
+  --shm-size=128g \
+  --device=/dev/infiniband/uverbs0 \
+  --device=/dev/infiniband/rdma_cm \
+  -e MC_STORE_USE_HUGEPAGE=1 \
+  -e MC_STORE_HUGEPAGE_SIZE=2MB \
+  -e MOONCAKE_GLOBAL_SEGMENT_SIZE=8gb \
+  -e MC_MMAP_ARENA_POOL_SIZE=56gb \
+  -t -i mooncake:from-source /bin/bash
+```
+
+The `64gb` / `56gb` values above are tuned examples for large HiCache deployments, not defaults. The arena remains disabled unless you explicitly enable it, and if you enable it via gflag without an env override the default pool size is `8gb`. On smaller hosts, start with `8gb` or `16gb` and size upward with the helper. When you want the baseline direct-`mmap()` path instead of the arena, set `MC_DISABLE_MMAP_ARENA=1` (also accepts `true`, `yes`, or `on`) and omit `MC_MMAP_ARENA_POOL_SIZE`. Set it before the first Mooncake mmap-buffer allocation in the process. If you build the image from source with `docker/mooncake.Dockerfile`, that source-built image also installs the helper as `mooncake-hicache-sizing`.
+Without `MC_STORE_USE_HUGEPAGE=1`, the arena may opportunistically try hugepages and then retry on regular pages if HugeTLB is unavailable. When `MC_STORE_USE_HUGEPAGE=1` is set, both the arena path and the direct-`mmap()` fallback path require HugeTLB pages. Mooncake will not silently degrade that explicit hugepage request to regular pages.
+
 ## 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.
@@ -195,6 +262,14 @@ The following options can be used during `cmake ..` to specify whether to compil
 - `-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_HYGON=[ON|OFF]`: Enable Hygon DCU support via DTK SDK. **Default: OFF.** Uses CUDA-compatible runtime.
+- `-DDTK_ROOT=/path/to/dtk`: Override the default DTK SDK root used when `-DUSE_HYGON=ON`. If unset, Mooncake uses `DTK_HOME` or `/opt/dtk`.
+- `-DDTK_INCLUDE_DIR=/path/to/include`: Override the DTK include directory when `-DUSE_HYGON=ON`.
+- `-DDTK_LIB_DIR=/path/to/lib64`: Override the DTK library directory when `-DUSE_HYGON=ON`.
+- `-DUSE_COREX=[ON|OFF]`: Enable Iluvatar CoreX GPU support. **Default: OFF.** Uses CUDA-compatible runtime.
+- `-DCOREX_ROOT=/path/to/corex`: Override the default CoreX SDK root used when `-DUSE_COREX=ON`. If unset, Mooncake uses `COREX_HOME` or `/usr/local/corex`.
+- `-DCOREX_INCLUDE_DIR=/path/to/include`: Override the CoreX include directory when `-DUSE_COREX=ON`.
+- `-DCOREX_LIB_DIR=/path/to/lib`: Override the CoreX library directory when `-DUSE_COREX=ON`.
 - `-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`.
@@ -204,12 +279,16 @@ The following options can be used during `cmake ..` to specify whether to compil
 - `-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_RUST_EXAMPLE=[ON|OFF]`: Build the Transfer Engine Rust interface and sample code. **Default: OFF.**
+- `-DWITH_STORE_RUST=[ON|OFF]`: Build Mooncake Store Rust bindings and CMake Rust targets. **Default: ON.**
 - `-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_REDIS=[ON|OFF]`: Enable Redis-based metadata service for the Transfer Engine, require hiredis
 - `-DUSE_HTTP=[ON|OFF]`: Enable Http-based metadata service
 - `-DUSE_ETCD=[ON|OFF]`: Enable etcd-based metadata service, require go 1.23+
 - `-DSTORE_USE_ETCD=[ON|OFF]`: Enable etcd-based failover for Mooncake Store, require go 1.23+. **Note:** `-DUSE_ETCD` and `-DSTORE_USE_ETCD` are two independent options. Enabling `-DSTORE_USE_ETCD` does **not** depend on `-DUSE_ETCD`
+- `-DSTORE_USE_REDIS=[ON|OFF]`: Enable Redis-based failover for Mooncake Store, require hiredis. **Default: OFF.** **Note:** `-DUSE_REDIS` and `-DSTORE_USE_REDIS` are two independent options. Enabling `-DSTORE_USE_REDIS` does **not** depend on `-DUSE_REDIS`.
 - `-DBUILD_SHARED_LIBS=[ON|OFF]`: Build Transfer Engine as shared library, default is OFF
 - `-DBUILD_UNIT_TESTS=[ON|OFF]`: Build unit tests, default is ON
 - `-DBUILD_EXAMPLES=[ON|OFF]`: Build examples, default is ON
+- `-DUSE_ASCEND_DIRECT=[ON|OFF]`: Enable Ascend Direct transport and HCCS support via the ADXL engine (**recommended**).
+- `-DUSE_UBSHMEM=[ON|OFF]`: Enable Huawei Ascend NPU shared memory transport via CANN VMM APIs.
diff --git a/docs/source/getting_started/examples/lmcache-integration.md b/docs/source/getting_started/examples/lmcache-integration.md
index 7dbcef46..3df78327 100644
--- a/docs/source/getting_started/examples/lmcache-integration.md
+++ b/docs/source/getting_started/examples/lmcache-integration.md
@@ -1,4 +1,6 @@
-# Mooncake x LMCache: Unite to Pioneer KVCache-Centric LLM Serving System
+# Mooncake x LMCache Integration
+
+Mooncake and LMCache unite to pioneer KVCache-centric LLM serving systems.
 
 Mooncake and LMCache have announced a strategic collaboration aimed at pioneering a KVCache-centric Large Language Model (LLM) serving system. This partnership seeks to significantly enhance the efficiency, scalability, and responsiveness of LLM applications.
 
@@ -23,7 +25,12 @@ This strategic alliance delivers substantial mutual technological advancements:
 
 For a complete deployment guide with step-by-step instructions, see:
 
-👉 **[vLLM V1 Disaggregated Serving with Mooncake Store and LMCache](vllm-integration/vllmv1-lmcache-integration.md)**
+:::{toctree}
+:maxdepth: 1
+
+../examples/vllm-integration/vllmv1-lmcache-integration
+../examples/vllm-integration/vllmv1-lmcache-mp-integration
+:::
 
 ## Performance Benchmarking and Results
 
diff --git a/docs/source/getting_started/examples/lmdeploy-integration-v0.9.md b/docs/source/getting_started/examples/lmdeploy-integration-v0.9.md
index 5de2a669..ac6fd6fb 100644
--- a/docs/source/getting_started/examples/lmdeploy-integration-v0.9.md
+++ b/docs/source/getting_started/examples/lmdeploy-integration-v0.9.md
@@ -16,17 +16,17 @@ pip install mooncake-transfer-engine
 
 Note:
 
--   If any `.so` file is missing, uninstall the pip package with `pip3 uninstall mooncake-transfer-engine`, and build the binaries manually from source following the [build instructions](https://github.com/kvcache-ai/Mooncake/blob/main/doc/en/build.md).
+-   If any `.so` file is missing, uninstall the pip package with `pip3 uninstall mooncake-transfer-engine`, and build the binaries manually from source following the [build instructions](../build.md).
 
 ### Install the latest version of LMDeploy
 
-##### 1. Clone LMDeploy from the official repo
+#### 1. Clone LMDeploy from the official repo
 
 ```bash
 git clone https://github.com/InternLM/lmdeploy.git
 ```
 
-##### 2. Build
+#### 2. Build
 
 ##### 2.1 Build from source
 
@@ -74,7 +74,7 @@ lmdeploy serve api_server Qwen/Qwen3-8B \
 -   The `--role` parameter sets the node role in the disaggregated system (`Prefill` for token embedding and KV cache generation).
 -   The `--proxy-url` parameter connects the worker instance back to the proxy for coordination.
 -   The `--backend` parameter specifies the model execution backend (e.g., `pytorch`, `turbomind`).
--   The `--migration-backend` parameter defines the KV cache transport mechanism (e.g., `Mooncake` and `DlSllime`).
+-   The `--migration-backend` parameter defines the KV cache transport mechanism (e.g., `Mooncake` and `DLSlime`).
 
 #### Decode:
 
@@ -164,4 +164,4 @@ curl -X POST "http://192.168.0.147:8000/v1/completions" \
 -   You can specify multiple prefill or decode instances with distinct `--server-port` and different GPUs using `CUDA_VISIBLE_DEVICES`.
 -   MooncakeTransferEngine supports both intra-node (PCIe) and inter-node (RDMA/CXL) transfer, and device selection is automatic or customizable via config.
 -   When using HF models that timeout during prefill, consider setting model path to `~/Qwen3-8B` to accelerate loading from localhost.
--   Use `--log-level DEBUG` to get detailed runtime logs for troubleshooting.
\ No newline at end of file
+-   Use `--log-level DEBUG` to get detailed runtime logs for troubleshooting.
diff --git a/docs/source/getting_started/examples/sglang-integration-v1.md b/docs/source/getting_started/examples/sglang-integration-v1.md
index 3c6053ac..4338001d 100644
--- a/docs/source/getting_started/examples/sglang-integration-v1.md
+++ b/docs/source/getting_started/examples/sglang-integration-v1.md
@@ -1,175 +1,206 @@
 # SGLang Disaggregated Serving with MooncakeTransferEngine
 
 ## Overview
-This is the latest version of the MooncakeTransferEngine integration doc with the SGLang project based on [PR 4654](https://github.com/sgl-project/sglang/pull/4654) and [PR 4880](https://github.com/sgl-project/sglang/pull/4880) to support KVCache transfer for intra-node and inter-node disaggregated serving scenarios.
 
+SGLang uses Mooncake's Transfer Engine to enable disaggregated prefill-decode (PD) serving across nodes via RDMA, with support for EP and EPD backends. This integration is based on [PR 4654](https://github.com/sgl-project/sglang/pull/4654) and [PR 4880](https://github.com/sgl-project/sglang/pull/4880).
 
-**_Please note that this is still an experimental version and will be modified anytime based on feedback from the SGLang community._**
+In benchmarks, PD disaggregation with Mooncake achieves **~30% lower ITL** while maintaining comparable throughput ([details](../../performance/sglang-benchmark-results-v1)).
+
+```
+  +-----------+   Transfer Engine (RDMA)    +-----------+
+  | SGLang    | ◄━━━━━━━━━━━━━━━━━━━━━━► | SGLang     |
+  | Prefill   |     KV cache blocks        | Decode     |
+  +-----------+                            +-----------+
+```
+
+## Prerequisites
 
-## Installation
-### Prerequisite
 ```bash
 pip3 install mooncake-transfer-engine
 ```
 
-Note: If you encounter problems such as missing `lib*.so`, you should uninstall this package by `pip3 uninstall mooncake-transfer-engine`, and build the binaries manually according to the [instructions](build.md).
+If you encounter missing `lib*.so` errors, uninstall and build from source instead:
 
-### Install the latest version of SGLang
-#### 1. Clone SGLang from official repo
 ```bash
-git clone git@github.com:sgl-project/sglang.git
+pip3 uninstall mooncake-transfer-engine
+# Build from source — see build instructions
 ```
-#### 2. Build
-##### 2.1 Build from source
+
+### Install SGLang
+
+It is recommended to use uv for faster installation:
+
 ```bash
-cd sglang
 pip install --upgrade pip
-pip install -e "python[all]" --find-links https://flashinfer.ai/whl/cu124/torch2.5/flashinfer-python
+pip install uv
+uv pip install sglang
 ```
 
-#####  If running on AMD GPU, tried below steps to install sglang for rocm or just run in rocm docker from https://hub.docker.com/r/rocm/sgl-dev directly
+The major version of Cuda is 13 by default. To install sglang under Cuda 12 with pip or uv, please try the following commands:
+
 ```bash
-
 pip install --upgrade pip
-cd sgl-kernel
-python setup_rocm.py install
-cd ..
-pip install -e "python[all_hip]"
-
-or
-
-docker run -it --rm --network=host \
-    --device=/dev/kfd --device=/dev/dri \
-    --ipc=host --shm-size 16G \
-    --group-add video \
-    --cap-add=SYS_PTRACE \
-    --security-opt seccomp=unconfined \
-    -v /home/workspace:/workspace \
-   rocm/sgl-dev:20250707
-
-pip3 install mooncake-transfer-engine
-sudo apt update
-sudo apt install libibverbs1 libibverbs-dev -y
-sudo apt-get install lsof net-tools iputils-ping -y
-
+pip install uv
+uv pip install sglang
+uv pip install --force-reinstall  torch==2.11.0 torchaudio==2.11.0 torchvision --index-url https://download.pytorch.org/whl/cu129
+uv pip install --force-reinstall sglang-kernel --index-url https://docs.sglang.ai/whl/cu129/
+uv pip install --force-reinstall sgl-deep-gemm --index-url https://docs.sglang.ai/whl/cu129/ --no-deps
 ```
 
- - If you encounter any problems that you cannot solve, please refer to the [SGLang official compilation guide](https://docs.sglang.ai/start/install.html).
+See the [SGLang official compilation guide](https://docs.sglang.ai/start/install.html) if you encounter issues.
 
 ## Configuration
 
- - **Update(Apr 10, 2025)** Good news: The configuration file requirement has been removed since [PR 5460](https://github.com/sgl-project/sglang/pull/5460). There is no need to prepare the _mooncake.json_ file anymore.
+Key arguments:
 
+| Argument | Description |
+|----------|-------------|
+| `--disaggregation-mode` | `prefill` or `decode` — node role |
+| `--disaggregation-ib-device` | RDMA device(s). Auto-detected, comma-separated for multi-NIC |
+| `--tp-size` | Tensor parallelism size (optional) |
+| `--base-gpu-id` | Starting GPU index for same-node deployments |
+| `--host` / `--port` | SGLang service address |
 
+## Run PD Disaggregation
 
-### To run prefill instance and decode instance on different node
+### Multi-Node
 
-Prefill: 
 ```bash
-python -m sglang.launch_server --model-path Qwen/Qwen2.5-7B-Instruct-GPTQ-Int4 --disaggregation-mode prefill --port 30000 --host 192.168.0.137 --tp-size 2
-```
- - The `--model-path` parameter specifies the model to use.
- - The `--host` parameter specifies the SGLang service host.
- - The `--port` parameter specifies the SGLang service port on which to listen.
- - The `--disaggregation-mode` is the node's role, either 'prefill' or 'decode'.
- - The `--disaggregation-ib-device` is the device to be used for data transmission, it is optional since we will detect this config automatically. Or you can still explicitly specify devices if needed. If multiple NIC devices are used, they can be separated by commas, such as "erdma_0,erdma_1". Please note that there are no spaces between them.
- - Option `--tp-size` is supported. Example: append `--tp-size 2` to the run command to run SGLang with multiple GPUs.
+# Prefill node (192.168.0.137)
+python -m sglang.launch_server \
+  --model-path Qwen/Qwen2.5-7B-Instruct-GPTQ-Int4 \
+  --disaggregation-mode prefill \
+  --port 30000 --host 192.168.0.137 --tp-size 2
 
-Decode:
-```bash
-python -m sglang.launch_server --model-path Qwen/Qwen2.5-7B-Instruct-GPTQ-Int4 --disaggregation-mode decode --port 30001 --host 192.168.0.140 --tp-size 2
-```
+# Decode node (192.168.0.140)
+python -m sglang.launch_server \
+  --model-path Qwen/Qwen2.5-7B-Instruct-GPTQ-Int4 \
+  --disaggregation-mode decode \
+  --port 30001 --host 192.168.0.140 --tp-size 2
 
-Proxy:
-```bash
-python3 -m sglang.srt.disaggregation.mini_lb --prefill http://192.168.0.137:30000 --decode http://192.168.0.140:30001 --host 0.0.0.0 --port 8000
+# Router
+python3 -m sglang_router.launch_router \
+  --pd-disaggregation \
+  --prefill "http://192.168.0.137:30000" 8998 \
+  --decode  "http://192.168.0.140:30001" \
+  --policy round_robin \
+  --host 0.0.0.0 --port 8000
 ```
 
 Test:
+
 ```bash
-curl -X POST http://127.0.0.1:8000/generate -H "Content-Type: application/json" -d '{
-  "text": "Let me tell you a long story ",
-  "sampling_params": {
-    "temperature": 0
-  }
-}'
+curl -X POST http://127.0.0.1:8000/generate \
+  -H "Content-Type: application/json" \
+  -d '{"text": "Tell me a long story", "sampling_params": {"temperature": 0}}'
 ```
 
-### To run prefill instance and decode instance on the same node
+### Single-Node
 
-Prefill: 
-```bash
-python -m sglang.launch_server --model-path Qwen/Qwen2.5-7B-Instruct-GPTQ-Int4 --disaggregation-mode prefill --port 30000 --host 192.168.0.137 --tp-size 2
-```
-
-Decode:
-```bash
-python -m sglang.launch_server --model-path Qwen/Qwen2.5-7B-Instruct-GPTQ-Int4 --disaggregation-mode decode --port 30001 --base-gpu-id 2 --host 192.168.0.137 --tp-size 2
-```
- - The function of `--base-gpu-id` is similar to env var`CUDA_VISIBLE_DEVICES`, which is used to avoid reusing the 0th GPU card. The difference is that it is used to specify the starting number of the GPU. If it is set to 2, the first and second cards will be skipped, and the third card will be used directly.
-
-Proxy:
-```bash
-python3 -m sglang.srt.disaggregation.mini_lb --prefill http://192.168.0.137:30000 --decode http://192.168.0.137:30001 --host 0.0.0.0 --port 8000
-```
-
-Test:
-```bash
-curl -X POST http://127.0.0.1:8000/generate -H "Content-Type: application/json" -d '{
-  "text": "Let me tell you a long story ",
-  "sampling_params": {
-    "temperature": 0
-  }
-}'
-```
-
-Note:
- - TP is supported but not required, you can remove `--tp-size 2` if you want.
- - The `--disaggregation-ib-device` is the device to be used for data transmission, it is optional since we will detect this config automatically. Or you can still explicitly specify devices if needed. If multiple NIC devices are used, they can be separated by commas, such as "erdma_0,erdma_1". Please note that there are no spaces between them.
- - XpYd is supported. It is ok to run one prefill and multiple decode instances on the same node, however, multiple prefill instances on the same node are not supported due to the port conflict of bootstrap server.
-   - e.g., `python3 -m sglang.srt.disaggregation.mini_lb --prefill http://192.168.0.137:30000,http://192.168.0.140:30000 --decode http://192.168.0.137:30001,http://192.168.0.140:30001 --host 0.0.0.0 --port 8000`
- - HuggingFace timeout can be addressed by `export SGLANG_USE_MODELSCOPE=true`
-
-### To enable Mooncake EP Backend
-
-
-Prefill:
-```bash
-python -m sglang.launch_server --model-path deepseek-ai/DeepSeek-V3-0324 --disaggregation-mode prefill --port 30000 --host 192.168.0.137 --tp-size 8 --dp-size 8 --elastic-ep-backend mooncake --mooncake-ib-device <mooncake-ib-device> --moe-a2a-backend mooncake
-```
-
-Decode:
-```bash
-python -m sglang.launch_server --model-path deepseek-ai/DeepSeek-V3-0324 --disaggregation-mode decode --port 30001 --host 192.168.0.140 --tp-size 8 --dp-size 8 --elastic-ep-backend mooncake --mooncake-ib-device <mooncake-ib-device> --moe-a2a-backend mooncake
-```
-- Set `--elastic-ep-backend` and `--moe-a2a-backend` to "mooncake" to enable Mooncake EP Backend.
-- The value of `--mooncake-ib-device` should be the same as `--disaggregation-ib-device`.
-
-
-### To enable Mooncake EPD Backend
-
-Encoder:
 ```bash
+# Prefill
 python -m sglang.launch_server \
-    --model-path $MODEL \
-    --encoder-only \
-    --encoder-transfer-backend mooncake \
-    --port $PORT
-```
+  --model-path Qwen/Qwen2.5-7B-Instruct-GPTQ-Int4 \
+  --disaggregation-mode prefill \
+  --port 30000 --host 192.168.0.137 --tp-size 2
 
-Prefill:
-```bash
+# Decode (skip first 2 GPUs, use cards 2+)
 python -m sglang.launch_server \
-    --model-path $MODEL \
-    --disaggregation-mode prefill \
-    --disaggregation-transfer-backend mooncake \
-    --encoder-transfer-backend mooncake \
-    --tp $TP \
-    --mem-fraction-static $MEM_FRACTION \
-    --chunked-prefill-size $CHUNK_SIZE \
-    --language-only \
-    --encoder-urls http://127.0.0.1:30002 http://127.0.0.1:30003 http://127.0.0.1:30004 http://127.0.0.1:30005 http://127.0.0.1:30006 http://127.0.0.1:30007 \
-    --port $PORT
+  --model-path Qwen/Qwen2.5-7B-Instruct-GPTQ-Int4 \
+  --disaggregation-mode decode \
+  --port 30001 --base-gpu-id 2 --host 192.168.0.137 --tp-size 2
+
+# Router
+python3 -m sglang_router.launch_router \
+  --pd-disaggregation \
+  --prefill "http://192.168.0.137:30000" 8998 \
+  --decode  "http://192.168.0.137:30001" \
+  --policy round_robin \
+  --host 0.0.0.0 --port 8000
+
+# Router for multi node.
+# Here is an example for 2 decode node running on 192.168.0.137 and 192.168.0.140
+python3 -m sglang_router.launch_router \
+  --pd-disaggregation \
+  --prefill "http://192.168.0.137:30000" 8998 \
+  --decode "http://192.168.0.137:30001,http://192.168.0.140:30001" \
+  --policy round_robin \
+  --host 0.0.0.0 --port 8000
 ```
 
-- Set `--encoder-transfer-backend` to "mooncake" to enable Mooncake Backend.
+TP is supported but not required — omit `--tp-size 2` for single-GPU setups.
+
+Multiple decode instances per prefill are supported. Multiple prefills on the same node are not supported due to bootstrap server port conflicts.
+
+```{tip}
+If you encounter HuggingFace timeouts, set `export SGLANG_USE_MODELSCOPE=true`.
+```
+
+## Advanced Backends
+
+### EP Backend — Expert Parallelism for MoE Models
+
+For Mixture-of-Experts models (e.g., DeepSeek-V3), different nodes hold different expert weights. During inference, Mooncake transfers expert activations between prefill and decode nodes via RDMA, replacing the NCCL-based all-to-all with a faster, disaggregation-aware transfer path.
+
+```bash
+# Prefill
+python -m sglang.launch_server \
+  --model-path deepseek-ai/DeepSeek-V3-0324 \
+  --disaggregation-mode prefill \
+  --port 30000 --host 192.168.0.137 \
+  --tp-size 8 --dp-size 8 \
+  --elastic-ep-backend mooncake \
+  --moe-a2a-backend mooncake
+
+# Decode
+python -m sglang.launch_server \
+  --model-path deepseek-ai/DeepSeek-V3-0324 \
+  --disaggregation-mode decode \
+  --port 30001 --host 192.168.0.140 \
+  --tp-size 8 --dp-size 8 \
+  --elastic-ep-backend mooncake \
+  --moe-a2a-backend mooncake
+```
+
+Set `--mooncake-ib-device` to the same value as `--disaggregation-ib-device` if explicit device specification is needed.
+
+### EPD Backend — Encoder-Prefill-Decode for Multimodal Models
+
+For multimodal models (e.g., LLaVA, InternVL), the **encoder** (Vision Transformer) processes images, the **prefill** node handles text + encoded visual tokens, and the **decode** node generates output. Mooncake transfers encoder outputs (visual embeddings) between these nodes via RDMA.
+
+The three node roles are:
+
+| Role | Flag | Responsibility |
+|------|------|----------------|
+| Encoder | `--encoder-only` | Processes images, produces visual embeddings. Mooncake sends embeddings to prefill via `--encoder-transfer-backend mooncake` |
+| Prefill | `--disaggregation-mode prefill` | Receives encoder embeddings, runs prefill with text + visual context. Mooncake transfers KV cache to decode |
+| Decode | `--disaggregation-mode decode` | Receives KV cache from prefill, generates output tokens |
+
+```bash
+# Encoder-only node
+python -m sglang.launch_server \
+  --model-path $MODEL \
+  --encoder-only \
+  --encoder-transfer-backend mooncake \
+  --port 30002
+
+# Prefill node
+python -m sglang.launch_server \
+  --model-path $MODEL \
+  --disaggregation-mode prefill \
+  --disaggregation-transfer-backend mooncake \
+  --encoder-transfer-backend mooncake \
+  --tp $TP \
+  --mem-fraction-static $MEM_FRACTION \
+  --chunked-prefill-size $CHUNK_SIZE \
+  --language-only \
+  --encoder-urls http://127.0.0.1:30002 http://127.0.0.1:30003 \
+  --port $PORT
+
+# Decode node
+python -m sglang.launch_server \
+  --model-path $MODEL \
+  --disaggregation-mode decode \
+  --disaggregation-transfer-backend mooncake \
+  --port $PORT
+```
diff --git a/docs/source/getting_started/examples/sglang-integration/hicache-integration-v1.md b/docs/source/getting_started/examples/sglang-integration/hicache-integration-v1.md
index f47d4886..4a6e011c 100644
--- a/docs/source/getting_started/examples/sglang-integration/hicache-integration-v1.md
+++ b/docs/source/getting_started/examples/sglang-integration/hicache-integration-v1.md
@@ -6,22 +6,6 @@ This document describes how to use Mooncake as the storage backend for SGLang Hi
 
 ## Introduction
 
-### About Mooncake
-
-Mooncake aims to enhance the inference efficiency of large language models (LLMs), especially in slow object storage environments, by constructing a multi-level caching pool on high-speed interconnected DRAM/SSD resources. Compared to traditional caching systems, Mooncake utilizes (GPUDirect) RDMA technology to transfer data directly in a zero-copy manner, while maximizing the use of multi-NIC resources on a single machine.
-
-For more details about Mooncake, please refer to [Mooncake project](https://github.com/kvcache-ai/Mooncake) and [Mooncake documents](https://kvcache-ai.github.io/Mooncake/).
-
-### About SGLang HiCache
-
-SGLang HiCache is a hierarchical KV caching system that extends SGLang's RadixAttention with advanced multi-tier memory management capabilities. It implements a scalable hierarchical storage architecture that spans GPU memory, CPU memory, and external storage layers, delivering significant performance improvements for large language model inference.
-
-HiCache introduces a **HiRadixTree** that acts as a page table for referencing KV caches across different memory tiers: **GPU Memory (L1)**, **CPU Memory (L2)**, **Mooncake and other Storage Backends (L3)**.
-
-The system includes an intelligent cache controller that automatically manages data movement between tiers, implementing optimized prefetching strategies and multiple write policies (write-through, write-through-selective, and write-back).
-
-For more details about SGLang HiCache, please refer to [HiCache system design document](https://docs.sglang.ai/advanced_features/hicache_design.html) and [this blog](https://lmsys.org/blog/2025-09-10-sglang-hicache/).
-
 ### Mooncake & SGLang HiCache
 
 Mooncake serves as a high-performance L3 storage backend for SGLang HiCache, enabling distributed KV cache storage across multiple servers with RDMA-accelerated data transfer. This integration addresses the capacity limitations of traditional GPU-only or GPU+CPU caching by providing virtually unlimited cache storage through a distributed memory pool.
@@ -42,21 +26,26 @@ This integration is particularly valuable for production deployments involving l
 
 ### Install SGLang
 
-1. Clone SGLang from official repo
+It is recommended to use uv for faster installation:
 
 ```bash
-git clone git@github.com:sgl-project/sglang.git
-```
-
-2. Build
-
-```bash
-cd sglang
 pip install --upgrade pip
-pip install -e "python[all]"
+pip install uv
+uv pip install sglang
 ```
 
-For more details, please refer to [SGLang official installation guide](https://docs.sglang.ai/get_started/install.html).
+The major version of Cuda is 13 by default. To install sglang under Cuda 12 with pip or uv, please try the following commands:
+
+```bash
+pip install --upgrade pip
+pip install uv
+uv pip install sglang
+uv pip install --force-reinstall  torch==2.11.0 torchaudio==2.11.0 torchvision --index-url https://download.pytorch.org/whl/cu129
+uv pip install --force-reinstall sglang-kernel --index-url https://docs.sglang.ai/whl/cu129/
+uv pip install --force-reinstall sgl-deep-gemm --index-url https://docs.sglang.ai/whl/cu129/ --no-deps
+```
+
+See the [SGLang official compilation guide](https://docs.sglang.ai/start/install.html) if you encounter issues.
 
 ### Install Mooncake
 
@@ -66,37 +55,7 @@ For more details, please refer to [SGLang official installation guide](https://d
 pip install mooncake-transfer-engine
 ```
 
-**Method 2: from source**
-
-Clone Mooncake project:
-
-```bash
-git clone https://github.com/kvcache-ai/Mooncake --recursive
-```
-
-Install dependencies:
-
-```bash
-cd Mooncake
-bash dependencies.sh
-```
-
-Build the project:
-
-```bash
-mkdir build
-cd build
-cmake ..
-make -j
-```
-
-Install Mooncake:
-
-```bash
-sudo make install
-```
-
-For more details, please refer to [Mooncake official installation guide](https://kvcache-ai.github.io/Mooncake/getting_started/build.html).
+If you want to build from source or using some other advanced features which not contained in prebuilt pip package, please refer to [Mooncake official installation guide](https://kvcache-ai.github.io/Mooncake/getting_started/build.html).
 
 ## Deployment
 
@@ -365,54 +324,12 @@ python -m sglang.launch_server \
 
 ### Prefill/Decode Disaggregation
 
-In **PD disaggregation**, the configurations for the `metadata service`, `mooncake master`, and the optional `store service` remain the same as described above. The difference is that SGLang introduces three distinct roles: `prefill worker`, `decode worker`, and `router`.
+Mooncake HiCache works with SGLang's **PD disaggregation** mode. The `master service`, `metadata service`, and optional `store service` configurations are the same as described above.
 
-Among these, the `prefill worker` supports enabling **HiCache**. To run with PD disaggregation, start from the [PD configuration](https://kvcache-ai.github.io/Mooncake/getting_started/examples/sglang-integration-v1.html), and add the HiCache-related parameters (as previously described for the `SGLang server`) to the `prefill worker`.
+1. Follow the [PD Disaggregation Guide](../sglang-integration-v1) to set up the prefill, decode, and router workers.
+2. Add the HiCache-related parameters (`--enable-hierarchical-cache`, `--hicache-storage-backend mooncake`, `--hicache-storage-prefetch-policy`, etc.) to the **prefill worker** only, as described in the HiCache sections above.
 
-In the example below, one `prefill worker`, one `decode worker`, and one `router` are launched. HiCache is enabled on the `prefill worker` to optimize prefill performance.
-
-**Prefill worker**:
-
-```bash
-MOONCAKE_TE_META_DATA_SERVER="http://127.0.0.1:8080/metadata" \
-MOONCAKE_MASTER=127.0.0.1:50051 \
-MOONCAKE_PROTOCOL="rdma" \
-MOONCAKE_DEVICE="mlx5_1" \
-MOONCAKE_GLOBAL_SEGMENT_SIZE=4294967296 \
-python -m sglang.launch_server \
-    --model-path [model_path] \
-    --page-size 64 \
-    --enable-hierarchical-cache \
-    --hicache-storage-prefetch-policy timeout \
-    --hicache-storage-backend mooncake \
-    --disaggregation-mode prefill \
-    --disaggregation-ib-device "mlx5_1" \
-    --base-gpu-id 0 \
-    --port 30000
-```
-
-**Decode worker**:
-
-```bash
-python -m sglang.launch_server \
-    --model-path [model_path] \
-    --page-size 64 \
-    --disaggregation-mode decode \
-    --disaggregation-ib-device "mlx5_1" \
-    --base-gpu-id 1 \
-    --port 30001
-```
-
-**Router**:
-
-```bash
-python -m sglang_router.launch_router \
-    --pd-disaggregation \
-    --prefill "http://127.0.0.1:30000" \
-    --decode "http://127.0.0.1:30001" \
-    --host 0.0.0.0 \
-    --port 8000
-```
+The Mooncake and HiCache configuration (environment variables or JSON config) is applied identically to the prefill worker — no changes are needed on the decode worker or router.
 
 ## Troubleshooting
 
@@ -429,6 +346,72 @@ If the model is small but the GPU memory is large — especially in multi-TP (te
 
 In such cases, you should manually configure an appropriate L2 cache size based on your hardware. This can be done by setting `--hicache-ratio` or `--hicache-size`.
 
+**HugeTLB Bring-up Checklist:**
+
+Before enabling HugeTLB-backed HiCache memory, size the host's hugepage pool against your planned `--hicache-size`, `MOONCAKE_GLOBAL_SEGMENT_SIZE`, and `MC_MMAP_ARENA_POOL_SIZE` values.
+
+From a source checkout:
+
+```bash
+python3 scripts/check_hicache_hugepage_requirements.py \
+    --tp-size 4 \
+    --hicache-size 64gb \
+    --global-segment-size 8gb \
+    --arena-pool-size 56gb \
+    --available-hugetlb 512gb
+```
+
+From the source-built Docker image:
+
+```bash
+mooncake-hicache-sizing \
+    --tp-size 4 \
+    --hicache-size 64gb \
+    --global-segment-size 8gb \
+    --arena-pool-size 56gb
+```
+
+The `64gb` / `56gb` values above are tuned examples for large multi-GPU HiCache deployments, not defaults. Arena enablement remains opt-in, and the default pool is `8gb` only when you enable the arena via gflag without an explicit env override. On smaller hosts, start with `8gb` or `16gb` and size upward with the helper.
+
+The helper reports two numbers:
+
+* **Baseline floor** = `hicache-size + MOONCAKE_GLOBAL_SEGMENT_SIZE` per rank. Falling below this usually means startup or allocation failure.
+* **Clean arena target** = baseline floor + `MC_MMAP_ARENA_POOL_SIZE` per rank. Reaching this target makes arena-backed launches less likely to spill onto the regular-page fallback path.
+
+Treat the helper as a planning tool rather than a hard proof: runs can still succeed between the baseline floor and the clean arena target if some arena allocations fall back to regular pages.
+
+Reserve HugeTLB pages on the host before launching SGLang:
+
+```bash
+sudo sysctl -w vm.nr_hugepages=262144
+grep -E 'HugePages_Total|HugePages_Free|Hugepagesize' /proc/meminfo
+```
+
+With `2 MiB` pages, `262144` pages equals `512 GiB`; `49152` pages equals `96 GiB`. Persist the setting with `/etc/sysctl.d/90-mooncake-hugepages.conf` if you need it across reboots.
+
+**Memory Allocator Tuning:**
+
+Mooncake's mmap arena is opt-in for HiCache host KV allocations. Setting `MC_MMAP_ARENA_POOL_SIZE` explicitly enables the arena and sizes the pool; the arena then pre-allocates a hugepage-backed pool and serves subsequent allocations via atomic bump pointer, reducing per-allocation latency from ~1000ns to ~50ns. The `56gb` example below is a benchmark-scale tuning value, not the allocator default.
+
+For HugeTLB-backed runs, export the hugepage and allocator settings together:
+
+```bash
+export MC_STORE_USE_HUGEPAGE=1
+export MC_STORE_HUGEPAGE_SIZE=2MB
+export MOONCAKE_GLOBAL_SEGMENT_SIZE=8gb
+export MC_MMAP_ARENA_POOL_SIZE=56gb
+```
+
+To disable the arena and fall back to direct `mmap()` while keeping the hugepage-backed baseline path, set the flag before the first Mooncake mmap-buffer allocation in the process:
+
+```bash
+export MC_DISABLE_MMAP_ARENA=1
+```
+
+Without `MC_STORE_USE_HUGEPAGE=1`, the arena may opportunistically try hugepages and then retry on regular pages if HugeTLB is unavailable. If `MC_STORE_USE_HUGEPAGE=1` is set, Mooncake treats hugepages as a hard requirement for both the arena path and the direct-`mmap()` baseline path. It will not silently retry those host-buffer allocations on regular pages.
+
+If the helper reports `baseline_fits_arena_may_fallback`, either increase `vm.nr_hugepages` or reduce `MC_MMAP_ARENA_POOL_SIZE`. For containerized launches, pass the same environment variables through `docker run -e ...` and use `--ipc=host --ulimit memlock=-1 --shm-size=128g`.
+
 **More Information:**
 
 Additional troubleshooting information can be found [here](https://kvcache-ai.github.io/Mooncake/troubleshooting/troubleshooting.html).
diff --git a/docs/source/getting_started/examples/sglang-integration/hicache-quick-start.md b/docs/source/getting_started/examples/sglang-integration/hicache-quick-start.md
index e234349b..a338ce29 100644
--- a/docs/source/getting_started/examples/sglang-integration/hicache-quick-start.md
+++ b/docs/source/getting_started/examples/sglang-integration/hicache-quick-start.md
@@ -1,6 +1,6 @@
 # Quick Start: SGLang HiCache with Mooncake Backend
 
-Follow this streamlined workflow to get SGLang HiCache running with Mooncake as the L3 storage backend.
+Follow this streamlined workflow to get SGLang HiCache running with Mooncake as the L3 storage backend. In benchmarks, pre-populated Mooncake achieves **best TTFT** across all tiers, maintaining high cache hit rates as conversation rounds grow ([details](../../../performance/sglang-hicache-benchmark-results-v1)).
 
 > Need more background or tuning options? See the [Complete Guide](hicache-integration-v1.md).
 
@@ -10,21 +10,43 @@ Before you begin, make sure that:
 
 - SGLang is installed with HiCache support on the machine hosting your SGLang server. Refer to the [official installation guide](https://docs.sglang.ai/get_started/install.html) if needed.
 - Mooncake is installed and accessible as the hierarchical cache backend. Detailed build steps live in the [Mooncake documentation](https://kvcache-ai.github.io/Mooncake/getting_started/build.html).
+- If you want HugeTLB-backed host KV memory, reserve the hugepages on the host before launching SGLang. The helper `python3 scripts/check_hicache_hugepage_requirements.py ...` from a source checkout, or `mooncake-hicache-sizing ...` inside the source-built Docker image, can size the required budget from your `--hicache-size`, `MOONCAKE_GLOBAL_SEGMENT_SIZE`, and `MC_MMAP_ARENA_POOL_SIZE` values.
 - The router package is installed to provide the `sglang_router` entrypoint:
 
 ```bash
 pip install sglang-router
 ```
 
-### 1. Launch the Mooncake master service
+### 1. Size and reserve HugeTLB on the host
+
+```bash
+python3 scripts/check_hicache_hugepage_requirements.py \
+    --tp-size 4 \
+    --hicache-size 64gb \
+    --global-segment-size 8gb \
+    --arena-pool-size 56gb \
+    --available-hugetlb 512gb
+
+sudo sysctl -w vm.nr_hugepages=262144
+grep -E 'HugePages_Total|HugePages_Free|Hugepagesize' /proc/meminfo
+```
+
+`262144` is an example for `512 GiB` of `2 MiB` pages. The `64gb` / `56gb` sizing inputs above are tuned examples for large multi-GPU runs, not defaults. Adjust them, and the resulting `vm.nr_hugepages`, to match your planned HiCache budget. On smaller hosts, start with an arena pool of `8gb` or `16gb`.
+
+### 2. Launch the Mooncake master service
 
 ```bash
 mooncake_master --enable_http_metadata_server=true
 ```
 
-### 2. Launch SGLang with Mooncake L3 storage
+### 3. Launch SGLang with Mooncake L3 storage
 
 ```bash
+export MC_STORE_USE_HUGEPAGE="1"
+export MC_STORE_HUGEPAGE_SIZE="2MB"
+export MOONCAKE_GLOBAL_SEGMENT_SIZE="8gb"
+export MC_MMAP_ARENA_POOL_SIZE="56gb"
+
 MOONCAKE_MASTER=127.0.0.1:50051 python -m sglang.launch_server \
     --model-path [model_path] \
     --page-size 64 \
@@ -35,6 +57,8 @@ MOONCAKE_MASTER=127.0.0.1:50051 python -m sglang.launch_server \
 
 **Key flag:** `--hicache-storage-prefetch-policy {best_effort,wait_complete,timeout}` determines when prefetching from storage should stop. `timeout` usually offers the best balance when Mooncake is the backend.
 
+**Memory allocator:** Mooncake's mmap arena is opt-in. The `56gb` example above is a benchmark-scale tuning value, not the default. Setting `MC_MMAP_ARENA_POOL_SIZE` enables it and chooses the pool size; if you enable it via gflag instead, the default pool size is `8gb`. To force the baseline direct-`mmap()` path, set `MC_DISABLE_MMAP_ARENA=1` (also accepts `true`, `yes`, or `on`) before the first Mooncake mmap-buffer allocation in the process. See the [Complete Guide](hicache-integration-v1.md) for the full HugeTLB sizing and troubleshooting flow.
+
 ## Prefill/Decode Disaggregation
 
 The disaggregated setup runs three processes—prefill worker, decode worker, and router. Launch each command below in its own terminal window.
@@ -94,6 +118,7 @@ curl -X POST http://127.0.0.1:8000/generate \
 
 - `--disaggregation-ib-device` is optional—SGLang autodetects devices, but you can set it explicitly (comma-separated, no spaces) when multiple NICs are available.
 - Use `--tp-size` to enable tensor-parallel execution across GPUs if required.
+- For Docker launches, use `--ipc=host`, `--ulimit memlock=-1`, and a large `--shm-size`, then pass the same `MC_STORE_USE_HUGEPAGE`, `MC_STORE_HUGEPAGE_SIZE`, `MOONCAKE_GLOBAL_SEGMENT_SIZE`, and `MC_MMAP_ARENA_POOL_SIZE` values through `docker run -e ...`.
 - Optional flags to experiment with once you need them (some still have compatibility gaps):
   - `--disaggregation-decode-enable-offload-kvcache` writes the decode worker's outputs back into Mooncake; enable it when you want decoded KV to persist in L3.
 - Launch dedicated Mooncake `store service` nodes when you want to scale L3 capacity beyond what the SGLang servers contribute.
diff --git a/docs/source/getting_started/examples/sglang-integration/index.md b/docs/source/getting_started/examples/sglang-integration/index.md
index 62809df1..885c444f 100644
--- a/docs/source/getting_started/examples/sglang-integration/index.md
+++ b/docs/source/getting_started/examples/sglang-integration/index.md
@@ -1,8 +1,54 @@
-# SGLang HiCache with Mooncake Backend
+# Mooncake x SGLang Integration
+
+Mooncake integrates with SGLang through two paths — **PD Disaggregation** for cross-instance KV cache transfer via the Transfer Engine, and **HiCache L3 Backend** for hierarchical KV cache storage with Mooncake Store.
+
+---
+
+## PD Disaggregation
+
+SGLang uses Mooncake's Transfer Engine for direct zero-copy KV cache transfer between prefill and decode instances over RDMA, with support for EP and EPD backends. In benchmarks, PD disaggregation with Mooncake achieves **~30% lower ITL** while maintaining comparable throughput.
+
+```
+  +-----------+   Transfer Engine (RDMA)   +------------+
+  | SGLang    | ◄━━━━━━━━━━━━━━━━━━━━━━►   | SGLang     |
+  | Prefill   |     KV cache blocks        | Decode     |
+  +-----------+                            +------------+
+```
+
+**Related:** [Full PD Disaggregation Guide](../sglang-integration-v1) — installation, cross-node/same-node setup, XpYd topology, EP backend for MoE models, and EPD backend for multimodal models.
+
+**Benchmark:** [PD Disaggregation Performance](../../../performance/sglang-benchmark-results-v1) — compares 1P1D disaggregation with regular SGLang instances.
+
+---
+
+## HiCache with Mooncake Store
+
+HiCache extends SGLang's RadixAttention with three memory tiers, using Mooncake Store as the distributed L3 backend. When local cache misses, HiCache automatically prefetches KV blocks from remote storage via RDMA.
+
+```
+  +----------------------------------------------+
+  | SGLang + HiCache                             |
+  |  ┌─────────┐  ┌─────────┐  ┌──────────────┐  |
+  |  │ L1(GPU) │  │ L2(CPU) │  │ L3(Mooncake) │  |
+  |  └─────────┘  └─────────┘  └──────┬───────┘  |
+  +-----------------------------------+----------+
+                                      |
+                             +--------+---------+
+                             | Mooncake Store   |
+                             | Distributed Pool |
+                             +------------------+
+```
+
+**Related:**
+- [HiCache Quick Start](hicache-quick-start) — minimal setup steps with hugepage sizing
+- [HiCache Complete Guide](hicache-integration-v1) — full deployment, prefetch strategies, memory tuning, and architecture deep dive
+- [SGLang Performance Benchmarks](../../../performance/sglang/index) — benchmark overview for PD disaggregation and HiCache with Mooncake
 
 ::::{toctree}
 :maxdepth: 1
+:hidden:
 
+../sglang-integration-v1
 hicache-quick-start
 hicache-integration-v1
 ::::
diff --git a/docs/source/getting_started/examples/vllm-integration/disagg-prefill-decode.md b/docs/source/getting_started/examples/vllm-integration/disagg-prefill-decode.md
new file mode 100644
index 00000000..c8f0ab24
--- /dev/null
+++ b/docs/source/getting_started/examples/vllm-integration/disagg-prefill-decode.md
@@ -0,0 +1,337 @@
+# Disaggregated Prefill-Decode with MooncakeConnector
+
+## Overview
+
+This guide demonstrates how to use `MooncakeConnector` with vLLM for disaggregated Prefill-Decode (PD) serving. `MooncakeConnector` enables direct cross-node KV cache transfer between prefill and decode instances using RDMA technology, achieving up to **142.25 GB/s** peak bandwidth (71.1% utilization of 8x RoCE).
+
+For more details about Mooncake, please refer to [Mooncake project](https://github.com/kvcache-ai/Mooncake) and [Mooncake documents](https://kvcache-ai.github.io/Mooncake/).
+
+---
+
+## Choose Your vLLM Backend
+
+| Backend | vLLM Version | Status | Guide |
+|---------|-------------|--------|-------|
+| **vLLM V1** | Latest | Recommended | [Jump to V1 guide](#using-vllm-v1-recommended) |
+| **vLLM V0** | ≤ v0.6.4.post1 | Legacy | [Jump to V0 guide](#using-vllm-v0-legacy) |
+
+```{admonition} New Users
+:class: tip
+If you are starting a new deployment, use the **vLLM V1** backend. V0 support is maintained for existing deployments only.
+```
+
+---
+
+## Using vLLM V1 (Recommended)
+
+This section covers `MooncakeConnector` integration with vLLM V1 backend. The integration enables efficient cross-node KV cache transfer via RDMA.
+
+### Installation
+
+#### Prerequisites
+
+Install mooncake-transfer-engine through pip:
+
+```bash
+pip install mooncake-transfer-engine
+```
+
+```{note}
+If you encounter problems such as missing `lib*.so`, uninstall this package by `pip3 uninstall mooncake-transfer-engine`, and build the binaries manually according to the [instructions](../../build.md).
+```
+
+#### Install vLLM
+
+Refer to [vLLM official installation guide](https://docs.vllm.ai/en/latest/getting_started/installation.html) for the latest installation instructions.
+
+### Usage
+
+#### Basic Setup (Different Nodes)
+
+**Prefiller Node** (192.168.0.2):
+
+```bash
+vllm serve Qwen/Qwen2.5-7B-Instruct \
+  --port 8010 \
+  --kv-transfer-config '{"kv_connector":"MooncakeConnector","kv_role":"kv_producer"}'
+```
+
+**Decoder Node** (192.168.0.3):
+
+```bash
+vllm serve Qwen/Qwen2.5-7B-Instruct \
+  --port 8020 \
+  --kv-transfer-config '{"kv_connector":"MooncakeConnector","kv_role":"kv_consumer"}'
+```
+
+**Proxy Server:**
+
+```bash
+# In vllm root directory.
+python tests/v1/kv_connector/nixl_integration/toy_proxy_server.py \
+  --prefiller-host 192.168.0.2 --prefiller-port 8010 \
+  --decoder-host 192.168.0.3 --decoder-port 8020
+```
+
+> NOTE: The Mooncake Connector currently uses the proxy from nixl_integration. This will be replaced with a self-developed proxy in the future.
+
+Now you can send requests to the proxy server through port 8000.
+
+**Test:**
+
+```bash
+curl http://127.0.0.1:8000/v1/chat/completions \
+  -H "Content-Type: application/json" \
+  -d '{
+    "model": "Qwen/Qwen2.5-7B-Instruct",
+    "messages": [
+      {"role": "user", "content": "Tell me a long story about artificial intelligence."}
+    ]
+  }'
+```
+
+#### Advanced Configuration
+
+**With Tensor Parallelism:**
+
+Prefiller:
+
+```bash
+CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 \
+vllm serve Qwen/Qwen2.5-7B-Instruct \
+  --port 8010 \
+  --tensor-parallel-size 8 \
+  --kv-transfer-config '{"kv_connector":"MooncakeConnector","kv_role":"kv_producer"}'
+```
+
+Decoder:
+
+```bash
+CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 \
+vllm serve Qwen/Qwen2.5-7B-Instruct \
+  --port 8020 \
+  --tensor-parallel-size 8 \
+  --kv-transfer-config '{"kv_connector":"MooncakeConnector","kv_role":"kv_consumer"}'
+```
+
+#### Configuration Parameters
+
+- `--kv-transfer-config`: JSON string to configure the KV transfer connector
+  - `kv_connector`: Set to `"MooncakeConnector"`
+  - `kv_role`: Role of the instance
+    - `kv_producer`: For prefiller instances that generate KV caches
+    - `kv_consumer`: For decoder instances that consume KV caches
+    - `kv_both`: Enables symmetric functionality (experimental)
+  - `num_workers`: Thread pool size in each prefiller worker to send kvcache (default 10)
+
+### Environment Variables
+
+- `VLLM_MOONCAKE_BOOTSTRAP_PORT`: Port for Mooncake bootstrap server (default: 8998)
+  - Required only for prefiller instances
+  - Each vLLM worker needs a unique port on its host
+  - For TP/DP deployments, each worker's port is computed as: `base_port + dp_rank * tp_size + tp_rank`
+- `VLLM_MOONCAKE_ABORT_REQUEST_TIMEOUT`: Timeout (in seconds) for automatically releasing KV cache (default: 480)
+  - Used when a request is aborted to prevent holding resources indefinitely
+
+### Performance
+
+For detailed performance benchmarks and results, see the [vLLM Benchmark](../../../performance/vllm-v1-support-benchmark.md) documentation.
+
+---
+
+## Using vLLM V0 (Legacy)
+
+```{admonition} Legacy Backend
+:class: warning
+This section is for vLLM V0 backend (≤ v0.6.4.post1). For new deployments, use the [V1 backend](#using-vllm-v1-recommended) above.
+```
+
+This integration is based on [PR 10502](https://github.com/vllm-project/vllm/pull/10502) and [PR 10884](https://github.com/vllm-project/vllm/pull/10884). Preview benchmark results are available at [vLLM Benchmark Results V0.2](../../../performance/vllm-benchmark-results-v0.2.md).
+
+### Installation
+
+#### Prerequisite
+
+```bash
+pip3 install mooncake-transfer-engine
+```
+
+```{note}
+- If you encounter problems such as missing `lib*.so`, uninstall this package by `pip3 uninstall mooncake-transfer-engine`, and build the binaries manually according to the [instructions](../../build.md).
+- For vLLM version ≤ v0.8.4, it requires `mooncake-transfer-engine ≤ 0.3.3.post2`. In the latest release, the interface `mooncake_vllm_adaptor` has been deprecated.
+```
+
+#### Install vLLM
+
+**1. Clone vLLM from official repo:**
+
+```bash
+git clone git@github.com:vllm-project/vllm.git
+```
+
+**2. Build from source (Include C++ and CUDA code):**
+
+```bash
+cd vllm
+pip3 uninstall vllm -y
+pip3 install -e .
+```
+
+```{tip}
+If the build fails, try upgrading cmake: `pip3 install cmake --upgrade`.
+```
+
+If you encounter any problems, refer to the [vLLM official compilation guide](https://docs.vllm.ai/en/v0.6.4.post1/getting_started/installation.html#install-the-latest-code).
+
+### Configuration
+
+#### Prepare configuration file over RDMA
+
+Create a `mooncake.json` file for both Prefill and Decode instances. Use the identical config file on both sides.
+
+```json
+{
+  "prefill_url": "192.168.0.137:13003",
+  "decode_url": "192.168.0.139:13003",
+  "metadata_server": "192.168.0.139:2379",
+  "metadata_backend": "etcd",
+  "protocol": "rdma",
+  "device_name": "erdma_0"
+}
+```
+
+- `prefill_url`: The IP address and port of the Prefill node (port is used to communicate with metadata server).
+- `decode_url`: The IP address and port of the Decode node. If running prefill and decode on the same node, set a different port (at least 50 apart from `prefill_url` port) to avoid conflicts.
+- `metadata_server`: The metadata server address. Supports `etcd`, `redis`, and `http` backends. Example: `"etcd://192.168.0.137:2379"`, `"redis://192.168.0.137:6379"`, `"http://192.168.0.137:8080/metadata"`.
+- `metadata_backend`: Currently supports `"etcd"`, `"redis"`, and `"http"`. If absent and `metadata_server` has no prefix, defaults to `"etcd"`. This parameter will be deprecated in a future version.
+- `protocol`: `"rdma"` or `"tcp"`.
+- `device_name`: Required when protocol is `"rdma"`. Multiple NICs can be separated by commas (`"erdma_0,erdma_1"`).
+
+#### Prepare configuration file over TCP
+
+```json
+{
+  "prefill_url": "192.168.0.137:13003",
+  "decode_url": "192.168.0.139:13003",
+  "metadata_server": "192.168.0.139:2379",
+  "metadata_backend": "etcd",
+  "protocol": "tcp",
+  "device_name": ""
+}
+```
+
+### Run Example
+
+Change the IP addresses and ports according to your environment.
+
+```bash
+# Begin from root of your cloned repo!
+
+# 1. Start the etcd server
+etcd --listen-client-urls http://0.0.0.0:2379 --advertise-client-urls http://localhost:2379
+
+# 2. Run on the prefilling side (producer role)
+MOONCAKE_CONFIG_PATH=./mooncake.json VLLM_USE_MODELSCOPE=True python3 -m vllm.entrypoints.openai.api_server \
+    --model Qwen/Qwen2.5-7B-Instruct-GPTQ-Int4 \
+    --port 8100 \
+    --max-model-len 10000 \
+    --gpu-memory-utilization 0.8 \
+    --kv-transfer-config '{"kv_connector":"MooncakeConnector","kv_role":"kv_producer","kv_rank":0,"kv_parallel_size":2,"kv_buffer_size":2e9}'
+
+# 3. Run on the decoding side (consumer role)
+MOONCAKE_CONFIG_PATH=./mooncake.json VLLM_USE_MODELSCOPE=True python3 -m vllm.entrypoints.openai.api_server \
+    --model Qwen/Qwen2.5-7B-Instruct-GPTQ-Int4 \
+    --port 8200 \
+    --max-model-len 10000 \
+    --gpu-memory-utilization 0.8 \
+    --kv-transfer-config '{"kv_connector":"MooncakeConnector","kv_role":"kv_consumer","kv_rank":1,"kv_parallel_size":2,"kv_buffer_size":2e9}'
+```
+
+**Key parameters:**
+- `MOONCAKE_CONFIG_PATH`: Path to the mooncake.json configuration file.
+- `VLLM_USE_MODELSCOPE`: Optional. Remove if you have HuggingFace access.
+- `--kv-transfer-config`: Connector configuration
+  - `kv_connector`: `"MooncakeConnector"`
+  - `kv_role`: `"kv_producer"` or `"kv_consumer"`
+  - `kv_rank`: 0 for producer, 1 for consumer
+  - `kv_parallel_size`: Fixed to 2 currently
+  - `kv_buffer_size`: KVCache lookup buffer size; increase for longer prompts. If OOM occurs, decrease `--gpu-memory-utilization`.
+  - `kv_ip` and `kv_port`: Used to specify the IP address and port of the master node for `"PyNcclConnector"` distributed setup. Not used for `"MooncakeConnector"` currently. Instead, `"MooncakeConnector"` uses a config file to set up the distributed connection.
+- `--tensor-parallel-size` / `-tp`: Supported. If running on the same node, set different `CUDA_VISIBLE_DEVICES`.
+
+```{note}
+If running prefill and decode on the same node, set a different port for `decode_url`. To avoid port conflicts, ensure the decode port differs by at least 50 from the `prefill_url` port (e.g., `"decode_url": "192.168.0.137:13103"`). If the same URL is set for both, the port of `decode_url` will be automatically incremented by 100.
+```
+
+**Proxy Server:**
+
+```bash
+python3 proxy_server.py
+```
+
+```python
+# proxy_server.py
+import os
+import aiohttp
+from quart import Quart, make_response, request
+
+AIOHTTP_TIMEOUT = aiohttp.ClientTimeout(total=6 * 60 * 60)
+app = Quart(__name__)
+
+async def forward_request(url, data):
+    async with aiohttp.ClientSession(timeout=AIOHTTP_TIMEOUT) as session:
+        headers = {"Authorization": f"Bearer {os.environ.get('OPENAI_API_KEY')}"}
+        async with session.post(url=url, json=data, headers=headers) as response:
+            if response.status == 200:
+                async for chunk_bytes in response.content.iter_chunked(1024):
+                    yield chunk_bytes
+
+@app.route('/v1/completions', methods=['POST'])
+async def handle_request():
+    try:
+        original_request_data = await request.get_json()
+        prefill_request = original_request_data.copy()
+        prefill_request['max_tokens'] = 1  # prefill only
+        async for _ in forward_request('http://localhost:8100/v1/completions', prefill_request):
+            continue
+        generator = forward_request('http://192.168.0.139:8200/v1/completions',  # Change IP
+                                    original_request_data)
+        response = await make_response(generator)
+        response.timeout = None
+        return response
+    except Exception as e:
+        import sys, traceback
+        exc_info = sys.exc_info()
+        print("Error occurred in disagg prefill proxy server")
+        print(e)
+        print("".join(traceback.format_exception(*exc_info)))
+
+if __name__ == '__main__':
+    app.run(host="0.0.0.0", port=8000)
+```
+
+> Be sure to change the IP address in the proxy server code.
+
+### Test
+
+```bash
+curl -s http://localhost:8000/v1/completions \
+  -H "Content-Type: application/json" \
+  -d '{
+    "model": "Qwen/Qwen2.5-7B-Instruct-GPTQ-Int4",
+    "prompt": "San Francisco is a",
+    "max_tokens": 1000
+  }'
+```
+
+---
+
+## Troubleshooting
+
+- If you encounter connection issues, check that:
+  - All nodes can reach each other over the network
+  - Firewall rules allow traffic on the specified ports
+  - RDMA devices are properly configured and listed in `device_name`
+- For missing library errors, rebuild `mooncake-transfer-engine` from source
+- Enable debug logging with `VLLM_LOGGING_LEVEL=DEBUG` for detailed diagnostics
+- For production deployments, consider using a more robust proxy solution
diff --git a/docs/source/getting_started/examples/vllm-integration/index.md b/docs/source/getting_started/examples/vllm-integration/index.md
index c1ac0c84..d6107e50 100644
--- a/docs/source/getting_started/examples/vllm-integration/index.md
+++ b/docs/source/getting_started/examples/vllm-integration/index.md
@@ -1,12 +1,55 @@
-# vLLM Disaggregated Serving
+# Mooncake x vLLM Integration
+
+## Overview
+
+Mooncake integrates with vLLM to accelerate large language model serving through high-performance KV cache transfer and shared storage. The integration supports two primary scenarios:
+
+- **Disaggregated Prefill-Decode Serving**: Seamlessly split prefill and decode across nodes using `MooncakeConnector`, with RDMA-powered cross-node KV cache transfer achieving up to **142.25 GB/s** peak bandwidth (71.1% utilization of 8x RoCE). Transfer overhead is negligible &mdash; for 32K-token prompts (4.50 GB of KV data), transfer takes only **31.65 ms**, accounting for just **4.2%** of total TTFT.
+- **KV Cache Storage & Sharing**: Extend effective KV cache capacity via `MooncakeStore` / `MooncakeStoreConnector`, with hash-based prefix caching that enables multiple vLLM instances to share cached KV blocks. Supports CPU/Disk offloading and dynamic XpYd topologies at runtime.
+
+| Scenario | Guide | vLLM Backend |
+|----------|-------|-------------|
+| PD Disaggregation (KV transfer) | [Disaggregated Prefill-Decode](disagg-prefill-decode) | V1 ✅ / V0 ⚠️ |
+| KV Cache Storage & Sharing | [KV Cache Storage with MooncakeStore](kv-cache-storage) | V1 ✅ / V0 ⚠️ |
+
+For detailed benchmark coverage across these scenarios, see [vLLM Integration Performance Benchmarks](../../../performance/vllm/index).
+
+```{admonition} New to Mooncake + vLLM?
+:class: tip
+Start with the V1 guides above. Legacy V0 documentation is available for existing deployments only.
+```
+
+---
+
+## Getting Started
+
+### Disaggregated Prefill-Decode
+
+Direct KV cache transfer between prefill and decode nodes via `MooncakeConnector` using RDMA.
 
 ::::{toctree}
 :maxdepth: 1
 
-vllmv1-lmcache-integration
-vllm-integration-v0.2
-vllm-integration-v0.3
-vllm-integration-v1.0
+disagg-prefill-decode
 ::::
 
+### KV Cache Storage & Sharing
 
+Distributed KV cache storage via `MooncakeStore` / `MooncakeStoreConnector` for offloading, prefix caching, and cross-instance sharing.
+
+::::{toctree}
+:maxdepth: 1
+
+kv-cache-storage
+::::
+
+---
+
+## Archived Documentation
+
+The following pages are from earlier versions of the integration and are no longer maintained. All content has been consolidated into the scenario-based guides above.
+
+- [MooncakeStoreConnector (Original)](vllm-mooncakestoreconnector)
+- [vLLM V0 PD Disaggregation Demo (Original)](vllm-integration-v0.2)
+- [vLLM V0 MooncakeStore (Original)](vllm-integration-v0.3)
+- [vLLM V1 PD Disaggregation (Original)](vllm-integration-v1.0)
diff --git a/docs/source/getting_started/examples/vllm-integration/kv-cache-storage.md b/docs/source/getting_started/examples/vllm-integration/kv-cache-storage.md
new file mode 100644
index 00000000..c25427ad
--- /dev/null
+++ b/docs/source/getting_started/examples/vllm-integration/kv-cache-storage.md
@@ -0,0 +1,409 @@
+# KV Cache Storage & Sharing with MooncakeStore
+
+## Overview
+
+This guide demonstrates how to use `MooncakeStore` / `MooncakeStoreConnector` with vLLM to build a distributed KV cache storage pool. It enables KV cache offloading to CPU/SSD, hash-based prefix caching across multiple vLLM instances, and flexible XpYd disaggregated deployment — where you can dynamically adjust prefill and decode group sizes at runtime.
+
+Compared to Redis-based backends, MooncakeStore achieves significantly lower TTFT (e.g., **~32% improvement** in mean TTFT for 2P2D tp=2 under RDMA). See [benchmark results](../../../performance/vllm-benchmark-results-v1.md) for details.
+
+---
+
+## Choose Your vLLM Backend
+
+| Backend | Connector | vLLM Version | Status | Guide |
+|---------|-----------|-------------|--------|-------|
+| **vLLM V1** | `MooncakeStoreConnector` | Latest | Recommended | [Jump to V1 guide](#using-vllm-v1-recommended) |
+| **vLLM V0** | `MooncakeStore` | ≤ v0.6.4.post1 | Legacy | [Jump to V0 guide](#using-vllm-v0-legacy) |
+
+```{admonition} New Users
+:class: tip
+If you are starting a new deployment, use the **vLLM V1** backend with `MooncakeStoreConnector`. V0 support is maintained for existing deployments only.
+```
+
+Key differences from v0.x to v1:
+- **XpYd support and orchestration**: Dynamically change the population of prefill and decode groups
+- **More stable and fault-tolerant**: A sudden crash of a single vLLM instance is tolerable; instance-to-instance connections are removed, so each instance works as a vanilla vLLM instance capable of handling requests independently
+
+---
+
+## Using vLLM V1 (Recommended)
+
+This section covers `MooncakeStoreConnector` — the new vLLM KV connector that uses `MooncakeDistributedStore` as a shared KV cache pool. It enables:
+
+- **CPU/Disk offloading**: Extend effective KV cache capacity by offloading to CPU memory or SSD via Mooncake's transfer engine.
+- **Hash-based prefix caching across instances**: Multiple vLLM instances share cached KV blocks through the store using block-hash deduplication.
+- **Flexible deployment**: Works as a single-node KV cache extension (`kv_both`), or in disaggregated prefill-decode setups (`kv_producer` / `kv_consumer`).
+
+### Deployment
+
+#### 1. Prerequisites
+
+- [vLLM](https://github.com/vllm-project/vllm) is installed
+- [Mooncake](https://github.com/kvcache-ai/Mooncake) is installed
+
+Refer to the [vLLM official repository](https://github.com/vllm-project/vllm) and [Mooncake official repository](https://github.com/kvcache-ai/Mooncake) for installation instructions and building from source.
+
+#### 2. Mooncake Master Server
+
+**Start:**
+
+```shell
+mooncake_master --port 50063
+```
+
+**Configure Mooncake**: Create a JSON configuration file (e.g., `mooncake_config.json`):
+
+```json
+{
+  "metadata_server": "http://127.0.0.1:8092/metadata",
+  "master_server_address": "127.0.0.1:50063",
+  "global_segment_size": "0",
+  "local_buffer_size": "2147483648",
+  "protocol": "rdma",
+  "device_name": ""
+}
+```
+
+**Set environment variable:**
+
+```shell
+export MOONCAKE_CONFIG_PATH=/path/to/mooncake_config.json
+```
+
+#### 3. Usage
+
+**3.1 Single-Node KV Cache Offloading** (`kv_both`):
+
+```shell
+MOONCAKE_CONFIG_PATH=mooncake_config.json \
+vllm serve meta-llama/Llama-3.1-8B-Instruct \
+    --kv-transfer-config '{"kv_connector":"MooncakeStoreConnector","kv_role":"kv_both"}'
+```
+
+**3.2 XpYd Disaggregated Prefill-Decode** (`kv_producer/kv_consumer`):
+
+Prefill Node:
+
+```shell
+MOONCAKE_CONFIG_PATH=mooncake_config.json \
+VLLM_MOONCAKE_BOOTSTRAP_PORT=50052 \
+vllm serve meta-llama/Llama-3.1-8B-Instruct \
+    --port 8100 \
+    --kv-transfer-config '{
+        "kv_connector": "MultiConnector",
+        "kv_role": "kv_producer",
+        "kv_connector_extra_config": {
+            "connectors": [
+                {
+                    "kv_connector": "MooncakeConnector",
+                    "kv_role": "kv_producer"
+                },
+                {
+                    "kv_connector": "MooncakeStoreConnector",
+                    "kv_role": "kv_producer"
+                }
+            ]
+        }
+    }'
+```
+
+Decode Node:
+
+```shell
+MOONCAKE_CONFIG_PATH=mooncake_config.json \
+VLLM_MOONCAKE_BOOTSTRAP_PORT=50053 \
+vllm serve meta-llama/Llama-3.1-8B-Instruct \
+    --port 8200 \
+    --kv-transfer-config '{
+        "kv_connector": "MultiConnector",
+        "kv_role": "kv_consumer",
+        "kv_connector_extra_config": {
+            "connectors": [
+                {
+                    "kv_connector": "MooncakeConnector",
+                    "kv_role": "kv_consumer"
+                },
+                {
+                    "kv_connector": "MooncakeStoreConnector",
+                    "kv_role": "kv_consumer"
+                }
+            ]
+        }
+    }'
+```
+
+Proxy:
+
+```shell
+python examples/disaggregated/disaggregated_serving/mooncake_connector/mooncake_connector_proxy.py \
+    --prefill http://192.168.0.2:8100 \
+    --decode http://192.168.0.3:8200
+```
+
+> When running with data parallelism, set a fixed `PYTHONHASHSEED` so that block hashes are consistent across DP ranks:
+>
+> ```shell
+> PYTHONHASHSEED=0 vllm serve ...
+> ```
+>
+> Without this, identical prompts may produce different block hashes on different DP ranks, preventing cross-instance prefix cache hits.
+
+---
+
+## Using vLLM V0 (Legacy)
+
+```{admonition} Legacy Backend
+:class: warning
+This section is for vLLM V0 backend with `MooncakeStore`. For new deployments, use the [V1 backend with `MooncakeStoreConnector`](#using-vllm-v1-recommended) above.
+```
+
+This integration is based on [PR 10502](https://github.com/vllm-project/vllm/pull/10502) and [PR 12957](https://github.com/vllm-project/vllm/pull/12957) to support KVCache transfer for intra-node and inter-node disaggregated serving.
+
+### Installation
+
+#### Prerequisite
+
+```bash
+pip3 install mooncake-transfer-engine
+```
+
+```{note}
+- If you encounter problems such as missing `lib*.so`, uninstall by `pip3 uninstall mooncake-transfer-engine`, and build manually according to the [instructions](../../build.md).
+- For vLLM version ≤ v0.8.4, it requires `mooncake-transfer-engine ≤ 0.3.3.post2`. The interface `mooncake_vllm_adaptor` has been deprecated in the latest release.
+```
+
+#### Install vLLM
+
+**1. Clone vLLM:**
+
+```bash
+git clone git@github.com:vllm-project/vllm.git
+```
+
+**2. Build from source:**
+
+```bash
+cd vllm
+pip3 install -e .
+```
+
+If you encounter problems, refer to the [vLLM official compilation guide](https://docs.vllm.ai/en/latest/getting_started/installation/index.html).
+
+### Configuration
+
+#### Prepare configuration for RDMA
+
+Create a `mooncake.json` file:
+
+```json
+{
+    "local_hostname": "192.168.0.137",
+    "metadata_server": "etcd://192.168.0.137:2379",
+    "protocol": "rdma",
+    "device_name": "erdma_0",
+    "master_server_address": "192.168.0.137:50001"
+}
+```
+
+- `local_hostname`: The IP address of the current node. All prefill and decode instances on the same node can share this config.
+- `metadata_server`: The metadata server. Supports `etcd`, `redis`, and `http` backends.
+- `protocol`: `"rdma"` or `"tcp"`.
+- `device_name`: Required for RDMA. Multiple NICs separated by commas (`"erdma_0,erdma_1"`).
+- `master_server_address`: The IP address and port of the MooncakeStore master daemon.
+
+#### Prepare configuration for TCP
+
+```json
+{
+    "local_hostname": "192.168.0.137",
+    "metadata_server": "etcd://192.168.0.137:2379",
+    "protocol": "tcp",
+    "device_name": "",
+    "master_server_address": "192.168.0.137:50001"
+}
+```
+
+### Run Example
+
+Change the IP addresses and ports according to your environment. `VLLM_USE_V1=0` is required for vLLM V0 backend.
+
+```bash
+# Begin from root of your cloned repo!
+
+# 1. Start the etcd server
+etcd --listen-client-urls http://0.0.0.0:2379 --advertise-client-urls http://localhost:2379
+# You may need to terminate other etcd processes before running the above command
+
+# 2. Start the mooncake_master server
+mooncake_master --port 50001
+# If some vllm instances exit unexpectedly, some connection metadata will be
+# corrupted since they are not properly cleaned. In that case, we recommend
+# you restart the mooncake_master before running another test.
+
+# 3. Run multiple vllm instances
+# kv_producer role
+MOONCAKE_CONFIG_PATH=./mooncake.json VLLM_USE_V1=0 python3 -m vllm.entrypoints.openai.api_server \
+    --model Qwen/Qwen2.5-7B-Instruct-GPTQ-Int4 \
+    --port 8100 \
+    --max-model-len 10000 \
+    --gpu-memory-utilization 0.8 \
+    --kv-transfer-config '{"kv_connector":"MooncakeStoreConnector","kv_role":"kv_producer"}'
+
+CUDA_VISIBLE_DEVICES=1 MOONCAKE_CONFIG_PATH=./mooncake.json VLLM_USE_V1=0 python3 -m vllm.entrypoints.openai.api_server \
+    --model Qwen/Qwen2.5-7B-Instruct-GPTQ-Int4 \
+    --port 8101 \
+    --max-model-len 10000 \
+    --gpu-memory-utilization 0.8 \
+    --kv-transfer-config '{"kv_connector":"MooncakeStoreConnector","kv_role":"kv_producer"}'
+
+CUDA_VISIBLE_DEVICES=2 MOONCAKE_CONFIG_PATH=./mooncake.json VLLM_USE_V1=0 python3 -m vllm.entrypoints.openai.api_server \
+    --model Qwen/Qwen2.5-7B-Instruct-GPTQ-Int4 \
+    --port 8102 \
+    --max-model-len 10000 \
+    --gpu-memory-utilization 0.8 \
+    --kv-transfer-config '{"kv_connector":"MooncakeStoreConnector","kv_role":"kv_producer"}'
+
+CUDA_VISIBLE_DEVICES=3 MOONCAKE_CONFIG_PATH=./mooncake.json VLLM_USE_V1=0 python3 -m vllm.entrypoints.openai.api_server \
+    --model Qwen/Qwen2.5-7B-Instruct-GPTQ-Int4 \
+    --port 8103 \
+    --max-model-len 10000 \
+    --gpu-memory-utilization 0.8 \
+    --kv-transfer-config '{"kv_connector":"MooncakeStoreConnector","kv_role":"kv_producer"}'
+
+# kv_consumer role
+CUDA_VISIBLE_DEVICES=4 MOONCAKE_CONFIG_PATH=./mooncake.json VLLM_USE_V1=0 python3 -m vllm.entrypoints.openai.api_server \
+    --model Qwen/Qwen2.5-7B-Instruct-GPTQ-Int4 \
+    --port 8200 \
+    --max-model-len 10000 \
+    --gpu-memory-utilization 0.8 \
+    --kv-transfer-config '{"kv_connector":"MooncakeStoreConnector","kv_role":"kv_consumer"}'
+
+CUDA_VISIBLE_DEVICES=5 MOONCAKE_CONFIG_PATH=./mooncake.json VLLM_USE_V1=0 python3 -m vllm.entrypoints.openai.api_server \
+    --model Qwen/Qwen2.5-7B-Instruct-GPTQ-Int4 \
+    --port 8201 \
+    --max-model-len 10000 \
+    --gpu-memory-utilization 0.8 \
+    --kv-transfer-config '{"kv_connector":"MooncakeStoreConnector","kv_role":"kv_consumer"}'
+
+CUDA_VISIBLE_DEVICES=6 MOONCAKE_CONFIG_PATH=./mooncake.json VLLM_USE_V1=0 python3 -m vllm.entrypoints.openai.api_server \
+    --model Qwen/Qwen2.5-7B-Instruct-GPTQ-Int4 \
+    --port 8202 \
+    --max-model-len 10000 \
+    --gpu-memory-utilization 0.8 \
+    --kv-transfer-config '{"kv_connector":"MooncakeStoreConnector","kv_role":"kv_consumer"}'
+
+CUDA_VISIBLE_DEVICES=7 MOONCAKE_CONFIG_PATH=./mooncake.json VLLM_USE_V1=0 python3 -m vllm.entrypoints.openai.api_server \
+    --model Qwen/Qwen2.5-7B-Instruct-GPTQ-Int4 \
+    --port 8203 \
+    --max-model-len 10000 \
+    --gpu-memory-utilization 0.8 \
+    --kv-transfer-config '{"kv_connector":"MooncakeStoreConnector","kv_role":"kv_consumer"}'
+```
+
+**Key parameters:**
+- `MOONCAKE_CONFIG_PATH`: Path to the mooncake.json configuration file.
+- `VLLM_USE_MODELSCOPE`: Optional. Remove if you have HuggingFace access.
+- `VLLM_USE_V1=0`: Required since the disaggregated feature is currently only supported on V0 vLLM. You can also `export` this configuration to the env instead of putting it in front of every command.
+- `--model`: The model to use.
+- `--port`: The vllm service port on which to listen.
+- `--max-model-len`: The maximum length of the model.
+- `--tensor-parallel-size` / `-tp`: Supported. All instances should have the same tensor_parallel_size. If running prefill and decode on the same node, set different `CUDA_VISIBLE_DEVICES` (e.g., `CUDA_VISIBLE_DEVICES=0,1` for prefill and `CUDA_VISIBLE_DEVICES=2,3` for decode).
+- `--kv-transfer-config`: Set `kv_connector` to `"MooncakeStoreConnector"`, `kv_role` to `"kv_producer"`, `"kv_consumer"`, or `"kv_both"`.
+- If some vLLM instances exit unexpectedly, connection metadata may be corrupted. Restart `mooncake_master` before another test.
+
+```bash
+# 5. Start the proxy server
+cd vllm
+python3 examples/online_serving/disagg_examples/disagg_proxy_demo.py \
+    --model Qwen/Qwen2.5-7B-Instruct-GPTQ-Int4 \
+    --prefill localhost:8100 localhost:8101 \
+    --decode localhost:8200 localhost:8201 \
+    --port 8000
+```
+
+- `--model`: The model and tokenizer used by the proxy server.
+- `--port`: The proxy server port on which to listen.
+- `--prefill` / `-p`: IP and port of the vllm prefill instances.
+- `--decode` / `-d`: IP and port of the vllm decode instances.
+
+#### Dynamic XpYd Adjustment
+
+To dynamically adjust prefill and decode instances at runtime:
+
+```bash
+export ADMIN_API_KEY="xxxxxxxx"
+
+# or add it before the command:
+ADMIN_API_KEY="xxxxxxxx" python3 vllm/examples/online_serving/disagg_examples/disagg_demo.py \
+    --model Qwen/Qwen2.5-7B-Instruct-GPTQ-Int4 \
+    --prefill localhost:8100 localhost:8101 \
+    --decode localhost:8200 localhost:8201 \
+    --port 8000 \
+    --scheduling round_robin
+
+# Add instances to groups dynamically
+curl -X POST "http://localhost:8000/instances/add" \
+  -H "Content-Type: application/json" \
+  -H "X-API-Key: $ADMIN_API_KEY" \
+  -d '{"type": "prefill", "instance": "localhost:8102"}'
+
+curl -X POST "http://localhost:8000/instances/add" \
+  -H "Content-Type: application/json" \
+  -H "X-API-Key: $ADMIN_API_KEY" \
+  -d '{"type": "prefill", "instance": "localhost:8103"}'
+
+curl -X POST "http://localhost:8000/instances/add" \
+  -H "Content-Type: application/json" \
+  -H "X-API-Key: $ADMIN_API_KEY" \
+  -d '{"type": "decode", "instance": "localhost:8202"}'
+
+curl -X POST "http://localhost:8000/instances/add" \
+  -H "Content-Type: application/json" \
+  -H "X-API-Key: $ADMIN_API_KEY" \
+  -d '{"type": "decode", "instance": "localhost:8203"}'
+
+# Get proxy status
+curl localhost:8000/status | jq
+```
+
+```{note}
+Mooncake team provides this simple round-robin proxy as a demo. In production, you can implement custom global proxy strategies.
+```
+
+**Be sure to change the IP address in the commands.**
+
+### Test
+
+```bash
+curl -s http://localhost:8000/v1/completions \
+  -H "Content-Type: application/json" \
+  -d '{
+    "model": "Qwen/Qwen2.5-7B-Instruct-GPTQ-Int4",
+    "prompt": "San Francisco is a",
+    "max_tokens": 1000
+  }'
+```
+
+- If you are not testing on the proxy server, change `localhost` to the IP address of the proxy server.
+
+---
+
+## Performance
+
+| Scenario | Document |
+|----------|----------|
+| V1 MooncakeStoreConnector vs Redis | [Benchmark V1](../../../performance/vllm-benchmark-results-v1.md) |
+| V0 MooncakeStore vs Redis | [Benchmark V0](../../../performance/vllm-benchmark-results-v0.2.md) |
+
+---
+
+## Troubleshooting
+
+- If you encounter connection issues, check that:
+  - All nodes can reach each other over the network
+  - Firewall rules allow traffic on the specified ports
+  - RDMA devices are properly configured
+  - `mooncake_master` is running and reachable
+- For missing library errors, rebuild `mooncake-transfer-engine` from source
+- If vLLM instances exit unexpectedly, restart `mooncake_master` to clean up corrupted metadata
+- Enable debug logging with `VLLM_LOGGING_LEVEL=DEBUG` for detailed diagnostics
diff --git a/docs/source/getting_started/examples/vllm-integration/vllm-integration-v0.2.md b/docs/source/getting_started/examples/vllm-integration/vllm-integration-v0.2.md
index 1fc0ad17..6a567dd7 100644
--- a/docs/source/getting_started/examples/vllm-integration/vllm-integration-v0.2.md
+++ b/docs/source/getting_started/examples/vllm-integration/vllm-integration-v0.2.md
@@ -1,5 +1,10 @@
 # vLLM V0 Disaggregated Serving Demo
 
+```{admonition} Archived
+:class: warning
+This page has been **consolidated** into the unified [Disaggregated Prefill-Decode](disagg-prefill-decode) guide (see the V0 Legacy section). Please use that guide for up-to-date information.
+```
+
 ## Overview
 This is the latest version of mooncake-transfer-engine integration doc with the vLLM project based on [PR 10502](https://github.com/vllm-project/vllm/pull/10502) and [PR 10884](https://github.com/vllm-project/vllm/pull/10884) (vllm version: v0.6.4.post1/main) to accelerate KVCache transfer for inter-node disaggregated serving scenario. We have run some experiments to obtain some [preview benchmark results](../../../performance/vllm-benchmark-results-v0.2.md). More benchmark results will be released in due time.
 
@@ -16,7 +21,7 @@ pip3 install mooncake-transfer-engine
 ```
 
 Note:
-  - If you encounter problems such as missing `lib*.so`, you should uninstall this package by `pip3 uninstall mooncake-transfer-engine`, and build the binaries manually according to the [instructions](build.md).
+  - If you encounter problems such as missing `lib*.so`, you should uninstall this package by `pip3 uninstall mooncake-transfer-engine`, and build the binaries manually according to the [instructions](../../build.md).
   - For vLLM version <= v0.8.4, it requires mooncake-transfer-engine <= 0.3.3.post2. In the latest release, interface `mooncake_vllm_adaptor` has been deprecated.
 
 ### Install the latest version of vLLM
@@ -89,10 +94,20 @@ etcd --listen-client-urls http://0.0.0.0:2379 --advertise-client-urls http://loc
 # You may need to terminate other etcd processes before running the above command
 
 # 2. Run on the prefilling side (producer role)
-MOONCAKE_CONFIG_PATH=./mooncake.json VLLM_USE_MODELSCOPE=True python3 -m vllm.entrypoints.openai.api_server --model Qwen/Qwen2.5-7B-Instruct-GPTQ-Int4 --port 8100 --max-model-len 10000 --gpu-memory-utilization 0.8 --kv-transfer-config '{"kv_connector":"MooncakeConnector","kv_role":"kv_producer","kv_rank":0,"kv_parallel_size":2,"kv_buffer_size":2e9}'
+MOONCAKE_CONFIG_PATH=./mooncake.json VLLM_USE_MODELSCOPE=True python3 -m vllm.entrypoints.openai.api_server \
+    --model Qwen/Qwen2.5-7B-Instruct-GPTQ-Int4 \
+    --port 8100 \
+    --max-model-len 10000 \
+    --gpu-memory-utilization 0.8 \
+    --kv-transfer-config '{"kv_connector":"MooncakeConnector","kv_role":"kv_producer","kv_rank":0,"kv_parallel_size":2,"kv_buffer_size":2e9}'
 
 # 3. Run on the decoding side (consumer role)
-MOONCAKE_CONFIG_PATH=./mooncake.json VLLM_USE_MODELSCOPE=True python3 -m vllm.entrypoints.openai.api_server --model Qwen/Qwen2.5-7B-Instruct-GPTQ-Int4 --port 8200 --max-model-len 10000 --gpu-memory-utilization 0.8 --kv-transfer-config '{"kv_connector":"MooncakeConnector","kv_role":"kv_consumer","kv_rank":1,"kv_parallel_size":2,"kv_buffer_size":2e9}'
+MOONCAKE_CONFIG_PATH=./mooncake.json VLLM_USE_MODELSCOPE=True python3 -m vllm.entrypoints.openai.api_server \
+    --model Qwen/Qwen2.5-7B-Instruct-GPTQ-Int4 \
+    --port 8200 \
+    --max-model-len 10000 \
+    --gpu-memory-utilization 0.8 \
+    --kv-transfer-config '{"kv_connector":"MooncakeConnector","kv_role":"kv_consumer","kv_rank":1,"kv_parallel_size":2,"kv_buffer_size":2e9}'
 ```
 
 - `MOONCAKE_CONFIG_PATH` is the path to the mooncake.json configuration file.
diff --git a/docs/source/getting_started/examples/vllm-integration/vllm-integration-v0.3.md b/docs/source/getting_started/examples/vllm-integration/vllm-integration-v0.3.md
index 82168865..42e34106 100644
--- a/docs/source/getting_started/examples/vllm-integration/vllm-integration-v0.3.md
+++ b/docs/source/getting_started/examples/vllm-integration/vllm-integration-v0.3.md
@@ -1,5 +1,10 @@
 # vLLM V0 Disaggregated Serving with MooncakeStore
 
+```{admonition} Archived
+:class: warning
+This page has been **consolidated** into the unified [KV Cache Storage & Sharing](kv-cache-storage) guide (see the V0 Legacy section). Please use that guide for up-to-date information.
+```
+
 ## Overview
 This is the latest version of the MooncakeStore integration doc with the vLLM project based on [PR 10502](https://github.com/vllm-project/vllm/pull/10502) and [PR 12957](https://github.com/vllm-project/vllm/pull/12957) to support KVCache transfer for intra-node and inter-node disaggregated serving scenario. Benchmark results will be released soon.
 
@@ -24,7 +29,7 @@ pip3 install mooncake-transfer-engine
 ```
 
 Note:
-  - If you encounter problems such as missing `lib*.so`, you should uninstall this package by `pip3 uninstall mooncake-transfer-engine`, and build the binaries manually according to the [instructions](build.md).
+  - If you encounter problems such as missing `lib*.so`, you should uninstall this package by `pip3 uninstall mooncake-transfer-engine`, and build the binaries manually according to the [instructions](../../build.md).
 - For vLLM version <= v0.8.4, it requires mooncake-transfer-engine <= 0.3.3.post2. In the latest release, interface `mooncake_vllm_adaptor` has been deprecated.
 
 ### Install the latest version of vLLM
@@ -91,22 +96,62 @@ mooncake_master --port 50001
 
 # 3. Run multiple vllm instances
 # kv_producer role
-MOONCAKE_CONFIG_PATH=./mooncake.json VLLM_USE_V1=0 python3 -m vllm.entrypoints.openai.api_server --model Qwen/Qwen2.5-7B-Instruct-GPTQ-Int4 --port 8100 --max-model-len 10000 --gpu-memory-utilization 0.8 --kv-transfer-config '{"kv_connector":"MooncakeStoreConnector","kv_role":"kv_producer"}'
+MOONCAKE_CONFIG_PATH=./mooncake.json VLLM_USE_V1=0 python3 -m vllm.entrypoints.openai.api_server \
+    --model Qwen/Qwen2.5-7B-Instruct-GPTQ-Int4 \
+    --port 8100 \
+    --max-model-len 10000 \
+    --gpu-memory-utilization 0.8 \
+    --kv-transfer-config '{"kv_connector":"MooncakeStoreConnector","kv_role":"kv_producer"}'
 
-CUDA_VISIBLE_DEVICES=1 MOONCAKE_CONFIG_PATH=./mooncake.json VLLM_USE_V1=0 python3 -m vllm.entrypoints.openai.api_server --model Qwen/Qwen2.5-7B-Instruct-GPTQ-Int4 --port 8101 --max-model-len 10000 --gpu-memory-utilization 0.8 --kv-transfer-config '{"kv_connector":"MooncakeStoreConnector","kv_role":"kv_producer"}'
+CUDA_VISIBLE_DEVICES=1 MOONCAKE_CONFIG_PATH=./mooncake.json VLLM_USE_V1=0 python3 -m vllm.entrypoints.openai.api_server \
+    --model Qwen/Qwen2.5-7B-Instruct-GPTQ-Int4 \
+    --port 8101 \
+    --max-model-len 10000 \
+    --gpu-memory-utilization 0.8 \
+    --kv-transfer-config '{"kv_connector":"MooncakeStoreConnector","kv_role":"kv_producer"}'
 
-CUDA_VISIBLE_DEVICES=2 MOONCAKE_CONFIG_PATH=./mooncake.json VLLM_USE_V1=0 python3 -m vllm.entrypoints.openai.api_server --model Qwen/Qwen2.5-7B-Instruct-GPTQ-Int4 --port 8102 --max-model-len 10000 --gpu-memory-utilization 0.8 --kv-transfer-config '{"kv_connector":"MooncakeStoreConnector","kv_role":"kv_producer"}'
+CUDA_VISIBLE_DEVICES=2 MOONCAKE_CONFIG_PATH=./mooncake.json VLLM_USE_V1=0 python3 -m vllm.entrypoints.openai.api_server \
+    --model Qwen/Qwen2.5-7B-Instruct-GPTQ-Int4 \
+    --port 8102 \
+    --max-model-len 10000 \
+    --gpu-memory-utilization 0.8 \
+    --kv-transfer-config '{"kv_connector":"MooncakeStoreConnector","kv_role":"kv_producer"}'
 
-CUDA_VISIBLE_DEVICES=3 MOONCAKE_CONFIG_PATH=./mooncake.json VLLM_USE_V1=0 python3 -m vllm.entrypoints.openai.api_server --model Qwen/Qwen2.5-7B-Instruct-GPTQ-Int4 --port 8103 --max-model-len 10000 --gpu-memory-utilization 0.8 --kv-transfer-config '{"kv_connector":"MooncakeStoreConnector","kv_role":"kv_producer"}'
+CUDA_VISIBLE_DEVICES=3 MOONCAKE_CONFIG_PATH=./mooncake.json VLLM_USE_V1=0 python3 -m vllm.entrypoints.openai.api_server \
+    --model Qwen/Qwen2.5-7B-Instruct-GPTQ-Int4 \
+    --port 8103 \
+    --max-model-len 10000 \
+    --gpu-memory-utilization 0.8 \
+    --kv-transfer-config '{"kv_connector":"MooncakeStoreConnector","kv_role":"kv_producer"}'
 
 # kv_consumer role
-CUDA_VISIBLE_DEVICES=4 MOONCAKE_CONFIG_PATH=./mooncake.json VLLM_USE_V1=0 python3 -m vllm.entrypoints.openai.api_server --model Qwen/Qwen2.5-7B-Instruct-GPTQ-Int4 --port 8200 --max-model-len 10000 --gpu-memory-utilization 0.8 --kv-transfer-config '{"kv_connector":"MooncakeStoreConnector","kv_role":"kv_consumer"}'
+CUDA_VISIBLE_DEVICES=4 MOONCAKE_CONFIG_PATH=./mooncake.json VLLM_USE_V1=0 python3 -m vllm.entrypoints.openai.api_server \
+    --model Qwen/Qwen2.5-7B-Instruct-GPTQ-Int4 \
+    --port 8200 \
+    --max-model-len 10000 \
+    --gpu-memory-utilization 0.8 \
+    --kv-transfer-config '{"kv_connector":"MooncakeStoreConnector","kv_role":"kv_consumer"}'
 
-CUDA_VISIBLE_DEVICES=5 MOONCAKE_CONFIG_PATH=./mooncake.json VLLM_USE_V1=0 python3 -m vllm.entrypoints.openai.api_server --model Qwen/Qwen2.5-7B-Instruct-GPTQ-Int4 --port 8201 --max-model-len 10000 --gpu-memory-utilization 0.8 --kv-transfer-config '{"kv_connector":"MooncakeStoreConnector","kv_role":"kv_consumer"}'
+CUDA_VISIBLE_DEVICES=5 MOONCAKE_CONFIG_PATH=./mooncake.json VLLM_USE_V1=0 python3 -m vllm.entrypoints.openai.api_server \
+    --model Qwen/Qwen2.5-7B-Instruct-GPTQ-Int4 \
+    --port 8201 \
+    --max-model-len 10000 \
+    --gpu-memory-utilization 0.8 \
+    --kv-transfer-config '{"kv_connector":"MooncakeStoreConnector","kv_role":"kv_consumer"}'
 
-CUDA_VISIBLE_DEVICES=6 MOONCAKE_CONFIG_PATH=./mooncake.json VLLM_USE_V1=0 python3 -m vllm.entrypoints.openai.api_server --model Qwen/Qwen2.5-7B-Instruct-GPTQ-Int4 --port 8202 --max-model-len 10000 --gpu-memory-utilization 0.8 --kv-transfer-config '{"kv_connector":"MooncakeStoreConnector","kv_role":"kv_consumer"}'
+CUDA_VISIBLE_DEVICES=6 MOONCAKE_CONFIG_PATH=./mooncake.json VLLM_USE_V1=0 python3 -m vllm.entrypoints.openai.api_server \
+    --model Qwen/Qwen2.5-7B-Instruct-GPTQ-Int4 \
+    --port 8202 \
+    --max-model-len 10000 \
+    --gpu-memory-utilization 0.8 \
+    --kv-transfer-config '{"kv_connector":"MooncakeStoreConnector","kv_role":"kv_consumer"}'
 
-CUDA_VISIBLE_DEVICES=7 MOONCAKE_CONFIG_PATH=./mooncake.json VLLM_USE_V1=0 python3 -m vllm.entrypoints.openai.api_server --model Qwen/Qwen2.5-7B-Instruct-GPTQ-Int4 --port 8203 --max-model-len 10000 --gpu-memory-utilization 0.8 --kv-transfer-config '{"kv_connector":"MooncakeStoreConnector","kv_role":"kv_consumer"}'
+CUDA_VISIBLE_DEVICES=7 MOONCAKE_CONFIG_PATH=./mooncake.json VLLM_USE_V1=0 python3 -m vllm.entrypoints.openai.api_server \
+    --model Qwen/Qwen2.5-7B-Instruct-GPTQ-Int4 \
+    --port 8203 \
+    --max-model-len 10000 \
+    --gpu-memory-utilization 0.8 \
+    --kv-transfer-config '{"kv_connector":"MooncakeStoreConnector","kv_role":"kv_consumer"}'
 ```
 
 - `MOONCAKE_CONFIG_PATH` is the path to the mooncake.json configuration file.
@@ -127,7 +172,11 @@ CUDA_VISIBLE_DEVICES=7 MOONCAKE_CONFIG_PATH=./mooncake.json VLLM_USE_V1=0 python
 ```bash
 # 4. Start the proxy server
 cd vllm
-python3 examples/online_serving/disagg_examples/disagg_proxy_demo.py --model Qwen/Qwen2.5-7B-Instruct-GPTQ-Int4 --prefill localhost:8100 localhost:8101  --decode localhost:8200 localhost:8201  --port 8000
+python3 examples/online_serving/disagg_examples/disagg_proxy_demo.py \
+    --model Qwen/Qwen2.5-7B-Instruct-GPTQ-Int4 \
+    --prefill localhost:8100 localhost:8101 \
+    --decode localhost:8200 localhost:8201 \
+    --port 8000
 ```
 
 - The `--model` parameter specifies the model to use, also specifies the tokenizer used by the proxy server.
@@ -139,7 +188,12 @@ python3 examples/online_serving/disagg_examples/disagg_proxy_demo.py --model Qwe
 # If you want to dynamically adjust the instances of p-nodes and d-nodes during runtime, you need to configure this environment variables.
 export ADMIN_API_KEY="xxxxxxxx"
 # or add it before the command:
-ADMIN_API_KEY="xxxxxxxx" python3 vllm/examples/online_serving/disagg_examples/disagg_demo.py --model Qwen/Qwen2.5-7B-Instruct-GPTQ-Int4 --prefill localhost:8100 localhost:8101  --decode localhost:8200 localhost:8201  --port 8000 --scheduling round_robin
+ADMIN_API_KEY="xxxxxxxx" python3 vllm/examples/online_serving/disagg_examples/disagg_demo.py \
+    --model Qwen/Qwen2.5-7B-Instruct-GPTQ-Int4 \
+    --prefill localhost:8100 localhost:8101 \
+    --decode localhost:8200 localhost:8201 \
+    --port 8000 \
+    --scheduling round_robin
 
 # Then use this command to add instances into prefill group or decode group
 curl -X POST "http://localhost:8000/instances/add" -H "Content-Type: application/json" -H "X-API-Key: $ADMIN_API_KEY" -d '{"type": "prefill", "instance": "localhost:8102"}'
diff --git a/docs/source/getting_started/examples/vllm-integration/vllm-integration-v1.0.md b/docs/source/getting_started/examples/vllm-integration/vllm-integration-v1.0.md
index 3b140df6..6e7ad79e 100644
--- a/docs/source/getting_started/examples/vllm-integration/vllm-integration-v1.0.md
+++ b/docs/source/getting_started/examples/vllm-integration/vllm-integration-v1.0.md
@@ -1,5 +1,10 @@
 # vLLM v1 backend Disaggregated Serving with MooncakeConnector
 
+```{admonition} Archived
+:class: warning
+This page has been **consolidated** into the unified [Disaggregated Prefill-Decode](disagg-prefill-decode) guide. Please use that guide for up-to-date information.
+```
+
 ## Overview
 
 This guide demonstrates how to use the MooncakeConnector with vLLM v1 backend for disaggregated serving in Prefill-Decode separation architecture. The integration enables efficient cross-node KV cache transfer using RDMA technology.
@@ -16,7 +21,7 @@ Install mooncake-transfer-engine through pip:
 pip install mooncake-transfer-engine
 ```
 
-Note: If you encounter problems such as missing `lib*.so`, you should uninstall this package by `pip3 uninstall mooncake-transfer-engine`, and build the binaries manually according to the [instructions](../build.md).
+Note: If you encounter problems such as missing `lib*.so`, you should uninstall this package by `pip3 uninstall mooncake-transfer-engine`, and build the binaries manually according to the [instructions](../../build.md).
 
 ### Install vLLM
 
@@ -100,7 +105,7 @@ vllm serve Qwen/Qwen2.5-7B-Instruct \
     - `kv_producer`: For prefiller instances that generate KV caches
     - `kv_consumer`: For decoder instances that consume KV caches
     - `kv_both`: Enables symmetric functionality (experimental)
-    - `num_workers`: Thread pool size in each prefiller worker to send kvcache (default 10)
+  - `num_workers`: Thread pool size in each prefiller worker to send kvcache (default 10)
 
 ## Environment Variables
 
@@ -118,7 +123,7 @@ The following environment variables can be used to customize Mooncake behavior:
 
 ## Performance
 
-For detailed performance benchmarks and results, see the [vLLM Benchmark](../../performance/vllm-v1-support-benchmark.md) documentation.
+For detailed performance benchmarks and results, see the [vLLM Benchmark](../../../performance/vllm-v1-support-benchmark.md) documentation.
 
 ## Notes
 
diff --git a/docs/source/getting_started/examples/vllm-integration/vllm-mooncakestoreconnector.md b/docs/source/getting_started/examples/vllm-integration/vllm-mooncakestoreconnector.md
new file mode 100644
index 00000000..d76f5646
--- /dev/null
+++ b/docs/source/getting_started/examples/vllm-integration/vllm-mooncakestoreconnector.md
@@ -0,0 +1,135 @@
+# Guide: vLLM MooncakeStoreConnector
+
+```{admonition} Archived
+:class: warning
+This page has been **consolidated** into the unified [KV Cache Storage & Sharing](kv-cache-storage) guide (see the V1 Recommended section). Please use that guide for up-to-date information.
+```
+
+## Overview
+
+This document describes how to deploy vLLM's `MooncakeStoreConnector`. `MooncakeStoreConnector`  is a new vLLM's KV connector that uses `MooncakeDistributedStore` as a shared KV cache pool. It enables:
+
+* **CPU/Disk offloading**: Extend effective KV cache capacity by offloading to CPU memory or SSD via Mooncake's transfer engine.
+* **Hash-based prefix caching across instances**: Multiple vLLM instances share cached KV blocks through the store using block-hash deduplication.
+* **Flexible deployment**: Works as a single-node KV cache extension (`kv_both`), or in disaggregated prefill-decode setups (`kv_producer` / `kv_consumer`).
+
+
+
+## Deployment
+
+### 1. Prerequisites
+
+Before you begin, make sure that:
+
+* [vLLM](https://github.com/vllm-project/vllm) is installed, [Mooncake](https://github.com/kvcache-ai/Mooncake) is installed. Refer to the [vLLM official repository](https://github.com/vllm-project/vllm) and [Mooncake official repository](https://github.com/kvcache-ai/Mooncake) for more installation instructions and building from source.
+
+
+
+### 2. Mooncake Master Server
+
+**Start:**
+
+```shell
+mooncake_master --port 50063
+```
+
+**Configure Mooncake** : Create a JSON configuration file (e.g., `mooncake_config.json`):
+
+```json
+{
+  "metadata_server": "http://127.0.0.1:8092/metadata",
+  "master_server_address": "127.0.0.1:50063",
+  "global_segment_size": "0",
+  "local_buffer_size": "2147483648",
+  "protocol": "rdma",
+  "device_name": "",
+}
+```
+
+**Set environment variable:**
+
+```shell
+export MOONCAKE_CONFIG_PATH=/path/to/mooncake_config.json
+```
+
+
+
+### 3. Usage
+
+**3.1** **Single-Node KV Cache Offloading** (i.e., `kv_both`)
+
+```shell
+MOONCAKE_CONFIG_PATH=mooncake_config.json \
+vllm serve meta-llama/Llama-3.1-8B-Instruct \
+    --kv-transfer-config '{"kv_connector":"MooncakeStoreConnector","kv_role":"kv_both"}'
+```
+
+**3.2** **XpYd Disaggregated Prefill-Decode** (i.e., `kv_producer/kv_consumer`)
+
+Prefill Node:
+
+```shell
+MOONCAKE_CONFIG_PATH=mooncake_config.json \
+VLLM_MOONCAKE_BOOTSTRAP_PORT=50052 \
+vllm serve meta-llama/Llama-3.1-8B-Instruct \
+    --port 8100 \
+    --kv-transfer-config '{
+        "kv_connector": "MultiConnector",
+        "kv_role": "kv_producer",
+        "kv_connector_extra_config": {
+            "connectors": [
+                {
+                    "kv_connector": "MooncakeConnector",
+                    "kv_role": "kv_producer"
+                },
+                {
+                    "kv_connector": "MooncakeStoreConnector",
+                    "kv_role": "kv_producer"
+                }
+            ]
+        }
+    }'
+```
+
+Decode Node:
+
+```shell
+MOONCAKE_CONFIG_PATH=mooncake_config.json \
+VLLM_MOONCAKE_BOOTSTRAP_PORT=50053 \
+vllm serve meta-llama/Llama-3.1-8B-Instruct \
+    --port 8200 \
+    --kv-transfer-config '{
+        "kv_connector": "MultiConnector",
+        "kv_role": "kv_consumer",
+        "kv_connector_extra_config": {
+            "connectors": [
+                {
+                    "kv_connector": "MooncakeConnector",
+                    "kv_role": "kv_consumer"
+                },
+                {
+                    "kv_connector": "MooncakeStoreConnector",
+                    "kv_role": "kv_consumer"
+                }
+            ]
+        }
+    }'
+```
+
+Proxy:
+
+```shell
+python examples/disaggregated/disaggregated_serving/mooncake_connector/mooncake_connector_proxy.py \
+    --prefill http://192.168.0.2:8100 \
+    --decode http://192.168.0.3:8200
+```
+
+
+
+> When running with data parallelism, set a fixed `PYTHONHASHSEED` so that block hashes are consistent across DP ranks:
+>
+> ```shell
+> PYTHONHASHSEED=0 vllm serve ...
+> ```
+>
+> Without this, identical prompts may produce different block hashes on different DP ranks, preventing cross-instance prefix cache hits.
diff --git a/docs/source/getting_started/examples/vllm-integration/vllmv1-lmcache-integration.md b/docs/source/getting_started/examples/vllm-integration/vllmv1-lmcache-integration.md
index 676841e8..024bd760 100644
--- a/docs/source/getting_started/examples/vllm-integration/vllmv1-lmcache-integration.md
+++ b/docs/source/getting_started/examples/vllm-integration/vllmv1-lmcache-integration.md
@@ -2,43 +2,73 @@
 
 ## Overview
 
-The vLLM v1 version has been released with support for PD disaggregation. The detailed design document can be found [here](https://docs.google.com/document/d/1uPGdbEXksKXeN4Q9nUm9hzotqEjQhYmnpAhidLuAsjk). LMCache immediately implemented the corresponding connector to support storage, transmission, and loading of KVCache, enabling collaborative operation with PD nodes. Mooncake, as LMCache's backend storage engine, has undergone extensive optimizations in usability, performance, and stability. This document explains how to deploy a PD disaggregated serving demo using LMCache + Mooncake.
+This guide shows a two-machine 1-prefill/1-decode deployment using vLLM V1,
+LMCache's non-MP `LMCacheConnectorV1`, and Mooncake Store as LMCache's remote
+storage backend.
+
+LMCache supports both non-MP mode and MP mode with Mooncake Store. This page
+covers the non-MP path, where each vLLM instance loads an LMCache YAML config
+through `LMCACHE_CONFIG_FILE` and connects directly to Mooncake Store with
+`remote_url: "mooncakestore://..."`. For the LMCache multiprocess server path,
+see [vLLM V1 Disaggregated Serving with Mooncake Store and LMCache \[MP\]](vllmv1-lmcache-mp-integration.md).
+
+The examples below use:
+
+- Machine A: Mooncake master and vLLM decoder
+- Machine B: vLLM prefiller
+- Mooncake master RPC address: `{IP of Machine A}:50051`
+- Mooncake HTTP metadata endpoint: `http://{IP of Machine A}:8080/metadata`
+- RDMA device: `{RDMA device}`
+
+Replace these placeholders with the IP addresses, hostname, and RDMA device for
+your environment.
+
+## Prerequisites
+
+Install Mooncake, vLLM, and LMCache on both machines. For installation details,
+refer to the official documentation of each project:
+
+- [Mooncake build guide](../../build.md)
+- [LMCache installation](https://docs.lmcache.ai/getting_started/installation.html)
+- [vLLM installation](https://docs.vllm.ai/en/latest/getting_started/installation/)
 
 ## Deployment
 
-1. First, you need to prepare two GPU-equipped machines, which we will refer to as Machine A and Machine B. Install [vLLM](https://docs.vllm.ai/en/latest/getting_started/quickstart.html), [Mooncake](https://kvcache-ai.github.io/Mooncake/getting_started/build.html) and [LMCache](https://docs.lmcache.ai/getting_started/installation.html) on both Machine A and Machine B. For specific installation instructions, please refer to the official documentation of each repository.
+### 1. Start Mooncake Master on Machine A
 
-2. Start the Mooncake Master node on Machine A:
 ```bash
-mooncake_master -port 50052 -max_threads 64 -metrics_port 9004 \
+mooncake_master -v=1 \
+  --rpc_port=50051 \
+  --metrics_port=9003 \
   --enable_http_metadata_server=true \
   --http_metadata_server_host=0.0.0.0 \
   --http_metadata_server_port=8080
 ```
 
-3. Launch the Decoder instance on machine A
-- Modify the vllm/examples/others/lmcache/disagg_prefill_lmcache_v1/disagg_vllm_launcher.sh file.
+### 2. Configure and Start the vLLM Decoder on Machine A
+
+Modify the vLLM disaggregated prefill launcher to use a Mooncake-backed LMCache
+config for the decoder.
+
+The decoder should continue to use `LMCacheConnectorV1` with `kv_role` set to
+`kv_consumer`.
+
 ```diff
 diff --git a/examples/lmcache/disagg_prefill_lmcache_v1/disagg_vllm_launcher.sh b/examples/lmcache/disagg_prefill_lmcache_v1/disagg_vllm_launcher.sh
-index 831ef0bb5..a2ff0744c 100644
 --- a/examples/lmcache/disagg_prefill_lmcache_v1/disagg_vllm_launcher.sh
 +++ b/examples/lmcache/disagg_prefill_lmcache_v1/disagg_vllm_launcher.sh
+@@
  elif [[ $1 == "decoder" ]]; then
      # Decoder listens on port 8200
 -    decode_config_file=$SCRIPT_DIR/configs/lmcache-decoder-config.yaml
 +    decode_config_file=$SCRIPT_DIR/configs/mooncake-decoder-config.yaml
- 
-     UCX_TLS=cuda_ipc,cuda_copy,tcp \
-         LMCACHE_CONFIG_FILE=$decode_config_file \
-         LMCACHE_USE_EXPERIMENTAL=True \
-         VLLM_ENABLE_V1_MULTIPROCESSING=1 \
-         VLLM_WORKER_MULTIPROC_METHOD=spawn \
-         CUDA_VISIBLE_DEVICES=1 \
 ```
-- Add the `mooncake-decoder-config.yaml` file
+
+Create `configs/mooncake-decoder-config.yaml`:
+
 ```yaml
 chunk_size: 256
-remote_url: "mooncakestore://{IP of Machine A}:50052/"
+remote_url: "mooncakestore://{IP of Machine A}:50051/"
 remote_serde: "naive"
 local_cpu: False
 max_local_cpu_size: 100
@@ -47,38 +77,43 @@ extra_config:
   local_hostname: "{IP of Machine A}"
   metadata_server: "http://{IP of Machine A}:8080/metadata"
   protocol: "rdma"
-  device_name: "mlx5_0" # Multiple RDMA devices can be specified as comma-separated list
-  master_server_address: "{IP of Machine A}:50052"
+  device_name: "{RDMA device}"
+  master_server_address: "{IP of Machine A}:50051"
   global_segment_size: 32212254720 # 30GB
   local_buffer_size: 1073741824 # 1GB
   transfer_timeout: 1
   save_chunk_meta: False
 ```
 
-- Launch the Decoder instance using command 
+Launch the decoder:
+
 ```bash
 bash disagg_vllm_launcher.sh decoder Qwen/Qwen2.5-7B-Instruct-GPTQ-Int4
 ```
 
-4. Launch the Prefiller instance on machine B
-- Modify the vllm/examples/others/lmcache/disagg_prefill_lmcache_v1/disagg_vllm_launcher.sh file.
+### 3. Configure and Start the vLLM Prefiller on Machine B
+
+Modify the launcher to use a Mooncake-backed LMCache config for the prefiller.
+
+The prefiller should continue to use `LMCacheConnectorV1` with `kv_role` set to
+`kv_producer`.
+
 ```diff
 diff --git a/examples/lmcache/disagg_prefill_lmcache_v1/disagg_vllm_launcher.sh b/examples/lmcache/disagg_prefill_lmcache_v1/disagg_vllm_launcher.sh
-index 831ef0bb5..9e5a3f044 100644
 --- a/examples/lmcache/disagg_prefill_lmcache_v1/disagg_vllm_launcher.sh
 +++ b/examples/lmcache/disagg_prefill_lmcache_v1/disagg_vllm_launcher.sh
-@@ -18,12 +18,14 @@ fi
- 
+@@
  if [[ $1 == "prefiller" ]]; then
      # Prefiller listens on port 8100
 -    prefill_config_file=$SCRIPT_DIR/configs/lmcache-prefiller-config.yaml
 +    prefill_config_file=$SCRIPT_DIR/configs/mooncake-prefiller-config.yaml
 ```
 
-- Add the `mooncake-prefiller-config.yaml` file
+Create `configs/mooncake-prefiller-config.yaml`:
+
 ```yaml
 chunk_size: 256
-remote_url: "mooncakestore://{IP of Machine A}:50052/"
+remote_url: "mooncakestore://{IP of Machine A}:50051/"
 remote_serde: "naive"
 local_cpu: False
 max_local_cpu_size: 100
@@ -87,33 +122,70 @@ extra_config:
   local_hostname: "{IP of Machine B}"
   metadata_server: "http://{IP of Machine A}:8080/metadata"
   protocol: "rdma"
-  device_name: "mlx5_0" # Multiple RDMA devices can be specified as comma-separated list
-  master_server_address: "{IP of Machine A}:50052"
+  device_name: "{RDMA device}"
+  master_server_address: "{IP of Machine A}:50051"
   global_segment_size: 32212254720 # 30GB
   local_buffer_size: 1073741824 # 1GB
   transfer_timeout: 1
   save_chunk_meta: False
 ```
 
-- Launch the Prefiller instance using command 
+Launch the prefiller:
+
 ```bash
 bash disagg_vllm_launcher.sh prefiller Qwen/Qwen2.5-7B-Instruct-GPTQ-Int4
 ```
 
-5. Prepare the router `disagg_proxy_server`
+### 4. Start the Disaggregated Proxy
 
-We use the [disagg_proxy_server](https://github.com/LMCache/LMCache/blob/dev/examples/disagg_prefill/disagg_proxy_server.py) provided by LMCache. According to [LMCache/LMCache#1342](https://github.com/LMCache/LMCache/issues/1342), when using Mooncake Store as the backend, you need to comment out `wait_decode_kv_ready(req_id)` in the proxy code.
-
-6. Launch the `disagg_proxy_server` using command
+Use the LMCache [`disagg_proxy_server.py`](https://github.com/LMCache/LMCache/blob/dev/examples/disagg_prefill/disagg_proxy_server.py) to route requests between the prefiller and decoder. According to [LMCache/LMCache#1342](https://github.com/LMCache/LMCache/issues/1342), when using Mooncake Store as the backend, comment out the `wait_decode_kv_ready(...)` call in the proxy before starting it.
 
 ```bash
-python3 disagg_proxy_server.py --host localhost --port 9000 --prefiller-host IP_of_Machine_B --prefiller-port 8100 --decoder-host IP_of_Machine_A --decoder-port 8200 
+python3 disagg_proxy_server.py \
+  --host 0.0.0.0 \
+  --port 9000 \
+  --prefiller-host {IP of Machine B} \
+  --prefiller-port 8100 \
+  --decoder-host {IP of Machine A} \
+  --decoder-port 8200
 ```
 
-7. Now we can send the requests to the `disagg_proxy_server` to test PD disaggregated serving.
+### 5. Send a Test Request
+
+Send traffic to the proxy, not directly to either vLLM instance.
+
+```bash
+curl -N http://{Proxy IP}:9000/v1/chat/completions \
+  -H "Content-Type: application/json" \
+  -d '{
+    "model": "Qwen/Qwen2.5-7B-Instruct-GPTQ-Int4",
+    "messages": [
+      {
+        "role": "user",
+        "content": "Explain how KV cache reuse helps long-context serving."
+      }
+    ],
+    "max_tokens": 128,
+    "temperature": 0.7
+  }'
+```
+
+## Port and Configuration Checklist
+
+When changing ports away from these defaults, update all dependent settings
+together:
+
+- Mooncake master `--rpc_port` must match `remote_url` and
+  `extra_config.master_server_address`.
+- `extra_config.metadata_server` must point to the Mooncake HTTP metadata
+  endpoint when HTTP metadata is used.
+- Decoder and prefiller `device_name`, `protocol`, `global_segment_size`, and
+  `local_buffer_size` should be set for the local hardware and workload.
+- Proxy `--prefiller-port` and `--decoder-port` must match the two vLLM
+  instance ports.
 
 ## Additional Resources
 
-* [Mooncake x LMCache: Unite to Pioneer KVCache-Centric LLM Serving System](../../../getting_started/examples/lmcache-integration.md)
+* [Mooncake x LMCache: Unite to Pioneer KVCache-Centric LLM Serving System](../lmcache-integration.md)
 * [Using Mooncake in LMCache](https://docs.lmcache.ai/kv_cache/storage_backends/mooncake.html)
-* [Using LMCache in vLLM](https://github.com/vllm-project/vllm/tree/main/examples/others/lmcache)
\ No newline at end of file
+* [Using LMCache in vLLM](https://github.com/vllm-project/vllm/tree/main/examples/others/lmcache)
diff --git a/docs/source/getting_started/examples/vllm-integration/vllmv1-lmcache-mp-integration.md b/docs/source/getting_started/examples/vllm-integration/vllmv1-lmcache-mp-integration.md
new file mode 100644
index 00000000..d728753a
--- /dev/null
+++ b/docs/source/getting_started/examples/vllm-integration/vllmv1-lmcache-mp-integration.md
@@ -0,0 +1,193 @@
+# vLLM V1 Disaggregated Serving with Mooncake Store and LMCache [MP]
+
+## Overview
+
+This guide shows a single-machine 1-prefill/1-decode deployment using vLLM V1,
+LMCache's multiprocess server, and Mooncake Store as the LMCache L2 backend.
+
+LMCache supports both non-MP mode and MP mode with Mooncake Store. This page
+covers the MP path, where vLLM instances connect to an LMCache server through
+`LMCacheMPConnector`, and the LMCache server connects to Mooncake Store through
+the `mooncake_store` L2 adapter. For the non-MP `LMCacheConnectorV1` path, see
+[vLLM V1 Disaggregated Serving with Mooncake Store and LMCache](vllmv1-lmcache-integration.md).
+
+In this setup, one machine runs Mooncake master, one LMCache MP server, the
+disaggregated proxy, the prefiller vLLM instance, and the decoder vLLM
+instance. The prefiller and decoder should use different GPUs.
+
+This example uses `"metadata_server":"P2PHANDSHAKE"` for Mooncake transfer
+metadata, so the Mooncake HTTP metadata server is not needed. If you switch to
+HTTP metadata, remember that Mooncake's HTTP metadata server also defaults to
+`8080`, which conflicts with LMCache's HTTP API on a single host.
+
+## Prerequisites
+
+Install Mooncake, vLLM, and LMCache on the machine. The example assumes an RDMA
+deployment and uses:
+
+- local host address: `{IP of Machine}`
+- RDMA device: `{RDMA device}`
+- LMCache checkout path: `/path/to/LMCache`
+
+Replace these values with the local hostname/IP, RDMA device, and LMCache
+checkout path for your environment.
+
+LMCache requirement: This example requires LMCache v0.4.5 or later. LMCache
+must also be built from source with Mooncake support enabled, because the
+`mooncake_store` MP L2 adapter depends on the optional
+`lmcache.lmcache_mooncake` C++ extension. 
+
+The standard prebuilt LMCache wheels currently include the Python Mooncake
+adapter files, but do not include the optional `lmcache.lmcache_mooncake`
+native extension.
+
+```bash
+BUILD_MOONCAKE=1 \
+MOONCAKE_INCLUDE_DIR=/path/to/mooncake/include \
+MOONCAKE_LIB_DIR=/path/to/mooncake/lib \
+pip install -e /path/to/LMCache --verbose
+```
+
+## Deployment
+
+### 1. Start Mooncake Master
+
+```bash
+mooncake_master -v=1 \
+  --rpc_port=50051 \
+  --metrics_port=9003
+```
+
+### 2. Start the LMCache Multiprocess Server
+
+Start one LMCache MP server and configure Mooncake Store as the L2 adapter.
+
+```bash
+lmcache server \
+  --host 127.0.0.1 \
+  --port 5555 \
+  --http-host 127.0.0.1 \
+  --http-port 8080 \
+  --l1-size-gb 32 \
+  --eviction-policy LRU \
+  --no-l1-use-lazy \
+  --l2-adapter '{
+    "type": "mooncake_store",
+    "local_hostname": "{IP of Machine}",
+    "metadata_server": "P2PHANDSHAKE",
+    "protocol": "rdma",
+    "rdma_devices": "{RDMA device}",
+    "global_segment_size": 32212254720,
+    "local_buffer_size": 1073741824,
+    "master_server_addr": "127.0.0.1:50051"
+  }'
+```
+
+### 3. Start the Disaggregated Proxy
+
+The proxy receives client requests, sends prefill requests to the prefiller,
+sends decode requests to the decoder, and receives LMCache request telemetry
+from the prefiller.
+
+```bash
+python /path/to/LMCache/examples/disagg_prefill_mp/disagg_proxy_server.py \
+  --host 127.0.0.1 \
+  --port 8000 \
+  --prefiller-host 127.0.0.1 \
+  --prefiller-port 8100 \
+  --decoder-host 127.0.0.1 \
+  --decoder-port 8200 \
+  --telemetry-port 5768
+```
+
+### 4. Start the vLLM Prefiller
+
+The prefiller reports request telemetry back to the proxy so the proxy knows
+when KV cache storage has completed.
+
+```bash
+CUDA_VISIBLE_DEVICES=0 \
+LMCACHE_REQUEST_TELEMETRY_TYPE=fastapi \
+LMCACHE_REQUEST_TELEMETRY_ENDPOINT=http://127.0.0.1:5768/api/v1/telemetry \
+vllm serve Qwen/Qwen3-4B \
+  --host 127.0.0.1 \
+  --port 8100 \
+  --gpu-memory-utilization 0.8 \
+  --no-enable-log-requests \
+  --no-enable-prefix-caching \
+  --kv-transfer-config '{
+    "kv_connector": "LMCacheMPConnector",
+    "kv_role": "kv_both",
+    "kv_connector_extra_config": {
+      "lmcache.mp.host": "tcp://127.0.0.1",
+      "lmcache.mp.port": 5555
+    }
+  }'
+```
+
+### 5. Start the vLLM Decoder
+
+The decoder connects to the same local LMCache MP server. It does not need
+request telemetry environment variables; only the prefiller reports the "KV
+cache is stored" event back to the proxy.
+
+```bash
+CUDA_VISIBLE_DEVICES=1 \
+vllm serve Qwen/Qwen3-4B \
+  --host 127.0.0.1 \
+  --port 8200 \
+  --gpu-memory-utilization 0.8 \
+  --no-enable-log-requests \
+  --no-enable-prefix-caching \
+  --kv-transfer-config '{
+    "kv_connector": "LMCacheMPConnector",
+    "kv_role": "kv_both",
+    "kv_connector_extra_config": {
+      "lmcache.mp.host": "tcp://127.0.0.1",
+      "lmcache.mp.port": 5555
+    }
+  }'
+```
+
+### 6. Send a Test Request
+
+Send traffic to the proxy, not directly to either vLLM instance.
+
+```bash
+curl -N http://127.0.0.1:8000/v1/chat/completions \
+  -H "Content-Type: application/json" \
+  -d '{
+    "model": "Qwen/Qwen3-4B",
+    "messages": [
+      {
+        "role": "user",
+        "content": "Explain how KV cache reuse helps long-context serving."
+      }
+    ],
+    "max_tokens": 128,
+    "temperature": 0.7
+  }'
+```
+
+## Port and Configuration Checklist
+
+When changing ports away from these defaults, update all dependent settings
+together:
+
+- Mooncake master `--rpc_port` must match LMCache `master_server_addr`.
+- The vLLM prefiller and decoder should connect to the local LMCache MP server
+  via `kv_connector_extra_config.lmcache.mp.host` and
+  `kv_connector_extra_config.lmcache.mp.port`.
+- Proxy `--prefiller-port` and `--decoder-port` must match the two vLLM
+  `--port` values.
+- `LMCACHE_REQUEST_TELEMETRY_ENDPOINT` on the prefiller must point to the proxy
+  telemetry endpoint.
+- If `metadata_server` is changed from `P2PHANDSHAKE` to an HTTP metadata URL,
+  enable Mooncake HTTP metadata server and make sure its port does not conflict
+  with LMCache `--http-port`.
+
+## Additional Resources
+
+* [Mooncake x LMCache: Unite to Pioneer KVCache-Centric LLM Serving System](../lmcache-integration.md)
+* [LMCache MP `mooncake_store` L2 adapter](https://docs.lmcache.ai/mp/l2_storage.html#mooncake-store-mooncake-store-native-connector)
+* [LMCache multiprocess disaggregated prefill example](https://github.com/LMCache/LMCache/tree/dev/examples/disagg_prefill_mp)
diff --git a/docs/source/getting_started/observability.md b/docs/source/getting_started/observability.md
new file mode 100644
index 00000000..fd0e3922
--- /dev/null
+++ b/docs/source/getting_started/observability.md
@@ -0,0 +1,167 @@
+# Observability
+
+This document describes how to monitor and observe a running Mooncake Store deployment.
+
+## Master Metrics Log
+
+When started, the Mooncake master periodically prints a metrics summary log every 10 seconds (configurable via `kMetricReportIntervalSeconds`). This log provides a comprehensive snapshot of the master's runtime state.
+
+### Log Format
+
+```
+I0512 15:03:30.321475 239489 rpc_service.cpp:269] Master Admin Metrics: role=leader, state=serving, service_ready=true, master={...}, ha={...}, leader=127.0.0.1:50051, view_version=1
+```
+
+Each log line contains:
+- **GLog header**: timestamp, thread ID, source file and line number
+- **role**: HA role — `leader` or `standby`
+- **state**: HA runtime state — `serving`, `starting`, `stopping`, etc.
+- **service_ready**: whether the gRPC service is accepting requests
+- **master**: master metrics block (see below)
+- **ha**: HA metrics block
+- **leader**: (only when available) the current leader address and view version
+
+### Master Metrics Block
+
+A typical `master={...}` block looks like this:
+
+```
+Mem Storage: 94.09 MB / 100.00 MB (94.1%) | SSD Storage: 0 B / 0 B | Keys: 16058 (soft-pinned: 0) | Clients: 1 | Requests (Success/Total per sec): PutStart=0.00/0.00, PutEnd=0.00/0.00, PutRevoke=0.00/0.00, Get=0.00/0.00, Exist=0.00/0.00, Del=0.00/0.00, DelAll=0.00/0.00, Ping=1.00/1.00, CopyStart=0.00/0.00, CopyEnd=0.00/0.00, CopyRevoke=0.00/0.00, MoveStart=0.00/0.00, MoveEnd=0.00/0.00, MoveRevoke=0.00/0.00, EvictDiskReplica=0.00/0.00 | Batch Requests (per sec, Req=Success/PartialSuccess/Total, Item=Success/Total): PutStart:(Req=0.00/0.00/0.00, Item=0.00/0.00), PutEnd:(Req=0.00/0.00/0.00, Item=0.00/0.00), PutRevoke:(Req=0.00/0.00/0.00, Item=0.00/0.00), Get:(Req=0.00/0.00/0.00, Item=0.00/0.00), ExistKey:(Req=0.00/0.00/0.00, Item=0.00/0.00), QueryIp:(Req=0.00/0.00/0.00, Item=0.00/0.00), Clear:(Req=0.00/0.00/0.00, Item=0.00/0.00), CreateMoveTask:(Req=0.00/0.00), CreateCopyTask:(Req=0.00/0.00), QueryTask:(Req=0.00/0.00), FetchTasks:(Req=0.00/0.00), MarkTaskToComplete:(Req=0.00/0.00) | Eviction: Success/Attempts=0/0, AllocFail=0, keys=0, size=0 B | Discard: Released/Total=0/0, StagingSize=0 B | Snapshots: Success=0, Fail=0
+```
+
+Request counters are reported as **rates per second** over the time window between two consecutive log outputs (10 seconds by default). Real-time state values (storage, key count, client count, discard staging size) are not rate-limited and reflect the current value at log time.
+
+The metrics block consists of the following sections:
+
+#### Storage
+
+| Field | Description |
+|-------|-------------|
+| `Mem Storage` | Current memory usage / total memory capacity, with percentage |
+| `SSD Storage` | Current SSD-backed storage usage / total SSD capacity |
+
+#### Keys and Clients
+
+| Field | Description |
+|-------|-------------|
+| `Keys` | Total number of keys managed by the master |
+| `soft-pinned` | Number of keys with active soft-pin leases (protected from eviction) |
+| `Clients` | Number of currently connected clients |
+
+#### Requests (Success/Total per sec)
+
+Rate counters for individual (non-batch) RPC requests over the last time window. Each shows `<success_rate>/<total_rate>` in requests per second:
+
+| Counter | Description |
+|---------|-------------|
+| `PutStart` | Put object allocation requests |
+| `PutEnd` | Put object commit requests |
+| `PutRevoke` | Put object cancellation requests |
+| `Get` | Get replica list requests |
+| `Exist` | Key existence check requests |
+| `Del` | Single key deletion requests |
+| `DelAll` | Delete-all objects requests |
+| `Ping` | Client heartbeat/ping requests |
+| `CopyStart` | Copy object allocation requests |
+| `CopyEnd` | Copy object commit requests |
+| `CopyRevoke` | Copy object cancellation requests |
+| `MoveStart` | Move object allocation requests |
+| `MoveEnd` | Move object commit requests |
+| `MoveRevoke` | Move object cancellation requests |
+| `EvictDiskReplica` | Evict disk replica requests |
+
+#### Batch Requests (per sec)
+
+Batch operations aggregate multiple items into a single RPC. Rates are per second over the last time window. Format: `Req=<success>/<partial_success>/<total>`, `Item=<success_items>/<total_items>`:
+
+| Counter | Description |
+|---------|-------------|
+| `PutStart` | Batch put object allocation requests |
+| `PutEnd` | Batch put object commit requests |
+| `PutRevoke` | Batch put object cancellation requests |
+| `Get` | Batch get replica list requests |
+| `ExistKey` | Batch key existence check requests |
+| `QueryIp` | Batch query IP requests |
+| `Clear` | Batch replica clear requests |
+
+A request is considered "partial success" when it succeeds for some items but not all.
+
+#### Task Operations
+
+| Counter | Description |
+|---------|-------------|
+| `CreateMoveTask` | Move task creation requests |
+| `CreateCopyTask` | Copy task creation requests |
+| `QueryTask` | Task status query requests |
+| `FetchTasks` | Pending task fetch requests (polled by store clients) |
+| `MarkTaskToComplete` | Task completion acknowledgement requests |
+
+#### Eviction & Discard
+
+Eviction counters are **deltas** between two consecutive log outputs — they show what happened in the time window, not cumulative totals.
+
+| Field | Description |
+|-------|-------------|
+| `Eviction: Success/Attempts` | Eviction rounds that succeeded at least partially vs. total attempts in this window |
+| `AllocFail` | Number of PutStart/UpsertStart failures caused by replica allocation failure (triggers eviction) in this window |
+| `keys` | Number of keys evicted in this window |
+| `size` | Total size of evicted data in this window |
+| `Discard: Released/Total` | Released (cleaned up) vs. total discarded PutStart staging replicas (live values) |
+| `StagingSize` | Current size of discarded but not-yet-released staging buffers (live value) |
+
+## Prometheus Metrics Endpoint
+
+Mooncake master exposes Prometheus-format metrics at the HTTP admin endpoint. This allows integration with Prometheus, Grafana, or any Prometheus-compatible monitoring stack.
+
+### Endpoints
+
+The admin HTTP server runs on `metrics_port` (default: **9003**) and exposes the following endpoints:
+
+| Endpoint | Content-Type | Description |
+|----------|-------------|-------------|
+| `GET /metrics` | `text/plain; version=0.0.4` | All metrics in Prometheus exposition format |
+| `GET /metrics/summary` | `text/plain; version=0.0.4` | Human-readable summary (same content as the periodic log) |
+| `GET /health` | `application/json` | Health check with role, HA state, and service readiness |
+| `GET /role` | `text/plain` | Current HA role (`leader` / `standby`) |
+| `GET /ha_status` | `text/plain` | Current HA runtime state (`serving` / `starting` / etc.) |
+
+### Usage
+
+**Scrape the /metrics endpoint with Prometheus:**
+
+Add a scrape config to your `prometheus.yml`:
+
+```yaml
+scrape_configs:
+  - job_name: 'mooncake-master'
+    static_configs:
+      - targets: ['<master-host>:9003']
+    metrics_path: '/metrics'
+```
+
+**Quick check with curl:**
+
+```bash
+# Get Prometheus metrics
+curl http://<master-host>:9003/metrics
+
+# Get human-readable summary
+curl http://<master-host>:9003/metrics/summary
+
+# Check health
+curl http://<master-host>:9003/health
+```
+
+### Configuration
+
+The admin HTTP server is configured in the master config file (`master.json` or `master.yaml`):
+
+```text
+{
+  "enable_metric_reporting": true,
+  "metrics_port": 9003,
+  ...
+}
+```
+
+Set `enable_metric_reporting` to `false` to disable the periodic metrics log. HTTP endpoints (`/metrics`, `/health`, etc.) remain available regardless of this setting.
diff --git a/docs/source/getting_started/plugin-usage/3FS-USRBIO-Plugin.md b/docs/source/getting_started/plugin-usage/3FS-USRBIO-Plugin.md
index a542b05a..5566bc27 100644
--- a/docs/source/getting_started/plugin-usage/3FS-USRBIO-Plugin.md
+++ b/docs/source/getting_started/plugin-usage/3FS-USRBIO-Plugin.md
@@ -1,4 +1,8 @@
-# Mooncake HF3FS Plugin
+# Mooncake HF3FS Plugin (Experimental)
+
+```{warning}
+**Experimental / incomplete.** The HF3FS (3FS USRBIO) integration is under development and is not yet considered production-ready. Behavior, build flags, and configuration may change without notice. Use only for evaluation and testing.
+```
 
 This plugin implements 3FS native API (USRBIO) as a high-performance storage backend for Mooncake.
 
@@ -6,9 +10,9 @@ This plugin implements 3FS native API (USRBIO) as a high-performance storage bac
 
 ### 1. 3FS Installation
 - Build and install [3FS](https://github.com/deepseek-ai/3FS/)
-- Required library: `libhf3fs_api_shared.so` (Default location: `3FS_PATH/build/src/lib/api`)  
+- Required library: `libhf3fs_api_shared.so` (Default location: `3FS_PATH/build/src/lib/api`)
   → Install to: `/usr/lib/`
-- Required header: `hf3fs_usrbio.h` (Default location: `3FS_PATH/src/lib/api`)  
+- Required header: `hf3fs_usrbio.h` (Default location: `3FS_PATH/src/lib/api`)
   → Install to: `/usr/include/`
 
 ### 2. Mooncake Configuration
@@ -27,10 +31,10 @@ Start master server and specify the 3FS mount point:
 ```bash
 
 ./build/mooncake-store/src/mooncake_master \
-    --root_fs_dir=/path/to/3fs_mount_point 
+    --root_fs_dir=/path/to/3fs_mount_point
 ```
 ### Important Notes
-1. The specified directory **must** be a 3FS mount point  
+1. The specified directory **must** be a 3FS mount point
    - If not, the system will automatically fall back to POSIX API
 2. For optimal performance:
    - Ensure proper permissions on the 3FS mount point
@@ -40,4 +44,4 @@ Start master server and specify the 3FS mount point:
 ```bash
 
 ROLE=prefill MOONCAKE_STORAGE_ROOT_DIR=/mnt/3fs python3 ./stress_cluster_benchmark.py
-```
\ No newline at end of file
+```
diff --git a/docs/source/getting_started/quick-start.md b/docs/source/getting_started/quick-start.md
index 49ca87ec..0e42df1a 100644
--- a/docs/source/getting_started/quick-start.md
+++ b/docs/source/getting_started/quick-start.md
@@ -2,23 +2,50 @@
 
 This document describes how to quickly start using Mooncake Transfer Engine and Mooncake Store.
 
+## Before using Mooncake
+
+Install the following prerequisites before running any Mooncake component:
+- Python 3.10 or later; a virtual environment is recommended.
+- RDMA driver and SDK (for example, Mellanox OFED), if you plan to use RDMA for data transfer.
+- CUDA 12.1 or later, if the package is built with `-DUSE_CUDA` (disabled by default). For most CUDA-enabled use cases, such as RDMA-based KV cache transfer between GPUs or between GPU and DRAM, NVIDIA GPUDirect support is also required. *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`.
+- Hygon DTK SDK, if the package is built with `-DUSE_HYGON`. By default Mooncake looks for DTK under `DTK_HOME` or `/opt/dtk`.
+- Iluvatar CoreX SDK, if the package is built with `-DUSE_COREX`. By default Mooncake looks for CoreX under `COREX_HOME` or `/usr/local/corex`.
+
 ## Installation
 
 Install the Mooncake Transfer Engine package from PyPI, which includes both Mooncake Transfer Engine and Mooncake Store Python bindings:
 
 **For CUDA-enabled systems:**
+
+- CUDA < 13.0
 ```bash
-pip install mooncake-transfer-engine
+pip install mooncake-transfer-engine numpy pyzmq
+```
+
+- CUDA >= 13.0
+```bash
+pip install mooncake-transfer-engine-cuda13 numpy pyzmq
 ```
-📦 **Package Details**: [https://pypi.org/project/mooncake-transfer-engine/](https://pypi.org/project/mooncake-transfer-engine/)
 
 **For non-CUDA systems:**
 ```bash
-pip install mooncake-transfer-engine-non-cuda
+pip install mooncake-transfer-engine-non-cuda numpy pyzmq
 ```
-📦 **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.
+**For NPU systems:**
+```bash
+pip install mooncake-transfer-engine-npu
+```
+
+> **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, but it still needs system runtime libraries such as `libcurl4`, `libibverbs1`, `rdma-core`, `librdmacm1`, `libnuma1`, and `liburing2` on Ubuntu. In a fresh environment, run `sudo apt-get update` before installing them:
+>   ```bash
+>   sudo apt-get update && sudo apt-get install -y libcurl4 libibverbs1 rdma-core librdmacm1 libnuma1 liburing2
+>   ```
+> - 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`, first install the corresponding system runtime libraries. If the issue persists, uninstall the package and build the binaries manually.
 
 ## Transfer Engine Quick Start
 
@@ -40,9 +67,9 @@ def main():
 
     HOSTNAME = "localhost" # localhost for simple demo
     METADATA_SERVER = "P2PHANDSHAKE" # [ETCD_SERVER_URL, P2PHANDSHAKE, ...]
-    PROTOCOL = "rdma" # [rdma, tcp, ...]
+    PROTOCOL = "tcp" # use "rdma" on machines with RDMA devices configured
     DEVICE_NAME = "" # auto discovery if empty
-    
+
     # Initialize server engine
     server_engine = TransferEngine()
     server_engine.initialize(
@@ -52,22 +79,21 @@ def main():
         DEVICE_NAME
     )
     session_id = f"{HOSTNAME}:{server_engine.get_rpc_port()}"
-    
+
     # Allocate memory on server side (1MB buffer)
     server_buffer = np.zeros(1024 * 1024, dtype=np.uint8)
     server_ptr = server_buffer.ctypes.data
     server_len = server_buffer.nbytes
-    
-    # Register memory with Mooncake
-    if PROTOCOL == "rdma":
-        ret_value = server_engine.register_memory(server_ptr, server_len)
-        if ret_value != 0:
-            print("Mooncake memory registration failed.")
-            raise RuntimeError("Mooncake memory registration failed.")
+
+    # Register memory with Mooncake so the target address is advertised
+    ret_value = server_engine.register_memory(server_ptr, server_len)
+    if ret_value != 0:
+        print("Mooncake memory registration failed.")
+        raise RuntimeError("Mooncake memory registration failed.")
 
     print(f"Server initialized with session ID: {session_id}")
     print(f"Server buffer address: {server_ptr}, length: {server_len}")
-    
+
     # Send buffer info to client
     buffer_info = {
         "session_id": session_id,
@@ -76,7 +102,7 @@ def main():
     }
     socket.send_json(buffer_info)
     print("Buffer information sent to client")
-    
+
     # Keep server running
     try:
         while True:
@@ -85,18 +111,17 @@ def main():
         print("\nShutting down server...")
     finally:
         # Cleanup
-        if PROTOCOL == "rdma":
-            ret_value = server_engine.unregister_memory(server_ptr)
-            if ret_value != 0:
-                print("Mooncake memory deregistration failed.")
-                raise RuntimeError("Mooncake memory deregistration failed.")
+        ret_value = server_engine.unregister_memory(server_ptr)
+        if ret_value != 0:
+            print("Mooncake memory deregistration failed.")
+            raise RuntimeError("Mooncake memory deregistration failed.")
 
         socket.close()
         context.term()
 
 if __name__ == "__main__":
-    main() 
- 
+    main()
+
 ```
 
 ### Start Transfer Engine Sender (Client)
@@ -113,7 +138,7 @@ def main():
     context = zmq.Context()
     socket = context.socket(zmq.PULL)
     socket.connect(f"tcp://localhost:5555")
-    
+
     # Wait for buffer info from server
     print("Waiting for server buffer information...")
     buffer_info = socket.recv_json()
@@ -122,11 +147,11 @@ def main():
     server_len = buffer_info["len"]
     print(f"Received server info - Session ID: {server_session_id}")
     print(f"Server buffer address: {server_ptr}, length: {server_len}")
-    
+
     # Initialize client engine
     HOSTNAME = "localhost" # localhost for simple demo
     METADATA_SERVER = "P2PHANDSHAKE" # [ETCD_SERVER_URL, P2PHANDSHAKE, ...]
-    PROTOCOL = "rdma" # [rdma, tcp, ...]
+    PROTOCOL = "tcp" # use "rdma" on machines with RDMA devices configured
     DEVICE_NAME = "" # auto discovery if empty
 
     client_engine = TransferEngine()
@@ -137,15 +162,14 @@ def main():
         DEVICE_NAME
     )
     session_id = f"{HOSTNAME}:{client_engine.get_rpc_port()}"
-    
+
     # Allocate and initialize client buffer (1MB)
     client_buffer = np.ones(1024 * 1024, dtype=np.uint8)  # Fill with ones
     client_ptr = client_buffer.ctypes.data
     client_len = client_buffer.nbytes
-    
-    # Register memory with Mooncake
-    if PROTOCOL == "rdma":
-        ret_value = client_engine.register_memory(client_ptr, client_len)
+
+    # Register memory with Mooncake so the source address is advertised
+    ret_value = client_engine.register_memory(client_ptr, client_len)
     if ret_value != 0:
         print("Mooncake memory registration failed.")
         raise RuntimeError("Mooncake memory registration failed.")
@@ -161,24 +185,23 @@ def main():
             server_ptr,
             min(client_len, server_len)  # Transfer minimum of both lengths
         )
-    
+
         if ret >= 0:
             print("Transfer successful!")
         else:
             print("Transfer failed!")
-    
+
     # Cleanup
-    if PROTOCOL == "rdma":
-        ret_value = client_engine.unregister_memory(client_ptr)
-        if ret_value != 0:
-            print("Mooncake memory deregistration failed.")
-            raise RuntimeError("Mooncake memory deregistration failed.")
+    ret_value = client_engine.unregister_memory(client_ptr)
+    if ret_value != 0:
+        print("Mooncake memory deregistration failed.")
+        raise RuntimeError("Mooncake memory deregistration failed.")
 
     socket.close()
     context.term()
 
 if __name__ == "__main__":
-    main() 
+    main()
 
 ```
 
@@ -245,4 +268,25 @@ 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 & Tuning Guide](../deployment/mooncake-store-deployment-guide.md) for more examples and documentation.
+
+## Skills for AI Coding Assistants
+
+Mooncake ships a set of **built-in skills** under [`.claude/skills`](https://github.com/kvcache-ai/Mooncake/tree/main/.claude/skills) — reusable, task-focused playbooks that an AI coding assistant (such as Claude Code) invokes automatically when your request matches, or that you can run as a slash command:
+
+| Skill | Description |
+|-------|-------------|
+| `/mooncake-troubleshoot` | Diagnose Mooncake deployment and runtime issues (services, RDMA, env vars, logs). |
+| `/mooncake-ci-local` | Run pre-PR local validation via `scripts/run_ci_test.sh`. |
+| `/mooncake-api` | Work with the Mooncake Store, Transfer Engine, and EP/Backend Python APIs. |
+
+Install them without cloning the repository via the [Claude Code plugin marketplace](https://code.claude.com/docs/en/plugin-marketplaces):
+
+```text
+/plugin marketplace add kvcache-ai/Mooncake --sparse .claude-plugin
+/plugin install mooncake-troubleshoot@mooncake
+/plugin install mooncake-ci-local@mooncake
+/plugin install mooncake-api@mooncake
+```
+
+The `--sparse .claude-plugin` flag fetches only the marketplace catalog, and each plugin is published as a `git-subdir` source, so installing one fetches only that single skill directory — never the whole repo. If you are already working inside a Mooncake checkout, the skills under `.claude/skills/` load automatically with no setup.
diff --git a/docs/source/http-api-reference/http-service.md b/docs/source/http-api-reference/http-service.md
index b77c78dc..cbd79e00 100644
--- a/docs/source/http-api-reference/http-service.md
+++ b/docs/source/http-api-reference/http-service.md
@@ -10,6 +10,10 @@ The HTTP service serves multiple purposes:
 - **Data Inspection**: Examine stored objects and their replicas
 - **Health Checks**: Service availability and status verification
 
+The Python `mooncake.mooncake_store_service` module also provides a lightweight
+Store REST API for data operations and standalone segment mount/unmount
+workflows. Unless configured otherwise, it listens on port `8080`.
+
 ## HTTP Endpoints
 
 ### Metrics Endpoints
@@ -45,19 +49,44 @@ Retrieve replica information for a specific key, including memory locations and
 
 **Method**: `GET`
 **Parameters**: `key` (query parameter) - The object key to query
-**Content-Type**: `text/plain; version=0.0.4`
-**Response**: JSON-formatted replica descriptors for memory replicas
+**Content-Type**: `application/json; charset=utf-8`
+**Response**: JSON object with success status and replica data array
 
 **Example**:
 ```bash
 curl "http://localhost:8080/query_key?key=my_object"
 ```
 
-**Response Format**:
+**Success Response** (HTTP 200):
 ```json
 {
-  "transport_endpoint_": "hostname:port",
-  "buffer_descriptors": [...]
+  "success": true,
+  "data": [
+    {
+      "size_": 1073741824,
+      "buffer_address_": 140732000000000,
+      "protocol_": "rdma",
+      "transport_endpoint_": "192.168.1.100:12345"
+    }
+  ]
+}
+```
+
+**Error Response** (key not found, HTTP 404):
+```json
+{
+  "success": false,
+  "error_code": -704,
+  "error_message": "OBJECT_NOT_FOUND"
+}
+```
+
+**Error Response** (service unavailable, HTTP 503):
+```json
+{
+  "success": false,
+  "error_code": -1011,
+  "error_message": "service plane is not active"
 }
 ```
 
@@ -75,7 +104,7 @@ curl "http://localhost:8080/batch_query_keys?keys=key1,key2,key3"
 ```
 
 **Response Format**:
-```json
+```text
 {
   "success": true,
   "data": {
@@ -142,6 +171,57 @@ Used(bytes): 1073741824
 Capacity(bytes): 4294967296
 ```
 
+#### `/get_segments_detail`
+Get detailed information of all segments in JSON format, including segment metadata, allocator usage, and status.
+
+**Method**: `GET`
+**Content-Type**: `application/json; charset=utf-8`
+**Response**: JSON object containing an array of segment details
+
+**Example**:
+```bash
+curl http://localhost:8080/get_segments_detail
+```
+
+**Response Format**:
+```json
+{
+  "total_segments": 2,
+  "segments": [
+    {
+      "segment_name": "segment_0",
+      "segment_id": "00000000-0000-0000-0000-000000000001",
+      "client_id": "00000000-0000-0000-0000-000000000002",
+      "base_address": "0x300000000",
+      "size_bytes": 17179869184,
+      "size_human": "16 GiB",
+      "te_endpoint": "192.168.1.1:12345",
+      "protocol": "rdma",
+      "status": "MOUNTED",
+      "allocator_used_bytes": 1073741824,
+      "allocator_capacity_bytes": 17179869184,
+      "allocator_usage_percent": 6.25
+    }
+  ]
+}
+```
+
+**Fields**:
+- `total_segments` (integer): Total number of segments in the cluster
+- `segments` (array): Array of segment detail objects
+  - `segment_name` (string): Name of the segment
+  - `segment_id` (string): UUID of the segment
+  - `client_id` (string): UUID of the client that owns the segment
+  - `base_address` (string): Base memory address in hex
+  - `size_bytes` (integer): Segment size in bytes
+  - `size_human` (string): Human-readable segment size
+  - `te_endpoint` (string): Transport endpoint address
+  - `protocol` (string): Transfer protocol (e.g., rdma, tcp)
+  - `status` (string): Current segment status
+  - `allocator_used_bytes` (integer): Bytes currently allocated
+  - `allocator_capacity_bytes` (integer): Total allocator capacity in bytes
+  - `allocator_usage_percent` (number): Percentage of allocator capacity used
+
 ### Health Check Endpoints
 
 #### `/health`
@@ -150,7 +230,7 @@ Basic health check endpoint for service availability verification.
 **Method**: `GET`
 **Content-Type**: `text/plain; version=0.0.4`
 **Response**: `OK` when service is healthy
-**Status Codes**: 
+**Status Codes**:
 - `200 OK`: Service is healthy
 - Other: Service may be experiencing issues
 
@@ -159,3 +239,183 @@ Basic health check endpoint for service availability verification.
 curl http://localhost:8080/health
 ```
 
+## Store REST API Endpoints
+
+The following endpoints are served by the Python store REST service, which wraps
+`MooncakeDistributedStore` with an aiohttp service. The HTTP handlers live in
+Python, while mount and unmount operations are delegated to the underlying store
+binding. Start the service with:
+
+```bash
+python -m mooncake.mooncake_store_service \
+  --config /path/to/mooncake_config.json \
+  --port 8080
+```
+
+If the wheel console scripts are installed, the equivalent command is:
+
+```bash
+mc_store_rest_server --config /path/to/mooncake_config.json --port 8080
+```
+
+### `/api/mount_shm`
+Mount a named shared memory object as one or more Mooncake store segments. If
+the requested size exceeds the maximum registration size, the service may split
+the region and return multiple segment ids.
+
+**Method**: `POST`
+**Content-Type**: `application/json`
+
+**Request Body**:
+```json
+{
+  "name": "mooncake_segment",
+  "size": 16777216,
+  "offset": 0,
+  "protocol": "tcp",
+  "location": ""
+}
+```
+
+**Fields**:
+- `name` (string, required): Named shared memory object name. A leading `/` is
+  accepted, but path separators are not.
+- `size` (integer, required): Number of bytes to mount.
+- `offset` (integer, optional): File offset in bytes. Defaults to `0`.
+- `protocol` (string, optional): Transfer protocol. Defaults to the service
+  configuration protocol.
+- `location` (string, optional): Device or locality hint. Defaults to an empty
+  string.
+
+**Success Response**:
+```json
+{
+  "status": "success",
+  "segment_ids": ["00000000-0000-0000-0000-000000000001"]
+}
+```
+
+**Example**:
+```bash
+curl -X POST http://localhost:8080/api/mount_shm \
+  -H "Content-Type: application/json" \
+  -d '{
+        "name": "mooncake_segment",
+        "size": 16777216,
+        "offset": 0,
+        "protocol": "tcp",
+        "location": ""
+      }'
+```
+
+### `/api/unmount_shm`
+Unmount one or more segment ids previously returned by `/api/mount_shm`.
+
+**Method**: `POST`
+**Content-Type**: `application/json`
+
+**Request Body**:
+```json
+{
+  "segment_ids": ["00000000-0000-0000-0000-000000000001"],
+  "grace_period_seconds": 0
+}
+```
+
+`segment_ids` may also be provided as a single string for one segment.
+`grace_period_seconds` is optional and defaults to `0`, which keeps the
+existing immediate unmount behavior. When set to a positive value, the master
+keeps the segment readable for that grace period while preventing new
+allocations, then completes the unmount.
+
+**Success Response**:
+```json
+{
+  "status": "success"
+}
+```
+
+**Example**:
+```bash
+curl -X POST http://localhost:8080/api/unmount_shm \
+  -H "Content-Type: application/json" \
+  -d '{"segment_ids": ["00000000-0000-0000-0000-000000000001"],
+       "grace_period_seconds": 30}'
+```
+
+### `/api/mount`
+Allocate memory inside the store process and mount it as one or more Mooncake
+store segments. If the requested size exceeds the maximum registration size,
+the service may split it and return multiple segment ids. The response includes
+the actual allocated size after alignment.
+
+**Method**: `POST`
+**Content-Type**: `application/json`
+
+**Request Body**:
+```json
+{
+  "size": 16777216,
+  "protocol": "tcp",
+  "location": ""
+}
+```
+
+**Fields**:
+- `size` (integer, required): Number of bytes requested. Must be positive.
+- `protocol` (string, optional): Transfer protocol. Defaults to the service
+  configuration protocol.
+- `location` (string, optional): Device or locality hint. Defaults to an empty
+  string.
+
+**Success Response**:
+```json
+{
+  "status": "success",
+  "segment_ids": ["00000000-0000-0000-0000-000000000002"],
+  "allocated_size": 16777216
+}
+```
+
+**Example**:
+```bash
+curl -X POST http://localhost:8080/api/mount \
+  -H "Content-Type: application/json" \
+  -d '{"size": 16777216, "protocol": "tcp", "location": ""}'
+```
+
+### `/api/unmount`
+Unmount one or more segment ids previously returned by `/api/mount` and free
+the memory allocated by the store process.
+
+**Method**: `POST`
+**Content-Type**: `application/json`
+
+**Request Body**:
+```json
+{
+  "segment_ids": ["00000000-0000-0000-0000-000000000002"],
+  "grace_period_seconds": 0
+}
+```
+
+`segment_ids` may also be provided as a single string for one segment.
+`grace_period_seconds` is optional and defaults to `0`, which keeps the
+existing immediate unmount-and-free behavior. When set to a positive value, the
+master keeps the segment readable for that grace period while preventing new
+allocations, then the store releases the local allocated memory after cleanup.
+
+**Success Response**:
+```json
+{
+  "status": "success"
+}
+```
+
+**Example**:
+```bash
+curl -X POST http://localhost:8080/api/unmount \
+  -H "Content-Type: application/json" \
+  -d '{"segment_ids": ["00000000-0000-0000-0000-000000000002"],
+       "grace_period_seconds": 30}'
+```
diff --git a/docs/source/index.md b/docs/source/index.md
index a22dab44..3e9428d7 100644
--- a/docs/source/index.md
+++ b/docs/source/index.md
@@ -25,35 +25,59 @@ Mooncake is the serving platform for <a href="https://kimi.ai/">Kimi</a>, a lead
 Now both the Transfer Engine and Mooncake Store are open-sourced!
 This repository also hosts its technical report and the open-sourced traces.
 
+<h2 id="overview">🎉 Overview</h2>
+
+Mooncake features a KVCache-centric disaggregated architecture that separates prefill and decode clusters. It also leverages underutilized CPU, DRAM, and SSD resources in GPU clusters to build a disaggregated KVCache pool.
+
+:::{figure} ./image/architecture.png
+:align: center
+:alt: Mooncake architecture
+:width: 88%
+:::
+
+At the center of Mooncake is a KVCache-centric scheduler that balances effective throughput with latency-related Service Level Objectives (SLOs). In highly overloaded scenarios, Mooncake uses prediction-based early rejection to avoid wasting computation on requests that cannot meet their SLOs. Experiments show that Mooncake is especially effective for long-context workloads, achieving up to a 525% throughput increase in simulated scenarios while meeting SLOs. Under real workloads, Mooncake's architecture enables <a href="https://kimi.ai/">Kimi</a> to handle 75% more requests.
+
 <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.
- - **Dec 19, 2025**: Mooncake Transfer Engine has been [integrated into TensorRT LLM](https://github.com/NVIDIA/TensorRT-LLM/tree/main/cpp/tensorrt_llm/executor/cache_transmission/mooncake_utils) for KVCache transfer in PD-disaggregated inference.
- - **Dec 19, 2025**: Mooncake Transfer Engine has been directly integrated into vLLM v1 as a [KV Connector](https://docs.vllm.ai/en/latest/features/mooncake_connector_usage/) in PD-disaggregated setups.
- - **Nov 07, 2025**: [RBG + SGLang HiCache + Mooncake](https://github.com/sgl-project/rbg/blob/main/keps/74-mooncake-integration/README.md), a role-based out-of-the-box solution for cloud native deployment, which is elastic, scalable, and high-performance.
- - **Sept 18, 2025**: Mooncake Store empowers vLLM Ascend by serving as [the distributed KV cache pool backend](https://docs.vllm.ai/projects/ascend/zh-cn/main/user_guide/feature_guide/kv_pool.html).
- - **Sept 10, 2025**: SGLang officially supports Mooncake Store as a [hierarchical KV caching storage backend](https://lmsys.org/blog/2025-09-10-sglang-hicache/). The integration extends RadixAttention with multi-tier KV cache storage across device, host, and remote storage layers.
- - **Sept 10, 2025**: The official & high-performance version of Mooncake P2P Store is open-sourced as [checkpoint-engine](https://github.com/MoonshotAI/checkpoint-engine/). It has been successfully applied in K1.5 and K2 production training, updating Kimi-K2 model (1T parameters) across thousands of GPUs in ~20s.
- - **Aug 23, 2025**: [xLLM](https://github.com/jd-opensource/xllm) high-performance inference engine builds hybrid KV cache management based on Mooncake, supporting global KV cache management with intelligent offloading and prefetching.
- - **Aug 18, 2025**: vLLM-Ascend [integrates Mooncake Transfer Engine](https://docs.vllm.ai/projects/ascend/en/latest/developer_guide/feature_guide/disaggregated_prefill.html) for KV cache register and disaggregate prefill, enabling efficient distributed inference on Ascend NPUs.
- - **Jul 20, 2025**: Mooncake powers [the deployment of Kimi K2](https://lmsys.org/blog/2025-07-20-k2-large-scale-ep/) on 128 H200 GPUs with PD disaggregation and large-scale expert parallelism, achieving 224k tokens/sec prefill throughput and 288k tokens/sec decode throughput.
- - **Jun 20, 2025**: Mooncake becomes a PD disaggregation [backend](https://kvcache-ai.github.io/Mooncake/getting_started/examples/lmdeploy-integration-v0.9.html) for LMDeploy.
- - **May 9, 2025**: NIXL officially supports Mooncake Transfer Engine as [a backend plugin](https://github.com/ai-dynamo/nixl/blob/main/src/plugins/mooncake/README.md).
- - **May 8, 2025**: Mooncake x LMCache <a href="https://github.com/kvcache-ai/Mooncake/blob/main/doc/en/lmcache-integration.md" target="_blank">unite</a> to pioneer KVCache-centric LLM serving system.
- - **May 5, 2025**: Supported by Mooncake Team, SGLang release <a href="https://lmsys.org/blog/2025-05-05-large-scale-ep/" target="_blank">guidance</a> to deploy DeepSeek with PD Disaggregation on 96 H100 GPUs.
- - **Apr 22, 2025**: LMCache officially supports Mooncake Store as a <a href="https://blog.lmcache.ai/2025-04-22-tencent/" target="_blank">remote connector</a>.
- - **Apr 10, 2025**: SGLang officially supports Mooncake Transfer Engine for disaggregated prefilling and KV cache transfer.
- - **Mar 7, 2025**: We open-sourced the Mooncake Store, a distributed KVCache based on Transfer Engine. vLLM's xPyD disaggregated prefilling & decoding based on Mooncake Store will be released soon.
- - **Feb 25, 2025**: Mooncake receives the **Best Paper Award** at **FAST 2025**!
- - **Feb 21, 2025**: The updated <a href="https://github.com/kvcache-ai/Mooncake/tree/main/FAST25-release/traces" target="_blank">traces</a> used in our FAST'25 paper have been released.
- - **Dec 16, 2024**: vLLM officially supports Mooncake Transfer Engine for disaggregated prefilling and KV cache transfer.
- - **Nov 28, 2024**: We open-sourced the Transfer Engine, the central component of Mooncake. We also provide two demonstrations of Transfer Engine: a P2P Store and vLLM integration.
- - **July 9, 2024**: We open-sourced the trace as a <a href="https://github.com/kvcache-ai/Mooncake/blob/main/FAST25-release/arxiv-trace/mooncake_trace.jsonl" target="_blank">JSONL file</a>.
- - **June 27, 2024**: We present a series of Chinese blogs with more discussions on <a href="https://zhuanlan.zhihu.com/p/705754254">zhihu 1</a>, <a href="https://zhuanlan.zhihu.com/p/705910725">2</a>, <a href="https://zhuanlan.zhihu.com/p/706204757">3</a>, <a href="https://zhuanlan.zhihu.com/p/707997501">4</a>, <a href="https://zhuanlan.zhihu.com/p/9461861451">5</a>, <a href="https://zhuanlan.zhihu.com/p/1939988652114580803">6</a>, <a href="https://zhuanlan.zhihu.com/p/1959366095443064318">7</a>.
- - **June 26, 2024**: Initial technical report release.
+- **May 7, 2026**: 🚀 [vLLM officially features Mooncake Store](https://vllm.ai/blog/mooncake-store) — a deep dive into how Mooncake's distributed KVCache engine supercharges vLLM inference with high-throughput, memory-efficient, cross-instance KV cache sharing!
+- **Apr 29, 2026**: SGLang introduces [RDMA-based P2P weight transfer for large-scale distributed RL](https://lmsys.org/blog/2026-04-29-p2p-update/) using Mooncake TransferEngine, achieving 7x faster weight updates for the 1T-parameter Kimi-K2 model (53s → 7.2s) with zero-copy RDMA transfer across thousands of GPUs.
+- **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.
+
+:::{dropdown} More
+:animate: fade-in
+
+- **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).
+- **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.
+- **Dec 19, 2025**: Mooncake Transfer Engine has been [integrated into TensorRT LLM](https://github.com/NVIDIA/TensorRT-LLM/tree/main/cpp/tensorrt_llm/executor/cache_transmission/mooncake_utils) for KVCache transfer in PD-disaggregated inference.
+- **Dec 19, 2025**: Mooncake Transfer Engine has been directly integrated into vLLM v1 as a [KV Connector](https://docs.vllm.ai/en/latest/features/mooncake_connector_usage/) in PD-disaggregated setups.
+- **Nov 07, 2025**: [RBG + SGLang HiCache + Mooncake](https://github.com/sgl-project/rbg/blob/main/keps/74-mooncake-integration/README.md), a role-based out-of-the-box solution for cloud native deployment, which is elastic, scalable, and high-performance.
+- **Sept 18, 2025**: Mooncake Store empowers vLLM Ascend by serving as [the distributed KV cache pool backend](https://docs.vllm.ai/projects/ascend/zh-cn/main/user_guide/feature_guide/kv_pool.html).
+- **Sept 10, 2025**: SGLang officially supports Mooncake Store as a [hierarchical KV caching storage backend](https://lmsys.org/blog/2025-09-10-sglang-hicache/). The integration extends RadixAttention with multi-tier KV cache storage across device, host, and remote storage layers.
+- **Sept 10, 2025**: The official & high-performance version of Mooncake P2P Store is open-sourced as [checkpoint-engine](https://github.com/MoonshotAI/checkpoint-engine/). It has been successfully applied in K1.5 and K2 production training, updating Kimi-K2 model (1T parameters) across thousands of GPUs in ~20s.
+- **Aug 23, 2025**: [xLLM](https://github.com/jd-opensource/xllm) high-performance inference engine builds hybrid KV cache management based on Mooncake, supporting global KV cache management with intelligent offloading and prefetching.
+- **Aug 18, 2025**: vLLM-Ascend [integrates Mooncake Transfer Engine](https://docs.vllm.ai/projects/ascend/en/latest/developer_guide/feature_guide/disaggregated_prefill.html) for KV cache register and disaggregate prefill, enabling efficient distributed inference on Ascend NPUs.
+- **Jul 20, 2025**: Mooncake powers [the deployment of Kimi K2](https://lmsys.org/blog/2025-07-20-k2-large-scale-ep/) on 128 H200 GPUs with PD disaggregation and large-scale expert parallelism, achieving 224k tokens/sec prefill throughput and 288k tokens/sec decode throughput.
+- **Jun 20, 2025**: Mooncake becomes a PD disaggregation [backend](getting_started/examples/lmdeploy-integration-v0.9) for LMDeploy.
+- **May 9, 2025**: NIXL officially supports Mooncake Transfer Engine as [a backend plugin](https://github.com/ai-dynamo/nixl/blob/main/src/plugins/mooncake/README.md).
+- **May 8, 2025**: [Mooncake x LMCache](getting_started/examples/lmcache-integration) unite to pioneer KVCache-centric LLM serving system.
+- **May 5, 2025**: Supported by Mooncake Team, SGLang release <a href="https://lmsys.org/blog/2025-05-05-large-scale-ep/" target="_blank">guidance</a> to deploy DeepSeek with PD Disaggregation on 96 H100 GPUs.
+- **Apr 22, 2025**: LMCache officially supports Mooncake Store as a <a href="https://blog.lmcache.ai/2025-04-22-tencent/" target="_blank">remote connector</a>.
+- **Apr 10, 2025**: SGLang officially supports Mooncake Transfer Engine for disaggregated prefilling and KV cache transfer.
+- **Mar 7, 2025**: We open-sourced the Mooncake Store, a distributed KVCache based on Transfer Engine. vLLM's xPyD disaggregated prefilling & decoding based on Mooncake Store will be released soon.
+- **Feb 25, 2025**: Mooncake receives the **Best Paper Award** at **FAST 2025**!
+- **Feb 21, 2025**: The updated <a href="https://github.com/kvcache-ai/Mooncake/tree/main/FAST25-release/traces" target="_blank">traces</a> used in our FAST'25 paper have been released.
+- **Dec 16, 2024**: vLLM officially supports Mooncake Transfer Engine for disaggregated prefilling and KV cache transfer.
+- **Nov 28, 2024**: We open-sourced the Transfer Engine, the central component of Mooncake. We also provide two demonstrations of Transfer Engine: a P2P Store and vLLM integration.
+- **July 9, 2024**: We open-sourced the trace as a <a href="https://github.com/kvcache-ai/Mooncake/blob/main/FAST25-release/arxiv-trace/mooncake_trace.jsonl" target="_blank">JSONL file</a>.
+- **June 27, 2024**: We present a series of Chinese blogs with more discussions on <a href="https://zhuanlan.zhihu.com/p/705754254">zhihu 1</a>, <a href="https://zhuanlan.zhihu.com/p/705910725">2</a>, <a href="https://zhuanlan.zhihu.com/p/706204757">3</a>, <a href="https://zhuanlan.zhihu.com/p/707997501">4</a>, <a href="https://zhuanlan.zhihu.com/p/9461861451">5</a>, <a href="https://zhuanlan.zhihu.com/p/1939988652114580803">6</a>, <a href="https://zhuanlan.zhihu.com/p/1959366095443064318">7</a>.
+- **June 26, 2024**: Initial technical report release.
+
+:::
 
 ## Documentation
 
@@ -65,40 +89,32 @@ This repository also hosts its technical report and the open-sourced traces.
 
 getting_started/build
 getting_started/quick-start
-getting_started/supported-protocols
-getting_started/plugin-usage/3FS-USRBIO-Plugin
-getting_started/examples/lmcache-integration
-getting_started/examples/lmdeploy-integration-v0.9
-getting_started/examples/sglang-integration-v1
+
+:::
+
+% Deployment docs
+
+:::{toctree}
+:caption: Deployment
+:maxdepth: 2
+
+deployment/mooncake-store-deployment-guide
 getting_started/examples/sglang-integration/index
 getting_started/examples/vllm-integration/index
+Mooncake x LMCache Integration<getting_started/examples/lmcache-integration>
+Mooncake x LMDeploy Integration<getting_started/examples/lmdeploy-integration-v0.9>
 :::
 
+
 % Making the most out of Mooncake
 
 :::{toctree}
 :caption: Performance
 :maxdepth: 1
 
-performance/sglang-benchmark-results-v1
-performance/vllm-benchmark-results-v0.2
-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
-
-:::{toctree}
-:caption: Python API Reference
-:maxdepth: 1
-
-python-api-reference/mooncake-store
-python-api-reference/transfer-engine
-http-api-reference/http-service
-python-api-reference/ep-backend
+performance/vllm/index
+performance/sglang/index
+performance/mooncake/index
 :::
 
 % Explanation of Mooncake internals
@@ -110,10 +126,28 @@ python-api-reference/ep-backend
 design/architecture
 design/mooncake-store
 design/p2p-store
+design/mooncake-backend-pg
+design/mooncake-ep
 design/transfer-engine/index
+design/hicache-design
+design/engram
+design/unified-parallel-tensor-io
 design/tent/overview
 design/tent/tebench
-design/hicache-design
+design/conductor/conductor-architecture-design
+design/ssd-free-ratio-first-allocation
+:::
+
+% API Documentation
+
+:::{toctree}
+:caption: API Reference
+:maxdepth: 2
+
+api-reference/python/index
+api-reference/cpp/index
+api-reference/http/index
+api-reference/rust/index
 :::
 
 % Q&A for Mooncake
@@ -124,16 +158,9 @@ design/hicache-design
 
 troubleshooting/error-code
 troubleshooting/troubleshooting
+troubleshooting/pg-ep-troubleshooting
 :::
 
-% Deployment docs
-
-:::{toctree}
-:caption: Deployment
-:maxdepth: 2
-
-deployment/mooncake-store-deployment-guide
-:::
 
 % Community
 
@@ -143,3 +170,15 @@ deployment/mooncake-store-deployment-guide
 
 community/governance
 :::
+
+% Archived content
+
+:::{toctree}
+:caption: Archived
+:maxdepth: 1
+
+getting_started/examples/vllm-integration/vllm-mooncakestoreconnector
+getting_started/examples/vllm-integration/vllm-integration-v0.2
+getting_started/examples/vllm-integration/vllm-integration-v0.3
+getting_started/examples/vllm-integration/vllm-integration-v1.0
+:::
diff --git a/docs/source/performance/allocation-strategy-benchmark-result.md b/docs/source/performance/allocation-strategy-benchmark-result.md
new file mode 100644
index 00000000..55d0d7ee
--- /dev/null
+++ b/docs/source/performance/allocation-strategy-benchmark-result.md
@@ -0,0 +1,191 @@
+# AllocationStrategy Performance
+
+`allocation_strategy_bench` evaluates the allocation-routing overhead and behavior of `RandomAllocationStrategy` and `FreeRatioFirstAllocationStrategy` across multiple Mooncake Store segments.
+
+## Execution
+
+```bash
+./mooncake-store/benchmarks/allocation_strategy_bench --workload=dsa
+```
+
+## Workload
+
+The benchmark compares two workload types:
+
+- **KV-only**: a general single-object-size allocation workload without indexer objects. This is used as the non-DSA baseline.
+- **Paired KV+Indexer**: a DSA-style workload where KV cache objects and indexer objects are allocated together.
+- KV cache object size: 3,274,752 B (3.12 MB)
+- Indexer object size: 658,432 B (643 KB)
+- Each round samples `N` uniformly from `[1, 128]`.
+- In KV-only mode, each round issues `N` KV allocations.
+- In paired KV+Indexer mode, each round issues `N` KV allocations followed by `N` indexer allocations.
+- Segment capacity: 1 TB per segment by default.
+- Segment counts: 1, 2, 4, 8, 16.
+- Replica counts: 1, 2, 3 when enough segments are available.
+- On allocation failure, the benchmark randomly evicts 5% of all live objects and retries.
+- `UtilRatio_*` is sampled only when eviction is triggered, immediately before evicting objects.
+
+`skewed=no` uses uniform segment capacity. `skewed=yes` alternates segment capacity between 1.5x and 0.5x the base capacity.
+
+## Summary
+
+For KV-only, utilization reaches about 96.1%. For the DSA paired KV+indexer workload, both strategies reach about 95.2% utilization. The difference mainly comes from mixed object sizes in the paired workload. `RandomAllocationStrategy` generally has higher throughput and lower latency, while `FreeRatioFirstAllocationStrategy` tends to keep segment utilization more balanced.
+
+## Result
+
+KV-only is the non-DSA baseline. It represents the common case where Mooncake Store allocates only one object size.
+
+### KV-only / Skewed=no / Random
+
+
+| Replica | Segments | AllocSize | Cluster(GB) | UtilRatio_min | UtilRatio_p99 | UtilRatio_p90 | UtilRatio_p50 | UtilRatio_max | UtilRatio_avg | Throughput | Avg(ns) | P50(ns) | P90(ns) | P99(ns) | UtilStdDev | Evictions(round) | Succ/Total        |
+| ------- | -------- | --------- | ----------- | ------------- | ------------- | ------------- | ------------- | ------------- | ------------- | ---------- | ------- | ------- | ------- | ------- | ---------- | ---------------- | ----------------- |
+| 1       | 1        | 3198KB    | 1024.0      | 0.960935      | 0.960935      | 0.960935      | 0.960935      | 0.960935      | 0.960935      | 1015964    | 953     | 231     | 1132    | 1694    | 0.0000     | 63               | 1007262/1007262   |
+| 1       | 2        | 3198KB    | 2048.0      | 0.960935      | 0.960935      | 0.960935      | 0.960935      | 0.960935      | 0.960935      | 857559     | 1135    | 331     | 1633    | 2324    | 0.0001     | 63               | 2014525/2014525   |
+| 2       | 2        | 3198KB    | 2048.0      | 0.960935      | 0.960935      | 0.960935      | 0.960935      | 0.960935      | 0.960935      | 555934     | 1767    | 681     | 912     | 4388    | 0.0000     | 63               | 1007262/1007262   |
+| 1       | 4        | 3198KB    | 4096.0      | 0.960935      | 0.960935      | 0.960935      | 0.960935      | 0.960935      | 0.960935      | 809246     | 1204    | 340     | 1994    | 2435    | 0.0006     | 63               | 4029050/4029050   |
+| 2       | 4        | 3198KB    | 4096.0      | 0.960935      | 0.960935      | 0.960935      | 0.960935      | 0.960935      | 0.960935      | 503007     | 1957    | 772     | 1002    | 7545    | 0.0004     | 63               | 2014525/2014525   |
+| 3       | 4        | 3198KB    | 4096.0      | 0.960935      | 0.960935      | 0.960935      | 0.960935      | 0.960935      | 0.960935      | 377095     | 2620    | 1112    | 1412    | 6613    | 0.0002     | 63               | 1343016/1343016   |
+| 1       | 8        | 3198KB    | 8192.0      | 0.960935      | 0.960935      | 0.960935      | 0.960935      | 0.960935      | 0.960935      | 751926     | 1299    | 351     | 2124    | 2534    | 0.0004     | 63               | 8058100/8058100   |
+| 2       | 8        | 3198KB    | 8192.0      | 0.960935      | 0.960935      | 0.960935      | 0.960935      | 0.960935      | 0.960935      | 471589     | 2089    | 831     | 1053    | 7063    | 0.0002     | 63               | 4029050/4029050   |
+| 3       | 8        | 3198KB    | 8192.0      | 0.960935      | 0.960935      | 0.960935      | 0.960935      | 0.960935      | 0.960935      | 346496     | 2855    | 1203    | 1513    | 7133    | 0.0008     | 63               | 2686033/2686033   |
+| 1       | 16       | 3198KB    | 16384.0     | 0.960935      | 0.960935      | 0.960935      | 0.960935      | 0.960935      | 0.960935      | 691790     | 1409    | 371     | 2274    | 2835    | 0.0003     | 63               | 16116200/16116200 |
+| 2       | 16       | 3198KB    | 16384.0     | 0.960935      | 0.960935      | 0.960935      | 0.960935      | 0.960935      | 0.960935      | 434370     | 2271    | 892     | 1163    | 8135    | 0.0003     | 63               | 8058100/8058100   |
+| 3       | 16       | 3198KB    | 16384.0     | 0.960935      | 0.960935      | 0.960935      | 0.960935      | 0.960935      | 0.960935      | 322012     | 3074    | 1292    | 1613    | 8326    | 0.0006     | 63               | 5372067/5372067   |
+
+
+### KV-only / Skewed=no / FreeRatioFirst
+
+
+| Replica | Segments | AllocSize | Cluster(GB) | UtilRatio_min | UtilRatio_p99 | UtilRatio_p90 | UtilRatio_p50 | UtilRatio_max | UtilRatio_avg | Throughput | Avg(ns) | P50(ns) | P90(ns) | P99(ns) | UtilStdDev | Evictions(round) | Succ/Total        |
+| ------- | -------- | --------- | ----------- | ------------- | ------------- | ------------- | ------------- | ------------- | ------------- | ---------- | ------- | ------- | ------- | ------- | ---------- | ---------------- | ----------------- |
+| 1       | 1        | 3198KB    | 1024.0      | 0.960935      | 0.960935      | 0.960935      | 0.960935      | 0.960935      | 0.960935      | 887098     | 1096    | 361     | 1293    | 1894    | 0.0000     | 63               | 1007262/1007262   |
+| 1       | 2        | 3198KB    | 2048.0      | 0.960935      | 0.960935      | 0.960935      | 0.960935      | 0.960935      | 0.960935      | 768834     | 1269    | 461     | 1803    | 2475    | 0.0000     | 63               | 2014525/2014525   |
+| 2       | 2        | 3198KB    | 2048.0      | 0.960935      | 0.960935      | 0.960935      | 0.960935      | 0.960935      | 0.960935      | 499411     | 1971    | 882     | 1152    | 4219    | 0.0000     | 63               | 1007262/1007262   |
+| 1       | 4        | 3198KB    | 4096.0      | 0.960935      | 0.960935      | 0.960935      | 0.960935      | 0.960935      | 0.960935      | 680538     | 1438    | 561     | 2234    | 2685    | 0.0000     | 63               | 4029050/4029050   |
+| 2       | 4        | 3198KB    | 4096.0      | 0.960935      | 0.960935      | 0.960935      | 0.960935      | 0.960935      | 0.960935      | 441341     | 2235    | 1052    | 1272    | 8336    | 0.0000     | 63               | 2014525/2014525   |
+| 3       | 4        | 3198KB    | 4096.0      | 0.960935      | 0.960935      | 0.960935      | 0.960935      | 0.960935      | 0.960935      | 328024     | 3017    | 1533    | 1844    | 6682    | 0.0000     | 63               | 1343016/1343016   |
+| 1       | 8        | 3198KB    | 8192.0      | 0.960935      | 0.960935      | 0.960935      | 0.960935      | 0.960935      | 0.960935      | 608814     | 1611    | 661     | 2435    | 2985    | 0.0000     | 63               | 8058100/8058100   |
+| 2       | 8        | 3198KB    | 8192.0      | 0.960935      | 0.960935      | 0.960935      | 0.960935      | 0.960935      | 0.960935      | 387271     | 2551    | 1292    | 1554    | 5621    | 0.0000     | 63               | 4029050/4029050   |
+| 3       | 8        | 3198KB    | 8192.0      | 0.960935      | 0.960935      | 0.960935      | 0.960935      | 0.960935      | 0.960935      | 295263     | 3355    | 1743    | 2014    | 7685    | 0.0000     | 63               | 2686033/2686033   |
+| 1       | 16       | 3198KB    | 16384.0     | 0.960935      | 0.960935      | 0.960935      | 0.960935      | 0.960935      | 0.960935      | 571007     | 1720    | 671     | 2645    | 3186    | 0.0000     | 63               | 16116200/16116200 |
+| 2       | 16       | 3198KB    | 16384.0     | 0.960935      | 0.960935      | 0.960935      | 0.960935      | 0.960935      | 0.960935      | 343893     | 2876    | 1503    | 1773    | 7825    | 0.0000     | 63               | 8058100/8058100   |
+| 3       | 16       | 3198KB    | 16384.0     | 0.960935      | 0.960935      | 0.960935      | 0.960935      | 0.960935      | 0.960935      | 251734     | 3941    | 2124    | 2524    | 8346    | 0.0000     | 63               | 5372067/5372067   |
+
+
+### KV-only / Skewed=yes / Random
+
+
+| Replica | Segments | AllocSize | Cluster(GB) | UtilRatio_min | UtilRatio_p99 | UtilRatio_p90 | UtilRatio_p50 | UtilRatio_max | UtilRatio_avg | Throughput | Avg(ns) | P50(ns) | P90(ns) | P99(ns) | UtilStdDev | Evictions(round) | Succ/Total        |
+| ------- | -------- | --------- | ----------- | ------------- | ------------- | ------------- | ------------- | ------------- | ------------- | ---------- | ------- | ------- | ------- | ------- | ---------- | ---------------- | ----------------- |
+| 1       | 1        | 3198KB    | 1536.0      | 0.960937      | 0.960937      | 0.960937      | 0.960937      | 0.960937      | 0.960937      | 925104     | 1050    | 301     | 1393    | 2124    | 0.0000     | 63               | 1510893/1510893   |
+| 1       | 2        | 3198KB    | 2048.0      | 0.960936      | 0.960936      | 0.960936      | 0.960936      | 0.960936      | 0.960936      | 846763     | 1149    | 331     | 1693    | 2354    | 0.0081     | 63               | 2014525/2014525   |
+| 2       | 2        | 3198KB    | 2048.0      | 0.960936      | 0.960936      | 0.960936      | 0.960936      | 0.960936      | 0.960936      | 763883     | 1278    | 512     | 741     | 4468    | 0.0119     | 36               | 1007262/1007262   |
+| 1       | 4        | 3198KB    | 4096.0      | 0.960936      | 0.960936      | 0.960936      | 0.960936      | 0.960936      | 0.960936      | 793128     | 1230    | 350     | 2014    | 2465    | 0.0078     | 63               | 4029050/4029050   |
+| 2       | 4        | 3198KB    | 4096.0      | 0.960936      | 0.960936      | 0.960936      | 0.960936      | 0.960936      | 0.960936      | 495155     | 1988    | 841     | 1092    | 6463    | 0.0043     | 63               | 2014525/2014525   |
+| 3       | 4        | 3198KB    | 4096.0      | 0.960936      | 0.960936      | 0.960936      | 0.960936      | 0.960936      | 0.960936      | 442797     | 2227    | 1001    | 1292    | 6593    | 0.0116     | 47               | 1343016/1343016   |
+| 1       | 8        | 3198KB    | 8192.0      | 0.960936      | 0.960936      | 0.960936      | 0.960936      | 0.960936      | 0.960936      | 742956     | 1315    | 370     | 2135    | 2565    | 0.0083     | 63               | 8058100/8058100   |
+| 2       | 8        | 3198KB    | 8192.0      | 0.960936      | 0.960936      | 0.960936      | 0.960936      | 0.960936      | 0.960936      | 449694     | 2192    | 912     | 1162    | 7835    | 0.0062     | 63               | 4029050/4029050   |
+| 3       | 8        | 3198KB    | 8192.0      | 0.960936      | 0.960936      | 0.960936      | 0.960936      | 0.960936      | 0.960936      | 335572     | 2949    | 1333    | 1663    | 6973    | 0.0037     | 63               | 2686033/2686033   |
+| 1       | 16       | 3198KB    | 16384.0     | 0.960936      | 0.960936      | 0.960936      | 0.960936      | 0.960936      | 0.960936      | 680546     | 1438    | 381     | 2295    | 2876    | 0.0081     | 63               | 16116200/16116200 |
+| 2       | 16       | 3198KB    | 16384.0     | 0.960936      | 0.960936      | 0.960936      | 0.960936      | 0.960936      | 0.960936      | 419847     | 2350    | 952     | 1223    | 8747    | 0.0069     | 63               | 8058100/8058100   |
+| 3       | 16       | 3198KB    | 16384.0     | 0.960936      | 0.960936      | 0.960936      | 0.960936      | 0.960936      | 0.960936      | 316027     | 3133    | 1403    | 1714    | 7764    | 0.0057     | 63               | 5372067/5372067   |
+
+
+### KV-only / Skewed=yes / FreeRatioFirst
+
+
+| Replica | Segments | AllocSize | Cluster(GB) | UtilRatio_min | UtilRatio_p99 | UtilRatio_p90 | UtilRatio_p50 | UtilRatio_max | UtilRatio_avg | Throughput | Avg(ns) | P50(ns) | P90(ns) | P99(ns) | UtilStdDev | Evictions(round) | Succ/Total        |
+| ------- | -------- | --------- | ----------- | ------------- | ------------- | ------------- | ------------- | ------------- | ------------- | ---------- | ------- | ------- | ------- | ------- | ---------- | ---------------- | ----------------- |
+| 1       | 1        | 3198KB    | 1536.0      | 0.960937      | 0.960937      | 0.960937      | 0.960937      | 0.960937      | 0.960937      | 828717     | 1176    | 431     | 1494    | 2265    | 0.0000     | 63               | 1510893/1510893   |
+| 1       | 2        | 3198KB    | 2048.0      | 0.960936      | 0.960936      | 0.960936      | 0.960936      | 0.960936      | 0.960936      | 770115     | 1267    | 461     | 1774    | 2465    | 0.0000     | 63               | 2014525/2014525   |
+| 2       | 2        | 3198KB    | 2048.0      | 0.960936      | 0.960936      | 0.960936      | 0.960936      | 0.960936      | 0.960936      | 655606     | 1492    | 731     | 981     | 4649    | 0.0119     | 36               | 1007262/1007262   |
+| 1       | 4        | 3198KB    | 4096.0      | 0.960936      | 0.960936      | 0.960936      | 0.960936      | 0.960936      | 0.960936      | 691428     | 1415    | 541     | 2234    | 2675    | 0.0000     | 63               | 4029050/4029050   |
+| 2       | 4        | 3198KB    | 4096.0      | 0.960936      | 0.960936      | 0.960936      | 0.960936      | 0.960936      | 0.960936      | 444332     | 2219    | 1062    | 1323    | 5570    | 0.0000     | 63               | 2014525/2014525   |
+| 3       | 4        | 3198KB    | 4096.0      | 0.960936      | 0.960936      | 0.960936      | 0.960936      | 0.960936      | 0.960936      | 347967     | 2842    | 1493    | 1793    | 7003    | 0.0004     | 54               | 1343016/1343016   |
+| 1       | 8        | 3198KB    | 8192.0      | 0.960936      | 0.960936      | 0.960936      | 0.960936      | 0.960936      | 0.960936      | 611508     | 1604    | 661     | 2445    | 2976    | 0.0000     | 63               | 8058100/8058100   |
+| 2       | 8        | 3198KB    | 8192.0      | 0.960936      | 0.960936      | 0.960936      | 0.960936      | 0.960936      | 0.960936      | 383486     | 2576    | 1273    | 1513    | 8295    | 0.0000     | 63               | 4029050/4029050   |
+| 3       | 8        | 3198KB    | 8192.0      | 0.960936      | 0.960936      | 0.960936      | 0.960936      | 0.960936      | 0.960936      | 292720     | 3385    | 1773    | 2054    | 7694    | 0.0000     | 63               | 2686033/2686033   |
+| 1       | 16       | 3198KB    | 16384.0     | 0.960936      | 0.960936      | 0.960936      | 0.960936      | 0.960936      | 0.960936      | 573570     | 1712    | 671     | 2654    | 3196    | 0.0000     | 63               | 16116200/16116200 |
+| 2       | 16       | 3198KB    | 16384.0     | 0.960936      | 0.960936      | 0.960936      | 0.960936      | 0.960936      | 0.960936      | 346109     | 2858    | 1503    | 1763    | 9328    | 0.0000     | 63               | 8058100/8058100   |
+| 3       | 16       | 3198KB    | 16384.0     | 0.960936      | 0.960936      | 0.960936      | 0.960936      | 0.960936      | 0.960936      | 252816     | 3924    | 2144    | 2495    | 8496    | 0.0000     | 63               | 5372067/5372067   |
+
+
+Paired KV+Indexer is the DSA workload. It allocates two object sizes in the same segment pool.
+
+### Paired (KV + Indexer) / Skewed=no / Random
+
+
+| Replica | Segments | AllocSize  | Cluster(GB) | UtilRatio_min | UtilRatio_p99 | UtilRatio_p90 | UtilRatio_p50 | UtilRatio_max | UtilRatio_avg | Throughput | Avg(ns) | P50(ns) | P90(ns) | P99(ns) | UtilStdDev | Evictions(round) | Succ/Total        |
+| ------- | -------- | ---------- | ----------- | ------------- | ------------- | ------------- | ------------- | ------------- | ------------- | ---------- | ------- | ------- | ------- | ------- | ---------- | ---------------- | ----------------- |
+| 1       | 1        | 3198+643KB | 1024.0      | 0.952089      | 0.952089      | 0.952183      | 0.952498      | 0.952621      | 0.952437      | 921322     | 1049    | 310     | 1412    | 2164    | 0.0000     | 62               | 1677285/1677285   |
+| 1       | 2        | 3198+643KB | 2048.0      | 0.952162      | 0.952162      | 0.952250      | 0.952451      | 0.952627      | 0.952426      | 810735     | 1197    | 340     | 1964    | 2435    | 0.0001     | 64               | 3354570/3354570   |
+| 2       | 2        | 3198+643KB | 2048.0      | 0.952089      | 0.952089      | 0.952183      | 0.952498      | 0.952621      | 0.952437      | 517602     | 1901    | 761     | 1042    | 5049    | 0.0000     | 62               | 1677285/1677285   |
+| 1       | 4        | 3198+643KB | 4096.0      | 0.952316      | 0.952316      | 0.952360      | 0.952464      | 0.952611      | 0.952467      | 701903     | 1387    | 371     | 2164    | 3025    | 0.0003     | 63               | 6709139/6709139   |
+| 2       | 4        | 3198+643KB | 4096.0      | 0.952195      | 0.952195      | 0.952286      | 0.952501      | 0.952607      | 0.952459      | 429714     | 2292    | 852     | 1132    | 5821    | 0.0005     | 63               | 3354570/3354570   |
+| 3       | 4        | 3198+643KB | 4096.0      | 0.952201      | 0.952201      | 0.952322      | 0.952515      | 0.952627      | 0.952488      | 360298     | 2744    | 1173    | 1473    | 7574    | 0.0002     | 63               | 2236380/2236380   |
+| 1       | 8        | 3198+643KB | 8192.0      | 0.952320      | 0.952320      | 0.952351      | 0.952421      | 0.952586      | 0.952428      | 701440     | 1389    | 361     | 2234    | 2846    | 0.0005     | 63               | 13418278/13418278 |
+| 2       | 8        | 3198+643KB | 8192.0      | 0.952322      | 0.952322      | 0.952345      | 0.952424      | 0.952618      | 0.952449      | 435206     | 2266    | 882     | 1222    | 6001    | 0.0005     | 63               | 6709139/6709139   |
+| 3       | 8        | 3198+643KB | 8192.0      | 0.952278      | 0.952278      | 0.952318      | 0.952444      | 0.952618      | 0.952456      | 330833     | 2991    | 1272    | 1574    | 8035    | 0.0004     | 63               | 4472760/4472760   |
+| 1       | 16       | 3198+643KB | 16384.0     | 0.952246      | 0.952246      | 0.952310      | 0.952412      | 0.952580      | 0.952395      | 643565     | 1517    | 380     | 2534    | 3196    | 0.0004     | 64               | 26836557/26836557 |
+| 2       | 16       | 3198+643KB | 16384.0     | 0.952338      | 0.952338      | 0.952386      | 0.952448      | 0.952580      | 0.952449      | 405947     | 2432    | 942     | 1313    | 6803    | 0.0005     | 63               | 13418278/13418278 |
+| 3       | 16       | 3198+643KB | 16384.0     | 0.952169      | 0.952169      | 0.952259      | 0.952434      | 0.952622      | 0.952423      | 313732     | 3156    | 1323    | 1654    | 9017    | 0.0005     | 63               | 8945519/8945519   |
+
+
+### Paired (KV + Indexer) / Skewed=no / FreeRatioFirst
+
+
+| Replica | Segments | AllocSize  | Cluster(GB) | UtilRatio_min | UtilRatio_p99 | UtilRatio_p90 | UtilRatio_p50 | UtilRatio_max | UtilRatio_avg | Throughput | Avg(ns) | P50(ns) | P90(ns) | P99(ns) | UtilStdDev | Evictions(round) | Succ/Total        |
+| ------- | -------- | ---------- | ----------- | ------------- | ------------- | ------------- | ------------- | ------------- | ------------- | ---------- | ------- | ------- | ------- | ------- | ---------- | ---------------- | ----------------- |
+| 1       | 1        | 3198+643KB | 1024.0      | 0.952089      | 0.952089      | 0.952183      | 0.952498      | 0.952621      | 0.952437      | 823450     | 1183    | 431     | 1584    | 2344    | 0.0000     | 62               | 1677285/1677285   |
+| 1       | 2        | 3198+643KB | 2048.0      | 0.952172      | 0.952172      | 0.952349      | 0.952462      | 0.952627      | 0.952452      | 740063     | 1320    | 472     | 2084    | 2575    | 0.0000     | 63               | 3354570/3354570   |
+| 2       | 2        | 3198+643KB | 2048.0      | 0.952089      | 0.952089      | 0.952183      | 0.952498      | 0.952621      | 0.952437      | 470878     | 2093    | 982     | 1242    | 6462    | 0.0000     | 62               | 1677285/1677285   |
+| 1       | 4        | 3198+643KB | 4096.0      | 0.952345      | 0.952345      | 0.952395      | 0.952481      | 0.952575      | 0.952471      | 659232     | 1484    | 551     | 2305    | 2785    | 0.0000     | 64               | 6709139/6709139   |
+| 2       | 4        | 3198+643KB | 4096.0      | 0.952323      | 0.952323      | 0.952393      | 0.952493      | 0.952627      | 0.952490      | 429940     | 2295    | 1092    | 1313    | 7514    | 0.0000     | 63               | 3354570/3354570   |
+| 3       | 4        | 3198+643KB | 4096.0      | 0.952180      | 0.952180      | 0.952356      | 0.952509      | 0.952606      | 0.952481      | 318057     | 3113    | 1554    | 1864    | 7564    | 0.0000     | 63               | 2236380/2236380   |
+| 1       | 8        | 3198+643KB | 8192.0      | 0.952316      | 0.952316      | 0.952349      | 0.952459      | 0.952594      | 0.952455      | 580198     | 1692    | 671     | 2585    | 3126    | 0.0002     | 64               | 13418278/13418278 |
+| 2       | 8        | 3198+643KB | 8192.0      | 0.952221      | 0.952221      | 0.952316      | 0.952486      | 0.952606      | 0.952450      | 371743     | 2659    | 1332    | 1593    | 6593    | 0.0002     | 64               | 6709139/6709139   |
+| 3       | 8        | 3198+643KB | 8192.0      | 0.952215      | 0.952215      | 0.952316      | 0.952470      | 0.952618      | 0.952455      | 282562     | 3508    | 1814    | 2114    | 8115    | 0.0000     | 63               | 4472760/4472760   |
+| 1       | 16       | 3198+643KB | 16384.0     | 0.952328      | 0.952328      | 0.952368      | 0.952423      | 0.952570      | 0.952436      | 546995     | 1797    | 682     | 2825    | 3466    | 0.0000     | 63               | 26836557/26836557 |
+| 2       | 16       | 3198+643KB | 16384.0     | 0.952288      | 0.952288      | 0.952343      | 0.952419      | 0.952589      | 0.952435      | 325312     | 3043    | 1583    | 1894    | 10300   | 0.0000     | 63               | 13418278/13418278 |
+| 3       | 16       | 3198+643KB | 16384.0     | 0.952282      | 0.952282      | 0.952355      | 0.952458      | 0.952615      | 0.952453      | 242750     | 4088    | 2245    | 2585    | 10350   | 0.0000     | 63               | 8945519/8945519   |
+
+
+### Paired (KV + Indexer) / Skewed=yes / Random
+
+
+| Replica | Segments | AllocSize  | Cluster(GB) | UtilRatio_min | UtilRatio_p99 | UtilRatio_p90 | UtilRatio_p50 | UtilRatio_max | UtilRatio_avg | Throughput | Avg(ns) | P50(ns) | P90(ns) | P99(ns) | UtilStdDev | Evictions(round) | Succ/Total        |
+| ------- | -------- | ---------- | ----------- | ------------- | ------------- | ------------- | ------------- | ------------- | ------------- | ---------- | ------- | ------- | ------- | ------- | ---------- | ---------------- | ----------------- |
+| 1       | 1        | 3198+643KB | 1536.0      | 0.952196      | 0.952196      | 0.952295      | 0.952475      | 0.952628      | 0.952470      | 857346     | 1135    | 311     | 1643    | 2354    | 0.0000     | 63               | 2515927/2515927   |
+| 1       | 2        | 3198+643KB | 2048.0      | 0.952270      | 0.952270      | 0.952310      | 0.952441      | 0.952614      | 0.952444      | 805462     | 1210    | 341     | 1984    | 2474    | 0.0009     | 63               | 3354570/3354570   |
+| 2       | 2        | 3198+643KB | 2048.0      | 0.952421      | 0.952421      | 0.952443      | 0.952507      | 0.952630      | 0.952517      | 720532     | 1357    | 561     | 851     | 5120    | 0.0004     | 37               | 1677285/1677285   |
+| 1       | 4        | 3198+643KB | 4096.0      | 0.952181      | 0.952181      | 0.952313      | 0.952406      | 0.952634      | 0.952417      | 717110     | 1363    | 361     | 2235    | 2856    | 0.0001     | 63               | 6709139/6709139   |
+| 2       | 4        | 3198+643KB | 4096.0      | 0.952268      | 0.952268      | 0.952364      | 0.952494      | 0.952617      | 0.952484      | 458283     | 2151    | 901     | 1192    | 5511    | 0.0051     | 63               | 3354570/3354570   |
+| 3       | 4        | 3198+643KB | 4096.0      | 0.952307      | 0.952307      | 0.952409      | 0.952500      | 0.952630      | 0.952499      | 420616     | 2346    | 1052    | 1382    | 7073    | 0.0070     | 48               | 2236380/2236380   |
+| 1       | 8        | 3198+643KB | 8192.0      | 0.952310      | 0.952310      | 0.952340      | 0.952376      | 0.952612      | 0.952405      | 694136     | 1410    | 371     | 2244    | 2786    | 0.0003     | 64               | 13418278/13418278 |
+| 2       | 8        | 3198+643KB | 8192.0      | 0.952353      | 0.952353      | 0.952384      | 0.952445      | 0.952594      | 0.952456      | 433396     | 2276    | 942     | 1272    | 5931    | 0.0017     | 63               | 6709139/6709139   |
+| 3       | 8        | 3198+643KB | 8192.0      | 0.952324      | 0.952324      | 0.952357      | 0.952463      | 0.952608      | 0.952456      | 329256     | 3006    | 1363    | 1673    | 7003    | 0.0047     | 63               | 4472760/4472760   |
+| 1       | 16       | 3198+643KB | 16384.0     | 0.952264      | 0.952264      | 0.952299      | 0.952343      | 0.952590      | 0.952375      | 641829     | 1521    | 381     | 2565    | 3256    | 0.0003     | 63               | 26836557/26836557 |
+| 2       | 16       | 3198+643KB | 16384.0     | 0.952246      | 0.952246      | 0.952336      | 0.952423      | 0.952608      | 0.952426      | 399689     | 2470    | 1012    | 1403    | 6652    | 0.0011     | 63               | 13418278/13418278 |
+| 3       | 16       | 3198+643KB | 16384.0     | 0.952300      | 0.952300      | 0.952360      | 0.952451      | 0.952606      | 0.952447      | 303763     | 3261    | 1433    | 1763    | 9107    | 0.0021     | 63               | 8945519/8945519   |
+
+
+### Paired (KV + Indexer) / Skewed=yes / FreeRatioFirst
+
+
+| Replica | Segments | AllocSize  | Cluster(GB) | UtilRatio_min | UtilRatio_p99 | UtilRatio_p90 | UtilRatio_p50 | UtilRatio_max | UtilRatio_avg | Throughput | Avg(ns) | P50(ns) | P90(ns) | P99(ns) | UtilStdDev | Evictions(round) | Succ/Total        |
+| ------- | -------- | ---------- | ----------- | ------------- | ------------- | ------------- | ------------- | ------------- | ------------- | ---------- | ------- | ------- | ------- | ------- | ---------- | ---------------- | ----------------- |
+| 1       | 1        | 3198+643KB | 1536.0      | 0.952196      | 0.952196      | 0.952295      | 0.952475      | 0.952628      | 0.952470      | 775651     | 1258    | 441     | 1783    | 2525    | 0.0000     | 63               | 2515927/2515927   |
+| 1       | 2        | 3198+643KB | 2048.0      | 0.952206      | 0.952206      | 0.952284      | 0.952539      | 0.952758      | 0.952507      | 732598     | 1334    | 481     | 2064    | 2565    | 0.0000     | 64               | 3354570/3354570   |
+| 2       | 2        | 3198+643KB | 2048.0      | 0.952421      | 0.952421      | 0.952443      | 0.952507      | 0.952630      | 0.952517      | 624654     | 1570    | 791     | 1042    | 5070    | 0.0004     | 37               | 1677285/1677285   |
+| 1       | 4        | 3198+643KB | 4096.0      | 0.952378      | 0.952378      | 0.952465      | 0.952601      | 0.952898      | 0.952610      | 667880     | 1466    | 561     | 2304    | 2735    | 0.0000     | 63               | 6709139/6709139   |
+| 2       | 4        | 3198+643KB | 4096.0      | 0.952241      | 0.952241      | 0.952341      | 0.952545      | 0.952758      | 0.952532      | 424387     | 2325    | 1112    | 1373    | 5290    | 0.0006     | 64               | 3354570/3354570   |
+| 3       | 4        | 3198+643KB | 4096.0      | 0.952294      | 0.952294      | 0.952346      | 0.952524      | 0.952630      | 0.952503      | 332752     | 2974    | 1543    | 1864    | 7164    | 0.0075     | 55               | 2236380/2236380   |
+| 1       | 8        | 3198+643KB | 8192.0      | 0.952463      | 0.952463      | 0.952498      | 0.952691      | 0.953072      | 0.952725      | 584455     | 1679    | 671     | 2576    | 3116    | 0.0004     | 64               | 13418278/13418278 |
+| 2       | 8        | 3198+643KB | 8192.0      | 0.952329      | 0.952329      | 0.952433      | 0.952605      | 0.952899      | 0.952602      | 371745     | 2659    | 1323    | 1592    | 9427    | 0.0000     | 64               | 6709139/6709139   |
+| 3       | 8        | 3198+643KB | 8192.0      | 0.952365      | 0.952365      | 0.952410      | 0.952588      | 0.952800      | 0.952578      | 281189     | 3525    | 1853    | 2154    | 7875    | 0.0000     | 63               | 4472760/4472760   |
+| 1       | 16       | 3198+643KB | 16384.0     | 0.952303      | 0.952303      | 0.952374      | 0.952830      | 0.953388      | 0.952796      | 540381     | 1814    | 691     | 2875    | 3527    | 0.0003     | 64               | 26836557/26836557 |
+| 2       | 16       | 3198+643KB | 16384.0     | 0.952518      | 0.952518      | 0.952539      | 0.952636      | 0.953128      | 0.952681      | 320795     | 3086    | 1603    | 1974    | 7744    | 0.0005     | 64               | 13418278/13418278 |
+| 3       | 16       | 3198+643KB | 16384.0     | 0.952450      | 0.952450      | 0.952512      | 0.952592      | 0.952953      | 0.952614      | 240086     | 4134    | 2274    | 2625    | 10470   | 0.0000     | 63               | 8945519/8945519   |
+
+
diff --git a/docs/source/performance/allocator-benchmark-result.md b/docs/source/performance/allocator-benchmark-result.md
index 6a04323b..fd7c9a5d 100644
--- a/docs/source/performance/allocator-benchmark-result.md
+++ b/docs/source/performance/allocator-benchmark-result.md
@@ -1,5 +1,14 @@
 # Allocator Performance
 
+Mooncake Store uses two allocators at different layers:
+
+- **MmapArena** — backs the mmap buffer path used by SGLang HiCache for host KV cache memory. When explicitly enabled, it pre-allocates a hugepage-backed pool at startup and serves allocations via lock-free atomic bump pointer (~50ns per allocation). Configured via `MC_MMAP_ARENA_POOL_SIZE` (explicit opt-in) and `MC_DISABLE_MMAP_ARENA`.
+- **OffsetAllocator** — manages object-level allocation within Mooncake Store segments for KV cache entries. Evaluated below.
+
+These are complementary: MmapArena handles the backing memory, OffsetAllocator manages logical object placement within that memory.
+
+## OffsetAllocator
+
 We evaluated the performance of [OffsetAllocator](https://github.com/sebbbi/OffsetAllocator), the default memory allocator in Mooncake Store. This allocator is responsible for allocating memory from mounted segments to store the KV cache.
 
 In this context, the most important metric is **memory utilization**, defined as the ratio between the amount of memory that can be successfully allocated and the total available memory. A higher utilization means that more KV tensors can be cached, thereby accelerating LLM tasks. However, due to memory fragmentation, allocation may fail even when the allocated memory is well below the total available capacity.
@@ -185,4 +194,33 @@ avg alloc time: 145.575738 ns/op
 util ratio (min / p99 / p90 / p50 / max / avg):
 0.569255 / 0.712076 / 0.781224 / 0.855046 / 0.976057 / 0.848873
 avg alloc time: 142.508508 ns/op
-```
\ No newline at end of file
+```
+
+**OffsetAllocator (Before Optimization)**
+
+```
+util ratio (min / p99 / p90 / p50 / max / avg):
+0.569255 / 0.712076 / 0.781224 / 0.855046 / 0.976057 / 0.848873
+avg alloc time: 142.508508 ns/op
+```
+
+### Paired KV/Indexer Allocation Benchmark (DSA)
+
+In the DSA scenario, Mooncake Store stores both KV cache objects and indexer objects.
+We evaluated OffsetAllocator under a paired allocation workload using object sizes derived from a GLM-5.1-FP8 DSA serving configuration.
+
+In this configuration, Mooncake Store stores data at the page granularity, where each page contains 64 tokens. Therefore, the object sizes observed by the allocator are the packed page-level sizes rather than the per-token sizes:
+
+- KV cache object size: 3.12 MB
+- Indexer object size: 643 KB
+
+Each benchmark round samples `N` uniformly from `[1, 128]`, then issues `N` KV cache allocations followed by `N` indexer allocations. The pool size is 1024 GB. On allocation failure, the benchmark randomly evicts 5% of all live objects and retries. The utilization ratio below is sampled only when eviction is triggered, immediately before evicting objects.
+
+**OffsetAllocator**
+
+```
+util ratio (min / p99 / p90 / p50 / max / avg):
+0.951814 / 0.951879 / 0.952077 / 0.952382 / 0.952621 / 0.952343
+eviction-trigger samples: 243
+avg alloc time: 352.207986 ns/op
+```
diff --git a/docs/source/performance/mooncake/index.md b/docs/source/performance/mooncake/index.md
new file mode 100644
index 00000000..25f2b988
--- /dev/null
+++ b/docs/source/performance/mooncake/index.md
@@ -0,0 +1,20 @@
+# Mooncake Performance Benchmarks
+
+Benchmarks evaluating Mooncake Store's core storage, allocation, and cache hierarchy behavior.
+
+| Document | Area | Key Findings |
+|----------|------|---------------|
+| [Storage Benchmark](../storage-benchmark) | Mooncake Store storage | Measures end-to-end storage performance across Mooncake Store operations and deployment configurations. |
+| [Allocator Benchmark](../allocator-benchmark-result) | Segment allocation | The optimized OffsetAllocator significantly improves utilization for uniform-size LLM KV cache allocation patterns. |
+| [Allocation Strategy Benchmark](../allocation-strategy-benchmark-result) | Allocation routing | Compares random and free-ratio-first allocation across segments, replicas, skewed capacity, and DSA-style KV+indexer workloads. |
+| [SSD Offload Benchmark](../ssd-offload-benchmark-results) | Cache hierarchy | SSD offload extends the KV cache hierarchy with NVMe, reducing the performance cliff after DRAM cache capacity is exhausted in long multi-turn conversations. |
+
+:::{toctree}
+:maxdepth: 1
+:hidden:
+
+../storage-benchmark
+../allocator-benchmark-result
+../allocation-strategy-benchmark-result
+../ssd-offload-benchmark-results
+:::
diff --git a/docs/source/performance/sglang/index.md b/docs/source/performance/sglang/index.md
new file mode 100644
index 00000000..3a4c1804
--- /dev/null
+++ b/docs/source/performance/sglang/index.md
@@ -0,0 +1,16 @@
+# SGLang Integration Performance Benchmarks
+
+Benchmarks evaluating Mooncake's integration with SGLang across PD disaggregation and HiCache hierarchical KV cache storage.
+
+| Document | Scenario | Key Findings |
+|----------|----------|---------------|
+| [PD Disaggregation Performance](../sglang-benchmark-results-v1) | SGLang PD disaggregation with Mooncake Transfer Engine | 1P1D PD disaggregation achieves approximately **30% lower ITL** while maintaining comparable throughput against two regular instances. |
+| [HiCache with Mooncake Backend Benchmark](../sglang-hicache-benchmark-results-v1) | SGLang HiCache using Mooncake Store as L3 storage | Mooncake-backed HiCache improves prefill performance in multi-turn workloads by maintaining higher KV cache hit rates as conversation rounds grow. |
+
+:::{toctree}
+:maxdepth: 1
+:hidden:
+
+../sglang-benchmark-results-v1
+../sglang-hicache-benchmark-results-v1
+:::
diff --git a/docs/source/performance/vllm-benchmark-results-v0.2.md b/docs/source/performance/vllm-benchmark-results-v0.2.md
index 685ac934..4f3bd612 100644
--- a/docs/source/performance/vllm-benchmark-results-v0.2.md
+++ b/docs/source/performance/vllm-benchmark-results-v0.2.md
@@ -3,7 +3,7 @@ Here are some preview mooncake benchmark results on A10 with up to 2 RDMA NICs.
 
 In addition, we are also coordinating resources to integrate some machines with more RDMA NICs and more advanced GPUs. The official benchmark results will be released in due time.
 
-### Varying tp (input length = 1024, qps = 2, output length =6)
+## Varying tp (input length = 1024, qps = 2, output length =6)
 | Setting | num_rdma_nic | Successful Requests | Duration (s) | Total Input Tokens | Total Generated Tokens | Req Throughput (req/s) | Output Token Throughput (tok/s) | Total Token Throughput (tok/s) | Mean TTFT (ms) | Median TTFT (ms) | P99 TTFT (ms) | Mean TPOT (ms) | Median TPOT (ms) | P99 TPOT (ms) | Mean ITL (ms) | Median ITL (ms) | P99 ITL (ms) |
 |-----------------|--------------|---------------------|--------------|--------------------|------------------------|------------------------|---------------------------------|-------------------------------|----------------|-----------------|--------------|---------------|------------------|--------------|--------------|----------------|-------------|
 | tp = 1          | 2            | 200                 | 99.47        | 201995             | 1200                   | 2.01                   | 12.06                           | 2042.74                      | 1056.76        | 635.00           | 4006.59      | 97.08         | 26.94            | 781.91       | 97.01         | 14.05           | 2205.51     |
@@ -16,7 +16,7 @@ In addition, we are also coordinating resources to integrate some machines with
 | tp = 2          | TCP          | 200                 | 98.97        | 201995             | 1200                   | 2.02                   | 12.12                           | 2053.03                      | 333.74         | 251.32           | 954.63       | 28.74         | 15.49            | 161.24       | 28.70         | 15.35           | 393.52      |
 | tp = 4          | TCP          | 200                 | 98.78        | 201995             | 1200                   | 2.02                   | 12.15                           | 2056.94                      | 205.37         | 162.92           | 463.70       | 21.54         | 16.51            | 94.04        | 21.51         | 16.56           | 170.54      |
 
-### Varying qps (length = 1024, tp = 4, output length =6)
+## Varying qps (length = 1024, tp = 4, output length =6)
 |Setting | num_rdma_nic | Successful Requests | Duration (s) | Total Input Tokens | Total Generated Tokens | Req Throughput (req/s) | Output Token Throughput (tok/s) | Total Token Throughput (tok/s) | Mean TTFT (ms) | Median TTFT (ms) | P99 TTFT (ms) | Mean TPOT (ms) | Median TPOT (ms) | P99 TPOT (ms) | Mean ITL (ms) | Median ITL (ms) | P99 ITL (ms) |
 |-----------------|--------------|---------------------|--------------|--------------------|------------------------|------------------------|---------------------------------|-------------------------------|----------------|-----------------|--------------|---------------|------------------|--------------|--------------|----------------|-------------|
 | qps = 2         | 2            | 200                 | 98.77        | 201995             | 1200                   | 2.02                   | 12.15                           | 2057.33                      | 200.64         | 156.62           | 478.22       | 22.63         | 17.35            | 99.61        | 22.60         | 17.08           | 186.25      |
@@ -32,7 +32,7 @@ In addition, we are also coordinating resources to integrate some machines with
 | qps = 6         | TCP          | 200                 | 33.49        | 201995             | 1200                   | 5.97                   | 35.83                           | 6067.71                      | 907.74         | 514.85           | 3253.93      | 122.75        | 45.51            | 648.40       | 122.56        | 18.09           | 2282.92     |
 | qps = 8         | TCP          | 200                 | 28.39        | 201995             | 1200                   | 7.04                   | 42.26                           | 7156.09                      | 6714.57        | 7885.09          | 11787.51     | 1116.06       | 408.32           | 4645.25      | 1114.29       | 46.87           | 21898.03    |
 
-### Varying input length (tp = 4, qps = 2, output length =6)
+## Varying input length (tp = 4, qps = 2, output length =6)
 | Setting | num_rdma_nic | Successful Requests | Duration (s) | Total Input Tokens | Total Generated Tokens | Req Throughput (req/s) | Output Token Throughput (tok/s) | Total Token Throughput (tok/s) | Mean TTFT (ms) | Median TTFT (ms) | P99 TTFT (ms) | Mean TPOT (ms) | Median TPOT (ms) | P99 TPOT (ms) | Mean ITL (ms) | Median ITL (ms) | P99 ITL (ms) |
 |-----------------|--------------|---------------------|--------------|--------------------|------------------------|------------------------|---------------------------------|-------------------------------|----------------|-----------------|--------------|---------------|------------------|--------------|--------------|----------------|-------------|
 | 1024   | 2            | 200                 | 98.77        | 201995             | 1200                   | 2.02                   | 12.15                           | 2057.32                      | 195.47         | 151.55           | 482.84       | 22.83         | 19.27            | 96.55        | 22.81         | 18.12           | 158.16      |
diff --git a/docs/source/performance/vllm-benchmark-results-v1.md b/docs/source/performance/vllm-benchmark-results-v1.md
index 6a5acc57..05032748 100644
--- a/docs/source/performance/vllm-benchmark-results-v1.md
+++ b/docs/source/performance/vllm-benchmark-results-v1.md
@@ -2,7 +2,7 @@
 Here are some preview MooncakeStore benchmark results on A10 with "Qwen/Qwen2.5-7B-Instruct-GPTQ-Int4".
 
 
-### Varying PD ratio (input length = 1024, qps = 2, output length =6, num of requests = 200)
+## Varying PD ratio (input length = 1024, qps = 2, output length =6, num of requests = 200)
 | Configuration              | Backend              | Duration (s) | Output Token Throughput (tok/s) | Total Token Throughput (tok/s) | Mean TTFT (ms) | Median TTFT (ms) | P99 TTFT (ms) | Mean TPOT (ms) | Median TPOT (ms) | P99 TPOT (ms) | Mean ITL (ms) | Median ITL (ms) | P99 ITL (ms) |
 |----------------------------|----------------------|--------------|---------------------------------|-------------------------------|----------------|-----------------|--------------|---------------|------------------|--------------|--------------|----------------|-------------|
 | 2P2D tp = 1       | Redis                | 99.47        | 12.06                           | 2042.75                      | 844.28         | 666.84           | 2270.91      | 16.88         | 11.57            | 104.83       | 16.84         | 11.56           | 239.67      |
diff --git a/docs/source/performance/vllm/index.md b/docs/source/performance/vllm/index.md
new file mode 100644
index 00000000..b32561c7
--- /dev/null
+++ b/docs/source/performance/vllm/index.md
@@ -0,0 +1,18 @@
+# vLLM Integration Performance Benchmarks
+
+Benchmarks evaluating Mooncake's integration with vLLM across different backends and scenarios.
+
+| Document | Backend | Key Findings |
+|----------|---------|---------------|
+| [vLLM V1 + MooncakeConnector](../vllm-v1-support-benchmark) | vLLM V1 | 1P1D PD disaggregation on H800 with 8x RoCE: **142.25 GB/s** peak transfer bandwidth (71.1% of theoretical), KV transfer overhead just **4.2%** of total TTFT at 32K tokens |
+| [vLLM V1 + MooncakeStore vs Redis](../vllm-benchmark-results-v1) | vLLM V1 | MooncakeStore RDMA consistently outperforms Redis across all XpYd topologies — e.g., **~32% lower** mean TTFT in 2P2D tp=2 |
+| [vLLM V0 + MooncakeConnector (Legacy)](../vllm-benchmark-results-v0.2) | vLLM V0 | TP=4 reduces TTFT by ~80% vs TP=1; RDMA provides significant latency advantage over TCP across varying QPS and input lengths |
+
+:::{toctree}
+:maxdepth: 1
+:hidden:
+
+../vllm-v1-support-benchmark
+../vllm-benchmark-results-v1
+../vllm-benchmark-results-v0.2
+:::
diff --git a/docs/source/python-api-reference/dataproto-structured-object-transfer.md b/docs/source/python-api-reference/dataproto-structured-object-transfer.md
new file mode 100644
index 00000000..29b19da8
--- /dev/null
+++ b/docs/source/python-api-reference/dataproto-structured-object-transfer.md
@@ -0,0 +1,168 @@
+# DataProto structured object usage
+
+Mooncake can store DataProto-like objects as structured objects so callers can pass a lightweight handle between stages and materialize only the fields they need.
+
+A DataProto-like object is any object with these mapping-like attributes:
+
+- `batch`: tensor or ndarray fields indexed by batch row.
+- `non_tensor_batch`: per-row non-tensor fields.
+- `meta_info`: small metadata for the whole batch.
+
+Plain dictionaries are also accepted. A dictionary with only `batch`, `non_tensor_batch`, and `meta_info` keys is treated as an envelope; other dictionaries are treated as `batch` fields.
+
+## Public API
+
+```python
+from mooncake.structured_object_store import (
+    BundleTransferPolicy,
+    MooncakeBundleTransfer,
+    export_dataproto_ref,
+    import_dataproto_ref,
+    tensor_object_buffer,
+)
+
+transfer = MooncakeBundleTransfer(store, key_prefix="rl")
+
+ref = transfer.put_dataproto(
+    data,
+    namespace="rollout",
+    partition="step-1",
+    stage="rollout",
+)
+
+ref = transfer.append_dataproto_fields(
+    ref,
+    logprob_data,
+    stage="old_log_prob",
+)
+
+subset = transfer.get_dataproto(
+    ref,
+    fields=["input_ids", "old_log_probs"],
+    meta_info_keys=["step"],
+)
+
+handle = export_dataproto_ref(ref)
+ref = import_dataproto_ref(handle)
+transfer.cleanup_dataproto(ref)
+```
+
+## Lightweight handles
+
+`MooncakeDataProtoRef` contains only DataProto-level routing information:
+
+- `stage_refs`: stage name to structured object reference.
+- `field_index`: field name to `(stage, member, section)`.
+- `batch_size`, `namespace`, `partition`, `meta_info`, and optional `global_indexes`.
+
+It does not duplicate dtype, shape, chunk layout, or range metadata. Those details remain in the structured object manifest. Use `dataproto_manifest_view(ref)` when a caller needs an introspection view derived from the manifests.
+
+For process boundaries, use `export_dataproto_ref(ref)`. The exported handle is JSON-safe and contains manifest keys instead of embedded manifest payloads. `get_dataproto()`, `append_dataproto_fields()`, `dataproto_manifest_view()`, and `cleanup_dataproto()` accept either an in-memory ref or an exported handle.
+
+## Writing fields
+
+`put_dataproto()` writes one structured object for the requested stage. `append_dataproto_fields()` writes another structured object and updates the handle. Existing fields are not rewritten.
+
+Duplicate field names are rejected by default. Use `overwrite=True` only when replacing all fields from an existing stage; after the new stage object is written successfully, the old stage object is removed.
+
+Field names are global within a ref. A `batch` field and a `non_tensor_batch` field cannot use the same name.
+
+## Reading fields
+
+`get_dataproto()` supports:
+
+- `fields`: mixed batch and non-tensor field selection.
+- `batch_fields`: batch-only selection.
+- `non_tensor_fields`: non-tensor-only selection.
+- `meta_info_keys`: metadata selection.
+- `data_cls`: return a DataProto-like class instead of a plain dict.
+- `destinations`: caller-provided output buffers.
+- `rows`: row selection with a Python `slice`, `StructuredMemberSlice`, or an integer row-index sequence.
+
+Use `rows` to materialize the same batch rows across all selected `batch` and `non_tensor_batch` fields:
+
+```python
+subset = transfer.get_dataproto(
+    ref,
+    fields=["input_ids", "text", "rewards"],
+    rows=slice(128, 256),
+)
+
+gathered = transfer.get_dataproto(
+    ref,
+    batch_fields=["input_ids"],
+    rows=[7, 3, 9],
+)
+```
+
+Row selection supports `axis=0`. Tensor and ndarray batch fields, including native Mooncake tensor fields, are read by byte ranges from the stored payload. Structured object `non_tensor_batch` fields read only the selected row metadata and payload ranges. `destinations` may be combined with `rows` for fields that support caller-provided output buffers. Use `raw_destination(ptr, size, owner, pre_registered=True)` for BufferPool or otherwise pre-registered destination memory.
+
+The result is a plain dictionary when `data_cls` is omitted:
+
+```python
+{
+    "batch": {...},
+    "non_tensor_batch": {...},
+    "meta_info": {...},
+}
+```
+
+If `data_cls` is provided, Mooncake first tries `data_cls.from_dict(batch, non_tensor_batch, meta_info=meta_info)`, then falls back to `data_cls(batch=..., non_tensor_batch=..., meta_info=...)`.
+
+## Tensor and ndarray behavior
+
+Tensor fields are stored through the best available Mooncake path:
+
+1. `tensor_object_buffer(ptr, size, owner, batch_size=...)` with `copy_mode="zero_copy"` uses `put_tensor_from()` directly.
+2. Torch tensors use the store tensor API when available.
+3. If a tensor-native path is unavailable, Mooncake falls back to a serialized tensor payload.
+4. Scalar tensors use a correctness fallback until the native tensor codec preserves zero-dimensional shape.
+
+Numeric numpy arrays are stored as structured ndarray members. Contiguous arrays can be passed through without copying; non-contiguous arrays are made contiguous before storage. Row slices are materialized through structured range reads when the backend supports them.
+
+`non_tensor_batch` object arrays are structured-encoded according to their contents. Numeric scalar object arrays, ragged tensors, bytes, strings, JSON-like values, and selected media payloads have explicit codecs. These fields are serialized by design and should not be treated as zero-copy tensor data.
+
+## Materializing into caller buffers
+
+`destinations` can reuse caller-provided buffers:
+
+```python
+dst = np.empty((rows, width), dtype=np.int64)
+result = transfer.get_dataproto(ref, batch_fields=["input_ids"], destinations={"input_ids": dst})
+assert result["batch"]["input_ids"] is dst
+```
+
+For tensor payloads stored as Mooncake tensor objects or tensor-object buffers, pass a `tensor_object_buffer` destination:
+
+```python
+lease = pool.acquire(nbytes)
+result = transfer.get_dataproto(
+    ref,
+    batch_fields=["hidden_states"],
+    destinations={
+        "hidden_states": tensor_object_buffer(
+            lease.ptr,
+            lease.size,
+            lease,
+            batch_size=batch_size,
+        )
+    },
+)
+```
+
+The destination owner or lease must remain alive for as long as the materialized data may be used.
+
+## Copy policy
+
+Control-plane metadata and manifests are small and always use the copy path. Tensor and ndarray payloads use the configured payload policy:
+
+```python
+policy = BundleTransferPolicy(copy_mode="zero_copy")
+transfer.put_structured_object(payload, policy=policy)
+```
+
+`copy_mode="zero_copy"` requires tensor payloads to be provided as `tensor_object_buffer`; plain torch tensors are rejected because they do not expose a registered tensor-object buffer.
+
+## Cleanup
+
+`cleanup_dataproto(ref)` removes all stage objects referenced by the handle. It accepts both in-memory refs and exported transport handles.
diff --git a/docs/source/python-api-reference/ep-backend.md b/docs/source/python-api-reference/ep-backend.md
index 411348a5..5619349f 100644
--- a/docs/source/python-api-reference/ep-backend.md
+++ b/docs/source/python-api-reference/ep-backend.md
@@ -1,60 +1,62 @@
-# Mooncake EP & Mooncake Backend
+# Mooncake EP & Mooncake Backend (PG)
 
 ## Overview
 
-Mooncake EP is an adaption of [DeepEP](https://github.com/deepseek-ai/DeepEP) that supports **fault tolerance** and fast data transfer with **IBGDA**, designed as a critical component for large-scale, latency-sensitive MoE (Mixture of Experts) inference. Mooncake EP aims to retain full compatibility with the DeepEP API, with the addition of an `active_ranks` tensor passed to both the `dispatch` and `combine` functions to capture information about rank activeness. By integrating with the EPLB module, Mooncake EP ensures fault tolerance during MoE inference, enabling robust performance even in large-scale, fault-prone environments.
+Mooncake provides two closely related components for fault-tolerant MoE
+inference:
 
-Mooncake Backend is a PyTorch distributed backend (a replacement for NCCL and Gloo) that provides **fault-tolerant collective communication primitives** and can be seamlessly integrated into machine learning systems. Built with the [Transfer Engine](../design/transfer-engine/index.md), Mooncake Backend ensures that collective communications can continue even in the event of rank failures. Furthermore, it reports these failures to the upper layers of the system, allowing for graceful error handling without disrupting ongoing operations.
+- **Mooncake Backend (PG)** is a `torch.distributed` ProcessGroup backend. It
+  registers the `mooncake` accelerator backend and the `mooncake-cpu` backend,
+  implements common collective and point-to-point APIs, tracks active ranks, and
+  exposes elastic recovery helpers.
+- **Mooncake EP** is an expert-parallel dispatch/combine runtime for
+  latency-sensitive MoE inference. It follows the DeepEP low-latency programming
+  model while adding rank activeness awareness and Mooncake transport support.
 
-## Usage
+The usual integration pattern is to initialize a Mooncake process group first,
+then construct a Mooncake EP `Buffer` from that group. The process group is used
+both for regular collectives and for exchanging EP bootstrap metadata.
 
-### Mooncake EP
+For implementation details, see the [Mooncake Backend (PG) design guide](../design/mooncake-backend-pg.md)
+and the [Mooncake EP design guide](../design/mooncake-ep.md).
 
-> **Note:** Mooncake EP currently supports only the low-latency transfer mode.
+## Installation and build notes
 
-The API is largely consistent with DeepEP's, with only minor differences in a few parameters. Mooncake EP exposes a `Buffer` that can be imported from `mooncake.mooncake_ep_buffer`. For example, refer to `mooncake-wheel/tests/test_mooncake_ep.py`.
+Mooncake EP and PG are included in CUDA-enabled Mooncake wheels. When building
+from source, enable the EP/PG extensions with:
 
-#### Buffer.get_buffer_size_hint()
-
-**Signature:**
-
-```python
-@staticmethod
-def get_ep_buffer_size_hint(num_max_dispatch_tokens_per_rank: int, hidden: int, num_ranks: int, num_experts: int) -> int
+```bash
+cmake .. -DWITH_EP=ON
 ```
 
-Calculates the number of bytes to pre-allocate for data transfer.
+The extensions are compiled against a specific PyTorch version. At import time,
+`mooncake.pg` and `mooncake.ep` load version-suffixed extension modules that
+match the active `torch.__version__`. If the current PyTorch version does not
+match a built extension, import will fail with a message such as
+`Mooncake PG was not built against torch==...`.
 
-#### Buffer.\_\_init\_\_()
+## Mooncake Backend (PG) quick start
 
-**Signature:**
+### CUDA backend
 
 ```python
-def __init__(self, group: dist.ProcessGroup, num_ep_buffer_bytes: int = 0)
-```
+import os
 
-The constructor. Ensure that only one instance is created.
-
-- **group**: Must be a Mooncake Backend process group.
-- **num_ep_buffer_bytes**: The number of bytes acquired with `Buffer.get_buffer_size_hint()`
-
-#### Buffer.dispatch/Buffer.combine
-
-**Signature:** Similar to DeepEP's `low_latency_dispatch`/`low_latency_combine`, with two additional parameters:
-
-- **active_ranks**: A tensor of shape `(num_ranks,)` containing values of 0 or 1. The indices of the broken ranks will be set to 0.
-- **timeout_us**: The timeout in microseconds for a rank to be considered broken. Set to -1 for infinite timeout.
-
-### Mooncake Backend
-
-Basic usage:
-
-```python
 import torch
 import torch.distributed as dist
 from mooncake import pg
 
-active_ranks = torch.ones((world_size,), dtype=torch.int32, device="cuda")
+
+rank = int(os.environ["RANK"])
+world_size = int(os.environ["WORLD_SIZE"])
+local_rank = int(os.environ.get("LOCAL_RANK", rank))
+
+torch.cuda.set_device(local_rank)
+device = torch.device("cuda", local_rank)
+
+# Backend-level active-rank mask. Use int32 and place it on the backend device.
+active_ranks = torch.ones(world_size, dtype=torch.int32, device=device)
+
 dist.init_process_group(
     backend="mooncake",
     rank=rank,
@@ -62,46 +64,363 @@ dist.init_process_group(
     pg_options=pg.MooncakeBackendOptions(active_ranks),
 )
 
-dist.all_gather(...)           # Standard API usage
-assert active_ranks.all()  # Verify that no ranks are broken
+x = torch.tensor([rank + 1], dtype=torch.int32, device=device)
+dist.all_reduce(x, op=dist.ReduceOp.SUM)
+print(f"rank={rank}, all_reduce={int(x.cpu())}")
 ```
 
-For a full example, see `mooncake-wheel/tests/test_mooncake_backend.py`.
+Run it with the usual PyTorch launcher, for example:
 
----
+```bash
+torchrun --nproc-per-node=2 pg_quickstart.py
+```
 
-Recover usage (e.g., wants to recover rank #2):
+### CPU backend
+
+Use `backend="mooncake-cpu"` and put `active_ranks` on CPU:
 
 ```python
-# For the healthy processes, execute:
-import torch
-import torch.distributed as dist
-from mooncake import pg
-
-...
-
-broken_rank = 2
-backend = dist.group.WORLD._get_backend(torch.device("cpu"))
-while True:
-    (peer_state,) = pg.get_peer_state(backend, [broken_rank])
-    if peer_state:
-        pg.recover_ranks(backend, [broken_rank])
-        break
-    else:
-        # Handle ongoing logic, like inference
-        pass
-
-# For the new process, execute:
+active_ranks = torch.ones(world_size, dtype=torch.int32)
 dist.init_process_group(
     backend="mooncake-cpu",
-    rank=broken_rank,
-    world_size=num_processes,
-    pg_options=pg.MooncakeBackendOptions(
-        torch.ones((num_processes,), dtype=torch.int32),
-        is_extension=True,  # Must set this option to True
-    ),
+    rank=rank,
+    world_size=world_size,
+    pg_options=pg.MooncakeBackendOptions(active_ranks),
 )
 ```
 
-For a full example, see `mooncake-wheel/tests/test_mooncake_backend_elastic.py`.
+### Selecting network devices
 
+To explicitly restrict Mooncake to a list of NIC / HCA devices, call
+`pg.set_device_filter(...)` before `init_process_group()`:
+
+```python
+from mooncake import pg
+
+pg.set_device_filter(["mlx5_1", "mlx5_2"])
+```
+
+For test and benchmark commands, the same setting is commonly passed through
+`MOONCAKE_PGTEST_DEVICE_FILTERS=mlx5_1,mlx5_2`.
+
+## Mooncake Backend (PG) API reference
+
+### `MooncakeBackendOptions`
+
+```python
+pg.MooncakeBackendOptions(active_ranks)
+pg.MooncakeBackendOptions(active_ranks, is_extension)
+pg.MooncakeBackendOptions(active_ranks, is_extension, max_world_size)
+```
+
+Arguments:
+
+- `active_ranks`: `torch.int32` tensor used as the backend-level rank-health
+  mask. For `mooncake`, it must be on the accelerator device; for
+  `mooncake-cpu`, it must be on CPU. When `max_world_size` is set, size this
+  tensor to `max_world_size`, not the current visible world size.
+- `is_extension`: set to `True` for a replacement or joining process that will
+  enter an existing group through `join_group()`.
+- `max_world_size`: optional upper bound for reserved rank slots. It lets
+  healthy ranks reserve inactive future ranks while keeping
+  `dist.get_world_size()` equal to the current active size.
+
+### Utility functions
+
+| Function | Purpose | Notes |
+| --- | --- | --- |
+| `pg.set_host_ip(host_ip)` | Override the host IP used by the backend. | Call before `init_process_group()`. |
+| `pg.set_device_filter(filters)` | Restrict NIC/HCA selection. | Call before `init_process_group()`. |
+| `pg.set_transfer_engine(engine)` | Reuse an external `TransferEngine`. | The engine must outlive all process groups. |
+| `pg.get_preferred_hca(backend, location)` | Query topology-preferred HCA for a location. | Useful for topology-aware placement/debugging. |
+| `pg.get_active_ranks(backend)` | Return the backend active-rank tensor. | Used by EP fallback and recovery paths. |
+| `pg.get_num_synced_ranks(backend)` | Return the number of ranks synchronized by the backend. | Diagnostic helper. |
+| `pg.extend_group_size_to(backend, size)` | Reserve additional inactive ranks. | Newly extended ranks do not participate until recovered. |
+| `pg.get_peer_state(backend, ranks)` | Check whether candidate ranks have published peer metadata. | Collective among healthy ranks. |
+| `pg.recover_ranks(backend, ranks)` | Activate ready ranks and publish extension state. | Requires peer metadata to be ready. |
+| `pg.join_group(backend)` | Joiner-side blocking call for extension ranks. | Used after `is_extension=True` initialization. |
+
+### Supported distributed operations
+
+Mooncake Backend implements the following `torch.distributed` APIs. Support may
+depend on device type, dtype, PyTorch version, and whether the current backend is
+`mooncake` or `mooncake-cpu`; run the PG tests on the target environment before
+production use.
+
+| API family | Examples | Notes |
+| --- | --- | --- |
+| Collectives | `all_reduce`, `broadcast`, `all_gather`, `all_gather_into_tensor`, `reduce_scatter_tensor`, `all_to_all`, `barrier`, `reduce`, `gather`, `scatter` | Active ranks participate; inactive ranks are skipped by backend internals. |
+| Async work | `dist.all_reduce(..., async_op=True)` | Wait on the returned work object, then synchronize the device stream as needed. |
+| P2P | `isend`, `irecv`, `batch_isend_irecv` | Single-tensor P2P is routed through the Mooncake P2P shim. |
+
+## Elastic recovery protocol
+
+Mooncake PG supports a two-sided recovery protocol. Existing healthy ranks poll
+for replacement rank readiness, then activate those ranks. Replacement ranks
+start in extension mode, publish metadata, and wait until healthy ranks recover
+them.
+
+### Healthy-rank side
+
+```python
+from mooncake import pg
+
+active_ranks = torch.tensor([1, 1, 0], dtype=torch.int32, device=device)
+dist.init_process_group(
+    backend="mooncake",
+    rank=rank,
+    world_size=2,
+    pg_options=pg.MooncakeBackendOptions(
+        active_ranks,
+        False,             # is_extension
+        3,                 # max_world_size
+    ),
+)
+
+backend = dist.group.WORLD
+join_ranks = [2]
+
+while not all(pg.get_peer_state(backend, join_ranks)):
+    # Continue serving, back off, or poll according to your scheduler policy.
+    pass
+
+pg.recover_ranks(backend, join_ranks)
+```
+
+### Joining-rank side
+
+```python
+from mooncake import pg
+
+active_ranks = torch.tensor([1, 1, 1], dtype=torch.int32, device=device)
+dist.init_process_group(
+    backend="mooncake",
+    rank=2,
+    world_size=3,
+    pg_options=pg.MooncakeBackendOptions(
+        active_ranks,
+        True,              # is_extension
+        3,                 # max_world_size
+    ),
+)
+
+backend = dist.group.WORLD
+pg.join_group(backend)
+```
+
+Important semantics:
+
+- `get_peer_state()` is collective among the current healthy ranks. Call it in a
+  consistent order across those ranks.
+- New ranks are inactive after `extend_group_size_to()` and become collective
+  participants only after `recover_ranks()`.
+- A joining rank initialized with `is_extension=True` starts with local-only
+  behavior and blocks in `join_group()` until the corresponding healthy ranks
+  publish recovery state.
+- Subgroups must be created in the same order on healthy and joining processes,
+  following PyTorch `new_group()` ordering rules.
+
+## Mooncake EP quick start
+
+Mooncake EP exposes `Buffer` from `mooncake.mooncake_ep_buffer`. Initialize it
+with a Mooncake process group and a workspace size computed from the expected
+dispatch shape.
+
+```python
+import torch
+import torch.distributed as dist
+from mooncake import pg
+from mooncake.mooncake_ep_buffer import Buffer
+
+
+# Assume dist.init_process_group(..., backend="mooncake", ...) has completed.
+group = dist.group.WORLD
+rank = dist.get_rank(group)
+world_size = dist.get_world_size(group)
+
+num_tokens = 128
+hidden = 7168
+num_experts = 288
+top_k = 8
+max_tokens_per_rank = 128
+
+x = torch.randn(num_tokens, hidden, dtype=torch.bfloat16, device="cuda")
+scores = torch.randn(num_tokens, num_experts, dtype=torch.float32, device="cuda")
+topk_idx = torch.topk(scores, top_k, dim=-1).indices
+topk_weights = torch.softmax(
+    torch.randn(num_tokens, top_k, dtype=torch.float32, device="cuda"), dim=-1
+)
+
+num_ep_buffer_bytes = Buffer.get_ep_buffer_size_hint(
+    max_tokens_per_rank,
+    hidden,
+    world_size,
+    num_experts,
+)
+buffer = Buffer(group, num_ep_buffer_bytes)
+
+# EP-level rank-health tensor. Kernels may update it to 0 when timeout_us
+# detects a failed source rank.
+active_ranks = torch.ones(world_size, dtype=torch.int32, device="cuda")
+
+recv_x, recv_count, handle, event, hook = buffer.dispatch(
+    x,
+    topk_idx,
+    active_ranks,
+    num_max_dispatch_tokens_per_rank=max_tokens_per_rank,
+    num_experts=num_experts,
+    timeout_us=-1,
+    use_fp8=True,
+    async_finish=False,
+    return_recv_hook=False,
+)
+event.current_stream_wait()
+
+# Run local experts on recv_x here. If use_fp8=True, recv_x is a
+# (data, scales) tuple; dequantize or feed it into an FP8-aware expert kernel.
+expert_out = run_local_experts(recv_x, recv_count)
+
+combined_x, event, hook = buffer.combine(
+    expert_out,
+    topk_idx,
+    topk_weights,
+    active_ranks,
+    timeout_us=-1,
+    handle=handle,
+)
+event.current_stream_wait()
+```
+
+## Mooncake EP API reference
+
+### `Buffer.get_ep_buffer_size_hint(...)`
+
+```python
+Buffer.get_ep_buffer_size_hint(
+    num_max_dispatch_tokens_per_rank: int,
+    hidden: int,
+    num_ranks: int,
+    num_experts: int,
+) -> int
+```
+
+Returns the workspace size in bytes for the EP buffer. Use the maximum number of
+tokens a rank may dispatch in one step. Underestimating this value can cause
+buffer overflow or incorrect dispatch results.
+
+### `Buffer(group, num_ep_buffer_bytes=0)`
+
+Creates the EP runtime for a Mooncake process group. The constructor exchanges
+RDMA and IPC metadata through the group, initializes fast-path transports when
+available, and falls back to the Python implementation if the fast path is not
+usable.
+
+### `Buffer.dispatch(...)`
+
+```python
+recv_x, recv_count, handle, event, hook = buffer.dispatch(
+    x,
+    topk_idx,
+    active_ranks,
+    num_max_dispatch_tokens_per_rank,
+    num_experts,
+    timeout_us,
+    use_fp8=True,
+    async_finish=False,
+    return_recv_hook=False,
+)
+```
+
+Arguments:
+
+- `x`: local token hidden states, shape `[num_tokens, hidden]`, typically BF16
+  on CUDA.
+- `topk_idx`: selected expert IDs, shape `[num_tokens, top_k]`. Use `-1` to mark
+  masked selections.
+- `active_ranks`: EP-level rank-health tensor, shape `[num_ranks]`, dtype
+  `torch.int32`. Timeout detection may set failed source ranks to `0`.
+- `num_max_dispatch_tokens_per_rank`: workspace capacity per rank. It should be
+  at least the maximum local `num_tokens` across ranks for the current step.
+- `num_experts`: global expert count. It must be divisible by `num_ranks`.
+- `timeout_us`: timeout in microseconds. Use `-1` to disable timeout detection.
+- `use_fp8`: when `True`, dispatch returns FP8 data plus scales.
+- `async_finish`: when `True`, returned tensors are associated with the returned
+  event for stream-lifetime management.
+- `return_recv_hook`: when `True`, call the returned `hook()` to complete receive
+  synchronization; otherwise use `event.current_stream_wait()`.
+
+Returns:
+
+- `recv_x`: packed local-expert inputs. If `use_fp8=True`, this is
+  `(packed_data, packed_scales)`; otherwise it is a BF16 tensor.
+- `recv_count`: number of tokens received by each local expert.
+- `handle`: opaque metadata required by `combine()` and
+  `get_next_combine_buffer()`.
+- `event`: `EventOverlap` helper; call `event.current_stream_wait()` before using
+  outputs when no hook is used.
+- `hook`: optional synchronization hook used when `return_recv_hook=True`.
+
+### `Buffer.combine(...)`
+
+```python
+combined_x, event, hook = buffer.combine(
+    x,
+    topk_idx,
+    topk_weights,
+    active_ranks,
+    timeout_us,
+    handle,
+    zero_copy=False,
+    async_finish=False,
+    return_recv_hook=False,
+    out=None,
+)
+```
+
+Arguments:
+
+- `x`: local expert outputs packed in the layout returned by `dispatch()`.
+- `topk_idx` and `topk_weights`: routing metadata for combining expert outputs
+  back to local tokens.
+- `active_ranks`: same EP-level rank-health tensor used by `dispatch()`.
+- `timeout_us`: timeout in microseconds; use `-1` to disable timeout detection.
+- `handle`: the handle returned by the matching `dispatch()` call.
+- `zero_copy`: when `True`, write expert outputs into
+  `buffer.get_next_combine_buffer(handle)` and pass that tensor to `combine()`.
+- `out`: optional output tensor for the combined result.
+
+### `Buffer.get_next_combine_buffer(handle)`
+
+Returns the next combine buffer for zero-copy expert output. Use it only with the
+matching dispatch `handle` and pass the resulting tensor back to `combine()` with
+`zero_copy=True`.
+
+### `Buffer.update_ep_member()`
+
+Reconnects EP peers after backend membership changes. Call it after PG recovery
+updates rank activeness so EP transport metadata and QPs can be refreshed.
+
+## Active-rank tensors: PG vs EP
+
+There are two active-rank tensors in the API surface:
+
+- **PG active-rank mask**: passed to `pg.MooncakeBackendOptions`. This is the
+  backend-level health mask used by collective and recovery logic.
+- **EP active-rank tensor**: passed to `Buffer.dispatch()` and `Buffer.combine()`.
+  It is also rank-level (`[num_ranks]`, `torch.int32`) and may be updated by EP
+  kernels when timeout detection marks a peer as failed.
+
+In simple integrations these tensors often carry the same health information,
+but they are passed through different API layers. Keep their dtype, device, and
+shape consistent with the process group world size or reserved `max_world_size`.
+
+## Tests and examples
+
+- PG collectives: `mooncake-pg/tests/test_pg_collectives.py`
+- PG elastic recovery and subgroup extension: `mooncake-pg/tests/test_pg_elastic.py`
+- PG benchmark harness: `mooncake-pg/benchmark/README.md`
+- EP correctness and failure simulation: `mooncake-ep/tests/test_ep_grid.py`
+- Wheel-level EP example: `mooncake-wheel/tests/test_mooncake_ep.py`
+
+See [PG/EP troubleshooting](../troubleshooting/pg-ep-troubleshooting.md) for
+common setup and runtime issues.
diff --git a/docs/source/python-api-reference/mooncake-store.md b/docs/source/python-api-reference/mooncake-store.md
index bb3610b1..c738f5ef 100644
--- a/docs/source/python-api-reference/mooncake-store.md
+++ b/docs/source/python-api-reference/mooncake-store.md
@@ -12,9 +12,13 @@ pip install mooncake-transfer-engine
 📦 **Package Details**: [https://pypi.org/project/mooncake-transfer-engine/](https://pypi.org/project/mooncake-transfer-engine/)
 
 ### Required Service
-Only one service is required now:
+The only always-required service is:
 
-- `mooncake_master` — Master service which now embeds the HTTP metadata server
+- `mooncake_master` — Master service for cluster membership and object placement
+
+For Transfer Engine metadata, use the `P2PHANDSHAKE` connection string for
+decentralized peer discovery, enable the master's embedded HTTP metadata
+server, or provide an external metadata service.
 
 ## Quick Start
 
@@ -60,13 +64,17 @@ print(data.decode())  # Output: Hello, Mooncake Store!
 store.close()
 ```
 
-**RDMA device selection**: Leave `rdma_devices` as `""` to auto-select RDMA NICs. Provide a comma-separated list (e.g. `"mlx5_0,mlx5_1"`) to pin to specific hardware.
+**RDMA device selection**: For `protocol="rdma"` or `protocol="efa"`, leave
+`rdma_devices` as `""` to auto-discover NICs. Set `MC_MS_AUTO_DISC=0` when you
+want auto-discovery disabled, then provide a comma-separated list such as
+`"mlx5_0,mlx5_1"` to pin specific hardware.
 
 Mooncake selects available ports internally at `setup() `, so you do not need to fix specific port numbers in these examples. Internally, ports are chosen from a dynamic range (currently 12300–14300).
 
-#### P2P Hello World (preview)
+#### P2P Hello World
 
-The following setup uses the new P2P handshake and does not require an HTTP metadata server. This feature is not released yet; use only if you’re testing the latest code.
+The following setup uses P2P handshake and does not require an HTTP metadata
+server. Pass the literal `P2PHANDSHAKE` value as the metadata server.
 
 ```python
 from mooncake.store import MooncakeDistributedStore
@@ -130,6 +138,134 @@ store.close()
 
 </details>
 
+## Structured Object Store Helper
+
+`mooncake.structured_object_store` provides a higher-level helper for one logical object that contains multiple named members.
+It is designed for cases such as rollout / batch transfer where callers want to keep their own object semantics locally while using Mooncake for fast payload movement.
+
+The helper separates two concepts:
+
+- **structured object path**: named members with metadata-aware materialization;
+- **generic bundle path**: manifest + named payloads when the caller only needs raw grouped objects.
+
+### Main types
+
+```python
+from mooncake.structured_object_store import (
+    MooncakeBundleTransfer,
+    StructuredMemberSlice,
+    StructuredObjectPayload,
+)
+```
+
+- `MooncakeBundleTransfer`: public helper facade built on a `MooncakeDistributedStore`.
+- `StructuredObjectPayload`: structured object to write. Members are passed in `buffers`, and optional object metadata is passed in `metadata`.
+- `StructuredMemberSlice`: slice selection for one structured member during reads.
+
+### Structured object write and full read
+
+Use `put_structured_object()` to write one structured object. The default read path is `read_spec(ref)`, and full-object materialization is just the default case of the partial-read API.
+
+```python
+import numpy as np
+from mooncake.store import MooncakeDistributedStore
+from mooncake.structured_object_store import MooncakeBundleTransfer, StructuredObjectPayload
+
+store = MooncakeDistributedStore()
+transfer = MooncakeBundleTransfer(store, key_prefix="demo/structured")
+
+payload = StructuredObjectPayload(
+    metadata={"step": 7, "layout": "rollout"},
+    buffers={
+        "tokens": np.array(range(24), dtype=np.int32).reshape(6, 4),
+        "mask": np.ones((6, 4), dtype=np.int8),
+        "prompt_ids": b"sample-ids",
+    },
+)
+
+ref = transfer.put_structured_object(payload)
+result = transfer.materialize(transfer.read_spec(ref))
+
+tokens = result.objects["tokens"]
+prompt_ids = result.objects["prompt_ids"]
+metadata = result.metadata
+```
+
+### Structured object transfer policy
+
+By default, structured object writes use `BundleTransferPolicy(copy_mode="auto")`. This keeps the existing behavior: Mooncake uses the zero-copy `batch_put_from` fast path when buffer registration is available, and falls back to ordinary `store.put` when the fast path is unavailable.
+
+For very large tensors, buffer registration capacity can be exhausted. If the upper layer does not handle this failure mode yet, force the non-zero-copy path with `copy_mode="copy"`:
+
+```python
+from mooncake.structured_object_store import BundleTransferPolicy
+
+ref = transfer.put_structured_object(
+    payload,
+    policy=BundleTransferPolicy(copy_mode="copy"),
+)
+```
+
+Available modes:
+
+- `auto`: prefer zero-copy and fall back to regular `store.put` when zero-copy support is unavailable;
+- `copy`: force the non-zero-copy `store.put` path;
+- `zero_copy`: require the zero-copy `batch_put_from` path and raise an error if it is unavailable.
+
+If a caller has already registered a structured member buffer, pass `pre_registered_buffers={"member_name": True}` to avoid duplicate registration. The buffer must be the same writable, contiguous buffer used for transfer; read-only, non-contiguous, or internally copied buffers are rejected.
+
+### Partial reads
+
+Read narrowing happens on top of `read_spec(ref)`:
+
+- `select_members([...])` keeps only selected members;
+- `slice_member(name, axis=0, start=..., end=...)` slices one ndarray member;
+- `materialize(spec)` returns newly materialized objects.
+
+```python
+spec = (
+    transfer.read_spec(ref)
+    .select_members(["tokens"])
+    .slice_member("tokens", axis=0, start=2, end=5)
+)
+result = transfer.materialize(spec)
+
+selected_tokens = result.objects["tokens"]
+```
+
+Current scope:
+
+- byte members support full-member reads;
+- ndarray members support full reads and sliced reads;
+- full read is the default `read_spec(ref)` case.
+
+### Reusing caller-owned destinations
+
+Use `materialize_into()` when the caller already owns the destination ndarray buffers and wants Mooncake to fill them directly.
+
+```python
+destination = np.empty((3, 4), dtype=np.int32)
+spec = (
+    transfer.read_spec(ref)
+    .select_members(["tokens"])
+    .slice_member("tokens", axis=0, start=2, end=5)
+)
+result = transfer.materialize_into(spec, {"tokens": destination})
+
+assert result.objects["tokens"] is destination
+```
+
+`materialize_into()` is only for members whose destination layout is already known to the caller. For byte members or default object reconstruction, use `materialize()`.
+
+### Generic bundle fallback
+
+If the caller does not need structured member semantics, the same helper also supports raw named bundles:
+
+- `put_bundle(...)`
+- `remove_bundle(...)`
+
+Use the bundle path when the object is just a manifest plus named payloads, and use the structured object path when callers want member selection, slicing, and ndarray-aware materialization.
+
 ## Zero-Copy API (Advanced Performance)
 
 For maximum performance, especially with RDMA networks, use the zero-copy API. This allows direct memory access without intermediate copies.
@@ -138,7 +274,7 @@ For maximum performance, especially with RDMA networks, use the zero-copy API. T
 
 ⚠️ **Important**: `register_buffer` is required for zero-copy RDMA operations. Without proper buffer registration, undefined behavior and memory corruption may occur.
 
-Zero-copy operations require registering memory buffers with the store:
+Zero-copy operations require registered memory buffers. For repeated reads and writes, prefer the Python `BufferPool` helper described below so leases come from the store's setup-time local buffer instead of registering and unregistering memory for every operation.
 
 #### register_buffer()
 Register a memory buffer for direct RDMA access.
@@ -172,6 +308,45 @@ store.unregister_buffer(buffer_ptr)
 
 </details>
 
+#### BufferPool Helper
+
+`BufferPool` leases scratch buffers from the store's setup-time local buffer for repeated zero-copy operations. It is useful when a caller repeatedly needs temporary memory, for example as the destination buffer for `get_into()` or `get_into_ranges()`.
+
+```python
+from mooncake.buffer_pool import BufferPool
+
+pool = BufferPool(store)
+
+with pool.buffer(1024 * 1024) as lease:
+    n = store.get_into("my_key", lease.ptr, lease.size)
+    view = lease.buffer[:n]
+    # Consume view directly, or wrap it with np.frombuffer(view, dtype=...).
+    # Copy only if the data must outlive the lease: data = bytes(view)
+
+pool.close()
+```
+
+`acquire(size)` and `buffer(size)` return a lease object. A lease exposes:
+
+- `ptr`: the local-buffer address to pass to zero-copy APIs.
+- `size`: the requested logical size.
+- `buffer`: a Python `memoryview` over the logical requested size.
+- `release()`: returns the local-buffer allocation to the store allocator.
+
+Behavior and lifecycle rules:
+
+- Leases prefer the store local buffer shared with internal Store staging paths.
+- If local-buffer allocation is temporarily exhausted, `BufferPool` can allocate and register a short-lived overflow buffer; the overflow buffer is unregistered when the lease is released.
+- `max_regions` can limit the number of concurrently active external leases.
+- `max_bytes` bounds total active local-buffer and overflow leases; the default allows one local-buffer-sized overflow burst.
+- `acquire(size, block=False)` raises when both local and overflow capacity are exhausted instead of waiting.
+- `acquire(size, timeout=...)` can wait for another lease to be released.
+- Do not keep `lease.ptr` or a `memoryview` after releasing the lease.
+- `release()` fails while exported views are alive. Delete those views first, then release.
+- `close()` fails if leases are still active. Release all leases before closing the pool.
+
+Legacy code may still import `RegisteredBufferPool`, but new examples should prefer `mooncake.buffer_pool.BufferPool`.
+
 ---
 
 ### Zero-Copy Operations
@@ -288,7 +463,7 @@ from object `all_keys[i][j]`, then writes it into destination buffer
 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.
+- `buffer_ptrs`: Memory addresses of pre-allocated destination buffers. Each buffer must resolve to Store-managed registered memory, either from `BufferPool`/the setup-time local buffer or from an explicit `register_buffer()` call.
 - `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.
@@ -481,8 +656,261 @@ config.preferred_segment = self.get_hostname()
 
 ```python
 config = ReplicateConfig()
-config.prefer_alloc_in_same_node = "True
+config.prefer_alloc_in_same_node = "True"
 ```
+
+#### group_ids
+**Type:** `List[str] | None`
+**Default:** `None`
+**Description:** Optionally assigns object metadata to routing groups during writes. When this field is unset, Mooncake Store preserves the default ungrouped behavior. When it is set, each group ID maps to the object at the same position in the write request. Empty string (`""`) explicitly stores that object as ungrouped.
+
+For batch write APIs, the number of group IDs must match the number of keys:
+
+```python
+config = ReplicateConfig()
+config.group_ids = ["session-a", "", "session-b"]
+
+store.put_batch(
+    ["key-a", "key-b", "key-c"],
+    [b"value-a", b"value-b", b"value-c"],
+    config,
+)
+```
+
+For a single-object write, provide one group ID:
+
+```python
+config = ReplicateConfig()
+config.group_ids = ["session-a"]
+
+store.put("key-a", b"value-a", config)
+```
+
+---
+
+## Unified Parallel Tensor IO API
+
+Mooncake Store also provides a unified tensor IO family for tensors that are stored either as full objects or as explicitly identified parallel shards.
+
+This API family is the long-term interface for TP / DP / EP / PP-aware tensor IO:
+
+- write and upsert use `TensorParallelism`
+- reads use `ReadTarget`
+- legacy TP-only APIs remain available as compatibility wrappers
+
+### ParallelAxis
+
+`ParallelAxis` describes one axis in a shard identity.
+
+```python
+axis = mooncake.store.ParallelAxis()
+axis.kind = "tp"      # one of: "tp", "dp", "ep", "pp"
+axis.rank = 0
+axis.size = 8
+axis.split_dim = 1     # used for layout-sharding axes such as TP
+axis.expert_id = 3     # optional, for EP
+axis.stage_id = 1      # optional, for PP
+```
+
+**Fields:**
+- `kind`: Parallelism axis kind.
+- `rank`: Current shard rank on that axis.
+- `size`: Total number of shards on that axis.
+- `split_dim`: Optional tensor split dimension for layout-sharding axes.
+- `expert_id`: Optional expert identifier for EP layouts.
+- `stage_id`: Optional pipeline stage identifier for PP layouts.
+
+### TensorParallelism
+
+`TensorParallelism` is an ordered list of axes that identifies the stored or requested shard.
+
+```python
+parallelism = mooncake.store.TensorParallelism()
+parallelism.axes = [
+    tp_axis,
+]
+```
+
+Examples:
+- TP shard: `axes=[TP(...)]`
+- DP + TP shard: `axes=[DP(...), TP(...)]`
+- PP + TP shard: `axes=[PP(...), TP(...)]`
+- EP shard: `axes=[EP(...)]`
+
+### ReadTarget
+
+`ReadTarget` tells Mooncake whether the caller wants the stored form, a specific shard view, or the reconstructed full tensor.
+
+```python
+target = mooncake.store.ReadTarget()
+target.mode = "full"          # one of: "as_stored", "shard", "full"
+target.parallelism = None      # required for target shard reads
+```
+
+**Fields:**
+- `mode`: Read materialization mode.
+- `parallelism`: Optional `TensorParallelism`. Required when `mode="shard"`.
+
+### put_tensor_with_parallelism()
+
+Store a tensor using the unified parallelism model.
+
+```python
+def put_tensor_with_parallelism(
+    self,
+    key: str,
+    tensor,
+    parallelism: mooncake.store.TensorParallelism | None = None,
+    config: ReplicateConfig | None = None,
+    writer_partition = None,
+) -> int
+```
+
+Use `parallelism=None` to store a full tensor object. Provide `TensorParallelism` to store a shard-scoped object.
+
+`writer_partition` is an optional write-side shorthand for full-tensor inputs that should be stored as one shard. It describes the writer's `(rank, size, split_dim)` and is mutually exclusive with `parallelism`; do not provide both in one call.
+
+For TP-containing multi-axis layouts, the caller may pass the full source tensor; Mooncake derives and persists the uniform shard selected by the requested TP rank/layout. That applies to layouts such as `dp_tp`, `pp_tp`, and `ep_tp`.
+
+Plain single-axis TP remains shard-input for compatibility.
+
+Pure DP still does not imply a split axis by itself.
+
+### batch_put_tensor_with_parallelism()
+
+Batch version of unified tensor writes.
+
+```python
+def batch_put_tensor_with_parallelism(
+    self,
+    keys: list[str],
+    tensors: list,
+    parallelisms: list[mooncake.store.TensorParallelism | None] | None = None,
+    config: ReplicateConfig | None = None,
+    writer_partitions = None,
+) -> list[int]
+```
+
+`writer_partitions` is an optional write-side convenience input for batch full-tensor writes that should be partitioned into stored shards. Each entry describes the target shard write as `(rank, size, split_dim)`.
+
+Use `writer_partitions` when the caller has full tensors and wants Mooncake to derive the stored shard objects from writer-side partition info instead of constructing full `TensorParallelism` objects per element. TP-containing `parallelisms` can now express the same full-tensor-input behavior too; `writer_partitions` remains the lighter explicit write-side shorthand.
+
+### get_tensor_with_parallelism()
+
+Read a tensor through the unified read path.
+
+```python
+def get_tensor_with_parallelism(
+    self,
+    key: str,
+    target: mooncake.store.ReadTarget | None = None,
+)
+```
+
+Typical modes:
+- `target=None` or `mode="as_stored"`: return the stored local object.
+- `mode="shard"`: return the target shard described by `target.parallelism`.
+- `mode="full"`: reconstruct and return the full tensor.
+
+### batch_get_tensor_with_parallelism()
+
+Batch version of unified tensor reads.
+
+```python
+def batch_get_tensor_with_parallelism(
+    self,
+    keys: list[str],
+    targets: list[mooncake.store.ReadTarget | None] | None = None,
+) -> list
+```
+
+### get_tensor_with_parallelism_into() / batch_get_tensor_with_parallelism_into()
+
+Zero-copy unified read forms. The destination buffers must resolve to Store-managed registered memory, either from `BufferPool`/the setup-time local buffer or from an explicit `register_buffer()` call.
+
+```python
+def get_tensor_with_parallelism_into(
+    self,
+    key: str,
+    buffer_ptr: int,
+    size: int,
+    target: mooncake.store.ReadTarget | None = None,
+)
+```
+
+```python
+def batch_get_tensor_with_parallelism_into(
+    self,
+    keys: list[str],
+    buffer_ptrs: list[int],
+    sizes: list[int],
+    targets: list[mooncake.store.ReadTarget | None] | None = None,
+) -> list
+```
+
+### upsert_tensor_with_parallelism()
+
+Unified upsert form for tensor objects.
+
+```python
+def upsert_tensor_with_parallelism(
+    self,
+    key: str,
+    tensor,
+    parallelism: mooncake.store.TensorParallelism | None = None,
+    config: ReplicateConfig | None = None,
+    writer_partition = None,
+) -> int
+```
+
+The write semantics match `put_tensor_with_parallelism()`, including full-tensor input for TP-containing layouts and the mutually exclusive `writer_partition` shorthand.
+
+### batch_upsert_tensor_with_parallelism()
+
+Batch unified upsert form.
+
+```python
+def batch_upsert_tensor_with_parallelism(
+    self,
+    keys: list[str],
+    tensors: list,
+    parallelisms: list[mooncake.store.TensorParallelism | None] | None = None,
+    config: ReplicateConfig | None = None,
+    writer_partitions = None,
+) -> list[int]
+```
+
+The write semantics match `put_tensor_with_parallelism()`, including full-tensor input for TP-containing layouts.
+
+### *_from zero-copy write variants
+
+The unified write and upsert family also has `_from` variants for registered-memory inputs, including:
+
+- `put_tensor_with_parallelism_from(...)`
+- `batch_put_tensor_with_parallelism_from(...)`
+- `upsert_tensor_with_parallelism_from(...)`
+- `batch_upsert_tensor_with_parallelism_from(...)`
+
+These APIs accept Store-managed registered buffer pointers that contain serialized tensor objects in the current Mooncake tensor format:
+
+```text
+[TensorObjectHeader + layout metadata][tensor data]
+```
+
+As with other zero-copy APIs, every source pointer must resolve to Store-managed registered memory, either from `BufferPool`/the setup-time local buffer or from an explicit `register_buffer()` call.
+
+### Compatibility wrappers
+
+Legacy TP-only methods such as:
+
+- `put_tensor_with_tp(...)`
+- `batch_put_tensor_with_tp(...)`
+- `get_tensor_with_tp(...)`
+- `batch_get_tensor_with_tp(...)`
+- corresponding `_into`, `_from`, and upsert variants
+
+remain supported for compatibility, but they are wrapper-style APIs around the unified parallel tensor IO model. Prefer the unified `*_with_parallelism` family for new code and new documentation examples.
+
 ---
 
 ## Non-Zero-Copy API (Simple Usage)
@@ -540,11 +968,14 @@ print("Retrieved all keys successfully:", retrieved == values)
 
 ## Topology & Devices
 
-- Auto-discovery: Disabled by default. For `protocol="rdma"`, you must specify RDMA devices.
-- Enable auto-discovery (optional):
-  - `MC_MS_AUTO_DISC=1` enables auto-discovery; then `rdma_devices` is not required.
-  - Optionally restrict candidates with `MC_MS_FILTERS`, a comma-separated whitelist of NIC names, e.g. `MC_MS_FILTERS=mlx5_0,mlx5_2`.
-  - If `MC_MS_AUTO_DISC` is not set or set to `0`, auto-discovery remains disabled and `rdma_devices` is required for RDMA.
+- Auto-discovery: Enabled by default for `protocol="rdma"` or `protocol="efa"`
+  when `rdma_devices` is empty.
+- Discovery controls:
+  - `MC_MS_AUTO_DISC=1` forces auto-discovery; then `rdma_devices` is ignored.
+  - `MC_MS_AUTO_DISC=0` disables auto-discovery; then `rdma_devices` is required
+    for RDMA/EFA.
+  - `MC_MS_FILTERS` restricts auto-discovery to a comma-separated whitelist of
+    NIC names, e.g. `MC_MS_FILTERS=mlx5_0,mlx5_2`.
 
 Examples:
 
@@ -630,17 +1061,25 @@ def setup(
     protocol: str = "tcp",
     rdma_devices: str = "",
     master_server_addr: str,
+    engine: Optional[TransferEngine] = None,
+    enable_ssd_offload: bool = False,
+    ssd_offload_path: str = "",
+    tenant_id: str = "default",
 ) -> int
 ```
 
 **Parameters:**
 - `local_hostname` (str): **Required**. Local hostname and port (e.g., "localhost" or "localhost:12345")
-- `metadata_server` (str): **Required**. Metadata server address (e.g., "http://localhost:8080/metadata")
-- `global_segment_size` (int): Memory segment size in bytes for mounting (default: 16MB = 16777216)
-- `local_buffer_size` (int): Local buffer size in bytes (default: 1GB = 1073741824)
-- `protocol` (str): Network protocol - "tcp" or "rdma" (default: "tcp")
-- `rdma_devices` (str): RDMA device name(s), e.g. `"mlx5_0"` or `"mlx5_0,mlx5_1"`. Leave empty to auto-select NICs. Provide device names to pin the NICs. Always empty for TCP.
+- `metadata_server` (str): **Required**. Metadata connection string, e.g. `"P2PHANDSHAKE"` or `"http://localhost:8080/metadata"`.
+- `global_segment_size` (int): Memory segment size in bytes for mounting.
+- `local_buffer_size` (int): Local buffer size in bytes.
+- `protocol` (str): Network protocol, usually `"tcp"`, `"rdma"`, `"efa"`, `"cxl"`, or `"ascend"` depending on the build.
+- `rdma_devices` (str): RDMA/EFA device name(s), e.g. `"mlx5_0"` or `"mlx5_0,mlx5_1"`. Leave empty to auto-discover NICs unless `MC_MS_AUTO_DISC=0`; always empty for TCP.
 - `master_server_addr` (str): **Required**. Master server address (e.g., "localhost:50051")
+- `engine` (Optional[TransferEngine]): Existing Transfer Engine instance to reuse. Defaults to `None`.
+- `enable_ssd_offload` (bool): Enable client-side SSD offload support. Defaults to `False`.
+- `ssd_offload_path` (str): SSD offload directory. When provided, overrides the storage path environment configuration.
+- `tenant_id` (str): Tenant namespace for object keys. Defaults to `"default"`.
 
 **Returns:**
 - `int`: Status code (0 = success, non-zero = error code)
@@ -1204,6 +1643,163 @@ def init_all(self, protocol: str, device_name: str, mount_segment_size: int = 16
 
 ---
 
+#### mount_segment()
+Mount a local file or shared-memory region as one or more Mooncake store
+segments.
+
+```python
+def mount_segment(
+    self,
+    path: str,
+    size: int,
+    offset: int = 0,
+    protocol: str = "tcp",
+    location: str = "",
+) -> dict
+```
+
+**Parameters:**
+- `path` (str): File or shared-memory path to map and mount.
+- `size` (int): Number of bytes to mount.
+- `offset` (int, optional): File offset in bytes. Defaults to `0`.
+- `protocol` (str, optional): Transfer protocol. Defaults to `"tcp"`.
+- `location` (str, optional): Device or locality hint. Defaults to an empty
+  string.
+
+**Returns:**
+- `dict`: A result dictionary with:
+  - `ret` (int): Status code (0 = success, non-zero = error code)
+  - `segment_ids` (List[str]): Segment ids created by the mount operation
+
+**Example:**
+```python
+result = store.mount_segment(
+    "/dev/shm/mooncake_segment",
+    16 * 1024 * 1024,
+    offset=0,
+    protocol="tcp",
+    location="",
+)
+
+if result["ret"] == 0:
+    segment_ids = list(result["segment_ids"])
+    print("Mounted segments:", segment_ids)
+else:
+    print("Mount failed:", result["ret"])
+```
+
+The corresponding HTTP endpoints are `/api/mount_shm` and
+`/api/unmount_shm`, but the HTTP API is intentionally narrower: it accepts a
+named shared memory object name instead of an arbitrary path.
+
+---
+
+#### unmount_segment()
+Unmount one or more file or shared-memory segments by segment id.
+
+```python
+def unmount_segment(
+    self,
+    segment_ids: List[str],
+    grace_period_seconds: int = 0,
+) -> int
+```
+
+**Parameters:**
+- `segment_ids` (List[str]): Segment ids returned by `mount_segment()`.
+- `grace_period_seconds` (int, optional): Grace period before the segment is
+  fully unmounted. Defaults to `0`, which keeps the existing immediate unmount
+  behavior. During a positive grace period, the segment remains readable but no
+  longer accepts new allocations.
+
+**Returns:**
+- `int`: Status code (0 = success, non-zero = error code)
+
+**Example:**
+```python
+ret = store.unmount_segment(segment_ids, grace_period_seconds=30)
+if ret != 0:
+    print("Unmount failed:", ret)
+```
+
+---
+
+#### allocate_and_mount_segment()
+Allocate memory inside the store process and mount it as one or more Mooncake
+store segments.
+
+```python
+def allocate_and_mount_segment(
+    self,
+    size: int,
+    protocol: str = "tcp",
+    location: str = "",
+) -> dict
+```
+
+**Parameters:**
+- `size` (int): Number of bytes requested. The allocated size may be rounded up
+  for alignment.
+- `protocol` (str, optional): Transfer protocol. Defaults to `"tcp"`.
+- `location` (str, optional): Device or locality hint. Defaults to an empty
+  string.
+
+**Returns:**
+- `dict`: A result dictionary with:
+  - `ret` (int): Status code (0 = success, non-zero = error code)
+  - `segment_ids` (List[str]): Segment ids created by the mount operation
+  - `allocated_size` (int): Actual allocated size in bytes
+
+**Example:**
+```python
+result = store.allocate_and_mount_segment(
+    16 * 1024 * 1024,
+    protocol="tcp",
+    location="",
+)
+
+if result["ret"] == 0:
+    segment_ids = list(result["segment_ids"])
+    allocated_size = result["allocated_size"]
+```
+
+The corresponding HTTP endpoints are `/api/mount` and `/api/unmount`.
+
+---
+
+#### unmount_and_free_segment()
+Unmount one or more internally allocated segments by segment id and free their
+local memory.
+
+```python
+def unmount_and_free_segment(
+    self,
+    segment_ids: List[str],
+    grace_period_seconds: int = 0,
+) -> int
+```
+
+**Parameters:**
+- `segment_ids` (List[str]): Segment ids returned by
+  `allocate_and_mount_segment()`.
+- `grace_period_seconds` (int, optional): Grace period before the segment is
+  fully unmounted and its local allocated memory is released. Defaults to `0`,
+  which keeps the existing immediate unmount-and-free behavior. During a
+  positive grace period, the segment remains readable but no longer accepts new
+  allocations.
+
+**Returns:**
+- `int`: Status code (0 = success, non-zero = error code)
+
+**Example:**
+```python
+ret = store.unmount_and_free_segment(segment_ids, grace_period_seconds=30)
+if ret != 0:
+    print("Unmount and free failed:", ret)
+```
+
+---
+
 #### get_hostname()
 Get the hostname of the current store instance.
 
@@ -1824,7 +2420,7 @@ def upsert_tensor(self, key: str, tensor: torch.Tensor) -> int
 #### 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
+`[TensorObjectHeader+layout metadata][tensor data]`, matching the layout used by
 `get_tensor_into()`.
 
 ```python
@@ -1844,7 +2440,7 @@ def upsert_tensor_from(self, key: str, buffer_ptr: int, size: int) -> int
 #### batch_upsert_tensor_from()
 
 Upsert multiple tensors directly from pre-allocated buffers. Each buffer must
-use layout `[TensorMetadata][tensor data]`.
+use layout `[TensorObjectHeader+layout metadata][tensor data]`.
 
 ```python
 def batch_upsert_tensor_from(self, keys: List[str], buffer_ptrs: List[int], sizes: List[int]) -> List[int]
@@ -1927,6 +2523,7 @@ def batch_upsert_pub_tensor(self, keys: List[str], tensors_list: List[torch.Tens
 
 **Note:** This function requires `torch` to be installed and available in the environment. Not supported for dummy client.
 
+
 ---
 
 ### PyTorch Tensor Operations (Zero Copy)
@@ -2127,7 +2724,7 @@ def batch_get_into(self, keys: List[str], buffer_ptrs: List[int], sizes: List[in
 **Returns:**
 - `List[int]`: List of bytes read for each operation (positive = success, negative = error)
 
-⚠️ **Buffer Registration Required**: All buffers must be registered before batch zero-copy operations.
+⚠️ **Store-managed Buffer Required**: All buffers must resolve to Store-managed registered memory before batch zero-copy operations.
 
 **Example:**
 
@@ -2203,7 +2800,7 @@ List[int]
 **Returns:**
 - `List[int]`: List of bytes read for each operation (positive = success, negative = error)
 
-⚠️ **Buffer Registration Required**: All buffers must be registered before batch zero-copy operations.
+⚠️ **Store-managed Buffer Required**: All buffers must resolve to Store-managed registered memory before batch zero-copy operations.
 
 **Example:**
 
@@ -2332,8 +2929,6 @@ bind_to_numa_node(0)
 
 ---
 
----
-
 ## Error Handling
 
 Most methods return integer status codes:
@@ -2356,5 +2951,3 @@ For methods that return data (`get`, `get_batch`, `get_buffer`, `get_tensor`):
 4. **Configure replication** appropriately - more replicas provide better availability but use more storage
 5. **Use soft pinning** for frequently accessed objects to keep them in memory
 6. **Choose RDMA protocol** when available for maximum performance
-
----
diff --git a/docs/source/python-api-reference/transfer-engine.md b/docs/source/python-api-reference/transfer-engine.md
index d23b70f5..59ab1793 100644
--- a/docs/source/python-api-reference/transfer-engine.md
+++ b/docs/source/python-api-reference/transfer-engine.md
@@ -164,10 +164,14 @@ Gets the address of the first buffer in a specified segment.
 
 ### Data Transfer Operations
 
+```{note}
+The optional `transport_hint` argument pins the request onto a named transport (`"rdma"`, `"tcp"`, ...), overriding policy-driven selection for that one call. **TENT backend only** (`MC_USE_TENT=1`); silently ignored on the classic backend. See [TENT transport selector](../design/tent/transport-selector.md) for more details.
+```
+
 #### transfer_sync_write()
 
 ```python
-transfer_sync_write(target_hostname, buffer, peer_buffer_address, length)
+transfer_sync_write(target_hostname, buffer, peer_buffer_address, length, transport_hint="")
 ```
 
 Performs a synchronous write operation to transfer data from local buffer to remote buffer.
@@ -177,6 +181,7 @@ Performs a synchronous write operation to transfer data from local buffer to rem
 - `buffer` (int): The local buffer address
 - `peer_buffer_address` (int): The remote buffer address
 - `length` (int): The number of bytes to transfer
+- `transport_hint` (str, optional): TENT-only per-request transport pin. See the note above. Default `""` (policy-driven).
 
 **Returns:**
 - `int`: 0 on success, negative value on failure
@@ -184,7 +189,7 @@ Performs a synchronous write operation to transfer data from local buffer to rem
 #### transfer_sync_read()
 
 ```python
-transfer_sync_read(target_hostname, buffer, peer_buffer_address, length)
+transfer_sync_read(target_hostname, buffer, peer_buffer_address, length, transport_hint="")
 ```
 
 Performs a synchronous read operation to transfer data from remote buffer to local buffer.
@@ -194,6 +199,7 @@ Performs a synchronous read operation to transfer data from remote buffer to loc
 - `buffer` (int): The local buffer address
 - `peer_buffer_address` (int): The remote buffer address
 - `length` (int): The number of bytes to transfer
+- `transport_hint` (str, optional): TENT-only per-request transport pin. See the note above. Default `""` (policy-driven).
 
 **Returns:**
 - `int`: 0 on success, negative value on failure
@@ -201,7 +207,7 @@ Performs a synchronous read operation to transfer data from remote buffer to loc
 #### transfer_sync()
 
 ```python
-transfer_sync(target_hostname, buffer, peer_buffer_address, length, opcode, notify=None)
+transfer_sync(target_hostname, buffer, peer_buffer_address, length, opcode, notify=None, transport_hint="")
 ```
 
 Performs a synchronous transfer operation with specified opcode and optional notification.
@@ -213,6 +219,7 @@ Performs a synchronous transfer operation with specified opcode and optional not
 - `length` (int): The number of bytes to transfer
 - `opcode` (TransferOpcode): The transfer operation type (READ or WRITE)
 - `notify` (TransferNotify, optional): Notification object to send after transfer completion
+- `transport_hint` (str, optional): TENT-only per-request transport pin. See the note above. Default `""` (policy-driven).
 
 **Returns:**
 - `int`: 0 on success, negative value on failure
@@ -220,7 +227,7 @@ Performs a synchronous transfer operation with specified opcode and optional not
 #### transfer_submit_write()
 
 ```python
-transfer_submit_write(target_hostname, buffer, peer_buffer_address, length)
+transfer_submit_write(target_hostname, buffer, peer_buffer_address, length, transport_hint="")
 ```
 
 Submits an asynchronous write operation and returns immediately.
@@ -230,9 +237,10 @@ Submits an asynchronous write operation and returns immediately.
 - `buffer` (int): The local buffer address
 - `peer_buffer_address` (int): The remote buffer address
 - `length` (int): The number of bytes to transfer
+- `transport_hint` (str, optional): TENT-only per-request transport pin. See the note above. Default `""` (policy-driven).
 
 **Returns:**
-- `int`: Batch ID for tracking the operation, or negative value on failure
+- `int`: Batch ID for tracking the operation, or 0 on failure
 
 #### transfer_check_status()
 
@@ -255,7 +263,7 @@ Checks the status of an asynchronous transfer operation.
 #### transfer_write_on_cuda()
 
 ```python
-transfer_write_on_cuda(target_hostname, buffer, peer_buffer_address, length, stream_ptr)
+transfer_write_on_cuda(target_hostname, buffer, peer_buffer_address, length, stream_ptr=0, transport_hint="")
 ```
 
 Performs a write operation to transfer data from local buffer to remote buffer on a given cuda stream.
@@ -266,6 +274,7 @@ Performs a write operation to transfer data from local buffer to remote buffer o
 - `peer_buffer_address` (int): The remote buffer address
 - `length` (int): The number of bytes to transfer
 - `stream_ptr` (int): The integer representation of a CUDA stream pointer (`cudaStream_t`). For example, from a PyTorch stream, this can be obtained via `stream.cuda_stream`.
+- `transport_hint` (str, optional): TENT-only per-request transport pin. See the note at the top of this section. Default `""` (policy-driven).
 
 **Returns:**
 - `None`: The function returns immediately after successfully scheduling the transfer callback.
@@ -279,7 +288,7 @@ Performs a write operation to transfer data from local buffer to remote buffer o
 #### transfer_read_on_cuda()
 
 ```python
-transfer_read_on_cuda(target_hostname, buffer, peer_buffer_address, length, stream_ptr)
+transfer_read_on_cuda(target_hostname, buffer, peer_buffer_address, length, stream_ptr=0, transport_hint="")
 ```
 
 Performs a read operation to transfer data from remote buffer to local buffer on a given cuda stream.
@@ -290,6 +299,7 @@ Performs a read operation to transfer data from remote buffer to local buffer on
 - `peer_buffer_address` (int): The remote buffer address
 - `length` (int): The number of bytes to transfer
 - `stream_ptr` (int): The integer representation of a CUDA stream pointer (`cudaStream_t`). For example, from a PyTorch stream, this can be obtained via `stream.cuda_stream`.
+- `transport_hint` (str, optional): TENT-only per-request transport pin. See the note at the top of this section. Default `""` (policy-driven).
 
 **Returns:**
 - `None`: The function returns immediately after successfully scheduling the transfer callback.
@@ -307,7 +317,7 @@ Performs a read operation to transfer data from remote buffer to local buffer on
 #### batch_transfer_sync_write()
 
 ```python
-batch_transfer_sync_write(target_hostname, buffers, peer_buffer_addresses, lengths)
+batch_transfer_sync_write(target_hostname, buffers, peer_buffer_addresses, lengths, transport_hint="")
 ```
 
 Performs a batch synchronous write operation to transfer multiple data chunks from local buffers to remote buffers.
@@ -317,6 +327,7 @@ Performs a batch synchronous write operation to transfer multiple data chunks fr
 - `buffers` (List[int]): List of local buffer addresses
 - `peer_buffer_addresses` (List[int]): List of remote buffer addresses
 - `lengths` (List[int]): List of byte lengths for each transfer
+- `transport_hint` (str, optional): TENT-only per-batch transport pin. Default `""` (policy-driven).
 
 **Returns:**
 - `int`: 0 on success, negative value on failure
@@ -324,7 +335,7 @@ Performs a batch synchronous write operation to transfer multiple data chunks fr
 #### batch_transfer_sync_read()
 
 ```python
-batch_transfer_sync_read(target_hostname, buffers, peer_buffer_addresses, lengths)
+batch_transfer_sync_read(target_hostname, buffers, peer_buffer_addresses, lengths, transport_hint="")
 ```
 
 Performs a batch synchronous read operation to transfer multiple data chunks from remote buffers to local buffers.
@@ -334,6 +345,7 @@ Performs a batch synchronous read operation to transfer multiple data chunks fro
 - `buffers` (List[int]): List of local buffer addresses
 - `peer_buffer_addresses` (List[int]): List of remote buffer addresses
 - `lengths` (List[int]): List of byte lengths for each transfer
+- `transport_hint` (str, optional): TENT-only per-batch transport pin. Default `""` (policy-driven).
 
 **Returns:**
 - `int`: 0 on success, negative value on failure
@@ -341,7 +353,7 @@ Performs a batch synchronous read operation to transfer multiple data chunks fro
 #### batch_transfer_sync()
 
 ```python
-batch_transfer_sync(target_hostname, buffers, peer_buffer_addresses, lengths, opcode, notify=None)
+batch_transfer_sync(target_hostname, buffers, peer_buffer_addresses, lengths, opcode, notify=None, transport_hint="")
 ```
 
 Performs a batch synchronous transfer operation with specified opcode and optional notification.
@@ -353,6 +365,7 @@ Performs a batch synchronous transfer operation with specified opcode and option
 - `lengths` (List[int]): List of byte lengths for each transfer
 - `opcode` (TransferOpcode): The transfer operation type (READ or WRITE)
 - `notify` (TransferNotify, optional): Notification object to send after transfer completion
+- `transport_hint` (str, optional): TENT-only per-batch transport pin. Default `""` (policy-driven).
 
 **Returns:**
 - `int`: 0 on success, negative value on failure
@@ -360,7 +373,7 @@ Performs a batch synchronous transfer operation with specified opcode and option
 #### batch_transfer_async_write()
 
 ```python
-batch_transfer_async_write(target_hostname, buffers, peer_buffer_addresses, lengths)
+batch_transfer_async_write(target_hostname, buffers, peer_buffer_addresses, lengths, transport_hint="")
 ```
 
 Submits a batch asynchronous write operation and returns immediately.
@@ -370,6 +383,7 @@ Submits a batch asynchronous write operation and returns immediately.
 - `buffers` (List[int]): List of local buffer addresses
 - `peer_buffer_addresses` (List[int]): List of remote buffer addresses
 - `lengths` (List[int]): List of byte lengths for each transfer
+- `transport_hint` (str, optional): TENT-only per-batch transport pin. Default `""` (policy-driven).
 
 **Returns:**
 - `int`: Batch ID for tracking the operation, or 0 on failure
@@ -377,7 +391,7 @@ Submits a batch asynchronous write operation and returns immediately.
 #### batch_transfer_async_read()
 
 ```python
-batch_transfer_async_read(target_hostname, buffers, peer_buffer_addresses, lengths)
+batch_transfer_async_read(target_hostname, buffers, peer_buffer_addresses, lengths, transport_hint="")
 ```
 
 Submits a batch asynchronous read operation and returns immediately.
@@ -387,6 +401,7 @@ Submits a batch asynchronous read operation and returns immediately.
 - `buffers` (List[int]): List of local buffer addresses
 - `peer_buffer_addresses` (List[int]): List of remote buffer addresses
 - `lengths` (List[int]): List of byte lengths for each transfer
+- `transport_hint` (str, optional): TENT-only per-batch transport pin. Default `""` (policy-driven).
 
 **Returns:**
 - `int`: Batch ID for tracking the operation, or 0 on failure
@@ -394,7 +409,7 @@ Submits a batch asynchronous read operation and returns immediately.
 #### batch_transfer_async()
 
 ```python
-batch_transfer_async(target_hostname, buffers, peer_buffer_addresses, lengths, opcode)
+batch_transfer_async(target_hostname, buffers, peer_buffer_addresses, lengths, opcode, transport_hint="")
 ```
 
 Submits a batch asynchronous transfer operation with specified opcode and returns immediately.
@@ -405,6 +420,7 @@ Submits a batch asynchronous transfer operation with specified opcode and return
 - `peer_buffer_addresses` (List[int]): List of remote buffer addresses
 - `lengths` (List[int]): List of byte lengths for each transfer
 - `opcode` (TransferOpcode): The transfer operation type (READ or WRITE)
+- `transport_hint` (str, optional): TENT-only per-batch transport pin. Default `""` (policy-driven).
 
 **Returns:**
 - `int`: Batch ID for tracking the operation, or 0 on failure
@@ -426,7 +442,7 @@ Waits for multiple batch asynchronous transfer operations to complete.
 #### batch_transfer_write_on_cuda()
 
 ```python
-batch_transfer_write_on_cuda(target_hostname, buffers, peer_buffer_addresses, lengths, stream_ptr)
+batch_transfer_write_on_cuda(target_hostname, buffers, peer_buffer_addresses, lengths, stream_ptr=0, transport_hint="")
 ```
 
 Performs a batch write operation to transfer multiple data chunks from local buffers to remote buffers on a given cuda stream.
@@ -437,6 +453,7 @@ Performs a batch write operation to transfer multiple data chunks from local buf
 - `peer_buffer_addresses` (List[int]): List of remote buffer addresses
 - `lengths` (List[int]): List of byte lengths for each transfer
 - `stream_ptr` (int): The integer representation of a CUDA stream pointer (`cudaStream_t`). For example, from a PyTorch stream, this can be obtained via `stream.cuda_stream`.
+- `transport_hint` (str, optional): TENT-only per-batch transport pin. Default `""` (policy-driven).
 
 **Returns:**
 - `None`: The function returns immediately after successfully scheduling the transfer callback.
@@ -450,7 +467,7 @@ Performs a batch write operation to transfer multiple data chunks from local buf
 #### batch_transfer_read_on_cuda()
 
 ```python
-batch_transfer_read_on_cuda(target_hostname, buffers, peer_buffer_addresses, lengths, stream_ptr)
+batch_transfer_read_on_cuda(target_hostname, buffers, peer_buffer_addresses, lengths, stream_ptr=0, transport_hint="")
 ```
 
 Performs a batch read operation to transfer multiple data chunks from remote buffers to local buffers on a given cuda stream.
@@ -461,6 +478,7 @@ Performs a batch read operation to transfer multiple data chunks from remote buf
 - `peer_buffer_addresses` (List[int]): List of remote buffer addresses
 - `lengths` (List[int]): List of byte lengths for each transfer
 - `stream_ptr` (int): The integer representation of a CUDA stream pointer (`cudaStream_t`). For example, from a PyTorch stream, this can be obtained via `stream.cuda_stream`.
+- `transport_hint` (str, optional): TENT-only per-batch transport pin. Default `""` (policy-driven).
 
 **Returns:**
 - `None`: The function returns immediately after successfully scheduling the transfer callback.
@@ -588,6 +606,25 @@ Gets the list of pending transfer notifications received from other nodes.
 **Returns:**
 - `List[TransferNotify]`: List of notification objects containing name and message
 
+#### send_probe()
+
+```python
+send_probe(peer_server_name)
+```
+
+Sends a lightweight JSON-RPC probe to a peer to verify reachability. Used to
+test whether a previously-unreachable peer has recovered (e.g. SGLang's
+`MooncakeKVManager` uses this to clear entries from its `failed_sessions`
+blacklist).
+
+**Parameters:**
+- `peer_server_name` (str): Peer hostname in `host:port` form, as registered
+  with the metadata server.
+
+**Returns:**
+- `int`: 0 on success, non-zero on failure (peer unknown, unreachable, or
+  RPC error).
+
 ## Environment Variables
 
 The Transfer Engine respects the following environment variables:
@@ -596,7 +633,9 @@ The Transfer Engine respects the following environment variables:
 - `MC_METADATA_SERVER`: Default metadata server address
 - `MC_LEGACY_RPC_PORT_BINDING`: Enables legacy RPC port binding behavior
 - `MC_TCP_BIND_ADDRESS`: Specifies the TCP bind address
+- `MC_RDMA_BIND_ADDRESS`: Specifies the RDMA bind address for NIC path construction in dual-NIC environments. When set, RDMA NIC paths use this address while TCP handshake uses the address from `local_hostname`. This is useful when TCP and RDMA traffic use separate network interfaces (e.g., `eth0` for TCP and `rdma-net1` for RDMA).
 - `MC_CUSTOM_TOPO_JSON`: Path to custom topology JSON file
+- `MC_TE_FILTERS`: Optional comma-separated whitelist of IB device names (e.g. `mlx5_0,mlx5_2`) for legacy Transfer Engine topology discovery. When unset, all available devices are discovered.
 - `MC_TE_METRIC`: Enables metrics reporting (set to "1", "true", "yes", or "on"). **Note:** Not supported when using Transfer Engine TENT.
 - `MC_TE_METRIC_INTERVAL_SECONDS`: Sets metrics reporting interval in seconds
 
@@ -605,6 +644,7 @@ The Transfer Engine respects the following environment variables:
 ### Basic Setup and Data Transfer
 
 ```python
+import numpy as np
 from mooncake.engine import TransferEngine
 import os
 
@@ -631,7 +671,7 @@ data_len = len(data)
 engine.register_memory(buffer_data, buffer_data_len)
 
 # Get Remote Addr from ZMQ or upper-layer inference framework
-remote_addr = ??
+remote_addr = REMOTE_ADDR
 
 # Transfer data to remote node
 ret = engine.transfer_sync_write(
@@ -796,3 +836,32 @@ The Transfer Engine Python API is thread-safe for most operations. However, it's
 2. **Transfer Failures**: Verify target hostname is correct and network connectivity is established
 3. **Memory Issues**: Ensure sufficient system memory and proper buffer alignment
 4. **Performance Issues**: Check RDMA device configuration and network topology
+
+## Compile-time Feature Support Attributes
+
+The `mooncake.engine` module provides boolean attributes that indicate whether specific features were enabled during compilation:
+
+### Module Attributes
+
+- `engine.SUPPORT_CUDA`: Whether CUDA support is enabled
+- `engine.SUPPORT_EFA`: Whether EFA (Elastic Fabric Adapter) support is enabled
+- `engine.SUPPORT_HIP`: Whether HIP (Heterogeneous-compute Interface for Portability) support is enabled
+- `engine.SUPPORT_MNNVL`: Whether MNNVL transport protocol support is enabled
+- `engine.SUPPORT_INTRA_NVLINK`: Whether intra-node NVLink support is enabled
+
+### Usage Example
+
+```python
+from mooncake import engine
+
+# Check if CUDA is supported
+if engine.SUPPORT_CUDA:
+    print("CUDA support is available")
+
+# Check all features
+print(f"CUDA: {engine.SUPPORT_CUDA}")
+print(f"EFA: {engine.SUPPORT_EFA}")
+print(f"HIP: {engine.SUPPORT_HIP}")
+print(f"MNNVL: {engine.SUPPORT_MNNVL}")
+print(f"Intra-NVLink: {engine.SUPPORT_INTRA_NVLINK}")
+```
diff --git a/docs/source/troubleshooting/pg-ep-troubleshooting.md b/docs/source/troubleshooting/pg-ep-troubleshooting.md
new file mode 100644
index 00000000..684e08f9
--- /dev/null
+++ b/docs/source/troubleshooting/pg-ep-troubleshooting.md
@@ -0,0 +1,259 @@
+# Mooncake PG/EP Troubleshooting
+
+This page covers common setup, import, runtime, and recovery issues for
+Mooncake Backend (PG) and Mooncake EP.
+
+## Import fails with a PyTorch version error
+
+Symptoms:
+
+```text
+Mooncake PG was not built against torch==...
+Mooncake EP was not built against torch==...
+```
+
+Cause:
+
+`mooncake.pg` and `mooncake.ep` load native extension modules whose names include
+the current PyTorch version. If the wheel or source build does not contain an
+extension for the active `torch.__version__`, import fails.
+
+Fixes:
+
+1. Verify the active PyTorch version:
+
+   ```bash
+   python - <<'PY'
+   import torch
+   print(torch.__version__)
+   print(torch.version.cuda)
+   PY
+   ```
+
+2. Install a Mooncake wheel built for that PyTorch/CUDA combination, or rebuild
+   from source with EP/PG enabled:
+
+   ```bash
+   cmake .. -DWITH_EP=ON
+   make -j
+   ```
+
+3. Make sure the Python environment used at runtime is the same one used for the
+   build.
+
+## `activeRanks must be int` or device mismatch
+
+Symptoms:
+
+```text
+activeRanks must be int.
+activeRanks must be on CPU.
+activeRanks must be on GPU.
+activeRanks must be sized to max_world_size when max_world_size is set
+```
+
+Causes and fixes:
+
+- Use `torch.int32`, not `bool`, `int64`, or floating-point types.
+- For `backend="mooncake"`, put `active_ranks` on the accelerator device.
+- For `backend="mooncake-cpu"`, put `active_ranks` on CPU.
+- If `max_world_size` is set, allocate `active_ranks` with length
+  `max_world_size`, even if the initial visible `world_size` is smaller.
+
+Examples:
+
+```python
+# CUDA / accelerator backend
+active_ranks = torch.ones(max_world_size, dtype=torch.int32, device="cuda")
+
+# CPU backend
+active_ranks = torch.ones(max_world_size, dtype=torch.int32)
+```
+
+## `dist.get_world_size()` is smaller than `max_world_size`
+
+This is expected. Mooncake PG distinguishes reserved capacity from visible active
+membership:
+
+- `max_world_size` reserves future rank slots.
+- `dist.get_world_size()` returns the visible active size.
+- Reserved ranks are inactive until `recover_ranks()` activates them.
+
+Use `pg.get_active_ranks(backend)` to inspect the current backend mask.
+
+## `get_peer_state()` or `join_group()` hangs
+
+Common causes:
+
+- Healthy ranks are not all calling `get_peer_state()` in the same order.
+- The joining rank did not initialize with `is_extension=True`.
+- `max_world_size` / rank numbering differs between healthy and joining ranks.
+- Subgroups were created in different orders on healthy and joining processes.
+- The joining process has not published peer metadata yet.
+- Network device filters differ across ranks.
+
+Debug checklist:
+
+1. Confirm all ranks use the same rendezvous address and store.
+2. Print rank, visible `world_size`, and `max_world_size` at initialization.
+3. Confirm `active_ranks` length and values on every rank.
+4. Verify healthy ranks poll the same `join_ranks` list.
+5. For subgroup recovery, confirm every rank calls `dist.new_group()` in the same
+   order.
+6. If using RDMA, set the same HCA whitelist on every rank.
+
+## Newly extended ranks participate too early
+
+After `pg.extend_group_size_to(backend, new_size)`, new ranks are reserved but
+inactive. They should not participate in collectives until healthy ranks call
+`pg.recover_ranks(backend, ranks)`.
+
+If a new rank appears to participate early, check whether:
+
+- `active_ranks` was initialized with `1` for future ranks without masking them
+  through the backend protocol;
+- the application called collectives on the joining process before
+  `pg.join_group()` returned;
+- different ranks used inconsistent `max_world_size` or rank IDs.
+
+## EP dispatch/combine timeout marks a rank inactive
+
+Mooncake EP kernels can update the EP-level `active_ranks` tensor when a source
+rank does not make progress before `timeout_us`.
+
+If this happens unexpectedly:
+
+- Increase `timeout_us` to rule out slow startup or transient scheduling delay.
+- Check whether the source rank exited, crashed, or skipped the matching
+  dispatch/combine call.
+- Confirm `num_experts % num_ranks == 0` and that every rank uses the same
+  `num_experts`, `top_k`, and buffer sizing assumptions.
+- Ensure all ranks call dispatch and combine in the same order.
+
+If timeout detection is not desired for a test, pass `timeout_us=-1`.
+
+## EP falls back instead of using the fast path
+
+The Python wrapper uses fallback when the native runtime cannot use IBGDA/RDMA or
+fully accessible P2P.
+
+Possible causes:
+
+- RDMA devices or drivers are not available inside the container.
+- HCA selection picked the wrong device.
+- GPUDirect RDMA / peer-memory support is missing.
+- CUDA IPC or peer access is unavailable between local GPUs.
+- The environment was built without the required accelerator support.
+
+Debug checklist:
+
+1. Restrict HCA selection with `pg.set_device_filter([...])` or
+   `MOONCAKE_PGTEST_DEVICE_FILTERS`.
+2. Check whether all intended ranks can see the same accelerator and RDMA
+   devices.
+3. Run PG collectives first; EP metadata exchange depends on a healthy process
+   group.
+4. Run EP tests with fallback enabled and disabled to separate correctness from
+   transport setup.
+
+## RDMA connection errors or severe latency from wrong HCA selection
+
+Symptoms can include connection setup failures, repeated transport timeouts,
+unexpected fallback, very low bandwidth, or RDMA retry-style errors when the
+backend auto-selects an unsuitable NIC/HCA. This commonly happens on machines
+with multiple RDMA devices where some devices are for management traffic, are on
+different fabrics, or are not reachable from peer ranks.
+
+Fix: set an explicit device filter before initializing Mooncake PG/EP so all
+ranks use the intended HCA list.
+
+For application code:
+
+```python
+from mooncake import pg
+
+pg.set_device_filter(["mlx5_1", "mlx5_2"])
+# call dist.init_process_group(...) after setting the filter
+```
+
+For PG tests and benchmarks:
+
+```bash
+MOONCAKE_PGTEST_DEVICE_FILTERS=mlx5_1,mlx5_2 \
+python -m unittest discover -s mooncake-pg/tests -k CUDA -v
+```
+
+Use the same filter on every rank. If the problem disappears after setting the
+filter, treat the original failure as a topology/device-selection issue rather
+than an EP kernel or collective correctness problem.
+
+## EP output mismatch or buffer overflow symptoms
+
+Common causes:
+
+- `num_max_dispatch_tokens_per_rank` is smaller than the actual per-rank token
+  count.
+- `num_experts` is not divisible by `num_ranks`.
+- `topk_idx` contains expert IDs outside `[0, num_experts)` except for masked
+  `-1` entries.
+- The `handle` from one dispatch call is reused with an unrelated combine call.
+- `zero_copy=True` is used without writing expert outputs into
+  `get_next_combine_buffer(handle)`.
+
+Fixes:
+
+- Size the buffer for peak traffic:
+
+  ```python
+  num_ep_buffer_bytes = Buffer.get_ep_buffer_size_hint(
+      max_tokens_per_rank, hidden, world_size, num_experts
+  )
+  ```
+
+- Keep dispatch and combine paired by using the matching `handle`.
+- Call `event.current_stream_wait()` or the returned `hook()` before consuming
+  operation outputs.
+
+## P2P send/recv errors
+
+Mooncake PG's PyTorch P2P path currently supports single-tensor send/recv
+through the backend shim. If `batch_isend_irecv()` fails:
+
+- verify each `P2POp` contains one tensor;
+- ensure peer ranks are active;
+- check that all ranks issue matching send/recv operations;
+- test CPU and accelerator backends separately to isolate device-specific
+  transport issues.
+
+## Useful smoke commands
+
+```bash
+# PG CPU sanity
+python -m unittest discover -s mooncake-pg/tests -k CPU -v
+
+# PG CUDA sanity
+python -m unittest discover -s mooncake-pg/tests -k CUDA -v
+
+# PG all-reduce benchmark smoke
+PYTHONPATH=mooncake-pg \
+python mooncake-pg/benchmark/pgbench.py \
+  --collective all_reduce --backend mooncake --device cuda -g 2 -b 8 -e 1M -f 2
+
+# EP grid test
+python mooncake-ep/tests/test_ep_grid.py
+```
+
+Adapt process counts, device filters, and launchers to the target cluster.
+
+## What to include in bug reports
+
+When reporting PG/EP issues, include:
+
+- Mooncake commit and installation method;
+- PyTorch version and CUDA/accelerator runtime version;
+- exact backend name (`mooncake` or `mooncake-cpu`);
+- world size, max world size, rank IDs, and subgroup layout;
+- active-rank tensor dtype/device/values;
+- HCA/device filters;
+- minimal command or script reproducing the issue;
+- logs from all ranks around initialization, failure detection, and recovery.
diff --git a/docs/source/troubleshooting/troubleshooting.md b/docs/source/troubleshooting/troubleshooting.md
index b85f55a7..3bf11c2e 100644
--- a/docs/source/troubleshooting/troubleshooting.md
+++ b/docs/source/troubleshooting/troubleshooting.md
@@ -11,7 +11,7 @@ This document lists common errors that may occur when using Mooncake Store and p
 
 ## Metadata and Out-of-Band Communication
 1. At startup, a `TransferMetadata` object is constructed according to the incoming `metadata_server` parameter. During program execution, this object is used to communicate with the etcd server to maintain internal data required for connection.
-2. At startup, the current node is registered with the cluster according to the incoming `connectable_name` parameter and `rpc_port` parameter, and the TCP port specified by the `rpc_port` parameter is listened. Before other nodes send the first read/write request to the current node, they will use the above information, resolve DNS, and initiate a connection through socket's `connect()` method. 
+2. At startup, the current node is registered with the cluster according to the incoming `connectable_name` parameter and `rpc_port` parameter, and the TCP port specified by the `rpc_port` parameter is listened. Before other nodes send the first read/write request to the current node, they will use the above information, resolve DNS, and initiate a connection through socket's `connect()` method.
 
 Errors in this part usually indicate that the error occurred within the `mooncake-transfer-engine/src/transfer_metadata.cpp` file.
 
@@ -110,9 +110,10 @@ Errors in this part usually indicate that the error occurred within the `mooncak
      *    hard    memlock    unlimited
      ```
 
-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:
+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. Possible ways 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
+   - If the leak persists under sustained peer failures (many `endpoint evicted` log lines accompanying the QP growth), update to a version that includes the fix for [issue #1845](https://github.com/kvcache-ai/Mooncake/issues/1845). Prior to that fix, the endpoint store's `waiting_list_` only drained when new endpoints were inserted, so evictions under failure load accumulated QPs until the driver limit was hit. The fix adds a periodic reclaim tick to `monitorWorker`.
 
 ## RDMA Transfer Period
 ### Recommended Troubleshooting Directions
@@ -123,6 +124,74 @@ Note: In most cases, the errors output, except for the first occurrence, are `wo
 
 In addition, if the error `Failed to get description of XXX` is displayed, it indicates that the Segment name input by the user when calling the `openSegment` interface cannot be found in the etcd database. For memory read/write scenarios, the Segment name needs to strictly match the `local_hostname` field filled in by the other node during initialization.
 
+## TCP Transport
+### Recommended Troubleshooting Directions
+
+1. If sustained, high-concurrency TCP traffic (for example, PD-disaggregated KV transfers over the TCP transport) fails with `connect: Cannot assign requested address`, the initiator side has exhausted its ephemeral port range. Each transfer opens a fresh short-lived socket, and ports held in `TIME_WAIT` accumulate faster than the kernel can reclaim them.
+
+   **Diagnostic Commands:**
+   ```bash
+   # Confirm large numbers of TIME_WAIT sockets to the peer
+   ss -tan state time-wait | wc -l
+
+   # Check the local ephemeral port range
+   sysctl net.ipv4.ip_local_port_range
+   ```
+
+   **Solutions:**
+   - Enable the TCP connection pool so that long-lived sockets are reused across transfers instead of being opened per transfer:
+     ```bash
+     export MC_TCP_ENABLE_CONNECTION_POOL=1
+     ```
+   - Widen the ephemeral port range if the workload genuinely needs many distinct connections:
+     ```bash
+     sysctl -w net.ipv4.ip_local_port_range="1024 65535"
+     ```
+   - As a last resort, enable `TIME_WAIT` reuse on the initiator. `tcp_tw_reuse` only affects outbound connections and requires TCP timestamps to be enabled on both sides:
+     ```bash
+     sysctl -w net.ipv4.tcp_tw_reuse=1
+     ```
+
+2. The TCP connection pool (`MC_TCP_ENABLE_CONNECTION_POOL=1`) has two known limitations in the current implementation. The pool is functional for most workloads, but operators running long-lived services should be aware of them:
+
+   * **Idle-expired connections are not always reclaimed.** `cleanupIdleConnections()` only inspects the tail of each endpoint's deque. When a newer `in_use` entry has been pushed behind an older idle-expired entry, the loop terminates at the tail and the idle entry is never removed. The pool's reported state diverges from reality: it still lists the connection as idle while the socket has been sitting past `kConnectionIdleTimeout`, may have been half-closed by the peer, and will fail on the next `getConnection()` that tries to reuse it. In long-running services this also leaks sockets and file descriptors until the process is restarted.
+
+     **Diagnostic Commands:**
+     ```bash
+     # Watch for unbounded growth of open sockets from the process
+     ls /proc/$PID/fd | wc -l
+     ss -tan | awk '$1=="ESTAB"' | wc -l
+     ```
+
+     **Workaround:** restart the process periodically if fd usage climbs without bound.
+
+   * **`asio::socket::close()` runs under `pool_mutex_`.** `close()` cancels outstanding async operations and posts completion handlers to the io_context. Handlers such as `returnConnection()` also acquire `pool_mutex_`, so the current code is one scheduling step away from a circular wait. No deadlock has been observed in practice, but if the transfer engine hangs with all worker threads stuck waiting on `pool_mutex_`, this is the first place to look.
+
+## Memory Allocator
+
+Mooncake Store's mmap arena is opt-in for mmap buffer allocations. Setting `MC_MMAP_ARENA_POOL_SIZE` explicitly enables it and pre-allocates a hugepage-backed pool; if it is enabled via gflag instead, the default pool size is `8gb`. If the arena cannot allocate hugepages, it falls back to regular pages automatically.
+
+If you encounter memory allocation issues related to the arena:
+
+- **HugeTLB pool is too small for the launch:** Run `python3 scripts/check_hicache_hugepage_requirements.py ...` from a source checkout, or `mooncake-hicache-sizing ...` inside the Docker image, using your `--hicache-size`, `MOONCAKE_GLOBAL_SEGMENT_SIZE`, and `MC_MMAP_ARENA_POOL_SIZE` values. If it reports `insufficient_for_baseline`, increase `vm.nr_hugepages`, reduce `--hicache-size`, or shrink `MOONCAKE_GLOBAL_SEGMENT_SIZE`.
+- **Arena pool too large for available hugepages:** Reduce the pool size with `MC_MMAP_ARENA_POOL_SIZE="8gb"` to fit within your hugepage budget, or leave `MC_MMAP_ARENA_POOL_SIZE` unset to stay on the baseline direct-`mmap()` path.
+- **Arena only partially fits:** If the helper reports `baseline_fits_arena_may_fallback`, the baseline should start but some arena allocations may spill onto regular pages. Either increase `vm.nr_hugepages` or reduce `MC_MMAP_ARENA_POOL_SIZE`.
+- **Need to disable the arena entirely:** Set `MC_DISABLE_MMAP_ARENA=1` (also accepts `true`, `yes`, or `on`) before the first Mooncake mmap-buffer allocation to fall back to per-call `mmap()`.
+- **Arena init failed once and the process stayed on the baseline path:** Arena initialization is lazy and one-shot per process. After a failed first attempt, Mooncake keeps using direct `mmap()` until the process restarts. Fix the env / hugepage budget, then restart the process before retrying.
+- **Arena OOM during serving:** The arena logs a warning and falls back to direct `mmap()` for that allocation. If this happens frequently, increase `MC_MMAP_ARENA_POOL_SIZE`.
+- **Direct hugepage mmap fails immediately:** Verify the host really has hugepages reserved with `grep -E 'HugePages_Total|HugePages_Free|Hugepagesize' /proc/meminfo`. When `MC_STORE_USE_HUGEPAGE=1` is set, direct `mmap()` allocations depend on the OS HugeTLB pool.
+
+To provision `2 MiB` hugepages on Linux:
+
+```bash
+sudo sysctl -w vm.nr_hugepages=49152
+grep -E 'HugePages_Total|HugePages_Free|Hugepagesize' /proc/meminfo
+printf 'vm.nr_hugepages=49152\n' | sudo tee /etc/sysctl.d/90-mooncake-hugepages.conf
+sudo sysctl --system
+```
+
+`49152` pages equals `96 GiB`; for `512 GiB` use `262144` pages.
+
 ## SGLang Common Questions
 
 ### Do I need RDMA to run SGLang and Mooncake?
diff --git a/docs/source/zh_archive/README.md b/docs/source/zh_archive/README.md
index 4ef7d418..32a8a705 100644
--- a/docs/source/zh_archive/README.md
+++ b/docs/source/zh_archive/README.md
@@ -5,4 +5,4 @@
 ## 当前状态
 - 已归档，不再活跃维护
 - 网站当前只显示英文文档
-- 如需查看/编辑最新文档，请访问docs目录
\ No newline at end of file
+- 如需查看/编辑最新文档，请访问docs目录
diff --git a/docs/source/zh_archive/ascend_direct_transport.md b/docs/source/zh_archive/ascend_direct_transport.md
index 50ed7dbc..dade6374 100644
--- a/docs/source/zh_archive/ascend_direct_transport.md
+++ b/docs/source/zh_archive/ascend_direct_transport.md
@@ -36,7 +36,7 @@ Transfer Engine RPC using <协议> listening on <IP>:<实际端口>，记录目
 
 | 变量名 | 描述                                   | 默认值 | 示例                                                                    |
 |--------|--------------------------------------|--------|-----------------------------------------------------------------------|
-| `ASCEND_AUTO_CONNECT` | 启用自动连接管理                             | 0（禁用） | `ASCEND_AUTO_CONNECT=1`                                               |
+| `ASCEND_AUTO_CONNECT` | 启用自动连接管理                             | 1（启用） | `ASCEND_AUTO_CONNECT=0`                                               |
 | `ASCEND_ENABLE_USE_FABRIC_MEM` | 在Mooncake Store中启用fabric内存传输模式（仅A3）  | 0（禁用） | `ASCEND_ENABLE_USE_FABRIC_MEM=1`                                      |
 | `ASCEND_USE_ASYNC_TRANSFER` | 启用异步传输模式                             | 0（禁用） | `ASCEND_USE_ASYNC_TRANSFER=1`                                         |
 | `ASCEND_GLOBAL_RESOURCE_CONFIG` | 全局资源配置                               | - | `ASCEND_GLOBAL_RESOURCE_CONFIG="{\"fabric_memory.max_capacity\":32}"` |
@@ -51,7 +51,7 @@ Transfer Engine RPC using <协议> listening on <IP>:<实际端口>，记录目
 | `HCCL_RDMA_RETRY_CNT` | RDMA数据包重传次数                          | - | `HCCL_RDMA_RETRY_CNT=7`                                               |
 
 详细说明：
-ASCEND_AUTO_CONNECT: 需要CANN升级到9.0之后的版本，所以默认值为0，在支持该功能的版本推荐启用：当对端异常下线后可以自动断链。
+ASCEND_AUTO_CONNECT: 需要CANN升级到9.0之后的版本，默认值为1：当对端异常下线后可以自动断链。如需关闭可设置`ASCEND_AUTO_CONNECT=0`。
 ASCEND_ENABLE_USE_FABRIC_MEM：需要CANN升级到9.0之后的版本，HDK升级到26.0之后的版本，在支持该功能的版本使用Mooncake Store时推荐启用：可显著提升传输性能。
 ASCEND_USE_ASYNC_TRANSFER: 需要CANN升级到8.5之后的版本，用于开启Hixl异步传输模式，默认为同步模式。
 ASCEND_GLOBAL_RESOURCE_CONFIG：配置Hixl的全局资源，具体查看hixl的文档关于OPTION_GLOBAL_RESOURCE_CONFIG的配置。
diff --git a/docs/source/zh_archive/build.md b/docs/source/zh_archive/build.md
index 5f9b94e0..62ff9d2e 100644
--- a/docs/source/zh_archive/build.md
+++ b/docs/source/zh_archive/build.md
@@ -16,7 +16,7 @@
      请切换安装如下 PyPI 源的 wheel 包:
      ```bash
       pip install mooncake-transfer-engine-non-cuda
-      ```
+     ```
      或者请附带 `-DUSE_CUDA=OFF` 使用源码编译安装。
 
 ## 自动安装
@@ -71,7 +71,7 @@
                    libnuma-dev \
                    libcurl4-openssl-dev \
                    libhiredis-dev
-
+    
     # For centos/alibaba linux os
     yum install cmake \
                 gflags-devel \
@@ -95,12 +95,12 @@
     ```
 
 2. 如果你要编译Nvidia GPUDirect 支持模块，首先需按照 https://docs.nvidia.com/cuda/cuda-installation-guide-linux/ 的指引安装 CUDA (确保启用 `nvidia-fs` 以正确编译 `cuFile` 模块)。之后:
-    1) 按照 https://docs.nvidia.com/cuda/gpudirect-rdma/ 的第 3.7 节说明安装 `nvidia-peermem` 以启用 GPU-Direct RDMA
-    2) 配置 `LIBRARY_PATH` 和 `LD_LIBRARY_PATH` 以确保编译过程期间链入 `cuFile`, `cudart` 等库:
+    1) 配置 `LIBRARY_PATH` 和 `LD_LIBRARY_PATH` 以确保编译过程期间链入 `cuFile`, `cudart` 等库:
     ```bash
     export LIBRARY_PATH=$LIBRARY_PATH:/usr/local/cuda/lib64
     export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/cuda/lib64
     ```
+    > **注意：** Mooncake 可以使用 DMA-BUF 路径进行 GPU-Direct RDMA，**无需** `nvidia-peermem` 内核模块。如需使用 DMA-BUF 路径，请在启动 Mooncake 前设置运行时环境变量 `WITH_NVIDIA_PEERMEM=0`。如果偏好依赖 `nvidia-peermem` 的传统 `ibv_reg_mr` 路径，请设置运行时环境变量 `WITH_NVIDIA_PEERMEM=1`。安装 `nvidia-peermem` 的说明见 https://docs.nvidia.com/cuda/gpudirect-rdma/ 第 3.7 节。
 
 3. 如果你要编译Moore Threads GPUDirect RDMA 支持模块，首先需按照 https://docs.mthreads.com/musa-sdk/musa-sdk-doc-online/install_guide 的指引安装 MUSA SDK。之后:
     1) 安装 `mthreads-peermem` 以启用 GPU-Direct RDMA
@@ -144,7 +144,20 @@
     - `-DMACA_LIB_DIR=/path/to/maca/lib64`
     - `-DMACA_RUNTIME_LIBS="mcruntime;mxc-runtime64;rt"`（分号分隔的 CMake 列表）
 
-6. 安装 yalantinglibs
+6. 若需编译华为 Ascend NPU 支持，请先按照 https://www.hiascend.com/document 的指引安装 Ascend CANN Toolkit。之后：
+    1) 在 CANN 安装目录下执行 `source set_env.sh` 配置编译环境（无需手动设置 `ASCEND_HOME_PATH` 等环境变量）。
+    2) Mooncake 提供两种 Ascend NPU 传输路径，按需选择其中一个：
+       - `-DUSE_ASCEND_DIRECT=ON`（**推荐**）：基于 ADXL 引擎的 Ascend Direct 传输（[版本配套关系参考](https://gitcode.com/cann/hixl/wiki/Mooncake%20+%20HIXL%20%E5%BF%AB%E9%80%9F%E4%B8%8A%E6%89%8B%E6%8C%87%E5%8D%97.md)）。
+       - `-DUSE_UBSHMEM=ON`：基于 CANN VMM接口传输（CANN版本 >= 9.0.0 、驱动版本 >= 26.0.0、灵衢版本 >= 1.5）
+
+    启用 Ascend NPU 编译示例：
+    ```bash
+    source /usr/local/Ascend/cann/set_env.sh
+    cmake .. -DUSE_ASCEND_DIRECT=ON
+    make -j
+    ```
+
+7. 安装 yalantinglibs
     ```bash
     git clone https://github.com/alibaba/yalantinglibs.git
     cd yalantinglibs
@@ -154,7 +167,7 @@
     make install
     ```
 
-7. 进入项目根目录，运行下列命令进行编译
+8. 进入项目根目录，运行下列命令进行编译
    ```bash
    mkdir build
    cd build
@@ -162,7 +175,7 @@
    make -j
    ```
 
-8. 安装 Mooncake python 包和 mooncake_master 可执行文件
+9. 安装 Mooncake python 包和 mooncake_master 可执行文件
    ```bash
    make install
    ```
@@ -180,10 +193,17 @@
 - `-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_HYGON=[ON|OFF]`: 通过 DTK SDK 启用对海光 DCU 的支持。默认 OFF；使用 CUDA 兼容运行时。
+- `-DDTK_ROOT=/path/to/dtk`: 在 `-DUSE_HYGON=ON` 时覆盖默认 DTK SDK 根路径；未设置时使用 `DTK_HOME` 或 `/opt/dtk`。
+- `-DDTK_INCLUDE_DIR=/path/to/include` / `-DDTK_LIB_DIR=/path/to/lib64`: 在 `-DUSE_HYGON=ON` 时覆盖 DTK 头文件与库目录。
+- `-DUSE_COREX=[ON|OFF]`: 启用对天数智芯 CoreX GPU 的支持。默认 OFF；使用 CUDA 兼容运行时。
+- `-DCOREX_ROOT=/path/to/corex`: 在 `-DUSE_COREX=ON` 时覆盖默认 CoreX SDK 根路径；未设置时使用 `COREX_HOME` 或 `/usr/local/corex`。
+- `-DCOREX_INCLUDE_DIR=/path/to/include` / `-DCOREX_LIB_DIR=/path/to/lib`: 在 `-DUSE_COREX=ON` 时覆盖 CoreX 头文件与库目录。
 - `-DUSE_CXL=[ON|OFF]`: 启用 CXL 支持
 - `-DWITH_STORE=[ON|OFF]`: 编译 Mooncake Store 组件
 - `-DWITH_P2P_STORE=[ON|OFF]`: 启用 Golang 支持并编译 P2P Store 组件，需要 go 1.23+
-- `-DWITH_WITH_RUST_EXAMPLE=[ON|OFF]`: 启用 Rust 支持
+- `-DWITH_RUST_EXAMPLE=[ON|OFF]`: 编译 Transfer Engine Rust 接口与示例代码，默认关闭。
+- `-DWITH_STORE_RUST=[ON|OFF]`: 编译 Mooncake Store Rust 绑定和 CMake Rust targets，默认开启。
 - `-DUSE_REDIS=[ON|OFF]`: 启用基于 Redis 的元数据服务
 - `-DUSE_HTTP=[ON|OFF]`: 启用基于 Http 的元数据服务
 - `-DUSE_ETCD=[ON|OFF]`: 启用基于 etcd 的元数据服务，需要 go 1.23+
@@ -191,7 +211,8 @@
 - `-DBUILD_SHARED_LIBS=[ON|OFF]`: 将 Transfer Engine 编译为共享库，默认为 OFF
 - `-DBUILD_UNIT_TESTS=[ON|OFF]`: 编译单元测试，默认为 ON
 - `-DBUILD_EXAMPLES=[ON|OFF]`: 编译示例程序，默认为 ON
-- `-DUSE_ASCEND_DIRECT=[ON|OFF]`: 启用 Ascend Direct RDMA 及 HCCS 支持
+- `-DUSE_ASCEND_DIRECT=[ON|OFF]`: 通过 ADXL 引擎启用 Ascend Direct 传输及 HCCS 支持（**NPU 推荐方式**）
+- `-DUSE_UBSHMEM=[ON|OFF]`: 通过 CANN VMM API 启用华为 Ascend NPU 共享内存传输
 - `-DUSE_MUSA=[ON|OFF]`: 启用Moore Threads GPUDirect RDMA
 
 ## 在 Docker 容器中使用 Mooncake
diff --git a/docs/source/zh_archive/mooncake-store.md b/docs/source/zh_archive/mooncake-store.md
index 3ddf3d83..768b310a 100644
--- a/docs/source/zh_archive/mooncake-store.md
+++ b/docs/source/zh_archive/mooncake-store.md
@@ -699,7 +699,10 @@ mooncake提供了DFS可用空间的配置，用户可以在启动master时指定
 
 启用持久化功能后，对于每次 `Put`或`BatchPut` 操作，都会发起一次同步的memory pool写入操作和一次异步的DFS持久化操作。之后执行 `Get`或 `BatchGet` 时，如果在memory pool中没有找到对应的kvcache，则会尝试从DFS中读取该文件数据，并返回给用户。
 
-#### 3FS USRBIO 插件
+#### 3FS USRBIO 插件（实验性 / 未完成）
+
+> **实验性功能：** HF3FS（3FS USRBIO）相关集成仍在开发中，尚未达到可视为生产就绪的程度；接口与行为可能变更，建议仅用于评估与测试。
+
 如需通过3FS原生接口（USRBIO）实现高性能持久化文件读写，请参阅本文档的配置说明。[3FS USRBIO 插件配置](/mooncake-store/src/hf3fs/README.md)。
 
 ### 内置元数据服务器
@@ -718,10 +721,12 @@ HTTP 元数据服务器可通过以下参数进行配置：
 
 #### 环境变量说明
 
-- **MC_STORE_CLUSTER_ID**: 在多集群复用 master 场景下标识元数据, 默认 'mooncake'
+- **MC_STORE_CLUSTER_ID**: 在多集群复用 master 场景下标识元数据, 默认 'mooncake_cluster'
 - **MC_STORE_MEMCPY**: 控制是否启用本地 memcpy 优化, 1/true 启用, 0/false 禁用
 - **MC_STORE_CLIENT_METRIC**: 启用客户端指标上报, 默认启用；设为 0/false 禁用
 - **MC_STORE_CLIENT_METRIC_INTERVAL**: 指标上报间隔(秒), 默认 0(仅收集不上报)
+- **MC_STORE_CLIENT_MIN_PORT**: 客户端使用的最小端口号, 默认 12300。端口范围必须在 1024–32767 或 61000–65535 之间（排除知名端口和临时端口）。无效值将回退为默认值。
+- **MC_STORE_CLIENT_MAX_PORT**: 客户端使用的最大端口号, 默认 14300。范围约束同上；必须 ≥ MC_STORE_CLIENT_MIN_PORT。
 - **MC_STORE_USE_HUGEPAGE**: 启用 hugepage 优化, 默认禁用, 设置为 1/true 启用
 - **MC_STORE_HUGEPAGE_SIZE**: hugepage 页大小, 默认 2M
 
diff --git a/docs/source/zh_archive/run-examples.md b/docs/source/zh_archive/run-examples.md
index bbec1f15..8c72fa97 100644
--- a/docs/source/zh_archive/run-examples.md
+++ b/docs/source/zh_archive/run-examples.md
@@ -34,13 +34,12 @@ Mooncake 目前仅支持 Linux 操作系统，并且依赖以下软件：
 
    首先按照 https://docs.nvidia.com/cuda/cuda-installation-guide-linux/ 中的说明安装 cuda（需要勾选 nvidia-fs 选项, 以便于 cufile 的正常使用），之后：
 
-   1) 参考 https://docs.nvidia.com/cuda/gpudirect-rdma/ 的 3.7 节，安装 nvidia-peermem 以启用 gpu-direct RDMA:
-
-   2) 配置 `LIBRARY_PATH` 和 `LD_LIBRARY_PATH`，用于编译和运行时链接 cufile, cudart 等库：
+   1) 配置 `LIBRARY_PATH` 和 `LD_LIBRARY_PATH`，用于编译和运行时链接 cufile, cudart 等库：
    ```bash
    export $LIBRARY_PATH=$LIBRARY_PATH:/usr/local/cuda/lib64
    export $LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/cuda/lib64
    ```
+   > **注意：** 默认情况下，Mooncake 使用 DMA-BUF 路径进行 GPU-Direct RDMA，**无需** `nvidia-peermem` 内核模块。如需使用依赖 `nvidia-peermem` 的传统 `ibv_reg_mr` 路径，请在启动 Mooncake 前设置运行时环境变量 `WITH_NVIDIA_PEERMEM=1`。安装 `nvidia-peermem` 的说明见 https://docs.nvidia.com/cuda/gpudirect-rdma/ 第 3.7 节。
 
 6. 编译 Mooncake 组件
    ```bash
@@ -56,7 +55,8 @@ Mooncake 支持在执行 `cmake` 命令期间添加下列高级编译选项：
 - `-DUSE_CXL=[ON|OFF]`：编译 Transfer Engine 时启用或关闭 CXL 协议的支持。默认关闭。
 - `-DWITH_P2P_STORE=[ON|OFF]`：编译 P2P Store 及示例程序，默认开启。
 - `-DWITH_ALLOCATOR=[ON|OFF]`：编译 Managed Store 所用的中心分配器模块，默认开启。
-- `-DWITH_WITH_RUST_EXAMPLE=[ON|OFF]`：编译 Transfer Engine 时启用或关闭 Rust 语言支持，默认关闭。
+- `-DWITH_RUST_EXAMPLE=[ON|OFF]`：编译 Transfer Engine Rust 接口与示例代码，默认关闭。
+- `-DWITH_STORE_RUST=[ON|OFF]`：编译 Mooncake Store Rust 绑定和 CMake Rust targets，默认开启。
 
 
 ## Transfer Engine Bench 使用方法
diff --git a/docs/source/zh_archive/sunrise_link_transport.md b/docs/source/zh_archive/sunrise_link_transport.md
new file mode 100644
index 00000000..c78dff4c
--- /dev/null
+++ b/docs/source/zh_archive/sunrise_link_transport.md
@@ -0,0 +1,133 @@
+# Sunrise Link Transport
+
+## 概述
+Sunrise Link 是曦望（Sunrise）开发的 GPU 间高速直连互连技术。在 Mooncake 中，它作为 GPU 传输后端集成于 TENT 传输框架。它通过 Tang Runtime（`tangrt`）提供设备内存分配、指针属性查询、Peer Copy 与 IPC 句柄相关能力，并在 TENT 中以 `SUNRISE_LINK` 传输类型注册。
+
+在运行时，当编译开启 `USE_SUNRISE` 且配置项 `transports/sunrise_link/enable=true` 时，`TransferEngineImpl` 会加载 Sunrise Link Transport。
+
+### 传输流水线
+
+1. **初始化**：`dlopen` 加载 `libptml_shared.so` + `libtangrt_shared.so`，初始化 PTML。
+2. **拓扑检测**：`ptmlPtlinkPhytopoDetect` 发现芯片间端口，构建 C2C 映射（`local_chipid, remote_chipid → local_port`）。
+3. **内存注册**：`tangIpcGetMemHandle` 生成 GPU 缓冲区 IPC 句柄。
+4. **传输执行**：同设备 → `tangMemcpy`；跨设备 → `tangMemcpyPeer` / `tangMemcpyPeer_v2`，回退 `tangDeviceGetPeerPointer` + C2C；远程 → `tangIpcOpenMemHandle` + peer copy；主机 ↔ 设备 → `tangMemcpy` + 方向。
+5. **异步路径**：传输大小超过 `async_memcpy_threshold` 时，使用 `tangMemcpyAsync` + 每设备 stream。
+
+---
+
+## 新增依赖
+Sunrise Link Transport 在 Mooncake 基础依赖外，额外依赖 Tang Runtime：
+
+- **头文件路径**：`/usr/local/tangrt/include`
+- **库路径**：`/usr/local/tangrt/lib/linux-x86_64`（或扁平布局 `/usr/local/tangrt/lib`）
+- **动态库**：`libtangrt_shared.so`（运行时按默认安装路径加载）
+
+建议确认 `tangrt` 与 `ptml` 相关运行库可被动态链接器找到（例如通过 `LD_LIBRARY_PATH` 或系统库路径配置）。
+
+---
+
+## 构建与编译
+
+**前置条件**
+
+- 已安装可用的 Tang Runtime（默认安装到 `/usr/local/tangrt`）
+- 编译环境可访问 Mooncake 及其基础依赖
+
+**CMake 配置**
+
+```bash
+# 克隆 Mooncake 仓库
+git clone https://github.com/kvcache-ai/Mooncake.git
+cd Mooncake
+
+# 启用 TENT + Sunrise Link
+mkdir build && cd build
+cmake .. -DUSE_TENT=ON -DUSE_SUNRISE=ON
+
+# 编译
+make -j$(nproc)
+```
+
+如果 Tang Runtime 安装在非默认路径：
+
+```bash
+cmake .. -DUSE_TENT=ON -DUSE_SUNRISE=ON -DMC_TANGRT_ROOT=/opt/tangrt
+```
+
+---
+
+## 运行与测试
+
+### 使用 `transfer_engine_bench`（经典 TE 后端）
+
+`transfer_engine_bench` 支持 `sunrise_link` 协议，可用于基本连通性与性能验证。
+
+```bash
+# 终端 1：目标端（Target）
+./transfer_engine_bench \
+  --mode=target \
+  --protocol=sunrise_link \
+  --local_server_name=10.0.0.2 \
+  --metadata_server=P2PHANDSHAKE \
+  --gpu_id=0
+
+# 终端 2：发起端（Initiator）
+./transfer_engine_bench \
+  --mode=initiator \
+  --protocol=sunrise_link \
+  --metadata_server=P2PHANDSHAKE \
+  --segment_id=10.0.0.2:$PORT \
+  --gpu_id=0 \
+  --block_size=8388608 \
+  --batch_size=32
+```
+
+> 说明：当 `metadata_server=P2PHANDSHAKE` 时，目标端实际监听端口可能为动态分配端口。请以目标端日志中打印的实际端口替换 `--segment_id` 中的 `$PORT`。
+
+### 使用 `tebench`（TENT 后端）
+
+角色由 `--target_seg_name` 决定：为空 → 目标端，否则 → 发起端。
+
+```bash
+# 终端 1：目标端
+./tebench --backend=tent --metadata_type=p2p --xport_type=sunrise_link \
+  --seg_type=VRAM --local_gpu_id=0
+
+# 终端 2：发起端
+./tebench --backend=tent --metadata_type=p2p --xport_type=sunrise_link \
+  --seg_type=VRAM --target_seg_name=192.168.172.52:168 \
+  --local_gpu_id=1 --target_gpu_id=0
+```
+
+---
+
+## 环境变量
+
+| 变量 | 描述 | 默认值 |
+|------|------|--------|
+| `MC_TANGRT_ROOT` | Tang Runtime 根目录 | `/usr/local/tangrt` |
+| `MC_TANGRT_LIB_DIR` | `.so` 目录路径，优先于 `MC_TANGRT_ROOT` + 架构查找 | （未设置） |
+| `MC_TANGRT_LIB_ARCH` | `<root>/lib/` 下的架构子目录 | `linux-x86_64` |
+
+---
+
+## 配置项说明（TENT）
+
+Sunrise Link 支持通过配置文件调整行为，常用项包括：
+
+| 选项 | 类型 | 描述 | 默认值 |
+|------|------|------|--------|
+| `transports/sunrise_link/enable` | bool | 是否启用该传输 | `true` |
+| `transports/sunrise_link/async_memcpy_threshold` | int | 异步 memcpy 最小传输大小（MiB）。设为 0 禁用。 | `0` |
+| `transports/sunrise_link/enable_port_striping` | bool | 为大型跨设备拷贝（>4 MiB）启用多端口条带化 | `false` |
+
+可根据业务负载与设备拓扑微调上述参数。
+
+---
+
+## 注意事项
+
+1. **设备上下文**：内部管理 `tangSetDevice`。直接调用 Tang API 时，需在查询指针属性前设置正确设备。
+2. **IPC 句柄生命周期**：IPC 句柄缓存复用，在 `uninstall()` 中释放。
+3. **并发 Peer Copy**：每设备互斥锁保护上下文操作。跨设备拷贝按固定顺序（较低设备 ID 优先）获取源/目标锁以防死锁。
+4. **映射主机指针规范化**：`tangHostAlloc` 的主机指针若有对应设备指针，自动规范化以选择正确的 peer copy 路径。
diff --git a/docs/source/zh_archive/transfer-engine-python.md b/docs/source/zh_archive/transfer-engine-python.md
index 18a98189..371545c1 100644
--- a/docs/source/zh_archive/transfer-engine-python.md
+++ b/docs/source/zh_archive/transfer-engine-python.md
@@ -186,7 +186,7 @@ transfer_submit_write(target_hostname, buffer, peer_buffer_address, length)
 - `length` (int): 要传输的字节数
 
 **返回值：**
-- `int`: 用于跟踪操作的批次ID，失败时返回负值
+- `int`: 用于跟踪操作的批次ID，失败时返回0
 
 #### transfer_check_status()
 
@@ -287,7 +287,9 @@ TransferOpcode.WRITE  # 写操作
 - `MC_METADATA_SERVER`: 默认元数据服务器地址
 - `MC_LEGACY_RPC_PORT_BINDING`: 启用传统RPC端口绑定行为
 - `MC_TCP_BIND_ADDRESS`: 指定TCP绑定地址
+- `MC_RDMA_BIND_ADDRESS`: 指定RDMA绑定地址，用于双网卡环境下的NIC路径构建。设置后，RDMA NIC路径使用此地址，而TCP握手使用 `local_hostname` 中的地址。适用于TCP和RDMA流量使用不同网络接口的场景（例如 `eth0` 用于TCP，`rdma-net1` 用于RDMA）。
 - `MC_CUSTOM_TOPO_JSON`: 自定义拓扑JSON文件路径
+- `MC_TE_FILTERS`: 可选，逗号分隔的 IB 设备名白名单（如 `mlx5_0,mlx5_2`），用于 Legacy Transfer Engine 拓扑自动发现；未设置时不限制。
 - `MC_TE_METRIC`: 启用指标报告（设置为"1"、"true"、"yes"或"on"）
 - `MC_TE_METRIC_INTERVAL_SECONDS`: 设置指标报告间隔（秒）
 
diff --git a/docs/source/zh_archive/transfer-engine.md b/docs/source/zh_archive/transfer-engine.md
index 20c0de97..180fa44f 100644
--- a/docs/source/zh_archive/transfer-engine.md
+++ b/docs/source/zh_archive/transfer-engine.md
@@ -396,6 +396,7 @@ int init(const std::string &metadata_conn_string,
 - `MC_IB_TC` 当使用`RDMA`通信协议时，在交换机和网卡默认配置不一致场景/需要流量规划场景下，可能需要修改 RDMA 网卡的 Traffic Class 配置，默认值 -1
 - `MC_IB_PCI_RELAXED_ORDERING` 将网络适配器的PCIe顺序设置为放宽有时会带来更好的性能。可设置 1 以启用RO功能，默认值 0
 - `MC_GID_INDEX` 每个设备实例使用的 GID 序号，默认值 3（或平台支持的最大值）
+- `MC_PKEY_INDEX` QP 转换到 INIT 状态时使用的 `pkey_index`（partition key 表索引）。有效范围：0 到 65535，默认值 0。当 fabric 所需的 partition key 不在 HCA pkey 表的 0 号位置时需要设置该值
 - `MC_MAX_CQE_PER_CTX` 每个设备实例中 CQ 缓冲区大小，默认值 4096
 - `MC_MAX_EP_PER_CTX` 每个设备实例中活跃 EndPoint 数量上限，默认值 65536。**注意：** 小于 0.3.7.post1 的版本，这里默认值是 256，但是无法手动设置为 65536，最大值支持 65535！请注意
 - `MC_NUM_QP_PER_EP` 每个 EndPoint 中 QP 数量，数量越多则细粒度 I/O 性能越好，默认值 2
@@ -406,6 +407,7 @@ int init(const std::string &metadata_conn_string,
 - `MC_WORKERS_PER_CTX` 每个设备实例对应的异步工作线程数量
 - `MC_SLICE_SIZE` Transfer Engine 中用户请求的切分粒度
 - `MC_RETRY_CNT` Transfer Engine 中最大重试次数
+- `MC_AUTO_GID_MAX_RETRIES` classic RDMA 握手恢复过程中自动重探测本地 GID 的最大重试次数，默认值 2。设置为 0 可以关闭自动 GID 重试。
 - `MC_LOG_LEVEL` 该选项可以设置成`TRACE`/`INFO`/`WARNING`/`ERROR`（详情见 [glog doc](https://github.com/google/glog/blob/master/docs/logging.md)），则在运行时会输出更详细的日志
 - `MC_HANDSHAKE_LISTEN_BACKLOG` 监听握手连接的 backlog 大小, 默认值 128
 - `MC_HANDSHAKE_MAX_LENGTH` P2P 模式下握手消息的最大长度（字节）。有效范围：1MB 到 128MB。默认值为 1MB (1048576 字节)。当单个 RDMA 实例注册大量内存缓冲区（>10,000）时，需要增大此值以避免握手失败。示例：设置为 10485760 表示 10MB
@@ -418,7 +420,9 @@ int init(const std::string &metadata_conn_string,
 - `MC_FORCE_MNNVL` 强制使用 Multi-Node NVLink 作为主要传输方式，无论是否安装了有效的 RDMA 网卡
 - `MC_INTRA_NVLINK` 指定使用Intra-Node NVLink 作为主要传输方式，同时注意该设置不能与MC_FORCE_MNNVL一起使用
 - `MC_FORCE_TCP` 强制使用 TCP 作为主要传输方式，无论是否安装了有效的 RDMA 网卡
-- `MC_MIN_PRC_PORT` 指定 RPC 服务使用的最小端口号。默认值为 15000。
-- `MC_MAX_PRC_PORT` 指定 RPC 服务使用的最大端口号。默认值为 17000。
+- `MC_MIN_RPC_PORT` 指定 RPC 服务使用的最小端口号。默认值为 15000。
+- `MC_MAX_RPC_PORT` 指定 RPC 服务使用的最大端口号。默认值为 17000。
 - `MC_PATH_ROUNDROBIN` 指定 RDMA 路径选择使用 Round Robin 模式，这对于传输大块数据可能有利。
+- `MC_TE_FILTERS` 可选，逗号分隔的 IB 设备名白名单（如 `mlx5_0,mlx5_2`），用于 Legacy Transfer Engine 拓扑自动发现时限制可见 RDMA 网卡；未设置时不限制。
+- `WITH_NVIDIA_PEERMEM` 设置为 `1`、`ON` 或 `TRUE` 时，Mooncake 使用 `ibv_reg_mr()` 直接注册 GPU 内存（需要 `nvidia-peermem` 内核模块）。默认情况下（未设置或为 `0`），Mooncake 使用不需要 `nvidia-peermem` 的 DMA-BUF 路径。
 - `MC_ENDPOINT_STORE_TYPE` 选择 FIFO Endpoint Store (`FIFO`) 或者 Sieve Endpoint Store (`SIEVE`)，模式是 `SIEVE`。
diff --git a/extern/yalantinglibs b/extern/yalantinglibs
index 73dea196..6a0e067d 160000
--- a/extern/yalantinglibs
+++ b/extern/yalantinglibs
@@ -1 +1 @@
-Subproject commit 73dea196d23ad8fcd4914c6ef1238f390b9a1c48
+Subproject commit 6a0e067d9a43492cf8e4e280b531924fbd724dbd
diff --git a/image/components.png b/image/components.png
index 576cbb4d..e5d05413 100644
Binary files a/image/components.png and b/image/components.png differ
diff --git a/image/hardwares/MetaX_logo.png b/image/hardwares/MetaX_logo.png
new file mode 100755
index 00000000..307e2344
Binary files /dev/null and b/image/hardwares/MetaX_logo.png differ
diff --git a/image/hardwares/T-Head_logo.png b/image/hardwares/T-Head_logo.png
new file mode 100755
index 00000000..0d001c9e
Binary files /dev/null and b/image/hardwares/T-Head_logo.png differ
diff --git a/image/hardwares/cambricon_logo.png b/image/hardwares/cambricon_logo.png
new file mode 100755
index 00000000..56588e74
Binary files /dev/null and b/image/hardwares/cambricon_logo.png differ
diff --git a/image/partners/aws-logo.png b/image/partners/aws-logo.png
new file mode 100755
index 00000000..bf4d497d
Binary files /dev/null and b/image/partners/aws-logo.png differ
diff --git a/image/partners/hygon_logo.png b/image/partners/hygon_logo.png
new file mode 100644
index 00000000..4d418152
Binary files /dev/null and b/image/partners/hygon_logo.png differ
diff --git a/image/partners/nvidia_logo.png b/image/partners/nvidia_logo.png
index 5d51569e..70ec4de2 100644
Binary files a/image/partners/nvidia_logo.png and b/image/partners/nvidia_logo.png differ
diff --git a/image/partners/sunrise_logo.png b/image/partners/sunrise_logo.png
new file mode 100644
index 00000000..c4b1f921
Binary files /dev/null and b/image/partners/sunrise_logo.png differ
diff --git a/mooncake-common/FindUrma.cmake b/mooncake-common/FindUrma.cmake
index d2d93cc3..0af8d1a7 100644
--- a/mooncake-common/FindUrma.cmake
+++ b/mooncake-common/FindUrma.cmake
@@ -4,7 +4,7 @@ include(FetchContent)
 FetchContent_Declare(
         urma
         GIT_REPOSITORY https://atomgit.com/openeuler/umdk.git
-        GIT_TAG        v25.12.0
+        GIT_TAG        v25.12.0.B081
 )
 
 FetchContent_MakeAvailable(urma)
diff --git a/mooncake-common/SetupPyTorchEnv.cmake b/mooncake-common/SetupPyTorchEnv.cmake
index 526a190a..0a00d1b7 100644
--- a/mooncake-common/SetupPyTorchEnv.cmake
+++ b/mooncake-common/SetupPyTorchEnv.cmake
@@ -23,6 +23,18 @@ function(install_pytorch_wheel _version _cuda_major _cuda_minor _module_prefix)
     # 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.12.0")
+    # PyTorch 2.12.0+ no longer publishes cu128 wheels.
+    # Use cu126 for all CUDA 12 builds.
+    if(_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()
+
   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.
diff --git a/mooncake-common/common.cmake b/mooncake-common/common.cmake
index 93bd9841..2426a6eb 100644
--- a/mooncake-common/common.cmake
+++ b/mooncake-common/common.cmake
@@ -1,12 +1,22 @@
 set(CMAKE_C_STANDARD 99)
 set(CMAKE_CXX_STANDARD 20)
+set(CMAKE_CUDA_STANDARD 20)
 
-set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -g -Wall -Wextra -Wno-unused-parameter -fPIC")
-set(CMAKE_C_FLAGS "${CMAKE_C_FLAGS} -g -Wall -Wextra -Wno-unused-parameter -fPIC")
+option(ENABLE_DEBUG_SYMBOLS "Include debug symbols (-g) in compilation" ON)
+
+set(CMAKE_CXX_FLAGS
+    "${CMAKE_CXX_FLAGS} -Wall -Wextra -Wno-unused-parameter -fPIC")
+set(CMAKE_C_FLAGS "${CMAKE_C_FLAGS} -Wall -Wextra -Wno-unused-parameter -fPIC")
+
+if(ENABLE_DEBUG_SYMBOLS)
+  set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -g")
+  set(CMAKE_C_FLAGS "${CMAKE_C_FLAGS} -g")
+endif()
 
 if(CMAKE_CXX_COMPILER_ID STREQUAL "GNU")
   set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -fcoroutines")
-  set(CMAKE_CXX_FLAGS_RELEASE "${CMAKE_CXX_FLAGS_RELEASE} -fno-tree-slp-vectorize")
+  set(CMAKE_CXX_FLAGS_RELEASE
+      "${CMAKE_CXX_FLAGS_RELEASE} -fno-tree-slp-vectorize")
 endif()
 
 set(CMAKE_C_FLAGS_RELEASE "-O3")
@@ -21,7 +31,7 @@ endif()
 
 option(ENABLE_ASAN "enable address sanitizer" OFF)
 
-if (ENABLE_ASAN)
+if(ENABLE_ASAN)
   set(CMAKE_C_FLAGS "${CMAKE_C_FLAGS} -fsanitize=leak")
   set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -fsanitize=leak")
   set(CMAKE_C_FLAGS "${CMAKE_C_FLAGS} -fsanitize=address")
@@ -29,7 +39,7 @@ if (ENABLE_ASAN)
 endif()
 
 # keep debuginfo by default
-if (NOT CMAKE_BUILD_TYPE)
+if(NOT CMAKE_BUILD_TYPE)
   set(CMAKE_BUILD_TYPE "RelWithDebInfo")
 endif()
 
@@ -44,7 +54,7 @@ set(CMAKE_EXPORT_COMPILE_COMMANDS ON)
 include(${CMAKE_CURRENT_LIST_DIR}/limit_jobs.cmake)
 
 option(ENABLE_SCCACHE "Whether to open sccache" OFF)
-if (ENABLE_SCCACHE)
+if(ENABLE_SCCACHE)
   find_program(SCCACHE sccache REQUIRED)
 endif()
 if(SCCACHE AND ENABLE_SCCACHE)
@@ -59,40 +69,51 @@ add_compile_options(-fno-tree-slp-vectorize)
 option(BUILD_EXAMPLES "Build examples" ON)
 
 option(BUILD_UNIT_TESTS "Build unit tests" ON)
+option(BUILD_BENCHMARK "Build benchmarks" 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_HYGON "option for enabling gpu features for Hygon DCU with DTK" OFF)
+option(USE_COREX "option for enabling gpu features for Iluvatar CoreX" OFF)
 option(USE_NVMEOF "option for using NVMe over Fabric" OFF)
 option(USE_TCP "option for using TCP transport" ON)
 option(USE_BAREX "option for using accl-barex transport" OFF)
 option(USE_ASCEND "option for using npu with HCCL" OFF)
 option(USE_ASCEND_DIRECT "option for using ascend npu with adxl engine" OFF)
 option(USE_UBSHMEM "option for using ascend npu with shmem" OFF)
-option(USE_ASCEND_HETEROGENEOUS "option for transferring between ascend npu and gpu" OFF)
+option(USE_ASCEND_HETEROGENEOUS
+       "option for transferring between ascend npu and gpu" OFF)
 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)
+option(USE_SUNRISE
+       "option for enabling gpu features for Sunrise GPU with Tang runtime" OFF)
 
-if (USE_UB)
+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)
+if(USE_EFA)
   # Find libfabric headers and library; default to AWS EFA installer path
-  find_path(LIBFABRIC_INCLUDE_DIR rdma/fabric.h
+  find_path(
+    LIBFABRIC_INCLUDE_DIR rdma/fabric.h
     HINTS /opt/amazon/efa/include
     PATH_SUFFIXES include)
-  find_library(LIBFABRIC_LIBRARY fabric
+  find_library(
+    LIBFABRIC_LIBRARY fabric
     HINTS /opt/amazon/efa/lib
     PATH_SUFFIXES lib lib64)
 
-  if (NOT LIBFABRIC_INCLUDE_DIR OR NOT LIBFABRIC_LIBRARY)
-    message(FATAL_ERROR "libfabric not found. Install AWS EFA or set LIBFABRIC_INCLUDE_DIR/LIBFABRIC_LIBRARY.")
+  if(NOT LIBFABRIC_INCLUDE_DIR OR NOT LIBFABRIC_LIBRARY)
+    message(
+      FATAL_ERROR
+        "libfabric not found. Install AWS EFA or set LIBFABRIC_INCLUDE_DIR/LIBFABRIC_LIBRARY."
+    )
   endif()
 
   get_filename_component(LIBFABRIC_LIB_DIR ${LIBFABRIC_LIBRARY} DIRECTORY)
@@ -103,126 +124,231 @@ if (USE_EFA)
   message(STATUS "  libfabric include: ${LIBFABRIC_INCLUDE_DIR}")
   message(STATUS "  libfabric library: ${LIBFABRIC_LIBRARY}")
 endif()
+if(USE_CXI)
+  # Find libfabric headers and library; default to AWS EFA installer path
+  find_path(LIBFABRIC_INCLUDE_DIR rdma/fabric.h PATH_SUFFIXES include)
+  find_library(LIBFABRIC_LIBRARY fabric PATH_SUFFIXES lib lib64)
+
+  if(NOT LIBFABRIC_INCLUDE_DIR OR NOT LIBFABRIC_LIBRARY)
+    message(
+      FATAL_ERROR
+        "libfabric not found. Install AWS EFA or set LIBFABRIC_INCLUDE_DIR/LIBFABRIC_LIBRARY."
+    )
+  endif()
+
+  get_filename_component(LIBFABRIC_LIB_DIR ${LIBFABRIC_LIBRARY} DIRECTORY)
+  include_directories(${LIBFABRIC_INCLUDE_DIR})
+  link_directories(${LIBFABRIC_LIB_DIR})
+  add_compile_definitions(USE_CXI)
+  message(STATUS "HPE CXI (libfabric) transport is enabled")
+  message(STATUS "  libfabric include: ${LIBFABRIC_INCLUDE_DIR}")
+  message(STATUS "  libfabric library: ${LIBFABRIC_LIBRARY}")
+endif()
 option(USE_ETCD "option for enable etcd as metadata server" OFF)
 option(USE_ETCD_LEGACY "option for enable etcd based on etcd-cpp-api-v3" OFF)
 option(USE_REDIS "option for enable redis as metadata server" OFF)
 option(USE_HTTP "option for enable http as metadata server" ON)
-option(WITH_RUST_EXAMPLE "build the Rust interface and sample code for the transfer engine" OFF)
+option(WITH_RUST_EXAMPLE
+       "build the Rust interface and sample code for the transfer engine" OFF)
 option(WITH_METRICS "enable metrics and metrics reporting thread" ON)
 option(USE_3FS "option for using 3FS storage backend" OFF)
-option(WITH_NVIDIA_PEERMEM "disable to support RDMA without nvidia-peermem. If WITH_NVIDIA_PEERMEM=OFF then USE_CUDA=ON is required." ON)
-option(USE_EVENT_DRIVEN_COMPLETION "option for using event-driven completion (store & transfer engine)" OFF)
+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")
+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)
+option(USE_MLX5DV
+       "enable mlx5 direct verbs (libmlx5) for QP UDP source port override" OFF)
 set(LRU_MAX_CAPACITY 1000)
 
-if (USE_LRU_MASTER)
+if(USE_LRU_MASTER)
   add_compile_definitions(USE_LRU_MASTER)
   add_compile_definitions(LRU_MAX_CAPACITY)
 endif()
 
-if (USE_EVENT_DRIVEN_COMPLETION)
+if(USE_EVENT_DRIVEN_COMPLETION)
   add_compile_definitions(USE_EVENT_DRIVEN_COMPLETION)
   message(STATUS "Event-driven completion is enabled")
 else()
   message(STATUS "Event-driven completion is disabled")
 endif()
 
-if (USE_NVMEOF)
+if(USE_NVMEOF)
   set(USE_CUDA ON)
   add_compile_definitions(USE_NVMEOF)
   message(STATUS "NVMe-oF support is enabled")
 endif()
 
-if (USE_MNNVL)
-  if (NOT USE_HIP AND NOT USE_MUSA AND NOT USE_MACA)
+if(USE_MNNVL)
+  if(NOT USE_HIP
+     AND NOT USE_MUSA
+     AND NOT USE_MACA)
     set(USE_CUDA ON)
   endif()
   add_compile_definitions(USE_MNNVL)
   message(STATUS "Multi-Node NVLink support is enabled")
 endif()
 
-if (USE_CUDA)
+if(USE_CUDA)
   add_compile_definitions(USE_CUDA)
   message(STATUS "CUDA support is enabled")
   include_directories(/usr/local/cuda/include)
-  link_directories(
-    /usr/local/cuda/lib
-    /usr/local/cuda/lib64
-  )
+  link_directories(/usr/local/cuda/lib /usr/local/cuda/lib64)
 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)
+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)
+    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 "")
+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 "")
+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)
+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)
+    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 "")
+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")
+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)
+if(USE_MLU)
   add_compile_definitions(USE_MLU)
   message(STATUS "MLU support is enabled")
   include_directories(${MLU_INCLUDE_DIR})
-  if (EXISTS "${MLU_LIB_DIR}")
+  if(EXISTS "${MLU_LIB_DIR}")
     link_directories(${MLU_LIB_DIR})
   endif()
 endif()
 
-if (USE_MACA)
+if(USE_MACA)
   add_compile_definitions(USE_MACA)
   message(STATUS "MACA support is enabled")
   include_directories(${MACA_INCLUDE_DIR})
-  if (EXISTS "${MACA_LIB_DIR}")
+  if(EXISTS "${MACA_LIB_DIR}")
     link_directories(${MACA_LIB_DIR})
   endif()
 endif()
 
-if (USE_MUSA)
+if(NOT DEFINED MC_TANGRT_ROOT OR MC_TANGRT_ROOT STREQUAL "")
+  if(DEFINED ENV{MC_TANGRT_ROOT} AND NOT "$ENV{MC_TANGRT_ROOT}" STREQUAL "")
+    set(MC_TANGRT_ROOT
+        "$ENV{MC_TANGRT_ROOT}"
+        CACHE PATH "Path to Tang runtime root" FORCE)
+  else()
+    set(MC_TANGRT_ROOT
+        "/usr/local/tangrt"
+        CACHE PATH "Path to Tang runtime root" FORCE)
+  endif()
+endif()
+
+if(USE_SUNRISE)
+  add_compile_definitions(USE_SUNRISE)
+  message(STATUS "Sunrise (Tang runtime) support is enabled")
+  include_directories(${MC_TANGRT_ROOT}/include)
+endif()
+
+if(USE_MUSA)
   add_compile_definitions(USE_MUSA)
   message(STATUS "MUSA support is enabled")
   include_directories(/usr/local/musa/include)
-  link_directories(
-    /usr/local/musa/lib
-  )
+  link_directories(/usr/local/musa/lib)
 endif()
 
-if (USE_HIP)
+if(USE_HYGON)
+  if(NOT DEFINED DTK_ROOT OR DTK_ROOT STREQUAL "")
+    if(DEFINED ENV{DTK_HOME} AND NOT "$ENV{DTK_HOME}" STREQUAL "")
+      set(DTK_ROOT
+          "$ENV{DTK_HOME}"
+          CACHE PATH "Path to Hygon DTK SDK" FORCE)
+    else()
+      set(DTK_ROOT
+          "/opt/dtk"
+          CACHE PATH "Path to Hygon DTK SDK" FORCE)
+    endif()
+  endif()
+
+  if(NOT DEFINED DTK_INCLUDE_DIR OR DTK_INCLUDE_DIR STREQUAL "")
+    set(DTK_INCLUDE_DIR "${DTK_ROOT}/cuda/cuda-11/include")
+  endif()
+
+  if(NOT DEFINED DTK_LIB_DIR OR DTK_LIB_DIR STREQUAL "")
+    set(DTK_LIB_DIR "${DTK_ROOT}/cuda/cuda-11/lib64")
+  endif()
+
+  add_compile_definitions(USE_HYGON)
+  message(STATUS "Hygon DCU/DTK support is enabled")
+  include_directories(${DTK_INCLUDE_DIR})
+  if(EXISTS "${DTK_LIB_DIR}")
+    link_directories(${DTK_LIB_DIR})
+  endif()
+endif()
+
+if(USE_COREX)
+  if(NOT DEFINED COREX_ROOT OR COREX_ROOT STREQUAL "")
+    if(DEFINED ENV{COREX_HOME} AND NOT "$ENV{COREX_HOME}" STREQUAL "")
+      set(COREX_ROOT
+          "$ENV{COREX_HOME}"
+          CACHE PATH "Path to Iluvatar CoreX SDK" FORCE)
+    else()
+      set(COREX_ROOT
+          "/usr/local/corex"
+          CACHE PATH "Path to Iluvatar CoreX SDK" FORCE)
+    endif()
+  endif()
+
+  if(NOT DEFINED COREX_INCLUDE_DIR OR COREX_INCLUDE_DIR STREQUAL "")
+    set(COREX_INCLUDE_DIR "${COREX_ROOT}/include")
+  endif()
+
+  if(NOT DEFINED COREX_LIB_DIR OR COREX_LIB_DIR STREQUAL "")
+    set(COREX_LIB_DIR "${COREX_ROOT}/lib")
+  endif()
+
+  add_compile_definitions(USE_COREX)
+  message(STATUS "Iluvatar CoreX support is enabled")
+  include_directories(${COREX_INCLUDE_DIR})
+  if(EXISTS "${COREX_LIB_DIR}")
+    link_directories(${COREX_LIB_DIR})
+  endif()
+endif()
+
+if(USE_HIP)
   list(APPEND CMAKE_PREFIX_PATH "/opt/rocm/lib/cmake")
   find_package(HIP REQUIRED)
   include_directories(${HIP_INCLUDE_DIRS})
@@ -231,52 +357,56 @@ if (USE_HIP)
 
   find_program(HIPIFY_PERL_EXECUTABLE hipify-perl)
   if(NOT HIPIFY_PERL_EXECUTABLE)
-    message(FATAL_ERROR
-            "hipify-perl not found.\n"
-            "Please ensure the ROCm or HIP SDK is installed and in your PATH.")
+    message(
+      FATAL_ERROR
+        "hipify-perl not found.\n"
+        "Please ensure the ROCm or HIP SDK is installed and in your PATH.")
   endif()
 endif()
 
-# This function converts given CUDA source files into HIP-compatible
-# files using hipify-perl, placing the outputs in the build directory for use in
-# project compilation. The file path changes to a new location after hipify.
+# This function converts given CUDA source files into HIP-compatible files using
+# hipify-perl, placing the outputs in the build directory for use in project
+# compilation. The file path changes to a new location after hipify.
 function(hipify_files input_var_name)
-    set(result_files)
+  set(result_files)
 
-    foreach(input_file IN LISTS ${input_var_name})
-        file(RELATIVE_PATH rel_path ${CMAKE_SOURCE_DIR} ${input_file})
-        set(output_file "${CMAKE_BINARY_DIR}/${rel_path}")
+  foreach(input_file IN LISTS ${input_var_name})
+    file(RELATIVE_PATH rel_path ${CMAKE_SOURCE_DIR} ${input_file})
+    set(output_file "${CMAKE_BINARY_DIR}/${rel_path}")
 
-        get_filename_component(output_dir ${output_file} DIRECTORY)
-        file(MAKE_DIRECTORY ${output_dir})
+    get_filename_component(output_dir ${output_file} DIRECTORY)
+    file(MAKE_DIRECTORY ${output_dir})
 
-        add_custom_command(
-            OUTPUT ${output_file}
-            COMMAND ${HIPIFY_PERL_EXECUTABLE} ${input_file} > ${output_file}
-            DEPENDS ${input_file}
-            COMMENT "HIPifying ${input_file} → ${output_file}"
-        )
+    add_custom_command(
+      OUTPUT ${output_file}
+      COMMAND ${HIPIFY_PERL_EXECUTABLE} ${input_file} > ${output_file}
+      DEPENDS ${input_file}
+      COMMENT "HIPifying ${input_file} → ${output_file}")
 
-        list(APPEND result_files ${output_file})
-    endforeach()
+    list(APPEND result_files ${output_file})
+  endforeach()
 
-    set(${input_var_name} ${result_files} PARENT_SCOPE)
+  set(${input_var_name}
+      ${result_files}
+      PARENT_SCOPE)
 endfunction()
 
-if (USE_CXL)
+if(USE_CXL)
   add_compile_definitions(USE_CXL)
   message(STATUS "CXL support is enabled")
 endif()
 
-if (USE_TCP)
+if(USE_TCP)
   add_compile_definitions(USE_TCP)
 endif()
 
-if (USE_BAREX)
+if(USE_BAREX)
   add_compile_definitions(USE_BAREX)
 endif()
 
-if (USE_ASCEND OR USE_ASCEND_DIRECT OR USE_UBSHMEM)
+if(USE_ASCEND
+   OR USE_ASCEND_DIRECT
+   OR USE_UBSHMEM)
   set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -DOPEN_BUILD_PROJECT ")
   set(CMAKE_C_FLAGS "${CMAKE_C_FLAGS} -DOPEN_BUILD_PROJECT ")
   string(TOLOWER "${CMAKE_SYSTEM_PROCESSOR}" CURRENT_CPU)
@@ -292,7 +422,8 @@ if (USE_ASCEND OR USE_ASCEND_DIRECT OR USE_UBSHMEM)
     message(STATUS "Use env ASCEND_HOME_PATH")
     file(GLOB ASCEND_TOOLKIT_ROOT "$ENV{ASCEND_HOME_PATH}/${CPU_ARCH}-linux")
   else()
-    file(GLOB ASCEND_TOOLKIT_ROOT "/usr/local/Ascend/ascend-toolkit/latest/${CPU_ARCH}-linux")
+    file(GLOB ASCEND_TOOLKIT_ROOT
+         "/usr/local/Ascend/ascend-toolkit/latest/${CPU_ARCH}-linux")
   endif()
   set(ASCEND_LIB_DIR "${ASCEND_TOOLKIT_ROOT}/lib64")
   set(ASCEND_INCLUDE_DIR "${ASCEND_TOOLKIT_ROOT}/include")
@@ -301,24 +432,25 @@ if (USE_ASCEND OR USE_ASCEND_DIRECT OR USE_UBSHMEM)
   link_directories(${ASCEND_LIB_DIR})
 endif()
 
-if (USE_ASCEND)
+if(USE_ASCEND)
   set(ASCEND_DEVLIB_DIR "${ASCEND_TOOLKIT_ROOT}/devlib")
   link_directories(${ASCEND_DEVLIB_DIR})
   add_compile_definitions(USE_ASCEND)
 endif()
 
-if (USE_ASCEND_DIRECT)
+if(USE_ASCEND_DIRECT)
   set(BUILD_SHARED_LIBS ON)
   add_compile_definitions(USE_ASCEND_DIRECT)
 endif()
 
-if (USE_UBSHMEM)
+if(USE_UBSHMEM)
   set(BUILD_SHARED_LIBS ON)
   add_compile_definitions(USE_UBSHMEM)
 endif()
 
-if (USE_ASCEND_HETEROGENEOUS)
-  file(GLOB ASCEND_TOOLKIT_ROOT "/usr/local/Ascend/ascend-toolkit/latest/*-linux")
+if(USE_ASCEND_HETEROGENEOUS)
+  file(GLOB ASCEND_TOOLKIT_ROOT
+       "/usr/local/Ascend/ascend-toolkit/latest/*-linux")
   set(ASCEND_LIB_DIR "${ASCEND_TOOLKIT_ROOT}/lib64")
   set(ASCEND_INCLUDE_DIR "${ASCEND_TOOLKIT_ROOT}/include")
   add_compile_definitions(USE_ASCEND_HETEROGENEOUS)
@@ -326,21 +458,26 @@ if (USE_ASCEND_HETEROGENEOUS)
   link_directories(${ASCEND_LIB_DIR})
 endif()
 
-if (USE_REDIS)
+if(USE_REDIS)
   add_compile_definitions(USE_REDIS)
   message(STATUS "Redis as metadata server support is enabled")
 endif()
 
-if (USE_HTTP)
+if(USE_HTTP)
   add_compile_definitions(USE_HTTP)
   message(STATUS "Http as metadata server support is enabled")
 endif()
 
-if (NOT USE_ETCD AND NOT USE_REDIS AND NOT USE_HTTP)
-  message(STATUS "None of USE_ETCD, USE_REDIS, USE_HTTP is selected, only \"P2PHANDSHAKE\" is supported as metadata server")
+if(NOT USE_ETCD
+   AND NOT USE_REDIS
+   AND NOT USE_HTTP)
+  message(
+    STATUS
+      "None of USE_ETCD, USE_REDIS, USE_HTTP is selected, only \"P2PHANDSHAKE\" is supported as metadata server"
+  )
 endif()
 
-if (WITH_METRICS)
+if(WITH_METRICS)
   add_compile_definitions(WITH_METRICS)
   message(STATUS "metrics is enabled")
 endif()
@@ -350,12 +487,23 @@ if(USE_3FS)
   message(STATUS "3FS storage backend is enabled")
 endif()
 
-if(WITH_NVIDIA_PEERMEM)
-  add_compile_definitions(WITH_NVIDIA_PEERMEM)
+if(EXISTS "/usr/include/boost1.78")
+  include_directories(SYSTEM "/usr/include/boost1.78")
+  link_directories("/usr/lib64/boost1.78")
 endif()
 
 set(GFLAGS_USE_TARGET_NAMESPACE "true")
 find_package(yaml-cpp REQUIRED)
 find_package(gflags REQUIRED)
+if(NOT TARGET gflags::gflags)
+  foreach(_gflags_target gflags-shared gflags_shared gflags)
+    if(TARGET ${_gflags_target})
+      add_library(gflags::gflags INTERFACE IMPORTED)
+      set_target_properties(gflags::gflags PROPERTIES INTERFACE_LINK_LIBRARIES
+                                                      ${_gflags_target})
+      break()
+    endif()
+  endforeach()
+endif()
 find_package(yalantinglibs CONFIG REQUIRED)
 add_compile_definitions(YLT_ENABLE_IBV)
diff --git a/mooncake-common/etcd/etcd_wrapper.go b/mooncake-common/etcd/etcd_wrapper.go
index 8a0eef49..6c1695d3 100644
--- a/mooncake-common/etcd/etcd_wrapper.go
+++ b/mooncake-common/etcd/etcd_wrapper.go
@@ -50,6 +50,10 @@ type prefixWatchInfo struct {
 	callbackContext unsafe.Pointer
 	// done is closed when the watch goroutine fully exits (no more callbacks).
 	done chan struct{}
+	// broken means this watch ended because reset/error made it untrustworthy.
+	broken bool
+	// brokenNotified prevents duplicate WATCH_BROKEN callbacks from exit races.
+	brokenNotified bool
 }
 
 // Use different etcd client so they are not affected by each other,
@@ -82,6 +86,39 @@ const (
 	snapshotTimeout    = 60 * time.Second   // 1 minute for large files
 )
 
+const (
+	storeDialKeepAliveTime    = 10 * time.Second
+	storeDialKeepAliveTimeout = 3 * time.Second
+)
+
+func newStoreClientConfig(validEndpoints []string) clientv3.Config {
+	return clientv3.Config{
+		Endpoints:            validEndpoints,
+		DialTimeout:          5 * time.Second,
+		DialKeepAliveTime:    storeDialKeepAliveTime,
+		DialKeepAliveTimeout: storeDialKeepAliveTimeout,
+		PermitWithoutStream:  true,
+	}
+}
+
+func parseEtcdEndpoints(endpoints *C.char) []string {
+	if endpoints == nil {
+		return nil
+	}
+	endpointStr := C.GoString(endpoints)
+	endpointStr = strings.ReplaceAll(endpointStr, ",", ";")
+	endpointList := strings.Split(endpointStr, ";")
+
+	var validEndpoints []string
+	for _, ep := range endpointList {
+		ep = strings.TrimSpace(ep)
+		if ep != "" {
+			validEndpoints = append(validEndpoints, ep)
+		}
+	}
+	return validEndpoints
+}
+
 //export NewEtcdClient
 func NewEtcdClient(endpoints *C.char, errMsg **C.char) int {
 	globalMutex.Lock()
@@ -197,6 +234,12 @@ func EtcdCloseWrapper() {
 	}
 }
 
+func getStoreClient() *clientv3.Client {
+	storeMutex.Lock()
+	defer storeMutex.Unlock()
+	return storeClient
+}
+
 //export NewStoreEtcdClient
 func NewStoreEtcdClient(endpoints *C.char, errMsg **C.char) int {
 	storeMutex.Lock()
@@ -206,29 +249,13 @@ func NewStoreEtcdClient(endpoints *C.char, errMsg **C.char) int {
 		return -2
 	}
 
-	endpointStr := C.GoString(endpoints)
-	// Support multiple endpoints separated by comma or semicolon.
-	endpointStr = strings.ReplaceAll(endpointStr, ",", ";")
-	endpointList := strings.Split(endpointStr, ";")
-
-	// Filter out any empty strings that might result from splitting
-	var validEndpoints []string
-	for _, ep := range endpointList {
-		ep = strings.TrimSpace(ep)
-		if ep != "" {
-			validEndpoints = append(validEndpoints, ep)
-		}
-	}
-
+	validEndpoints := parseEtcdEndpoints(endpoints)
 	if len(validEndpoints) == 0 {
 		*errMsg = C.CString("no valid endpoints provided")
 		return -1
 	}
 
-	cli, err := clientv3.New(clientv3.Config{
-		Endpoints:   validEndpoints,
-		DialTimeout: 5 * time.Second,
-	})
+	cli, err := clientv3.New(newStoreClientConfig(validEndpoints))
 
 	if err != nil {
 		*errMsg = C.CString(err.Error())
@@ -239,6 +266,36 @@ func NewStoreEtcdClient(endpoints *C.char, errMsg **C.char) int {
 	return 0
 }
 
+//export EtcdStoreResetClientWrapper
+func EtcdStoreResetClientWrapper(endpoints *C.char, errMsg **C.char) int {
+	validEndpoints := parseEtcdEndpoints(endpoints)
+	if len(validEndpoints) == 0 {
+		*errMsg = C.CString("no valid endpoints provided")
+		return -1
+	}
+
+	cli, err := clientv3.New(newStoreClientConfig(validEndpoints))
+	if err != nil {
+		*errMsg = C.CString(err.Error())
+		return -1
+	}
+
+	cancelAllStoreKeepAlives()
+	cancelAllStoreWatches()
+	cancelAllStorePrefixWatches()
+
+	storeMutex.Lock()
+	oldClient := storeClient
+	storeClient = cli
+	storeMutex.Unlock()
+
+	if oldClient != nil {
+		oldClient.Close()
+	}
+
+	return 0
+}
+
 //export NewSnapshotEtcdClient
 func NewSnapshotEtcdClient(endpoints *C.char, errMsg **C.char) int {
 	snapshotMutex.Lock()
@@ -286,14 +343,15 @@ func NewSnapshotEtcdClient(endpoints *C.char, errMsg **C.char) int {
 //export EtcdStoreGetWrapper
 func EtcdStoreGetWrapper(key *C.char, keySize C.int, value **C.char,
 	valueSize *C.int, revisionId *int64, errMsg **C.char) int {
-	if storeClient == nil {
+	cli := getStoreClient()
+	if cli == nil {
 		*errMsg = C.CString("etcd client not initialized")
 		return -1
 	}
 	k := C.GoStringN(key, keySize)
 	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
 	defer cancel()
-	resp, err := storeClient.Get(ctx, k)
+	resp, err := cli.Get(ctx, k)
 	if err != nil {
 		*errMsg = C.CString(err.Error())
 		return -1
@@ -312,13 +370,14 @@ func EtcdStoreGetWrapper(key *C.char, keySize C.int, value **C.char,
 
 //export EtcdStoreGrantLeaseWrapper
 func EtcdStoreGrantLeaseWrapper(ttl int64, leaseId *int64, errMsg **C.char) int {
-	if storeClient == nil {
+	cli := getStoreClient()
+	if cli == nil {
 		*errMsg = C.CString("etcd client not initialized")
 		return -1
 	}
 	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
 	defer cancel()
-	resp, err := storeClient.Grant(ctx, ttl)
+	resp, err := cli.Grant(ctx, ttl)
 	if err != nil {
 		*errMsg = C.CString(err.Error())
 		return -1
@@ -329,13 +388,14 @@ func EtcdStoreGrantLeaseWrapper(ttl int64, leaseId *int64, errMsg **C.char) int
 
 //export EtcdStoreRevokeLeaseWrapper
 func EtcdStoreRevokeLeaseWrapper(leaseId int64, errMsg **C.char) int {
-	if storeClient == nil {
+	cli := getStoreClient()
+	if cli == 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))
+	_, err := cli.Revoke(ctx, clientv3.LeaseID(leaseId))
 	if err != nil {
 		if errors.Is(err, rpctypes.ErrLeaseNotFound) {
 			return 0
@@ -349,7 +409,8 @@ func EtcdStoreRevokeLeaseWrapper(leaseId int64, errMsg **C.char) int {
 //export EtcdStoreCreateWithLeaseWrapper
 func EtcdStoreCreateWithLeaseWrapper(key *C.char, keySize C.int, value *C.char, valueSize C.int,
 	leaseId int64, revisionId *int64, errMsg **C.char) int {
-	if storeClient == nil {
+	cli := getStoreClient()
+	if cli == nil {
 		*errMsg = C.CString("etcd client not initialized")
 		return -1
 	}
@@ -359,7 +420,7 @@ func EtcdStoreCreateWithLeaseWrapper(key *C.char, keySize C.int, value *C.char,
 	defer cancel()
 
 	// Create a transaction
-	txn := storeClient.Txn(ctx)
+	txn := cli.Txn(ctx)
 
 	// Only put the key if it does not exist
 	resp, err := txn.If(clientv3.Compare(clientv3.CreateRevision(k), "=", 0)).
@@ -401,9 +462,81 @@ func cancelAndDeleteWatch(k string) int {
 	return -1
 }
 
+func cancelAllStoreKeepAlives() {
+	storeKeepAliveMutex.Lock()
+	cancels := make([]context.CancelFunc, 0, len(storeKeepAliveCtx))
+	for leaseId, cancel := range storeKeepAliveCtx {
+		cancels = append(cancels, cancel)
+		delete(storeKeepAliveCtx, leaseId)
+	}
+	storeKeepAliveMutex.Unlock()
+
+	for _, cancel := range cancels {
+		cancel()
+	}
+}
+
+func cancelAllStoreWatches() {
+	storeWatchMutex.Lock()
+	cancels := make([]context.CancelFunc, 0, len(storeWatchCtx))
+	for key, cancel := range storeWatchCtx {
+		cancels = append(cancels, cancel)
+		delete(storeWatchCtx, key)
+	}
+	storeWatchMutex.Unlock()
+
+	for _, cancel := range cancels {
+		cancel()
+	}
+}
+
+func cancelAllStorePrefixWatches() {
+	storePrefixWatchMutex.Lock()
+	cancels := make([]context.CancelFunc, 0, len(storePrefixWatchCtx))
+	for prefix, watchInfo := range storePrefixWatchCtx {
+		watchInfo.broken = true
+		storePrefixWatchCtx[prefix] = watchInfo
+		cancels = append(cancels, watchInfo.cancel)
+		// Do not delete map entries here; let goroutine defer clean up.
+	}
+	storePrefixWatchMutex.Unlock()
+
+	for _, cancel := range cancels {
+		cancel()
+	}
+}
+
+// notifyStorePrefixWatchBrokenOnce delivers a WATCH_BROKEN callback for the
+// prefix watch `p` at most once. If markBroken is true the watch is first
+// marked broken (used by reset/error exit paths); explicit cancel passes
+// markBroken=false so a normal shutdown does not generate a spurious
+// WATCH_BROKEN. The cgo callback is invoked outside storePrefixWatchMutex
+// to avoid deadlocks from cgo re-entry while holding a Go mutex.
+func notifyStorePrefixWatchBrokenOnce(p string, callbackContext unsafe.Pointer, callbackFunc unsafe.Pointer, markBroken bool) {
+	shouldNotify := false
+
+	storePrefixWatchMutex.Lock()
+	if watchInfo, exists := storePrefixWatchCtx[p]; exists {
+		if markBroken {
+			watchInfo.broken = true
+		}
+		if watchInfo.broken && !watchInfo.brokenNotified {
+			watchInfo.brokenNotified = true
+			shouldNotify = true
+		}
+		storePrefixWatchCtx[p] = watchInfo
+	}
+	storePrefixWatchMutex.Unlock()
+
+	if shouldNotify {
+		C.call_watch_cb(callbackFunc, callbackContext, nil, 0, nil, 0, C.int(2) /*WATCH_BROKEN*/, C.longlong(0))
+	}
+}
+
 //export EtcdStoreWatchUntilDeletedWrapper
 func EtcdStoreWatchUntilDeletedWrapper(key *C.char, keySize C.int, errMsg **C.char) int {
-	if storeClient == nil {
+	cli := getStoreClient()
+	if cli == nil {
 		*errMsg = C.CString("etcd client not initialized")
 		return -1
 	}
@@ -426,7 +559,7 @@ func EtcdStoreWatchUntilDeletedWrapper(key *C.char, keySize C.int, errMsg **C.ch
 	defer cancelAndDeleteWatch(k)
 
 	// Start watching the key
-	watchChan := storeClient.Watch(ctx, k)
+	watchChan := cli.Watch(ctx, k)
 
 	// Wait for the key to be deleted
 	for {
@@ -490,7 +623,8 @@ func hasKeepAliveContext(leaseId int64) bool {
 
 //export EtcdStoreKeepAliveWrapper
 func EtcdStoreKeepAliveWrapper(leaseId int64, errMsg **C.char) int {
-	if storeClient == nil {
+	cli := getStoreClient()
+	if cli == nil {
 		*errMsg = C.CString("etcd client not initialized")
 		return -1
 	}
@@ -511,7 +645,7 @@ func EtcdStoreKeepAliveWrapper(leaseId int64, errMsg **C.char) int {
 	defer cancelAndDeleteKeepAlive(leaseId)
 
 	// Start keep alive
-	keepAliveChan, err := storeClient.KeepAlive(ctx, clientv3.LeaseID(leaseId))
+	keepAliveChan, err := cli.KeepAlive(ctx, clientv3.LeaseID(leaseId))
 	if err != nil {
 		*errMsg = C.CString(err.Error())
 		return -1
@@ -564,7 +698,8 @@ func EtcdStoreWaitKeepAliveReadyWrapper(leaseId int64, timeoutMs int, errMsg **C
 
 //export EtcdStorePutWrapper
 func EtcdStorePutWrapper(key *C.char, keySize C.int, value *C.char, valueSize C.int, errMsg **C.char) int {
-	if storeClient == nil {
+	cli := getStoreClient()
+	if cli == nil {
 		*errMsg = C.CString("etcd client not initialized")
 		return -1
 	}
@@ -572,7 +707,7 @@ func EtcdStorePutWrapper(key *C.char, keySize C.int, value *C.char, valueSize C.
 	v := C.GoStringN(value, valueSize)
 	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
 	defer cancel()
-	_, err := storeClient.Put(ctx, k, v)
+	_, err := cli.Put(ctx, k, v)
 	if err != nil {
 		*errMsg = C.CString(err.Error())
 		return -1
@@ -588,7 +723,8 @@ func EtcdStorePutWrapper(key *C.char, keySize C.int, value *C.char, valueSize C.
 //
 //export EtcdStoreCreateWrapper
 func EtcdStoreCreateWrapper(key *C.char, keySize C.int, value *C.char, valueSize C.int, errMsg **C.char) int {
-	if storeClient == nil {
+	cli := getStoreClient()
+	if cli == nil {
 		*errMsg = C.CString("etcd client not initialized")
 		return -1
 	}
@@ -597,7 +733,7 @@ func EtcdStoreCreateWrapper(key *C.char, keySize C.int, value *C.char, valueSize
 	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
 	defer cancel()
 
-	txn := storeClient.Txn(ctx)
+	txn := cli.Txn(ctx)
 	resp, err := txn.If(clientv3.Compare(clientv3.CreateRevision(k), "=", 0)).
 		Then(clientv3.OpPut(k, v)).
 		Commit()
@@ -614,7 +750,8 @@ func EtcdStoreCreateWrapper(key *C.char, keySize C.int, value *C.char, valueSize
 
 //export EtcdStoreBatchCreateWrapper
 func EtcdStoreBatchCreateWrapper(keys **C.char, values **C.char, count C.int, errMsg **C.char) int {
-	if storeClient == nil {
+	cli := getStoreClient()
+	if cli == nil {
 		*errMsg = C.CString("etcd client not initialized")
 		return -1
 	}
@@ -642,7 +779,7 @@ func EtcdStoreBatchCreateWrapper(keys **C.char, values **C.char, count C.int, er
 	defer cancel()
 
 	// Use Txn to ensure atomicity of the batch
-	resp, err := storeClient.Txn(ctx).If(cmps...).Then(ops...).Commit()
+	resp, err := cli.Txn(ctx).If(cmps...).Then(ops...).Commit()
 	if err != nil {
 		*errMsg = C.CString(err.Error())
 		return -1
@@ -657,14 +794,15 @@ func EtcdStoreBatchCreateWrapper(keys **C.char, values **C.char, count C.int, er
 
 //export EtcdStoreGetWithPrefixWrapper
 func EtcdStoreGetWithPrefixWrapper(prefix *C.char, prefixSize C.int, keys **C.char, keySizes **C.int, values **C.char, valueSizes **C.int, count *C.int, errMsg **C.char) int {
-	if storeClient == nil {
+	cli := getStoreClient()
+	if cli == nil {
 		*errMsg = C.CString("etcd client not initialized")
 		return -1
 	}
 	p := C.GoStringN(prefix, prefixSize)
 	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
 	defer cancel()
-	resp, err := storeClient.Get(ctx, p, clientv3.WithPrefix(), clientv3.WithSort(clientv3.SortByKey, clientv3.SortAscend))
+	resp, err := cli.Get(ctx, p, clientv3.WithPrefix(), clientv3.WithSort(clientv3.SortByKey, clientv3.SortAscend))
 	if err != nil {
 		*errMsg = C.CString(err.Error())
 		return -1
@@ -702,7 +840,8 @@ func EtcdStoreGetWithPrefixWrapper(prefix *C.char, prefixSize C.int, keys **C.ch
 
 //export EtcdStoreGetRangeAsJsonWrapper
 func EtcdStoreGetRangeAsJsonWrapper(startKey *C.char, startKeySize C.int, endKey *C.char, endKeySize C.int, limit C.int, outJson **C.char, outJsonSize *C.int, revisionId *C.longlong, errMsg **C.char) int {
-	if storeClient == nil {
+	cli := getStoreClient()
+	if cli == nil {
 		*errMsg = C.CString("etcd client not initialized")
 		return -1
 	}
@@ -719,7 +858,7 @@ func EtcdStoreGetRangeAsJsonWrapper(startKey *C.char, startKeySize C.int, endKey
 	if limit > 0 {
 		opts = append(opts, clientv3.WithLimit(int64(limit)))
 	}
-	resp, err := storeClient.Get(ctx, start, opts...)
+	resp, err := cli.Get(ctx, start, opts...)
 	if err != nil {
 		*errMsg = C.CString(err.Error())
 		return -1
@@ -752,14 +891,15 @@ func EtcdStoreGetRangeAsJsonWrapper(startKey *C.char, startKeySize C.int, endKey
 
 //export EtcdStoreGetFirstKeyWithPrefixWrapper
 func EtcdStoreGetFirstKeyWithPrefixWrapper(prefix *C.char, prefixSize C.int, firstKey **C.char, firstKeySize *C.int, errMsg **C.char) int {
-	if storeClient == nil {
+	cli := getStoreClient()
+	if cli == nil {
 		*errMsg = C.CString("etcd client not initialized")
 		return -1
 	}
 	p := C.GoStringN(prefix, prefixSize)
 	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
 	defer cancel()
-	resp, err := storeClient.Get(ctx, p, clientv3.WithPrefix(), clientv3.WithSort(clientv3.SortByKey, clientv3.SortAscend), clientv3.WithLimit(1))
+	resp, err := cli.Get(ctx, p, clientv3.WithPrefix(), clientv3.WithSort(clientv3.SortByKey, clientv3.SortAscend), clientv3.WithLimit(1))
 	if err != nil {
 		*errMsg = C.CString(err.Error())
 		return -1
@@ -776,7 +916,8 @@ func EtcdStoreGetFirstKeyWithPrefixWrapper(prefix *C.char, prefixSize C.int, fir
 
 //export EtcdStoreGetLastKeyWithPrefixWrapper
 func EtcdStoreGetLastKeyWithPrefixWrapper(prefix *C.char, prefixSize C.int, lastKey **C.char, lastKeySize *C.int, errMsg **C.char) int {
-	if storeClient == nil {
+	cli := getStoreClient()
+	if cli == nil {
 		*errMsg = C.CString("etcd client not initialized")
 		return -1
 	}
@@ -784,7 +925,7 @@ func EtcdStoreGetLastKeyWithPrefixWrapper(prefix *C.char, prefixSize C.int, last
 	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
 	defer cancel()
 
-	resp, err := storeClient.Get(
+	resp, err := cli.Get(
 		ctx, p,
 		clientv3.WithPrefix(),
 		clientv3.WithSort(clientv3.SortByKey, clientv3.SortDescend),
@@ -806,7 +947,8 @@ func EtcdStoreGetLastKeyWithPrefixWrapper(prefix *C.char, prefixSize C.int, last
 
 //export EtcdStoreDeleteRangeWrapper
 func EtcdStoreDeleteRangeWrapper(startKey *C.char, startKeySize C.int, endKey *C.char, endKeySize C.int, errMsg **C.char) int {
-	if storeClient == nil {
+	cli := getStoreClient()
+	if cli == nil {
 		*errMsg = C.CString("etcd client not initialized")
 		return -1
 	}
@@ -814,7 +956,7 @@ func EtcdStoreDeleteRangeWrapper(startKey *C.char, startKeySize C.int, endKey *C
 	end := C.GoStringN(endKey, endKeySize)
 	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
 	defer cancel()
-	_, err := storeClient.Delete(ctx, start, clientv3.WithRange(end))
+	_, err := cli.Delete(ctx, start, clientv3.WithRange(end))
 	if err != nil {
 		*errMsg = C.CString(err.Error())
 		return -1
@@ -824,7 +966,8 @@ func EtcdStoreDeleteRangeWrapper(startKey *C.char, startKeySize C.int, endKey *C
 
 //export EtcdStoreWatchWithPrefixFromRevisionWrapper
 func EtcdStoreWatchWithPrefixFromRevisionWrapper(prefix *C.char, prefixSize C.int, startRevision C.longlong, callbackContext unsafe.Pointer, callbackFunc unsafe.Pointer, errMsg **C.char) int {
-	if storeClient == nil {
+	cli := getStoreClient()
+	if cli == nil {
 		*errMsg = C.CString("etcd client not initialized")
 		return -1
 	}
@@ -844,27 +987,53 @@ func EtcdStoreWatchWithPrefixFromRevisionWrapper(prefix *C.char, prefixSize C.in
 		return -1
 	}
 	doneCh := make(chan struct{})
+	// createdCh reports whether the server-side watch was successfully
+	// established. The goroutine sends exactly one value: true once etcd
+	// confirms the watch (clientv3.WithCreatedNotify), or false if the
+	// goroutine exits before that confirmation. The wrapper blocks on it before
+	// returning, so callers can rely on "watch is live server-side" the moment
+	// WatchWithPrefixFromRevision returns OK. This closes the race where the
+	// goroutine had not yet issued storeClient.Watch() when the caller
+	// proceeded to read the current view. Buffered (size 1) so the goroutine
+	// never blocks on the send even if the wrapper has already timed out.
+	createdCh := make(chan bool, 1)
 	storePrefixWatchCtx[p] = prefixWatchInfo{
 		cancel:          cancel,
 		callbackContext: callbackContext,
 		done:            doneCh,
+		broken:          false,
+		brokenNotified:  false,
 	}
 	storePrefixWatchMutex.Unlock()
 
-	go func(doneCh chan struct{}) {
+	go func(doneCh chan struct{}, createdCh chan bool) {
+		// createdSignalled guards against sending on createdCh more than once.
+		createdSignalled := false
+		signalCreated := func(ok bool) {
+			if !createdSignalled {
+				createdSignalled = true
+				createdCh <- ok
+			}
+		}
 		defer func() {
-			// Remove watch entry and signal completion
+			// Report failure if the watch was never confirmed (e.g. the
+			// goroutine exits before any response), then remove the watch
+			// entry and signal completion. Ordering matters: the failure
+			// signal is sent before `done` is closed.
+			signalCreated(false)
 			storePrefixWatchMutex.Lock()
 			delete(storePrefixWatchCtx, p)
 			storePrefixWatchMutex.Unlock()
 			close(doneCh)
 		}()
 
-		opts := []clientv3.OpOption{clientv3.WithPrefix()}
+		// WithCreatedNotify makes etcd send an initial response with
+		// Created == true as soon as the watch is registered server-side.
+		opts := []clientv3.OpOption{clientv3.WithPrefix(), clientv3.WithCreatedNotify()}
 		if startRevision > 0 {
 			opts = append(opts, clientv3.WithRev(int64(startRevision)))
 		}
-		watchChan := storeClient.Watch(ctx, p, opts...)
+		watchChan := cli.Watch(ctx, p, opts...)
 
 		for {
 			select {
@@ -873,20 +1042,27 @@ func EtcdStoreWatchWithPrefixFromRevisionWrapper(prefix *C.char, prefixSize C.in
 					// Channel closed. Check if context was cancelled.
 					select {
 					case <-ctx.Done():
+						notifyStorePrefixWatchBrokenOnce(p, callbackContext, callbackFunc, false)
 						return
 					default:
 						// Channel closed unexpectedly (not cancelled). Notify C++ watcher to reconnect.
-						C.call_watch_cb(callbackFunc, callbackContext, nil, 0, nil, 0, C.int(2) /*WATCH_BROKEN*/, C.longlong(0))
+						notifyStorePrefixWatchBrokenOnce(p, callbackContext, callbackFunc, true)
 						return
 					}
 				}
+				// The first response with Created == true confirms the
+				// server-side watch is established; release the waiter.
+				if watchResp.Created {
+					signalCreated(true)
+				}
 				if watchResp.Err() != nil {
 					// Watch error. Check if context was cancelled.
 					select {
 					case <-ctx.Done():
+						notifyStorePrefixWatchBrokenOnce(p, callbackContext, callbackFunc, false)
 						return
 					default:
-						C.call_watch_cb(callbackFunc, callbackContext, nil, 0, nil, 0, C.int(2) /*WATCH_BROKEN*/, C.longlong(0))
+						notifyStorePrefixWatchBrokenOnce(p, callbackContext, callbackFunc, true)
 						return
 					}
 				}
@@ -900,6 +1076,7 @@ func EtcdStoreWatchWithPrefixFromRevisionWrapper(prefix *C.char, prefixSize C.in
 				for _, event := range watchResp.Events {
 					select {
 					case <-ctx.Done():
+						notifyStorePrefixWatchBrokenOnce(p, callbackContext, callbackFunc, false)
 						return
 					default:
 					}
@@ -942,12 +1119,39 @@ func EtcdStoreWatchWithPrefixFromRevisionWrapper(prefix *C.char, prefixSize C.in
 					}
 				}
 			case <-ctx.Done():
+				notifyStorePrefixWatchBrokenOnce(p, callbackContext, callbackFunc, false)
 				return
 			}
 		}
-	}(doneCh)
+	}(doneCh, createdCh)
 
-	return 0
+	// Block until the server confirms the watch is established. Returning OK
+	// then guarantees the watch is live server-side, so a caller that arms the
+	// watch before reading state will not miss an event in between.
+	const watchCreatedTimeout = 5 * time.Second
+	select {
+	case created := <-createdCh:
+		if !created {
+			// The goroutine exited before the watch was confirmed (e.g. the
+			// Watch RPC failed). It has already torn itself down; wait for it
+			// to finish so the prefix entry is gone and no callback is in
+			// flight, then report failure.
+			cancel()
+			<-doneCh
+			*errMsg = C.CString("etcd watch goroutine exited before the watch was established")
+			return -1
+		}
+		return 0
+	case <-time.After(watchCreatedTimeout):
+		// The watch was not established within the timeout. Tear down the
+		// goroutine and report failure so the caller can fall back. Wait for
+		// the goroutine to fully exit so the prefix entry is removed and no
+		// callback can be in flight when we return.
+		cancel()
+		<-doneCh
+		*errMsg = C.CString("timeout waiting for etcd watch to be created")
+		return -1
+	}
 }
 
 func cancelAndDeletePrefixWatch(p string) int {
diff --git a/mooncake-common/etcd/go.mod b/mooncake-common/etcd/go.mod
index d1507b5d..ef5b3503 100644
--- a/mooncake-common/etcd/go.mod
+++ b/mooncake-common/etcd/go.mod
@@ -2,7 +2,7 @@ module github.com/kvcache-ai/Mooncake/mooncake-common/etcd
 
 go 1.25.0
 
-toolchain go1.25.9
+toolchain go1.25.10
 
 require (
 	go.etcd.io/etcd/api/v3 v3.5.21
@@ -18,7 +18,7 @@ require (
 	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.48.0 // indirect
+	golang.org/x/net v0.55.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
diff --git a/mooncake-common/include/environ.h b/mooncake-common/include/environ.h
index c5a72c8e..78f8f450 100644
--- a/mooncake-common/include/environ.h
+++ b/mooncake-common/include/environ.h
@@ -40,8 +40,8 @@ class Environ {
         return enable_dest_device_affinity_;
     }
     bool GetUseIpv6() const { return use_ipv6_; }
-    int GetMinPrcPort() const { return min_prc_port_; }
-    int GetMaxPrcPort() const { return max_prc_port_; }
+    int GetMinRpcPort() const { return min_rpc_port_; }
+    int GetMaxRpcPort() const { return max_rpc_port_; }
     int GetEnableParallelRegMr() const { return enable_parallel_reg_mr_; }
     std::string GetEndpointStoreType() const { return endpoint_store_type_; }
     bool GetForceTcp() const { return force_tcp_; }
@@ -49,9 +49,8 @@ class Environ {
     bool GetForceMnnvl() const { return force_mnnvl_; }
     bool GetIntraNvlink() const { return intra_nvlink_; }
     bool GetPathRoundrobin() const { return path_roundrobin_; }
-
-   private:
-    Environ();
+    bool GetWithNvidiaPeermem() const { return with_nvidia_peermem_; }
+    int GetEfaCqThreads() const { return efa_cq_threads_; }
 
     // Helper method to get int from env
     static int GetInt(const char* name, int default_value);
@@ -63,6 +62,9 @@ class Environ {
     static std::string GetString(const char* name,
                                  const std::string& default_value);
 
+   private:
+    Environ();
+
     // Member variables
     int num_cq_per_ctx_;
     int num_comp_channels_per_ctx_;
@@ -90,8 +92,8 @@ class Environ {
     int fragment_ratio_;
     bool enable_dest_device_affinity_;
     bool use_ipv6_;
-    int min_prc_port_;
-    int max_prc_port_;
+    int min_rpc_port_;
+    int max_rpc_port_;
     int enable_parallel_reg_mr_;
     std::string endpoint_store_type_;
     bool force_tcp_;
@@ -99,6 +101,8 @@ class Environ {
     bool force_mnnvl_;
     bool intra_nvlink_;
     bool path_roundrobin_;
+    bool with_nvidia_peermem_;
+    int efa_cq_threads_;
 };
 
 }  // namespace mooncake
diff --git a/mooncake-common/k8s-lease/CMakeLists.txt b/mooncake-common/k8s-lease/CMakeLists.txt
index 7ccc2ac6..ed555f17 100644
--- a/mooncake-common/k8s-lease/CMakeLists.txt
+++ b/mooncake-common/k8s-lease/CMakeLists.txt
@@ -1,5 +1,6 @@
 add_custom_command(
     OUTPUT ${CMAKE_CURRENT_BINARY_DIR}/libk8s_lease_wrapper.so
+           ${CMAKE_CURRENT_BINARY_DIR}/libk8s_lease_wrapper.h
     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"
diff --git a/mooncake-common/k8s-lease/k8s_lease_wrapper.go b/mooncake-common/k8s-lease/k8s_lease_wrapper.go
index 72f90731..b67caca5 100644
--- a/mooncake-common/k8s-lease/k8s_lease_wrapper.go
+++ b/mooncake-common/k8s-lease/k8s_lease_wrapper.go
@@ -20,6 +20,7 @@ import "C"
 
 import (
 	"context"
+	"encoding/json"
 	"fmt"
 	"os"
 	"sync"
@@ -29,6 +30,7 @@ import (
 	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/types"
 	"k8s.io/apimachinery/pkg/watch"
 	"k8s.io/client-go/kubernetes"
 	"k8s.io/client-go/rest"
@@ -486,4 +488,62 @@ func K8sLeaseCancelWatch(
 	return 0
 }
 
+// patchPodLabel sets or removes a label on a pod using a JSON merge patch.
+// If value is non-nil, the label is set; if nil, the label is removed.
+func patchPodLabel(namespace, podName, labelKey string, value interface{}) error {
+	if err := ensureClientInitialized(); err != nil {
+		return err
+	}
+
+	patch := map[string]interface{}{
+		"metadata": map[string]interface{}{
+			"labels": map[string]interface{}{
+				labelKey: value,
+			},
+		},
+	}
+	patchBytes, err := json.Marshal(patch)
+	if err != nil {
+		return fmt.Errorf("failed to marshal label patch: %w", err)
+	}
+
+	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+	defer cancel()
+
+	_, err = globalClient.CoreV1().Pods(namespace).Patch(
+		ctx, podName, types.MergePatchType, patchBytes, metav1.PatchOptions{})
+	if err != nil {
+		return fmt.Errorf("failed to patch pod label: %w", err)
+	}
+	return nil
+}
+
+//export K8sPatchPodLabel
+func K8sPatchPodLabel(
+	ns, podName, labelKey, labelValue *C.char,
+	errMsg **C.char,
+) C.int {
+	err := patchPodLabel(C.GoString(ns), C.GoString(podName),
+		C.GoString(labelKey), C.GoString(labelValue))
+	if err != nil {
+		*errMsg = C.CString(err.Error())
+		return -1
+	}
+	return 0
+}
+
+//export K8sRemovePodLabel
+func K8sRemovePodLabel(
+	ns, podName, labelKey *C.char,
+	errMsg **C.char,
+) C.int {
+	err := patchPodLabel(C.GoString(ns), C.GoString(podName),
+		C.GoString(labelKey), nil)
+	if err != nil {
+		*errMsg = C.CString(err.Error())
+		return -1
+	}
+	return 0
+}
+
 func main() {}
diff --git a/mooncake-common/k8s-lease/k8s_lease_wrapper_test.go b/mooncake-common/k8s-lease/k8s_lease_wrapper_test.go
index 705b579a..47b50e02 100644
--- a/mooncake-common/k8s-lease/k8s_lease_wrapper_test.go
+++ b/mooncake-common/k8s-lease/k8s_lease_wrapper_test.go
@@ -7,6 +7,7 @@ import (
 	"time"
 
 	coordinationv1 "k8s.io/api/coordination/v1"
+	corev1 "k8s.io/api/core/v1"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"k8s.io/client-go/kubernetes"
 	"k8s.io/client-go/kubernetes/fake"
@@ -22,6 +23,15 @@ func swapClient(newClient kubernetes.Interface) kubernetes.Interface {
 	return old
 }
 
+// swapInitClientFn replaces initClientFn and returns the old one.
+func swapInitClientFn(newFn func() error) func() error {
+	clientMutex.Lock()
+	defer clientMutex.Unlock()
+	old := initClientFn
+	initClientFn = newFn
+	return old
+}
+
 // TestGetHolderWithFakeClient tests getHolder using a fake K8s clientset.
 func TestGetHolderWithFakeClient(t *testing.T) {
 	holderID := "node-1:8080"
@@ -53,6 +63,61 @@ func TestGetHolderWithFakeClient(t *testing.T) {
 	}
 }
 
+func TestPatchPodLabelAutoInitializesClient(t *testing.T) {
+	fakeClient := fake.NewSimpleClientset(&corev1.Pod{
+		ObjectMeta: metav1.ObjectMeta{
+			Name:      "test-pod",
+			Namespace: "default",
+		},
+	})
+
+	oldClient := swapClient(nil)
+	defer swapClient(oldClient)
+
+	initCalls := 0
+	oldInitFn := swapInitClientFn(func() error {
+		initCalls++
+		clientMutex.Lock()
+		globalClient = fakeClient
+		clientMutex.Unlock()
+		return nil
+	})
+	defer swapInitClientFn(oldInitFn)
+
+	const labelKey = "mooncake.io/store-role"
+	if err := patchPodLabel("default", "test-pod", labelKey, "leader"); err != nil {
+		t.Fatalf("patchPodLabel(set) failed: %v", err)
+	}
+	if initCalls != 1 {
+		t.Fatalf("initClientFn calls = %d, want 1", initCalls)
+	}
+
+	pod, err := fakeClient.CoreV1().Pods("default").Get(
+		context.Background(), "test-pod", metav1.GetOptions{})
+	if err != nil {
+		t.Fatalf("get pod after label set failed: %v", err)
+	}
+	if got := pod.Labels[labelKey]; got != "leader" {
+		t.Fatalf("label %q = %q, want %q", labelKey, got, "leader")
+	}
+
+	if err := patchPodLabel("default", "test-pod", labelKey, nil); err != nil {
+		t.Fatalf("patchPodLabel(remove) failed: %v", err)
+	}
+	if initCalls != 1 {
+		t.Fatalf("initClientFn calls after second patch = %d, want 1", initCalls)
+	}
+
+	pod, err = fakeClient.CoreV1().Pods("default").Get(
+		context.Background(), "test-pod", metav1.GetOptions{})
+	if err != nil {
+		t.Fatalf("get pod after label removal failed: %v", err)
+	}
+	if _, exists := pod.Labels[labelKey]; exists {
+		t.Fatalf("label %q still present after removal", labelKey)
+	}
+}
+
 // TestGetHolderNotFound tests getHolder when the Lease does not exist.
 func TestGetHolderNotFound(t *testing.T) {
 	fakeClient := fake.NewSimpleClientset()
diff --git a/mooncake-common/src/default_config.cpp b/mooncake-common/src/default_config.cpp
index 65d7eeea..266e1430 100644
--- a/mooncake-common/src/default_config.cpp
+++ b/mooncake-common/src/default_config.cpp
@@ -28,9 +28,9 @@ void DefaultConfig::Load() {
         return path.size() >= ext.size() &&
                path.compare(path.size() - ext.size(), ext.size(), ext) == 0;
     };
-    if (check_extension(path_, ".yaml")) {
+    if (check_extension(path_, ".yaml") || check_extension(path_, ".yml")) {
         loadFromYAML();
-        type_ = ConfigType::YAML;  // YAML type
+        type_ = ConfigType::YAML;
     } else if (check_extension(path_, ".json")) {
         loadFromJSON();
         type_ = ConfigType::JSON;  // JSON type
diff --git a/mooncake-common/src/environ.cpp b/mooncake-common/src/environ.cpp
index 28cf5fb7..46aa1e9a 100644
--- a/mooncake-common/src/environ.cpp
+++ b/mooncake-common/src/environ.cpp
@@ -1,6 +1,10 @@
 #include "environ.h"
+#include <cerrno>
+#include <climits>
 #include <cstring>
 #include <algorithm>
+#include <cctype>
+#include <iostream>
 
 namespace mooncake {
 
@@ -12,7 +16,17 @@ Environ& Environ::Get() {
 int Environ::GetInt(const char* name, int default_value) {
     const char* val = std::getenv(name);
     if (val) {
-        return std::atoi(val);
+        char* endptr = nullptr;
+        errno = 0;
+        long result = std::strtol(val, &endptr, 10);
+        if (endptr == val || *endptr != '\0' || errno == ERANGE ||
+            result < INT_MIN || result > INT_MAX) {
+            std::cerr << "[Mooncake] Warning: invalid value '" << val
+                      << "' for env " << name << ", using default "
+                      << default_value << std::endl;
+            return default_value;
+        }
+        return static_cast<int>(result);
     }
     return default_value;
 }
@@ -20,7 +34,17 @@ int Environ::GetInt(const char* name, int default_value) {
 size_t Environ::GetSizeT(const char* name, size_t default_value) {
     const char* val = std::getenv(name);
     if (val) {
-        return static_cast<size_t>(std::strtoull(val, nullptr, 10));
+        char* endptr = nullptr;
+        errno = 0;
+        long long result = std::strtoll(val, &endptr, 10);
+        if (endptr == val || *endptr != '\0' || errno == ERANGE || result < 0 ||
+            static_cast<unsigned long long>(result) > SIZE_MAX) {
+            std::cerr << "[Mooncake] Warning: invalid value '" << val
+                      << "' for env " << name << ", using default "
+                      << default_value << std::endl;
+            return default_value;
+        }
+        return static_cast<size_t>(result);
     }
     return default_value;
 }
@@ -29,7 +53,8 @@ bool Environ::GetBool(const char* name, bool default_value) {
     const char* val = std::getenv(name);
     if (val) {
         std::string s(val);
-        std::transform(s.begin(), s.end(), s.begin(), ::tolower);
+        std::transform(s.begin(), s.end(), s.begin(),
+                       [](unsigned char c) { return std::tolower(c); });
         return s == "1" || s == "true" || s == "on" || s == "yes";
     }
     return default_value;
@@ -72,8 +97,8 @@ Environ::Environ() {
     enable_dest_device_affinity_ =
         GetBool("MC_ENABLE_DEST_DEVICE_AFFINITY", false);
     use_ipv6_ = GetBool("MC_USE_IPV6", false);
-    min_prc_port_ = GetInt("MC_MIN_PRC_PORT", 15000);
-    max_prc_port_ = GetInt("MC_MAX_PRC_PORT", 17000);
+    min_rpc_port_ = GetInt("MC_MIN_RPC_PORT", GetInt("MC_MIN_PRC_PORT", 15000));
+    max_rpc_port_ = GetInt("MC_MAX_RPC_PORT", GetInt("MC_MAX_PRC_PORT", 17000));
     enable_parallel_reg_mr_ = GetInt("MC_ENABLE_PARALLEL_REG_MR", -1);
     endpoint_store_type_ = GetString("MC_ENDPOINT_STORE_TYPE", "SIEVE");
     force_tcp_ = GetBool("MC_FORCE_TCP", false);
@@ -81,6 +106,8 @@ Environ::Environ() {
     force_mnnvl_ = GetBool("MC_FORCE_MNNVL", false);
     intra_nvlink_ = GetBool("MC_INTRA_NVLINK", false);
     path_roundrobin_ = GetBool("MC_PATH_ROUNDROBIN", false);
+    with_nvidia_peermem_ = GetBool("WITH_NVIDIA_PEERMEM", true);
+    efa_cq_threads_ = GetInt("MC_EFA_CQ_THREADS", 1);
 }
 
 }  // namespace mooncake
diff --git a/mooncake-common/tests/CMakeLists.txt b/mooncake-common/tests/CMakeLists.txt
index 34f38ff1..ec43cfa2 100644
--- a/mooncake-common/tests/CMakeLists.txt
+++ b/mooncake-common/tests/CMakeLists.txt
@@ -1,7 +1,11 @@
 add_executable(default_config_test default_config_test.cpp)
 target_link_libraries(default_config_test PUBLIC
-    mooncake_common 
+    mooncake_common
     gtest
-    pthread PUBLIC
+    pthread
 )
 add_test(NAME default_config_test COMMAND default_config_test)
+
+add_executable(environ_test environ_test.cpp)
+target_link_libraries(environ_test PUBLIC mooncake_common gtest pthread)
+add_test(NAME environ_test COMMAND environ_test)
diff --git a/mooncake-common/tests/environ_test.cpp b/mooncake-common/tests/environ_test.cpp
new file mode 100644
index 00000000..9a0b5340
--- /dev/null
+++ b/mooncake-common/tests/environ_test.cpp
@@ -0,0 +1,185 @@
+// Copyright 2024 KVCache.AI
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "environ.h"
+
+#include <gtest/gtest.h>
+
+#include <climits>
+#include <cstdlib>
+
+using mooncake::Environ;
+
+class EnvironTest : public ::testing::Test {
+   protected:
+    void SetUp() override { clearTestEnvVars(); }
+    void TearDown() override { clearTestEnvVars(); }
+
+    void clearTestEnvVars() {
+        unsetenv("MC_TEST_INT");
+        unsetenv("MC_TEST_SIZET");
+        unsetenv("MC_TEST_BOOL");
+        unsetenv("MC_TEST_STRING");
+    }
+};
+
+// --- GetInt ---
+
+TEST_F(EnvironTest, GetIntValidValue) {
+    setenv("MC_TEST_INT", "42", 1);
+    EXPECT_EQ(Environ::GetInt("MC_TEST_INT", 0), 42);
+}
+
+TEST_F(EnvironTest, GetIntNegativeValue) {
+    setenv("MC_TEST_INT", "-100", 1);
+    EXPECT_EQ(Environ::GetInt("MC_TEST_INT", 0), -100);
+}
+
+TEST_F(EnvironTest, GetIntZero) {
+    setenv("MC_TEST_INT", "0", 1);
+    EXPECT_EQ(Environ::GetInt("MC_TEST_INT", 99), 0);
+}
+
+TEST_F(EnvironTest, GetIntMissing) {
+    EXPECT_EQ(Environ::GetInt("MC_TEST_INT", 77), 77);
+}
+
+TEST_F(EnvironTest, GetIntEmpty) {
+    setenv("MC_TEST_INT", "", 1);
+    EXPECT_EQ(Environ::GetInt("MC_TEST_INT", 55), 55);
+}
+
+TEST_F(EnvironTest, GetIntNonNumeric) {
+    setenv("MC_TEST_INT", "abc", 1);
+    EXPECT_EQ(Environ::GetInt("MC_TEST_INT", 55), 55);
+}
+
+TEST_F(EnvironTest, GetIntTrailingGarbage) {
+    setenv("MC_TEST_INT", "123abc", 1);
+    EXPECT_EQ(Environ::GetInt("MC_TEST_INT", 55), 55);
+}
+
+TEST_F(EnvironTest, GetIntOverflow) {
+    setenv("MC_TEST_INT", "99999999999999999999", 1);
+    EXPECT_EQ(Environ::GetInt("MC_TEST_INT", 55), 55);
+}
+
+TEST_F(EnvironTest, GetIntMaxValue) {
+    setenv("MC_TEST_INT", std::to_string(INT_MAX).c_str(), 1);
+    EXPECT_EQ(Environ::GetInt("MC_TEST_INT", 0), INT_MAX);
+}
+
+TEST_F(EnvironTest, GetIntMinValue) {
+    setenv("MC_TEST_INT", std::to_string(INT_MIN).c_str(), 1);
+    EXPECT_EQ(Environ::GetInt("MC_TEST_INT", 0), INT_MIN);
+}
+
+// --- GetSizeT ---
+
+TEST_F(EnvironTest, GetSizeTValidValue) {
+    setenv("MC_TEST_SIZET", "65536", 1);
+    EXPECT_EQ(Environ::GetSizeT("MC_TEST_SIZET", 0), 65536u);
+}
+
+TEST_F(EnvironTest, GetSizeTZero) {
+    setenv("MC_TEST_SIZET", "0", 1);
+    EXPECT_EQ(Environ::GetSizeT("MC_TEST_SIZET", 99), 0u);
+}
+
+TEST_F(EnvironTest, GetSizeTMissing) {
+    EXPECT_EQ(Environ::GetSizeT("MC_TEST_SIZET", 4096), 4096u);
+}
+
+TEST_F(EnvironTest, GetSizeTNonNumeric) {
+    setenv("MC_TEST_SIZET", "bogus", 1);
+    EXPECT_EQ(Environ::GetSizeT("MC_TEST_SIZET", 4096), 4096u);
+}
+
+TEST_F(EnvironTest, GetSizeTTrailingGarbage) {
+    setenv("MC_TEST_SIZET", "100MB", 1);
+    EXPECT_EQ(Environ::GetSizeT("MC_TEST_SIZET", 4096), 4096u);
+}
+
+TEST_F(EnvironTest, GetSizeTLargeValue) {
+    setenv("MC_TEST_SIZET", "1099511627776", 1);  // 1 TiB
+    EXPECT_EQ(Environ::GetSizeT("MC_TEST_SIZET", 0), 1099511627776ull);
+}
+
+TEST_F(EnvironTest, GetSizeTNegativeValue) {
+    setenv("MC_TEST_SIZET", "-1", 1);
+    EXPECT_EQ(Environ::GetSizeT("MC_TEST_SIZET", 4096), 4096u);
+}
+
+TEST_F(EnvironTest, GetSizeTNegativeWithLeadingSpace) {
+    setenv("MC_TEST_SIZET", " -1", 1);
+    EXPECT_EQ(Environ::GetSizeT("MC_TEST_SIZET", 4096), 4096u);
+}
+
+TEST_F(EnvironTest, GetSizeTOverflow) {
+    setenv("MC_TEST_SIZET", "99999999999999999999999999", 1);
+    EXPECT_EQ(Environ::GetSizeT("MC_TEST_SIZET", 4096), 4096u);
+}
+
+// --- GetBool ---
+
+TEST_F(EnvironTest, GetBoolTrue) {
+    for (const char* v :
+         {"1", "true", "TRUE", "True", "on", "ON", "yes", "YES"}) {
+        setenv("MC_TEST_BOOL", v, 1);
+        EXPECT_TRUE(Environ::GetBool("MC_TEST_BOOL", false)) << "for: " << v;
+    }
+}
+
+TEST_F(EnvironTest, GetBoolFalse) {
+    for (const char* v : {"0", "false", "FALSE", "off", "no", "whatever"}) {
+        setenv("MC_TEST_BOOL", v, 1);
+        EXPECT_FALSE(Environ::GetBool("MC_TEST_BOOL", false)) << "for: " << v;
+    }
+}
+
+TEST_F(EnvironTest, GetBoolMissing) {
+    EXPECT_TRUE(Environ::GetBool("MC_TEST_BOOL", true));
+    EXPECT_FALSE(Environ::GetBool("MC_TEST_BOOL", false));
+}
+
+TEST_F(EnvironTest, GetBoolEmpty) {
+    setenv("MC_TEST_BOOL", "", 1);
+    EXPECT_FALSE(Environ::GetBool("MC_TEST_BOOL", true));
+}
+
+// --- GetString ---
+
+TEST_F(EnvironTest, GetStringValidValue) {
+    setenv("MC_TEST_STRING", "hello", 1);
+    EXPECT_EQ(Environ::GetString("MC_TEST_STRING", "default"), "hello");
+}
+
+TEST_F(EnvironTest, GetStringMissing) {
+    EXPECT_EQ(Environ::GetString("MC_TEST_STRING", "default"), "default");
+}
+
+TEST_F(EnvironTest, GetStringEmpty) {
+    setenv("MC_TEST_STRING", "", 1);
+    EXPECT_EQ(Environ::GetString("MC_TEST_STRING", "default"), "");
+}
+
+TEST_F(EnvironTest, GetStringWithSpaces) {
+    setenv("MC_TEST_STRING", "hello world", 1);
+    EXPECT_EQ(Environ::GetString("MC_TEST_STRING", ""), "hello world");
+}
+
+int main(int argc, char** argv) {
+    ::testing::InitGoogleTest(&argc, argv);
+    return RUN_ALL_TESTS();
+}
diff --git a/mooncake-ep/BuildEpExt.cmake b/mooncake-ep/BuildEpExt.cmake
index d8355936..c94529cb 100644
--- a/mooncake-ep/BuildEpExt.cmake
+++ b/mooncake-ep/BuildEpExt.cmake
@@ -10,6 +10,7 @@
 #   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
+#   EP_USE_MUSA         - set to "1" when building for MUSA (MTLink path)
 
 cmake_minimum_required(VERSION 3.16)
 
@@ -34,6 +35,11 @@ endif()
 set(ENV{MAKEFLAGS} "")
 set(ENV{MFLAGS} "")
 set(ENV{TORCH_CUDA_ARCH_LIST} "${TORCH_CUDA_ARCH_LIST}")
+if(EP_USE_MUSA)
+  set(ENV{MOONCAKE_EP_USE_MUSA} "1")
+else()
+  unset(ENV{MOONCAKE_EP_USE_MUSA})
+endif()
 
 # ---------------------------------------------------------------------------
 # 2. Ensure engine.so exists in mooncake-wheel/mooncake/ for setup.py linking.
diff --git a/mooncake-ep/include/mooncake_ep_api.cuh b/mooncake-ep/include/mooncake_ep_api.cuh
index 940b9d0d..2ab1c632 100644
--- a/mooncake-ep/include/mooncake_ep_api.cuh
+++ b/mooncake-ep/include/mooncake_ep_api.cuh
@@ -18,6 +18,18 @@ void dispatch(void* packed_recv_x, float* packed_recv_x_scales,
               void* workspace, cudaStream_t stream, int64_t timeout_ticks,
               int phases);
 
+void mark_phase_ack(void* mxa_buffer, const int32_t* nvlink_available,
+                    void* const* ipc_peer_ptrs, int* ack_buffer, int rank,
+                    int num_ranks, int epoch, cudaStream_t stream);
+
+void wait_phase_ack(int* ack_buffer, int rank, int num_ranks, int epoch,
+                    cudaStream_t stream, int64_t timeout_ticks);
+
+void mark_and_wait_phase_ack(void* mxa_buffer, const int32_t* nvlink_available,
+                             void* const* ipc_peer_ptrs, int* ack_buffer,
+                             int rank, int num_ranks, int epoch,
+                             cudaStream_t stream, int64_t timeout_ticks);
+
 void combine(void* combined_x, int32_t* active_ranks, void* mxa_buffer,
              int* rdma_send_signal_buffer, int* rdma_recv_signal_buffer,
              void* rdma_send_data_buffer, void* rdma_recv_data_buffer,
diff --git a/mooncake-ep/include/mooncake_ep_buffer.h b/mooncake-ep/include/mooncake_ep_buffer.h
index 3cbb5d0b..6088ce20 100644
--- a/mooncake-ep/include/mooncake_ep_buffer.h
+++ b/mooncake-ep/include/mooncake_ep_buffer.h
@@ -5,17 +5,20 @@
 #include <cuda_bf16.h>
 #include <cuda.h>
 #include <cuda_runtime.h>
-#include <fstream>
-#include <mooncake_ibgda/memheap.h>
-#include <mooncake_ibgda/mlx5gda.h>
+#include <memory>
 #include <mooncake_ep_api.cuh>
 #include <mooncake_ep_configs.cuh>
 #include <mooncake_ep_event.h>
 #include <mooncake_ep_exception.cuh>
 #include <torch/torch.h>
+#include <transport/device/device_transport.h>
 
 namespace mooncake {
 
+class TransferEngine;
+
+// MAX_QP_COUNT is defined in mooncake_ep_configs.cuh (shared with kernel code).
+
 struct BufferLayout {
     int* rdma_send_signal_buffer;
     int* rdma_recv_signal_buffer;
@@ -39,7 +42,7 @@ struct BufferPair {
         size_t signaling_buffer_bytes = num_experts * sizeof(int);
         size_t send_recv_buffer_bytes =
             num_experts * num_max_dispatch_tokens_per_rank *
-            (2 * sizeof(int4) + hidden * sizeof(nv_bfloat16));
+            (2 * sizeof(int4) + hidden * EP_BF16_SIZE);
         for (int i = 0; i < 2; ++i) {
             size_t rdma_base_offset = total_bytes +
                                       2 * i * signaling_buffer_bytes +
@@ -66,44 +69,27 @@ struct MooncakeEpBuffer {
     int rank, num_ranks;
     int clock_rate_khz;
 
-    // MXA Buffer
+    // GDR buffer — owned by p2p_transport_
     int buffer_idx{};
+    int phase_epochs[2]{};
     int64_t num_ep_buffer_bytes;
     void* gdr_buffer = nullptr;
 
-    // IBGDA
-    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.
-    ibv_mr* mr = nullptr;
-    std::vector<mlx5gda_qp*> qps;
-    ibv_gid gid;
-    void* raddrs = nullptr;
-    void* rkeys = nullptr;
-    void* qp_devctxs = nullptr;
-    std::string device_name;
-    bool is_roce_ = false;
+    // Device transports — own all platform-specific state.
+    // p2p_transport_: NVLink intra-node P2P.
+    // rdma_transport_: IBGDA inter-node RDMA.  nullptr when IBGDA unavailable.
+    device::P2pTransport* p2p_transport_ = nullptr;
+    device::RdmaTransport* rdma_transport_ = nullptr;
+    // When EP creates transports itself (no engine provided), ownership lives
+    // in these unique_ptrs.  When an engine is provided, the engine owns them
+    // and these remain null.
+    std::unique_ptr<device::P2pTransport> owned_p2p_transport_;
+    std::unique_ptr<device::RdmaTransport> owned_rdma_transport_;
+
     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;
-    void** ipc_peer_ptrs_host = nullptr;
-    void** ipc_peer_ptrs = nullptr;
-    bool p2p_ipc_all_enabled_ = false;
-
     // Stream for communication
     at::cuda::CUDAStream comm_stream;
 
@@ -111,8 +97,11 @@ struct MooncakeEpBuffer {
     void* workspace = nullptr;
 
    public:
+    // If engine is provided, EP gets P2pTransport/RdmaTransport from it
+    // (engine owns the transports).  Otherwise EP creates its own via the
+    // factory functions (EP owns them via owned_p2p_transport_ etc.).
     MooncakeEpBuffer(int rank, int num_ranks, int64_t num_ep_buffer_bytes,
-                     std::string device_name);
+                     TransferEngine* engine = nullptr);
 
     ~MooncakeEpBuffer() noexcept(false);
 
@@ -136,79 +125,62 @@ struct MooncakeEpBuffer {
     torch::Tensor get_next_combine_buffer(int num_max_dispatch_tokens_per_rank,
                                           int hidden, int num_experts);
 
-    int init_ibgda();
+    bool ibgda_disabled() const { return ibgda_disabled_; }
 
-    bool ibgda_disabled() { return ibgda_disabled_; }
-
-    bool is_roce() { return is_roce_; }
-
-    // Decide whether EP can safely run CUDA kernels (\"fast-path\").
-    //
-    // There are two independent ways EP kernels can work:
-    // - IBGDA RDMA path: requires successful IBGDA init (qps/mr/etc).
-    // - NVLink P2P+IPC path: requires full P2P+IPC across ranks on the same
-    // node.
-    //
-    // IMPORTANT INVARIANT:
-    // If `p2p_ipc_all_enabled_ == true`, `sync_nvlink_ipc_handles()` guarantees
-    // `nvlink_available[dst_rank] == 1` for every rank pair, so the CUDA
-    // kernels will never take the IBGDA branch and therefore do NOT require
-    // `qps`.
-    bool use_fast_path() {
-        if (!ibgda_disabled_) {
-            return true;  // IBGDA available
-        }
-        // IBGDA disabled: only allow fast-path if we can rely on NVLink
-        // P2P+IPC.
-        if (!p2p_ipc_all_enabled_) {
-            LOG(WARNING) << "Failed to initialize IBGDA. "
-                         << "Using fallback implementation. "
-                         << "Performance will be degraded.";
-        }
-        return p2p_ipc_all_enabled_;
+    bool is_roce() const {
+        return rdma_transport_ && rdma_transport_->isRoce();
     }
 
+    // Fast-path: IBGDA available, or all peers accessible via P2P.
+    bool use_fast_path() {
+        if (!ibgda_disabled_) return true;
+        bool p2p_all = p2p_transport_ && p2p_transport_->allPeersAccessible();
+        if (!p2p_all) {
+            LOG(WARNING) << "IBGDA unavailable and P2P not fully accessible. "
+                         << "Using fallback (degraded performance).";
+        }
+        return p2p_all;
+    }
+
+    // Recreate QPs (called when active_ranks changes).
     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<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<int>& active_ranks_mask);
+    // Connect IBGDA QPs to peers.  Unified entry point — handles both IB and
+    // RoCE based on is_roce().
+    void sync_ibgda_peers(const std::vector<int64_t>& remote_addrs,
+                          const std::vector<int32_t>& remote_keys,
+                          const std::vector<std::vector<int32_t>>& peer_qpns,
+                          const std::vector<std::vector<int32_t>>& peer_lids,
+                          const std::vector<int64_t>& subnet_prefixes,
+                          const std::vector<int64_t>& interface_ids,
+                          const std::vector<int>& active_ranks_mask);
 
+    // Metadata accessors for Python-level bootstrap exchange.
     std::tuple<int64_t, int32_t> get_mr_info() {
-        return {(int64_t)mr->addr, (int32_t)mr->rkey};
+        if (!rdma_transport_) return {0, 0};
+        auto m = rdma_transport_->localMetadata();
+        return {m.raddr, m.rkey};
     }
 
     std::tuple<int64_t, int64_t> get_gid() {
-        return {(int64_t)gid.global.subnet_prefix,
-                (int64_t)gid.global.interface_id};
+        if (!rdma_transport_) return {0, 0};
+        auto m = rdma_transport_->localMetadata();
+        return {m.subnet_prefix, m.interface_id};
     }
 
     std::vector<int32_t> get_local_qpns() {
-        std::vector<int32_t> local_qpns;
-        for (int i = 0; i < USE_QP_COUNT; ++i) {
-            local_qpns.push_back((int32_t)qps[i]->qpn);
-        }
-        return local_qpns;
+        if (!rdma_transport_) return {};
+        return rdma_transport_->localMetadata().qpns;
     }
 
     std::vector<int32_t> get_local_lids() {
-        std::vector<int32_t> local_lids;
-        for (int i = 0; i < USE_QP_COUNT; ++i) {
-            local_lids.push_back((int32_t)qps[i]->port_attr.lid);
-        }
-        return local_lids;
+        if (!rdma_transport_) return {};
+        return rdma_transport_->localMetadata().lids;
     }
 
+    // IPC handle for P2P (NVLink).
     std::vector<int32_t> get_ipc_handle();
+
     void sync_nvlink_ipc_handles(
         const std::vector<std::vector<int32_t>>& remote_handles,
         const std::vector<int>& active_ranks_mask);
diff --git a/mooncake-ep/include/mooncake_ep_configs.cuh b/mooncake-ep/include/mooncake_ep_configs.cuh
index 10e0c1c5..103f6b8f 100644
--- a/mooncake-ep/include/mooncake_ep_configs.cuh
+++ b/mooncake-ep/include/mooncake_ep_configs.cuh
@@ -2,6 +2,7 @@
 
 #define NUM_MAX_NVL_PEERS 8
 #define NUM_MAX_RDMA_PEERS 20
+#define MAX_QP_COUNT 256
 #define NUM_MAX_FIFO_SLOTS 32768
 #define NUM_WORKSPACE_BYTES (32 * 1024 * 1024)
 #define NUM_MAX_LOCAL_EXPERTS 1024
@@ -42,3 +43,14 @@
 #include <cuda_fp8.h>
 #include <cuda_runtime.h>
 #include <infiniband/mlx5dv.h>
+
+// torchada maps nv_bfloat16 → __mt_bfloat16 which is an incomplete type on
+// MUSA, so sizeof(__mt_bfloat16) fails.  mt_bfloat16 (the complete typedef in
+// musa_bf16.hpp) requires the MUSA device compiler (mcc) and cannot be
+// included from host .cpp files.  Use EP_BF16_SIZE: sizeof(nv_bfloat16) on
+// CUDA, hardcoded 2 on MUSA (both are 2 bytes).
+#ifdef MOONCAKE_EP_USE_MUSA
+#define EP_BF16_SIZE 2
+#else
+#define EP_BF16_SIZE sizeof(nv_bfloat16)
+#endif
diff --git a/mooncake-ep/include/mooncake_ep_device.h b/mooncake-ep/include/mooncake_ep_device.h
new file mode 100644
index 00000000..244751be
--- /dev/null
+++ b/mooncake-ep/include/mooncake_ep_device.h
@@ -0,0 +1,92 @@
+#pragma once
+// ============================================================================
+// mooncake_ep_device.h - Minimal platform-specific definitions
+// ============================================================================
+// Only contains items that torchada's source-to-source translation cannot
+// handle.  Everything else uses plain CUDA API names — torchada maps them
+// to MUSA equivalents at build time via SimplePorting text replacement.
+// ============================================================================
+
+#ifdef MOONCAKE_EP_USE_MUSA
+
+// -- FP8 types (MUSA uses different names; not in torchada mapping) ----------
+#include <musa_fp8.h>
+using ep_fp8_storage_t = __mt_fp8_storage_t;
+using ep_fp8x2_storage_t = __mt_fp8x2_storage_t;
+#if defined(__CUDACC__) || defined(__MCC__)
+__device__ __forceinline__ ep_fp8x2_storage_t ep_cvt_float2_to_fp8x2(float2 x) {
+    return __musa_cvt_float2_to_fp8x2(x, __MT_SATFINITE, __MT_E4M3);
+}
+#endif
+
+// -- Device intrinsics (MUSA doesn't have __ldg / __activemask) --------------
+#ifndef __ldg
+#define __ldg(ptr) (*(ptr))
+#endif
+#ifndef __activemask
+#define __activemask() (0xffffffff)
+#endif
+
+#if defined(__CUDACC__) || defined(__MCC__)
+__forceinline__ __device__ int get_lane_id() { return threadIdx.x % 32; }
+#endif
+
+// -- Kernel launch (MUSA: no __launch_bounds__, no cooperative launch) --------
+#define EP_LAUNCH_BOUNDS(max_threads, min_blocks)
+
+#define SETUP_LAUNCH_CONFIG(num_sms, num_threads, stream) \
+    dim3 _grid(num_sms);                                  \
+    dim3 _block(num_threads);                             \
+    cudaStream_t _stream = stream
+
+#define LAUNCH_KERNEL(config, kernel, ...)                     \
+    kernel<<<_grid, _block, 0, _stream>>>(__VA_ARGS__);        \
+    {                                                          \
+        auto _err = cudaGetLastError();                        \
+        if (_err != cudaSuccess) {                             \
+            fprintf(stderr, "[EP] kernel launch failed: %s\n", \
+                    cudaGetErrorString(_err));                 \
+        }                                                      \
+    }
+
+#else  // !MOONCAKE_EP_USE_MUSA
+
+// -- FP8 types (CUDA native names) -------------------------------------------
+#include <cuda_fp8.h>
+using ep_fp8_storage_t = __nv_fp8_storage_t;
+using ep_fp8x2_storage_t = __nv_fp8x2_storage_t;
+#if defined(__CUDACC__) || defined(__MCC__)
+__device__ __forceinline__ ep_fp8x2_storage_t ep_cvt_float2_to_fp8x2(float2 x) {
+    return __nv_cvt_float2_to_fp8x2(x, __NV_SATFINITE, __NV_E4M3);
+}
+#endif
+
+// -- Device intrinsics -------------------------------------------------------
+#if defined(__CUDACC__) || defined(__MCC__)
+__forceinline__ __device__ int get_lane_id() {
+    int lane_id;
+    asm("mov.s32 %0, %laneid;" : "=r"(lane_id));
+    return lane_id;
+}
+#endif
+
+// -- Kernel launch (CUDA: cooperative launch) --------------------------------
+#define EP_LAUNCH_BOUNDS(max_threads, min_blocks) \
+    __launch_bounds__(max_threads, min_blocks)
+
+#define SETUP_LAUNCH_CONFIG(num_sms, num_threads, stream) \
+    cudaLaunchConfig_t cfg = {                            \
+        (num_sms), (num_threads), 0, stream, nullptr, 0}; \
+    cudaLaunchAttribute attr[1];                          \
+    attr[0].id = cudaLaunchAttributeCooperative;          \
+    attr[0].val.cooperative = 1;                          \
+    cfg.attrs = attr;                                     \
+    cfg.numAttrs = 1
+
+#define LAUNCH_KERNEL(config, kernel, ...) \
+    CUDA_CHECK(cudaLaunchKernelEx(config, kernel, ##__VA_ARGS__))
+
+#endif  // MOONCAKE_EP_USE_MUSA
+
+// Both platforms need IB verbs
+#include <infiniband/mlx5dv.h>
diff --git a/mooncake-ep/include/mooncake_ep_exception.cuh b/mooncake-ep/include/mooncake_ep_exception.cuh
index 107c6129..06074459 100644
--- a/mooncake-ep/include/mooncake_ep_exception.cuh
+++ b/mooncake-ep/include/mooncake_ep_exception.cuh
@@ -42,12 +42,22 @@ class EPException : public std::exception {
 #endif
 
 #ifndef EP_DEVICE_ASSERT
+#ifdef MOONCAKE_EP_USE_MUSA
+// MUSA SDK 4.3.x can turn kernels that merely contain a device-side __trap()
+// branch into illegal memory accesses, even when the assertion condition is
+// true.  Keep these invariants as host/static checks on MUSA builds.
+#define EP_DEVICE_ASSERT(cond) \
+    do {                       \
+        (void)sizeof(cond);    \
+    } while (0)
+#else
 #define EP_DEVICE_ASSERT(cond)                                           \
     do {                                                                 \
         if (not(cond)) {                                                 \
             printf("Assertion failed: %s:%d, condition: %s\n", __FILE__, \
                    __LINE__, #cond);                                     \
-            asm("trap;");                                                \
+            __trap();                                                    \
         }                                                                \
     } while (0)
 #endif
+#endif
diff --git a/mooncake-ep/include/mooncake_ep_launch.cuh b/mooncake-ep/include/mooncake_ep_launch.cuh
index aa39a141..c146c6d7 100644
--- a/mooncake-ep/include/mooncake_ep_launch.cuh
+++ b/mooncake-ep/include/mooncake_ep_launch.cuh
@@ -2,21 +2,8 @@
 
 #include <mooncake_ep_configs.cuh>
 
-#ifndef SETUP_LAUNCH_CONFIG
-#define SETUP_LAUNCH_CONFIG(num_sms, num_threads, stream) \
-    cudaLaunchConfig_t cfg = {                            \
-        (num_sms), (num_threads), 0, stream, nullptr, 0}; \
-    cudaLaunchAttribute attr[1];                          \
-    attr[0].id = cudaLaunchAttributeCooperative;          \
-    attr[0].val.cooperative = 1;                          \
-    cfg.attrs = attr;                                     \
-    cfg.numAttrs = 1
-#endif
-
-#ifndef LAUNCH_KERNEL
-#define LAUNCH_KERNEL(config, kernel, ...) \
-    CUDA_CHECK(cudaLaunchKernelEx(config, kernel, ##__VA_ARGS__))
-#endif
+// SETUP_LAUNCH_CONFIG and LAUNCH_KERNEL are defined in mooncake_ep_device.h
+// (unified CUDA/MUSA launch configuration).
 
 #define SWITCH_RANKS(case_macro)                           \
     switch (num_ranks) {                                   \
diff --git a/mooncake-ep/include/mooncake_ep_utils.cuh b/mooncake-ep/include/mooncake_ep_utils.cuh
index db92a04a..90492f67 100644
--- a/mooncake-ep/include/mooncake_ep_utils.cuh
+++ b/mooncake-ep/include/mooncake_ep_utils.cuh
@@ -1,6 +1,6 @@
 #pragma once
 
-#include <mooncake_ep_exception.cuh>
+#include <mooncake_ep_device.h>
 
 #define UNROLLED_WARP_COPY(UNROLL_FACTOR, LANE_ID, N, DST, SRC, LD_FUNC,      \
                            ST_FUNC)                                           \
@@ -47,288 +47,8 @@ struct VecInt<16> {
     using vec_t = int4;
 };
 
-__device__ __forceinline__ void trap() { asm("trap;"); }
-
-__device__ __forceinline__ void memory_fence() {
-    asm volatile("fence.acq_rel.sys;" ::: "memory");
-}
-
-__device__ __forceinline__ void memory_fence_gpu() {
-    asm volatile("fence.acq_rel.gpu;" ::: "memory");
-}
-
-__device__ __forceinline__ void memory_fence_cta() {
-    asm volatile("fence.acq_rel.cta;" ::: "memory");
-}
-
-__device__ __forceinline__ void st_relaxed_sys_global(const int *ptr, int val) {
-    asm volatile("st.relaxed.sys.global.s32 [%0], %1;" ::"l"(ptr), "r"(val)
-                 : "memory");
-}
-
-__device__ __forceinline__ void st_release_sys_global(const int *ptr, int val) {
-    asm volatile("st.release.sys.global.s32 [%0], %1;" ::"l"(ptr), "r"(val)
-                 : "memory");
-}
-
-__device__ __forceinline__ void st_release_cta(const int *ptr, int val) {
-    asm volatile("st.release.cta.s32 [%0], %1;" ::"l"(ptr), "r"(val)
-                 : "memory");
-}
-
-__device__ __forceinline__ int ld_acquire_sys_global(const int *ptr) {
-    int ret;
-    asm volatile("ld.acquire.sys.global.s32 %0, [%1];" : "=r"(ret) : "l"(ptr));
-    return ret;
-}
-
-__device__ __forceinline__ uint64_t ld_acquire_sys_global(const uint64_t *ptr) {
-    uint64_t ret;
-    asm volatile("ld.acquire.sys.global.u64 %0, [%1];" : "=l"(ret) : "l"(ptr));
-    return ret;
-}
-
-__device__ __forceinline__ int ld_acquire_global(const int *ptr) {
-    int ret;
-    asm volatile("ld.acquire.gpu.global.s32 %0, [%1];" : "=r"(ret) : "l"(ptr));
-    return ret;
-}
-
-__device__ __forceinline__ int atomic_add_release_sys_global(const int *ptr,
-                                                             int value) {
-    int ret;
-    asm volatile("atom.add.release.sys.global.s32 %0, [%1], %2;"
-                 : "=r"(ret)
-                 : "l"(ptr), "r"(value));
-    return ret;
-}
-
-__device__ __forceinline__ int atomic_add_release_global(const int *ptr,
-                                                         int value) {
-    int ret;
-    asm volatile("atom.add.release.gpu.global.s32 %0, [%1], %2;"
-                 : "=r"(ret)
-                 : "l"(ptr), "r"(value));
-    return ret;
-}
-
-__device__ __forceinline__ int ld_acquire_cta(const int *ptr) {
-    int ret;
-    asm volatile("ld.acquire.cta.s32 %0, [%1];" : "=r"(ret) : "l"(ptr));
-    return ret;
-}
-
-__device__ __forceinline__ uint8_t ld_na_relaxed(const uint8_t *ptr) {
-    uint16_t ret;
-    asm volatile("ld.relaxed.gpu.global.L1::no_allocate.b8 %0, [%1];"
-                 : "=h"(ret)
-                 : "l"(ptr));
-    return static_cast<uint8_t>(ret);
-}
-
-__device__ __forceinline__ uint16_t ld_na_relaxed(const uint16_t *ptr) {
-    uint16_t ret;
-    asm volatile("ld.relaxed.gpu.global.L1::no_allocate.b16 %0, [%1];"
-                 : "=h"(ret)
-                 : "l"(ptr));
-    return ret;
-}
-
-__device__ __forceinline__ uint32_t ld_na_relaxed(const uint32_t *ptr) {
-    uint32_t ret;
-    asm volatile("ld.relaxed.gpu.global.L1::no_allocate.b32 %0, [%1];"
-                 : "=r"(ret)
-                 : "l"(ptr));
-    return ret;
-}
-
-__device__ __forceinline__ uint64_t ld_na_relaxed(const uint64_t *ptr) {
-    uint64_t ret;
-    asm volatile("ld.relaxed.gpu.global.L1::no_allocate.b64 %0, [%1];"
-                 : "=l"(ret)
-                 : "l"(ptr));
-    return ret;
-}
-
-__device__ __forceinline__ int ld_volatile_global(const int *ptr) {
-    int ret;
-    asm volatile("ld.volatile.global.s32 %0, [%1];" : "=r"(ret) : "l"(ptr));
-    return ret;
-}
-
-__device__ __forceinline__ float ld_volatile_global(const float *ptr) {
-    float ret;
-    asm volatile("ld.volatile.global.f32 %0, [%1];" : "=f"(ret) : "l"(ptr));
-    return ret;
-}
-
-__device__ __forceinline__ int64_t ld_volatile_global(const int64_t *ptr) {
-    int64_t ret;
-    asm volatile("ld.volatile.global.s64 %0, [%1];" : "=l"(ret) : "l"(ptr));
-    return ret;
-}
-
-__device__ __forceinline__ int64_t ld_volatile_global(const uint64_t *ptr) {
-    int64_t ret;
-    asm volatile("ld.volatile.global.u64 %0, [%1];" : "=l"(ret) : "l"(ptr));
-    return ret;
-}
-
-#ifndef DISABLE_AGGRESSIVE_PTX_INSTRS
-#define LD_NC_FUNC "ld.global.nc.L1::no_allocate.L2::256B"
-#else
-#define LD_NC_FUNC "ld.volatile.global"
-#endif
-
-// `ld.global.nc.L1::no_allocate` will be translated into
-// `LDG.E.NA.[width].CONSTANT` in SASS
-template <typename dtype_t>
-__device__ __forceinline__ dtype_t ld_nc_global(const dtype_t *ptr) {
-    auto ret = ld_nc_global(
-        reinterpret_cast<const typename VecInt<sizeof(dtype_t)>::vec_t *>(ptr));
-    return *reinterpret_cast<dtype_t *>(&ret);
-}
-
-template <>
-__device__ __forceinline__ uint8_t ld_nc_global(const uint8_t *ptr) {
-    uint16_t ret;
-    // NOTES: we must use `uint16_t` as inline ASM does not support 8-bit
-    // constraint letter (`h` below means unsigned 16-bit)
-    asm volatile(LD_NC_FUNC ".u8 %0, [%1];" : "=h"(ret) : "l"(ptr));
-    return static_cast<uint8_t>(ret);
-}
-
-template <>
-__device__ __forceinline__ int ld_nc_global(const int *ptr) {
-    int ret;
-    asm volatile(LD_NC_FUNC ".s32 %0, [%1];" : "=r"(ret) : "l"(ptr));
-    return ret;
-}
-
-template <>
-__device__ __forceinline__ int64_t ld_nc_global(const int64_t *ptr) {
-    int64_t ret;
-    asm volatile(LD_NC_FUNC ".s64 %0, [%1];" : "=l"(ret) : "l"(ptr));
-    return ret;
-}
-
-template <>
-__device__ __forceinline__ float ld_nc_global(const float *ptr) {
-    float ret;
-    asm volatile(LD_NC_FUNC ".f32 %0, [%1];" : "=f"(ret) : "l"(ptr));
-    return ret;
-}
-
-template <>
-__device__ __forceinline__ int2 ld_nc_global(const int2 *ptr) {
-    int2 ret;
-    asm volatile(LD_NC_FUNC ".v2.s32 {%0, %1}, [%2];"
-                 : "=r"(ret.x), "=r"(ret.y)
-                 : "l"(ptr));
-    return ret;
-}
-
-template <>
-__device__ __forceinline__ int4 ld_nc_global(const int4 *ptr) {
-    int4 ret;
-    asm volatile(LD_NC_FUNC ".v4.s32 {%0, %1, %2, %3}, [%4];"
-                 : "=r"(ret.x), "=r"(ret.y), "=r"(ret.z), "=r"(ret.w)
-                 : "l"(ptr));
-    return ret;
-}
-
-__device__ __forceinline__ void st_na_relaxed(const uint8_t *ptr, uint8_t val) {
-    asm volatile("st.relaxed.gpu.global.L1::no_allocate.b8 [%0], %1;"
-                 :
-                 : "l"(ptr), "h"(static_cast<uint16_t>(val)));
-}
-
-__device__ __forceinline__ void st_na_relaxed(const uint16_t *ptr,
-                                              uint16_t val) {
-    asm volatile("st.relaxed.gpu.global.L1::no_allocate.b16 [%0], %1;"
-                 :
-                 : "l"(ptr), "h"(val));
-}
-
-__device__ __forceinline__ void st_na_relaxed(const uint32_t *ptr,
-                                              uint32_t val) {
-    asm volatile("st.relaxed.gpu.global.L1::no_allocate.b32 [%0], %1;"
-                 :
-                 : "l"(ptr), "r"(val));
-}
-
-__device__ __forceinline__ void st_na_relaxed(const int *ptr, int val) {
-    asm volatile("st.relaxed.gpu.global.L1::no_allocate.b32 [%0], %1;"
-                 :
-                 : "l"(ptr), "r"(val));
-}
-
-__device__ __forceinline__ void st_na_relaxed(const int4 *ptr, int4 val) {
-    asm volatile(
-        "st.relaxed.gpu.global.L1::no_allocate.v4.s32 [%0], {%1, %2, %3, %4};"
-        :
-        : "l"(ptr), "r"(val.x), "r"(val.y), "r"(val.z), "r"(val.w));
-}
-
-__device__ __forceinline__ void st_na_release(const int *ptr, int val) {
-    asm volatile("st.release.gpu.global.L1::no_allocate.b32 [%0], %1;"
-                 :
-                 : "l"(ptr), "r"(val));
-}
-
-__device__ __forceinline__ void st_na_release(const uint32_t *ptr,
-                                              uint32_t val) {
-    asm volatile("st.release.gpu.global.L1::no_allocate.b32 [%0], %1;"
-                 :
-                 : "l"(ptr), "r"(val));
-}
-
-__device__ __forceinline__ void st_na_release(const uint64_t *ptr,
-                                              uint64_t val) {
-    asm volatile("st.release.gpu.global.L1::no_allocate.b64 [%0], %1;"
-                 :
-                 : "l"(ptr), "l"(val));
-}
-
-// `st.global.L1::no_allocate` will be translated into `ST.E.NA.[width]` in SASS
-#ifndef DISABLE_AGGRESSIVE_PTX_INSTRS
-#define ST_NA_FUNC "st.global.L1::no_allocate"
-#else
-#define ST_NA_FUNC "st.global"
-#endif
-
-template <typename dtype_t>
-__device__ __forceinline__ void st_na_global(const dtype_t *ptr,
-                                             const dtype_t &value) {
-    st_na_global(
-        reinterpret_cast<const typename VecInt<sizeof(dtype_t)>::vec_t *>(ptr),
-        *reinterpret_cast<const typename VecInt<sizeof(dtype_t)>::vec_t *>(
-            &value));
-}
-
-template <>
-__device__ __forceinline__ void st_na_global(const int *ptr, const int &value) {
-    asm volatile(ST_NA_FUNC ".s32 [%0], %1;" ::"l"(ptr), "r"(value));
-}
-
-template <>
-__device__ __forceinline__ void st_na_global(const int64_t *ptr,
-                                             const int64_t &value) {
-    asm volatile(ST_NA_FUNC ".s64 [%0], %1;" ::"l"(ptr), "l"(value));
-}
-
-template <>
-__device__ __forceinline__ void st_na_global(const float *ptr,
-                                             const float &value) {
-    asm volatile(ST_NA_FUNC ".f32 [%0], %1;" ::"l"(ptr), "f"(value));
-}
-
-template <>
-__device__ __forceinline__ void st_na_global(const int4 *ptr,
-                                             const int4 &value) {
-    asm volatile(ST_NA_FUNC ".v4.s32 [%0], {%1, %2, %3, %4};" ::"l"(ptr),
-                 "r"(value.x), "r"(value.y), "r"(value.z), "r"(value.w));
-}
+// ---- TMA / mbarrier helpers (CUDA only) ----
+#ifndef MOONCAKE_EP_USE_MUSA
 
 __device__ __forceinline__ void fence_view_async_shared() {
     asm volatile("fence.proxy.async.shared::cta; \n" ::);
@@ -416,6 +136,8 @@ __device__ __forceinline__ void tma_store_wait() {
     asm volatile("cp.async.bulk.wait_group.read %0;" ::"n"(N) : "memory");
 }
 
+#endif  // MOONCAKE_EP_USE_MUSA
+
 template <typename dtype_t>
 __host__ __device__ dtype_t cell_div(dtype_t a, dtype_t b) {
     return (a + b - 1) / b;
@@ -486,53 +208,4 @@ __forceinline__ __device__ float half_warp_reduce_max(float value) {
     return value;
 }
 
-__forceinline__ __device__ int get_lane_id() {
-    int lane_id;
-    asm("mov.s32 %0, %laneid;" : "=r"(lane_id));
-    return lane_id;
-}
-
-template <int kNumRanks>
-__forceinline__ __device__ void move_fifo_slots(int &head) {
-    head = (head + kNumRanks) % NUM_MAX_FIFO_SLOTS;
-}
-
-template <int kNumRanks>
-__device__ __forceinline__ bool not_finished(int *task, int expected) {
-    auto result = false;
-    auto lane_id = threadIdx.x % 32;
-    if (lane_id < kNumRanks)
-        result = ld_volatile_global(task + lane_id) != expected;
-    return __any_sync(0xffffffff, result);
-}
-
-template <int kNumRanks>
-__forceinline__ __device__ void timeout_check(int **task_fifo_ptrs, int head,
-                                              int rank, int expected,
-                                              int tag = 0) {
-    auto start_time = clock64();
-    while (not_finished<kNumRanks>(task_fifo_ptrs[rank] + head, expected)) {
-        if (clock64() - start_time > NUM_TIMEOUT_CYCLES and threadIdx.x == 0) {
-            printf("Timeout check failed: %d (rank = %d)\n", tag, rank);
-            trap();
-        }
-    }
-}
-
-template <int kNumRanks>
-__forceinline__ __device__ void barrier_device(int **task_fifo_ptrs, int head,
-                                               int rank, int tag = 0) {
-    auto thread_id = static_cast<int>(threadIdx.x);
-    EP_DEVICE_ASSERT(kNumRanks <= 32);
-
-    if (thread_id < kNumRanks) {
-        atomicAdd_system(task_fifo_ptrs[rank] + head + thread_id,
-                         FINISHED_SUM_TAG);
-        memory_fence();
-        atomicSub_system(task_fifo_ptrs[thread_id] + head + rank,
-                         FINISHED_SUM_TAG);
-    }
-    timeout_check<kNumRanks>(task_fifo_ptrs, head, rank, 0, tag);
-}
-
 }  // namespace mooncake
diff --git a/mooncake-ep/setup.py b/mooncake-ep/setup.py
index 75a2fe62..0345c548 100644
--- a/mooncake-ep/setup.py
+++ b/mooncake-ep/setup.py
@@ -3,7 +3,23 @@ import re
 
 from setuptools import setup
 import torch
-from torch.utils.cpp_extension import BuildExtension, CUDAExtension, CUDA_HOME
+
+use_musa = os.getenv("MOONCAKE_EP_USE_MUSA", "").upper() in {"1", "ON", "TRUE", "YES"}
+if use_musa:
+    try:
+        import torchada  # noqa: F401
+    except ImportError as e:
+        raise ImportError(
+            "torchada is required to build the MUSA EP extension. "
+            "Please install it first using 'pip install torchada'."
+        ) from e
+
+
+from torch.utils.cpp_extension import (  # noqa: E402
+    BuildExtension,
+    CUDAExtension,
+    CUDA_HOME,
+)
 
 
 torch_version = re.match(r"\d+(?:\.\d+)*", torch.__version__).group()
@@ -13,18 +29,43 @@ 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.
+abi_define = f"-D_GLIBCXX_USE_CXX11_ABI={abi_flag}"
+cxx_args = [abi_define, "-std=c++20", "-O3", "-g0"]
+
 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)
-
-
+if use_musa:
+    cuda_libraries = []
+    musa_defines = ["-DUSE_MUSA", "-DMOONCAKE_EP_USE_MUSA=1"]
+    cxx_args += musa_defines
+    # torchada maps the "nvcc" key to "mcc".
+    device_args = [
+        abi_define,
+        *musa_defines,
+        "-std=c++20",
+        "--cuda-gpu-arch=mp_21",
+        "--cuda-gpu-arch=mp_31",
+        "-O3",
+    ]
+else:
+    cxx_args.append("-DUSE_CUDA")
+    device_args = [
+        abi_define,
+        "-std=c++20",
+        "-DUSE_CUDA",
+        "-Xcompiler",
+        "-O3",
+        "-Xcompiler",
+        "-g0",
+    ]
+    # Link against the CUDA driver stub library if available.
+    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,
@@ -39,18 +80,16 @@ setup(
                 "src/ep_py.cpp",
                 "src/mooncake_ep_buffer.cpp",
                 "src/mooncake_ep_kernel.cu",
-                "src/mooncake_ibgda/mlx5gda.cpp",
             ],
-            extra_compile_args={
-                "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"],
-            },
+            extra_compile_args={"cxx": cxx_args, "nvcc": device_args},
             libraries=cuda_libraries,
             library_dirs=cuda_library_dirs,
             extra_link_args=[
                 "-Wl,-rpath,$ORIGIN",
                 "-L" + os.path.join(current_dir, "../mooncake-wheel/mooncake"),
+                "-Wl,--push-state,--no-as-needed",
                 "-l:engine.so",
+                "-Wl,--pop-state",
             ],
         )
     ],
diff --git a/mooncake-ep/src/CMakeLists.txt b/mooncake-ep/src/CMakeLists.txt
index 850f8e12..574ab514 100644
--- a/mooncake-ep/src/CMakeLists.txt
+++ b/mooncake-ep/src/CMakeLists.txt
@@ -1,4 +1,4 @@
-add_library(mooncake_ep ep_py.cpp mooncake_ep_buffer.cpp mooncake_ep_kernel.cu mooncake_ibgda/mlx5gda.cpp)
+add_library(mooncake_ep ep_py.cpp mooncake_ep_buffer.cpp mooncake_ep_kernel.cu)
 
 set_target_properties(mooncake_ep PROPERTIES POSITION_INDEPENDENT_CODE ON)
 target_link_libraries(mooncake_ep PUBLIC ${TORCH_LIBRARIES} transfer_engine ibverbs mlx5)
diff --git a/mooncake-ep/src/ep_py.cpp b/mooncake-ep/src/ep_py.cpp
index b7458366..bd3cf53a 100644
--- a/mooncake-ep/src/ep_py.cpp
+++ b/mooncake-ep/src/ep_py.cpp
@@ -21,13 +21,12 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
     m.attr("MAX_QP_COUNT") = pybind11::int_(MAX_QP_COUNT);
 
     py::class_<MooncakeEpBuffer>(m, "Buffer")
-        .def(py::init<int, int, int64_t, std::string>())
+        .def(py::init<int, int, int64_t>())
         .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)
+        .def("sync_ibgda_peers", &MooncakeEpBuffer::sync_ibgda_peers)
         .def("get_mr_info", &MooncakeEpBuffer::get_mr_info)
         .def("get_gid", &MooncakeEpBuffer::get_gid)
         .def("get_local_qpns", &MooncakeEpBuffer::get_local_qpns)
@@ -41,4 +40,4 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
              &MooncakeEpBuffer::get_next_combine_buffer);
 }
 
-}  // namespace mooncake
\ No newline at end of file
+}  // namespace mooncake
diff --git a/mooncake-ep/src/mooncake_ep_buffer.cpp b/mooncake-ep/src/mooncake_ep_buffer.cpp
index 555d0d10..e8819f36 100644
--- a/mooncake-ep/src/mooncake_ep_buffer.cpp
+++ b/mooncake-ep/src/mooncake_ep_buffer.cpp
@@ -1,66 +1,34 @@
 #include <mooncake_ep_buffer.h>
-#include <arpa/inet.h>
 #include <glog/logging.h>
+#include <sstream>
+#include <transfer_engine.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;
+// Initialize an RDMA transport: register memory, allocate control buffer,
+// create QPs.  Returns true on success, false if IBGDA is unavailable.
+static bool initRdmaTransport(device::RdmaTransport* t, void* gdr_buffer,
+                              int64_t num_ep_buffer_bytes, int num_ranks,
+                              int use_qp_count, void* stream) {
+    int ret = t->initialize("", num_ranks, use_qp_count);
+    if (ret == 0) {
+        ret = t->registerMemory(gdr_buffer, num_ep_buffer_bytes);
     }
-    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 &&
-            a->s6_addr32[2] == htonl(0x0000ffff));
-}
-
-// Dynamically find the best GID index (RoCE v2 + IPv4-mapped address, or IB)
-// Returns GID index on success, -1 on failure
-static int findBestGidIndex(ibv_context* ctx, uint8_t port,
-                            ibv_port_attr& port_attr) {
-    for (int i = 0; i < port_attr.gid_tbl_len; i++) {
-        ibv_gid_entry gid_entry;
-        int ret = ibv_query_gid_ex(ctx, port, i, &gid_entry, 0);
-        if (ret) {
-            continue;
-        }
-
-        bool is_v4mapped = ipv6_addr_v4mapped(
-            reinterpret_cast<const struct in6_addr*>(gid_entry.gid.raw));
-
-        // Look for IPv4-mapped address + RoCE v2, or IB type
-        if ((is_v4mapped && gid_entry.gid_type == IBV_GID_TYPE_ROCE_V2) ||
-            gid_entry.gid_type == IBV_GID_TYPE_IB) {
-            return i;
-        }
+    if (ret == 0) {
+        ret = t->allocateControlBuffer();
     }
-    return -1;
+    if (ret == 0) {
+        ret = t->createQueuePairs(stream);
+    }
+    return ret == 0;
 }
 
 MooncakeEpBuffer::MooncakeEpBuffer(int rank, int num_ranks,
                                    int64_t num_ep_buffer_bytes,
-                                   std::string device_name)
+                                   TransferEngine* engine)
     : rank(rank),
       num_ranks(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
@@ -68,108 +36,55 @@ MooncakeEpBuffer::MooncakeEpBuffer(int rank, int num_ranks,
     CUDA_CHECK(cudaDeviceGetAttribute(&clock_rate_khz, cudaDevAttrClockRate,
                                       device_id));
 
-    // 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;
+    // P2P transport — owns GDR buffer allocation and IPC handle exchange.
+    if (engine) {
+        p2p_transport_ = engine->getOrCreateP2pTransport(num_ranks);
     } else {
-        CUDA_CHECK(cudaMalloc(&gdr_buffer, num_ep_buffer_bytes));
+        owned_p2p_transport_ = device::createP2pDeviceTransport(num_ranks);
+        p2p_transport_ = owned_p2p_transport_.get();
     }
-    CUDA_CHECK(cudaMalloc(&raddrs, num_ranks * sizeof(uint64_t)));
-    CUDA_CHECK(cudaMalloc(&rkeys, num_ranks * sizeof(uint32_t)));
-    CUDA_CHECK(
-        cudaMalloc(&qp_devctxs, USE_QP_COUNT * sizeof(mlx5gda_qp_devctx)));
 
-    // Allocate NVLink P2P arrays
-    CUDA_CHECK(cudaMalloc(&nvlink_available, num_ranks * sizeof(int32_t)));
-    CUDA_CHECK(cudaMemset(nvlink_available, 0, num_ranks * sizeof(int32_t)));
-    CUDA_CHECK(cudaMallocHost(&ipc_peer_ptrs_host, num_ranks * sizeof(void*)));
-    CUDA_CHECK(cudaMalloc(&ipc_peer_ptrs, num_ranks * sizeof(void*)));
-    for (int i = 0; i < num_ranks; ++i) {
-        ipc_peer_ptrs_host[i] = nullptr;
+    gdr_buffer = p2p_transport_->allocateBuffer(num_ep_buffer_bytes);
+    if (!gdr_buffer) {
+        throw std::runtime_error("[EP] Failed to allocate GDR buffer");
     }
-    CUDA_CHECK(cudaMemset(ipc_peer_ptrs, 0, num_ranks * sizeof(void*)));
+    CUDA_CHECK(cudaMemset(gdr_buffer, 0, num_ep_buffer_bytes));
 
-    int ret = init_ibgda();
-    if (ret != 0) {
-        ibgda_disabled_ = true;
+    // RDMA transport — optional; disabled if init fails.
+    if (engine) {
+        rdma_transport_ = engine->getOrCreateRdmaTransport();
+        if (rdma_transport_) {
+            if (!initRdmaTransport(rdma_transport_, gdr_buffer,
+                                   num_ep_buffer_bytes, num_ranks, USE_QP_COUNT,
+                                   comm_stream.stream())) {
+                rdma_transport_ = nullptr;
+                ibgda_disabled_ = true;
+                LOG(INFO) << "[EP] IBGDA unavailable, using P2P-only path";
+            }
+        } else {
+            ibgda_disabled_ = true;
+        }
+    } else {
+        // Read optional NIC whitelist from env var (same convention as PG
+        // tests).
+        std::vector<std::string> device_filter;
+        if (const char* env = std::getenv("MOONCAKE_EP_DEVICE_FILTER")) {
+            std::string s(env);
+            std::istringstream ss(s);
+            std::string tok;
+            while (std::getline(ss, tok, ',')) {
+                if (!tok.empty()) device_filter.push_back(tok);
+            }
+        }
+        auto t = device::createIbgdaDeviceTransport(device_filter);
+        if (initRdmaTransport(t.get(), gdr_buffer, num_ep_buffer_bytes,
+                              num_ranks, USE_QP_COUNT, comm_stream.stream())) {
+            owned_rdma_transport_ = std::move(t);
+            rdma_transport_ = owned_rdma_transport_.get();
+        } else {
+            ibgda_disabled_ = true;
+            LOG(INFO) << "[EP] IBGDA unavailable, using P2P-only path";
+        }
     }
 
     // Create 32 MiB workspace
@@ -178,29 +93,28 @@ MooncakeEpBuffer::MooncakeEpBuffer(int rank, int num_ranks,
 }
 
 MooncakeEpBuffer::~MooncakeEpBuffer() noexcept(false) {
-    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);
-    if (nvlink_available) cudaFree(nvlink_available);
-    if (ipc_peer_ptrs) cudaFree(ipc_peer_ptrs);
-    if (ipc_peer_ptrs_host) {
-        // Close IPC handles
-        for (int i = 0; i < num_ranks; ++i) {
-            if (ipc_peer_ptrs_host[i] != nullptr &&
-                ipc_peer_ptrs_host[i] != gdr_buffer) {
-                cudaIpcCloseMemHandle(ipc_peer_ptrs_host[i]);
-            }
+    // When EP owns the rdma transport, destructor handles QP/MR/ctrl_buf
+    // teardown. When engine owns it, just clear the pointer.
+    owned_rdma_transport_.reset();
+    rdma_transport_ = nullptr;
+
+    // When EP owns the p2p transport, free the buffer and reset.
+    if (owned_p2p_transport_) {
+        if (gdr_buffer) {
+            owned_p2p_transport_->freeBuffer(gdr_buffer);
+            gdr_buffer = nullptr;
+        }
+        owned_p2p_transport_.reset();
+    } else {
+        // Engine-owned: just free the buffer via the engine's transport.
+        if (gdr_buffer && p2p_transport_) {
+            p2p_transport_->freeBuffer(gdr_buffer);
+            gdr_buffer = nullptr;
         }
-        cudaFreeHost(ipc_peer_ptrs_host);
     }
+    p2p_transport_ = nullptr;
+
+    if (workspace) cudaFree(workspace);
 }
 
 std::tuple<torch::Tensor, std::optional<torch::Tensor>, torch::Tensor,
@@ -234,8 +148,10 @@ MooncakeEpBuffer::dispatch(const torch::Tensor& x,
     BufferPair layout(gdr_buffer, num_max_dispatch_tokens_per_rank, hidden,
                       num_ranks, num_experts);
     EP_HOST_ASSERT(layout.total_bytes <= num_ep_buffer_bytes);
-    auto buffer = layout.buffers[buffer_idx];
+    int current_buffer_idx = buffer_idx;
+    auto buffer = layout.buffers[current_buffer_idx];
     auto next_buffer = layout.buffers[buffer_idx ^= 1];
+    int phase_epoch = ++phase_epochs[current_buffer_idx];
 
     // Wait previous tasks to be finished
     // NOTES: the hook mode will always use the default stream
@@ -251,12 +167,12 @@ MooncakeEpBuffer::dispatch(const torch::Tensor& x,
         x.options().dtype(use_fp8 ? torch::kFloat8_e4m3fn : torch::kBFloat16));
     auto packed_recv_src_info = torch::empty(
         {num_local_experts, num_ranks * num_max_dispatch_tokens_per_rank},
-        torch::dtype(torch::kInt32).device(torch::kCUDA));
+        torch::dtype(torch::kInt32).device(x.device()));
     auto packed_recv_layout_range =
         torch::empty({num_local_experts, num_ranks},
-                     torch::dtype(torch::kInt64).device(torch::kCUDA));
+                     torch::dtype(torch::kInt64).device(x.device()));
     auto packed_recv_count = torch::zeros(
-        {num_local_experts}, torch::dtype(torch::kInt32).device(torch::kCUDA));
+        {num_local_experts}, torch::dtype(torch::kInt32).device(x.device()));
 
     // Allocate column-majored scales
     auto packed_recv_x_scales = std::optional<torch::Tensor>();
@@ -268,7 +184,7 @@ MooncakeEpBuffer::dispatch(const torch::Tensor& x,
         packed_recv_x_scales =
             torch::empty({num_local_experts, num_scales,
                           num_ranks * num_max_dispatch_tokens_per_rank},
-                         torch::dtype(torch::kFloat32).device(torch::kCUDA));
+                         torch::dtype(torch::kFloat32).device(x.device()));
         packed_recv_x_scales =
             torch::transpose(packed_recv_x_scales.value(), 1, 2);
         packed_recv_x_scales_ptr = packed_recv_x_scales->data_ptr<float>();
@@ -278,6 +194,37 @@ MooncakeEpBuffer::dispatch(const torch::Tensor& x,
         timeout_us == -1 ? -1
                          : (int64_t)clock_rate_khz * (int64_t)timeout_us / 1000;
 
+    void* raddrs_ptr = rdma_transport_ ? rdma_transport_->raddrsPtr() : nullptr;
+    void* rkeys_ptr = rdma_transport_ ? rdma_transport_->rkeysPtr() : nullptr;
+    void* qp_devctxs_ptr =
+        rdma_transport_ ? rdma_transport_->qpDevCtxsPtr() : nullptr;
+    int32_t* nvlink_avail = p2p_transport_->availableTablePtr();
+    void** ipc_ptrs = p2p_transport_->peerPtrsTablePtr();
+
+    auto mark_send_done = [=]() {
+#ifdef MOONCAKE_EP_USE_MUSA
+        mooncake::mark_phase_ack(gdr_buffer, nvlink_avail, ipc_ptrs,
+                                 buffer.rdma_send_signal_buffer, rank,
+                                 num_ranks, phase_epoch, launch_stream);
+#endif
+    };
+
+    auto wait_peer_send_done = [=]() {
+#ifdef MOONCAKE_EP_USE_MUSA
+        mooncake::wait_phase_ack(buffer.rdma_send_signal_buffer, rank,
+                                 num_ranks, phase_epoch, launch_stream,
+                                 timeout_ticks);
+#endif
+    };
+
+    auto mark_and_wait_peer_send_done = [=]() {
+#ifdef MOONCAKE_EP_USE_MUSA
+        mooncake::mark_and_wait_phase_ack(
+            gdr_buffer, nvlink_avail, ipc_ptrs, buffer.rdma_send_signal_buffer,
+            rank, num_ranks, phase_epoch, launch_stream, timeout_ticks);
+#endif
+    };
+
     auto launcher = [=](int phases) {
         mooncake::dispatch(
             packed_recv_x.data_ptr(), packed_recv_x_scales_ptr,
@@ -286,16 +233,25 @@ MooncakeEpBuffer::dispatch(const torch::Tensor& x,
             packed_recv_count.data_ptr<int>(), active_ranks.data_ptr<int32_t>(),
             gdr_buffer, buffer.rdma_send_signal_buffer,
             buffer.rdma_recv_signal_buffer, buffer.rdma_send_data_buffer,
-            buffer.rdma_recv_data_buffer, nullptr, nullptr, raddrs, rkeys,
-            qp_devctxs, nvlink_available, ipc_peer_ptrs, x.data_ptr(),
+            buffer.rdma_recv_data_buffer, nullptr, nullptr, raddrs_ptr,
+            rkeys_ptr, qp_devctxs_ptr, nvlink_avail, ipc_ptrs, x.data_ptr(),
             topk_idx.data_ptr<int64_t>(), next_buffer.rdma_recv_signal_buffer,
             num_tokens, hidden, num_max_dispatch_tokens_per_rank, num_topk,
             num_experts, rank, num_ranks, use_fp8, workspace, launch_stream,
             timeout_ticks, phases);
     };
-    launcher(return_recv_hook
-                 ? LOW_LATENCY_SEND_PHASE
-                 : (LOW_LATENCY_SEND_PHASE | LOW_LATENCY_RECV_PHASE));
+    if (return_recv_hook) {
+        launcher(LOW_LATENCY_SEND_PHASE);
+        mark_send_done();
+    } else {
+#ifdef MOONCAKE_EP_USE_MUSA
+        launcher(LOW_LATENCY_SEND_PHASE);
+        mark_and_wait_peer_send_done();
+        launcher(LOW_LATENCY_RECV_PHASE);
+#else
+        launcher(LOW_LATENCY_SEND_PHASE | LOW_LATENCY_RECV_PHASE);
+#endif
+    }
 
     // Wait streams
     std::optional<EventHandle> event;
@@ -311,7 +267,10 @@ MooncakeEpBuffer::dispatch(const torch::Tensor& x,
     // Receiver callback
     std::optional<std::function<void()>> recv_hook = std::nullopt;
     if (return_recv_hook)
-        recv_hook = [=]() { launcher(LOW_LATENCY_RECV_PHASE); };
+        recv_hook = [=]() {
+            wait_peer_send_done();
+            launcher(LOW_LATENCY_RECV_PHASE);
+        };
 
     // Return values
     return {packed_recv_x,
@@ -363,8 +322,10 @@ MooncakeEpBuffer::combine(const torch::Tensor& x, const torch::Tensor& topk_idx,
     BufferPair layout(gdr_buffer, num_max_dispatch_tokens_per_rank, hidden,
                       num_ranks, num_experts);
     EP_HOST_ASSERT(layout.total_bytes <= num_ep_buffer_bytes);
-    auto buffer = layout.buffers[buffer_idx];
+    int current_buffer_idx = buffer_idx;
+    auto buffer = layout.buffers[current_buffer_idx];
     auto next_buffer = layout.buffers[buffer_idx ^= 1];
+    int phase_epoch = ++phase_epochs[current_buffer_idx];
 
     // Wait previous tasks to be finished
     // NOTES: the hook mode will always use the default stream
@@ -389,14 +350,45 @@ MooncakeEpBuffer::combine(const torch::Tensor& x, const torch::Tensor& topk_idx,
         timeout_us == -1 ? -1
                          : (int64_t)clock_rate_khz * (int64_t)timeout_us / 1000;
 
+    void* raddrs_ptr = rdma_transport_ ? rdma_transport_->raddrsPtr() : nullptr;
+    void* rkeys_ptr = rdma_transport_ ? rdma_transport_->rkeysPtr() : nullptr;
+    void* qp_devctxs_ptr =
+        rdma_transport_ ? rdma_transport_->qpDevCtxsPtr() : nullptr;
+    int32_t* nvlink_avail = p2p_transport_->availableTablePtr();
+    void** ipc_ptrs = p2p_transport_->peerPtrsTablePtr();
+
+    auto mark_send_done = [=]() {
+#ifdef MOONCAKE_EP_USE_MUSA
+        mooncake::mark_phase_ack(gdr_buffer, nvlink_avail, ipc_ptrs,
+                                 buffer.rdma_send_signal_buffer, rank,
+                                 num_ranks, phase_epoch, launch_stream);
+#endif
+    };
+
+    auto wait_peer_send_done = [=]() {
+#ifdef MOONCAKE_EP_USE_MUSA
+        mooncake::wait_phase_ack(buffer.rdma_send_signal_buffer, rank,
+                                 num_ranks, phase_epoch, launch_stream,
+                                 timeout_ticks);
+#endif
+    };
+
+    auto mark_and_wait_peer_send_done = [=]() {
+#ifdef MOONCAKE_EP_USE_MUSA
+        mooncake::mark_and_wait_phase_ack(
+            gdr_buffer, nvlink_avail, ipc_ptrs, buffer.rdma_send_signal_buffer,
+            rank, num_ranks, phase_epoch, launch_stream, timeout_ticks);
+#endif
+    };
+
     // Kernel launch
     auto launcher = [=](int phases) {
         mooncake::combine(
             combined_x.data_ptr(), active_ranks.data_ptr<int32_t>(), gdr_buffer,
             buffer.rdma_send_signal_buffer, buffer.rdma_recv_signal_buffer,
             buffer.rdma_send_data_buffer, buffer.rdma_recv_data_buffer, nullptr,
-            nullptr, raddrs, rkeys, qp_devctxs, nvlink_available, ipc_peer_ptrs,
-            x.data_ptr(), topk_idx.data_ptr<int64_t>(),
+            nullptr, raddrs_ptr, rkeys_ptr, qp_devctxs_ptr, nvlink_avail,
+            ipc_ptrs, x.data_ptr(), topk_idx.data_ptr<int64_t>(),
             topk_weights.data_ptr<float>(), src_info.data_ptr<int>(),
             layout_range.data_ptr<int64_t>(),
             next_buffer.rdma_recv_signal_buffer, num_combined_tokens, hidden,
@@ -404,9 +396,18 @@ MooncakeEpBuffer::combine(const torch::Tensor& x, const torch::Tensor& topk_idx,
             num_ranks, workspace, launch_stream, timeout_ticks, phases,
             zero_copy);
     };
-    launcher(return_recv_hook
-                 ? LOW_LATENCY_SEND_PHASE
-                 : (LOW_LATENCY_SEND_PHASE | LOW_LATENCY_RECV_PHASE));
+    if (return_recv_hook) {
+        launcher(LOW_LATENCY_SEND_PHASE);
+        mark_send_done();
+    } else {
+#ifdef MOONCAKE_EP_USE_MUSA
+        launcher(LOW_LATENCY_SEND_PHASE);
+        mark_and_wait_peer_send_done();
+        launcher(LOW_LATENCY_RECV_PHASE);
+#else
+        launcher(LOW_LATENCY_SEND_PHASE | LOW_LATENCY_RECV_PHASE);
+#endif
+    }
 
     // Wait streams
     std::optional<EventHandle> event;
@@ -422,7 +423,10 @@ MooncakeEpBuffer::combine(const torch::Tensor& x, const torch::Tensor& topk_idx,
     // Receiver callback
     std::optional<std::function<void()>> recv_hook = std::nullopt;
     if (return_recv_hook)
-        recv_hook = [=]() { launcher(LOW_LATENCY_RECV_PHASE); };
+        recv_hook = [=]() {
+            wait_peer_send_done();
+            launcher(LOW_LATENCY_RECV_PHASE);
+        };
 
     // Return values
     return {combined_x, event, recv_hook};
@@ -435,12 +439,11 @@ torch::Tensor MooncakeEpBuffer::get_next_combine_buffer(
 
     auto buffer = layout.buffers[buffer_idx];
     auto dtype = torch::kBFloat16;
-    size_t num_bytes_per_combine_msg = hidden * sizeof(nv_bfloat16);
-    auto num_msg_elems = static_cast<int>(num_bytes_per_combine_msg /
-                                          elementSize(torch::kBFloat16));
+    size_t num_bytes_per_combine_msg = hidden * EP_BF16_SIZE;
+    auto num_msg_elems =
+        static_cast<int>(num_bytes_per_combine_msg / elementSize(dtype));
 
-    EP_HOST_ASSERT(num_bytes_per_combine_msg % elementSize(torch::kBFloat16) ==
-                   0);
+    EP_HOST_ASSERT(num_bytes_per_combine_msg % elementSize(dtype) == 0);
     return torch::from_blob(
         buffer.rdma_send_data_buffer,
         {num_experts / num_ranks, num_ranks * num_max_dispatch_tokens_per_rank,
@@ -450,373 +453,57 @@ torch::Tensor MooncakeEpBuffer::get_next_combine_buffer(
         torch::TensorOptions().dtype(dtype).device(torch::kCUDA));
 }
 
-int MooncakeEpBuffer::init_ibgda() {
-    int num_devices;
-    ibv_device** dev_list = ibv_get_device_list(&num_devices);
-    int nic_id = -1;
-    for (int i = 0; i < num_devices; ++i) {
-        const char* name = ibv_get_device_name(dev_list[i]);
-        if (name && device_name == name) {
-            nic_id = i;
-            break;
-        }
-    }
-    if (nic_id == -1) {
-        throw std::runtime_error("Device matching name '" + device_name +
-                                 "' not found.");
-    }
-    LOG(INFO) << "[EP] GPU " << device_id << " uses NIC " << nic_id
-              << " out of " << num_devices << " NIC(s)";
-    ibv_context* ctx = ibv_open_device(dev_list[nic_id]);
-    if (!ctx) {
-        perror("Failed to open device");
-        return -1;
-    }
-
-    // Query port attributes to get GID table length
-    ibv_port_attr port_attr;
-    const uint8_t port_num = 1;
-    if (ibv_query_port(ctx, port_num, &port_attr)) {
-        perror("Failed to query port");
-        return -1;
-    }
-
-    // Dynamically find the best GID index (replaces hardcoded index 3)
-    gid_index_ = findBestGidIndex(ctx, port_num, port_attr);
-    if (gid_index_ < 0) {
-        LOG(ERROR) << "[EP] Failed to find a suitable GID index on "
-                   << device_name;
-        return -1;
-    }
-
-    if (ibv_query_gid(ctx, port_num, gid_index_, &gid)) {
-        perror("Failed to query gid");
-        return -1;
-    }
-    ibv_free_device_list(dev_list);
-
-    pd = ibv_alloc_pd(ctx);
-    if (!pd) {
-        perror("Failed to allocate protection domain");
-        return -1;
-    }
-    mlx5dv_obj dv_obj = {};
-    dv_obj.pd.in = pd;
-    dv_obj.pd.out = &mpd;
-    if (mlx5dv_init_obj(&dv_obj, MLX5DV_OBJ_PD)) {
-        perror("Failed to initialize mlx5dv object");
-    }
-    mr = ibv_reg_mr(pd, gdr_buffer, num_ep_buffer_bytes,
-                    IBV_ACCESS_LOCAL_WRITE | IBV_ACCESS_REMOTE_READ |
-                        IBV_ACCESS_REMOTE_WRITE | IBV_ACCESS_REMOTE_ATOMIC);
-    if (!mr) {
-        perror("Failed to reg mr");
-    }
-
-    // Allocate ctrl_buf without zero-initialization. Individual regions will be
-    // 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));
-    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,
-                "If the error is `Bad address`, probably because your GPU "
-                "does not support GPUDirect RDMA.\n");
-        // Keep internal state consistent: IBGDA init failed, so `mr` must not
-        // be treated as valid.
-        if (mr) {
-            ibv_dereg_mr(mr);
-            mr = nullptr;
-        }
-        return -1;
-    }
-    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 < 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 create QP");
-            return -1;
-        }
-        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");
-            return -1;
-        }
-        // 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.push_back(qp);
-    }
-    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;
+    if (!rdma_transport_) return;
+    int ret = rdma_transport_->recreateQueuePairs(comm_stream.stream());
+    if (ret != 0) {
+        ibgda_disabled_ = true;
+        LOG(ERROR) << "[EP] Failed to recreate QPs";
     }
 }
 
-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,
-                               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,
-        };
-        if (mlx5gda_modify_rc_qp_init2rtr(
-                qps[i], ah_attr, (uint32_t)remote_qpns[i], IBV_MTU_4096)) {
-            perror("Failed to mlx5gda_modify_rc_qp_init2rtr");
-            exit(1);
-        }
-        if (mlx5gda_modify_rc_qp_rtr2rts(qps[i])) {
-            perror("Failed to mlx5gda_modify_rc_qp_rtr2rts");
-            exit(1);
-        }
-    }
-    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),
-                   cudaMemcpyHostToDevice);
-        uint32_t rkey = i == rank ? mr->lkey : (uint32_t)remote_keys[i];
-        cudaMemcpy(rkeys + i * sizeof(uint32_t), &rkey, sizeof(uint32_t),
-                   cudaMemcpyHostToDevice);
-    }
-}
+void MooncakeEpBuffer::sync_ibgda_peers(
+    const std::vector<int64_t>& remote_addrs,
+    const std::vector<int32_t>& remote_keys,
+    const std::vector<std::vector<int32_t>>& peer_qpns,
+    const std::vector<std::vector<int32_t>>& peer_lids,
+    const std::vector<int64_t>& subnet_prefixes,
+    const std::vector<int64_t>& interface_ids,
+    const std::vector<int>& active_ranks_mask) {
+    if (!rdma_transport_) return;
 
-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,
-                                 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[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;
-        ah_attr.grh.sgid_index =
-            gid_index_;  // Use dynamically discovered GID index
-        ah_attr.grh.hop_limit = 1;
-        ah_attr.port_num = 1;
-        ah_attr.dlid = qps[i]->port_attr.lid | 0xC000;
-        if (mlx5gda_modify_rc_qp_init2rtr(
-                qps[i], ah_attr, (uint32_t)remote_qpns[i], IBV_MTU_4096)) {
-            perror("Failed to mlx5gda_modify_rc_qp_init2rtr");
-            exit(1);
-        }
-        if (mlx5gda_modify_rc_qp_rtr2rts(qps[i])) {
-            perror("Failed to mlx5gda_modify_rc_qp_rtr2rts");
-            exit(1);
+    // Flatten per-rank QPN/LID vectors into per-QP arrays expected by
+    // connectPeers.  Each rank owns USE_QP_COUNT/num_ranks QPs.
+    int qps_per_rank = USE_QP_COUNT / num_ranks;
+    std::vector<int32_t> flat_qpns, flat_lids;
+    flat_qpns.reserve(USE_QP_COUNT);
+    flat_lids.reserve(USE_QP_COUNT);
+    for (int r = 0; r < num_ranks; ++r) {
+        for (int q = 0; q < qps_per_rank; ++q) {
+            flat_qpns.push_back(peer_qpns[r][q]);
+            flat_lids.push_back(peer_lids[r][q]);
         }
     }
-    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),
-                   cudaMemcpyHostToDevice);
-        uint32_t rkey = i == rank ? mr->lkey : (uint32_t)remote_keys[i];
-        cudaMemcpy(rkeys + i * sizeof(uint32_t), &rkey, sizeof(uint32_t),
-                   cudaMemcpyHostToDevice);
+
+    int ret = rdma_transport_->connectPeers(
+        rank, rdma_transport_->isRoce(), remote_addrs, remote_keys, flat_qpns,
+        flat_lids, subnet_prefixes, interface_ids, active_ranks_mask);
+    if (ret != 0) {
+        ibgda_disabled_ = true;
+        LOG(WARNING)
+            << "[EP] IBGDA connectPeers failed, falling back to P2P-only path";
     }
 }
 
 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
-    const size_t handle_size = sizeof(cudaIpcMemHandle_t);
-    const size_t num_int32s =
-        (handle_size + sizeof(int32_t) - 1) / sizeof(int32_t);
-    std::vector<int32_t> handle_ints(num_int32s);
-    memcpy(handle_ints.data(), &handle, handle_size);
-    return handle_ints;
+    return p2p_transport_->exportIpcHandle(gdr_buffer);
 }
 
 void MooncakeEpBuffer::sync_nvlink_ipc_handles(
     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;
-
-    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);
-
-        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;
-                    }
-                }
-            }
-        }
-
-        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;
-            }
-        }
-    }
-
-    CUDA_CHECK(cudaMemcpy(nvlink_available, nvlink_array.data(),
-                          num_ranks * sizeof(int32_t), cudaMemcpyHostToDevice));
-    CUDA_CHECK(cudaMemcpy(ipc_peer_ptrs, ipc_peer_ptrs_host,
-                          num_ranks * sizeof(void*), cudaMemcpyHostToDevice));
+    p2p_transport_->importPeerHandles(gdr_buffer, rank, num_ranks,
+                                      remote_handles, active_ranks_mask);
 }
 
 }  // namespace mooncake
diff --git a/mooncake-ep/src/mooncake_ep_kernel.cu b/mooncake-ep/src/mooncake_ep_kernel.cu
index 51979fd9..e6162d43 100644
--- a/mooncake-ep/src/mooncake_ep_kernel.cu
+++ b/mooncake-ep/src/mooncake_ep_kernel.cu
@@ -1,157 +1,133 @@
 // clang-format off
 
-#include <cooperative_groups.h>
 #include <cstdio>
-#include <cuda/atomic>
 
 #include <mooncake_ep_configs.cuh>
 #include <mooncake_ep_exception.cuh>
 #include <mooncake_ep_launch.cuh>
-#include <mooncake_ibgda/mlx5gda.h>
+#include <transport/device/comm_device.cuh>
 #include <mooncake_ep_utils.cuh>
 
 namespace mooncake {
 
-static __device__ void device_mutex_lock_system(uint32_t *mutex) {
-    cuda::atomic_ref<uint32_t, cuda::thread_scope_system> lock(*mutex);
-    // Spin until the mutex is acquired
-    while (lock.exchange(1, cuda::memory_order_acquire) != 0);
-}
+using mooncake::device::CommCtx;
+using mooncake::device::make_comm_ctx;
+using mooncake::device::mc_route_put;
+using mooncake::device::mc_rdma_put;
+using mooncake::device::mc_signal;
+using mooncake::device::mc_red_add;
+using mooncake::device::mc_bar_sync;
+using mooncake::device::mc_grid_sync;
+using mooncake::device::mc_ld_nc;
+using mooncake::device::mc_ld_nc_s32;
+using mooncake::device::mc_ld_nc_f32;
+using mooncake::device::mc_st_na;
+using mooncake::device::mc_ld_acquire;
+using mooncake::device::mc_st_release;
+using mooncake::device::mc_atomic_add_release;
+using mooncake::device::mc_fence;
+using mooncake::device::mc_fence_barrier_fence;
 
-static __device__ void device_mutex_unlock_system(uint32_t *mutex) {
-    cuda::atomic_ref<uint32_t, cuda::thread_scope_system> lock(*mutex);
-    // Release the mutex
-    lock.store(0, cuda::memory_order_release);
-}
+__global__ void mark_phase_ack_kernel(void* mxa_buffer,
+                                      const int32_t* nvlink_available,
+                                      void* const* ipc_peer_ptrs,
+                                      int* ack_buffer, int rank,
+                                      int num_ranks, int epoch) {
+    const CommCtx comm_ctx = make_comm_ctx(
+        mxa_buffer, nvlink_available, ipc_peer_ptrs, nullptr, nullptr, nullptr,
+        ack_buffer, ack_buffer, rank, num_ranks, MAX_QP_COUNT);
 
-static __device__ uint16_t device_byteswap(uint16_t x) {
-    return __byte_perm(x, x, 0x2301);
-}
-
-static __device__ uint32_t device_byteswap(uint32_t x) {
-    return __byte_perm(x, x, 0x0123);
-}
-
-static __device__ uint64_t device_byteswap(uint64_t x) {
-    uint32_t hi = (uint32_t)(x >> 32);
-    uint32_t lo = (uint32_t)(x);
-
-    hi = __byte_perm(hi, hi, 0x0123);
-    lo = __byte_perm(lo, lo, 0x0123);
-
-    return ((uint64_t)lo << 32) | hi;
-}
-
-__device__ static inline uint16_t ptx_ld16_acq_sys_na(uint16_t *ptr) {
-    uint16_t val;
-    asm volatile("ld.acquire.sys.global.L1::no_allocate.b16 %0, [%1];" : "=h"(val) : "l"(ptr));
-    return val;
-}
-
-// must be called with mutex locked
-static __device__ void __mlx5gda_device_poll_cq(struct mlx5gda_qp_devctx *ctx, uint16_t expect) {
-    uint16_t wq_tail = ctx->wq_tail;
-    while ((int16_t)(wq_tail - expect) <= 0) {
-        uint16_t cq_wqe_counter_be = ptx_ld16_acq_sys_na(&ctx->cq->wqe_counter);
-        // printf("cq_wqe_counter_be=0x%x, expect=0x%x\n", cq_wqe_counter_be, expect);
-        uint8_t opcode = ctx->cq->op_own >> 4;
-        if (opcode == 0xD) {
-            printf("Requester_Error: syndrome = 0x%lx\n", ctx->cq->timestamp >> 56);
+    for (int peer = static_cast<int>(threadIdx.x); peer < num_ranks;
+         peer += static_cast<int>(blockDim.x)) {
+        if (peer == rank) {
+            mc_st_release(ack_buffer + rank, epoch);
+        } else {
+            void* dst = mc_route_put(comm_ctx, peer, ack_buffer + rank);
+            if (dst != nullptr)
+                mc_st_release(reinterpret_cast<int*>(dst), epoch);
         }
-        EP_DEVICE_ASSERT(opcode == 0x0 || opcode == 0xF);
-        wq_tail = device_byteswap(cq_wqe_counter_be) + 1;
-    }
-    if (wq_tail != ctx->wq_tail) {
-        ctx->wq_tail = wq_tail;
     }
 }
 
-__device__ static inline void ptx_st32_rel_sys_na(uint32_t *ptr, uint32_t val) {
-    asm volatile("st.release.sys.global.L1::no_allocate.b32 [%0], %1;" : : "l"(ptr), "r"(val));
+__global__ void wait_phase_ack_kernel(int* ack_buffer, int rank, int num_ranks,
+                                      int epoch, int64_t timeout_ticks) {
+    for (int peer = static_cast<int>(threadIdx.x); peer < num_ranks;
+         peer += static_cast<int>(blockDim.x)) {
+        if (peer == rank)
+            continue;
+
+        int64_t start_time = static_cast<int64_t>(clock64());
+        while (mc_ld_acquire(ack_buffer + peer) < epoch) {
+            int64_t end_time = static_cast<int64_t>(clock64());
+            if (timeout_ticks != -1 && end_time - start_time > timeout_ticks)
+                return;
+        }
+    }
 }
 
-__device__ static inline void ptx_st64_rel_sys_na(uint64_t *ptr, uint64_t val) {
-    asm volatile("st.release.sys.global.L1::no_allocate.b64 [%0], %1;" : : "l"(ptr), "l"(val));
+__global__ void mark_and_wait_phase_ack_kernel(
+        void* mxa_buffer, const int32_t* nvlink_available,
+        void* const* ipc_peer_ptrs, int* ack_buffer, int rank, int num_ranks,
+        int epoch, int64_t timeout_ticks) {
+    const CommCtx comm_ctx = make_comm_ctx(
+        mxa_buffer, nvlink_available, ipc_peer_ptrs, nullptr, nullptr, nullptr,
+        ack_buffer, ack_buffer, rank, num_ranks, MAX_QP_COUNT);
+
+    for (int peer = static_cast<int>(threadIdx.x); peer < num_ranks;
+         peer += static_cast<int>(blockDim.x)) {
+        if (peer == rank) {
+            mc_st_release(ack_buffer + rank, epoch);
+        } else {
+            void* dst = mc_route_put(comm_ctx, peer, ack_buffer + rank);
+            if (dst != nullptr)
+                mc_st_release(reinterpret_cast<int*>(dst), epoch);
+        }
+    }
+
+    __syncthreads();
+
+    for (int peer = static_cast<int>(threadIdx.x); peer < num_ranks;
+         peer += static_cast<int>(blockDim.x)) {
+        if (peer == rank)
+            continue;
+
+        int64_t start_time = static_cast<int64_t>(clock64());
+        while (mc_ld_acquire(ack_buffer + peer) < epoch) {
+            int64_t end_time = static_cast<int64_t>(clock64());
+            if (timeout_ticks != -1 && end_time - start_time > timeout_ticks)
+                return;
+        }
+    }
 }
 
-/**
- * Ring DB to post WQs up to last_posted_wqe.
- *
- * Must be called with mutex locked.
- */
-static __device__ void __mlx5gda_device_post_send_db(struct mlx5gda_qp_devctx *ctx) {
-    // 1. update dbr
-    uint32_t num_posted_wqe = (uint32_t)(ctx->wq_head);
-    ptx_st32_rel_sys_na(&ctx->dbr->send_counter, device_byteswap((uint32_t)ctx->wq_head));
-    // 2. ring db
-    struct mlx5gda_wqebb *last_wqe = ctx->wq + ((num_posted_wqe - 1) & ctx->wqeid_mask);
-    // printf("Last wqe=%lx -> bf=%p\n", *(uint64_t*)last_wqe, ctx->bf + ctx->bf_offset);
-    ptx_st64_rel_sys_na((uint64_t*)(ctx->bf + ctx->bf_offset), *(uint64_t*)last_wqe);
-    // 3. toggle bf
-    ctx->bf_offset ^= MLX5GDA_BF_SIZE;
+void mark_phase_ack(void* mxa_buffer, const int32_t* nvlink_available,
+                    void* const* ipc_peer_ptrs, int* ack_buffer, int rank,
+                    int num_ranks, int epoch, cudaStream_t stream) {
+    SETUP_LAUNCH_CONFIG(1, 32, stream);
+    LAUNCH_KERNEL(&cfg, mark_phase_ack_kernel, mxa_buffer, nvlink_available,
+                  ipc_peer_ptrs, ack_buffer, rank, num_ranks, epoch);
 }
 
-static __device__ void __mlx5gda_device_write_rdma_write_wqe(
-    struct mlx5gda_qp_devctx *ctx, uint64_t laddr, __be32 lkey,
-    uint64_t raddr, __be32 rkey, uint32_t bytes) {
-    struct mlx5gda_rdma_write_wqe *wqe = (mlx5gda_rdma_write_wqe *)(ctx->wq + (ctx->wq_head & ctx->wqeid_mask));
-    struct mlx5_wqe_ctrl_seg &ctrl_seg = wqe->ctrl;
-    struct mlx5_wqe_raddr_seg &raddr_seg = wqe->raddr;
-    struct mlx5_wqe_data_seg &data_seg = wqe->data;
-
-    ctrl_seg = {};
-    ctrl_seg.qpn_ds = device_byteswap((ctx->qpn << 8) | 3);
-    ctrl_seg.fm_ce_se = MLX5_WQE_CTRL_CQ_UPDATE;
-    ctrl_seg.opmod_idx_opcode = device_byteswap(((uint32_t)ctx->wq_head << 8) | MLX5_OPCODE_RDMA_WRITE);
-
-    raddr_seg.raddr = device_byteswap(raddr);
-    raddr_seg.rkey = rkey;
-    raddr_seg.reserved = 0;
-
-    data_seg.byte_count = device_byteswap(bytes);
-    data_seg.lkey = lkey;
-    data_seg.addr = device_byteswap(laddr);
-
-    ++ctx->wq_head;
+void wait_phase_ack(int* ack_buffer, int rank, int num_ranks, int epoch,
+                    cudaStream_t stream, int64_t timeout_ticks) {
+    SETUP_LAUNCH_CONFIG(1, 32, stream);
+    LAUNCH_KERNEL(&cfg, wait_phase_ack_kernel, ack_buffer, rank, num_ranks,
+                  epoch, timeout_ticks);
 }
 
-struct mlx5_wqe_atomic_add_32_seg {
-    __be32		add_data;
-    __be32		field_boundary;
-    __be64		compare;
-};
-
-static __device__ void __mlx5gda_device_write_rdma_atomic_add_wqe(
-    struct mlx5gda_qp_devctx *ctx, const int& value, uint64_t laddr,
-    __be32 lkey, uint64_t raddr, __be32 rkey) {
-    struct mlx5gda_rdma_atomic_wqe *wqe = (mlx5gda_rdma_atomic_wqe *)(ctx->wq + (ctx->wq_head & ctx->wqeid_mask));
-    struct mlx5_wqe_ctrl_seg &ctrl_seg = wqe->ctrl;
-    struct mlx5_wqe_raddr_seg &raddr_seg = wqe->raddr;
-    struct mlx5_wqe_data_seg &data_seg = wqe->data;
-
-    ctrl_seg = {};
-    ctrl_seg.qpn_ds = device_byteswap((ctx->qpn << 8) | 4);
-    ctrl_seg.fm_ce_se = MLX5_WQE_CTRL_CQ_UPDATE;
-    ctrl_seg.opmod_idx_opcode = device_byteswap(MLX5_OPCODE_ATOMIC_MASKED_FA | ((uint32_t)ctx->wq_head << 8) | 0x08000000);
-
-    raddr_seg.raddr = device_byteswap(raddr);
-    raddr_seg.rkey = rkey;
-    raddr_seg.reserved = 0;
-
-    auto atomic_add_32_seg = reinterpret_cast<mlx5_wqe_atomic_add_32_seg *>(&wqe->atomic);
-    atomic_add_32_seg->add_data = device_byteswap((uint32_t)value);
-    atomic_add_32_seg->field_boundary = 0;
-    atomic_add_32_seg->compare = 0;
-
-    data_seg.byte_count = device_byteswap((uint32_t)4);
-    data_seg.lkey = lkey;
-    data_seg.addr = device_byteswap(laddr);
-
-    ++ctx->wq_head;
+void mark_and_wait_phase_ack(void* mxa_buffer,
+                             const int32_t* nvlink_available,
+                             void* const* ipc_peer_ptrs, int* ack_buffer,
+                             int rank, int num_ranks, int epoch,
+                             cudaStream_t stream, int64_t timeout_ticks) {
+    SETUP_LAUNCH_CONFIG(1, 32, stream);
+    LAUNCH_KERNEL(&cfg, mark_and_wait_phase_ack_kernel, mxa_buffer,
+                  nvlink_available, ipc_peer_ptrs, ack_buffer, rank, num_ranks,
+                  epoch, timeout_ticks);
 }
 
 template <bool kUseFP8, int kNumWarpGroups, int kNumWarpsPerGroup, int kHidden>
-__global__ __launch_bounds__(kNumWarpGroups * kNumWarpsPerGroup * 32, 1) void
+__global__ EP_LAUNCH_BOUNDS(kNumWarpGroups * kNumWarpsPerGroup * 32, 1) void
 dispatch(void* packed_recv_x, float* packed_recv_x_scales,
          int* packed_recv_src_info, int64_t* packed_recv_layout_range,
          int* packed_recv_count, int32_t* active_ranks,
@@ -182,20 +158,22 @@ dispatch(void* packed_recv_x, float* packed_recv_x_scales,
     constexpr int kNumPerChannels = 128;
     constexpr float kFP8Margin = 1e-4, kFP8Amax = 448, kFP8AmaxInv = 1.0f / 448.0f;
     const int num_scales = kHidden / kNumPerChannels;
-    const size_t hidden_bytes = kHidden * (kUseFP8 ? sizeof(__nv_fp8_storage_t) : sizeof(nv_bfloat16));
+    const size_t hidden_bytes = kHidden * (kUseFP8 ? sizeof(ep_fp8_storage_t) : EP_BF16_SIZE);
     const size_t hidden_int4 = hidden_bytes / sizeof(int4);
 
     // Message package: hidden data, FP8 scales, index at source
     // NOTES: currently we have 3 reserved int fields for future use
     using vec_t = typename std::conditional<kUseFP8, int2, int4>::type;
-    const size_t num_bytes_per_msg = sizeof(int4) + (kUseFP8 ? (kHidden + num_scales * sizeof(float)) : (kHidden * sizeof(nv_bfloat16)));
+    const size_t num_bytes_per_msg = sizeof(int4) + (kUseFP8 ? (kHidden + num_scales * sizeof(float)) : (kHidden * EP_BF16_SIZE));
     const size_t num_int4_per_msg = num_bytes_per_msg / sizeof(int4);
     EP_DEVICE_ASSERT(num_bytes_per_msg % sizeof(int4) == 0);
 
-    // IBGDA
-    auto raddr_array = reinterpret_cast<uint64_t*>(raddrs);
-    auto rkey_array = reinterpret_cast<uint32_t*>(rkeys);
-    auto ctx_array = reinterpret_cast<mlx5gda_qp_devctx*>(qp_devctxs);
+    // Communication context — platform dispatch is inside comm_device.cuh
+    const CommCtx comm_ctx = make_comm_ctx(
+        mxa_buffer, nvlink_available, ipc_peer_ptrs,
+        raddrs, rkeys, qp_devctxs,
+        rdma_send_signal_buffer, rdma_recv_signal_buffer,
+        rank, num_ranks, MAX_QP_COUNT);
     const size_t num_qp_per_rank = MAX_QP_COUNT / num_ranks;
 
     // Sending phase
@@ -209,7 +187,7 @@ dispatch(void* packed_recv_x, float* packed_recv_x_scales,
     // 1. The first-kind warps for FP8 cast and sending top-k tokens
     // 2. The last warp for reading `topk_idx` and count for per-expert information
     if (warp_id < num_warps - 1) {
-        constexpr int kNumElemsPerRead = sizeof(int4) / sizeof(nv_bfloat16);
+        constexpr int kNumElemsPerRead = sizeof(int4) / EP_BF16_SIZE;
         EP_DEVICE_ASSERT(kHidden % kNumElemsPerRead == 0);
         EP_STATIC_ASSERT(kNumElemsPerRead * 32 % kNumPerChannels == 0, "Invalid vectorization");
         const auto num_threads = (num_warps - 1) * 32;
@@ -250,11 +228,11 @@ dispatch(void* packed_recv_x, float* packed_recv_x_scales,
 
                         // Cast into send buffer
                         vec_t int2_value;
-                        auto fp8x2_values = reinterpret_cast<__nv_fp8x2_storage_t*>(&int2_value);
+                        auto fp8x2_values = reinterpret_cast<ep_fp8x2_storage_t*>(&int2_value);
                         #pragma unroll
                         for (int j = 0; j < kNumElemsPerRead; j += 2) {
                             float2 fp32x2 = {fp32_values[j] * scale, fp32_values[j + 1] * scale};
-                            fp8x2_values[j / 2] = __nv_cvt_float2_to_fp8x2(fp32x2, __NV_SATFINITE, __NV_E4M3);
+                            fp8x2_values[j / 2] = ep_cvt_float2_to_fp8x2(fp32x2);
                         }
                         rdma_x_vec[i] = int2_value;
                     } else {
@@ -262,51 +240,49 @@ dispatch(void* packed_recv_x, float* packed_recv_x_scales,
                         rdma_x_vec[i] = *reinterpret_cast<vec_t*>(&int4_value);
                     }
                 }
-            asm volatile("bar.sync 1, %0;" :: "r"(num_threads));
+            mc_bar_sync(1, num_threads);
 
-            // Issue IBGDA sends
+            // Issue sends
             if (dst_expert_idx >= 0) {
                 int slot_idx = lane_id == 0 ? atomicAdd(atomic_counter_per_expert + dst_expert_idx, 1) : 0;
                 slot_idx = __shfl_sync(0xffffffff, slot_idx, 0);
                 const auto dst_rank = dst_expert_idx / num_local_experts;
                 const auto dst_expert_local_idx = dst_expert_idx % num_local_experts;
-                const auto src_ptr = reinterpret_cast<uint64_t>(rdma_x_src_idx);
-                const auto dst_ptr = reinterpret_cast<uint64_t>(rdma_recv_data_buffer) +
-                                     dst_expert_local_idx * num_ranks * num_max_dispatch_tokens_per_rank * num_bytes_per_msg +
-                                     rank * num_max_dispatch_tokens_per_rank * num_bytes_per_msg +
-                                     slot_idx * num_bytes_per_msg;
-                if (dst_rank != rank) {
-                    bool use_nvlink = nvlink_available[dst_rank] != 0;
-                    if (use_nvlink) {
-                        size_t offset = (char *)dst_ptr - (char *)(mxa_buffer);
-                        void* peer_dst_ptr = (char *)ipc_peer_ptrs[dst_rank] + offset;
-                        // NOTES: only 2 load iterations for 7K hidden with 8 unrolls
-                        const auto* src_int4_ptr = reinterpret_cast<const int4*>(src_ptr);
-                        const auto* dst_int4_ptr = reinterpret_cast<int4*>(peer_dst_ptr);
-                        UNROLLED_WARP_COPY(8, lane_id, num_int4_per_msg, dst_int4_ptr, src_int4_ptr, ld_nc_global, st_na_global);
-                    } else {
-                        if (lane_id == 0) {
-                            uint64_t req_rptr_actual = raddr_array[dst_rank] + ((char *)dst_ptr - (char *)(mxa_buffer));
-                            auto ctx = ctx_array + dst_rank * num_qp_per_rank + dst_expert_local_idx % num_qp_per_rank;
-                            device_mutex_lock_system(&ctx->mutex);
-                            __mlx5gda_device_write_rdma_write_wqe(ctx, src_ptr, device_byteswap(rkey_array[rank]), req_rptr_actual, device_byteswap(rkey_array[dst_rank]), num_bytes_per_msg);
-                            __mlx5gda_device_post_send_db(ctx);
-                            device_mutex_unlock_system(&ctx->mutex);
-                        }
-                    }
-                } else {
-                    // NOTES: only 2 load iterations for 7K hidden with 8 unrolls
+                const auto src_ptr = reinterpret_cast<const void*>(rdma_x_src_idx);
+                const auto dst_ptr = reinterpret_cast<void*>(
+                    reinterpret_cast<uint64_t>(rdma_recv_data_buffer) +
+                    dst_expert_local_idx * num_ranks * num_max_dispatch_tokens_per_rank * num_bytes_per_msg +
+                    rank * num_max_dispatch_tokens_per_rank * num_bytes_per_msg +
+                    slot_idx * num_bytes_per_msg);
+
+                void* write_dst = mc_route_put(comm_ctx, dst_rank, dst_ptr);
+                if (write_dst != nullptr) {
+                    // Local or P2P path — warp-cooperative copy
                     const auto* src_int4_ptr = reinterpret_cast<const int4*>(src_ptr);
-                    const auto* dst_int4_ptr = reinterpret_cast<int4*>(dst_ptr);
-                    UNROLLED_WARP_COPY(8, lane_id, num_int4_per_msg, dst_int4_ptr, src_int4_ptr, ld_nc_global, st_na_global);
+                    const auto* dst_int4_ptr = reinterpret_cast<int4*>(write_dst);
+                    mc_fence();
+                    UNROLLED_WARP_COPY(8, lane_id, num_int4_per_msg, dst_int4_ptr, src_int4_ptr, mc_ld_nc, mc_st_na);
+                    mc_fence();
+                } else {
+                    // IBGDA path — send directly from source buffer
+                    mc_rdma_put(comm_ctx, dst_expert_local_idx % num_qp_per_rank, dst_rank, num_qp_per_rank,
+                                      src_ptr, dst_ptr, num_bytes_per_msg, lane_id);
                 }
 
                 // Increase counter after finishing
                 __syncwarp();
-                lane_id == 0 ? atomic_add_release_global(atomic_finish_counter_per_expert + dst_expert_idx, 1) : 0;
+                lane_id == 0 ? mc_atomic_add_release(atomic_finish_counter_per_expert + dst_expert_idx, 1) : 0;
             }
         }
     } else if (warp_id == num_warps - 1) {
+#ifdef MOONCAKE_EP_USE_MUSA
+        // Participate in __syncthreads() barriers from data warps.
+        // Each token iteration in the send loop above calls
+        // __syncthreads() once; the count warp must match.
+        for (int token_idx = sm_id; token_idx < num_tokens; token_idx += num_sms) {
+            __syncthreads();
+        }
+#endif
         EP_DEVICE_ASSERT(num_sms > 1);
         if (sm_id == 0) {
             // The first SM is also responsible for cleaning the next buffer
@@ -318,7 +294,7 @@ dispatch(void* packed_recv_x, float* packed_recv_x_scales,
             __syncwarp();
             #pragma unroll
             for (int i = lane_id; i < num_experts; i += 32)
-                atomic_add_release_global(atomic_finish_counter_per_expert + i, FINISHED_SUM_TAG);
+                mc_atomic_add_release(atomic_finish_counter_per_expert + i, FINISHED_SUM_TAG);
         }
 
         // This SM should be responsible for some destination experts, read `topk_idx` for them
@@ -340,11 +316,11 @@ dispatch(void* packed_recv_x, float* packed_recv_x_scales,
             auto sum = warp_reduce_sum(expert_count[i - expert_begin_idx]);
             if (lane_id == 0) {
                 shared_num_tokens_sent_per_expert[i - expert_begin_idx] = sum;
-                atomic_add_release_global(atomic_finish_counter_per_expert + i, FINISHED_SUM_TAG - sum);
+                mc_atomic_add_release(atomic_finish_counter_per_expert + i, FINISHED_SUM_TAG - sum);
             }
         }
     }
-    __syncthreads();
+    mc_fence_barrier_fence();
 
     // Issue count sends
     if (responsible_expert_idx < num_experts and sub_warp_id == 0 and lane_id == 0) {
@@ -353,25 +329,13 @@ dispatch(void* packed_recv_x, float* packed_recv_x_scales,
         const auto num_tokens_sent = shared_num_tokens_sent_per_expert[responsible_expert_idx - sm_id * kNumWarpGroups];
 
         // Wait local sends issued and send expert counts
-        while (ld_acquire_global(atomic_finish_counter_per_expert + responsible_expert_idx) != FINISHED_SUM_TAG * 2);
+        while (mc_ld_acquire(atomic_finish_counter_per_expert + responsible_expert_idx) != FINISHED_SUM_TAG * 2);
         if (dst_rank != rank) {
-            bool use_nvlink = nvlink_available[dst_rank] != 0;
-            if (use_nvlink) {
-                int* signal_ptr = rdma_recv_signal_buffer + dst_expert_local_idx * num_ranks + rank;
-                size_t offset = (char *)signal_ptr - (char *)(mxa_buffer);
-                int* peer_signal_ptr = (int *)((char *)ipc_peer_ptrs[dst_rank] + offset);
-                st_na_release(peer_signal_ptr, -num_tokens_sent - 1);
-            } else {
-                uint64_t laddr = (uint64_t)((char *)(raddr_array[rank]) + ((char *)(rdma_send_signal_buffer + dst_expert_local_idx * num_ranks + rank) - (char *)(mxa_buffer)));
-                uint64_t rptr_actual = (uint64_t)((char *)(raddr_array[dst_rank]) + ((char *)(rdma_recv_signal_buffer + dst_expert_local_idx * num_ranks + rank) - (char *)(mxa_buffer)));
-                auto ctx = ctx_array + dst_rank * num_qp_per_rank + dst_expert_local_idx % num_qp_per_rank;
-                device_mutex_lock_system(&ctx->mutex);
-                __mlx5gda_device_write_rdma_atomic_add_wqe(ctx, -num_tokens_sent - 1, laddr, device_byteswap(rkey_array[rank]), rptr_actual, device_byteswap(rkey_array[dst_rank]));
-                __mlx5gda_device_post_send_db(ctx);
-                device_mutex_unlock_system(&ctx->mutex);
-            }
+            int* signal_ptr = rdma_recv_signal_buffer + dst_expert_local_idx * num_ranks + rank;
+            mc_red_add(comm_ctx, dst_rank, dst_expert_local_idx % num_qp_per_rank, num_qp_per_rank,
+                       signal_ptr, static_cast<int32_t>(-num_tokens_sent - 1));
         } else {
-            st_na_release(rdma_recv_signal_buffer + dst_expert_local_idx * num_ranks + rank, -num_tokens_sent - 1);
+            mc_st_release(rdma_recv_signal_buffer + dst_expert_local_idx * num_ranks + rank, -num_tokens_sent - 1);
         }
 
         // Clean workspace for next use
@@ -391,7 +355,7 @@ dispatch(void* packed_recv_x, float* packed_recv_x_scales,
 
     // For send-and-recv kernels, we need a grid sync for making `packed_recv_count` visible
     if (phases & LOW_LATENCY_SEND_PHASE)
-        cooperative_groups::this_grid().sync();
+        mc_grid_sync();
 
     // Receiving and packing
     if (responsible_expert_idx < num_experts) {
@@ -415,7 +379,7 @@ dispatch(void* packed_recv_x, float* packed_recv_x_scales,
         EP_STATIC_ASSERT(kNumWarpsPerGroup > 1, "Requires more than one warp per group");
         if (sub_warp_id == 1 and lane_id == 0) {
             unsigned long long start_time = clock64();
-            while ((num_recv_tokens = ld_acquire_sys_global(rdma_recv_signal_buffer + local_expert_idx * num_ranks + src_rank)) == 0) {
+            while ((num_recv_tokens = mc_ld_acquire(rdma_recv_signal_buffer + local_expert_idx * num_ranks + src_rank)) == 0) {
                 unsigned long long end_time = clock64();
                 if (timeout_ticks != -1 && end_time - start_time > timeout_ticks) {
                     active_ranks[src_rank] = 0;
@@ -431,36 +395,40 @@ dispatch(void* packed_recv_x, float* packed_recv_x_scales,
             shared_recv_token_begin_idx[warp_group_id] = recv_token_begin_idx;
             recv_range[src_rank] = pack2<int, int64_t>(num_recv_tokens, recv_token_begin_idx);
         }
-        asm volatile("bar.sync %0, %1;" :: "r"(warp_group_id + 2), "r"(kNumWarpsPerGroup * 32));
+        mc_bar_sync(warp_group_id + 2, kNumWarpsPerGroup * 32);
         num_recv_tokens = shared_num_recv_tokens[warp_group_id];
         recv_token_begin_idx = shared_recv_token_begin_idx[warp_group_id];
 
         // Copy tokens
         EP_DEVICE_ASSERT(num_scales <= 64);
+        mc_fence();
         for (int i = sub_warp_id; i < num_recv_tokens; i += kNumWarpsPerGroup) {
             // Copy source info
             const auto src_src_idx = reinterpret_cast<int*>(rdma_recv_x_uint8 + i * num_bytes_per_msg);
             if (lane_id == 0)
-                recv_src_info[recv_token_begin_idx + i] = ld_nc_global(src_src_idx);
+                recv_src_info[recv_token_begin_idx + i] = mc_ld_nc_s32(src_src_idx);
             __syncwarp();
 
             // Copy data
             // NOTES: only 2 load iterations for 7K hidden with 7 unrolls
             const auto src_data = reinterpret_cast<int4*>(reinterpret_cast<uint8_t*>(src_src_idx) + sizeof(int4));
             const auto dst_data = recv_x_int4 + (recv_token_begin_idx + i) * hidden_int4;
-            UNROLLED_WARP_COPY(7, lane_id, hidden_int4, dst_data, src_data, ld_nc_global, st_na_global);
+            mc_fence();
+            UNROLLED_WARP_COPY(7, lane_id, hidden_int4, dst_data, src_data, mc_ld_nc, mc_st_na);
 
             // Copy scales
             if (kUseFP8) {
                 const auto src_scales = reinterpret_cast<float*>(reinterpret_cast<uint8_t*>(src_data) + hidden_bytes);
                 const auto dst_scales = reinterpret_cast<float*>(recv_x_scales + recv_token_begin_idx + i);
                 const auto scale_stride = num_ranks * num_max_dispatch_tokens_per_rank;
-                auto scale_0 = lane_id < num_scales ? ld_nc_global(src_scales + lane_id) : 0;
-                auto scale_1 = (lane_id + 32) < num_scales ? ld_nc_global(src_scales + lane_id + 32) : 0;
+                auto scale_0 = lane_id < num_scales ? mc_ld_nc_f32(src_scales + lane_id) : 0;
+                auto scale_1 = (lane_id + 32) < num_scales ? mc_ld_nc_f32(src_scales + lane_id + 32) : 0;
                 lane_id < num_scales ? dst_scales[lane_id * scale_stride] = scale_0 : 0.0f;
                 (lane_id + 32) < num_scales ? dst_scales[(lane_id + 32) * scale_stride] = scale_1 : 0.0f;
             }
         }
+    } else {
+        mc_bar_sync(warp_group_id + 2, kNumWarpsPerGroup * 32);
     }
 }
 
@@ -480,7 +448,12 @@ void dispatch(void* packed_recv_x, float* packed_recv_x_scales,
               void* workspace, cudaStream_t stream, int64_t timeout_ticks, int phases) {
     constexpr int kNumMaxTopK = 11;
     constexpr int kNumWarpsPerGroup = 4;
+#ifdef MOONCAKE_EP_USE_MUSA
+    // MT S5000 benefits from slightly more CTAs while keeping enough warps for top-k<=11.
+    constexpr int kNumWarpGroups = 5;
+#else
     constexpr int kNumWarpGroups = 8;
+#endif
     EP_STATIC_ASSERT(kNumMaxTopK + 1 <= kNumWarpGroups * kNumWarpsPerGroup, "Too many top-k selections");
 
     const auto num_warps = kNumWarpGroups * kNumWarpsPerGroup;
@@ -517,7 +490,7 @@ LAUNCH_KERNEL(&cfg, dispatch_func, \
 }
 
 template <int kNumWarpGroups, int kNumWarpsPerGroup, int kHidden, int kNumMaxTopk>
-__global__ __launch_bounds__(kNumWarpGroups * kNumWarpsPerGroup * 32, 1) void
+__global__ EP_LAUNCH_BOUNDS(kNumWarpGroups * kNumWarpsPerGroup * 32, 1) void
 combine(void* combined_x, int32_t* active_ranks,
         void* mxa_buffer,
         int* rdma_send_signal_buffer, int* rdma_recv_signal_buffer,
@@ -545,17 +518,19 @@ combine(void* combined_x, int32_t* active_ranks,
     const auto responsible_expert_idx = sm_id * kNumWarpGroups + warp_group_id;
 
     // Data type staffs
-    constexpr int kNumElemsPerInt4 = sizeof(int4) / sizeof(nv_bfloat16);
+    constexpr int kNumElemsPerInt4 = sizeof(int4) / EP_BF16_SIZE;
     const size_t hidden_bf16_int4 = kHidden / kNumElemsPerInt4;
 
     // Message package
-    constexpr size_t num_bytes_per_slot = kHidden * sizeof(nv_bfloat16);
+    constexpr size_t num_bytes_per_slot = kHidden * EP_BF16_SIZE;
     EP_STATIC_ASSERT(num_bytes_per_slot % sizeof(int4) == 0, "Invalid vectorization");
 
-    // IBGDA
-    auto raddr_array = reinterpret_cast<uint64_t*>(raddrs);
-    auto rkey_array = reinterpret_cast<uint32_t*>(rkeys);
-    auto ctx_array = reinterpret_cast<mlx5gda_qp_devctx*>(qp_devctxs);
+    // Communication context — platform dispatch is inside comm_device.cuh
+    const CommCtx comm_ctx = make_comm_ctx(
+        mxa_buffer, nvlink_available, ipc_peer_ptrs,
+        raddrs, rkeys, qp_devctxs,
+        rdma_send_signal_buffer, rdma_recv_signal_buffer,
+        rank, num_ranks, MAX_QP_COUNT);
     const size_t num_qp_per_rank = MAX_QP_COUNT / num_ranks;
 
     // Sending phase
@@ -571,7 +546,7 @@ combine(void* combined_x, int32_t* active_ranks,
         // Notify before executing `int_p`
         __syncwarp();
         if (lane_id == 0)
-            atomic_add_release_global(atomic_clean_flag, num_experts);
+            mc_atomic_add_release(atomic_clean_flag, num_experts);
     }
 
     // Issue IBGDA sends
@@ -598,63 +573,43 @@ combine(void* combined_x, int32_t* active_ranks,
 
             // Copy directly to local rank, or copy to buffer and issue RDMA
             auto src_idx = __ldg(local_src_info + token_idx);
-            const auto buf_ptr = reinterpret_cast<int64_t>(rdma_send_x_vec_row);
-            const auto dst_ptr = reinterpret_cast<uint64_t>(rdma_recv_data_buffer) + (global_expert_idx * num_max_dispatch_tokens_per_rank + src_idx) * num_bytes_per_slot;
-            if (dst_rank == rank) {
-                const auto dst_int4_ptr = reinterpret_cast<int4*>(dst_ptr);
-                UNROLLED_WARP_COPY(7, lane_id, hidden_bf16_int4, dst_int4_ptr, x_int4, ld_nc_global, st_na_global);
-            } else {
-                bool use_nvlink = nvlink_available[dst_rank] != 0;
-                if (use_nvlink) {
-                    size_t offset = (char *)dst_ptr - (char *)(mxa_buffer);
-                    void* peer_dst_ptr = (char *)ipc_peer_ptrs[dst_rank] + offset;
-                    const auto dst_int4_ptr = reinterpret_cast<int4*>(peer_dst_ptr);
-                    UNROLLED_WARP_COPY(7, lane_id, hidden_bf16_int4, dst_int4_ptr, x_int4, ld_nc_global, st_na_global);
-                } else {
-                    const auto buf_int4_ptr = reinterpret_cast<int4*>(buf_ptr);
-                    if (not zero_copy)
-                        UNROLLED_WARP_COPY(7, lane_id, hidden_bf16_int4, buf_int4_ptr, x_int4, ld_nc_global, st_na_global);
-                    __syncwarp();
+            const auto buf_ptr = reinterpret_cast<void*>(rdma_send_x_vec_row);
+            const auto dst_ptr = reinterpret_cast<void*>(
+                reinterpret_cast<uint64_t>(rdma_recv_data_buffer) +
+                (global_expert_idx * num_max_dispatch_tokens_per_rank + src_idx) * num_bytes_per_slot);
 
-                    if (lane_id == 0) {
-                        uint64_t req_rptr_actual = raddr_array[dst_rank] + ((char *)dst_ptr - (char *)(mxa_buffer));
-                        auto ctx = ctx_array + dst_rank * num_qp_per_rank + local_expert_idx % num_qp_per_rank;
-                        device_mutex_lock_system(&ctx->mutex);
-                        __mlx5gda_device_write_rdma_write_wqe(ctx, (uint64_t) buf_ptr, device_byteswap(rkey_array[rank]), req_rptr_actual, device_byteswap(rkey_array[dst_rank]), num_bytes_per_slot);
-                        __mlx5gda_device_post_send_db(ctx);
-                        device_mutex_unlock_system(&ctx->mutex);
-                    }
-                }
+            void* write_dst = mc_route_put(comm_ctx, dst_rank, dst_ptr);
+            if (write_dst != nullptr) {
+                // Local or P2P path — warp-cooperative copy
+                const auto dst_int4_ptr = reinterpret_cast<int4*>(write_dst);
+                UNROLLED_WARP_COPY(7, lane_id, hidden_bf16_int4, dst_int4_ptr, x_int4, mc_ld_nc, mc_st_na);
+                mc_fence();
+            } else {
+                // IBGDA path — stage to send buffer then RDMA write
+                const auto buf_int4_ptr = reinterpret_cast<int4*>(buf_ptr);
+                if (not zero_copy)
+                    UNROLLED_WARP_COPY(7, lane_id, hidden_bf16_int4, buf_int4_ptr, x_int4, mc_ld_nc, mc_st_na);
+                __syncwarp();
+                mc_rdma_put(comm_ctx, local_expert_idx % num_qp_per_rank, dst_rank, num_qp_per_rank,
+                                  buf_ptr, dst_ptr, num_bytes_per_slot, lane_id);
             }
         }
-
         // Put finishing flag
         EP_STATIC_ASSERT(kNumWarpsPerGroup > 1, "Requires more than one warp per group");
-        asm volatile("bar.sync %0, %1;" :: "r"(warp_group_id + 1), "r"(kNumWarpsPerGroup * 32));
+        mc_bar_sync(warp_group_id + 1, kNumWarpsPerGroup * 32);
         if (sub_warp_id == 1 and lane_id == 0) {
-            while (ld_acquire_global(atomic_clean_flag) == 0);
+            while (mc_ld_acquire(atomic_clean_flag) == 0);
             if (dst_rank != rank) {
-                bool use_nvlink = nvlink_available[dst_rank] != 0;
-                if (use_nvlink) {
-                    int* signal_ptr = rdma_recv_signal_buffer + global_expert_idx;
-                    size_t offset = (char *)signal_ptr - (char *)(mxa_buffer);
-                    int* peer_signal_ptr = (int *)((char *)ipc_peer_ptrs[dst_rank] + offset);
-                    st_na_release(peer_signal_ptr, 1);
-                } else {
-                    uint64_t laddr = (uint64_t)((char *)(raddr_array[rank]) + ((char *)(rdma_send_signal_buffer + global_expert_idx) - (char *)(mxa_buffer)));
-                    uint64_t req_rptr_actual = (uint64_t)((char *)(raddr_array[dst_rank]) + ((char *)(rdma_recv_signal_buffer + global_expert_idx) - (char *)(mxa_buffer)));
-                    auto ctx = ctx_array + dst_rank * num_qp_per_rank + local_expert_idx % num_qp_per_rank;
-                    device_mutex_lock_system(&ctx->mutex);
-                    __mlx5gda_device_write_rdma_atomic_add_wqe(ctx, 1, laddr, device_byteswap(rkey_array[rank]), req_rptr_actual, device_byteswap(rkey_array[dst_rank]));
-                    __mlx5gda_device_post_send_db(ctx);
-                    device_mutex_unlock_system(&ctx->mutex);
-                }
+                int* signal_ptr = rdma_recv_signal_buffer + global_expert_idx;
+                mc_signal(comm_ctx, dst_rank, local_expert_idx % num_qp_per_rank, num_qp_per_rank, signal_ptr, 1);
             } else {
-                st_na_release(rdma_recv_signal_buffer + global_expert_idx, 1);
+                mc_st_release(rdma_recv_signal_buffer + global_expert_idx, 1);
             }
-            atomic_add_release_global(atomic_clean_flag, -1);
+            mc_atomic_add_release(atomic_clean_flag, -1);
         }
         __syncwarp();
+    } else {
+        mc_bar_sync(warp_group_id + 1, kNumWarpsPerGroup * 32);
     }
 
     // Receiving phase
@@ -668,7 +623,7 @@ combine(void* combined_x, int32_t* active_ranks,
         EP_STATIC_ASSERT(kNumWarpsPerGroup > 1, "Invalid number of warps per group");
         if (sub_warp_id == 0 and lane_id == 0) {
             unsigned long long start_time = clock64();
-            while (ld_acquire_sys_global(rdma_recv_signal_buffer + responsible_expert_idx) == 0) {
+            while (mc_ld_acquire(rdma_recv_signal_buffer + responsible_expert_idx) == 0) {
                 unsigned long long end_time = clock64();
                 if (timeout_ticks != -1 && end_time - start_time > timeout_ticks) {
                     active_ranks[src_rank] = 0;
@@ -679,13 +634,22 @@ combine(void* combined_x, int32_t* active_ranks,
             }
         }
     }
-    cooperative_groups::this_grid().sync();
+#ifdef MOONCAKE_EP_USE_MUSA
+    // mc_grid_sync() is a no-op on MUSA; use a block-wide fence/barrier before
+    // reduction so threads see peer writes.
+    __syncthreads();
+    mc_fence();
+    __syncthreads();
+#else
+    mc_grid_sync();
+#endif
 
     // Reduce tokens with FP8 cast
     EP_DEVICE_ASSERT(num_topk <= 32 and hidden_bf16_int4 <= num_threads);
     EP_STATIC_ASSERT(kHidden % (32 * kNumElemsPerInt4) == 0, "Invalid vectorization");
     if (thread_id < hidden_bf16_int4) {
         for (int token_idx = sm_id; token_idx < num_combined_tokens; token_idx += num_sms) {
+            mc_fence();
             // Read top-k indices and weights
             int reg_topk_idx[kNumMaxTopk];
             float reg_topk_weights[kNumMaxTopk];
@@ -698,12 +662,16 @@ combine(void* combined_x, int32_t* active_ranks,
             float combined_values[kNumElemsPerInt4] = {0.0f};
             #pragma unroll
             for (int i = 0; i < num_topk; ++ i) if (reg_topk_idx[i] >= 0) {
+                // Skip experts on inactive ranks (timed out during combine recv)
+                int expert_src_rank = reg_topk_idx[i] / num_local_experts;
+                if (!active_ranks[expert_src_rank])
+                    continue;
                 // Read from sources
                 auto rdma_buffer_type = reinterpret_cast<const int*>(reinterpret_cast<uint8_t*>(rdma_recv_data_buffer) + (reg_topk_idx[i] * num_max_dispatch_tokens_per_rank + token_idx) * num_bytes_per_slot);
                 auto rdma_buffer_row = reinterpret_cast<const uint8_t*>(rdma_buffer_type);
 
                 // Reduce
-                auto x_vec = ld_nc_global(reinterpret_cast<const int4*>(rdma_buffer_row) + thread_id);
+                auto x_vec = mc_ld_nc(reinterpret_cast<const int4*>(rdma_buffer_row) + thread_id);
                 const auto x_bf16 = reinterpret_cast<nv_bfloat16*>(&x_vec);
                 #pragma unroll
                 for (int j = 0; j < kNumElemsPerInt4; ++ j)
@@ -770,4 +738,4 @@ LAUNCH_KERNEL(&cfg, combine_func, \
 #undef COMBINE_LAUNCH_CASE
 }
 
-} // namespace mooncake
\ No newline at end of file
+} // namespace mooncake
diff --git a/mooncake-ep/tests/test_ep_grid.py b/mooncake-ep/tests/test_ep_grid.py
new file mode 100644
index 00000000..25b4b173
--- /dev/null
+++ b/mooncake-ep/tests/test_ep_grid.py
@@ -0,0 +1,502 @@
+import os
+import itertools
+import unittest
+
+import torch
+import torch.distributed as dist
+import torch.multiprocessing as mp
+import torch.testing as testing
+import faulthandler
+import traceback
+
+from mooncake.mooncake_ep_buffer import Buffer
+import mooncake.pg as pg
+
+
+def using_musa_backend() -> bool:
+    return os.getenv("MOONCAKE_EP_USE_MUSA", "").upper() in {
+        "1",
+        "ON",
+        "TRUE",
+        "YES",
+    }
+
+
+def import_torchada_if_needed():
+    if not using_musa_backend():
+        return
+    try:
+        import torchada  # noqa: F401 — maps torch.cuda.* to torch.musa.* on MUSA
+    except ImportError as exc:
+        raise unittest.SkipTest(
+            "torchada is required when MOONCAKE_EP_USE_MUSA is enabled"
+        ) from exc
+
+
+def dequantize_fp8(x_fp8: torch.Tensor, scales: torch.Tensor) -> torch.Tensor:
+    hidden = x_fp8.shape[-1]
+    x_view = x_fp8.reshape(-1, hidden // 128, 128).float()
+    scales_view = scales.reshape(-1, hidden // 128, 1).float()
+    dequantized = (x_view * scales_view).reshape(x_fp8.shape)
+    return dequantized.to(torch.bfloat16)
+
+
+def run_test_iteration(
+    group: dist.ProcessGroup,
+    rank: int,
+    num_ranks: int,
+    max_tokens: int,
+    hidden: int,
+    num_experts: int,
+    top_k: int,
+    use_fp8: bool,
+    zero_copy: bool,
+    async_finish: bool,
+    return_recv_hook: bool,
+    use_fallback: bool,
+    fail_rank: int,
+    buf: Buffer = None,
+):
+    assert not (async_finish and return_recv_hook), (
+        "Should be filtered out by generate_tests."
+    )
+
+    torch.manual_seed(2026 + rank)
+    scale = 1.0 - 0.05 * (rank / num_ranks)
+    num_tokens = int(max_tokens * scale)
+
+    # Prepare test data
+    x = torch.randn(num_tokens, hidden, dtype=torch.bfloat16)
+    scores = torch.randn((num_tokens, num_experts), dtype=torch.float32)
+    topk_idx = torch.topk(scores, top_k, dim=-1)[1]
+    topk_weights = torch.softmax(
+        torch.rand(num_tokens, top_k, dtype=torch.float32), dim=-1
+    )
+    active_ranks = torch.ones((num_ranks,), dtype=torch.int32)
+
+    # Prepare expected result
+    def get_mock_factor(expert_id):
+        return expert_id * 0.1 + 1.0
+
+    # Since `get_mock_factor` is simply *0.1 + 1, we can make it more efficient
+    # instead of the nested loops below:
+    #
+    #   for i in range(num_tokens):
+    #       for j in range(top_k):
+    #           expert_id = topk_idx[i, j].item()
+    #           w = topk_weights[i, j].item()
+    #           mock_expert_factor = get_mock_factor(expert_id)
+    #           if fail_rank != -1 and (expert_id // num_local_experts == fail_rank):
+    #               continue
+    #           expected_out[i] += (x[i] * mock_expert_factor) * w
+    num_local_experts = num_experts // num_ranks
+    factors = get_mock_factor(topk_idx)
+    if fail_rank != -1:
+        expert_owners = topk_idx // num_local_experts
+        valid_mask = expert_owners != fail_rank
+        factors = factors * valid_mask.to(factors.dtype)
+    sum_weights = (factors * topk_weights).sum(dim=1, keepdim=True)
+    expected_out = x * sum_weights
+    expected_out = expected_out.to(torch.bfloat16)
+
+    # Initialize the buffer.
+    num_ep_buffer_bytes = Buffer.get_ep_buffer_size_hint(
+        max_tokens, hidden, num_ranks, num_experts
+    )
+    if buf is None:
+        buf = Buffer(group, num_ep_buffer_bytes)
+
+    if use_fallback:
+        buf._use_fallback = True
+
+    # 5s timeout if we simulate a failed rank
+    timeout_us = 5 * 1_000_000 if fail_rank != -1 else -1
+
+    dist.barrier(group)
+
+    if rank == fail_rank:
+        os._exit(0)
+
+    # Dispatch
+    recv_x, recv_count, handle, event, hook = buf.dispatch(
+        x,
+        topk_idx,
+        active_ranks,
+        num_max_dispatch_tokens_per_rank=max_tokens,
+        num_experts=num_experts,
+        timeout_us=timeout_us,
+        use_fp8=use_fp8,
+        async_finish=async_finish,
+        return_recv_hook=return_recv_hook,
+    )
+
+    if return_recv_hook:
+        hook()
+    if async_finish:
+        event.current_stream_wait()
+
+    torch.cuda.synchronize()
+    # Fault-tolerance check
+    if fail_rank != -1:
+        assert active_ranks[fail_rank].item() == 0, (
+            f"[Rank {rank}] Failed rank {fail_rank} is not recorded in active_ranks. "
+            f"active_ranks: {active_ranks}"
+        )
+        assert active_ranks.sum().item() == active_ranks.numel() - 1, (
+            f"[Rank {rank}] Expected exactly one failed rank {fail_rank}, but found "
+            f"{active_ranks.numel() - active_ranks.sum().item()} inactive ranks. "
+            f"Maybe the timeout is too small? "
+            f"active_ranks: {active_ranks}"
+        )
+
+    # Mock expert forward
+    if use_fp8:
+        recv_bf16 = dequantize_fp8(recv_x[0], recv_x[1])
+    else:
+        recv_bf16 = recv_x
+
+    expert_out = torch.empty_like(recv_bf16)
+
+    for le in range(num_local_experts):
+        expert_id = rank * num_local_experts + le
+        mock_expert_factor = get_mock_factor(expert_id)
+        expert_out[le] = recv_bf16[le] * mock_expert_factor
+
+    expert_out = expert_out.to(torch.bfloat16)
+
+    # Combine
+    if zero_copy:
+        cb_buf = buf.get_next_combine_buffer(handle)
+        cb_buf.copy_(expert_out)
+        expert_to_pass = cb_buf.contiguous()
+    else:
+        expert_to_pass = expert_out.contiguous()
+
+    out_tensor = torch.zeros_like(x)
+    combined_x, event, hook = buf.combine(
+        expert_to_pass,
+        topk_idx,
+        topk_weights,
+        active_ranks,
+        timeout_us=timeout_us,
+        handle=handle,
+        zero_copy=zero_copy,
+        async_finish=async_finish,
+        return_recv_hook=return_recv_hook,
+        out=out_tensor,
+    )
+
+    if return_recv_hook:
+        hook()
+    if async_finish:
+        event.current_stream_wait()
+
+    torch.cuda.synchronize()
+
+    testing.assert_close(
+        combined_x,
+        expected_out,
+        rtol=0.15 if use_fp8 else 5e-2,
+        atol=5e-3 if use_fp8 else 1e-3,
+        msg=lambda msg: f"[Rank {rank}] Combine Mismatch. {msg}",
+    )
+
+    torch.cuda.synchronize()
+    if fail_rank == -1:
+        dist.barrier(group)
+
+
+def worker(rank, world_size, config_dict):
+    import_torchada_if_needed()
+
+    # Device filter
+    device_filter = [
+        f
+        for f in os.getenv("DEVICE_FILTER", "mlx5_1,mlx5_2,mlx5_3,mlx5_4").split(",")
+        if f
+    ]
+    if device_filter:
+        pg.set_device_filter(device_filter)
+
+    torch.cuda.set_device(rank)
+    torch.set_default_dtype(torch.bfloat16)
+    torch.set_default_device("cuda")
+
+    dist.init_process_group(backend="mooncake", rank=rank, world_size=world_size)
+    group = dist.group.WORLD
+
+    try:
+        run_test_iteration(
+            group=group,
+            rank=rank,
+            num_ranks=world_size,
+            **config_dict,
+        )
+    except Exception:
+        traceback.print_exc()
+        raise
+
+    os._exit(0)
+
+
+def run_stale_data_test(
+    group: dist.ProcessGroup,
+    rank: int,
+    num_ranks: int,
+    max_tokens: int,
+    hidden: int,
+    num_experts: int,
+    top_k: int,
+):
+    """Verify combine does not read stale data from a previous round when
+    a peer rank times out.
+
+    Two rounds share the same buffer:
+      Round 1 — normal dispatch + combine (fills rdma_recv_data_buffer).
+      Round 2 — both ranks dispatch; rank 1 exits before combine; rank 0
+                combines with a short timeout.
+
+    Without the fix the reduction loop reads stale Round-1 data for rank 1's
+    experts, producing silently wrong results.
+    """
+    num_local_experts = num_experts // num_ranks
+    fail_rank = 1
+
+    def get_mock_factor(expert_id):
+        return expert_id * 0.1 + 1.0
+
+    num_ep_buffer_bytes = Buffer.get_ep_buffer_size_hint(
+        max_tokens, hidden, num_ranks, num_experts
+    )
+    buf = Buffer(group, num_ep_buffer_bytes)
+
+    # Round 1: normal round to fill buffer with data
+    run_test_iteration(
+        group=group,
+        rank=rank,
+        num_ranks=num_ranks,
+        max_tokens=max_tokens,
+        hidden=hidden,
+        num_experts=num_experts,
+        top_k=top_k,
+        use_fp8=False,
+        zero_copy=False,
+        async_finish=True,
+        return_recv_hook=False,
+        use_fallback=False,
+        fail_rank=-1,
+        buf=buf,
+    )
+
+    # Round 2: dispatch, then rank 1 exits before combine
+    torch.manual_seed(200 + rank)
+    scale = 1.0 - 0.05 * (rank / num_ranks)
+    num_tokens = int(max_tokens * scale)
+    x = torch.randn(num_tokens, hidden, dtype=torch.bfloat16)
+    scores = torch.randn((num_tokens, num_experts), dtype=torch.float32)
+    topk_idx = torch.topk(scores, top_k, dim=-1)[1]
+    topk_weights = torch.softmax(
+        torch.rand(num_tokens, top_k, dtype=torch.float32), dim=-1
+    )
+    active_ranks = torch.ones((num_ranks,), dtype=torch.int32)
+
+    # Expected output excluding fail_rank's experts
+    factors = get_mock_factor(topk_idx)
+    valid_mask = (topk_idx // num_local_experts) != fail_rank
+    factors = factors * valid_mask.to(factors.dtype)
+    expected = (x * (factors * topk_weights).sum(dim=1, keepdim=True)).to(
+        torch.bfloat16
+    )
+
+    dist.barrier(group)
+
+    recv_x, _, handle, event, _ = buf.dispatch(
+        x,
+        topk_idx,
+        active_ranks,
+        num_max_dispatch_tokens_per_rank=max_tokens,
+        num_experts=num_experts,
+        timeout_us=-1,
+        use_fp8=False,
+        async_finish=True,
+    )
+    event.current_stream_wait()
+    torch.cuda.synchronize()
+
+    dist.barrier(group)
+
+    if rank == fail_rank:
+        os._exit(0)
+
+    expert_out = torch.empty_like(recv_x)
+    for le in range(num_local_experts):
+        eid = rank * num_local_experts + le
+        expert_out[le] = recv_x[le] * get_mock_factor(eid)
+    expert_out = expert_out.to(torch.bfloat16)
+
+    out = torch.zeros_like(x)
+    combined_x, event, _ = buf.combine(
+        expert_out,
+        topk_idx,
+        topk_weights,
+        active_ranks,
+        timeout_us=5_000,
+        handle=handle,
+        async_finish=True,
+        out=out,
+    )
+    event.current_stream_wait()
+    torch.cuda.synchronize()
+
+    assert active_ranks[fail_rank].item() == 0, (
+        f"[Rank {rank}] active_ranks[{fail_rank}] should be 0 after timeout"
+    )
+    assert not torch.isnan(combined_x).any().item()
+    testing.assert_close(
+        combined_x,
+        expected,
+        rtol=5e-2,
+        atol=1e-3,
+        msg=lambda m: f"[Rank {rank}] Stale data read in combine reduction. {m}",
+    )
+
+
+def stale_data_worker(rank, world_size):
+    import_torchada_if_needed()
+
+    device_filter = [
+        f
+        for f in os.getenv("DEVICE_FILTER", "mlx5_1,mlx5_2,mlx5_3,mlx5_4").split(",")
+        if f
+    ]
+    if device_filter:
+        pg.set_device_filter(device_filter)
+
+    torch.cuda.set_device(rank)
+    torch.set_default_dtype(torch.bfloat16)
+    torch.set_default_device("cuda")
+
+    dist.init_process_group(backend="mooncake", rank=rank, world_size=world_size)
+    group = dist.group.WORLD
+
+    try:
+        run_stale_data_test(
+            group=group,
+            rank=rank,
+            num_ranks=world_size,
+            max_tokens=256,
+            hidden=2048,
+            num_experts=288,
+            top_k=8,
+        )
+    except Exception:
+        traceback.print_exc()
+        raise
+
+    os._exit(0)
+
+
+class TestMooncakeEPBuffer(unittest.TestCase):
+    def setUp(self):
+        import_torchada_if_needed()
+        self.world_size = torch.cuda.device_count()
+        os.environ["MASTER_ADDR"] = "127.0.0.1"
+        os.environ["MASTER_PORT"] = "29500"
+
+    def run_single_config(self, config_dict):
+        mp.spawn(
+            worker,
+            args=(self.world_size, config_dict),
+            nprocs=self.world_size,
+            join=True,
+            daemon=False,
+        )
+
+    def test_combine_stale_data(self):
+        mp.spawn(
+            stale_data_worker,
+            args=(self.world_size,),
+            nprocs=self.world_size,
+            join=True,
+            daemon=False,
+        )
+
+
+def make_test_name(cfg):
+    parts = ["test_ep"]
+
+    flags = []
+    if cfg["use_fp8"]:
+        flags.append("fp8")
+    if cfg["zero_copy"]:
+        flags.append("0copy")
+    if cfg["async_finish"]:
+        flags.append("async")
+    if cfg["return_recv_hook"]:
+        flags.append("hook")
+    if cfg["use_fallback"]:
+        flags.append("fallback")
+
+    if flags:
+        parts.append("_".join(flags))
+    else:
+        parts.append("base")
+
+    fail_rank = cfg["fail_rank"]
+    if fail_rank == -1:
+        parts.append("nofail")
+    else:
+        parts.append(f"fail{fail_rank}")
+
+    parts.append(f"t{cfg['max_tokens']}")
+    parts.append(f"h{cfg['hidden']}")
+    parts.append(f"e{cfg['num_experts']}")
+    parts.append(f"k{cfg['top_k']}")
+
+    return "_".join(parts)
+
+
+def generate_tests():
+    test_grid = {
+        "use_fp8": [False, True],
+        "zero_copy": [False, True],
+        "async_finish": [False, True],
+        "return_recv_hook": [False, True],
+        "use_fallback": [False, True],
+        "fail_rank": [-1, 1],
+        "shapes": [
+            {"max_tokens": 256, "hidden": 2048, "num_experts": 288, "top_k": 8},
+        ],
+    }
+
+    keys = list(test_grid.keys())
+    config_tuples = list(itertools.product(*[test_grid[k] for k in keys]))
+
+    for t in config_tuples:
+        raw_dict = dict(zip(keys, t))
+
+        if raw_dict["async_finish"] and raw_dict["return_recv_hook"]:
+            continue
+
+        # Flatten
+        config_dict = {}
+        for k, v in raw_dict.items():
+            if k == "shapes":
+                config_dict.update(v)
+            else:
+                config_dict[k] = v
+
+        # Make a test name
+        test_name = make_test_name(config_dict)
+
+        def test_func(self, cfg=config_dict):
+            self.run_single_config(cfg)
+
+        test_func.__name__ = test_name
+        setattr(TestMooncakeEPBuffer, test_name, test_func)
+
+
+if __name__ == "__main__":
+    faulthandler.enable()
+    generate_tests()
+    unittest.main(verbosity=2)
diff --git a/mooncake-integration/CMakeLists.txt b/mooncake-integration/CMakeLists.txt
index 7ecfb156..735015a3 100644
--- a/mooncake-integration/CMakeLists.txt
+++ b/mooncake-integration/CMakeLists.txt
@@ -1,34 +1,40 @@
 file(GLOB SOURCES "*.cpp")
 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
-)
+  COMMAND ${PYTHON_EXECUTABLE} -c
+          "import sys; print([s for s in sys.path if 'packages' in s][0])"
+  OUTPUT_VARIABLE PYTHON_SYS_PATH)
 string(STRIP ${PYTHON_SYS_PATH} PYTHON_SYS_PATH)
 
 if("${PYTHON_SYS_PATH}" STREQUAL "")
-    message(FATAL_ERROR "Python path is empty! Please check the python env.")
+  message(FATAL_ERROR "Python path is empty! Please check the python env.")
 endif()
 
-if (WITH_STORE)
-    include_directories("../mooncake-store/include")
-    include_directories("../mooncake-store/include/cachelib_memory_allocator")
+if(WITH_STORE)
+  include_directories("../mooncake-store/include")
+  include_directories("../mooncake-store/include/cachelib_memory_allocator")
+  include_directories("../mooncake-store/include/engram")
 
-    include_directories("../mooncake-store/include/cachelib_memory_allocator/include")
-    include_directories("../mooncake-store/include/cachelib_memory_allocator/fake_include")
+  include_directories(
+    "../mooncake-store/include/cachelib_memory_allocator/include")
+  include_directories(
+    "../mooncake-store/include/cachelib_memory_allocator/fake_include")
 endif()
 
 include_directories("/usr/include/jsoncpp")
 include_directories("./")
 include_directories("../mooncake-transfer-engine/include")
 
-find_package(Python3 COMPONENTS Interpreter Development REQUIRED)
+find_package(
+  Python3
+  COMPONENTS Interpreter Development.Module
+  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
-)
+  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)
@@ -37,179 +43,183 @@ message("${PYTHON_SYS_PATH}")
 
 set(PYTHON_PACKAGE_NAME "mooncake")
 
-if (WITH_TE)
-    pybind11_add_module(engine ${SOURCES} ${CACHE_ALLOCATOR_SOURCES}
-        transfer_engine/transfer_engine_py.cpp
-    )
-    set_target_properties(engine PROPERTIES
-        INSTALL_RPATH "$ORIGIN"
-    )
+if(WITH_TE)
+  pybind11_add_module(engine ${SOURCES} ${CACHE_ALLOCATOR_SOURCES}
+                      transfer_engine/transfer_engine_py.cpp)
+  set_target_properties(engine PROPERTIES INSTALL_RPATH "$ORIGIN")
 
-    # Propagate EFA compile definition to engine target
-    if(USE_EFA)
-        target_compile_definitions(engine PRIVATE USE_EFA)
-    endif()
+  # Propagate EFA compile definition to engine target
+  if(USE_EFA)
+    target_compile_definitions(engine PRIVATE USE_EFA)
+  endif()
+  if(USE_CXI)
+    target_compile_definitions(engine PRIVATE USE_CXI)
+  endif()
 
-    target_link_libraries(engine PRIVATE
-        $<TARGET_OBJECTS:rpc_communicator>
-    )
+  # Propagate USE_TENT compile definition to engine target
+  if(USE_TENT)
+    target_compile_definitions(engine PRIVATE USE_TENT)
+  endif()
 
-    target_link_libraries(engine PRIVATE ${Python3_LIBRARIES})
-    target_include_directories(engine PRIVATE ${Python3_INCLUDE_DIRS})
-    if (USE_ASCEND_DIRECT)
-        target_link_libraries(engine PUBLIC
-            ascendcl
-            transfer_engine
-            glog::glog
-            gflags::gflags
-        )
-    else()
-        target_link_libraries(engine PUBLIC
-            transfer_engine
-            glog::glog
-            gflags::gflags
-        )
-    endif()
+  target_link_libraries(engine PRIVATE $<TARGET_OBJECTS:rpc_communicator>)
 
-    if (USE_CUDA)
-        find_package(CUDAToolkit REQUIRED)
-        target_link_libraries(engine PRIVATE CUDA::cudart)
-    endif()
+  target_include_directories(engine PRIVATE ${Python3_INCLUDE_DIRS})
+  if(USE_ASCEND_DIRECT)
+    target_link_libraries(engine PUBLIC ascendcl transfer_engine glog::glog
+                                        gflags::gflags)
+  else()
+    target_link_libraries(engine PUBLIC transfer_engine glog::glog
+                                        gflags::gflags)
+  endif()
+
+  if(USE_CUDA)
+    find_package(CUDAToolkit REQUIRED)
+    target_link_libraries(engine PRIVATE CUDA::cudart)
+  endif()
 endif()
 
-
-set(ALLOCATOR_SO_PATH "${CMAKE_BINARY_DIR}/mooncake-transfer-engine/nvlink-allocator/nvlink_allocator.so")
-set(UBSHMEM_ALLOCATOR_SO_PATH "${CMAKE_BINARY_DIR}/mooncake-transfer-engine/ubshmem-allocator/ubshmem_fabric_allocator.so")
+set(ALLOCATOR_SO_PATH
+    "${CMAKE_BINARY_DIR}/mooncake-transfer-engine/nvlink-allocator/nvlink_allocator.so"
+)
+set(UBSHMEM_ALLOCATOR_SO_PATH
+    "${CMAKE_BINARY_DIR}/mooncake-transfer-engine/ubshmem-allocator/ubshmem_fabric_allocator.so"
+)
 
 if(USE_MNNVL)
-    message(STATUS "USE_MNNVL is enabled, nvlink_allocator.so will be installed in the Python package")
-    install(FILES
-        "${ALLOCATOR_SO_PATH}"
-        DESTINATION ${PYTHON_SYS_PATH}/${PYTHON_PACKAGE_NAME}
-    )
+  message(
+    STATUS
+      "USE_MNNVL is enabled, nvlink_allocator.so will be installed in the Python package"
+  )
+  install(FILES "${ALLOCATOR_SO_PATH}"
+          DESTINATION ${PYTHON_SYS_PATH}/${PYTHON_PACKAGE_NAME})
 endif()
 
 if(USE_UBSHMEM)
-    message(STATUS "USE_UBSHMEM is enabled, ubshmem_fabric_allocator.so will be installed in the Python package")
-    install(FILES
-        "${UBSHMEM_ALLOCATOR_SO_PATH}"
-        DESTINATION ${PYTHON_SYS_PATH}/${PYTHON_PACKAGE_NAME}
-    )
+  message(
+    STATUS
+      "USE_UBSHMEM is enabled, ubshmem_fabric_allocator.so will be installed in the Python package"
+  )
+  install(FILES "${UBSHMEM_ALLOCATOR_SO_PATH}"
+          DESTINATION ${PYTHON_SYS_PATH}/${PYTHON_PACKAGE_NAME})
 endif()
 
-if (WITH_STORE)
-    pybind11_add_module(store ${SOURCES} ${CACHE_ALLOCATOR_SOURCES}
-        store/store_py.cpp
-        integration_utils.h
-    )
-    set_target_properties(store PROPERTIES
-        INSTALL_RPATH "$ORIGIN"
-    )
-    if (USE_ASCEND_DIRECT)
-        target_link_libraries(store PUBLIC
-            ascendcl
-            transfer_engine
-            glog::glog
-            gflags::gflags
-            mooncake_store
-            cachelib_memory_allocator
-        )
-    else()
-        target_link_libraries(store PUBLIC
-            transfer_engine
-            glog::glog
-            gflags::gflags
-            mooncake_store
-            cachelib_memory_allocator
-        )
-    endif()
+if(WITH_STORE)
+  pybind11_add_module(
+    store
+    ${SOURCES}
+    ${CACHE_ALLOCATOR_SOURCES}
+    store/store_py.cpp
+    store/buffer_pool.cpp
+    store/engram_store_py.cpp
+    integration_utils.h)
+  set_target_properties(store PROPERTIES INSTALL_RPATH "$ORIGIN")
+  if(USE_ASCEND_DIRECT)
+    target_link_libraries(
+      store PUBLIC ascendcl transfer_engine glog::glog gflags::gflags
+                   mooncake_store cachelib_memory_allocator)
+  else()
+    target_link_libraries(
+      store PUBLIC transfer_engine glog::glog gflags::gflags mooncake_store
+                   cachelib_memory_allocator)
+  endif()
 endif()
 
 message("${PYTHON_SYS_PATH}")
 
 file(WRITE ${CMAKE_CURRENT_BINARY_DIR}/${PYTHON_PACKAGE_NAME}/__init__.py
-    "# Auto-generated by CMake\n"
-)
+     "# Auto-generated by CMake\n")
 
-if (USE_MNNVL)
-    message(STATUS "allocator.py will be installed in the Python package")
-    install(FILES
-        "${CMAKE_CURRENT_SOURCE_DIR}/allocator.py"
-        DESTINATION ${PYTHON_SYS_PATH}/${PYTHON_PACKAGE_NAME}
-    )
+if(USE_MNNVL)
+  message(STATUS "allocator.py will be installed in the Python package")
+  install(FILES "${CMAKE_CURRENT_SOURCE_DIR}/allocator.py"
+          DESTINATION ${PYTHON_SYS_PATH}/${PYTHON_PACKAGE_NAME})
 endif()
 
-if (USE_UBSHMEM)
-    message(STATUS "allocator_ascend_npu.py will be installed in the Python package")
-    install(FILES
-        "${CMAKE_CURRENT_SOURCE_DIR}/allocator_ascend_npu.py"
-        DESTINATION ${PYTHON_SYS_PATH}/${PYTHON_PACKAGE_NAME}
-    )
+if(USE_UBSHMEM)
+  message(
+    STATUS "allocator_ascend_npu.py will be installed in the Python package")
+  install(FILES "${CMAKE_CURRENT_SOURCE_DIR}/allocator_ascend_npu.py"
+          DESTINATION ${PYTHON_SYS_PATH}/${PYTHON_PACKAGE_NAME})
+endif()
+
+if(USE_MNNVL OR USE_UBSHMEM)
+  message(
+    STATUS "fabric_allocator_utils.py will be installed in the Python package")
+  install(FILES "${CMAKE_CURRENT_SOURCE_DIR}/fabric_allocator_utils.py"
+          DESTINATION ${PYTHON_SYS_PATH}/${PYTHON_PACKAGE_NAME})
 endif()
 
 if(USE_INTRA_NVLINK)
-    message(STATUS "USE_INTRA_NVLINK is enabled, IntraNode nvlink is now activated")
+  message(
+    STATUS "USE_INTRA_NVLINK is enabled, IntraNode nvlink is now activated")
 endif()
 
-install(FILES "${CMAKE_CURRENT_SOURCE_DIR}/store/async_store.py" DESTINATION ${PYTHON_SYS_PATH}/${PYTHON_PACKAGE_NAME})
+install(FILES "${CMAKE_CURRENT_SOURCE_DIR}/store/async_store.py"
+        DESTINATION ${PYTHON_SYS_PATH}/${PYTHON_PACKAGE_NAME})
 
-#Install Python scripts from mooncake - wheel / mooncake / directory
-install(FILES
+# Install Python scripts from mooncake - wheel / mooncake / directory
+install(
+  FILES
     "${CMAKE_CURRENT_SOURCE_DIR}/../mooncake-wheel/mooncake/http_metadata_server.py"
     "${CMAKE_CURRENT_SOURCE_DIR}/../mooncake-wheel/mooncake/cli_bench.py"
     "${CMAKE_CURRENT_SOURCE_DIR}/../mooncake-wheel/mooncake/transfer_engine_topology_dump.py"
     "${CMAKE_CURRENT_SOURCE_DIR}/../mooncake-wheel/mooncake/mooncake_connector_v1.py"
     "${CMAKE_CURRENT_SOURCE_DIR}/../mooncake-wheel/mooncake/vllm_v1_proxy_server.py"
-    DESTINATION ${PYTHON_SYS_PATH}/${PYTHON_PACKAGE_NAME}
-)
+  DESTINATION ${PYTHON_SYS_PATH}/${PYTHON_PACKAGE_NAME})
 
-if (WITH_STORE)
-    install(FILES
-        "${CMAKE_CURRENT_SOURCE_DIR}/../mooncake-wheel/mooncake/mooncake_store_service.py"
-        "${CMAKE_CURRENT_SOURCE_DIR}/../mooncake-wheel/mooncake/mooncake_config.py"
-        "${CMAKE_CURRENT_SOURCE_DIR}/../mooncake-wheel/mooncake/cli.py"
-        DESTINATION ${PYTHON_SYS_PATH}/${PYTHON_PACKAGE_NAME}
-    )
+if(WITH_STORE)
+  install(
+    FILES
+      "${CMAKE_CURRENT_SOURCE_DIR}/../mooncake-wheel/mooncake/mooncake_store_service.py"
+      "${CMAKE_CURRENT_SOURCE_DIR}/../mooncake-wheel/mooncake/mooncake_config.py"
+      "${CMAKE_CURRENT_SOURCE_DIR}/../mooncake-wheel/mooncake/cli.py"
+    DESTINATION ${PYTHON_SYS_PATH}/${PYTHON_PACKAGE_NAME})
+
+  if(USE_NOF)
+    install(
+      FILES
+        "${CMAKE_CURRENT_SOURCE_DIR}/../mooncake-wheel/mooncake/mooncake_ssd_register.py"
+        "${CMAKE_CURRENT_SOURCE_DIR}/../mooncake-wheel/mooncake/mooncake_ssd_unregister.py"
+        "${CMAKE_CURRENT_SOURCE_DIR}/../mooncake-wheel/mooncake/spdk_tgt_create.py"
+      DESTINATION ${PYTHON_SYS_PATH}/${PYTHON_PACKAGE_NAME})
+  endif()
 endif()
 
-install(
-    DIRECTORY
-        ${CMAKE_CURRENT_BINARY_DIR}/${PYTHON_PACKAGE_NAME}/
-    DESTINATION
-        ${PYTHON_SYS_PATH}/${PYTHON_PACKAGE_NAME}
-)
+install(DIRECTORY ${CMAKE_CURRENT_BINARY_DIR}/${PYTHON_PACKAGE_NAME}/
+        DESTINATION ${PYTHON_SYS_PATH}/${PYTHON_PACKAGE_NAME})
 
 install(
-    CODE "
+  CODE "
         execute_process(COMMAND chmod 767 \"${PYTHON_SYS_PATH}/${PYTHON_PACKAGE_NAME}\")
         execute_process(COMMAND chmod 766 \"${PYTHON_SYS_PATH}/${PYTHON_PACKAGE_NAME}/__init__.py\")
-    "
-)
+    ")
 
-if (WITH_STORE)
-    install(TARGETS store DESTINATION ${PYTHON_SYS_PATH}/${PYTHON_PACKAGE_NAME})
+if(WITH_STORE)
+  install(TARGETS store DESTINATION ${PYTHON_SYS_PATH}/${PYTHON_PACKAGE_NAME})
 endif()
 
-if (WITH_TE)
-    install(TARGETS engine DESTINATION ${PYTHON_SYS_PATH}/${PYTHON_PACKAGE_NAME})
-    install(TARGETS asio_shared DESTINATION ${PYTHON_SYS_PATH}/${PYTHON_PACKAGE_NAME})
+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 "
+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\"
diff --git a/mooncake-integration/allocator.py b/mooncake-integration/allocator.py
index 77fa1c98..f8804a1f 100644
--- a/mooncake-integration/allocator.py
+++ b/mooncake-integration/allocator.py
@@ -3,20 +3,22 @@ import logging
 import os
 import threading
 from importlib import resources
-from typing import Dict, Final, Optional
+from typing import Dict, Final
 from enum import IntEnum
 
 from torch import device as torch_device
 from torch.cuda.memory import CUDAPluggableAllocator
 
+from .fabric_allocator_utils import get_mooncake_so_path, probe_allocator_backend
+
 logger = logging.getLogger(__name__)
 
 
 class MemoryBackend(IntEnum):
-    USE_CUDAMALLOC  = 0
+    USE_CUDAMALLOC = 0
     USE_CUMEMCREATE = 1
-    UNKNOWN         = -1
-    UNSUPPORTED     = -2
+    UNKNOWN = -1
+    UNSUPPORTED = -2
 
 
 class NVLinkAllocator:
@@ -27,65 +29,34 @@ class NVLinkAllocator:
 
     @classmethod
     def _get_so_path(cls) -> str:
-        """Dynamically locate nvlink_allocator.so in the mooncake package installation"""
-        try:
-            # Attempt to locate package resource
-            with resources.path("mooncake", "nvlink_allocator.so") as so_path:
-                if so_path.exists():
-                    return str(so_path)
-        except (ImportError, FileNotFoundError, TypeError):
-            pass
-
-        # Fallback strategy: check in package location via import metadata
-        try:
-            import mooncake
-
-            base_path = os.path.dirname(os.path.abspath(mooncake.__file__))
-            so_path = os.path.join(base_path, "nvlink_allocator.so")
-            if os.path.exists(so_path):
-                return so_path
-        except (ImportError, FileNotFoundError, TypeError):
-            raise ImportError(
-                "SGLANG_MOONCAKE_CUSTOM_MEM_POOL require mooncake-transfer-engine >= 0.3.3.post2."
-            )
+        return get_mooncake_so_path(
+            "nvlink_allocator.so",
+            "SGLANG_MOONCAKE_CUSTOM_MEM_POOL require mooncake-transfer-engine >= 0.3.3.post2.",
+        )
 
     @classmethod
     def _probe_fabric_memory_support(cls, so_path: str) -> MemoryBackend:
-        """
-        Probe whether the system supports fabric memory by calling a C++ function
-        that attempts cuMemCreate with CU_MEM_HANDLE_TYPE_FABRIC.
-        We assume the shared library exports a symbol like:
-            extern "C" MemoryBackendType mc_probe_fabric_support(int device_id);
-        """
+        supported_type = probe_allocator_backend(
+            so_path,
+            "mc_allocator_probe",
+            ctypes.c_int,
+            int(MemoryBackend.UNSUPPORTED),
+        )
         try:
+            backend = MemoryBackend(supported_type)
+        except ValueError:
+            logger.info("Unknown Backend error")
+            return MemoryBackend.UNKNOWN
 
-            lib = ctypes.CDLL(so_path)
-
-            # Try to get the probe function
-            probe_func = lib.mc_probe_fabric_support
-            probe_func.argtypes = [ctypes.c_int]
-            probe_func.restype = ctypes.c_int
-
-            # Use device 0 for probing
-            dev_id = 0
-            supported_type = probe_func(dev_id)
-            if supported_type == MemoryBackend.USE_CUDAMALLOC:
-                logger.info(f"Use CudaMalloc fallback")
-            elif supported_type == MemoryBackend.USE_CUMEMCREATE:
-                logger.info(f"Supports Fabric Memory")
-            else:
-                logger.info("Unknown Backend error")
-            return supported_type
-
-        except AttributeError:
-            logger.warning(
-                "Symbol 'mc_probe_fabric_support' not found in nvlink_allocator.so. "
-                "Assuming fabric memory is NOT supported (you may need to update the library)."
-            )
-            return MemoryBackend.UNSUPPORTED
-        except Exception as e:
-            logger.warning(f"Failed to probe fabric memory support: {e}")
-            return MemoryBackend.UNSUPPORTED
+        if backend == MemoryBackend.USE_CUDAMALLOC:
+            logger.info("Use CudaMalloc fallback")
+        elif backend == MemoryBackend.USE_CUMEMCREATE:
+            logger.info("Supports Fabric Memory")
+        elif backend == MemoryBackend.UNSUPPORTED:
+            logger.info("Allocator backend probing is unsupported")
+        else:
+            logger.info("Unknown Backend error")
+        return backend
 
     @classmethod
     def detect_mem_backend(cls) -> MemoryBackend:
@@ -93,14 +64,15 @@ class NVLinkAllocator:
         if not cls._probe_done:
             with cls._lock:
                 if cls._probe_done:
-                    return
-                so_path = None
+                    return cls._supports_fabric
                 try:
-                    so_path = cls._get_so_path()
-                    # First try dedicated probe function
-                    cls._supports_fabric = cls._probe_fabric_memory_support(so_path)
+                    cls._supports_fabric = cls._probe_fabric_memory_support(
+                        cls._get_so_path()
+                    )
                 except Exception as e:
-                    logger.error(f"Critical error during fabric memory probe setup: {e}")
+                    logger.error(
+                        f"Critical error during fabric memory probe setup: {e}"
+                    )
                     cls._supports_fabric = MemoryBackend.UNSUPPORTED
 
                 cls._probe_done = True
@@ -112,7 +84,7 @@ class NVLinkAllocator:
             if device not in cls._instances:
                 so_path = cls._get_so_path()
                 cls._instances[device] = CUDAPluggableAllocator(
-                    so_path, "mc_nvlink_malloc", "mc_nvlink_free"
+                    so_path, "mc_allocator_malloc", "mc_allocator_free"
                 )
             return cls._instances[device]
 
diff --git a/mooncake-integration/allocator_ascend_npu.py b/mooncake-integration/allocator_ascend_npu.py
index c4f38fc3..bc4e0419 100644
--- a/mooncake-integration/allocator_ascend_npu.py
+++ b/mooncake-integration/allocator_ascend_npu.py
@@ -1,14 +1,13 @@
 import ctypes
 import logging
-import os
 import threading
-from importlib import resources
 from typing import Dict, Final
-import torch_npu
 
 from torch import device as torch_device
 from torch_npu.npu.memory import NPUPluggableAllocator
 
+from .fabric_allocator_utils import get_mooncake_so_path, probe_allocator_backend
+
 logger = logging.getLogger(__name__)
 
 
@@ -20,62 +19,26 @@ class UBShmemAllocator:
 
     @classmethod
     def _get_so_path(cls) -> str:
-        """Dynamically locate ubshmem_fabric_allocator.so in the mooncake package installation"""
-        try:
-            # Attempt to locate package resource
-            with resources.path("mooncake", "ubshmem_fabric_allocator.so") as so_path:
-                if so_path.exists():
-                    return str(so_path)
-        except (ImportError, FileNotFoundError, TypeError):
-            pass
-
-        # Fallback strategy: check in package location via import metadata
-        try:
-            import mooncake
-
-            base_path = os.path.dirname(os.path.abspath(mooncake.__file__))
-            so_path = os.path.join(base_path, "ubshmem_fabric_allocator.so")
-            if os.path.exists(so_path):
-                return so_path
-        except (ImportError, FileNotFoundError, TypeError):
-            raise ImportError(
-                "UBShmemAllocator require mooncake-transfer-engine with USE_UBSHMEM enabled."
-            )
+        return get_mooncake_so_path(
+            "ubshmem_fabric_allocator.so",
+            "UBShmemAllocator require mooncake-transfer-engine with USE_UBSHMEM enabled.",
+        )
 
     @classmethod
     def _probe_fabric_memory_support(cls, so_path: str) -> bool:
-        """
-        Probe whether the system supports fabric memory by calling a C++ function
-        that attempts aclrtMallocPhysical with fabric memory support.
-        The shared library exports a symbol like:
-            extern "C" bool mc_probe_ub_fabric_support(int device_id);
-        """
-        try:
-            lib = ctypes.CDLL(so_path)
-
-            # Try to get the probe function
-            probe_func = lib.mc_probe_ub_fabric_support
-            probe_func.argtypes = [ctypes.c_int]
-            probe_func.restype = ctypes.c_bool
-
-            # Use device 0 for probing
-            dev_id = 0
-            supported = probe_func(dev_id)
-            if supported:
-                logger.info(f"Supports Fabric Memory with aclMallocPhysical")
-            else:
-                logger.info("Fabric memory not supported")
-            return supported
-
-        except AttributeError:
-            logger.warning(
-                "Symbol 'mc_probe_ub_fabric_support' not found in ubshmem_fabric_allocator.so. "
-                "Assuming fabric memory is NOT supported (you may need to update the library)."
+        supported = bool(
+            probe_allocator_backend(
+                so_path,
+                "mc_allocator_probe",
+                ctypes.c_int,
+                0,
             )
-            return False
-        except Exception as e:
-            logger.warning(f"Failed to probe fabric memory support: {e}")
-            return False
+        )
+        if supported:
+            logger.info("Supports Fabric Memory with aclMallocPhysical")
+        else:
+            logger.info("Fabric memory not supported")
+        return supported
 
     @classmethod
     def detect_mem_backend(cls) -> bool:
@@ -84,13 +47,14 @@ class UBShmemAllocator:
             with cls._lock:
                 if cls._probe_done:
                     return cls._supports_fabric
-                so_path = None
                 try:
-                    so_path = cls._get_so_path()
-                    # First try dedicated probe function
-                    cls._supports_fabric = cls._probe_fabric_memory_support(so_path)
+                    cls._supports_fabric = cls._probe_fabric_memory_support(
+                        cls._get_so_path()
+                    )
                 except Exception as e:
-                    logger.error(f"Critical error during fabric memory probe setup: {e}")
+                    logger.error(
+                        f"Critical error during fabric memory probe setup: {e}"
+                    )
                     cls._supports_fabric = False
 
                 cls._probe_done = True
@@ -102,6 +66,6 @@ class UBShmemAllocator:
             if device not in cls._instances:
                 so_path = cls._get_so_path()
                 cls._instances[device] = NPUPluggableAllocator(
-                    so_path, "mc_ub_fabric_malloc", "mc_ub_fabric_free"
+                    so_path, "mc_allocator_malloc", "mc_allocator_free"
                 )
             return cls._instances[device]
diff --git a/mooncake-integration/fabric_allocator_utils.py b/mooncake-integration/fabric_allocator_utils.py
new file mode 100644
index 00000000..03292aa5
--- /dev/null
+++ b/mooncake-integration/fabric_allocator_utils.py
@@ -0,0 +1,58 @@
+import ctypes
+import logging
+import os
+from importlib import resources
+from typing import Any, TypeVar
+
+logger = logging.getLogger(__name__)
+
+ProbeValue = TypeVar("ProbeValue")
+
+
+def get_mooncake_so_path(so_name: str, import_error_message: str) -> str:
+    try:
+        with resources.path("mooncake", so_name) as so_path:
+            if so_path.exists():
+                return str(so_path)
+    except (ImportError, FileNotFoundError, TypeError):
+        pass
+
+    try:
+        import mooncake
+
+        base_path = os.path.dirname(os.path.abspath(mooncake.__file__))
+        so_path = os.path.join(base_path, so_name)
+        if os.path.exists(so_path):
+            return so_path
+    except (ImportError, FileNotFoundError, TypeError):
+        pass
+
+    raise ImportError(import_error_message)
+
+
+def probe_allocator_backend(
+    so_path: str,
+    probe_symbol: str,
+    restype: Any,
+    unsupported_value: ProbeValue,
+) -> ProbeValue:
+    try:
+        lib = ctypes.CDLL(so_path)
+        probe_func = getattr(lib, probe_symbol)
+        probe_func.argtypes = [ctypes.c_int]
+        probe_func.restype = restype
+        return probe_func(0)
+    except AttributeError:
+        logger.warning(
+            "Symbol '%s' not found in %s. Assuming allocator probing is unsupported.",
+            probe_symbol,
+            os.path.basename(so_path),
+        )
+        return unsupported_value
+    except Exception as exc:
+        logger.warning(
+            "Failed to probe allocator backend from %s: %s",
+            os.path.basename(so_path),
+            exc,
+        )
+        return unsupported_value
diff --git a/mooncake-integration/integration_utils.h b/mooncake-integration/integration_utils.h
index 8179b72c..7ba7b6a7 100644
--- a/mooncake-integration/integration_utils.h
+++ b/mooncake-integration/integration_utils.h
@@ -4,8 +4,16 @@
 #include <pybind11/pybind11.h>
 #include <pybind11/stl.h>
 
+#include <algorithm>
 #include <array>
+#include <cstddef>
+#include <cstdint>
+#include <cstring>
 #include <functional>
+#include <limits>
+#include <optional>
+#include <span>
+#include <vector>
 
 namespace py = pybind11;
 
@@ -68,6 +76,32 @@ static const std::array<ArrayCreatorFunc, 15> array_creators = {{
     create_typed_array<uint8_t>,  // FLOAT8_E5M2 = 14 (using uint8_t as storage)
 }};
 
+inline std::optional<size_t> TensorDtypeElementSize(int32_t dtype) {
+    switch (static_cast<TensorDtype>(dtype)) {
+        case TensorDtype::FLOAT64:
+        case TensorDtype::INT64:
+        case TensorDtype::UINT64:
+            return size_t{8};
+        case TensorDtype::FLOAT32:
+        case TensorDtype::INT32:
+        case TensorDtype::UINT32:
+            return size_t{4};
+        case TensorDtype::INT16:
+        case TensorDtype::UINT16:
+        case TensorDtype::FLOAT16:
+        case TensorDtype::BFLOAT16:
+            return size_t{2};
+        case TensorDtype::INT8:
+        case TensorDtype::UINT8:
+        case TensorDtype::BOOL:
+        case TensorDtype::FLOAT8_E4M3:
+        case TensorDtype::FLOAT8_E5M2:
+            return size_t{1};
+        default:
+            return std::nullopt;
+    }
+}
+
 inline TensorDtype get_tensor_dtype(py::object dtype_obj) {
     if (dtype_obj.is_none()) {
         return TensorDtype::UNKNOWN;
@@ -96,10 +130,272 @@ inline TensorDtype get_tensor_dtype(py::object dtype_obj) {
     return TensorDtype::UNKNOWN;
 }
 
-struct TensorMetadata {
-    int32_t dtype;
-    int32_t ndim;
-    int64_t shape[4];
+inline py::object tensor_dtype_to_torch_dtype(TensorDtype dtype) {
+    auto torch = torch_module();
+    switch (dtype) {
+        case TensorDtype::FLOAT32:
+            return torch.attr("float32");
+        case TensorDtype::FLOAT64:
+            return torch.attr("float64");
+        case TensorDtype::INT8:
+            return torch.attr("int8");
+        case TensorDtype::UINT8:
+            return torch.attr("uint8");
+        case TensorDtype::INT16:
+            return torch.attr("int16");
+        case TensorDtype::UINT16:
+            return torch.attr("uint16");
+        case TensorDtype::INT32:
+            return torch.attr("int32");
+        case TensorDtype::UINT32:
+            return torch.attr("uint32");
+        case TensorDtype::INT64:
+            return torch.attr("int64");
+        case TensorDtype::UINT64:
+            return torch.attr("uint64");
+        case TensorDtype::BOOL:
+            return torch.attr("bool");
+        case TensorDtype::FLOAT16:
+            return torch.attr("float16");
+        case TensorDtype::BFLOAT16:
+            return torch.attr("bfloat16");
+        case TensorDtype::FLOAT8_E4M3:
+            return torch.attr("float8_e4m3fn");
+        case TensorDtype::FLOAT8_E5M2:
+            return torch.attr("float8_e5m2");
+        default:
+            return py::none();
+    }
+}
+
+constexpr uint32_t kTensorObjectMagic = 0x4d4f4f4e;
+constexpr uint16_t kTensorObjectVersion = 1;
+constexpr size_t kMaxTensorDims = 8;
+constexpr size_t kMaxLayoutAxes = 4;
+
+enum class TensorLayoutKind : uint32_t {
+    FULL = 0,
+    SHARD = 1,
 };
 
+enum class LayoutAxisKind : int32_t {
+    DP = 0,
+    TP = 1,
+    EP = 2,
+    PP = 3,
+    RESERVED = 4,
+};
+
+struct TensorShape {
+    int64_t dims[kMaxTensorDims];
+};
+
+struct LayoutAxis {
+    int32_t kind;
+    int32_t axis_index;
+    int32_t shard_rank;
+    int32_t shard_count;
+    int32_t split_dim;
+    int32_t reserved0;
+    int64_t reserved1;
+};
+
+struct TensorLayoutMetadata {
+    TensorShape global_shape;
+    TensorShape local_shape;
+    uint32_t axis_count;
+    uint32_t reserved0;
+    LayoutAxis axes[kMaxLayoutAxes];
+};
+
+struct TensorObjectHeader {
+    uint32_t magic;
+    uint16_t version;
+    uint16_t header_size;
+    int32_t dtype;
+    int32_t ndim;
+    uint32_t layout_kind;
+    uint32_t reserved_flags;
+    uint64_t data_offset;
+    uint64_t data_bytes;
+};
+
+struct TensorMetadata {
+    TensorObjectHeader header;
+    TensorLayoutMetadata layout;
+};
+
+struct ParsedTensorMetadata {
+    TensorMetadata metadata;
+    size_t data_offset;
+    size_t data_bytes;
+};
+
+inline TensorShape MakeTensorShape(const std::vector<int64_t> &dims) {
+    TensorShape shape{};
+    std::fill(std::begin(shape.dims), std::end(shape.dims), -1);
+    for (size_t i = 0; i < dims.size() && i < kMaxTensorDims; ++i) {
+        shape.dims[i] = dims[i];
+    }
+    return shape;
+}
+
+inline std::vector<int64_t> TensorShapeToVector(const TensorShape &shape,
+                                                int32_t ndim) {
+    std::vector<int64_t> dims;
+    dims.reserve(ndim);
+    for (int32_t i = 0; i < ndim; ++i) {
+        dims.push_back(shape.dims[i]);
+    }
+    return dims;
+}
+
+inline TensorMetadata BuildTensorMetadata(
+    int32_t dtype, const std::vector<int64_t> &global_shape,
+    const std::vector<int64_t> &local_shape,
+    TensorLayoutKind layout_kind = TensorLayoutKind::FULL,
+    std::span<const LayoutAxis> axes = {}) {
+    TensorMetadata metadata{};
+    metadata.header.magic = kTensorObjectMagic;
+    metadata.header.version = kTensorObjectVersion;
+    metadata.header.header_size = sizeof(TensorMetadata);
+    metadata.header.dtype = dtype;
+    metadata.header.ndim = static_cast<int32_t>(global_shape.size());
+    metadata.header.layout_kind = static_cast<uint32_t>(layout_kind);
+    metadata.header.reserved_flags = 0;
+    metadata.header.data_offset = sizeof(TensorMetadata);
+    metadata.header.data_bytes = 0;
+
+    metadata.layout.global_shape = MakeTensorShape(global_shape);
+    metadata.layout.local_shape = MakeTensorShape(local_shape);
+    metadata.layout.axis_count =
+        static_cast<uint32_t>(std::min(axes.size(), kMaxLayoutAxes));
+    metadata.layout.reserved0 = 0;
+    for (size_t i = 0; i < metadata.layout.axis_count; ++i) {
+        metadata.layout.axes[i] = axes[i];
+    }
+    return metadata;
+}
+
+inline std::optional<size_t> TensorMetadataExpectedDataBytes(
+    const TensorMetadata &metadata) {
+    auto element_size = TensorDtypeElementSize(metadata.header.dtype);
+    if (!element_size.has_value()) {
+        return std::nullopt;
+    }
+
+    size_t numel = 1;
+    for (int32_t i = 0; i < metadata.header.ndim; ++i) {
+        const int64_t dim = metadata.layout.local_shape.dims[i];
+        if (dim < 0) {
+            return std::nullopt;
+        }
+        const auto dim_size = static_cast<size_t>(dim);
+        if (dim_size != 0 &&
+            numel > std::numeric_limits<size_t>::max() / dim_size) {
+            return std::nullopt;
+        }
+        numel *= dim_size;
+    }
+
+    if (*element_size != 0 &&
+        numel > std::numeric_limits<size_t>::max() / *element_size) {
+        return std::nullopt;
+    }
+    return numel * *element_size;
+}
+
+inline bool ValidateTensorMetadata(const TensorMetadata &metadata,
+                                   size_t total_length) {
+    if (metadata.header.magic != kTensorObjectMagic ||
+        metadata.header.version != kTensorObjectVersion ||
+        metadata.header.header_size != sizeof(TensorMetadata)) {
+        return false;
+    }
+
+    if (metadata.header.dtype < 0 ||
+        metadata.header.dtype >= static_cast<int32_t>(TensorDtype::NR_DTYPES)) {
+        return false;
+    }
+
+    if (metadata.header.ndim < 0 ||
+        metadata.header.ndim > static_cast<int32_t>(kMaxTensorDims)) {
+        return false;
+    }
+
+    if (metadata.layout.axis_count > kMaxLayoutAxes) {
+        return false;
+    }
+
+    if (metadata.header.layout_kind >
+        static_cast<uint32_t>(TensorLayoutKind::SHARD)) {
+        return false;
+    }
+
+    if (metadata.header.data_offset < sizeof(TensorMetadata) ||
+        metadata.header.data_offset > total_length) {
+        return false;
+    }
+
+    if (metadata.header.data_bytes !=
+        total_length - metadata.header.data_offset) {
+        return false;
+    }
+
+    for (int32_t i = 0; i < metadata.header.ndim; ++i) {
+        if (metadata.layout.global_shape.dims[i] < 0 ||
+            metadata.layout.local_shape.dims[i] < 0) {
+            return false;
+        }
+    }
+
+    auto expected_data_bytes = TensorMetadataExpectedDataBytes(metadata);
+    if (!expected_data_bytes.has_value() ||
+        metadata.header.data_bytes != *expected_data_bytes) {
+        return false;
+    }
+
+    for (size_t i = metadata.header.ndim; i < kMaxTensorDims; ++i) {
+        if (metadata.layout.global_shape.dims[i] != -1 ||
+            metadata.layout.local_shape.dims[i] != -1) {
+            return false;
+        }
+    }
+
+    for (size_t i = 0; i < metadata.layout.axis_count; ++i) {
+        const auto &axis = metadata.layout.axes[i];
+        if (axis.kind < static_cast<int32_t>(LayoutAxisKind::DP) ||
+            axis.kind > static_cast<int32_t>(LayoutAxisKind::RESERVED)) {
+            return false;
+        }
+        if (axis.shard_count <= 0 || axis.shard_rank < 0 ||
+            axis.shard_rank >= axis.shard_count) {
+            return false;
+        }
+        if (axis.split_dim < -1 || axis.split_dim >= metadata.header.ndim) {
+            return false;
+        }
+    }
+
+    return true;
+}
+
+inline std::optional<ParsedTensorMetadata> ParseTensorMetadata(
+    const char *data, size_t total_length) {
+    if (!data || total_length < sizeof(TensorMetadata)) {
+        return std::nullopt;
+    }
+
+    ParsedTensorMetadata parsed{};
+    std::memcpy(&parsed.metadata, data, sizeof(TensorMetadata));
+    if (!ValidateTensorMetadata(parsed.metadata, total_length)) {
+        return std::nullopt;
+    }
+
+    parsed.data_offset =
+        static_cast<size_t>(parsed.metadata.header.data_offset);
+    parsed.data_bytes = static_cast<size_t>(parsed.metadata.header.data_bytes);
+    return parsed;
+}
+
 }  // namespace mooncake
diff --git a/mooncake-integration/store/buffer_pool.cpp b/mooncake-integration/store/buffer_pool.cpp
new file mode 100644
index 00000000..e5eb82ce
--- /dev/null
+++ b/mooncake-integration/store/buffer_pool.cpp
@@ -0,0 +1,591 @@
+#include "buffer_pool.h"
+
+#include <pybind11/gil.h>
+
+#include <atomic>
+#include <chrono>
+#include <condition_variable>
+#include <cstring>
+#include <cstdint>
+#include <cstdlib>
+#include <limits>
+#include <memory>
+#include <mutex>
+#include <optional>
+#include <stdexcept>
+#include <unordered_map>
+#include <unordered_set>
+#include <utility>
+
+#include "pyclient.h"
+#include "real_client.h"
+#include "types.h"
+
+namespace py = pybind11;
+
+namespace mooncake {
+namespace {
+
+constexpr char kPyClientCapsuleName[] = "mooncake.PyClient.shared_ptr";
+constexpr char kPyClientCapsuleMethod[] = "_get_pyclient_capsule";
+
+std::shared_ptr<PyClient> unwrap_pyclient_capsule(py::object capsule) {
+    if (capsule.is_none()) {
+        return nullptr;
+    }
+    if (!PyCapsule_CheckExact(capsule.ptr())) {
+        throw std::runtime_error(
+            "store wrapper returned a non-capsule PyClient handle");
+    }
+    py::capsule py_client_capsule(capsule);
+    const char *capsule_name = py_client_capsule.name();
+    if (capsule_name == nullptr ||
+        std::strcmp(capsule_name, kPyClientCapsuleName) != 0) {
+        throw std::runtime_error(
+            "store wrapper returned an unexpected PyClient capsule type");
+    }
+    auto *ptr = static_cast<std::shared_ptr<PyClient> *>(
+        py_client_capsule.get_pointer());
+    if (ptr == nullptr) {
+        throw std::runtime_error(
+            "store wrapper returned an empty PyClient capsule");
+    }
+    return *ptr;
+}
+
+std::shared_ptr<PyClient> unwrap_store(py::object store_obj) {
+    if (store_obj.is_none()) {
+        return nullptr;
+    }
+    try {
+        return store_obj.cast<std::shared_ptr<PyClient>>();
+    } catch (const py::cast_error &) {
+    }
+    if (!py::hasattr(store_obj, kPyClientCapsuleMethod)) {
+        throw std::runtime_error(
+            "BufferPool store parameter must be a PyClient or store wrapper "
+            "that "
+            "implements _get_pyclient_capsule()");
+    }
+    return unwrap_pyclient_capsule(store_obj.attr(kPyClientCapsuleMethod)());
+}
+
+}  // namespace
+
+class BufferPoolNative;
+class BufferLeaseNative;
+
+class BufferLeaseViewNative {
+   public:
+    BufferLeaseViewNative(std::shared_ptr<BufferLeaseNative> lease,
+                          size_t size);
+    ~BufferLeaseViewNative();
+
+    BufferLeaseViewNative(const BufferLeaseViewNative &) = delete;
+    BufferLeaseViewNative &operator=(const BufferLeaseViewNative &) = delete;
+
+    py::buffer_info buffer_info();
+
+   private:
+    std::shared_ptr<BufferLeaseNative> lease_;
+    size_t size_ = 0;
+};
+
+class BufferLeaseNative
+    : public std::enable_shared_from_this<BufferLeaseNative> {
+   public:
+    BufferLeaseNative(std::shared_ptr<BufferPoolNative> pool,
+                      std::shared_ptr<BufferHandle> handle,
+                      size_t requested_size);
+    ~BufferLeaseNative();
+
+    BufferLeaseNative(const BufferLeaseNative &) = delete;
+    BufferLeaseNative &operator=(const BufferLeaseNative &) = delete;
+
+    uintptr_t ptr() const;
+    size_t size() const { return requested_size_; }
+    py::object buffer();
+    py::buffer_info buffer_info(size_t size);
+    void add_export();
+    void release_export();
+    BufferLeaseNative &enter() { return *this; }
+    void exit(const py::object &, const py::object &, const py::object &) {
+        release();
+    }
+    void release();
+
+   private:
+    void release_lease(bool check_exported_views);
+
+    std::shared_ptr<BufferPoolNative> pool_;
+    std::shared_ptr<BufferHandle> handle_;
+    size_t requested_size_ = 0;
+    bool closed_ = false;
+    std::atomic<size_t> exports_{0};
+};
+
+class BufferPoolNative : public std::enable_shared_from_this<BufferPoolNative> {
+   public:
+    BufferPoolNative(const py::object &store, size_t max_bytes,
+                     size_t min_size_class, py::object max_size_class,
+                     size_t alignment, bool block_on_exhaustion,
+                     py::object default_timeout, py::object max_regions);
+    ~BufferPoolNative() {
+        try {
+            close();
+        } catch (...) {
+        }
+    }
+
+    std::shared_ptr<BufferLeaseNative> acquire(size_t size, py::object block,
+                                               py::object timeout);
+    std::shared_ptr<BufferLeaseNative> buffer(size_t size, py::object block,
+                                              py::object timeout) {
+        return acquire(size, block, timeout);
+    }
+    void prewarm(size_t size, size_t count);
+    void close();
+
+   private:
+    struct Region {
+        std::shared_ptr<BufferHandle> handle;
+        size_t size = 0;
+        bool overflow_registered = false;
+    };
+
+    friend class BufferLeaseNative;
+
+    std::shared_ptr<BufferLeaseNative> make_lease_locked(Region region,
+                                                         size_t requested_size);
+    std::optional<Region> try_acquire_locked(std::unique_lock<std::mutex> &lock,
+                                             size_t requested_size);
+    std::shared_ptr<BufferHandle> allocate_overflow_unlocked(size_t size);
+    void unregister_overflow_unlocked(const Region &region);
+    void release_internal(std::shared_ptr<BufferHandle> &handle);
+    void reserve_locked(size_t size_class);
+    void unreserve_locked(size_t size_class);
+    std::pair<size_t, bool> allocation_size(size_t size) const;
+    bool has_capacity_for_locked(size_t size_class) const;
+    void raise_if_open_locked() const;
+
+    py::object store_obj_;
+    std::shared_ptr<PyClient> py_client_;
+    size_t local_buffer_capacity_ = 0;
+    size_t max_bytes_ = 0;
+    size_t min_size_class_ = 0;
+    size_t max_size_class_ = 0;
+    size_t alignment_ = 0;
+    bool block_on_exhaustion_ = true;
+    std::optional<double> default_timeout_;
+    std::optional<size_t> max_regions_;
+
+    mutable std::mutex mutex_;
+    std::condition_variable condition_;
+    std::unordered_map<uintptr_t, Region> regions_;
+    std::unordered_set<uintptr_t> in_use_;
+    bool closed_ = false;
+    bool closing_ = false;
+    size_t total_bytes_ = 0;
+    size_t reserved_bytes_ = 0;
+    size_t reserved_regions_ = 0;
+    size_t acquire_count_ = 0;
+    size_t allocate_count_ = 0;
+    size_t oversize_allocate_count_ = 0;
+    size_t wait_count_ = 0;
+    size_t release_count_ = 0;
+};
+
+BufferLeaseViewNative::BufferLeaseViewNative(
+    std::shared_ptr<BufferLeaseNative> lease, size_t size)
+    : lease_(std::move(lease)), size_(size) {
+    lease_->add_export();
+}
+
+BufferLeaseViewNative::~BufferLeaseViewNative() { lease_->release_export(); }
+
+py::buffer_info BufferLeaseViewNative::buffer_info() {
+    return lease_->buffer_info(size_);
+}
+
+BufferLeaseNative::BufferLeaseNative(std::shared_ptr<BufferPoolNative> pool,
+                                     std::shared_ptr<BufferHandle> handle,
+                                     size_t requested_size)
+    : pool_(std::move(pool)),
+      handle_(std::move(handle)),
+      requested_size_(requested_size) {}
+
+BufferLeaseNative::~BufferLeaseNative() {
+    try {
+        release_lease(false);
+    } catch (...) {
+    }
+}
+
+uintptr_t BufferLeaseNative::ptr() const {
+    if (!handle_) throw std::runtime_error("buffer lease is closed");
+    return reinterpret_cast<uintptr_t>(handle_->ptr());
+}
+
+py::object BufferLeaseNative::buffer() {
+    auto view = std::make_shared<BufferLeaseViewNative>(shared_from_this(),
+                                                        requested_size_);
+    return py::memoryview(py::cast(view));
+}
+
+py::buffer_info BufferLeaseNative::buffer_info(size_t size) {
+    if (!handle_) throw std::runtime_error("buffer lease is closed");
+    return py::buffer_info(handle_->ptr(), sizeof(uint8_t),
+                           py::format_descriptor<uint8_t>::format(), 1, {size},
+                           {sizeof(uint8_t)});
+}
+
+void BufferLeaseNative::add_export() {
+    exports_.fetch_add(1, std::memory_order_relaxed);
+}
+
+void BufferLeaseNative::release_export() {
+    exports_.fetch_sub(1, std::memory_order_relaxed);
+}
+
+void BufferLeaseNative::release() { release_lease(true); }
+
+void BufferLeaseNative::release_lease(bool check_exported_views) {
+    if (closed_) return;
+    if (check_exported_views && exports_.load(std::memory_order_acquire) != 0) {
+        throw std::runtime_error(
+            "cannot release buffer while exported views exist");
+    }
+    auto pool = pool_;
+    if (pool && handle_) {
+        pool->release_internal(handle_);
+    }
+    pool_.reset();
+    closed_ = true;
+}
+
+BufferPoolNative::BufferPoolNative(const py::object &store, size_t max_bytes,
+                                   size_t min_size_class,
+                                   py::object max_size_class, size_t alignment,
+                                   bool block_on_exhaustion,
+                                   py::object default_timeout,
+                                   py::object max_regions)
+    : store_obj_(store),
+      py_client_(unwrap_store(store)),
+      max_bytes_(max_bytes),
+      min_size_class_(min_size_class),
+      alignment_(alignment),
+      block_on_exhaustion_(block_on_exhaustion) {
+    if (store_obj_.is_none() || !py_client_) {
+        throw std::runtime_error("MooncakeDistributedStore is not initialized");
+    }
+    if (min_size_class_ == 0 || alignment_ == 0) {
+        throw std::runtime_error(
+            "min_size_class and alignment must be positive");
+    }
+    if (alignment_ < sizeof(void *) || (alignment_ & (alignment_ - 1)) != 0) {
+        throw std::runtime_error(
+            "alignment must be a power of two and at least sizeof(void*)");
+    }
+    if (!py_client_->client_buffer_allocator_ ||
+        py_client_->client_buffer_allocator_->size() == 0) {
+        throw std::runtime_error(
+            "BufferPool requires a store configured with a local buffer");
+    }
+    local_buffer_capacity_ = py_client_->client_buffer_allocator_->size();
+    if (max_bytes == 0) {
+        if (local_buffer_capacity_ >
+            std::numeric_limits<size_t>::max() - local_buffer_capacity_) {
+            max_bytes_ = local_buffer_capacity_;
+        } else {
+            max_bytes_ = local_buffer_capacity_ * 2;
+        }
+    } else {
+        max_bytes_ = std::max(max_bytes, local_buffer_capacity_);
+    }
+    max_size_class_ =
+        max_size_class.is_none() ? max_bytes_ : max_size_class.cast<size_t>();
+    if (!default_timeout.is_none()) {
+        default_timeout_ = default_timeout.cast<double>();
+    }
+    if (!max_regions.is_none()) {
+        max_regions_ = max_regions.cast<size_t>();
+    }
+}
+
+std::shared_ptr<BufferLeaseNative> BufferPoolNative::acquire(
+    size_t size, py::object block, py::object timeout) {
+    bool should_block =
+        block.is_none() ? block_on_exhaustion_ : block.cast<bool>();
+    std::optional<double> timeout_s = default_timeout_;
+    if (!timeout.is_none()) timeout_s = timeout.cast<double>();
+    if (timeout_s.has_value() && *timeout_s < 0) {
+        throw std::runtime_error("timeout must be non-negative");
+    }
+    auto deadline = timeout_s.has_value()
+                        ? std::chrono::steady_clock::now() +
+                              std::chrono::duration<double>(*timeout_s)
+                        : std::chrono::steady_clock::time_point::max();
+
+    auto [size_class, _] = allocation_size(size);
+    if (size_class > max_bytes_) {
+        throw std::runtime_error("requested buffer size exceeds pool capacity");
+    }
+    py::gil_scoped_release release_gil;
+    std::unique_lock<std::mutex> lock(mutex_);
+    while (true) {
+        raise_if_open_locked();
+        auto region = try_acquire_locked(lock, size);
+        if (region.has_value())
+            return make_lease_locked(std::move(*region), size);
+        if (!should_block) throw std::runtime_error("buffer pool is exhausted");
+        ++wait_count_;
+        if (timeout_s.has_value()) {
+            if (std::chrono::steady_clock::now() >= deadline) {
+                throw std::runtime_error("timed out waiting for buffer");
+            }
+            condition_.wait_until(lock, deadline);
+        } else {
+            condition_.wait(lock);
+        }
+    }
+}
+
+void BufferPoolNative::prewarm(size_t size, size_t count) {
+    (void)size;
+    (void)count;
+}
+
+void BufferPoolNative::close() {
+    std::unique_lock<std::mutex> lock(mutex_);
+    if (closed_) return;
+    closing_ = true;
+    lock.unlock();
+    {
+        py::gil_scoped_release release_gil;
+        lock.lock();
+        condition_.wait(lock, [&] { return reserved_regions_ == 0; });
+        lock.unlock();
+    }
+    lock.lock();
+    if (!in_use_.empty()) {
+        closing_ = false;
+        condition_.notify_all();
+        throw std::runtime_error("cannot close buffer pool with active leases");
+    }
+    closed_ = true;
+    closing_ = false;
+    condition_.notify_all();
+}
+
+std::shared_ptr<BufferLeaseNative> BufferPoolNative::make_lease_locked(
+    Region region, size_t requested_size) {
+    auto lease = std::make_shared<BufferLeaseNative>(
+        shared_from_this(), region.handle, requested_size);
+    uintptr_t ptr = reinterpret_cast<uintptr_t>(region.handle->ptr());
+    in_use_.insert(ptr);
+    ++acquire_count_;
+    return lease;
+}
+
+std::optional<BufferPoolNative::Region> BufferPoolNative::try_acquire_locked(
+    std::unique_lock<std::mutex> &lock, size_t requested_size) {
+    const size_t allocation_size = std::max<size_t>(requested_size, 1);
+    if (!has_capacity_for_locked(allocation_size)) {
+        return std::nullopt;
+    }
+    reserve_locked(allocation_size);
+    lock.unlock();
+    std::shared_ptr<BufferHandle> handle;
+    bool overflow_registered = false;
+    try {
+        if (auto alloc_result =
+                py_client_->client_buffer_allocator_->allocate(allocation_size);
+            alloc_result.has_value()) {
+            handle = std::make_shared<BufferHandle>(std::move(*alloc_result));
+        } else {
+            handle = allocate_overflow_unlocked(allocation_size);
+            overflow_registered = handle != nullptr;
+        }
+    } catch (...) {
+        lock.lock();
+        unreserve_locked(allocation_size);
+        throw;
+    }
+    lock.lock();
+    unreserve_locked(allocation_size);
+    if (!handle) {
+        return std::nullopt;
+    }
+    if (closing_ || closed_) {
+        Region region{handle, allocation_size, overflow_registered};
+        lock.unlock();
+        if (overflow_registered) {
+            unregister_overflow_unlocked(region);
+        }
+        handle.reset();
+        lock.lock();
+        throw std::runtime_error("buffer pool is closing");
+    }
+    const uintptr_t ptr = reinterpret_cast<uintptr_t>(handle->ptr());
+    Region region{handle, allocation_size, overflow_registered};
+    regions_.emplace(ptr, region);
+    total_bytes_ += allocation_size;
+    ++allocate_count_;
+    if (requested_size > max_size_class_) {
+        ++oversize_allocate_count_;
+    }
+    return region;
+}
+
+std::shared_ptr<BufferHandle> BufferPoolNative::allocate_overflow_unlocked(
+    size_t size) {
+    void *ptr = nullptr;
+    if (posix_memalign(&ptr, alignment_, size) != 0) {
+        return nullptr;
+    }
+    auto data = std::shared_ptr<void>(ptr, std::free);
+    int ret = -1;
+    {
+        py::gil_scoped_acquire acquire_gil;
+        ret = py::cast<int>(store_obj_.attr("register_buffer")(
+            reinterpret_cast<uintptr_t>(ptr), size));
+    }
+    if (ret != 0) {
+        data.reset();
+        throw std::runtime_error("overflow buffer registration failed");
+    }
+    return std::make_shared<BufferHandle>(
+        ptr, size, [data = std::move(data)]() mutable { data.reset(); });
+}
+
+void BufferPoolNative::unregister_overflow_unlocked(const Region &region) {
+    py::gil_scoped_acquire acquire_gil;
+    int ret = py::cast<int>(store_obj_.attr("unregister_buffer")(
+        reinterpret_cast<uintptr_t>(region.handle->ptr())));
+    if (ret != 0) {
+        throw std::runtime_error("overflow buffer unregistration failed");
+    }
+}
+
+void BufferPoolNative::release_internal(std::shared_ptr<BufferHandle> &handle) {
+    if (!handle) throw std::runtime_error("buffer lease is closed");
+    uintptr_t ptr = reinterpret_cast<uintptr_t>(handle->ptr());
+    Region region;
+    {
+        std::lock_guard<std::mutex> lock(mutex_);
+        if (!in_use_.count(ptr)) {
+            throw std::runtime_error("buffer lease is not active");
+        }
+        region = regions_.at(ptr);
+        regions_.erase(ptr);
+        if (total_bytes_ >= region.size) total_bytes_ -= region.size;
+        in_use_.erase(ptr);
+        ++release_count_;
+    }
+    if (region.overflow_registered) {
+        unregister_overflow_unlocked(region);
+    }
+    region.handle.reset();
+    handle.reset();
+    condition_.notify_all();
+}
+
+void BufferPoolNative::reserve_locked(size_t size_class) {
+    raise_if_open_locked();
+    if (!has_capacity_for_locked(size_class)) {
+        throw std::runtime_error("buffer pool capacity exceeded");
+    }
+    reserved_bytes_ += size_class;
+    ++reserved_regions_;
+}
+
+void BufferPoolNative::unreserve_locked(size_t size_class) {
+    reserved_bytes_ -= size_class;
+    --reserved_regions_;
+    condition_.notify_all();
+}
+
+std::pair<size_t, bool> BufferPoolNative::allocation_size(size_t size) const {
+    size = std::max<size_t>(size, 1);
+    return {size, size > max_size_class_};
+}
+
+bool BufferPoolNative::has_capacity_for_locked(size_t size_class) const {
+    if (size_class > max_bytes_) {
+        return false;
+    }
+    if (total_bytes_ > max_bytes_ - reserved_bytes_) {
+        return false;
+    }
+    if (size_class > max_bytes_ - total_bytes_ - reserved_bytes_) {
+        return false;
+    }
+    return !max_regions_.has_value() ||
+           in_use_.size() + reserved_regions_ < *max_regions_;
+}
+
+void BufferPoolNative::raise_if_open_locked() const {
+    if (closed_) {
+        throw std::runtime_error("buffer pool is closed");
+    }
+    if (closing_) {
+        throw std::runtime_error("buffer pool is closing");
+    }
+}
+
+void bind_buffer_pool(py::module &m) {
+    py::class_<BufferLeaseViewNative, std::shared_ptr<BufferLeaseViewNative>>(
+        m, "BufferLeaseView", py::buffer_protocol())
+        .def_property_readonly(
+            "buffer",
+            [](const py::object &self) { return py::memoryview(self); })
+        .def_buffer(&BufferLeaseViewNative::buffer_info);
+
+    py::class_<BufferLeaseNative, std::shared_ptr<BufferLeaseNative>>(
+        m, "BufferLease")
+        .def_property_readonly("ptr", &BufferLeaseNative::ptr)
+        .def_property_readonly("size", &BufferLeaseNative::size)
+        .def_property_readonly("buffer", &BufferLeaseNative::buffer)
+        .def("release", &BufferLeaseNative::release)
+        .def("__enter__", &BufferLeaseNative::enter,
+             py::return_value_policy::reference_internal)
+        .def("__exit__", &BufferLeaseNative::exit);
+
+    py::class_<BufferPoolNative, std::shared_ptr<BufferPoolNative>>(
+        m, "BufferPool")
+        .def(py::init([](const py::object &store, size_t max_bytes,
+                         size_t min_size_class, py::object max_size_class,
+                         size_t alignment, bool block_on_exhaustion,
+                         py::object default_timeout, py::object max_regions,
+                         py::object prewarm_size, size_t prewarm_count) {
+                 auto pool = std::make_shared<BufferPoolNative>(
+                     store, max_bytes, min_size_class, max_size_class,
+                     alignment, block_on_exhaustion, default_timeout,
+                     max_regions);
+                 if (!prewarm_size.is_none() && prewarm_count > 0) {
+                     pool->prewarm(prewarm_size.cast<size_t>(), prewarm_count);
+                 }
+                 return pool;
+             }),
+             py::arg("store"), py::arg("max_bytes") = 0,
+             py::arg("min_size_class") = 64 * 1024,
+             py::arg("max_size_class") = py::none(),
+             py::arg("alignment") = 8 * 1024 * 1024,
+             py::arg("block_on_exhaustion") = true,
+             py::arg("default_timeout") = py::none(),
+             py::arg("max_regions") = py::none(),
+             py::arg("prewarm_size") = py::none(), py::arg("prewarm_count") = 0)
+        .def("acquire", &BufferPoolNative::acquire, py::arg("size"),
+             py::arg("block") = py::none(), py::arg("timeout") = py::none())
+        .def("buffer", &BufferPoolNative::buffer, py::arg("size"),
+             py::arg("block") = py::none(), py::arg("timeout") = py::none())
+        .def("prewarm", &BufferPoolNative::prewarm, py::arg("size"),
+             py::arg("count"))
+        .def("close", &BufferPoolNative::close);
+
+    m.attr("RegisteredBufferPool") = m.attr("BufferPool");
+    m.attr("RegisteredBufferLease") = m.attr("BufferLease");
+    m.attr("RegisteredBufferLeaseView") = m.attr("BufferLeaseView");
+}
+
+}  // namespace mooncake
diff --git a/mooncake-integration/store/buffer_pool.h b/mooncake-integration/store/buffer_pool.h
new file mode 100644
index 00000000..7de258db
--- /dev/null
+++ b/mooncake-integration/store/buffer_pool.h
@@ -0,0 +1,11 @@
+#pragma once
+
+#include <pybind11/pybind11.h>
+
+namespace py = pybind11;
+
+namespace mooncake {
+
+void bind_buffer_pool(py::module &m);
+
+}  // namespace mooncake
diff --git a/mooncake-integration/store/engram_store_py.cpp b/mooncake-integration/store/engram_store_py.cpp
new file mode 100644
index 00000000..3f434736
--- /dev/null
+++ b/mooncake-integration/store/engram_store_py.cpp
@@ -0,0 +1,257 @@
+#include <pybind11/numpy.h>
+#include <pybind11/pybind11.h>
+#include <pybind11/stl.h>
+
+#include <cstring>
+
+#include "engram/engram_store.h"
+#include "engram/engram_store_config.h"
+#include "pyclient.h"
+
+namespace py = pybind11;
+using namespace mooncake;
+using namespace mooncake::engram;
+
+namespace {
+
+constexpr char kPyClientCapsuleName[] = "mooncake.PyClient.shared_ptr";
+constexpr char kPyClientCapsuleMethod[] = "_get_pyclient_capsule";
+
+std::vector<std::vector<std::vector<int64_t>>> py_obj_to_vec3d(py::object obj) {
+    std::vector<std::vector<std::vector<int64_t>>> result;
+    for (auto batch_item : obj) {
+        std::vector<std::vector<int64_t>> batch_vec;
+        for (auto token_item : batch_item) {
+            std::vector<int64_t> token_vec;
+            for (auto head_item : token_item) {
+                token_vec.push_back(head_item.cast<int64_t>());
+            }
+            batch_vec.push_back(std::move(token_vec));
+        }
+        result.push_back(std::move(batch_vec));
+    }
+    return result;
+}
+
+std::shared_ptr<PyClient> unwrap_pyclient_capsule(py::object capsule) {
+    if (capsule.is_none()) {
+        return nullptr;
+    }
+
+    if (!PyCapsule_CheckExact(capsule.ptr())) {
+        throw std::runtime_error(
+            "store wrapper returned a non-capsule PyClient handle");
+    }
+
+    py::capsule py_client_capsule(capsule);
+    const char* capsule_name = py_client_capsule.name();
+    if (capsule_name == nullptr ||
+        std::strcmp(capsule_name, kPyClientCapsuleName) != 0) {
+        throw std::runtime_error(
+            "store wrapper returned an unexpected PyClient capsule type");
+    }
+
+    auto* ptr = static_cast<std::shared_ptr<PyClient>*>(
+        py_client_capsule.get_pointer());
+    if (ptr == nullptr) {
+        throw std::runtime_error(
+            "store wrapper returned an empty PyClient capsule");
+    }
+    return *ptr;
+}
+
+std::shared_ptr<PyClient> unwrap_store(py::object store_obj) {
+    if (store_obj.is_none()) {
+        return nullptr;
+    }
+
+    try {
+        return store_obj.cast<std::shared_ptr<PyClient>>();
+    } catch (const py::cast_error&) {
+    }
+
+    if (!py::hasattr(store_obj, kPyClientCapsuleMethod)) {
+        throw std::runtime_error(
+            "EngramStore store parameter must be a PyClient or store wrapper "
+            "that "
+            "implements _get_pyclient_capsule()");
+    }
+
+    try {
+        py::object capsule = store_obj.attr(kPyClientCapsuleMethod)();
+        return unwrap_pyclient_capsule(capsule);
+    } catch (const py::error_already_set& e) {
+        throw std::runtime_error(
+            "Failed to unwrap store wrapper for EngramStore: " +
+            std::string(e.what()));
+    }
+}
+
+py::array_t<float> require_embedding_buffer(py::handle buf,
+                                            int64_t expected_rows,
+                                            int expected_cols) {
+    if (!py::isinstance<py::array>(buf)) {
+        throw std::runtime_error("embedding_buffers must be NumPy arrays");
+    }
+
+    auto arr =
+        py::array_t<float, py::array::c_style | py::array::forcecast>::ensure(
+            buf);
+    if (!arr) {
+        throw std::runtime_error("embedding_buffers must be float32 arrays");
+    }
+
+    auto req = arr.request();
+    if (req.ndim != 2) {
+        throw std::runtime_error("each embedding buffer must be 2D [N_h, D]");
+    }
+    if (req.shape[0] != expected_rows || req.shape[1] != expected_cols) {
+        throw std::runtime_error(
+            "embedding buffer shape does not match "
+            "get_table_vocab_sizes()/get_embedding_dim()");
+    }
+    return arr;
+}
+
+py::array_t<float> lookup_to_numpy(
+    EngramStore& self,
+    const std::vector<std::vector<std::vector<int64_t>>>& row_ids) {
+    if (row_ids.empty() || row_ids[0].empty()) {
+        throw std::runtime_error("row_ids must not be empty");
+    }
+
+    const int B = static_cast<int>(row_ids.size());
+    const int L = static_cast<int>(row_ids[0].size());
+    const int H = self.get_num_heads();
+    const int D = self.get_embedding_dim();
+    py::array_t<float> output({B, L, H, D});
+    auto out_buf = output.request();
+
+    int ret =
+        self.lookup_rows(row_ids, out_buf.ptr, out_buf.size * sizeof(float));
+    if (ret != 0) {
+        throw std::runtime_error("EngramStore lookup failed");
+    }
+    return output;
+}
+
+py::array_t<float> lookup_array_to_numpy(
+    EngramStore& self,
+    const py::array_t<int64_t, py::array::c_style | py::array::forcecast>&
+        row_ids) {
+    auto req = row_ids.request();
+    if (req.ndim != 3) {
+        throw std::runtime_error("row_ids array must have shape [B, L, H]");
+    }
+
+    const int B = static_cast<int>(req.shape[0]);
+    const int L = static_cast<int>(req.shape[1]);
+    const int H = static_cast<int>(req.shape[2]);
+    if (H != self.get_num_heads()) {
+        throw std::runtime_error("row_ids last dimension must match num_heads");
+    }
+
+    const int D = self.get_embedding_dim();
+    py::array_t<float> output({B, L, H, D});
+    auto out_buf = output.request();
+
+    int ret =
+        self.lookup_rows_contiguous(static_cast<const int64_t*>(req.ptr), B, L,
+                                    out_buf.ptr, out_buf.size * sizeof(float));
+    if (ret != 0) {
+        throw std::runtime_error("EngramStore lookup failed");
+    }
+    return output;
+}
+
+}  // namespace
+
+namespace mooncake {
+namespace engram {
+
+void bind_engram_store(py::module& m) {
+    py::class_<EngramStoreConfig>(m, "EngramStoreConfig")
+        .def(py::init<>())
+        .def_readwrite("table_vocab_sizes",
+                       &EngramStoreConfig::table_vocab_sizes)
+        .def_readwrite("embedding_dim", &EngramStoreConfig::embedding_dim);
+
+    py::class_<EngramStore>(m, "EngramStore")
+        .def("get_table_vocab_sizes", &EngramStore::get_table_vocab_sizes)
+        .def("get_store_keys", &EngramStore::get_store_keys)
+        .def("get_num_heads", &EngramStore::get_num_heads)
+        .def("get_embedding_dim", &EngramStore::get_embedding_dim)
+        .def(
+            "remove_from_store",
+            [](EngramStore& self, bool force) {
+                int ret = self.remove_from_store(force);
+                if (ret < 0) {
+                    throw std::runtime_error("remove_from_store failed, rc=" +
+                                             std::to_string(ret));
+                }
+                return ret;
+            },
+            py::arg("force") = false,
+            "Remove all Mooncake Store tables owned by this EngramStore layer. "
+            "Returns the number of removed head tables; missing keys are "
+            "ignored.")
+        .def(py::init([](int layer_id, const EngramStoreConfig& cfg,
+                         py::object store_obj) {
+                 std::shared_ptr<PyClient> store = unwrap_store(store_obj);
+                 return new EngramStore(layer_id, cfg, store);
+             }),
+             py::arg("layer_id"), py::arg("config"),
+             py::arg("store") = py::none())
+        .def(
+            "lookup",
+            [](EngramStore& self, py::object row_ids_obj) {
+                if (py::isinstance<py::array>(row_ids_obj)) {
+                    auto row_ids_array = py::array_t<
+                        int64_t, py::array::c_style |
+                                     py::array::forcecast>::ensure(row_ids_obj);
+                    if (!row_ids_array) {
+                        throw std::runtime_error(
+                            "row_ids array must be convertible to int64");
+                    }
+                    return lookup_array_to_numpy(self, row_ids_array);
+                }
+                return lookup_to_numpy(self, py_obj_to_vec3d(row_ids_obj));
+            },
+            py::arg("row_ids"),
+            "Lookup embeddings by precomputed row IDs. Returns [B, L, H, D].")
+        .def(
+            "populate",
+            [](EngramStore& self, py::list embedding_buffers) {
+                const std::vector<int64_t> vocab_sizes =
+                    self.get_table_vocab_sizes();
+                const int embed_dim = self.get_embedding_dim();
+                if (static_cast<size_t>(py::len(embedding_buffers)) !=
+                    vocab_sizes.size()) {
+                    throw std::runtime_error(
+                        "embedding_buffers size must match num_heads");
+                }
+
+                std::vector<py::array_t<float>> arrays;
+                std::vector<void*> bufs;
+                std::vector<size_t> sizes;
+                arrays.reserve(vocab_sizes.size());
+                bufs.reserve(vocab_sizes.size());
+                sizes.reserve(vocab_sizes.size());
+                for (size_t i = 0; i < vocab_sizes.size(); ++i) {
+                    auto arr = require_embedding_buffer(
+                        embedding_buffers[i], vocab_sizes[i], embed_dim);
+                    auto req = arr.request();
+                    arrays.push_back(arr);
+                    bufs.push_back(req.ptr);
+                    sizes.push_back(req.size * sizeof(float));
+                }
+                int ret = self.populate(bufs, sizes);
+                if (ret != 0) {
+                    throw std::runtime_error("populate failed");
+                }
+            },
+            py::arg("embedding_buffers"));
+}
+
+}  // namespace engram
+}  // namespace mooncake
diff --git a/mooncake-integration/store/store_py.cpp b/mooncake-integration/store/store_py.cpp
index 8872e4ed..fa279cb3 100644
--- a/mooncake-integration/store/store_py.cpp
+++ b/mooncake-integration/store/store_py.cpp
@@ -2,14 +2,33 @@
 #include <pybind11/stl.h>
 #include <numa.h>
 
+#include <functional>
+#include <limits>
+#include <numeric>
+#include <unordered_map>
+#include <unordered_set>
+
 #include "pyclient.h"
 #include "dummy_client.h"
 #include "real_client.h"
 #include "types.h"
+#include "memory_alloc.h"
+#include "ssd_register_client.h"
 
 #include <cstdlib>  // for atexit
+#include <memory>
 
 #include "integration_utils.h"
+#include "buffer_pool.h"
+
+// Forward declaration for EngramStore bindings
+namespace mooncake {
+namespace engram {
+void bind_engram_store(py::module &m);
+}
+}  // namespace mooncake
+#include "engram/engram_store.h"
+#include "engram/engram_store_config.h"
 
 namespace py = pybind11;
 
@@ -20,54 +39,62 @@ struct PyTensorInfo {
     uintptr_t data_ptr;
     size_t tensor_size;
     TensorMetadata metadata;
+    py::object owner;
 
-    // Check validity
     bool valid() const {
         // Basic size check
-        if (tensor_size == 0 || data_ptr == 0) {
+        if (data_ptr == 0 && tensor_size != 0) {
             return false;
         }
 
         // Validate metadata
+        TensorMetadata validated = metadata;
+        validated.header.data_bytes = tensor_size;
+
         // Check dtype is within valid range (0 to TensorDtype::NR_DTYPES,
         // excluding UNKNOWN=-1)
-        if (metadata.dtype < 0 ||
-            metadata.dtype >= static_cast<int32_t>(TensorDtype::NR_DTYPES)) {
+        if (validated.header.dtype >=
+            static_cast<int32_t>(TensorDtype::NR_DTYPES)) {
             return false;
         }
 
         // Check ndim is within valid range (0 to shape array size)
-        const int kMaxDims = std::size(metadata.shape);
-        if (metadata.ndim < 0 || metadata.ndim > kMaxDims) {
+        const int kMaxDims = kMaxTensorDims;
+        if (validated.header.ndim < 0 || validated.header.ndim > kMaxDims) {
             return false;
         }
 
         // Validate shape array
         // For valid dimensions (0 to ndim-1), shape should be >= 0
-        for (int i = 0; i < metadata.ndim; ++i) {
-            if (metadata.shape[i] <= 0) {
+        auto shape = TensorShapeToVector(validated.layout.local_shape,
+                                         validated.header.ndim);
+        for (int i = 0; i < validated.header.ndim; ++i) {
+            if (shape[i] < 0) {
                 return false;  // Invalid dimension size
             }
         }
 
-        // For padding dimensions (ndim to kMaxDims-1), shape should be -1
-        // (placeholder)
-        for (int i = metadata.ndim; i < kMaxDims; ++i) {
-            if (metadata.shape[i] != -1) {
-                return false;  // Padding dimensions should be -1
-            }
-        }
-
-        return true;
+        return ValidateTensorMetadata(
+            validated, validated.header.data_offset + tensor_size);
     }
 };
 
+TensorMetadata build_full_tensor_metadata(const py::handle &tensor,
+                                          TensorDtype dtype_enum,
+                                          size_t tensor_size);
+std::optional<ParsedTensorMetadata> parse_tensor_metadata_from_buffer(
+    BufferHandle *buffer_handle, char *usr_buffer, int64_t data_length,
+    bool *take_ownership, char **exported_data, size_t *total_length);
+std::pair<int64_t, int64_t> calculate_shard_range(int64_t dim_size, int rank,
+                                                  int shard_count);
+
 PyTensorInfo extract_tensor_info(const py::object &tensor,
                                  const std::string &key_name = "") {
     PyTensorInfo info = {
         0,
         0,
         {},
+        py::none(),
     };
 
     if (!(tensor.attr("__class__")
@@ -80,9 +107,12 @@ PyTensorInfo extract_tensor_info(const py::object &tensor,
     }
 
     try {
-        info.data_ptr = tensor.attr("data_ptr")().cast<uintptr_t>();
-        size_t numel = tensor.attr("numel")().cast<size_t>();
-        size_t element_size = tensor.attr("element_size")().cast<size_t>();
+        py::object contiguous_tensor = tensor.attr("contiguous")();
+        info.owner = contiguous_tensor;
+        info.data_ptr = contiguous_tensor.attr("data_ptr")().cast<uintptr_t>();
+        size_t numel = contiguous_tensor.attr("numel")().cast<size_t>();
+        size_t element_size =
+            contiguous_tensor.attr("element_size")().cast<size_t>();
         info.tensor_size = numel * element_size;
 
         pybind11::object shape_obj = tensor.attr("shape");
@@ -92,44 +122,56 @@ PyTensorInfo extract_tensor_info(const py::object &tensor,
         if (dtype_enum == TensorDtype::UNKNOWN) {
             LOG(ERROR) << "Unsupported tensor dtype"
                        << (key_name.empty() ? "" : " for " + key_name);
-            return {0, 0, {}};
+            return {0, 0, {}, py::none()};
         }
 
         pybind11::tuple shape_tuple =
             pybind11::cast<pybind11::tuple>(shape_obj);
         int32_t ndim = static_cast<int32_t>(shape_tuple.size());
 
-        if (ndim > 4) {
-            LOG(ERROR) << "Tensor has more than 4 dimensions: " << ndim;
-            return {0, 0, {}};
+        if (ndim > static_cast<int32_t>(kMaxTensorDims)) {
+            LOG(ERROR) << "Tensor has more than " << kMaxTensorDims
+                       << " dimensions: " << ndim;
+            return {0, 0, {}, py::none()};
         }
 
-        info.metadata.dtype = static_cast<int32_t>(dtype_enum);
-        info.metadata.ndim = ndim;
-
-        for (int i = 0; i < 4; i++) {
-            info.metadata.shape[i] =
-                (i < ndim) ? shape_tuple[i].cast<int64_t>() : -1;
-        }
+        info.metadata =
+            build_full_tensor_metadata(tensor, dtype_enum, info.tensor_size);
     } catch (const std::exception &e) {
         LOG(ERROR) << "Error extracting tensor info: " << e.what();
-        return {0, 0, {}};
+        return {0, 0, {}, py::none()};
     }
 
     return info;
 }
 
+py::tuple tensor_shape_tuple(const TensorMetadata &metadata) {
+    std::vector<int64_t> shape_vec;
+    shape_vec.reserve(metadata.header.ndim);
+    for (int i = 0; i < metadata.header.ndim; i++) {
+        shape_vec.push_back(metadata.layout.local_shape.dims[i]);
+    }
+    return py::cast(shape_vec);
+}
+
+void append_tensor_payload_span(std::vector<std::span<const char>> &values,
+                                uintptr_t data_ptr, size_t tensor_size) {
+    if (tensor_size == 0) return;
+    values.emplace_back(reinterpret_cast<const char *>(data_ptr), tensor_size);
+}
+
 pybind11::object buffer_to_tensor(BufferHandle *buffer_handle, char *usr_buffer,
-                                  int64_t data_length) {
+                                  int64_t data_length,
+                                  bool own_user_buffer = false) {
     if (!buffer_handle && !usr_buffer) return pybind11::none();
     if (buffer_handle && usr_buffer) return pybind11::none();
 
-    bool take_ownership = !!buffer_handle;
+    bool take_ownership = !!buffer_handle || own_user_buffer;
     size_t total_length;
     char *exported_data;
-    if (take_ownership) {
+    if (buffer_handle) {
         total_length = buffer_handle->size();
-        if (total_length <= sizeof(TensorMetadata)) {
+        if (total_length < sizeof(TensorMetadata)) {
             LOG(ERROR) << "Invalid data format: insufficient data for metadata";
             return pybind11::none();
         }
@@ -143,12 +185,18 @@ pybind11::object buffer_to_tensor(BufferHandle *buffer_handle, char *usr_buffer,
     } else {
         exported_data = usr_buffer;
         if (data_length < 0) {
+            if (take_ownership) {
+                delete[] exported_data;
+            }
             LOG(ERROR) << "Get tensor into failed with error code: "
                        << data_length;
             return pybind11::none();
         }
         total_length = static_cast<size_t>(data_length);
-        if (total_length <= sizeof(TensorMetadata)) {
+        if (total_length < sizeof(TensorMetadata)) {
+            if (take_ownership) {
+                delete[] exported_data;
+            }
             LOG(ERROR) << "Invalid data format: insufficient data for metadata";
             return pybind11::none();
         }
@@ -156,18 +204,30 @@ pybind11::object buffer_to_tensor(BufferHandle *buffer_handle, char *usr_buffer,
     TensorMetadata metadata;
     memcpy(&metadata, exported_data, sizeof(TensorMetadata));
 
-    if (metadata.ndim < 0 || metadata.ndim > 4) {
+    auto parsed = ParseTensorMetadata(exported_data, total_length);
+    if (!parsed.has_value()) {
         if (take_ownership) {
             delete[] exported_data;
         }
-        LOG(ERROR) << "Invalid tensor metadata: ndim=" << metadata.ndim;
+        LOG(ERROR) << "Invalid tensor metadata";
         return pybind11::none();
     }
 
-    TensorDtype dtype_enum = static_cast<TensorDtype>(metadata.dtype);
-    size_t tensor_size = total_length - sizeof(TensorMetadata);
+    int ndim = metadata.header.ndim;
+    if (ndim < 0 || ndim > static_cast<int>(kMaxTensorDims)) {
+        if (take_ownership) {
+            delete[] exported_data;
+        }
+        LOG(ERROR) << "Invalid tensor metadata: ndim=" << ndim
+                   << ", max supported=" << kMaxTensorDims;
+        return pybind11::none();
+    }
 
-    if (tensor_size == 0 || dtype_enum == TensorDtype::UNKNOWN) {
+    TensorDtype dtype_enum = static_cast<TensorDtype>(metadata.header.dtype);
+    size_t tensor_size = parsed->data_bytes;
+    size_t data_offset = parsed->data_offset;
+
+    if (dtype_enum == TensorDtype::UNKNOWN) {
         if (take_ownership) {
             delete[] exported_data;
         }
@@ -185,20 +245,28 @@ pybind11::object buffer_to_tensor(BufferHandle *buffer_handle, char *usr_buffer,
         return pybind11::none();
     }
 
+    std::unique_ptr<char[]> exported_data_guard(take_ownership ? exported_data
+                                                               : nullptr);
     try {
+        if (tensor_size == 0) {
+            py::object dtype = tensor_dtype_to_torch_dtype(dtype_enum);
+            if (dtype.is_none()) {
+                LOG(ERROR) << "Unsupported dtype enum: " << dtype_index;
+                return pybind11::none();
+            }
+            py::object tensor = torch_module().attr("empty")(
+                tensor_shape_tuple(metadata), py::arg("dtype") = dtype);
+            return tensor;
+        }
+
         // Construct numpy array
-        bool take_ownership = !!buffer_handle;
         py::object np_array = array_creators[dtype_index](
-            exported_data, sizeof(TensorMetadata), tensor_size, take_ownership);
+            exported_data, data_offset, tensor_size, take_ownership);
+        exported_data_guard.release();
 
         // Reshape
-        if (metadata.ndim > 0) {
-            std::vector<uint64_t> shape_vec;
-            for (int i = 0; i < metadata.ndim; i++) {
-                shape_vec.push_back(metadata.shape[i]);
-            }
-            py::tuple shape_tuple = py::cast(shape_vec);
-            np_array = np_array.attr("reshape")(shape_tuple);
+        if (ndim > 0) {
+            np_array = np_array.attr("reshape")(tensor_shape_tuple(metadata));
         }
 
         // Convert to torch tensor
@@ -218,13 +286,43 @@ pybind11::object buffer_to_tensor(BufferHandle *buffer_handle, char *usr_buffer,
 
     } catch (const std::exception &e) {
         LOG(ERROR) << "Failed to convert buffer to tensor: " << e.what();
-        if (take_ownership) {
-            delete[] exported_data;
-        }
         return pybind11::none();
     }
 }
 
+py::tuple serialize_tensor_metadata(py::object tensor) {
+    PyTensorInfo info = extract_tensor_info(tensor);
+    if (!info.valid()) {
+        throw py::value_error(
+            "unsupported tensor for Mooncake tensor serialization");
+    }
+
+    TensorMetadata metadata = info.metadata;
+    metadata.header.data_bytes = info.tensor_size;
+    return py::make_tuple(py::bytes(reinterpret_cast<const char *>(&metadata),
+                                    sizeof(TensorMetadata)),
+                          info.data_ptr, info.tensor_size, info.owner);
+}
+
+size_t tensor_metadata_size() { return sizeof(TensorMetadata); }
+
+py::object deserialize_tensor_from_bytes(py::bytes payload) {
+    std::string data = payload;
+    if (data.size() >
+        static_cast<size_t>(std::numeric_limits<int64_t>::max())) {
+        throw py::value_error("serialized tensor payload is too large");
+    }
+
+    char *owned_data = new char[data.size()];
+    memcpy(owned_data, data.data(), data.size());
+    py::object tensor = buffer_to_tensor(
+        nullptr, owned_data, static_cast<int64_t>(data.size()), true);
+    if (tensor.is_none()) {
+        throw py::value_error("invalid serialized Mooncake tensor payload");
+    }
+    return tensor;
+}
+
 std::vector<std::vector<void *>> CastAddrs2Ptrs(
     const std::vector<std::vector<uintptr_t>> &all_buffer_ptrs) {
     std::vector<std::vector<void *>> all_buffers;
@@ -245,11 +343,15 @@ std::vector<std::vector<void *>> CastAddrs2Ptrs(
 inline int to_py_ret(ErrorCode error_code) {
     return static_cast<int>(error_code);
 }
+
+#include "store_py_internal.h"
+
 }  // namespace
 // Python-specific wrapper functions that handle GIL and return pybind11 types
 class MooncakeStorePyWrapper {
    public:
     std::shared_ptr<PyClient> store_{nullptr};
+    std::shared_ptr<RealClient> real_client_{nullptr};
     bool use_dummy_client_{false};
 
     MooncakeStorePyWrapper() = default;
@@ -260,6 +362,7 @@ class MooncakeStorePyWrapper {
         auto real_client = RealClient::create();
         use_dummy_client_ = false;
         store_ = real_client;
+        real_client_ = real_client;
         resource_tracker.registerInstance(
             std::static_pointer_cast<PyClient>(store_));
         return real_client;
@@ -276,7 +379,85 @@ class MooncakeStorePyWrapper {
         return store_->health_check();
     }
 
-    std::string get_tp_key_name(const std::string &base_key, int rank) {
+    py::dict mount_segment(const std::string &path, size_t size, size_t offset,
+                           const std::string &protocol,
+                           const std::string &location) {
+        py::dict result;
+        result["ret"] = -1;
+        result["segment_ids"] = py::list();
+
+        auto real_client = std::dynamic_pointer_cast<RealClient>(store_);
+        if (!real_client) {
+            LOG(ERROR) << "mount_segment requires RealClient";
+            return result;
+        }
+        std::vector<std::string> segment_ids;
+        int ret;
+        {
+            py::gil_scoped_release release;
+            ret = real_client->mountSegment(path, offset, size, protocol,
+                                            location, segment_ids);
+        }
+        result["ret"] = ret;
+        py::list ids;
+        for (const auto &id : segment_ids) {
+            ids.append(id);
+        }
+        result["segment_ids"] = ids;
+        return result;
+    }
+
+    int unmount_segment(const std::vector<std::string> &segment_ids,
+                        uint64_t grace_period_seconds = 0) {
+        auto real_client = std::dynamic_pointer_cast<RealClient>(store_);
+        if (!real_client) {
+            LOG(ERROR) << "unmount_segment requires RealClient";
+            return -1;
+        }
+        py::gil_scoped_release release;
+        return real_client->unmountSegment(segment_ids, grace_period_seconds);
+    }
+
+    py::dict allocate_and_mount_segment(size_t size,
+                                        const std::string &protocol,
+                                        const std::string &location) {
+        py::dict result;
+        result["ret"] = -1;
+        result["segment_ids"] = py::list();
+        result["allocated_size"] = 0;
+
+        auto real_client = std::dynamic_pointer_cast<RealClient>(store_);
+        if (!real_client) {
+            LOG(ERROR) << "allocate_and_mount_segment requires RealClient";
+            return result;
+        }
+        std::vector<std::string> segment_ids;
+        size_t allocated_size = 0;
+        int ret;
+        {
+            py::gil_scoped_release release;
+            ret = real_client->allocateAndMountSegment(
+                size, protocol, location, segment_ids, &allocated_size);
+        }
+        result["ret"] = ret;
+        result["segment_ids"] = py::cast(segment_ids);
+        result["allocated_size"] = allocated_size;
+        return result;
+    }
+
+    int unmount_and_free_segment(const std::vector<std::string> &segment_ids,
+                                 uint64_t grace_period_seconds = 0) {
+        auto real_client = std::dynamic_pointer_cast<RealClient>(store_);
+        if (!real_client) {
+            LOG(ERROR) << "unmount_and_free_segment requires RealClient";
+            return -1;
+        }
+        py::gil_scoped_release release;
+        return real_client->unmountAndFreeSegment(segment_ids,
+                                                  grace_period_seconds);
+    }
+
+    std::string get_tp_key_name(const std::string &base_key, int rank) const {
         return base_key + "_tp_" + std::to_string(rank);
     }
 
@@ -541,8 +722,7 @@ class MooncakeStorePyWrapper {
         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);
+        append_tensor_payload_span(values, info.data_ptr, info.tensor_size);
 
         // Store (GIL Released)
         py::gil_scoped_release release_gil;
@@ -563,24 +743,74 @@ class MooncakeStorePyWrapper {
                                ReplicateConfig{});  // Default config
     }
 
+    ReplicateConfig MakeIndexedConfig(
+        const ReplicateConfig &config,
+        const std::vector<size_t> &original_indices) const {
+        if (!config.group_ids.has_value()) {
+            return config;
+        }
+
+        ReplicateConfig indexed_config = config;
+        std::vector<std::string> group_ids;
+        group_ids.reserve(original_indices.size());
+        for (size_t index : original_indices) {
+            group_ids.push_back(config.group_ids->at(index));
+        }
+        indexed_config.group_ids = std::move(group_ids);
+        return indexed_config;
+    }
+
+    ReplicateConfig MakeRepeatedIndexedConfig(
+        const ReplicateConfig &config,
+        const std::vector<size_t> &original_indices, int repeat_count) const {
+        if (!config.group_ids.has_value()) {
+            return config;
+        }
+
+        ReplicateConfig indexed_config = config;
+        std::vector<std::string> group_ids;
+        group_ids.reserve(original_indices.size() *
+                          static_cast<size_t>(repeat_count));
+        for (size_t index : original_indices) {
+            for (int i = 0; i < repeat_count; ++i) {
+                group_ids.push_back(config.group_ids->at(index));
+            }
+        }
+        indexed_config.group_ids = std::move(group_ids);
+        return indexed_config;
+    }
+
+    std::vector<int> ValidateGroupIdsForBatchConfig(
+        const ReplicateConfig &config, size_t key_count,
+        const char *operation_name) const {
+        if (config.group_ids.has_value() &&
+            config.group_ids->size() != key_count) {
+            LOG(ERROR) << operation_name
+                       << ": group_ids size must match keys size";
+            return std::vector<int>(key_count,
+                                    to_py_ret(ErrorCode::INVALID_PARAMS));
+        }
+        return {};
+    }
+
     int put_tensor_with_tp_impl(
         const std::string &key, pybind11::object tensor,
         const ReplicateConfig &config = ReplicateConfig{}, int tp_rank = 0,
         int tp_size = 1, int split_dim = 0) {
         try {
-            py::tuple chunks =
-                tensor.attr("chunk")(tp_size, split_dim).cast<py::tuple>();
-            if (static_cast<int>(chunks.size()) != tp_size) {
-                LOG(ERROR) << "Chunking failed: got " << chunks.size()
-                           << " chunks, expected " << tp_size;
+            std::vector<std::string> tp_keys;
+            tp_keys.reserve(tp_size);
+            for (int rank = 0; rank < tp_size; ++rank) {
+                tp_keys.push_back(get_tp_key_name(key, rank));
+            }
+            auto infos = build_tp_shard_infos(
+                tensor, tp_size, split_dim,
+                [this, &key](int rank) { return get_tp_key_name(key, rank); });
+            if (!infos.has_value()) {
                 return to_py_ret(ErrorCode::INVALID_PARAMS);
             }
-
-            for (int rank = 0; rank < tp_size; ++rank) {
-                pybind11::object chunk = chunks[rank].attr("contiguous")();
-                std::string tp_key = get_tp_key_name(key, rank);
-
-                int ret = put_tensor_impl(tp_key, chunk, config);
+            auto results = batch_put_tensor_infos_impl(tp_keys, *infos, config);
+            for (int ret : results) {
                 if (ret != 0) return ret;
             }
             return 0;
@@ -605,77 +835,30 @@ class MooncakeStorePyWrapper {
                                        tp_size, split_dim);
     }
 
+    std::vector<int> batch_put_tensor_infos_impl(
+        const std::vector<std::string> &keys,
+        const std::vector<PyTensorInfo> &infos,
+        const ReplicateConfig &config = ReplicateConfig{}) {
+        return batch_write_tensor_impl(
+            keys, infos, config, "put",
+            [this](const std::vector<std::string> &write_keys,
+                   const std::vector<void *> &buffer_ptrs,
+                   const std::vector<size_t> &buffer_sizes,
+                   const ReplicateConfig &write_config) {
+                return store_->batch_put_from(write_keys, buffer_ptrs,
+                                              buffer_sizes, write_config);
+            });
+    }
+
     std::vector<int> batch_put_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;
-
-            // Temporary containers for the batch operation
-            std::vector<std::string> valid_keys;
-            std::vector<void *> buffer_ptrs;
-            std::vector<size_t> buffer_sizes;
-            std::vector<size_t> original_indices;  // Map back to results
-
-            // Note: In batch mode, we need contiguous memory for Metadata +
-            // Data.
-            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);
-
-                // Transfer ownership to temp_allocations so it survives until
-                // batch_put_from returns
-                temp_allocations.push_back(
-                    std::make_unique<BufferHandle>(std::move(*alloc_result)));
-            }
-
-            if (!valid_keys.empty()) {
-                std::vector<int> op_results = store_->batch_put_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;
+        return batch_put_tensor_infos_impl(keys, infos, config);
     }
 
     std::vector<int> batch_put_tensor(const std::vector<std::string> &keys,
@@ -699,25 +882,30 @@ class MooncakeStorePyWrapper {
         const ReplicateConfig &config = ReplicateConfig{}, int tp_rank = 0,
         int tp_size = 1, int split_dim = 0) {
         std::vector<std::string> all_chunk_keys;
-        py::list all_chunks_list;
+        std::vector<PyTensorInfo> all_chunk_infos;
         std::vector<size_t> processed_indices;
         std::vector<int> final_results(base_keys.size(),
                                        to_py_ret(ErrorCode::INVALID_PARAMS));
-
+        auto group_ids_error = ValidateGroupIdsForBatchConfig(
+            config, base_keys.size(), "batch_put_tensor_with_tp");
+        if (!group_ids_error.empty()) {
+            return group_ids_error;
+        }
         try {
-            // Chunking phase (GIL Held)
             for (size_t i = 0; i < base_keys.size(); ++i) {
                 py::object tensor = tensors_list[i];
-                // Quick validation
                 if (tensor.is_none() ||
                     !tensor.attr("shape").cast<py::tuple>()) {
                     final_results[i] = to_py_ret(ErrorCode::INVALID_PARAMS);
                     continue;
                 }
 
-                py::tuple chunks =
-                    tensor.attr("chunk")(tp_size, split_dim).cast<py::tuple>();
-                if (static_cast<int>(chunks.size()) != tp_size) {
+                auto infos = build_tp_shard_infos(
+                    tensor, tp_size, split_dim,
+                    [this, &base_keys, i](int rank) {
+                        return get_tp_key_name(base_keys[i], rank);
+                    });
+                if (!infos.has_value()) {
                     final_results[i] = to_py_ret(ErrorCode::INVALID_PARAMS);
                     continue;
                 }
@@ -726,25 +914,24 @@ class MooncakeStorePyWrapper {
                 for (int rank = 0; rank < tp_size; ++rank) {
                     all_chunk_keys.push_back(
                         get_tp_key_name(base_keys[i], rank));
-                    all_chunks_list.append(chunks[rank].attr(
-                        "contiguous")());  // Ensure contiguous here
+                    all_chunk_infos.push_back((*infos)[rank]);
                 }
             }
 
             if (all_chunk_keys.empty()) return final_results;
 
-            // Reuse the standard batch_put implementation
-            std::vector<int> chunk_results =
-                batch_put_tensor_impl(all_chunk_keys, all_chunks_list, config);
+            ReplicateConfig chunk_config =
+                MakeRepeatedIndexedConfig(config, processed_indices, tp_size);
+            std::vector<int> chunk_results = batch_put_tensor_infos_impl(
+                all_chunk_keys, all_chunk_infos, chunk_config);
 
-            // 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
+                        final_results[original_idx] = res;
                         all_ok = false;
                         break;
                     }
@@ -797,10 +984,15 @@ class MooncakeStorePyWrapper {
             LOG(ERROR) << "put_tensor_from is not supported for dummy client";
             return to_py_ret(ErrorCode::INVALID_PARAMS);
         }
-        if (size <= sizeof(TensorMetadata)) {
+        if (size < sizeof(TensorMetadata)) {
             LOG(ERROR) << "Buffer size too small for tensor metadata";
             return to_py_ret(ErrorCode::INVALID_PARAMS);
         }
+        if (!ParseTensorMetadata(reinterpret_cast<const char *>(buffer), size)
+                 .has_value()) {
+            LOG(ERROR) << "Invalid tensor metadata";
+            return to_py_ret(ErrorCode::INVALID_PARAMS);
+        }
         py::gil_scoped_release release_gil;
         return store_->put_from(key, buffer, size, ReplicateConfig{});
     }
@@ -836,12 +1028,19 @@ class MooncakeStorePyWrapper {
                 return std::vector<int>(keys.size(),
                                         to_py_ret(ErrorCode::INVALID_PARAMS));
             }
-            if (sizes[i] <= sizeof(TensorMetadata)) {
+            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));
             }
+            if (!ParseTensorMetadata(
+                     reinterpret_cast<const char *>(buffer_ptrs[i]), sizes[i])
+                     .has_value()) {
+                LOG(ERROR) << "Invalid tensor metadata at index " << i;
+                return std::vector<int>(keys.size(),
+                                        to_py_ret(ErrorCode::INVALID_PARAMS));
+            }
         }
         std::vector<void *> buffers;
         buffers.reserve(buffer_ptrs.size());
@@ -852,6 +1051,43 @@ class MooncakeStorePyWrapper {
         return store_->batch_put_from(keys, buffers, sizes, ReplicateConfig{});
     }
 
+    int put_tensor_info_impl(const std::string &key, const PyTensorInfo &info,
+                             const ReplicateConfig &config) {
+        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),
+                            info.metadata.header.data_offset);
+        append_tensor_payload_span(values, info.data_ptr, info.tensor_size);
+
+        py::gil_scoped_release release_gil;
+        int ret = store_->put_parts(key, values, config);
+        if (ret != 0)
+            LOG(ERROR) << "put_parts failed for key " << key << " with code "
+                       << ret;
+        return ret;
+    }
+
+    int put_manifest_impl(const std::string &key,
+                          const WriterShardManifest &manifest,
+                          const ReplicateConfig &config) {
+        return write_manifest_impl(
+            key, manifest, "put",
+            [this, &key, &config](std::span<const char> bytes) {
+                return store_->put(key, bytes, config);
+            });
+    }
+
+    int upsert_manifest_impl(const std::string &key,
+                             const WriterShardManifest &manifest,
+                             const ReplicateConfig &config) {
+        return write_manifest_impl(
+            key, manifest, "upsert",
+            [this, &key, &config](std::span<const char> bytes) {
+                return store_->upsert(key, bytes, config);
+            });
+    }
+
     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) {
@@ -868,7 +1104,7 @@ class MooncakeStorePyWrapper {
             LOG(ERROR) << "Buffer pointer cannot be null";
             return to_py_ret(ErrorCode::INVALID_PARAMS);
         }
-        if (size <= sizeof(TensorMetadata)) {
+        if (size < sizeof(TensorMetadata)) {
             LOG(ERROR) << "Buffer size too small for tensor metadata";
             return to_py_ret(ErrorCode::INVALID_PARAMS);
         }
@@ -928,7 +1164,7 @@ class MooncakeStorePyWrapper {
                 final_results[i] = to_py_ret(ErrorCode::INVALID_PARAMS);
                 continue;
             }
-            if (sizes[i] <= sizeof(TensorMetadata)) {
+            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);
@@ -967,8 +1203,83 @@ class MooncakeStorePyWrapper {
         return final_results;
     }
 
+#include "store_py_parallel_write.h"
+
+    std::shared_ptr<RealClient> get_real_client() const {
+        if (use_dummy_client_) {
+            return nullptr;
+        }
+        return real_client_;
+    }
+
+    std::optional<RealClient::WritableBufferRegion>
+    resolve_writable_buffer_region(uintptr_t buffer_ptr, size_t size,
+                                   const std::string &context) const {
+        auto real_client = get_real_client();
+        if (!real_client) {
+            LOG(ERROR) << context << ": real client is not available";
+            return std::nullopt;
+        }
+
+        auto region = real_client->resolve_writable_buffer_region(
+            reinterpret_cast<void *>(buffer_ptr));
+        if (!region.has_value()) {
+            LOG(ERROR) << context
+                       << ": buffer is not Store-managed writable memory";
+            return std::nullopt;
+        }
+
+        if (region->offset + size > region->size) {
+            LOG(ERROR) << context << ": buffer range exceeds writable region";
+            return std::nullopt;
+        }
+
+        return region;
+    }
+
+    std::optional<ParsedParallelismManifest> load_parallelism_manifest(
+        const std::string &key, const std::string &context) const {
+        std::shared_ptr<BufferHandle> manifest_handle;
+        {
+            py::gil_scoped_release release_gil;
+            manifest_handle =
+                store_->get_buffer(get_parallelism_manifest_key_name(key));
+        }
+        auto parsed_manifest =
+            parse_writer_shard_manifest(manifest_handle.get());
+        if (!parsed_manifest.has_value()) {
+            return std::nullopt;
+        }
+        if (parsed_manifest->manifest.header.ndim !=
+            static_cast<int32_t>(parsed_manifest->global_shape.size())) {
+            LOG(ERROR) << context << ": invalid parallelism manifest shape";
+            return std::nullopt;
+        }
+        return parsed_manifest;
+    }
+
+#include "store_py_parallel_read.h"
+
     // --- Upsert tensor methods ---
 
+    int upsert_tensor_info_impl(const std::string &key,
+                                const PyTensorInfo &info,
+                                const ReplicateConfig &config) {
+        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),
+                            info.metadata.header.data_offset);
+        append_tensor_payload_span(values, 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_impl(const std::string &key, pybind11::object tensor,
                            const ReplicateConfig &config) {
         auto info = extract_tensor_info(tensor, key);
@@ -977,8 +1288,7 @@ class MooncakeStorePyWrapper {
         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);
+        append_tensor_payload_span(values, info.data_ptr, info.tensor_size);
 
         py::gil_scoped_release release_gil;
         int ret = store_->upsert_parts(key, values, config);
@@ -997,28 +1307,29 @@ class MooncakeStorePyWrapper {
         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";
+    int upsert_tensor_with_tp_impl(
+        const std::string &key, pybind11::object tensor,
+        const ReplicateConfig &config = ReplicateConfig{}, int tp_size = 1,
+        int split_dim = 0) {
+        try {
+            auto infos = build_tp_shard_infos(
+                tensor, tp_size, split_dim,
+                [this, &key](int rank) { return get_tp_key_name(key, rank); });
+            if (!infos.has_value()) {
+                return to_py_ret(ErrorCode::INVALID_PARAMS);
+            }
+
+            for (int rank = 0; rank < tp_size; ++rank) {
+                int ret = upsert_tensor_info_impl(get_tp_key_name(key, rank),
+                                                  (*infos)[rank], config);
+                if (ret != 0) return ret;
+            }
+            return 0;
+
+        } catch (const std::exception &e) {
+            LOG(ERROR) << "Failed to upsert tensor with tp: " << e.what();
             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(
@@ -1052,12 +1363,19 @@ class MooncakeStorePyWrapper {
                 return std::vector<int>(keys.size(),
                                         to_py_ret(ErrorCode::INVALID_PARAMS));
             }
-            if (sizes[i] <= sizeof(TensorMetadata)) {
+            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));
             }
+            if (!ParseTensorMetadata(
+                     reinterpret_cast<const char *>(buffer_ptrs[i]), sizes[i])
+                     .has_value()) {
+                LOG(ERROR) << "Invalid tensor metadata at index " << i;
+                return std::vector<int>(keys.size(),
+                                        to_py_ret(ErrorCode::INVALID_PARAMS));
+            }
         }
         std::vector<void *> buffers;
         buffers.reserve(buffer_ptrs.size());
@@ -1073,6 +1391,12 @@ class MooncakeStorePyWrapper {
         const std::vector<std::string> &keys,
         const pybind11::list &tensors_list,
         const ReplicateConfig &config = ReplicateConfig{}) {
+        auto group_ids_error = ValidateGroupIdsForBatchConfig(
+            config, keys.size(), "batch_upsert_tensor");
+        if (!group_ids_error.empty()) {
+            return group_ids_error;
+        }
+
         std::vector<PyTensorInfo> infos(keys.size());
         std::vector<int> results(keys.size(), 0);
 
@@ -1112,9 +1436,11 @@ class MooncakeStorePyWrapper {
                 // 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);
+                if (infos[i].tensor_size > 0) {
+                    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());
@@ -1126,8 +1452,10 @@ class MooncakeStorePyWrapper {
             }
 
             if (!valid_keys.empty()) {
+                ReplicateConfig write_config =
+                    MakeIndexedConfig(config, original_indices);
                 std::vector<int> op_results = store_->batch_upsert_from(
-                    valid_keys, buffer_ptrs, buffer_sizes, config);
+                    valid_keys, buffer_ptrs, buffer_sizes, write_config);
                 for (size_t i = 0; i < op_results.size(); ++i) {
                     results[original_indices[i]] = op_results[i];
                 }
@@ -1152,6 +1480,48 @@ class MooncakeStorePyWrapper {
         return batch_upsert_tensor_impl(keys, tensors_list, 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);
+        }
+        if (!ParseTensorMetadata(reinterpret_cast<const char *>(buffer), size)
+                 .has_value()) {
+            LOG(ERROR) << "Invalid tensor metadata";
+            return to_py_ret(ErrorCode::INVALID_PARAMS);
+        }
+        py::gil_scoped_release release_gil;
+        return store_->upsert_from(key, buffer, size, ReplicateConfig{});
+    }
+
+    int validate_replicate_config(
+        const ReplicateConfig &config = ReplicateConfig{}) {
+        if (!config.preferred_segments.empty() &&
+            config.preferred_segments.size() != config.replica_num) {
+            LOG(ERROR) << "Preferred segments size ("
+                       << config.preferred_segments.size()
+                       << ") must match replica_num (" << config.replica_num
+                       << ")";
+            return to_py_ret(ErrorCode::INVALID_PARAMS);
+        }
+        return 0;
+    }
+
     int upsert_pub_tensor(const std::string &key, pybind11::object tensor,
                           const ReplicateConfig &config = ReplicateConfig{}) {
         if (!is_client_initialized() || use_dummy_client_) {
@@ -1190,19 +1560,6 @@ class MooncakeStorePyWrapper {
     }
 
     // --- End Upsert tensor methods ---
-
-    int validate_replicate_config(
-        const ReplicateConfig &config = ReplicateConfig{}) {
-        if (!config.preferred_segments.empty() &&
-            config.preferred_segments.size() != config.replica_num) {
-            LOG(ERROR) << "Preferred segments size ("
-                       << config.preferred_segments.size()
-                       << ") must match replica_num (" << config.replica_num
-                       << ")";
-            return to_py_ret(ErrorCode::INVALID_PARAMS);
-        }
-        return 0;
-    }
     int pub_tensor(const std::string &key, pybind11::object tensor,
                    const ReplicateConfig &config = ReplicateConfig{}) {
         if (!is_client_initialized() || use_dummy_client_) {
@@ -1391,18 +1748,61 @@ class MooncakeHostMemAllocatorPyWrapper {
     ~MooncakeHostMemAllocatorPyWrapper() { shm_helper_ = nullptr; }
 };
 
+uintptr_t get_alloc_func_addr() {
+    return reinterpret_cast<uintptr_t>(&hugepage_memory_alloc);
+}
+
+uintptr_t get_free_func_addr() {
+    return reinterpret_cast<uintptr_t>(&hugepage_memory_free);
+}
+
+class MooncakeDistributedNoFRegisterPyWrapper {
+   public:
+    std::shared_ptr<NoFRegisterClient> register_{nullptr};
+
+    MooncakeDistributedNoFRegisterPyWrapper() = default;
+};
+
 PYBIND11_MODULE(store, m) {
+    m.def("_serialize_tensor", &serialize_tensor_metadata,
+          "Inspect a torch tensor as Mooncake tensor metadata, data pointer, "
+          "size, and owner.");
+    m.def("_tensor_metadata_size", &tensor_metadata_size,
+          "Return the serialized Mooncake tensor metadata size.");
+    m.def("_deserialize_tensor", &deserialize_tensor_from_bytes,
+          "Deserialize Mooncake tensor metadata plus payload bytes.");
+
+    // Object data type classification
+    py::enum_<ObjectDataType>(m, "ObjectDataType")
+        .value("UNKNOWN", ObjectDataType::UNKNOWN)
+        .value("KVCACHE", ObjectDataType::KVCACHE)
+        .value("TENSOR", ObjectDataType::TENSOR)
+        .value("WEIGHT", ObjectDataType::WEIGHT)
+        .value("SAMPLE", ObjectDataType::SAMPLE)
+        .value("ACTIVATION", ObjectDataType::ACTIVATION)
+        .value("GRADIENT", ObjectDataType::GRADIENT)
+        .value("OPTIMIZER_STATE", ObjectDataType::OPTIMIZER_STATE)
+        .value("METADATA", ObjectDataType::METADATA)
+        .value("GENERAL", ObjectDataType::GENERAL)
+        .value("HIDDEN_STATE", ObjectDataType::HIDDEN_STATE)
+        .export_values();
+
     // Define the ReplicateConfig class
     py::class_<ReplicateConfig>(m, "ReplicateConfig")
         .def(py::init<>())
         .def_readwrite("replica_num", &ReplicateConfig::replica_num)
+        .def_readwrite("nof_replica_num", &ReplicateConfig::nof_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_nof_segments",
+                       &ReplicateConfig::preferred_nof_segments)
         .def_readwrite("preferred_segment", &ReplicateConfig::preferred_segment)
         .def_readwrite("prefer_alloc_in_same_node",
                        &ReplicateConfig::prefer_alloc_in_same_node)
+        .def_readwrite("data_type", &ReplicateConfig::data_type)
+        .def_readwrite("group_ids", &ReplicateConfig::group_ids)
         .def("__str__", [](const ReplicateConfig &config) {
             std::ostringstream oss;
             oss << config;
@@ -1434,6 +1834,9 @@ PYBIND11_MODULE(store, m) {
         .def("is_disk_replica",
              static_cast<bool (Replica::Descriptor::*)() const noexcept>(
                  &Replica::Descriptor::is_disk_replica))
+        .def("is_local_disk_replica",
+             static_cast<bool (Replica::Descriptor::*)() const noexcept>(
+                 &Replica::Descriptor::is_local_disk_replica))
         .def(
             "get_memory_descriptor",
             static_cast<const MemoryDescriptor &(Replica::Descriptor::*)()
@@ -1535,6 +1938,8 @@ PYBIND11_MODULE(store, m) {
             }
         });
 
+    bind_buffer_pool(m);
+
     py::class_<MooncakeHostMemAllocatorPyWrapper>(m, "MooncakeHostMemAllocator")
         .def(py::init<>())
         .def("alloc",
@@ -1549,10 +1954,118 @@ PYBIND11_MODULE(store, m) {
                  return self.shm_helper_->free(reinterpret_cast<void *>(ptr));
              });
 
+    py::class_<ParallelAxisSpec>(m, "ParallelAxis")
+        .def(py::init<>())
+        .def_readwrite("kind", &ParallelAxisSpec::kind)
+        .def_readwrite("rank", &ParallelAxisSpec::rank)
+        .def_readwrite("size", &ParallelAxisSpec::size)
+        .def_readwrite("split_dim", &ParallelAxisSpec::split_dim)
+        .def_readwrite("expert_id", &ParallelAxisSpec::expert_id)
+        .def_readwrite("stage_id", &ParallelAxisSpec::stage_id);
+
+    auto make_pyclient_capsule =
+        [](const std::shared_ptr<PyClient> &store) -> py::object {
+        if (!store) {
+            return py::none();
+        }
+        auto ptr = std::make_unique<std::shared_ptr<PyClient>>(store);
+        std::shared_ptr<PyClient> *raw_ptr = ptr.get();
+        py::object cap =
+            py::capsule(raw_ptr, "mooncake.PyClient.shared_ptr", [](void *p) {
+                delete static_cast<std::shared_ptr<PyClient> *>(p);
+            });
+        (void)ptr.release();  // Transfer ownership to capsule
+        return cap;
+    };
+
+    py::class_<TensorParallelismSpec>(m, "TensorParallelism")
+        .def(py::init<>())
+        .def_readwrite("axes", &TensorParallelismSpec::axes);
+
+    py::class_<ReadTargetSpec>(m, "ReadTarget")
+        .def(py::init<>())
+        .def_property(
+            "mode",
+            [](const ReadTargetSpec &self) {
+                switch (self.mode) {
+                    case ReadTargetMode::AS_STORED:
+                        return std::string("as_stored");
+                    case ReadTargetMode::SHARD:
+                        return std::string("shard");
+                    case ReadTargetMode::FULL:
+                        return std::string("full");
+                }
+                return std::string("as_stored");
+            },
+            [](ReadTargetSpec &self, const std::string &mode) {
+                auto parsed = parse_read_target_mode(py::str(mode));
+                if (!parsed.has_value()) {
+                    throw std::runtime_error("Unsupported ReadTarget mode");
+                }
+                self.mode = *parsed;
+            })
+        .def_readwrite("parallelism", &ReadTargetSpec::parallelism);
+
+    py::class_<MooncakeDistributedNoFRegisterPyWrapper>(
+        m, "MooncakeDistributedNoFRegister")
+        .def(py::init<>())
+        .def("real_register",
+             [](MooncakeDistributedNoFRegisterPyWrapper &self,
+                const std::string &nqn = "", size_t nsid = 1,
+                const std::string &traddr = "", size_t trsvcid = 4420,
+                uintptr_t base = 0x0, size_t size = 1024,
+                const std::string &master_server_addr = "127.0.0.1:50051") {
+                 self.register_ = std::make_shared<NoFRegisterClient>();
+                 return self.register_->set_register(nqn, nsid, traddr, trsvcid,
+                                                     base, size,
+                                                     master_server_addr);
+             })
+        .def("real_unregister_by_endpoint",
+             [](MooncakeDistributedNoFRegisterPyWrapper &self,
+                const std::string &nqn = "", size_t nsid = 1,
+                const std::string &traddr = "", size_t trsvcid = 4420,
+                const std::string &master_server_addr = "127.0.0.1:50051") {
+                 self.register_ = std::make_shared<NoFRegisterClient>();
+                 return self.register_->set_unregister_by_endpoint(
+                     nqn, nsid, traddr, trsvcid, master_server_addr);
+             });
     // Create a wrapper that exposes DistributedObjectStore with Python-specific
     // methods
+    // Helper function to extract PyClient shared_ptr from
+    // MooncakeStorePyWrapper This is used by EngramStore to get the underlying
+    // PyClient We return it as a Python capsule to avoid type registration
+    // issues
+    m.def(
+        "_get_pyclient_from_wrapper",
+        [](MooncakeStorePyWrapper &wrapper) -> py::object {
+            if (!wrapper.store_) {
+                return py::none();
+            }
+            // Return as a capsule containing the shared_ptr
+            // The caller (engram_store_py.cpp) will extract it
+            // Use unique_ptr for RAII: if py::capsule throws, the pointer is
+            // freed; otherwise release() transfers ownership to the capsule.
+            auto ptr =
+                std::make_unique<std::shared_ptr<PyClient>>(wrapper.store_);
+            std::shared_ptr<PyClient> *raw_ptr = ptr.get();
+            py::object cap = py::capsule(raw_ptr, [](void *p) {
+                delete static_cast<std::shared_ptr<PyClient> *>(p);
+            });
+            (void)ptr.release();  // Transfer ownership to capsule
+            return cap;
+        },
+        py::arg("wrapper"),
+        "Get PyClient from MooncakeDistributedStore (internal use, returns "
+        "capsule)");
+
     py::class_<MooncakeStorePyWrapper>(m, "MooncakeDistributedStore")
         .def(py::init<>())
+        .def(
+            "_get_pyclient_capsule",
+            [make_pyclient_capsule](MooncakeStorePyWrapper &self)
+                -> py::object { return make_pyclient_capsule(self.store_); },
+            "Internal use: expose the underlying PyClient handle as a typed "
+            "capsule")
         .def(
             "setup",
             [](MooncakeStorePyWrapper &self, const std::string &local_hostname,
@@ -1564,7 +2077,8 @@ PYBIND11_MODULE(store, m) {
                const std::string &master_server_addr = "127.0.0.1:50051",
                const py::object &engine = py::none(),
                bool enable_ssd_offload = false,
-               const std::string &ssd_offload_path = "") {
+               const std::string &ssd_offload_path = "",
+               const std::string &tenant_id = "default") {
                 auto real_client = self.init_real_client();
                 std::shared_ptr<mooncake::TransferEngine> transfer_engine =
                     nullptr;
@@ -1576,14 +2090,14 @@ PYBIND11_MODULE(store, m) {
                     local_hostname, metadata_server, global_segment_size,
                     local_buffer_size, protocol, rdma_devices,
                     master_server_addr, transfer_engine, "", enable_ssd_offload,
-                    ssd_offload_path);
+                    ssd_offload_path, tenant_id);
             },
             py::arg("local_hostname"), py::arg("metadata_server"),
             py::arg("global_segment_size"), py::arg("local_buffer_size"),
             py::arg("protocol"), py::arg("rdma_devices"),
             py::arg("master_server_addr"), py::arg("engine") = py::none(),
             py::arg("enable_ssd_offload") = false,
-            py::arg("ssd_offload_path") = "")
+            py::arg("ssd_offload_path") = "", py::arg("tenant_id") = "default")
         .def(
             "setup",
             [](MooncakeStorePyWrapper &self, const py::dict &config_dict) {
@@ -1622,6 +2136,7 @@ PYBIND11_MODULE(store, m) {
                 auto &resource_tracker = ResourceTracker::getInstance();
                 self.use_dummy_client_ = true;
                 self.store_ = std::make_shared<DummyClient>();
+                self.real_client_.reset();
                 resource_tracker.registerInstance(
                     std::static_pointer_cast<PyClient>(self.store_));
                 auto [ip, port] = parseHostNameWithPort(server_address);
@@ -1638,6 +2153,18 @@ PYBIND11_MODULE(store, m) {
                  return self.store_->initAll(protocol, device_name,
                                              mount_segment_size);
              })
+        .def("mount_segment", &MooncakeStorePyWrapper::mount_segment,
+             py::arg("path"), py::arg("size"), py::arg("offset") = 0,
+             py::arg("protocol") = "tcp", py::arg("location") = "")
+        .def("unmount_segment", &MooncakeStorePyWrapper::unmount_segment,
+             py::arg("segment_ids"), py::arg("grace_period_seconds") = 0)
+        .def("allocate_and_mount_segment",
+             &MooncakeStorePyWrapper::allocate_and_mount_segment,
+             py::arg("size"), py::arg("protocol") = "tcp",
+             py::arg("location") = "")
+        .def("unmount_and_free_segment",
+             &MooncakeStorePyWrapper::unmount_and_free_segment,
+             py::arg("segment_ids"), py::arg("grace_period_seconds") = 0)
         .def("alloc_from_mem_pool",
              [](MooncakeStorePyWrapper &self, size_t size) {
                  py::gil_scoped_release release;
@@ -1645,6 +2172,13 @@ PYBIND11_MODULE(store, m) {
              })
         .def("get", &mooncake::MooncakeStorePyWrapper::get)
         .def("get_batch", &mooncake::MooncakeStorePyWrapper::get_batch)
+        .def("get_offload_rpc_read_count",
+             [](MooncakeStorePyWrapper &self) -> int64_t {
+                 auto real_client =
+                     std::dynamic_pointer_cast<RealClient>(self.store_);
+                 return real_client ? real_client->get_offload_rpc_read_count()
+                                    : 0;
+             })
         .def(
             "get_buffer",
             [](MooncakeStorePyWrapper &self, const std::string &key) {
@@ -1834,7 +2368,8 @@ PYBIND11_MODULE(store, m) {
             "  size: The size of buffer.\n"
             "  tp_rank: The current tensor parallel rank (default 0).\n"
             "  tp_size: The total tensor parallel size (default 1).\n"
-            "  split_dim: The dimension to split the tensor along (default 0).")
+            "  split_dim: The tensor dimension used for TP sharding "
+            "(default 0).")
         .def(
             "batch_get_tensor_with_tp_into",
             &MooncakeStorePyWrapper::batch_get_tensor_with_tp_into,
@@ -1893,6 +2428,91 @@ PYBIND11_MODULE(store, m) {
              py::arg("tensors_list"), py::arg("config") = ReplicateConfig{},
              "Batch upsert PyTorch tensors with configurable replication "
              "settings")
+        .def("get_tensor_with_parallelism",
+             &MooncakeStorePyWrapper::get_tensor_with_parallelism,
+             py::arg("key"), py::arg("target") = py::none(),
+             "Get a PyTorch tensor from the store using a ReadTarget request.")
+        .def("batch_get_tensor_with_parallelism",
+             &MooncakeStorePyWrapper::batch_get_tensor_with_parallelism,
+             py::arg("keys"), py::arg("targets") = py::none(),
+             "Get a batch of PyTorch tensors from the store using ReadTarget "
+             "requests.")
+        .def("get_tensor_with_parallelism_into",
+             &MooncakeStorePyWrapper::get_tensor_with_parallelism_into,
+             py::arg("key"), py::arg("buffer_ptr"), py::arg("size"),
+             py::arg("target") = py::none(),
+             "Get a PyTorch tensor from the store directly into a "
+             "pre-allocated buffer using a ReadTarget request.")
+        .def("batch_get_tensor_with_parallelism_into",
+             &MooncakeStorePyWrapper::batch_get_tensor_with_parallelism_into,
+             py::arg("keys"), py::arg("buffer_ptrs"), py::arg("sizes"),
+             py::arg("targets") = py::none(),
+             "Get a batch of PyTorch tensors into pre-allocated buffers using "
+             "ReadTarget requests.")
+        .def("put_tensor_with_parallelism",
+             &MooncakeStorePyWrapper::put_tensor_with_parallelism,
+             py::arg("key"), py::arg("tensor"),
+             py::arg("parallelism") = py::none(),
+             py::arg("config") = ReplicateConfig{},
+             py::arg("writer_partition") = py::none(),
+             "Put a PyTorch tensor into the store using a TensorParallelism "
+             "request or a writer_partition request.")
+        .def("batch_put_tensor_with_parallelism",
+             &MooncakeStorePyWrapper::batch_put_tensor_with_parallelism,
+             py::arg("keys"), py::arg("tensors_list"),
+             py::arg("parallelisms") = py::none(),
+             py::arg("config") = ReplicateConfig{},
+             py::arg("writer_partitions") = py::none(),
+             "Put a batch of PyTorch tensors into the store using "
+             "TensorParallelism requests or writer_partition requests.")
+        .def("put_tensor_with_parallelism_from",
+             &MooncakeStorePyWrapper::put_tensor_with_parallelism_from,
+             py::arg("key"), py::arg("buffer_ptr"), py::arg("size"),
+             py::arg("parallelism") = py::none(),
+             py::arg("config") = ReplicateConfig{},
+             py::arg("writer_partition") = py::none(),
+             "Put a tensor directly from a pre-allocated buffer using a "
+             "TensorParallelism request or a writer_partition request.")
+        .def("batch_put_tensor_with_parallelism_from",
+             &MooncakeStorePyWrapper::batch_put_tensor_with_parallelism_from,
+             py::arg("keys"), py::arg("buffer_ptrs"), py::arg("sizes"),
+             py::arg("parallelisms") = py::none(),
+             py::arg("config") = ReplicateConfig{},
+             py::arg("writer_partitions") = py::none(),
+             "Put a batch of tensors directly from pre-allocated buffers using "
+             "TensorParallelism requests or writer_partition requests.")
+        .def("upsert_tensor_with_parallelism",
+             &MooncakeStorePyWrapper::upsert_tensor_with_parallelism,
+             py::arg("key"), py::arg("tensor"),
+             py::arg("parallelism") = py::none(),
+             py::arg("config") = ReplicateConfig{},
+             py::arg("writer_partition") = py::none(),
+             "Upsert a PyTorch tensor into the store using a TensorParallelism "
+             "request or a writer_partition request.")
+        .def("upsert_tensor_with_parallelism_from",
+             &MooncakeStorePyWrapper::upsert_tensor_with_parallelism_from,
+             py::arg("key"), py::arg("buffer_ptr"), py::arg("size"),
+             py::arg("parallelism") = py::none(),
+             py::arg("config") = ReplicateConfig{},
+             py::arg("writer_partition") = py::none(),
+             "Upsert a tensor directly from a pre-allocated buffer using a "
+             "TensorParallelism request or a writer_partition request.")
+        .def("batch_upsert_tensor_with_parallelism_from",
+             &MooncakeStorePyWrapper::batch_upsert_tensor_with_parallelism_from,
+             py::arg("keys"), py::arg("buffer_ptrs"), py::arg("sizes"),
+             py::arg("parallelisms") = py::none(),
+             py::arg("config") = ReplicateConfig{},
+             py::arg("writer_partitions") = py::none(),
+             "Upsert a batch of tensors directly from pre-allocated buffers "
+             "using TensorParallelism requests or writer_partition requests.")
+        .def("batch_upsert_tensor_with_parallelism",
+             &MooncakeStorePyWrapper::batch_upsert_tensor_with_parallelism,
+             py::arg("keys"), py::arg("tensors_list"),
+             py::arg("parallelisms") = py::none(),
+             py::arg("config") = ReplicateConfig{},
+             py::arg("writer_partitions") = py::none(),
+             "Upsert a batch of PyTorch tensors into the store using "
+             "TensorParallelism requests or writer_partition requests.")
         .def(
             "upsert_from",
             [](MooncakeStorePyWrapper &self, const std::string &key,
@@ -2383,6 +3003,12 @@ PYBIND11_MODULE(store, m) {
         py::arg("node"),
         "Bind the current thread and memory allocation preference to the "
         "specified NUMA node");
+
+    m.def("get_alloc_func_addr", &get_alloc_func_addr);
+    m.def("get_free_func_addr", &get_free_func_addr);
+
+    // Add EngramStore bindings
+    mooncake::engram::bind_engram_store(m);
 }
 
 }  // namespace mooncake
diff --git a/mooncake-integration/store/store_py_internal.h b/mooncake-integration/store/store_py_internal.h
new file mode 100644
index 00000000..a7d030a4
--- /dev/null
+++ b/mooncake-integration/store/store_py_internal.h
@@ -0,0 +1,1500 @@
+std::vector<int64_t> tensor_shape_to_vector(const py::handle &tensor) {
+    py::tuple shape_tuple = py::cast<py::tuple>(tensor.attr("shape"));
+    std::vector<int64_t> shape;
+    shape.reserve(shape_tuple.size());
+    for (const auto &dim : shape_tuple) {
+        shape.push_back(dim.cast<int64_t>());
+    }
+    return shape;
+}
+
+enum class ReadTargetMode {
+    AS_STORED,
+    SHARD,
+    FULL,
+};
+
+struct ParallelAxisSpec {
+    std::string kind;
+    int rank{0};
+    int size{1};
+    std::optional<int> split_dim;
+    std::optional<int> expert_id;
+    std::optional<int> stage_id;
+};
+
+struct TensorParallelismSpec {
+    std::vector<ParallelAxisSpec> axes;
+};
+
+struct ReadTargetSpec {
+    ReadTargetMode mode{ReadTargetMode::AS_STORED};
+    std::optional<TensorParallelismSpec> parallelism;
+};
+
+struct WriterPartitionSpec {
+    int rank{0};
+    int size{1};
+    int split_dim{0};
+};
+
+struct WriterShardManifestHeader {
+    uint32_t magic;
+    uint16_t version;
+    uint16_t header_size;
+    int32_t dtype;
+    int32_t ndim;
+    int32_t split_dim;
+    int32_t shard_count;
+    int32_t reserved0;
+};
+
+struct WriterShardManifest {
+    WriterShardManifestHeader header;
+    TensorShape global_shape;
+};
+
+struct ParsedWriterShardManifest {
+    WriterShardManifest manifest;
+    std::vector<int64_t> global_shape;
+};
+
+using ParallelismManifest = WriterShardManifest;
+using ParsedParallelismManifest = ParsedWriterShardManifest;
+
+struct WriterShardTensorInfo {
+    PyTensorInfo info;
+    std::string shard_key;
+    WriterShardManifest manifest;
+};
+
+struct ParallelismShardTensorInfo {
+    PyTensorInfo info;
+    ParallelismManifest manifest;
+};
+
+struct RawTensorShardWritePlan {
+    TensorMetadata metadata;
+    std::vector<std::pair<size_t, size_t>> data_ranges;
+};
+
+struct TensorIntoFragment {
+    std::string read_key;
+    size_t dst_offset{0};
+    size_t src_offset{0};
+    size_t size{0};
+};
+
+struct PlannedQueryResult {
+    std::string read_key;
+    std::optional<CachedQueryResultResponse> cached_query_result;
+};
+
+struct TensorIntoRegularFullFormulaPlan {
+    std::vector<std::string> read_keys;
+    std::vector<int64_t> global_shape;
+    int split_dim{0};
+    size_t element_size{0};
+    size_t data_offset{sizeof(TensorMetadata)};
+};
+
+struct TensorIntoPlan {
+    uintptr_t user_buffer_ptr{0};
+    uintptr_t registered_buffer_ptr{0};
+    size_t registered_buffer_size{0};
+    size_t total_length{0};
+    std::vector<TensorIntoFragment> fragments;
+    std::optional<TensorIntoRegularFullFormulaPlan> regular_full_formula;
+    std::vector<PlannedQueryResult> query_results;
+    std::optional<TensorMetadata> materialized_metadata;
+};
+
+struct ResolvedTensorRead {
+    std::string read_key;
+    ParsedTensorMetadata metadata;
+    std::shared_ptr<BufferHandle> buffer_handle;
+};
+
+enum class ParallelismShardReadStorageRoute {
+    PARALLELISM_SHARD_KEY,
+    WRITER_PARTITION_COMPATIBLE_SHARD,
+};
+
+struct ReconstructedShardSource {
+    std::string read_key;
+    ParsedTensorMetadata metadata;
+    std::optional<CachedQueryResultResponse> cached_query_result;
+};
+
+struct FullTensorReconstructionSources {
+    std::vector<ReconstructedShardSource> sources;
+    std::vector<int64_t> global_shape;
+    int split_dim{0};
+    int32_t dtype{0};
+    bool allow_empty_fragments{false};
+};
+
+std::optional<LayoutAxisKind> parse_layout_axis_kind(const std::string &kind);
+std::optional<TensorMetadata> build_shard_metadata_from_shapes(
+    int32_t dtype, const std::vector<int64_t> &global_shape,
+    const std::vector<int64_t> &local_shape,
+    const std::vector<ParallelAxisSpec> &axes, size_t tensor_size);
+std::optional<TensorMetadata> build_shard_metadata_from_parallelism(
+    const py::handle &full_tensor, const py::handle &shard_tensor,
+    TensorDtype dtype_enum, const std::vector<ParallelAxisSpec> &axes,
+    size_t tensor_size);
+std::optional<size_t> find_tp_axis_index(
+    const std::vector<ParallelAxisSpec> &axes);
+bool validate_uniform_shard_request(const std::vector<int64_t> &shape,
+                                    int split_dim, int shard_count,
+                                    const std::string &error_context);
+std::optional<py::object> materialize_shard_tensor(const py::handle &tensor,
+                                                   int split_dim, int rank,
+                                                   int shard_count);
+std::string get_parallelism_key_name(const std::string &base_key,
+                                     const TensorParallelismSpec &parallelism);
+bool parallelism_matches_metadata(const TensorParallelismSpec &parallelism,
+                                  const TensorMetadata &metadata);
+std::optional<TensorParallelismSpec> parallelism_from_metadata(
+    const TensorMetadata &metadata);
+std::optional<TensorParallelismSpec> parse_tensor_parallelism_spec(
+    const py::object &parallelism_obj);
+std::optional<TensorParallelismSpec> validate_parallelism_spec(
+    const std::optional<TensorParallelismSpec> &parallelism,
+    const std::string &error_context, bool allow_empty);
+std::optional<TensorParallelismSpec> canonicalize_parallelism_spec(
+    const TensorParallelismSpec &parallelism);
+std::optional<WriterPartitionSpec> parse_writer_partition_spec(
+    const py::object &obj, const std::string &error_context);
+
+constexpr uint32_t kWriterShardManifestMagic = 0x574d414e;
+constexpr uint16_t kWriterShardManifestVersion = 1;
+
+std::string get_writer_manifest_key_name(const std::string &base_key) {
+    return base_key + "__writer_manifest";
+}
+
+std::string get_parallelism_manifest_key_name(const std::string &base_key) {
+    return base_key + "__parallelism_manifest";
+}
+
+std::string get_writer_shard_key_name(const std::string &base_key,
+                                      const WriterPartitionSpec &writer) {
+    return base_key + "__writer_" + std::to_string(writer.rank) + "of" +
+           std::to_string(writer.size) + "_sd" +
+           std::to_string(writer.split_dim);
+}
+
+std::optional<ParsedWriterShardManifest> parse_writer_shard_manifest(
+    BufferHandle *buffer_handle) {
+    if (!buffer_handle || buffer_handle->size() < sizeof(WriterShardManifest)) {
+        return std::nullopt;
+    }
+
+    ParsedWriterShardManifest parsed{};
+    std::memcpy(&parsed.manifest, buffer_handle->ptr(),
+                sizeof(WriterShardManifest));
+    if (parsed.manifest.header.magic != kWriterShardManifestMagic ||
+        parsed.manifest.header.version != kWriterShardManifestVersion ||
+        parsed.manifest.header.header_size != sizeof(WriterShardManifest) ||
+        parsed.manifest.header.ndim < 0 ||
+        parsed.manifest.header.ndim > static_cast<int32_t>(kMaxTensorDims) ||
+        parsed.manifest.header.shard_count <= 0 ||
+        parsed.manifest.header.split_dim < 0 ||
+        parsed.manifest.header.split_dim >= parsed.manifest.header.ndim) {
+        return std::nullopt;
+    }
+
+    parsed.global_shape = TensorShapeToVector(parsed.manifest.global_shape,
+                                              parsed.manifest.header.ndim);
+    for (auto dim : parsed.global_shape) {
+        if (dim < 0) {
+            return std::nullopt;
+        }
+    }
+    return parsed;
+}
+
+WriterShardManifest build_writer_shard_manifest_from_shape(
+    const std::vector<int64_t> &shape, int32_t dtype,
+    const WriterPartitionSpec &writer) {
+    WriterShardManifest manifest{};
+    manifest.header.magic = kWriterShardManifestMagic;
+    manifest.header.version = kWriterShardManifestVersion;
+    manifest.header.header_size = sizeof(WriterShardManifest);
+    manifest.header.dtype = dtype;
+    manifest.header.ndim = static_cast<int32_t>(shape.size());
+    manifest.header.split_dim = writer.split_dim;
+    manifest.header.shard_count = writer.size;
+    manifest.header.reserved0 = 0;
+    manifest.global_shape = MakeTensorShape(shape);
+    return manifest;
+}
+
+ParallelismManifest build_parallelism_manifest_from_shape(
+    const std::vector<int64_t> &shape, int32_t dtype, int split_dim,
+    int shard_count) {
+    ParallelismManifest manifest{};
+    manifest.header.magic = kWriterShardManifestMagic;
+    manifest.header.version = kWriterShardManifestVersion;
+    manifest.header.header_size = sizeof(ParallelismManifest);
+    manifest.header.dtype = dtype;
+    manifest.header.ndim = static_cast<int32_t>(shape.size());
+    manifest.header.split_dim = split_dim;
+    manifest.header.shard_count = shard_count;
+    manifest.header.reserved0 = 0;
+    manifest.global_shape = MakeTensorShape(shape);
+    return manifest;
+}
+
+WriterShardManifest build_writer_shard_manifest(
+    const py::handle &tensor, TensorDtype dtype_enum,
+    const WriterPartitionSpec &writer) {
+    return build_writer_shard_manifest_from_shape(
+        tensor_shape_to_vector(tensor), static_cast<int32_t>(dtype_enum),
+        writer);
+}
+
+TensorMetadata build_writer_shard_tensor_metadata(
+    const py::handle &full_tensor, const py::handle &shard_tensor,
+    TensorDtype dtype_enum, const WriterPartitionSpec &writer,
+    size_t tensor_size) {
+    ParallelAxisSpec axis_spec{"tp",         writer.rank,
+                               writer.size,  writer.split_dim,
+                               std::nullopt, std::nullopt};
+    auto metadata = build_shard_metadata_from_parallelism(
+        full_tensor, shard_tensor, dtype_enum, {axis_spec}, tensor_size);
+    return metadata.value_or(TensorMetadata{});
+}
+
+TensorMetadata build_full_tensor_metadata(const py::handle &tensor,
+                                          TensorDtype dtype_enum,
+                                          size_t tensor_size) {
+    auto shape = tensor_shape_to_vector(tensor);
+    TensorMetadata metadata = BuildTensorMetadata(
+        static_cast<int32_t>(dtype_enum), shape, shape, TensorLayoutKind::FULL);
+    metadata.header.data_bytes = tensor_size;
+    return metadata;
+}
+
+std::optional<LayoutAxis> build_layout_axis_from_spec(
+    const ParallelAxisSpec &axis_spec, size_t axis_index,
+    int64_t local_split_extent = -1) {
+    auto kind = parse_layout_axis_kind(axis_spec.kind);
+    if (!kind.has_value()) {
+        return std::nullopt;
+    }
+
+    LayoutAxis axis{};
+    axis.kind = static_cast<int32_t>(*kind);
+    axis.axis_index = static_cast<int32_t>(axis_index);
+    axis.shard_rank = axis_spec.rank;
+    axis.shard_count = axis_spec.size;
+    axis.split_dim = axis_spec.split_dim.value_or(-1);
+    axis.reserved0 = 0;
+    axis.reserved1 = 0;
+
+    if (*kind == LayoutAxisKind::EP && axis_spec.expert_id.has_value()) {
+        axis.reserved0 = axis_spec.expert_id.value();
+    }
+    if (*kind == LayoutAxisKind::PP && axis_spec.stage_id.has_value()) {
+        axis.reserved0 = axis_spec.stage_id.value();
+    }
+    if (local_split_extent >= 0) {
+        axis.reserved1 = local_split_extent;
+    }
+    return axis;
+}
+
+std::optional<TensorMetadata> build_shard_metadata_from_shapes(
+    int32_t dtype, const std::vector<int64_t> &global_shape,
+    const std::vector<int64_t> &local_shape,
+    const std::vector<ParallelAxisSpec> &axes, size_t tensor_size) {
+    if (axes.empty() || axes.size() > kMaxLayoutAxes ||
+        global_shape.size() != local_shape.size()) {
+        return std::nullopt;
+    }
+
+    std::vector<LayoutAxis> layout_axes;
+    layout_axes.reserve(axes.size());
+    for (size_t i = 0; i < axes.size(); ++i) {
+        int64_t local_split_extent = -1;
+        if (axes[i].split_dim.has_value()) {
+            int split_dim = axes[i].split_dim.value();
+            if (split_dim < 0 ||
+                split_dim >= static_cast<int>(local_shape.size())) {
+                return std::nullopt;
+            }
+            local_split_extent = local_shape[split_dim];
+        }
+        auto layout_axis =
+            build_layout_axis_from_spec(axes[i], i, local_split_extent);
+        if (!layout_axis.has_value()) {
+            return std::nullopt;
+        }
+        layout_axes.push_back(*layout_axis);
+    }
+
+    TensorMetadata metadata = BuildTensorMetadata(
+        dtype, global_shape, local_shape, TensorLayoutKind::SHARD,
+        std::span<const LayoutAxis>(layout_axes.data(), layout_axes.size()));
+    metadata.header.data_bytes = tensor_size;
+    return metadata;
+}
+
+std::optional<TensorMetadata> build_shard_metadata_from_parallelism(
+    const py::handle &full_tensor, const py::handle &shard_tensor,
+    TensorDtype dtype_enum, const std::vector<ParallelAxisSpec> &axes,
+    size_t tensor_size) {
+    return build_shard_metadata_from_shapes(
+        static_cast<int32_t>(dtype_enum), tensor_shape_to_vector(full_tensor),
+        tensor_shape_to_vector(shard_tensor), axes, tensor_size);
+}
+
+bool validate_uniform_shard_request(const std::vector<int64_t> &shape,
+                                    int split_dim, int shard_count,
+                                    const std::string &error_context);
+
+TensorMetadata build_tp_shard_metadata(const py::handle &full_tensor,
+                                       const py::handle &shard_tensor,
+                                       TensorDtype dtype_enum, int tp_rank,
+                                       int tp_size, int split_dim,
+                                       size_t tensor_size) {
+    ParallelAxisSpec axis_spec{"tp",      tp_rank,      tp_size,
+                               split_dim, std::nullopt, std::nullopt};
+    auto metadata = build_shard_metadata_from_parallelism(
+        full_tensor, shard_tensor, dtype_enum, {axis_spec}, tensor_size);
+    return metadata.value_or(TensorMetadata{});
+}
+
+std::optional<size_t> find_tp_axis_index(
+    const std::vector<ParallelAxisSpec> &axes) {
+    for (size_t i = 0; i < axes.size(); ++i) {
+        auto kind = parse_layout_axis_kind(axes[i].kind);
+        if (kind == std::optional<LayoutAxisKind>(LayoutAxisKind::TP)) {
+            return i;
+        }
+    }
+    return std::nullopt;
+}
+
+std::optional<PyTensorInfo> build_direct_parallelism_shard_info(
+    const py::handle &tensor, const std::vector<ParallelAxisSpec> &axes,
+    const std::string &key_name) {
+    TensorDtype dtype_enum = get_tensor_dtype(tensor.attr("dtype"));
+    if (dtype_enum == TensorDtype::UNKNOWN) {
+        return std::nullopt;
+    }
+
+    auto info = extract_tensor_info(py::reinterpret_borrow<py::object>(tensor),
+                                    key_name);
+    auto metadata = build_shard_metadata_from_parallelism(
+        tensor, tensor, dtype_enum, axes, info.tensor_size);
+    if (!metadata.has_value()) {
+        return std::nullopt;
+    }
+    info.metadata = *metadata;
+    if (!info.valid()) {
+        return std::nullopt;
+    }
+    return info;
+}
+
+std::optional<ParallelismShardTensorInfo> build_direct_parallelism_shard_info(
+    const py::handle &shard_tensor, const std::vector<ParallelAxisSpec> &axes,
+    const std::string &key_name, bool infer_global_shape) {
+    auto tp_axis_index = find_tp_axis_index(axes);
+    if (!tp_axis_index.has_value()) {
+        return std::nullopt;
+    }
+
+    const auto &tp_axis = axes[*tp_axis_index];
+    if (!tp_axis.split_dim.has_value()) {
+        return std::nullopt;
+    }
+
+    auto local_shape = tensor_shape_to_vector(shard_tensor);
+    int split_dim = tp_axis.split_dim.value();
+    if (split_dim < 0 || split_dim >= static_cast<int>(local_shape.size())) {
+        return std::nullopt;
+    }
+
+    auto global_shape = local_shape;
+    if (infer_global_shape) {
+        global_shape[split_dim] *= tp_axis.size;
+    }
+    if (!validate_uniform_shard_request(
+            global_shape, split_dim, tp_axis.size,
+            "build_direct_parallelism_shard_info")) {
+        return std::nullopt;
+    }
+
+    TensorDtype dtype_enum = get_tensor_dtype(shard_tensor.attr("dtype"));
+    if (dtype_enum == TensorDtype::UNKNOWN) {
+        return std::nullopt;
+    }
+
+    ParallelismShardTensorInfo shard_info;
+    shard_info.info = extract_tensor_info(
+        py::reinterpret_borrow<py::object>(shard_tensor), key_name);
+    auto metadata = build_shard_metadata_from_parallelism(
+        shard_tensor, shard_tensor, dtype_enum, axes,
+        shard_info.info.tensor_size);
+    if (!metadata.has_value()) {
+        return std::nullopt;
+    }
+    shard_info.info.metadata = *metadata;
+    shard_info.info.metadata.layout.global_shape =
+        MakeTensorShape(global_shape);
+    if (!shard_info.info.valid()) {
+        return std::nullopt;
+    }
+    shard_info.manifest = build_parallelism_manifest_from_shape(
+        global_shape, static_cast<int32_t>(dtype_enum), split_dim,
+        tp_axis.size);
+    return shard_info;
+}
+
+std::optional<ParallelismShardTensorInfo>
+build_requested_parallelism_shard_info(const py::handle &tensor,
+                                       const TensorParallelismSpec &parallelism,
+                                       const std::string &key_name,
+                                       const std::string &error_context) {
+    auto tp_axis_index = find_tp_axis_index(parallelism.axes);
+    if (!tp_axis_index.has_value()) {
+        return std::nullopt;
+    }
+
+    const auto &tp_axis = parallelism.axes[*tp_axis_index];
+    if (!tp_axis.split_dim.has_value()) {
+        return std::nullopt;
+    }
+
+    const int split_dim = tp_axis.split_dim.value();
+    auto global_shape = tensor_shape_to_vector(tensor);
+    if (!validate_uniform_shard_request(global_shape, split_dim, tp_axis.size,
+                                        error_context)) {
+        return std::nullopt;
+    }
+
+    TensorDtype dtype_enum = get_tensor_dtype(tensor.attr("dtype"));
+    if (dtype_enum == TensorDtype::UNKNOWN) {
+        return std::nullopt;
+    }
+
+    auto shard_tensor =
+        materialize_shard_tensor(tensor, split_dim, tp_axis.rank, tp_axis.size);
+    if (!shard_tensor.has_value()) {
+        return std::nullopt;
+    }
+
+    ParallelismShardTensorInfo shard_info;
+    shard_info.info = extract_tensor_info(*shard_tensor, key_name);
+    auto metadata = build_shard_metadata_from_parallelism(
+        tensor, *shard_tensor, dtype_enum, parallelism.axes,
+        shard_info.info.tensor_size);
+    if (!metadata.has_value()) {
+        return std::nullopt;
+    }
+    shard_info.info.metadata = *metadata;
+    if (!shard_info.info.valid()) {
+        return std::nullopt;
+    }
+    shard_info.manifest = build_parallelism_manifest_from_shape(
+        global_shape, static_cast<int32_t>(dtype_enum), split_dim,
+        tp_axis.size);
+    return shard_info;
+}
+
+bool is_uniform_shardable_dim(int64_t dim_size, int shard_count) {
+    return dim_size >= 0 && shard_count > 0 && dim_size % shard_count == 0;
+}
+
+bool validate_uniform_shard_request(const std::vector<int64_t> &shape,
+                                    int split_dim, int shard_count,
+                                    const std::string &error_context) {
+    if (split_dim < 0 || split_dim >= static_cast<int>(shape.size()) ||
+        shard_count <= 0) {
+        LOG(ERROR) << error_context << ": invalid shard parameters";
+        return false;
+    }
+    if (!is_uniform_shardable_dim(shape[split_dim], shard_count)) {
+        LOG(ERROR) << error_context << ": only uniform sharding is supported";
+        return false;
+    }
+    return true;
+}
+
+std::optional<py::object> materialize_shard_tensor(const py::handle &tensor,
+                                                   int split_dim, int rank,
+                                                   int shard_count) {
+    auto shape = tensor_shape_to_vector(tensor);
+    if (rank < 0 || rank >= shard_count ||
+        !validate_uniform_shard_request(shape, split_dim, shard_count,
+                                        "materialize_shard_tensor")) {
+        return std::nullopt;
+    }
+
+    py::object torch_tensor = py::reinterpret_borrow<py::object>(tensor);
+    const auto [start, size] =
+        calculate_shard_range(shape[split_dim], rank, shard_count);
+    return torch_tensor.attr("narrow")(split_dim, start, size)
+        .attr("contiguous")();
+}
+
+[[maybe_unused]] std::optional<std::vector<PyTensorInfo>> build_tp_shard_infos(
+    const py::handle &tensor, int tp_size, int split_dim,
+    const std::function<std::string(int)> &key_for_rank,
+    const std::vector<ParallelAxisSpec> &axes = {}) {
+    TensorDtype dtype_enum = get_tensor_dtype(tensor.attr("dtype"));
+    if (dtype_enum == TensorDtype::UNKNOWN) {
+        return std::nullopt;
+    }
+
+    auto shape = tensor_shape_to_vector(tensor);
+    if (!validate_uniform_shard_request(shape, split_dim, tp_size,
+                                        "build_tp_shard_infos")) {
+        return std::nullopt;
+    }
+
+    auto tp_axis_index = find_tp_axis_index(axes);
+    std::vector<PyTensorInfo> infos;
+    infos.reserve(tp_size);
+
+    for (int rank = 0; rank < tp_size; ++rank) {
+        auto chunk = materialize_shard_tensor(tensor, split_dim, rank, tp_size);
+        if (!chunk.has_value()) {
+            return std::nullopt;
+        }
+
+        auto info = extract_tensor_info(*chunk, key_for_rank(rank));
+        if (axes.empty()) {
+            info.metadata =
+                build_tp_shard_metadata(tensor, *chunk, dtype_enum, rank,
+                                        tp_size, split_dim, info.tensor_size);
+        } else {
+            auto shard_axes = axes;
+            if (!tp_axis_index.has_value()) {
+                return std::nullopt;
+            }
+            shard_axes[*tp_axis_index].rank = rank;
+            auto metadata = build_shard_metadata_from_parallelism(
+                tensor, *chunk, dtype_enum, shard_axes, info.tensor_size);
+            if (!metadata.has_value()) {
+                return std::nullopt;
+            }
+            info.metadata = *metadata;
+        }
+        if (!info.valid()) {
+            return std::nullopt;
+        }
+        infos.push_back(info);
+    }
+    return infos;
+}
+
+bool is_shard_tensor_metadata(const TensorMetadata &metadata) {
+    return metadata.header.layout_kind ==
+           static_cast<uint32_t>(TensorLayoutKind::SHARD);
+}
+
+const LayoutAxis *find_layout_axis(const TensorMetadata &metadata,
+                                   LayoutAxisKind kind) {
+    for (size_t i = 0; i < metadata.layout.axis_count; ++i) {
+        if (metadata.layout.axes[i].kind == static_cast<int32_t>(kind)) {
+            return &metadata.layout.axes[i];
+        }
+    }
+    return nullptr;
+}
+
+std::optional<LayoutAxisKind> parse_layout_axis_kind(const std::string &kind) {
+    std::string normalized = kind;
+    std::transform(normalized.begin(), normalized.end(), normalized.begin(),
+                   [](unsigned char c) { return std::toupper(c); });
+    if (normalized == "TP") return LayoutAxisKind::TP;
+    if (normalized == "DP") return LayoutAxisKind::DP;
+    if (normalized == "EP") return LayoutAxisKind::EP;
+    if (normalized == "PP") return LayoutAxisKind::PP;
+    LOG(ERROR) << "Unsupported parallel axis kind: " << kind;
+    return std::nullopt;
+}
+
+bool is_single_axis_parallelism_kind(const TensorParallelismSpec &parallelism,
+                                     LayoutAxisKind kind) {
+    return parallelism.axes.size() == 1 &&
+           parse_layout_axis_kind(parallelism.axes[0].kind) ==
+               std::optional<LayoutAxisKind>(kind);
+}
+
+bool uses_legacy_tp_storage_key(const TensorParallelismSpec &parallelism) {
+    return is_single_axis_parallelism_kind(parallelism, LayoutAxisKind::TP);
+}
+
+bool should_use_legacy_single_tp_write_route(
+    const TensorParallelismSpec &parallelism) {
+    return false;
+}
+
+bool can_read_single_tp_request_from_writer_partition_storage(
+    const TensorParallelismSpec &parallelism) {
+    return uses_legacy_tp_storage_key(parallelism);
+}
+
+std::vector<ParallelismShardReadStorageRoute>
+resolve_parallelism_shard_read_routes(
+    const TensorParallelismSpec &parallelism) {
+    std::vector<ParallelismShardReadStorageRoute> routes{
+        ParallelismShardReadStorageRoute::PARALLELISM_SHARD_KEY,
+    };
+    if (can_read_single_tp_request_from_writer_partition_storage(parallelism)) {
+        routes.push_back(ParallelismShardReadStorageRoute::
+                             WRITER_PARTITION_COMPATIBLE_SHARD);
+    }
+    return routes;
+}
+
+std::optional<TensorParallelismSpec> writer_partition_parallelism_from_metadata(
+    const TensorMetadata &metadata) {
+    auto parallelism = parallelism_from_metadata(metadata);
+    if (!parallelism.has_value() || !uses_legacy_tp_storage_key(*parallelism)) {
+        return std::nullopt;
+    }
+    return parallelism;
+}
+
+enum class ParallelismWriteStorageRoute {
+    DIRECT_FULL_OBJECT,
+    WRITER_PARTITION_SHARD,
+    LEGACY_SINGLE_TP,
+    TP_SHARDED_PARALLELISM,
+    GENERIC_PARALLELISM_SHARD,
+};
+
+struct ResolvedParallelismWriteRequest {
+    ParallelismWriteStorageRoute route;
+    std::optional<TensorParallelismSpec> parallelism;
+    std::optional<WriterPartitionSpec> writer_partition;
+};
+
+struct TensorWriteStoreOps {
+    const char *parts_operation_name;
+};
+
+std::optional<ResolvedParallelismWriteRequest>
+resolve_parallelism_write_request(const py::object &parallelism_obj,
+                                  const py::object &writer_partition_obj,
+                                  const std::string &error_context) {
+    if (!writer_partition_obj.is_none()) {
+        if (!parallelism_obj.is_none()) {
+            LOG(ERROR)
+                << error_context
+                << ": writer_partition cannot be combined with parallelism";
+            return std::nullopt;
+        }
+        auto writer_partition =
+            parse_writer_partition_spec(writer_partition_obj, error_context);
+        if (!writer_partition.has_value()) {
+            return std::nullopt;
+        }
+        return ResolvedParallelismWriteRequest{
+            ParallelismWriteStorageRoute::WRITER_PARTITION_SHARD,
+            std::nullopt,
+            *writer_partition,
+        };
+    }
+    if (parallelism_obj.is_none()) {
+        return ResolvedParallelismWriteRequest{
+            ParallelismWriteStorageRoute::DIRECT_FULL_OBJECT,
+            std::nullopt,
+            std::nullopt,
+        };
+    }
+
+    auto parallelism = validate_parallelism_spec(
+        parse_tensor_parallelism_spec(parallelism_obj), error_context, false);
+    if (!parallelism.has_value()) {
+        return std::nullopt;
+    }
+    if (should_use_legacy_single_tp_write_route(*parallelism)) {
+        return ResolvedParallelismWriteRequest{
+            ParallelismWriteStorageRoute::LEGACY_SINGLE_TP,
+            *parallelism,
+            std::nullopt,
+        };
+    }
+    for (const auto &axis : parallelism->axes) {
+        if (parse_layout_axis_kind(axis.kind) ==
+            std::optional<LayoutAxisKind>(LayoutAxisKind::TP)) {
+            return ResolvedParallelismWriteRequest{
+                ParallelismWriteStorageRoute::TP_SHARDED_PARALLELISM,
+                *parallelism,
+                std::nullopt,
+            };
+        }
+    }
+    return ResolvedParallelismWriteRequest{
+        ParallelismWriteStorageRoute::GENERIC_PARALLELISM_SHARD,
+        *parallelism,
+        std::nullopt,
+    };
+}
+
+std::optional<py::list> validate_batch_request_list(
+    const py::object &request_list_obj, size_t expected_size,
+    const std::string &error_context, const char *request_name) {
+    if (!py::isinstance<py::list>(request_list_obj)) {
+        LOG(ERROR) << error_context << ": " << request_name
+                   << " must be a list or None";
+        return std::nullopt;
+    }
+
+    py::list request_list = py::cast<py::list>(request_list_obj);
+    if (request_list.size() != expected_size) {
+        LOG(ERROR) << error_context << ": keys and " << request_name
+                   << " must have the same length";
+        return std::nullopt;
+    }
+    return request_list;
+}
+
+std::optional<py::list> validate_batch_parallelism_list(
+    const py::object &parallelisms, size_t expected_size,
+    const std::string &error_context) {
+    return validate_batch_request_list(parallelisms, expected_size,
+                                       error_context, "parallelisms");
+}
+
+std::optional<py::list> validate_batch_writer_partition_list(
+    const py::object &writer_partitions, size_t expected_size,
+    const std::string &error_context) {
+    return validate_batch_request_list(writer_partitions, expected_size,
+                                       error_context, "writer_partitions");
+}
+
+bool axis_specs_equal(const ParallelAxisSpec &lhs,
+                      const ParallelAxisSpec &rhs) {
+    auto lhs_kind = parse_layout_axis_kind(lhs.kind);
+    auto rhs_kind = parse_layout_axis_kind(rhs.kind);
+    return lhs_kind.has_value() && rhs_kind.has_value() &&
+           lhs_kind == rhs_kind && lhs.rank == rhs.rank &&
+           lhs.size == rhs.size && lhs.split_dim == rhs.split_dim &&
+           lhs.expert_id == rhs.expert_id && lhs.stage_id == rhs.stage_id;
+}
+
+int canonical_axis_kind_order(LayoutAxisKind kind) {
+    switch (kind) {
+        case LayoutAxisKind::DP:
+            return 0;
+        case LayoutAxisKind::TP:
+            return 1;
+        case LayoutAxisKind::EP:
+            return 2;
+        case LayoutAxisKind::PP:
+            return 3;
+        default:
+            return 4;
+    }
+}
+
+std::optional<TensorParallelismSpec> canonicalize_parallelism_spec(
+    const TensorParallelismSpec &parallelism) {
+    TensorParallelismSpec canonical = parallelism;
+    std::sort(canonical.axes.begin(), canonical.axes.end(),
+              [](const ParallelAxisSpec &lhs, const ParallelAxisSpec &rhs) {
+                  auto lhs_kind = parse_layout_axis_kind(lhs.kind);
+                  auto rhs_kind = parse_layout_axis_kind(rhs.kind);
+                  if (!lhs_kind.has_value() || !rhs_kind.has_value()) {
+                      return lhs.kind < rhs.kind;
+                  }
+                  int lhs_order = canonical_axis_kind_order(*lhs_kind);
+                  int rhs_order = canonical_axis_kind_order(*rhs_kind);
+                  if (lhs_order != rhs_order) {
+                      return lhs_order < rhs_order;
+                  }
+                  return lhs.kind < rhs.kind;
+              });
+    return canonical;
+}
+
+const ParallelAxisSpec *find_axis_spec_by_kind(
+    const TensorParallelismSpec &parallelism, LayoutAxisKind kind) {
+    for (const auto &axis : parallelism.axes) {
+        auto axis_kind = parse_layout_axis_kind(axis.kind);
+        if (axis_kind == std::optional<LayoutAxisKind>(kind)) {
+            return &axis;
+        }
+    }
+    return nullptr;
+}
+
+bool parallelism_specs_equal_by_kind(const TensorParallelismSpec &lhs,
+                                     const TensorParallelismSpec &rhs,
+                                     bool allow_tp_rank_mismatch = false) {
+    if (lhs.axes.size() != rhs.axes.size()) {
+        return false;
+    }
+
+    for (const auto &lhs_axis : lhs.axes) {
+        auto lhs_kind = parse_layout_axis_kind(lhs_axis.kind);
+        if (!lhs_kind.has_value()) {
+            return false;
+        }
+        const auto *rhs_axis = find_axis_spec_by_kind(rhs, *lhs_kind);
+        if (!rhs_axis) {
+            return false;
+        }
+        if (*lhs_kind == LayoutAxisKind::TP && allow_tp_rank_mismatch) {
+            if (lhs_axis.size != rhs_axis->size ||
+                lhs_axis.split_dim != rhs_axis->split_dim ||
+                lhs_axis.expert_id != rhs_axis->expert_id ||
+                lhs_axis.stage_id != rhs_axis->stage_id) {
+                return false;
+            }
+            continue;
+        }
+        if (!axis_specs_equal(lhs_axis, *rhs_axis)) {
+            return false;
+        }
+    }
+    return true;
+}
+
+bool is_default_replicate_config(const ReplicateConfig &config) {
+    return config.replica_num == 1 && !config.with_soft_pin &&
+           !config.with_hard_pin && config.preferred_segments.empty() &&
+           config.preferred_segment.empty() &&
+           !config.prefer_alloc_in_same_node && !config.group_ids.has_value();
+}
+
+std::optional<ParallelAxisSpec> parse_parallel_axis_spec(
+    const py::handle &obj) {
+    if (!py::hasattr(obj, "kind") || !py::hasattr(obj, "rank") ||
+        !py::hasattr(obj, "size")) {
+        LOG(ERROR) << "ParallelAxis must provide kind, rank, and size";
+        return std::nullopt;
+    }
+
+    ParallelAxisSpec axis;
+    axis.kind = py::cast<std::string>(obj.attr("kind"));
+    if (!parse_layout_axis_kind(axis.kind).has_value()) {
+        return std::nullopt;
+    }
+    axis.rank = py::cast<int>(obj.attr("rank"));
+    axis.size = py::cast<int>(obj.attr("size"));
+    py::object split_dim = obj.attr("split_dim");
+    if (!split_dim.is_none()) axis.split_dim = py::cast<int>(split_dim);
+    py::object expert_id = obj.attr("expert_id");
+    if (!expert_id.is_none()) axis.expert_id = py::cast<int>(expert_id);
+    py::object stage_id = obj.attr("stage_id");
+    if (!stage_id.is_none()) axis.stage_id = py::cast<int>(stage_id);
+    return axis;
+}
+
+std::optional<TensorParallelismSpec> parse_tensor_parallelism_spec(
+    const py::object &obj) {
+    if (obj.is_none()) {
+        return std::nullopt;
+    }
+    if (!py::hasattr(obj, "axes")) {
+        LOG(ERROR) << "TensorParallelism must provide axes";
+        return std::nullopt;
+    }
+
+    TensorParallelismSpec parallelism;
+    py::list axes = py::cast<py::list>(obj.attr("axes"));
+    for (const auto &axis_obj : axes) {
+        auto axis = parse_parallel_axis_spec(axis_obj);
+        if (!axis.has_value()) {
+            return std::nullopt;
+        }
+        parallelism.axes.push_back(*axis);
+    }
+    return parallelism;
+}
+
+std::optional<ReadTargetMode> parse_read_target_mode(const py::object &obj) {
+    std::string mode = py::cast<std::string>(obj);
+    std::transform(mode.begin(), mode.end(), mode.begin(),
+                   [](unsigned char c) { return std::toupper(c); });
+    if (mode == "AS_STORED") return ReadTargetMode::AS_STORED;
+    if (mode == "SHARD") return ReadTargetMode::SHARD;
+    if (mode == "FULL") return ReadTargetMode::FULL;
+    LOG(ERROR) << "Unsupported ReadTarget mode: " << mode;
+    return std::nullopt;
+}
+
+std::optional<WriterPartitionSpec> parse_writer_partition_spec(
+    const py::object &obj, const std::string &error_context) {
+    if (obj.is_none()) {
+        return std::nullopt;
+    }
+    if (!py::hasattr(obj, "rank") || !py::hasattr(obj, "size") ||
+        !py::hasattr(obj, "split_dim")) {
+        LOG(ERROR)
+            << error_context
+            << ": writer_partition must provide rank, size, and split_dim";
+        return std::nullopt;
+    }
+
+    WriterPartitionSpec writer;
+    writer.rank = py::cast<int>(obj.attr("rank"));
+    writer.size = py::cast<int>(obj.attr("size"));
+    writer.split_dim = py::cast<int>(obj.attr("split_dim"));
+    if (writer.size <= 0 || writer.rank < 0 || writer.rank >= writer.size ||
+        writer.split_dim < 0) {
+        LOG(ERROR) << error_context << ": invalid writer partition";
+        return std::nullopt;
+    }
+    return writer;
+}
+
+bool is_valid_writer_partition(const WriterPartitionSpec &writer,
+                               const std::vector<int64_t> &shape,
+                               const std::string &error_context) {
+    if (writer.size <= 0 || writer.rank < 0 || writer.rank >= writer.size) {
+        LOG(ERROR) << error_context << ": invalid writer rank/size";
+        return false;
+    }
+    if (writer.split_dim < 0 ||
+        writer.split_dim >= static_cast<int>(shape.size())) {
+        LOG(ERROR) << error_context << ": split_dim out of range";
+        return false;
+    }
+    if (!validate_uniform_shard_request(shape, writer.split_dim, writer.size,
+                                        error_context)) {
+        return false;
+    }
+    return true;
+}
+
+std::optional<WriterShardTensorInfo> build_writer_shard_tensor_info(
+    const std::string &key, const py::object &tensor,
+    const WriterPartitionSpec &writer, const std::string &error_context) {
+    auto shape = tensor_shape_to_vector(tensor);
+    if (!is_valid_writer_partition(writer, shape, error_context)) {
+        return std::nullopt;
+    }
+
+    TensorDtype dtype_enum = get_tensor_dtype(tensor.attr("dtype"));
+    if (dtype_enum == TensorDtype::UNKNOWN) {
+        LOG(ERROR) << error_context << ": unsupported tensor dtype";
+        return std::nullopt;
+    }
+
+    auto shard_tensor = materialize_shard_tensor(tensor, writer.split_dim,
+                                                 writer.rank, writer.size);
+    if (!shard_tensor.has_value()) {
+        LOG(ERROR) << error_context << ": failed to materialize writer shard";
+        return std::nullopt;
+    }
+
+    WriterShardTensorInfo writer_info;
+    writer_info.shard_key = get_writer_shard_key_name(key, writer);
+    writer_info.info =
+        extract_tensor_info(*shard_tensor, writer_info.shard_key);
+    writer_info.info.metadata = build_writer_shard_tensor_metadata(
+        tensor, *shard_tensor, dtype_enum, writer,
+        writer_info.info.tensor_size);
+    if (!writer_info.info.valid()) {
+        LOG(ERROR) << error_context << ": invalid writer shard tensor info";
+        return std::nullopt;
+    }
+    writer_info.manifest =
+        build_writer_shard_manifest(tensor, dtype_enum, writer);
+    return writer_info;
+}
+
+std::optional<ReadTargetSpec> parse_read_target_spec(const py::object &obj) {
+    if (obj.is_none()) {
+        return ReadTargetSpec{};
+    }
+    if (!py::hasattr(obj, "mode") || !py::hasattr(obj, "parallelism")) {
+        LOG(ERROR) << "ReadTarget must provide mode and parallelism";
+        return std::nullopt;
+    }
+
+    ReadTargetSpec target;
+    auto mode = parse_read_target_mode(obj.attr("mode"));
+    if (!mode.has_value()) {
+        return std::nullopt;
+    }
+    target.mode = *mode;
+    auto parallelism = parse_tensor_parallelism_spec(obj.attr("parallelism"));
+    if (obj.attr("parallelism").is_none()) {
+        target.parallelism = std::nullopt;
+    } else if (!parallelism.has_value()) {
+        return std::nullopt;
+    } else {
+        target.parallelism = *parallelism;
+    }
+    return target;
+}
+
+std::optional<ResolvedTensorRead> probe_parallelism_shard_read(
+    const std::function<std::optional<ParsedTensorMetadata>(
+        const std::string &, std::shared_ptr<BufferHandle> *)> &load_metadata,
+    const std::string &key, const TensorParallelismSpec &parallelism,
+    ParallelismShardReadStorageRoute route) {
+    switch (route) {
+        case ParallelismShardReadStorageRoute::PARALLELISM_SHARD_KEY: {
+            std::string read_key = get_parallelism_key_name(key, parallelism);
+            std::shared_ptr<BufferHandle> buffer_handle;
+            auto metadata = load_metadata(read_key, &buffer_handle);
+            if (!metadata.has_value() || !parallelism_matches_metadata(
+                                             parallelism, metadata->metadata)) {
+                return std::nullopt;
+            }
+            return ResolvedTensorRead{read_key, *metadata,
+                                      std::move(buffer_handle)};
+        }
+        case ParallelismShardReadStorageRoute::
+            WRITER_PARTITION_COMPATIBLE_SHARD: {
+            const auto &axis = parallelism.axes[0];
+            WriterPartitionSpec writer{
+                .rank = axis.rank,
+                .size = axis.size,
+                .split_dim = axis.split_dim.value_or(0),
+            };
+            std::string read_key = get_writer_shard_key_name(key, writer);
+            std::shared_ptr<BufferHandle> buffer_handle;
+            auto metadata = load_metadata(read_key, &buffer_handle);
+            if (!metadata.has_value()) {
+                return std::nullopt;
+            }
+
+            auto writer_parallelism =
+                writer_partition_parallelism_from_metadata(metadata->metadata);
+            if (!writer_parallelism.has_value() ||
+                !parallelism_matches_metadata(parallelism,
+                                              metadata->metadata)) {
+                return std::nullopt;
+            }
+            return ResolvedTensorRead{read_key, *metadata,
+                                      std::move(buffer_handle)};
+        }
+    }
+    return std::nullopt;
+}
+
+std::optional<ResolvedTensorRead> resolve_parallelism_shard_read(
+    const std::function<std::optional<ParsedTensorMetadata>(
+        const std::string &, std::shared_ptr<BufferHandle> *)> &load_metadata,
+    const std::string &key, const TensorParallelismSpec &parallelism) {
+    for (auto route : resolve_parallelism_shard_read_routes(parallelism)) {
+        auto resolved = probe_parallelism_shard_read(load_metadata, key,
+                                                     parallelism, route);
+        if (resolved.has_value()) {
+            return resolved;
+        }
+    }
+    return std::nullopt;
+}
+
+std::optional<TensorParallelismSpec> validate_parallelism_spec(
+    const std::optional<TensorParallelismSpec> &parallelism,
+    const std::string &error_context, bool allow_empty = false) {
+    if (!parallelism.has_value()) {
+        if (allow_empty) {
+            return parallelism;
+        }
+        LOG(ERROR) << error_context << ": parallelism is required";
+        return std::nullopt;
+    }
+
+    if (parallelism->axes.empty()) {
+        if (allow_empty) {
+            return parallelism;
+        }
+        LOG(ERROR) << error_context << ": parallelism axes cannot be empty";
+        return std::nullopt;
+    }
+    if (parallelism->axes.size() > kMaxLayoutAxes) {
+        LOG(ERROR) << error_context
+                   << ": axis count exceeds max supported axes";
+        return std::nullopt;
+    }
+
+    std::unordered_set<int> seen_kinds;
+    for (const auto &axis : parallelism->axes) {
+        auto kind = parse_layout_axis_kind(axis.kind);
+        if (!kind.has_value()) {
+            LOG(ERROR) << error_context << ": unsupported axis kind "
+                       << axis.kind;
+            return std::nullopt;
+        }
+        if (axis.size <= 0 || axis.rank < 0 || axis.rank >= axis.size) {
+            LOG(ERROR) << error_context << ": invalid rank/size for axis "
+                       << axis.kind;
+            return std::nullopt;
+        }
+        if (!seen_kinds.insert(static_cast<int>(*kind)).second) {
+            LOG(ERROR) << error_context
+                       << ": duplicate axis kinds are not supported";
+            return std::nullopt;
+        }
+
+        switch (*kind) {
+            case LayoutAxisKind::TP:
+                if (!axis.split_dim.has_value()) {
+                    LOG(ERROR)
+                        << error_context << ": TP axis requires split_dim";
+                    return std::nullopt;
+                }
+                break;
+            case LayoutAxisKind::DP:
+                if (axis.split_dim.has_value()) {
+                    LOG(ERROR) << error_context
+                               << ": DP axis must not provide split_dim";
+                    return std::nullopt;
+                }
+                break;
+            case LayoutAxisKind::EP:
+                if (!axis.expert_id.has_value()) {
+                    LOG(ERROR)
+                        << error_context << ": EP axis requires expert_id";
+                    return std::nullopt;
+                }
+                if (axis.split_dim.has_value()) {
+                    LOG(ERROR)
+                        << error_context
+                        << ": EP axis must not provide split_dim in this phase";
+                    return std::nullopt;
+                }
+                break;
+            case LayoutAxisKind::PP:
+                if (!axis.stage_id.has_value()) {
+                    LOG(ERROR)
+                        << error_context << ": PP axis requires stage_id";
+                    return std::nullopt;
+                }
+                if (axis.split_dim.has_value()) {
+                    LOG(ERROR) << error_context
+                               << ": PP axis must not provide split_dim";
+                    return std::nullopt;
+                }
+                break;
+            default:
+                LOG(ERROR) << error_context << ": unsupported axis kind";
+                return std::nullopt;
+        }
+    }
+
+    return parallelism;
+}
+
+std::string normalize_axis_kind_string(LayoutAxisKind kind) {
+    switch (kind) {
+        case LayoutAxisKind::TP:
+            return "tp";
+        case LayoutAxisKind::DP:
+            return "dp";
+        case LayoutAxisKind::EP:
+            return "ep";
+        case LayoutAxisKind::PP:
+            return "pp";
+        default:
+            return "unknown";
+    }
+}
+
+std::string encode_axis_key_suffix(const ParallelAxisSpec &axis) {
+    auto kind = parse_layout_axis_kind(axis.kind);
+    if (!kind.has_value()) {
+        return "invalid";
+    }
+
+    std::string suffix = normalize_axis_kind_string(*kind) + "_" +
+                         std::to_string(axis.rank) + "of" +
+                         std::to_string(axis.size);
+    if (axis.split_dim.has_value()) {
+        suffix += "_sd" + std::to_string(axis.split_dim.value());
+    }
+    if (axis.expert_id.has_value()) {
+        suffix += "_eid" + std::to_string(axis.expert_id.value());
+    }
+    if (axis.stage_id.has_value()) {
+        suffix += "_sid" + std::to_string(axis.stage_id.value());
+    }
+    return suffix;
+}
+
+std::string get_parallelism_key_name(const std::string &base_key,
+                                     const TensorParallelismSpec &parallelism) {
+    if (uses_legacy_tp_storage_key(parallelism)) {
+        return base_key + "_tp_" + std::to_string(parallelism.axes[0].rank);
+    }
+
+    auto canonical = canonicalize_parallelism_spec(parallelism);
+    if (!canonical.has_value()) {
+        return base_key;
+    }
+
+    std::string key = base_key;
+    for (const auto &axis : canonical->axes) {
+        key += "__" + encode_axis_key_suffix(axis);
+    }
+    return key;
+}
+
+std::optional<TensorParallelismSpec> parallelism_from_metadata(
+    const TensorMetadata &metadata) {
+    if (!is_shard_tensor_metadata(metadata)) {
+        return std::nullopt;
+    }
+
+    TensorParallelismSpec parallelism;
+    parallelism.axes.reserve(metadata.layout.axis_count);
+    for (size_t i = 0; i < metadata.layout.axis_count; ++i) {
+        const auto &stored_axis = metadata.layout.axes[i];
+        auto kind = static_cast<LayoutAxisKind>(stored_axis.kind);
+        ParallelAxisSpec axis;
+        axis.kind = normalize_axis_kind_string(kind);
+        axis.rank = stored_axis.shard_rank;
+        axis.size = stored_axis.shard_count;
+        if (stored_axis.split_dim >= 0) {
+            axis.split_dim = stored_axis.split_dim;
+        }
+        if (kind == LayoutAxisKind::EP) {
+            axis.expert_id = stored_axis.reserved0;
+        }
+        if (kind == LayoutAxisKind::PP) {
+            axis.stage_id = stored_axis.reserved0;
+        }
+        parallelism.axes.push_back(axis);
+    }
+    return canonicalize_parallelism_spec(parallelism);
+}
+
+bool parallelism_matches_metadata(const TensorParallelismSpec &parallelism,
+                                  const TensorMetadata &metadata) {
+    auto requested_parallelism = canonicalize_parallelism_spec(parallelism);
+    auto stored_parallelism = parallelism_from_metadata(metadata);
+    if (!requested_parallelism.has_value() || !stored_parallelism.has_value() ||
+        stored_parallelism->axes.size() != requested_parallelism->axes.size()) {
+        return false;
+    }
+
+    for (size_t i = 0; i < requested_parallelism->axes.size(); ++i) {
+        if (!axis_specs_equal(requested_parallelism->axes[i],
+                              stored_parallelism->axes[i])) {
+            return false;
+        }
+    }
+    return true;
+}
+
+std::optional<TensorParallelismSpec>
+resolve_tp_compatible_parallelism_from_metadata(
+    const TensorParallelismSpec &parallelism, const TensorMetadata &metadata,
+    const std::string &error_context) {
+    auto stored_parallelism = parallelism_from_metadata(metadata);
+    if (!stored_parallelism.has_value()) {
+        LOG(ERROR) << error_context << ": missing shard parallelism metadata";
+        return std::nullopt;
+    }
+    if (stored_parallelism->axes.size() != parallelism.axes.size()) {
+        LOG(ERROR) << error_context << ": axis count mismatch";
+        return std::nullopt;
+    }
+
+    auto request_tp_axis_index = find_tp_axis_index(parallelism.axes);
+    auto stored_tp_axis_index = find_tp_axis_index(stored_parallelism->axes);
+    if (!request_tp_axis_index.has_value() ||
+        !stored_tp_axis_index.has_value()) {
+        LOG(ERROR) << error_context << ": reconstruction requires a TP axis";
+        return std::nullopt;
+    }
+    if (*request_tp_axis_index != *stored_tp_axis_index) {
+        LOG(ERROR) << error_context << ": TP axis position mismatch";
+        return std::nullopt;
+    }
+
+    for (size_t i = 0; i < parallelism.axes.size(); ++i) {
+        if (i == *request_tp_axis_index) {
+            const auto request_kind =
+                parse_layout_axis_kind(parallelism.axes[i].kind);
+            const auto stored_kind =
+                parse_layout_axis_kind(stored_parallelism->axes[i].kind);
+            if (!request_kind.has_value() || !stored_kind.has_value() ||
+                request_kind != stored_kind ||
+                parallelism.axes[i].split_dim !=
+                    stored_parallelism->axes[i].split_dim ||
+                parallelism.axes[i].expert_id !=
+                    stored_parallelism->axes[i].expert_id ||
+                parallelism.axes[i].stage_id !=
+                    stored_parallelism->axes[i].stage_id) {
+                LOG(ERROR) << error_context << ": TP axis metadata mismatch";
+                return std::nullopt;
+            }
+            continue;
+        }
+        if (!axis_specs_equal(parallelism.axes[i],
+                              stored_parallelism->axes[i])) {
+            LOG(ERROR) << error_context << ": non-TP axis metadata mismatch";
+            return std::nullopt;
+        }
+    }
+    return stored_parallelism;
+}
+
+std::pair<int64_t, int64_t> calculate_shard_range(int64_t dim_size, int rank,
+                                                  int shard_count) {
+    if (dim_size < 0 || shard_count <= 0 || rank < 0 || rank >= shard_count) {
+        return {0, 0};
+    }
+
+    const int64_t start = (dim_size * static_cast<int64_t>(rank)) / shard_count;
+    const int64_t end =
+        (dim_size * static_cast<int64_t>(rank + 1)) / shard_count;
+    return {start, end - start};
+}
+
+[[maybe_unused]] std::optional<ParsedTensorMetadata>
+parse_tensor_metadata_from_buffer(BufferHandle *buffer_handle, char *usr_buffer,
+                                  int64_t data_length, bool *take_ownership,
+                                  char **exported_data, size_t *total_length) {
+    if (!buffer_handle && !usr_buffer) return std::nullopt;
+    if (buffer_handle && usr_buffer) return std::nullopt;
+
+    *take_ownership = !!buffer_handle;
+    if (*take_ownership) {
+        *total_length = buffer_handle->size();
+        if (*total_length < sizeof(TensorMetadata)) {
+            LOG(ERROR) << "Invalid data format: insufficient data for metadata";
+            return std::nullopt;
+        }
+
+        *exported_data = new char[*total_length];
+        if (!*exported_data) return std::nullopt;
+        memcpy(*exported_data, buffer_handle->ptr(), *total_length);
+    } else {
+        *exported_data = usr_buffer;
+        if (data_length < 0) {
+            LOG(ERROR) << "Get tensor into failed with error code: "
+                       << data_length;
+            return std::nullopt;
+        }
+        *total_length = static_cast<size_t>(data_length);
+        if (*total_length < sizeof(TensorMetadata)) {
+            LOG(ERROR) << "Invalid data format: insufficient data for metadata";
+            return std::nullopt;
+        }
+    }
+
+    auto parsed = ParseTensorMetadata(*exported_data, *total_length);
+    if (!parsed.has_value()) {
+        if (*take_ownership) {
+            delete[] *exported_data;
+            *exported_data = nullptr;
+        }
+        LOG(ERROR) << "Invalid tensor metadata";
+        return std::nullopt;
+    }
+
+    return parsed;
+}
+
+std::optional<ParsedTensorMetadata> parse_tensor_metadata_from_raw_buffer(
+    uintptr_t buffer_ptr, size_t size, const char *operation_name) {
+    if (buffer_ptr == 0) {
+        LOG(ERROR) << operation_name << ": buffer pointer cannot be null";
+        return std::nullopt;
+    }
+    auto parsed =
+        ParseTensorMetadata(reinterpret_cast<const char *>(buffer_ptr), size);
+    if (!parsed.has_value()) {
+        LOG(ERROR) << operation_name << ": invalid tensor object metadata";
+        return std::nullopt;
+    }
+    return parsed;
+}
+
+std::optional<RawTensorShardWritePlan> build_raw_tensor_shard_write_plan(
+    const ParsedTensorMetadata &parsed, int split_dim, int rank,
+    int shard_count, const char *operation_name) {
+    const auto global_shape = TensorShapeToVector(
+        parsed.metadata.layout.global_shape, parsed.metadata.header.ndim);
+    const auto local_shape = TensorShapeToVector(
+        parsed.metadata.layout.local_shape, parsed.metadata.header.ndim);
+    if (global_shape.size() != local_shape.size() || rank < 0 ||
+        rank >= shard_count ||
+        !validate_uniform_shard_request(global_shape, split_dim, shard_count,
+                                        operation_name)) {
+        return std::nullopt;
+    }
+
+    int64_t global_numel = 1;
+    int64_t local_numel = 1;
+    for (size_t dim = 0; dim < global_shape.size(); ++dim) {
+        if (global_shape[dim] < 0 || local_shape[dim] < 0) {
+            LOG(ERROR) << operation_name << ": invalid tensor shape";
+            return std::nullopt;
+        }
+        global_numel *= global_shape[dim];
+        local_numel *= local_shape[dim];
+        if (static_cast<int>(dim) != split_dim &&
+            global_shape[dim] != local_shape[dim]) {
+            LOG(ERROR) << operation_name
+                       << ": unsupported non-split local shape";
+            return std::nullopt;
+        }
+    }
+    if (local_numel < 0 || global_numel < 0) {
+        LOG(ERROR) << operation_name << ": invalid tensor numel";
+        return std::nullopt;
+    }
+    if ((local_numel == 0 && parsed.data_bytes != 0) ||
+        (local_numel > 0 &&
+         parsed.data_bytes % static_cast<size_t>(local_numel) != 0)) {
+        LOG(ERROR) << operation_name << ": invalid tensor byte size";
+        return std::nullopt;
+    }
+
+    const size_t element_size =
+        local_numel == 0 ? 0
+                         : parsed.data_bytes / static_cast<size_t>(local_numel);
+    const auto [shard_start, shard_extent] =
+        calculate_shard_range(global_shape[split_dim], rank, shard_count);
+
+    std::vector<int64_t> shard_shape = global_shape;
+    shard_shape[split_dim] = shard_extent;
+    int64_t shard_numel = 1;
+    for (auto dim : shard_shape) {
+        shard_numel *= dim;
+    }
+    const size_t shard_bytes = static_cast<size_t>(shard_numel) * element_size;
+
+    TensorMetadata metadata =
+        BuildTensorMetadata(parsed.metadata.header.dtype, global_shape,
+                            shard_shape, TensorLayoutKind::SHARD);
+    metadata.header.data_bytes = shard_bytes;
+
+    RawTensorShardWritePlan plan;
+    plan.metadata = metadata;
+    if (shard_extent == 0 || element_size == 0) {
+        return plan;
+    }
+
+    int64_t elements_before = 1;
+    for (int i = 0; i < split_dim; ++i) {
+        elements_before *= global_shape[i];
+    }
+    int64_t elements_after = 1;
+    for (size_t i = split_dim + 1; i < global_shape.size(); ++i) {
+        elements_after *= global_shape[i];
+    }
+
+    const size_t row_bytes = static_cast<size_t>(shard_extent) *
+                             static_cast<size_t>(elements_after) * element_size;
+    plan.data_ranges.reserve(static_cast<size_t>(elements_before));
+    for (int64_t slice_idx = 0; slice_idx < elements_before; ++slice_idx) {
+        const size_t src_offset =
+            parsed.data_offset +
+            static_cast<size_t>(slice_idx * global_shape[split_dim] +
+                                shard_start) *
+                static_cast<size_t>(elements_after) * element_size;
+        plan.data_ranges.emplace_back(src_offset, row_bytes);
+    }
+    return plan;
+}
diff --git a/mooncake-integration/store/store_py_parallel_read.h b/mooncake-integration/store/store_py_parallel_read.h
new file mode 100644
index 00000000..06207cdf
--- /dev/null
+++ b/mooncake-integration/store/store_py_parallel_read.h
@@ -0,0 +1,1764 @@
+std::optional<ParsedTensorMetadata> get_tensor_metadata(
+    const std::string &key,
+    std::shared_ptr<BufferHandle> *buffer_handle_out = nullptr) {
+    if (!is_client_initialized()) {
+        LOG(ERROR) << "Client not initialized";
+        return std::nullopt;
+    }
+
+    std::shared_ptr<BufferHandle> buffer_handle;
+    {
+        py::gil_scoped_release release_gil;
+        buffer_handle = store_->get_buffer(key);
+    }
+    if (!buffer_handle) {
+        return std::nullopt;
+    }
+
+    if (buffer_handle_out) {
+        *buffer_handle_out = buffer_handle;
+    }
+    return ParseTensorMetadata(static_cast<const char *>(buffer_handle->ptr()),
+                               buffer_handle->size());
+}
+
+std::optional<ParsedTensorMetadata> parse_tensor_metadata_from_prefix(
+    const void *buffer_ptr, const std::string &context,
+    const std::string &key) {
+    if (buffer_ptr == nullptr) {
+        LOG(ERROR) << context << ": invalid metadata prefix for key " << key;
+        return std::nullopt;
+    }
+
+    TensorMetadata metadata{};
+    std::memcpy(&metadata, buffer_ptr, sizeof(TensorMetadata));
+    auto parsed = ParseTensorMetadata(
+        static_cast<const char *>(buffer_ptr),
+        static_cast<size_t>(metadata.header.data_offset) +
+            static_cast<size_t>(metadata.header.data_bytes));
+    if (!parsed.has_value()) {
+        LOG(ERROR) << context << ": invalid tensor metadata for key " << key;
+    }
+    return parsed;
+}
+
+std::vector<CachedQueryResultResponse> batch_query_for_reuse(
+    const std::vector<std::string> &keys) {
+    if (auto real_client = get_real_client()) {
+        return real_client->batch_get_query_results(keys);
+    }
+
+    auto query_results = store_->batch_query(keys);
+    std::vector<CachedQueryResultResponse> cached_results;
+    cached_results.reserve(query_results.size());
+    auto now = std::chrono::steady_clock::now();
+    for (const auto &query_result : query_results) {
+        cached_results.push_back(
+            to_cached_query_result_response(query_result, now));
+    }
+    return cached_results;
+}
+
+std::optional<std::vector<ReconstructedShardSource>>
+load_reconstructed_shard_sources_batch(const std::vector<std::string> &keys,
+                                       const std::string &context) {
+    if (!is_client_initialized()) {
+        LOG(ERROR) << context << ": client is not initialized";
+        return std::nullopt;
+    }
+    if (keys.empty()) {
+        return std::vector<ReconstructedShardSource>{};
+    }
+
+    std::vector<CachedQueryResultResponse> cached_query_results;
+    {
+        py::gil_scoped_release release_gil;
+        cached_query_results = batch_query_for_reuse(keys);
+    }
+    if (cached_query_results.size() != keys.size()) {
+        LOG(ERROR) << context << ": BatchQuery result size mismatch";
+        return std::nullopt;
+    }
+
+    mooncake::PyClient::QueryResultCache query_result_cache;
+    auto now = std::chrono::steady_clock::now();
+    std::vector<size_t> metadata_key_indices;
+    metadata_key_indices.reserve(keys.size());
+    std::vector<std::optional<CachedQueryResultResponse>>
+        reusable_query_results(keys.size());
+
+    for (size_t i = 0; i < keys.size(); ++i) {
+        auto query_result =
+            from_cached_query_result_response(cached_query_results[i], now);
+        if (!query_result || query_result->IsLeaseExpired(now)) {
+            continue;
+        }
+
+        query_result_cache.emplace(keys[i], std::move(query_result));
+        reusable_query_results[i] = cached_query_results[i];
+        metadata_key_indices.push_back(i);
+    }
+
+    std::vector<std::optional<ParsedTensorMetadata>> prefix_metadata(
+        keys.size());
+    if (!metadata_key_indices.empty()) {
+        const size_t scratch_size =
+            std::max<size_t>(metadata_key_indices.size(), 1) *
+            sizeof(TensorMetadata);
+        auto scratch_buffer = std::make_unique<char[]>(scratch_size);
+        if (store_->register_buffer(scratch_buffer.get(), scratch_size) != 0) {
+            LOG(ERROR) << context
+                       << ": failed to register metadata scratch buffer";
+            return std::nullopt;
+        }
+
+        char *scratch_base = scratch_buffer.get();
+        std::vector<void *> metadata_buffers;
+        metadata_buffers.reserve(metadata_key_indices.size());
+        std::vector<std::vector<std::string>> metadata_all_keys(
+            metadata_key_indices.size());
+        std::vector<std::vector<std::vector<size_t>>> metadata_all_dst_offsets(
+            metadata_key_indices.size());
+        std::vector<std::vector<std::vector<size_t>>> metadata_all_src_offsets(
+            metadata_key_indices.size(), {{0}});
+        std::vector<std::vector<std::vector<size_t>>> metadata_all_sizes(
+            metadata_key_indices.size(), {{sizeof(TensorMetadata)}});
+        for (size_t i = 0; i < metadata_key_indices.size(); ++i) {
+            // Each metadata read is modeled as a separate request, but they all
+            // write into non-overlapping offsets in the same registered scratch
+            // buffer.
+            metadata_buffers.push_back(scratch_base);
+            metadata_all_keys[i] = {keys[metadata_key_indices[i]]};
+            metadata_all_dst_offsets[i] = {{i * sizeof(TensorMetadata)}};
+        }
+
+        std::vector<std::vector<std::vector<int64_t>>> metadata_results;
+        {
+            py::gil_scoped_release release_gil;
+            metadata_results = store_->get_into_ranges(
+                metadata_buffers, metadata_all_keys, metadata_all_dst_offsets,
+                metadata_all_src_offsets, metadata_all_sizes,
+                &query_result_cache);
+        }
+        for (size_t i = 0;
+             i < metadata_results.size() && i < metadata_key_indices.size();
+             ++i) {
+            const size_t key_index = metadata_key_indices[i];
+            if (metadata_results[i].size() == 1 &&
+                metadata_results[i][0].size() == 1 &&
+                metadata_results[i][0][0] ==
+                    static_cast<int64_t>(sizeof(TensorMetadata))) {
+                prefix_metadata[key_index] = parse_tensor_metadata_from_prefix(
+                    scratch_base + i * sizeof(TensorMetadata), context,
+                    keys[key_index]);
+            }
+        }
+        store_->unregister_buffer(scratch_buffer.get());
+    }
+
+    std::vector<ReconstructedShardSource> sources;
+    sources.reserve(keys.size());
+    for (size_t i = 0; i < keys.size(); ++i) {
+        auto parsed = std::move(prefix_metadata[i]);
+        auto reusable_query_result = reusable_query_results[i];
+        if (!parsed.has_value()) {
+            parsed = get_tensor_metadata(keys[i]);
+            reusable_query_result.reset();
+        }
+        if (!parsed.has_value()) {
+            LOG(ERROR) << context
+                       << ": missing reconstructed shard source for key "
+                       << keys[i];
+            return std::nullopt;
+        }
+        sources.push_back(
+            ReconstructedShardSource{keys[i], *parsed, reusable_query_result});
+    }
+    return sources;
+}
+
+std::optional<TensorIntoPlan> build_tensor_into_plan(
+    const std::string &read_key, uintptr_t buffer_ptr, size_t size,
+    const std::string &context,
+    const std::optional<ParsedTensorMetadata> &metadata = std::nullopt) {
+    std::optional<ParsedTensorMetadata> resolved_metadata = metadata;
+    if (!resolved_metadata.has_value()) {
+        resolved_metadata = get_tensor_metadata(read_key);
+    }
+    if (!resolved_metadata.has_value()) {
+        return std::nullopt;
+    }
+
+    const auto total_length =
+        resolved_metadata->data_offset + resolved_metadata->data_bytes;
+    if (total_length > size) {
+        LOG(ERROR) << context << ": buffer too small for key " << read_key;
+        return std::nullopt;
+    }
+
+    auto region = resolve_writable_buffer_region(buffer_ptr, size, context);
+    if (!region.has_value()) {
+        return std::nullopt;
+    }
+    if (region->offset + total_length > region->size) {
+        LOG(ERROR) << context
+                   << ": resolved destination range exceeds writable region";
+        return std::nullopt;
+    }
+
+    TensorIntoPlan plan;
+    plan.user_buffer_ptr = buffer_ptr;
+    plan.registered_buffer_ptr = reinterpret_cast<uintptr_t>(region->base);
+    plan.registered_buffer_size = region->size;
+    plan.total_length = total_length;
+    plan.fragments.push_back(TensorIntoFragment{
+        .read_key = read_key,
+        .dst_offset = region->offset,
+        .src_offset = 0,
+        .size = total_length,
+    });
+    return plan;
+}
+
+std::optional<size_t> extract_reconstruction_element_size(
+    const std::vector<ReconstructedShardSource> &sources,
+    const std::string &context) {
+    for (const auto &source : sources) {
+        auto element_size =
+            TensorDtypeElementSize(source.metadata.metadata.header.dtype);
+        if (!element_size.has_value()) {
+            LOG(ERROR) << context << ": invalid shard tensor dtype";
+            return std::nullopt;
+        }
+
+        const auto expected_data_bytes =
+            TensorMetadataExpectedDataBytes(source.metadata.metadata);
+        if (!expected_data_bytes.has_value() ||
+            source.metadata.data_bytes != *expected_data_bytes) {
+            LOG(ERROR) << context << ": invalid shard tensor byte size";
+            return std::nullopt;
+        }
+        return *element_size;
+    }
+    return size_t{0};
+}
+
+std::optional<std::pair<int64_t, int64_t>> get_source_shard_range(
+    const ReconstructedShardSource &source,
+    const std::vector<int64_t> &global_shape, int split_dim,
+    const std::string &context) {
+    const auto local_shape =
+        TensorShapeToVector(source.metadata.metadata.layout.local_shape,
+                            source.metadata.metadata.header.ndim);
+    if (local_shape.size() != global_shape.size()) {
+        LOG(ERROR) << context << ": invalid shard shape for key "
+                   << source.read_key;
+        return std::nullopt;
+    }
+    for (size_t dim = 0; dim < global_shape.size(); ++dim) {
+        if (static_cast<int>(dim) == split_dim) {
+            continue;
+        }
+        if (local_shape[dim] != global_shape[dim]) {
+            LOG(ERROR) << context << ": shard shape mismatch for key "
+                       << source.read_key;
+            return std::nullopt;
+        }
+    }
+
+    const LayoutAxis *tp_axis =
+        find_layout_axis(source.metadata.metadata, LayoutAxisKind::TP);
+    if (!is_shard_tensor_metadata(source.metadata.metadata) || !tp_axis) {
+        LOG(ERROR) << context << ": missing TP shard metadata for key "
+                   << source.read_key;
+        return std::nullopt;
+    }
+    if (tp_axis->split_dim != split_dim || tp_axis->shard_count <= 0 ||
+        tp_axis->shard_rank < 0 ||
+        tp_axis->shard_rank >= tp_axis->shard_count) {
+        LOG(ERROR) << context << ": invalid TP shard metadata for key "
+                   << source.read_key;
+        return std::nullopt;
+    }
+    if (!is_uniform_shardable_dim(global_shape[split_dim],
+                                  tp_axis->shard_count)) {
+        LOG(ERROR) << context << ": only uniform sharding is supported";
+        return std::nullopt;
+    }
+
+    const auto [shard_start, shard_extent] = calculate_shard_range(
+        global_shape[split_dim], tp_axis->shard_rank, tp_axis->shard_count);
+    if (local_shape[split_dim] != shard_extent) {
+        LOG(ERROR) << context << ": shard extent mismatch for key "
+                   << source.read_key;
+        return std::nullopt;
+    }
+    return std::pair<int64_t, int64_t>{shard_start, shard_extent};
+}
+
+std::optional<TensorIntoPlan> build_reconstructed_tensor_into_plan_from_sources(
+    uintptr_t buffer_ptr, size_t size,
+    std::vector<ReconstructedShardSource> sources,
+    const std::vector<int64_t> &global_shape, int split_dim,
+    const TensorMetadata &target_metadata, int64_t target_start,
+    int64_t target_extent, const std::string &context,
+    bool allow_empty_fragments = false) {
+    if (sources.empty()) {
+        LOG(ERROR) << context << ": missing reconstruction shard sources";
+        return std::nullopt;
+    }
+    if (split_dim < 0 || split_dim >= static_cast<int>(global_shape.size())) {
+        LOG(ERROR) << context << ": invalid split_dim";
+        return std::nullopt;
+    }
+    if (target_start < 0 || target_extent < 0 ||
+        target_start + target_extent > global_shape[split_dim]) {
+        LOG(ERROR) << context << ": invalid target shard range";
+        return std::nullopt;
+    }
+
+    auto region = resolve_writable_buffer_region(buffer_ptr, size, context);
+    if (!region.has_value()) {
+        return std::nullopt;
+    }
+
+    auto element_size = extract_reconstruction_element_size(sources, context);
+    if (!element_size.has_value()) {
+        return std::nullopt;
+    }
+
+    size_t target_tensor_numel = 1;
+    for (size_t dim = 0; dim < global_shape.size(); ++dim) {
+        const int64_t dim_extent = static_cast<int>(dim) == split_dim
+                                       ? target_extent
+                                       : global_shape[dim];
+        target_tensor_numel *= static_cast<size_t>(dim_extent);
+    }
+    const size_t target_tensor_bytes = target_tensor_numel * *element_size;
+    const size_t total_length = sizeof(TensorMetadata) + target_tensor_bytes;
+
+    if (total_length > size || region->offset + total_length > region->size) {
+        LOG(ERROR) << context << ": buffer too small for reconstructed tensor";
+        return std::nullopt;
+    }
+
+    TensorIntoPlan plan;
+    plan.user_buffer_ptr = buffer_ptr;
+    plan.registered_buffer_ptr = reinterpret_cast<uintptr_t>(region->base);
+    plan.registered_buffer_size = region->size;
+    plan.total_length = total_length;
+    plan.materialized_metadata = target_metadata;
+    plan.materialized_metadata->header.data_bytes = target_tensor_bytes;
+
+    int64_t elements_before = 1;
+    for (int i = 0; i < split_dim; ++i) {
+        elements_before *= global_shape[i];
+    }
+    int64_t elements_after = 1;
+    for (size_t i = split_dim + 1; i < global_shape.size(); ++i) {
+        elements_after *= global_shape[i];
+    }
+
+    std::vector<bool> covered(
+        static_cast<size_t>(target_extent > 0 ? target_extent : 0), false);
+    for (auto &source : sources) {
+        auto source_range =
+            get_source_shard_range(source, global_shape, split_dim, context);
+        if (!source_range.has_value()) {
+            return std::nullopt;
+        }
+        const auto [source_start, source_extent] = *source_range;
+        const int64_t overlap_start = std::max(source_start, target_start);
+        const int64_t overlap_end = std::min(source_start + source_extent,
+                                             target_start + target_extent);
+        if (overlap_end <= overlap_start) {
+            continue;
+        }
+        const int64_t overlap_extent = overlap_end - overlap_start;
+        const int64_t src_inner_offset = overlap_start - source_start;
+        const int64_t dst_inner_offset = overlap_start - target_start;
+        const size_t row_bytes = static_cast<size_t>(overlap_extent) *
+                                 static_cast<size_t>(elements_after) *
+                                 *element_size;
+        for (int64_t idx = dst_inner_offset;
+             idx < dst_inner_offset + overlap_extent; ++idx) {
+            covered[static_cast<size_t>(idx)] = true;
+        }
+
+        for (int64_t slice_idx = 0; slice_idx < elements_before; ++slice_idx) {
+            const size_t dst_offset =
+                region->offset + sizeof(TensorMetadata) +
+                static_cast<size_t>(slice_idx * target_extent +
+                                    dst_inner_offset) *
+                    static_cast<size_t>(elements_after) * *element_size;
+            const size_t src_offset =
+                source.metadata.data_offset +
+                static_cast<size_t>(slice_idx * source_extent +
+                                    src_inner_offset) *
+                    static_cast<size_t>(elements_after) * *element_size;
+            plan.fragments.push_back(TensorIntoFragment{
+                .read_key = source.read_key,
+                .dst_offset = dst_offset,
+                .src_offset = src_offset,
+                .size = row_bytes,
+            });
+        }
+    }
+
+    for (bool is_covered : covered) {
+        if (!is_covered) {
+            LOG(ERROR)
+                << context
+                << ": shard extents do not cover reconstructed dimension";
+            return std::nullopt;
+        }
+    }
+    if (plan.fragments.empty() && target_tensor_bytes != 0) {
+        LOG(ERROR) << context << ": no fragments planned for reconstruction";
+        return std::nullopt;
+    }
+
+    plan.query_results.reserve(sources.size());
+    for (auto &source : sources) {
+        if (source.cached_query_result.has_value()) {
+            plan.query_results.push_back(PlannedQueryResult{
+                source.read_key, std::move(source.cached_query_result)});
+        }
+    }
+    return plan;
+}
+
+std::optional<TensorIntoPlan> build_full_tensor_into_plan_from_sources(
+    uintptr_t buffer_ptr, size_t size,
+    std::vector<ReconstructedShardSource> sources,
+    const std::vector<int64_t> &global_shape, int split_dim, int32_t dtype,
+    const std::string &context, bool allow_empty_fragments = false) {
+    TensorMetadata full_metadata = BuildTensorMetadata(
+        dtype, global_shape, global_shape, TensorLayoutKind::FULL);
+    return build_reconstructed_tensor_into_plan_from_sources(
+        buffer_ptr, size, std::move(sources), global_shape, split_dim,
+        full_metadata, 0, global_shape[split_dim], context,
+        allow_empty_fragments);
+}
+
+std::string resolve_tp_read_key(const std::string &key, int tp_rank,
+                                int tp_size) const {
+    if (tp_size <= 1) return key;
+    return get_tp_key_name(key, tp_rank);
+}
+
+pybind11::object get_tensor_with_writer_shard_full(const std::string &key,
+                                                   const std::string &context) {
+    if (!is_client_initialized()) {
+        LOG(ERROR) << "Client is not initialized";
+        return pybind11::none();
+    }
+    if (use_dummy_client_) {
+        LOG(ERROR) << context << ": dummy client is not supported";
+        return pybind11::none();
+    }
+
+    auto reconstruction = load_writer_shard_reconstruction(key, context);
+    if (!reconstruction.has_value()) {
+        return get_tensor(key);
+    }
+
+    size_t total_tensor_numel = 1;
+    for (auto dim : reconstruction->global_shape) {
+        total_tensor_numel *= static_cast<size_t>(dim);
+    }
+
+    auto element_size =
+        extract_reconstruction_element_size(reconstruction->sources, context);
+    if (!element_size.has_value()) {
+        return py::none();
+    }
+
+    const size_t total_length =
+        sizeof(TensorMetadata) + total_tensor_numel * *element_size;
+    char *owned_buffer = new char[total_length];
+    if (store_->register_buffer(owned_buffer, total_length) != 0) {
+        LOG(ERROR) << context << ": failed to register reconstruction buffer";
+        delete[] owned_buffer;
+        return py::none();
+    }
+
+    auto plan = build_full_tensor_into_plan_from_sources(
+        reinterpret_cast<uintptr_t>(owned_buffer), total_length,
+        reconstruction->sources, reconstruction->global_shape,
+        reconstruction->split_dim, reconstruction->dtype, context,
+        reconstruction->allow_empty_fragments);
+    if (!plan.has_value()) {
+        store_->unregister_buffer(owned_buffer);
+        delete[] owned_buffer;
+        return py::none();
+    }
+
+    std::vector<TensorIntoPlan> plans;
+    plans.push_back(std::move(*plan));
+    auto success = execute_tensor_into_plan_transfers(plans);
+    if (success.empty() || !success[0]) {
+        store_->unregister_buffer(owned_buffer);
+        delete[] owned_buffer;
+        return py::none();
+    }
+
+    return buffer_to_tensor(
+        new BufferHandle(owned_buffer, total_length,
+                         [this, owned_buffer]() {
+                             store_->unregister_buffer(owned_buffer);
+                             delete[] owned_buffer;
+                         }),
+        nullptr, 0);
+}
+
+pybind11::object get_tensor_with_tp_full(
+    const std::string &key, int tp_rank, int tp_size, int split_dim,
+    const std::string &context,
+    const std::optional<TensorParallelismSpec> &parallelism = std::nullopt) {
+    if (!is_client_initialized()) {
+        LOG(ERROR) << "Client is not initialized";
+        return pybind11::none();
+    }
+    if (use_dummy_client_) {
+        LOG(ERROR) << context << ": dummy client is not supported";
+        return pybind11::none();
+    }
+    ParallelAxisSpec axis{
+        .kind = "tp",
+        .rank = tp_rank,
+        .size = tp_size,
+        .split_dim = split_dim,
+        .expert_id = std::nullopt,
+        .stage_id = std::nullopt,
+    };
+
+    auto reconstruction =
+        parallelism.has_value()
+            ? load_parallelism_full_reconstruction_sources(key, *parallelism,
+                                                           context)
+            : load_tp_full_reconstruction_sources(key, axis, context);
+    if (!reconstruction.has_value()) {
+        return pybind11::none();
+    }
+
+    const size_t total_tensor_numel =
+        std::accumulate(reconstruction->global_shape.begin(),
+                        reconstruction->global_shape.end(),
+                        static_cast<size_t>(1), std::multiplies<size_t>());
+    auto element_size =
+        extract_reconstruction_element_size(reconstruction->sources, context);
+    if (!element_size.has_value()) {
+        return pybind11::none();
+    }
+    const size_t total_length =
+        sizeof(TensorMetadata) + total_tensor_numel * *element_size;
+
+    char *owned_buffer = new char[total_length];
+    if (store_->register_buffer(owned_buffer, total_length) != 0) {
+        LOG(ERROR) << context << ": failed to register reconstruction buffer";
+        delete[] owned_buffer;
+        return pybind11::none();
+    }
+
+    auto plan = build_full_tensor_into_plan_from_sources(
+        reinterpret_cast<uintptr_t>(owned_buffer), total_length,
+        reconstruction->sources, reconstruction->global_shape,
+        reconstruction->split_dim, reconstruction->dtype, context,
+        reconstruction->allow_empty_fragments);
+    if (!plan.has_value()) {
+        store_->unregister_buffer(owned_buffer);
+        delete[] owned_buffer;
+        return pybind11::none();
+    }
+
+    std::vector<TensorIntoPlan> plans;
+    plans.push_back(std::move(*plan));
+    auto success = execute_tensor_into_plan_transfers(plans);
+    if (success.empty() || !success[0]) {
+        store_->unregister_buffer(owned_buffer);
+        delete[] owned_buffer;
+        return py::none();
+    }
+
+    return buffer_to_tensor(
+        new BufferHandle(owned_buffer, total_length,
+                         [this, owned_buffer]() {
+                             store_->unregister_buffer(owned_buffer);
+                             delete[] owned_buffer;
+                         }),
+        nullptr, 0);
+}
+
+std::optional<TensorIntoPlan> build_tp_full_tensor_into_plan(
+    const std::string &key, uintptr_t buffer_ptr, size_t size,
+    const ParallelAxisSpec &axis, const std::string &context,
+    const std::optional<TensorParallelismSpec> &parallelism = std::nullopt) {
+    auto reconstruction =
+        load_tp_full_reconstruction_sources(key, axis, context, parallelism);
+    if (!reconstruction.has_value()) {
+        return std::nullopt;
+    }
+    return build_full_tensor_into_plan_from_sources(
+        buffer_ptr, size, reconstruction->sources, reconstruction->global_shape,
+        reconstruction->split_dim, reconstruction->dtype, context,
+        reconstruction->allow_empty_fragments);
+}
+
+std::optional<FullTensorReconstructionSources>
+load_tp_full_reconstruction_sources(
+    const std::string &key, const ParallelAxisSpec &axis,
+    const std::string &context,
+    const std::optional<TensorParallelismSpec> &parallelism = std::nullopt) {
+    if (axis.size <= 0) {
+        LOG(ERROR) << context << ": tp_size must be positive";
+        return std::nullopt;
+    }
+
+    std::vector<std::string> read_keys;
+    read_keys.reserve(axis.size);
+    for (int shard_rank = 0; shard_rank < axis.size; ++shard_rank) {
+        if (parallelism.has_value()) {
+            auto shard_parallelism = *parallelism;
+            auto tp_axis_index = find_tp_axis_index(shard_parallelism.axes);
+            if (!tp_axis_index.has_value()) {
+                LOG(ERROR) << context
+                           << ": missing TP axis in full reconstruction";
+                return std::nullopt;
+            }
+            shard_parallelism.axes[*tp_axis_index].rank = shard_rank;
+            read_keys.push_back(
+                get_parallelism_key_name(key, shard_parallelism));
+        } else {
+            read_keys.push_back(
+                resolve_tp_read_key(key, shard_rank, axis.size));
+        }
+    }
+
+    auto ordered_sources =
+        load_reconstructed_shard_sources_batch(read_keys, context);
+    if (!ordered_sources.has_value()) {
+        return std::nullopt;
+    }
+
+    FullTensorReconstructionSources reconstruction;
+    reconstruction.sources.reserve(axis.size);
+    for (int shard_rank = 0; shard_rank < axis.size; ++shard_rank) {
+        auto &source = (*ordered_sources)[shard_rank];
+        const auto &metadata = source.metadata;
+        const LayoutAxis *tp_axis =
+            find_layout_axis(metadata.metadata, LayoutAxisKind::TP);
+        if (!is_shard_tensor_metadata(metadata.metadata) || !tp_axis ||
+            tp_axis->shard_rank != shard_rank ||
+            tp_axis->shard_count != axis.size) {
+            LOG(ERROR) << context << ": TP metadata mismatch for key "
+                       << source.read_key;
+            return std::nullopt;
+        }
+        if (parallelism.has_value()) {
+            auto stored_parallelism =
+                resolve_tp_compatible_parallelism_from_metadata(
+                    *parallelism, metadata.metadata, context);
+            if (!stored_parallelism.has_value()) {
+                return std::nullopt;
+            }
+        }
+        reconstruction.sources.push_back(std::move(source));
+    }
+
+    reconstruction.global_shape = TensorShapeToVector(
+        reconstruction.sources.front().metadata.metadata.layout.global_shape,
+        reconstruction.sources.front().metadata.metadata.header.ndim);
+    const LayoutAxis *stored_tp_axis = find_layout_axis(
+        reconstruction.sources.front().metadata.metadata, LayoutAxisKind::TP);
+    if (!stored_tp_axis) {
+        LOG(ERROR) << context << ": missing TP axis metadata";
+        return std::nullopt;
+    }
+    reconstruction.split_dim = stored_tp_axis->split_dim;
+    if (!is_uniform_shardable_dim(
+            reconstruction.global_shape[reconstruction.split_dim],
+            stored_tp_axis->shard_count)) {
+        LOG(ERROR) << context << ": only uniform sharding is supported";
+        return std::nullopt;
+    }
+    if (axis.split_dim.has_value() &&
+        axis.split_dim.value() != reconstruction.split_dim) {
+        LOG(ERROR) << context << ": split_dim mismatch";
+        return std::nullopt;
+    }
+    reconstruction.dtype =
+        reconstruction.sources.front().metadata.metadata.header.dtype;
+    return reconstruction;
+}
+
+std::optional<FullTensorReconstructionSources> load_writer_shard_reconstruction(
+    const std::string &key, const std::string &context) {
+    if (!is_client_initialized()) {
+        LOG(ERROR) << context << ": client is not initialized";
+        return std::nullopt;
+    }
+
+    std::shared_ptr<BufferHandle> manifest_handle;
+    {
+        py::gil_scoped_release release_gil;
+        manifest_handle = store_->get_buffer(get_writer_manifest_key_name(key));
+    }
+    auto parsed_manifest = parse_writer_shard_manifest(manifest_handle.get());
+    if (!parsed_manifest.has_value()) {
+        return std::nullopt;
+    }
+
+    const auto &manifest = *parsed_manifest;
+    const auto &global_shape = manifest.global_shape;
+    const int split_dim = manifest.manifest.header.split_dim;
+    const int shard_count = manifest.manifest.header.shard_count;
+    if (split_dim < 0 || split_dim >= static_cast<int>(global_shape.size())) {
+        LOG(ERROR) << context << ": invalid writer split_dim";
+        return std::nullopt;
+    }
+
+    std::vector<std::string> shard_keys;
+    shard_keys.reserve(shard_count);
+    for (int shard_rank = 0; shard_rank < shard_count; ++shard_rank) {
+        WriterPartitionSpec writer{
+            .rank = shard_rank,
+            .size = shard_count,
+            .split_dim = split_dim,
+        };
+        shard_keys.push_back(get_writer_shard_key_name(key, writer));
+    }
+
+    auto ordered_sources =
+        load_reconstructed_shard_sources_batch(shard_keys, context);
+    if (!ordered_sources.has_value()) {
+        return std::nullopt;
+    }
+
+    FullTensorReconstructionSources reconstruction;
+    reconstruction.sources.reserve(shard_count);
+    for (int shard_rank = 0; shard_rank < shard_count; ++shard_rank) {
+        auto &source = (*ordered_sources)[shard_rank];
+        auto writer_parallelism = writer_partition_parallelism_from_metadata(
+            source.metadata.metadata);
+        if (!writer_parallelism.has_value() ||
+            writer_parallelism->axes[0].rank != shard_rank ||
+            writer_parallelism->axes[0].size != shard_count ||
+            writer_parallelism->axes[0].split_dim != split_dim) {
+            LOG(ERROR) << context << ": writer shard metadata mismatch for key "
+                       << source.read_key;
+            return std::nullopt;
+        }
+        reconstruction.sources.push_back(std::move(source));
+    }
+    reconstruction.global_shape = manifest.global_shape;
+    reconstruction.split_dim = manifest.manifest.header.split_dim;
+    reconstruction.dtype = manifest.manifest.header.dtype;
+    reconstruction.allow_empty_fragments = true;
+    return reconstruction;
+}
+
+std::optional<TensorIntoPlan> build_writer_shard_full_tensor_into_plan(
+    const std::string &key, uintptr_t buffer_ptr, size_t size,
+    const std::string &context) {
+    auto reconstruction = load_writer_shard_reconstruction(key, context);
+    if (!reconstruction.has_value()) {
+        return build_tensor_into_plan(key, buffer_ptr, size, context);
+    }
+
+    return build_full_tensor_into_plan_from_sources(
+        buffer_ptr, size, reconstruction->sources, reconstruction->global_shape,
+        reconstruction->split_dim, reconstruction->dtype, context,
+        reconstruction->allow_empty_fragments);
+}
+
+std::vector<std::string> build_parallelism_shard_read_keys(
+    const std::string &key, const TensorParallelismSpec &canonical_parallelism,
+    size_t tp_axis_index, int shard_count,
+    std::optional<int> split_dim = std::nullopt) {
+    std::vector<std::string> shard_keys;
+    shard_keys.reserve(shard_count);
+    for (int shard_rank = 0; shard_rank < shard_count; ++shard_rank) {
+        auto shard_parallelism = canonical_parallelism;
+        shard_parallelism.axes[tp_axis_index].rank = shard_rank;
+        shard_parallelism.axes[tp_axis_index].size = shard_count;
+        if (split_dim.has_value()) {
+            shard_parallelism.axes[tp_axis_index].split_dim = *split_dim;
+        }
+        shard_keys.push_back(get_parallelism_key_name(key, shard_parallelism));
+    }
+    return shard_keys;
+}
+
+std::optional<FullTensorReconstructionSources>
+load_parallelism_manifest_reconstruction(
+    const std::string &key, const TensorParallelismSpec &parallelism,
+    const std::string &context) {
+    const ParallelAxisSpec *request_tp_axis =
+        find_axis_spec_by_kind(parallelism, LayoutAxisKind::TP);
+    if (!request_tp_axis) {
+        LOG(ERROR) << context << ": reconstruction requires a TP axis";
+        return std::nullopt;
+    }
+
+    auto parsed_manifest = load_parallelism_manifest(key, context);
+    if (!parsed_manifest.has_value()) {
+        return std::nullopt;
+    }
+
+    const auto &manifest = *parsed_manifest;
+    const auto &global_shape = manifest.global_shape;
+    const int split_dim = manifest.manifest.header.split_dim;
+    const int shard_count = manifest.manifest.header.shard_count;
+    if (split_dim < 0 || split_dim >= static_cast<int>(global_shape.size())) {
+        LOG(ERROR) << context << ": invalid parallelism split_dim";
+        return std::nullopt;
+    }
+    if (request_tp_axis->split_dim.has_value() &&
+        request_tp_axis->split_dim.value() != split_dim) {
+        LOG(ERROR) << context << ": split_dim mismatch";
+        return std::nullopt;
+    }
+    if (!is_uniform_shardable_dim(global_shape[split_dim],
+                                  request_tp_axis->size) ||
+        !is_uniform_shardable_dim(global_shape[split_dim], shard_count)) {
+        LOG(ERROR) << context << ": only uniform sharding is supported";
+        return std::nullopt;
+    }
+    auto canonical_parallelism = canonicalize_parallelism_spec(parallelism);
+    if (!canonical_parallelism.has_value()) {
+        LOG(ERROR) << context << ": failed to canonicalize parallelism";
+        return std::nullopt;
+    }
+    auto tp_axis_index = find_tp_axis_index(canonical_parallelism->axes);
+    if (!tp_axis_index.has_value()) {
+        LOG(ERROR) << context << ": missing TP axis in request";
+        return std::nullopt;
+    }
+
+    auto shard_keys = build_parallelism_shard_read_keys(
+        key, *canonical_parallelism, *tp_axis_index, shard_count);
+
+    auto ordered_sources =
+        load_reconstructed_shard_sources_batch(shard_keys, context);
+    if (!ordered_sources.has_value()) {
+        return std::nullopt;
+    }
+
+    FullTensorReconstructionSources reconstruction;
+    reconstruction.global_shape = global_shape;
+    reconstruction.split_dim = split_dim;
+    reconstruction.dtype = manifest.manifest.header.dtype;
+    reconstruction.sources.reserve(shard_count);
+    for (int shard_rank = 0; shard_rank < shard_count; ++shard_rank) {
+        auto shard_parallelism = *canonical_parallelism;
+        shard_parallelism.axes[*tp_axis_index].rank = shard_rank;
+        shard_parallelism.axes[*tp_axis_index].size = shard_count;
+        auto &source = (*ordered_sources)[shard_rank];
+        auto stored_parallelism =
+            resolve_tp_compatible_parallelism_from_metadata(
+                shard_parallelism, source.metadata.metadata, context);
+        if (!stored_parallelism.has_value() ||
+            !parallelism_specs_equal_by_kind(
+                shard_parallelism, *stored_parallelism,
+                true /* allow_tp_rank_mismatch */)) {
+            LOG(ERROR) << context << ": shard metadata mismatch for TP rank "
+                       << shard_rank;
+            return std::nullopt;
+        }
+        reconstruction.sources.push_back(std::move(source));
+    }
+    return reconstruction;
+}
+
+std::optional<FullTensorReconstructionSources>
+load_parallelism_full_reconstruction_sources(
+    const std::string &key, const TensorParallelismSpec &parallelism,
+    const std::string &context) {
+    if (auto manifest_reconstruction =
+            load_parallelism_manifest_reconstruction(key, parallelism, context);
+        manifest_reconstruction.has_value()) {
+        return manifest_reconstruction;
+    }
+
+    const ParallelAxisSpec *request_tp_axis =
+        find_axis_spec_by_kind(parallelism, LayoutAxisKind::TP);
+    if (!request_tp_axis) {
+        LOG(ERROR) << context << ": full reconstruction requires a TP axis";
+        return std::nullopt;
+    }
+
+    auto canonical_parallelism = canonicalize_parallelism_spec(parallelism);
+    if (!canonical_parallelism.has_value()) {
+        LOG(ERROR) << context << ": failed to canonicalize parallelism";
+        return std::nullopt;
+    }
+
+    auto tp_axis_index = find_tp_axis_index(canonical_parallelism->axes);
+    if (!tp_axis_index.has_value()) {
+        LOG(ERROR) << context << ": missing TP axis in request";
+        return std::nullopt;
+    }
+
+    FullTensorReconstructionSources reconstruction;
+    if (uses_legacy_tp_storage_key(*canonical_parallelism)) {
+        const std::string first_key =
+            resolve_tp_read_key(key, 0, request_tp_axis->size);
+        auto first_metadata = get_tensor_metadata(first_key);
+        if (!first_metadata.has_value()) {
+            LOG(ERROR) << context
+                       << ": no shard matched requested layout for TP rank 0";
+            return std::nullopt;
+        }
+        const LayoutAxis *stored_tp_axis =
+            find_layout_axis(first_metadata->metadata, LayoutAxisKind::TP);
+        if (!is_shard_tensor_metadata(first_metadata->metadata) ||
+            !stored_tp_axis) {
+            LOG(ERROR) << context << ": missing TP axis metadata";
+            return std::nullopt;
+        }
+        reconstruction.split_dim = stored_tp_axis->split_dim;
+        if (request_tp_axis->split_dim.has_value() &&
+            request_tp_axis->split_dim.value() != reconstruction.split_dim) {
+            LOG(ERROR) << context << ": split_dim mismatch";
+            return std::nullopt;
+        }
+        reconstruction.global_shape =
+            TensorShapeToVector(first_metadata->metadata.layout.global_shape,
+                                first_metadata->metadata.header.ndim);
+        if (!is_uniform_shardable_dim(
+                reconstruction.global_shape[reconstruction.split_dim],
+                request_tp_axis->size) ||
+            !is_uniform_shardable_dim(
+                reconstruction.global_shape[reconstruction.split_dim],
+                stored_tp_axis->shard_count)) {
+            LOG(ERROR) << context << ": only uniform sharding is supported";
+            return std::nullopt;
+        }
+        reconstruction.dtype = first_metadata->metadata.header.dtype;
+        std::vector<std::string> shard_keys;
+        shard_keys.reserve(stored_tp_axis->shard_count);
+        for (int shard_rank = 0; shard_rank < stored_tp_axis->shard_count;
+             ++shard_rank) {
+            shard_keys.push_back(resolve_tp_read_key(
+                key, shard_rank, stored_tp_axis->shard_count));
+        }
+        auto ordered_sources =
+            load_reconstructed_shard_sources_batch(shard_keys, context);
+        if (!ordered_sources.has_value()) {
+            return std::nullopt;
+        }
+        reconstruction.sources.reserve(stored_tp_axis->shard_count);
+        for (int shard_rank = 0; shard_rank < stored_tp_axis->shard_count;
+             ++shard_rank) {
+            auto &source = (*ordered_sources)[shard_rank];
+            const LayoutAxis *source_tp_axis =
+                find_layout_axis(source.metadata.metadata, LayoutAxisKind::TP);
+            if (!is_shard_tensor_metadata(source.metadata.metadata) ||
+                !source_tp_axis || source_tp_axis->shard_rank != shard_rank ||
+                source_tp_axis->shard_count != stored_tp_axis->shard_count ||
+                source_tp_axis->split_dim != reconstruction.split_dim) {
+                LOG(ERROR) << context
+                           << ": shard metadata mismatch for TP rank "
+                           << shard_rank;
+                return std::nullopt;
+            }
+            reconstruction.sources.push_back(std::move(source));
+        }
+        return reconstruction;
+    }
+
+    auto shard_keys = build_parallelism_shard_read_keys(
+        key, *canonical_parallelism, *tp_axis_index, request_tp_axis->size);
+    auto ordered_sources =
+        load_reconstructed_shard_sources_batch(shard_keys, context);
+    if (!ordered_sources.has_value()) {
+        return std::nullopt;
+    }
+
+    reconstruction.sources.reserve(request_tp_axis->size);
+    for (int shard_rank = 0; shard_rank < request_tp_axis->size; ++shard_rank) {
+        auto shard_parallelism = *canonical_parallelism;
+        shard_parallelism.axes[*tp_axis_index].rank = shard_rank;
+        auto &source = (*ordered_sources)[shard_rank];
+
+        auto stored_parallelism =
+            resolve_tp_compatible_parallelism_from_metadata(
+                shard_parallelism, source.metadata.metadata, context);
+        if (!stored_parallelism.has_value() ||
+            !parallelism_specs_equal_by_kind(
+                *canonical_parallelism, *stored_parallelism,
+                true /* allow_tp_rank_mismatch */)) {
+            LOG(ERROR) << context << ": shard metadata mismatch for TP rank "
+                       << shard_rank;
+            return std::nullopt;
+        }
+        reconstruction.sources.push_back(std::move(source));
+    }
+
+    const LayoutAxis *stored_tp_axis = find_layout_axis(
+        reconstruction.sources.front().metadata.metadata, LayoutAxisKind::TP);
+    if (!stored_tp_axis) {
+        LOG(ERROR) << context << ": missing TP axis metadata";
+        return std::nullopt;
+    }
+    reconstruction.split_dim = stored_tp_axis->split_dim;
+    if (request_tp_axis->split_dim.has_value() &&
+        request_tp_axis->split_dim.value() != reconstruction.split_dim) {
+        LOG(ERROR) << context << ": split_dim mismatch";
+        return std::nullopt;
+    }
+
+    reconstruction.global_shape = TensorShapeToVector(
+        reconstruction.sources.front().metadata.metadata.layout.global_shape,
+        reconstruction.sources.front().metadata.metadata.header.ndim);
+    if (!is_uniform_shardable_dim(
+            reconstruction.global_shape[reconstruction.split_dim],
+            request_tp_axis->size)) {
+        LOG(ERROR) << context << ": only uniform sharding is supported";
+        return std::nullopt;
+    }
+    reconstruction.dtype =
+        reconstruction.sources.front().metadata.metadata.header.dtype;
+    return reconstruction;
+}
+
+bool validate_regular_full_formula_sources(
+    const std::vector<ReconstructedShardSource> &sources,
+    const std::vector<int64_t> &global_shape, int split_dim, int32_t dtype,
+    int shard_count, const std::string &context) {
+    if (sources.size() != static_cast<size_t>(shard_count)) {
+        return false;
+    }
+    for (int shard_rank = 0; shard_rank < shard_count; ++shard_rank) {
+        const auto &source = sources[shard_rank];
+        if (source.metadata.metadata.header.dtype != dtype) {
+            LOG(ERROR) << context << ": shard dtype mismatch for key "
+                       << source.read_key;
+            return false;
+        }
+        const LayoutAxis *tp_axis =
+            find_layout_axis(source.metadata.metadata, LayoutAxisKind::TP);
+        if (!tp_axis || tp_axis->shard_rank != shard_rank ||
+            tp_axis->shard_count != shard_count ||
+            tp_axis->split_dim != split_dim) {
+            LOG(ERROR) << context << ": shard TP metadata mismatch for key "
+                       << source.read_key;
+            return false;
+        }
+        if (!get_source_shard_range(source, global_shape, split_dim, context)
+                 .has_value()) {
+            return false;
+        }
+    }
+    return true;
+}
+
+std::optional<TensorIntoPlan> build_parallelism_full_tensor_into_formula_plan(
+    const std::string &key, uintptr_t buffer_ptr, size_t size,
+    const TensorParallelismSpec &parallelism, const std::string &context) {
+    const ParallelAxisSpec *request_tp_axis =
+        find_axis_spec_by_kind(parallelism, LayoutAxisKind::TP);
+    if (!request_tp_axis) {
+        return std::nullopt;
+    }
+
+    auto manifest = load_parallelism_manifest(key, context);
+    if (!manifest.has_value()) {
+        return std::nullopt;
+    }
+    const auto &global_shape = manifest->global_shape;
+    const int split_dim = manifest->manifest.header.split_dim;
+    const int shard_count = manifest->manifest.header.shard_count;
+    if (split_dim < 0 || split_dim >= static_cast<int>(global_shape.size()) ||
+        shard_count <= 0) {
+        return std::nullopt;
+    }
+    if (request_tp_axis->split_dim.has_value() &&
+        request_tp_axis->split_dim.value() != split_dim) {
+        return std::nullopt;
+    }
+    if (!is_uniform_shardable_dim(global_shape[split_dim],
+                                  request_tp_axis->size) ||
+        !is_uniform_shardable_dim(global_shape[split_dim], shard_count)) {
+        return std::nullopt;
+    }
+    auto element_size = TensorDtypeElementSize(manifest->manifest.header.dtype);
+    if (!element_size.has_value()) {
+        return std::nullopt;
+    }
+
+    auto canonical_parallelism = canonicalize_parallelism_spec(parallelism);
+    if (!canonical_parallelism.has_value()) {
+        return std::nullopt;
+    }
+    auto tp_axis_index = find_tp_axis_index(canonical_parallelism->axes);
+    if (!tp_axis_index.has_value()) {
+        return std::nullopt;
+    }
+
+    size_t tensor_numel = 1;
+    for (auto dim : global_shape) {
+        if (dim < 0) {
+            return std::nullopt;
+        }
+        tensor_numel *= static_cast<size_t>(dim);
+    }
+    const size_t tensor_bytes = tensor_numel * *element_size;
+    const size_t total_length = sizeof(TensorMetadata) + tensor_bytes;
+
+    auto region = resolve_writable_buffer_region(buffer_ptr, size, context);
+    if (!region.has_value()) {
+        return std::nullopt;
+    }
+    if (total_length > size || region->offset + total_length > region->size) {
+        LOG(ERROR) << context << ": buffer too small for reconstructed tensor";
+        return std::nullopt;
+    }
+
+    TensorIntoPlan plan;
+    plan.user_buffer_ptr = buffer_ptr;
+    plan.registered_buffer_ptr = reinterpret_cast<uintptr_t>(region->base);
+    plan.registered_buffer_size = region->size;
+    plan.total_length = total_length;
+    plan.materialized_metadata =
+        BuildTensorMetadata(manifest->manifest.header.dtype, global_shape,
+                            global_shape, TensorLayoutKind::FULL);
+    plan.materialized_metadata->header.data_bytes = tensor_bytes;
+
+    TensorIntoRegularFullFormulaPlan formula;
+    formula.global_shape = global_shape;
+    formula.split_dim = split_dim;
+    formula.element_size = *element_size;
+    formula.data_offset = sizeof(TensorMetadata);
+    formula.read_keys = build_parallelism_shard_read_keys(
+        key, *canonical_parallelism, *tp_axis_index, shard_count, split_dim);
+
+    auto sources =
+        load_reconstructed_shard_sources_batch(formula.read_keys, context);
+    if (!sources.has_value() ||
+        !validate_regular_full_formula_sources(
+            *sources, global_shape, split_dim, manifest->manifest.header.dtype,
+            shard_count, context)) {
+        return std::nullopt;
+    }
+    plan.query_results.reserve(sources->size());
+    for (auto &source : *sources) {
+        if (!source.cached_query_result.has_value()) {
+            return std::nullopt;
+        }
+        plan.query_results.push_back(PlannedQueryResult{
+            source.read_key, std::move(*source.cached_query_result)});
+    }
+
+    plan.regular_full_formula = std::move(formula);
+    return plan;
+}
+
+std::optional<TensorIntoPlan> build_parallelism_full_tensor_into_plan(
+    const std::string &key, uintptr_t buffer_ptr, size_t size,
+    const TensorParallelismSpec &parallelism, const std::string &context,
+    bool allow_formula_plan = true) {
+    if (allow_formula_plan) {
+        if (auto formula_plan = build_parallelism_full_tensor_into_formula_plan(
+                key, buffer_ptr, size, parallelism, context);
+            formula_plan.has_value()) {
+            return formula_plan;
+        }
+    }
+
+    auto reconstruction =
+        load_parallelism_full_reconstruction_sources(key, parallelism, context);
+    if (!reconstruction.has_value()) {
+        return std::nullopt;
+    }
+    return build_full_tensor_into_plan_from_sources(
+        buffer_ptr, size, reconstruction->sources, reconstruction->global_shape,
+        reconstruction->split_dim, reconstruction->dtype, context,
+        reconstruction->allow_empty_fragments);
+}
+
+std::optional<TensorIntoPlan> build_parallelism_shard_tensor_into_plan(
+    const std::string &key, uintptr_t buffer_ptr, size_t size,
+    const TensorParallelismSpec &parallelism, const std::string &context) {
+    auto reconstruction =
+        load_parallelism_full_reconstruction_sources(key, parallelism, context);
+    if (!reconstruction.has_value()) {
+        return std::nullopt;
+    }
+
+    const ParallelAxisSpec *tp_axis =
+        find_axis_spec_by_kind(parallelism, LayoutAxisKind::TP);
+    if (!tp_axis) {
+        LOG(ERROR) << context << ": shard reconstruction requires a TP axis";
+        return std::nullopt;
+    }
+    const auto [target_start, target_extent] = calculate_shard_range(
+        reconstruction->global_shape[reconstruction->split_dim], tp_axis->rank,
+        tp_axis->size);
+    std::vector<int64_t> target_shape = reconstruction->global_shape;
+    target_shape[reconstruction->split_dim] = target_extent;
+    auto target_metadata = build_shard_metadata_from_shapes(
+        reconstruction->dtype, reconstruction->global_shape, target_shape,
+        parallelism.axes, 0);
+    if (!target_metadata.has_value()) {
+        LOG(ERROR) << context << ": failed to build target shard metadata";
+        return std::nullopt;
+    }
+    return build_reconstructed_tensor_into_plan_from_sources(
+        buffer_ptr, size, std::move(reconstruction->sources),
+        reconstruction->global_shape, reconstruction->split_dim,
+        *target_metadata, target_start, target_extent, context,
+        reconstruction->allow_empty_fragments);
+}
+
+pybind11::object get_tensor_with_parallelism_shard_full_materialized(
+    const std::string &key, const TensorParallelismSpec &parallelism,
+    const std::string &context) {
+    auto reconstruction =
+        load_parallelism_full_reconstruction_sources(key, parallelism, context);
+    if (!reconstruction.has_value()) {
+        return py::none();
+    }
+
+    const ParallelAxisSpec *tp_axis =
+        find_axis_spec_by_kind(parallelism, LayoutAxisKind::TP);
+    if (!tp_axis) {
+        LOG(ERROR) << context << ": shard reconstruction requires a TP axis";
+        return py::none();
+    }
+    const auto [target_start, target_extent] = calculate_shard_range(
+        reconstruction->global_shape[reconstruction->split_dim], tp_axis->rank,
+        tp_axis->size);
+
+    auto element_size =
+        extract_reconstruction_element_size(reconstruction->sources, context);
+    if (!element_size.has_value()) {
+        return py::none();
+    }
+
+    size_t target_tensor_numel = 1;
+    for (size_t dim = 0; dim < reconstruction->global_shape.size(); ++dim) {
+        const int64_t dim_extent =
+            static_cast<int>(dim) == reconstruction->split_dim
+                ? target_extent
+                : reconstruction->global_shape[dim];
+        target_tensor_numel *= static_cast<size_t>(dim_extent);
+    }
+    const size_t total_length =
+        sizeof(TensorMetadata) + target_tensor_numel * *element_size;
+
+    char *owned_buffer = new char[total_length];
+    if (store_->register_buffer(owned_buffer, total_length) != 0) {
+        LOG(ERROR) << context << ": failed to register reconstruction buffer";
+        delete[] owned_buffer;
+        return py::none();
+    }
+
+    auto plan = build_parallelism_shard_tensor_into_plan(
+        key, reinterpret_cast<uintptr_t>(owned_buffer), total_length,
+        parallelism, context);
+    if (!plan.has_value()) {
+        store_->unregister_buffer(owned_buffer);
+        delete[] owned_buffer;
+        return py::none();
+    }
+
+    std::vector<TensorIntoPlan> plans;
+    plans.push_back(std::move(*plan));
+    auto success = execute_tensor_into_plan_transfers(plans);
+    if (success.empty() || !success[0]) {
+        store_->unregister_buffer(owned_buffer);
+        delete[] owned_buffer;
+        return py::none();
+    }
+
+    return buffer_to_tensor(
+        new BufferHandle(owned_buffer, total_length,
+                         [this, owned_buffer]() {
+                             store_->unregister_buffer(owned_buffer);
+                             delete[] owned_buffer;
+                         }),
+        nullptr, 0);
+}
+
+std::vector<bool> execute_tensor_into_plan_transfers(
+    std::vector<TensorIntoPlan> &plans) {
+    std::vector<bool> success(plans.size(), false);
+    if (plans.empty()) {
+        return success;
+    }
+
+    std::vector<void *> buffers;
+    std::vector<std::vector<std::string>> all_keys;
+    std::vector<std::vector<std::vector<size_t>>> all_dst_offsets;
+    std::vector<std::vector<std::vector<size_t>>> all_src_offsets;
+    std::vector<std::vector<std::vector<size_t>>> all_sizes;
+    buffers.reserve(plans.size());
+    all_keys.reserve(plans.size());
+    all_dst_offsets.reserve(plans.size());
+    all_src_offsets.reserve(plans.size());
+    all_sizes.reserve(plans.size());
+
+    std::vector<size_t> transfer_plan_indices;
+    transfer_plan_indices.reserve(plans.size());
+
+    for (size_t plan_idx = 0; plan_idx < plans.size(); ++plan_idx) {
+        const auto &plan = plans[plan_idx];
+        std::unordered_map<std::string, size_t> key_to_index;
+        std::vector<std::string> keys;
+        std::vector<std::vector<size_t>> dst_offsets;
+        std::vector<std::vector<size_t>> src_offsets;
+        std::vector<std::vector<size_t>> sizes;
+
+        if (plan.regular_full_formula.has_value()) {
+            const auto &formula = *plan.regular_full_formula;
+            keys = formula.read_keys;
+            dst_offsets.resize(keys.size());
+            src_offsets.resize(keys.size());
+            sizes.resize(keys.size());
+
+            int64_t elements_before = 1;
+            for (int i = 0; i < formula.split_dim; ++i) {
+                elements_before *= formula.global_shape[i];
+            }
+            int64_t elements_after = 1;
+            for (size_t i = formula.split_dim + 1;
+                 i < formula.global_shape.size(); ++i) {
+                elements_after *= formula.global_shape[i];
+            }
+            const size_t base_dst_offset =
+                plan.user_buffer_ptr - plan.registered_buffer_ptr;
+            const int64_t split_extent =
+                formula.global_shape[formula.split_dim];
+            for (size_t shard_rank = 0; shard_rank < formula.read_keys.size();
+                 ++shard_rank) {
+                const auto [shard_start, shard_extent] = calculate_shard_range(
+                    split_extent, static_cast<int>(shard_rank),
+                    static_cast<int>(formula.read_keys.size()));
+                if (shard_extent <= 0) {
+                    continue;
+                }
+                const size_t row_bytes = static_cast<size_t>(shard_extent) *
+                                         static_cast<size_t>(elements_after) *
+                                         formula.element_size;
+                if (row_bytes == 0) {
+                    continue;
+                }
+                dst_offsets[shard_rank].reserve(
+                    static_cast<size_t>(elements_before));
+                src_offsets[shard_rank].reserve(
+                    static_cast<size_t>(elements_before));
+                sizes[shard_rank].reserve(static_cast<size_t>(elements_before));
+                for (int64_t slice_idx = 0; slice_idx < elements_before;
+                     ++slice_idx) {
+                    dst_offsets[shard_rank].push_back(
+                        base_dst_offset + sizeof(TensorMetadata) +
+                        static_cast<size_t>(slice_idx * split_extent +
+                                            shard_start) *
+                            static_cast<size_t>(elements_after) *
+                            formula.element_size);
+                    src_offsets[shard_rank].push_back(
+                        formula.data_offset +
+                        static_cast<size_t>(slice_idx * shard_extent) *
+                            static_cast<size_t>(elements_after) *
+                            formula.element_size);
+                    sizes[shard_rank].push_back(row_bytes);
+                }
+            }
+        } else {
+            key_to_index.reserve(plan.fragments.size());
+            keys.reserve(plan.fragments.size());
+            dst_offsets.reserve(plan.fragments.size());
+            src_offsets.reserve(plan.fragments.size());
+            sizes.reserve(plan.fragments.size());
+
+            for (const auto &fragment : plan.fragments) {
+                if (fragment.read_key.empty() || fragment.size == 0) {
+                    continue;
+                }
+                auto [it, inserted] =
+                    key_to_index.emplace(fragment.read_key, keys.size());
+                if (inserted) {
+                    keys.push_back(fragment.read_key);
+                    dst_offsets.emplace_back();
+                    src_offsets.emplace_back();
+                    sizes.emplace_back();
+                }
+                const size_t key_index = it->second;
+                dst_offsets[key_index].push_back(fragment.dst_offset);
+                src_offsets[key_index].push_back(fragment.src_offset);
+                sizes[key_index].push_back(fragment.size);
+            }
+        }
+
+        if (keys.empty()) {
+            if (plan.materialized_metadata.has_value()) {
+                std::memcpy(reinterpret_cast<void *>(plan.user_buffer_ptr),
+                            &*plan.materialized_metadata,
+                            sizeof(TensorMetadata));
+                success[plan_idx] = true;
+            }
+            continue;
+        }
+
+        buffers.push_back(reinterpret_cast<void *>(plan.registered_buffer_ptr));
+        transfer_plan_indices.push_back(plan_idx);
+        all_keys.push_back(std::move(keys));
+        all_dst_offsets.push_back(std::move(dst_offsets));
+        all_src_offsets.push_back(std::move(src_offsets));
+        all_sizes.push_back(std::move(sizes));
+    }
+
+    if (transfer_plan_indices.empty()) {
+        return success;
+    }
+
+    std::vector<std::vector<std::vector<int64_t>>> range_results;
+    {
+        py::gil_scoped_release release_gil;
+        mooncake::PyClient::QueryResultCache merged_query_result_cache;
+        auto now = std::chrono::steady_clock::now();
+        for (auto &plan : plans) {
+            for (auto &planned_query_result : plan.query_results) {
+                if (!planned_query_result.cached_query_result.has_value()) {
+                    continue;
+                }
+                auto query_result = from_cached_query_result_response(
+                    *planned_query_result.cached_query_result, now);
+                if (!query_result || query_result->IsLeaseExpired(now)) {
+                    continue;
+                }
+                merged_query_result_cache.emplace(planned_query_result.read_key,
+                                                  std::move(query_result));
+            }
+        }
+        range_results = store_->get_into_ranges(
+            buffers, all_keys, all_dst_offsets, all_src_offsets, all_sizes,
+            merged_query_result_cache.empty() ? nullptr
+                                              : &merged_query_result_cache);
+    }
+
+    for (size_t i = 0; i < transfer_plan_indices.size(); ++i) {
+        const size_t plan_idx = transfer_plan_indices[i];
+        if (i >= range_results.size() ||
+            range_results[i].size() != all_sizes[i].size()) {
+            continue;
+        }
+
+        success[plan_idx] = true;
+        for (size_t key_idx = 0;
+             key_idx < all_sizes[i].size() && success[plan_idx]; ++key_idx) {
+            if (range_results[i][key_idx].size() !=
+                all_sizes[i][key_idx].size()) {
+                success[plan_idx] = false;
+                break;
+            }
+            for (size_t frag_idx = 0; frag_idx < all_sizes[i][key_idx].size();
+                 ++frag_idx) {
+                if (range_results[i][key_idx][frag_idx] !=
+                    static_cast<int64_t>(all_sizes[i][key_idx][frag_idx])) {
+                    success[plan_idx] = false;
+                    break;
+                }
+            }
+        }
+        if (!success[plan_idx]) {
+            continue;
+        }
+        if (plans[plan_idx].materialized_metadata.has_value()) {
+            std::memcpy(
+                reinterpret_cast<void *>(plans[plan_idx].user_buffer_ptr),
+                &*plans[plan_idx].materialized_metadata,
+                sizeof(TensorMetadata));
+        }
+    }
+    return success;
+}
+
+py::list execute_tensor_into_plans(std::vector<TensorIntoPlan> plans) {
+    py::list results;
+    for (size_t i = 0; i < plans.size(); ++i) {
+        results.append(py::none());
+    }
+    auto success = execute_tensor_into_plan_transfers(plans);
+    for (size_t i = 0; i < plans.size(); ++i) {
+        if (!success[i]) {
+            continue;
+        }
+        results[i] = buffer_to_tensor(
+            NULL, reinterpret_cast<char *>(plans[i].user_buffer_ptr),
+            static_cast<int64_t>(plans[i].total_length));
+    }
+    return results;
+}
+
+std::optional<TensorIntoPlan> build_tensor_into_plan_for_target(
+    const std::string &key, uintptr_t buffer_ptr, size_t size,
+    const py::object &target, const std::string &context,
+    bool allow_formula_plan = true) {
+    auto parsed_target = parse_read_target_spec(target);
+    if (!parsed_target.has_value()) {
+        return std::nullopt;
+    }
+    if (parsed_target->mode == ReadTargetMode::AS_STORED &&
+        !parsed_target->parallelism.has_value()) {
+        return build_tensor_into_plan(key, buffer_ptr, size, context);
+    }
+
+    if (parsed_target->mode == ReadTargetMode::FULL &&
+        !parsed_target->parallelism.has_value()) {
+        return build_writer_shard_full_tensor_into_plan(key, buffer_ptr, size,
+                                                        context);
+    }
+
+    auto parallelism =
+        validate_parallelism_spec(parsed_target->parallelism, context, false);
+    if (!parallelism.has_value()) {
+        return std::nullopt;
+    }
+
+    if (parsed_target->mode == ReadTargetMode::SHARD) {
+        auto resolved = resolve_parallelism_shard_read(
+            [this](const std::string &read_key,
+                   std::shared_ptr<BufferHandle> *buffer_handle_out) {
+                return get_tensor_metadata(read_key, buffer_handle_out);
+            },
+            key, *parallelism);
+        if (resolved.has_value()) {
+            return build_tensor_into_plan(resolved->read_key, buffer_ptr, size,
+                                          context, resolved->metadata);
+        }
+        if (find_axis_spec_by_kind(*parallelism, LayoutAxisKind::TP)) {
+            return build_parallelism_shard_tensor_into_plan(
+                key, buffer_ptr, size, *parallelism, context);
+        }
+        LOG(ERROR) << context << ": parallelism metadata mismatch for key "
+                   << key;
+        return std::nullopt;
+    }
+
+    if (parsed_target->mode == ReadTargetMode::FULL) {
+        return build_parallelism_full_tensor_into_plan(
+            key, buffer_ptr, size, *parallelism, context, allow_formula_plan);
+    }
+
+    LOG(ERROR) << context << ": unsupported ReadTarget mode";
+    return std::nullopt;
+}
+
+pybind11::object get_tensor_with_parallelism(
+    const std::string &key, const py::object &target = py::none()) {
+    auto parsed_target = parse_read_target_spec(target);
+    if (!parsed_target.has_value()) {
+        return py::none();
+    }
+    if (parsed_target->mode == ReadTargetMode::AS_STORED &&
+        !parsed_target->parallelism.has_value()) {
+        return get_tensor(key);
+    }
+    if (parsed_target->mode == ReadTargetMode::FULL &&
+        !parsed_target->parallelism.has_value()) {
+        return get_tensor_with_writer_shard_full(key,
+                                                 "get_tensor_with_parallelism");
+    }
+
+    auto parallelism = validate_parallelism_spec(
+        parsed_target->parallelism, "get_tensor_with_parallelism", false);
+    if (!parallelism.has_value()) {
+        return py::none();
+    }
+
+    if (parsed_target->mode == ReadTargetMode::SHARD) {
+        auto resolved = resolve_parallelism_shard_read(
+            [this](const std::string &read_key,
+                   std::shared_ptr<BufferHandle> *buffer_handle_out) {
+                return get_tensor_metadata(read_key, buffer_handle_out);
+            },
+            key, *parallelism);
+        if (resolved.has_value()) {
+            return buffer_to_tensor(resolved->buffer_handle.get(), nullptr, 0);
+        }
+        if (find_axis_spec_by_kind(*parallelism, LayoutAxisKind::TP)) {
+            auto reconstructed =
+                get_tensor_with_parallelism_shard_full_materialized(
+                    key, *parallelism, "get_tensor_with_parallelism");
+            if (!reconstructed.is_none()) {
+                return reconstructed;
+            }
+        }
+        if (uses_legacy_tp_storage_key(*parallelism)) {
+            const auto *tp_axis =
+                find_axis_spec_by_kind(*parallelism, LayoutAxisKind::TP);
+            if (tp_axis) {
+                return get_tensor(
+                    resolve_tp_read_key(key, tp_axis->rank, tp_axis->size));
+            }
+        }
+        LOG(ERROR) << "get_tensor_with_parallelism"
+                   << ": parallelism metadata mismatch for key " << key;
+        return py::none();
+    }
+
+    if (parsed_target->mode == ReadTargetMode::FULL) {
+        const auto *tp_axis =
+            find_axis_spec_by_kind(*parallelism, LayoutAxisKind::TP);
+        if (!tp_axis) {
+            LOG(ERROR) << "get_tensor_with_parallelism"
+                       << ": full reconstruction requires a TP axis";
+            return py::none();
+        }
+        return get_tensor_with_tp_full(
+            key, tp_axis->rank, tp_axis->size, tp_axis->split_dim.value_or(0),
+            "get_tensor_with_parallelism", *parallelism);
+    }
+
+    LOG(ERROR) << "get_tensor_with_parallelism"
+               << ": unsupported ReadTarget mode";
+    return py::none();
+}
+
+pybind11::list batch_get_tensor_with_parallelism(
+    const std::vector<std::string> &keys,
+    const py::object &targets = py::none()) {
+    if (targets.is_none()) {
+        return batch_get_tensor(keys);
+    }
+
+    auto target_list = validate_batch_request_list(
+        targets, keys.size(), "batch_get_tensor_with_parallelism", "targets");
+    if (!target_list.has_value()) {
+        py::list empty;
+        for (size_t i = 0; i < keys.size(); ++i) {
+            empty.append(py::none());
+        }
+        return empty;
+    }
+
+    py::list results;
+    for (size_t i = 0; i < keys.size(); ++i) {
+        results.append(get_tensor_with_parallelism(
+            keys[i], py::reinterpret_borrow<py::object>((*target_list)[i])));
+    }
+    return results;
+}
+
+pybind11::object get_tensor_with_parallelism_into(
+    const std::string &key, uintptr_t buffer_ptr, size_t size,
+    const py::object &target = py::none()) {
+    auto plan = build_tensor_into_plan_for_target(
+        key, buffer_ptr, size, target, "get_tensor_with_parallelism_into");
+    if (!plan.has_value()) {
+        return py::none();
+    }
+    const bool used_formula_plan = plan->regular_full_formula.has_value();
+    auto results = execute_tensor_into_plans({*plan});
+    if (!results.empty() && !results[0].is_none()) {
+        return py::reinterpret_borrow<py::object>(results[0]);
+    }
+
+    if (used_formula_plan) {
+        LOG(WARNING)
+            << "get_tensor_with_parallelism_into"
+            << ": formula plan execution failed, falling back to generic path";
+        auto fallback_plan = build_tensor_into_plan_for_target(
+            key, buffer_ptr, size, target, "get_tensor_with_parallelism_into",
+            false /* allow_formula_plan */);
+        if (!fallback_plan.has_value()) {
+            return py::none();
+        }
+        auto fallback_results = execute_tensor_into_plans({*fallback_plan});
+        if (!fallback_results.empty()) {
+            return py::reinterpret_borrow<py::object>(fallback_results[0]);
+        }
+    }
+    return py::none();
+}
+
+pybind11::list batch_get_tensor_with_parallelism_into(
+    const std::vector<std::string> &keys,
+    const std::vector<uintptr_t> &buffer_ptrs, const std::vector<size_t> &sizes,
+    const py::object &targets = py::none()) {
+    py::list empty;
+    for (size_t i = 0; i < keys.size(); ++i) {
+        empty.append(py::none());
+    }
+    if (keys.size() != buffer_ptrs.size() || keys.size() != sizes.size()) {
+        LOG(ERROR)
+            << "batch_get_tensor_with_parallelism_into"
+            << ": keys, buffer_ptrs, and sizes must have the same length";
+        return empty;
+    }
+
+    std::optional<py::list> target_list = std::nullopt;
+    if (!targets.is_none()) {
+        target_list = validate_batch_request_list(
+            targets, keys.size(), "batch_get_tensor_with_parallelism_into",
+            "targets");
+        if (!target_list.has_value()) {
+            return empty;
+        }
+    }
+
+    std::vector<TensorIntoPlan> plans;
+    plans.reserve(keys.size());
+    std::vector<size_t> plan_indices;
+    plan_indices.reserve(keys.size());
+    std::vector<bool> plan_used_formula;
+    plan_used_formula.reserve(keys.size());
+    for (size_t i = 0; i < keys.size(); ++i) {
+        py::object target =
+            target_list.has_value()
+                ? py::reinterpret_borrow<py::object>((*target_list)[i])
+                : py::none();
+        auto plan = build_tensor_into_plan_for_target(
+            keys[i], buffer_ptrs[i], sizes[i], target,
+            "batch_get_tensor_with_parallelism_into");
+        if (!plan.has_value()) {
+            continue;
+        }
+        plan_indices.push_back(i);
+        plan_used_formula.push_back(plan->regular_full_formula.has_value());
+        plans.push_back(std::move(*plan));
+    }
+
+    auto results = execute_tensor_into_plans(std::move(plans));
+    std::vector<TensorIntoPlan> fallback_plans;
+    std::vector<size_t> fallback_indices;
+    fallback_plans.reserve(plan_indices.size());
+    fallback_indices.reserve(plan_indices.size());
+    for (size_t i = 0; i < plan_indices.size() && i < results.size(); ++i) {
+        if (!results[i].is_none() || !plan_used_formula[i]) {
+            empty[plan_indices[i]] = results[i];
+            continue;
+        }
+
+        const size_t original_index = plan_indices[i];
+        LOG(WARNING) << "batch_get_tensor_with_parallelism_into"
+                     << ": formula plan execution failed for key "
+                     << keys[original_index]
+                     << ", falling back to generic path";
+        py::object target = target_list.has_value()
+                                ? py::reinterpret_borrow<py::object>(
+                                      (*target_list)[original_index])
+                                : py::none();
+        auto fallback_plan = build_tensor_into_plan_for_target(
+            keys[original_index], buffer_ptrs[original_index],
+            sizes[original_index], target,
+            "batch_get_tensor_with_parallelism_into",
+            false /* allow_formula_plan */);
+        if (!fallback_plan.has_value()) {
+            continue;
+        }
+        fallback_indices.push_back(original_index);
+        fallback_plans.push_back(std::move(*fallback_plan));
+    }
+
+    auto fallback_results =
+        execute_tensor_into_plans(std::move(fallback_plans));
+    for (size_t i = 0;
+         i < fallback_indices.size() && i < fallback_results.size(); ++i) {
+        empty[fallback_indices[i]] = fallback_results[i];
+    }
+    return empty;
+}
diff --git a/mooncake-integration/store/store_py_parallel_write.h b/mooncake-integration/store/store_py_parallel_write.h
new file mode 100644
index 00000000..2bba450a
--- /dev/null
+++ b/mooncake-integration/store/store_py_parallel_write.h
@@ -0,0 +1,1487 @@
+template <typename ManifestWriteFn>
+int write_manifest_impl(const std::string &key,
+                        const WriterShardManifest &manifest,
+                        const char *operation_name,
+                        ManifestWriteFn &&write_manifest) {
+    std::span<const char> bytes(reinterpret_cast<const char *>(&manifest),
+                                sizeof(WriterShardManifest));
+    py::gil_scoped_release release_gil;
+    int ret = write_manifest(bytes);
+    if (ret != 0) {
+        LOG(ERROR) << operation_name << " manifest failed for key " << key
+                   << " with code " << ret;
+    }
+    return ret;
+}
+
+template <typename BatchWriteFromFn>
+std::vector<int> batch_write_tensor_impl(const std::vector<std::string> &keys,
+                                         const std::vector<PyTensorInfo> &infos,
+                                         const ReplicateConfig &config,
+                                         const char *operation_name,
+                                         BatchWriteFromFn &&batch_write_from) {
+    auto group_ids_error =
+        ValidateGroupIdsForBatchConfig(config, keys.size(), operation_name);
+    if (!group_ids_error.empty()) {
+        return group_ids_error;
+    }
+
+    std::vector<int> results(keys.size(), 0);
+
+    {
+        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()) {
+                results[i] = to_py_ret(ErrorCode::INVALID_PARAMS);
+                continue;
+            }
+
+            size_t total_size =
+                infos[i].metadata.header.data_offset + infos[i].tensor_size;
+            auto alloc_result =
+                store_->client_buffer_allocator_->allocate(total_size);
+
+            if (!alloc_result) {
+                LOG(ERROR) << "Failed to allocate buffer for " << operation_name
+                           << " key: " << keys[i];
+                results[i] = to_py_ret(ErrorCode::NO_AVAILABLE_HANDLE);
+                continue;
+            }
+
+            char *dst = static_cast<char *>(alloc_result->ptr());
+            std::memcpy(dst, &infos[i].metadata,
+                        infos[i].metadata.header.data_offset);
+            if (infos[i].tensor_size > 0) {
+                std::memcpy(dst + infos[i].metadata.header.data_offset,
+                            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()) {
+            ReplicateConfig write_config =
+                MakeIndexedConfig(config, original_indices);
+            std::vector<int> op_results = batch_write_from(
+                valid_keys, buffer_ptrs, buffer_sizes, write_config);
+            for (size_t i = 0; i < op_results.size(); ++i) {
+                results[original_indices[i]] = op_results[i];
+            }
+        }
+    }
+
+    return results;
+}
+
+bool ensure_tensor_write_supported(const char *operation_name) const {
+    if (!is_client_initialized() || use_dummy_client_) {
+        LOG(ERROR) << operation_name
+                   << ": client not initialized or dummy client not "
+                      "supported for tensors";
+        return false;
+    }
+    return true;
+}
+
+template <typename WriteFn>
+int execute_single_tensor_write(const char *operation_name,
+                                const ReplicateConfig &config,
+                                WriteFn &&write_fn) {
+    if (!ensure_tensor_write_supported(operation_name)) {
+        return to_py_ret(ErrorCode::INVALID_PARAMS);
+    }
+    int validate_result = validate_replicate_config(config);
+    if (validate_result) {
+        return validate_result;
+    }
+    return write_fn();
+}
+
+template <typename WriteFn>
+std::vector<int> execute_batch_tensor_write(
+    const char *operation_name, const char *size_error_context,
+    const std::vector<std::string> &keys, size_t value_count,
+    const ReplicateConfig &config, WriteFn &&write_fn) {
+    if (!ensure_tensor_write_supported(operation_name)) {
+        return std::vector<int>(keys.size(),
+                                to_py_ret(ErrorCode::INVALID_PARAMS));
+    }
+    if (keys.size() != value_count || keys.empty()) {
+        if (!keys.empty()) {
+            LOG(ERROR) << size_error_context;
+        }
+        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(), validate_result);
+    }
+    return write_fn();
+}
+
+bool validate_tensor_object_buffers(const std::vector<std::string> &keys,
+                                    const std::vector<uintptr_t> &buffer_ptrs,
+                                    const std::vector<size_t> &sizes,
+                                    const char *size_error_context,
+                                    const char *buffer_error_context) {
+    if (keys.size() != buffer_ptrs.size() || keys.size() != sizes.size()) {
+        LOG(ERROR) << size_error_context;
+        return false;
+    }
+    for (size_t i = 0; i < sizes.size(); ++i) {
+        if (!is_valid_tensor_object_buffer(buffer_ptrs[i], sizes[i],
+                                           std::string(buffer_error_context) +
+                                               " at index " +
+                                               std::to_string(i))) {
+            return false;
+        }
+    }
+    return true;
+}
+
+template <typename WriteFromFn>
+int execute_single_tensor_write_from(const char *operation_name,
+                                     const std::string &key,
+                                     uintptr_t buffer_ptr, size_t size,
+                                     const ReplicateConfig &config,
+                                     WriteFromFn &&write_from_fn) {
+    if (!ensure_tensor_write_supported(operation_name)) {
+        return to_py_ret(ErrorCode::INVALID_PARAMS);
+    }
+    if (!is_valid_tensor_object_buffer(buffer_ptr, size, operation_name)) {
+        return to_py_ret(ErrorCode::INVALID_PARAMS);
+    }
+    int validate_result = validate_replicate_config(config);
+    if (validate_result) {
+        return validate_result;
+    }
+    return write_from_fn(key, reinterpret_cast<void *>(buffer_ptr), size,
+                         config);
+}
+
+template <typename BatchWriteFromFn>
+std::vector<int> execute_batch_tensor_write_from(
+    const char *operation_name, const char *size_error_context,
+    const std::vector<std::string> &keys,
+    const std::vector<uintptr_t> &buffer_ptrs, const std::vector<size_t> &sizes,
+    const ReplicateConfig &config, BatchWriteFromFn &&batch_write_from_fn) {
+    if (!ensure_tensor_write_supported(operation_name)) {
+        return std::vector<int>(keys.size(),
+                                to_py_ret(ErrorCode::INVALID_PARAMS));
+    }
+    if (keys.empty()) {
+        return std::vector<int>();
+    }
+    if (!validate_tensor_object_buffers(keys, buffer_ptrs, sizes,
+                                        size_error_context,
+                                        "tensor object buffer")) {
+        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(), validate_result);
+    }
+    std::vector<void *> buffers;
+    buffers.reserve(buffer_ptrs.size());
+    for (uintptr_t ptr : buffer_ptrs) {
+        buffers.push_back(reinterpret_cast<void *>(ptr));
+    }
+    return batch_write_from_fn(keys, buffers, sizes, config);
+}
+
+template <typename DirectWriteFn, typename ParallelismWriteFn,
+          typename WriterPartitionWriteFn>
+std::vector<int> execute_batch_parallelism_write_requests(
+    const std::vector<std::string> &keys, size_t value_count,
+    const py::object &parallelisms, const py::object &writer_partitions,
+    const char *error_context, DirectWriteFn &&direct_write,
+    ParallelismWriteFn &&parallelism_write,
+    WriterPartitionWriteFn &&writer_partition_write) {
+    if (!parallelisms.is_none() && !writer_partitions.is_none()) {
+        LOG(ERROR)
+            << error_context
+            << ": writer_partitions cannot be combined with parallelisms";
+        return std::vector<int>(keys.size(),
+                                to_py_ret(ErrorCode::INVALID_PARAMS));
+    }
+    if (value_count != keys.size()) {
+        LOG(ERROR) << error_context
+                   << ": values and keys must have the same length";
+        return std::vector<int>(keys.size(),
+                                to_py_ret(ErrorCode::INVALID_PARAMS));
+    }
+    if (parallelisms.is_none() && writer_partitions.is_none()) {
+        return direct_write();
+    }
+
+    std::vector<int> results(keys.size(), to_py_ret(ErrorCode::INVALID_PARAMS));
+    if (!parallelisms.is_none()) {
+        auto parallelism_list = validate_batch_parallelism_list(
+            parallelisms, keys.size(), error_context);
+        if (!parallelism_list.has_value()) {
+            return std::vector<int>(keys.size(),
+                                    to_py_ret(ErrorCode::INVALID_PARAMS));
+        }
+        for (size_t i = 0; i < keys.size(); ++i) {
+            results[i] = parallelism_write(i, (*parallelism_list)[i]);
+        }
+        return results;
+    }
+
+    auto writer_partition_list = validate_batch_writer_partition_list(
+        writer_partitions, keys.size(), error_context);
+    if (!writer_partition_list.has_value()) {
+        return std::vector<int>(keys.size(),
+                                to_py_ret(ErrorCode::INVALID_PARAMS));
+    }
+    for (size_t i = 0; i < keys.size(); ++i) {
+        results[i] = writer_partition_write(i, (*writer_partition_list)[i]);
+    }
+    return results;
+}
+
+template <typename DirectWriteFn, typename WriterShardWriteFn,
+          typename LegacyTpWriteFn, typename TpParallelWriteFn,
+          typename GenericShardWriteFn>
+int execute_parallelism_tensor_write_route(
+    const std::string &key, pybind11::object tensor,
+    const ResolvedParallelismWriteRequest &request,
+    const ReplicateConfig &config, const char *error_context,
+    DirectWriteFn &&direct_write, WriterShardWriteFn &&writer_shard_write,
+    LegacyTpWriteFn &&legacy_tp_write, TpParallelWriteFn &&tp_parallel_write,
+    GenericShardWriteFn &&generic_shard_write) {
+    switch (request.route) {
+        case ParallelismWriteStorageRoute::DIRECT_FULL_OBJECT:
+            if (!ensure_tensor_write_supported(error_context)) {
+                return to_py_ret(ErrorCode::INVALID_PARAMS);
+            }
+            if (int validate_result = validate_replicate_config(config);
+                validate_result) {
+                return validate_result;
+            }
+            return direct_write(key, tensor, config);
+        case ParallelismWriteStorageRoute::WRITER_PARTITION_SHARD:
+            if (!request.writer_partition.has_value()) {
+                return to_py_ret(ErrorCode::INVALID_PARAMS);
+            }
+            return writer_shard_write(key, tensor, *request.writer_partition,
+                                      config);
+        case ParallelismWriteStorageRoute::LEGACY_SINGLE_TP:
+        case ParallelismWriteStorageRoute::TP_SHARDED_PARALLELISM:
+        case ParallelismWriteStorageRoute::GENERIC_PARALLELISM_SHARD: {
+            if (!request.parallelism.has_value()) {
+                return to_py_ret(ErrorCode::INVALID_PARAMS);
+            }
+            switch (request.route) {
+                case ParallelismWriteStorageRoute::LEGACY_SINGLE_TP: {
+                    const auto &axis = request.parallelism->axes[0];
+                    return legacy_tp_write(key, tensor, config, axis.rank,
+                                           axis.size,
+                                           axis.split_dim.value_or(0));
+                }
+                case ParallelismWriteStorageRoute::TP_SHARDED_PARALLELISM: {
+                    auto tp_axis_index =
+                        find_tp_axis_index(request.parallelism->axes);
+                    if (!tp_axis_index.has_value()) {
+                        return to_py_ret(ErrorCode::INVALID_PARAMS);
+                    }
+                    const auto &axis =
+                        request.parallelism->axes[*tp_axis_index];
+                    return tp_parallel_write(
+                        key, tensor, config, axis.rank, axis.size,
+                        axis.split_dim.value_or(0), request.parallelism->axes);
+                }
+                case ParallelismWriteStorageRoute::GENERIC_PARALLELISM_SHARD:
+                    return generic_shard_write(key, tensor,
+                                               *request.parallelism, config);
+                default:
+                    break;
+            }
+            break;
+        }
+    }
+    return to_py_ret(ErrorCode::INVALID_PARAMS);
+}
+
+template <typename WritePartsFn>
+int execute_tensor_parts_write(const std::string &key,
+                               const std::vector<std::span<const char>> &values,
+                               const ReplicateConfig &config,
+                               const TensorWriteStoreOps &ops,
+                               WritePartsFn &&write_parts) {
+    py::gil_scoped_release release_gil;
+    int ret = write_parts(key, values, config);
+    if (ret != 0) {
+        LOG(ERROR) << ops.parts_operation_name << " failed for key " << key
+                   << " with code " << ret;
+    }
+    return ret;
+}
+
+template <typename DirectWriteFn>
+int execute_parallelism_direct_tensor_write_from(const std::string &key,
+                                                 uintptr_t buffer_ptr,
+                                                 size_t size,
+                                                 const ReplicateConfig &config,
+                                                 const char *operation_name,
+                                                 DirectWriteFn &&direct_write) {
+    if (is_default_replicate_config(config)) {
+        return direct_write(key, buffer_ptr, size);
+    }
+    return execute_single_tensor_write_from(
+        operation_name, key, buffer_ptr, size, config,
+        [&](const std::string &write_key, void *buffer, size_t write_size,
+            const ReplicateConfig &write_config) {
+            return direct_write(write_key, reinterpret_cast<uintptr_t>(buffer),
+                                write_size, write_config);
+        });
+}
+
+pybind11::object decode_tensor_object_buffer(uintptr_t buffer_ptr, size_t size,
+                                             const char *operation_name) {
+    if (!is_valid_tensor_object_buffer(buffer_ptr, size, operation_name)) {
+        return py::none();
+    }
+    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 tensor buffer for " << operation_name;
+    }
+    return tensor;
+}
+
+std::string get_tp_write_shard_key(const std::string &base_key, int rank,
+                                   const std::vector<ParallelAxisSpec> &axes) {
+    if (axes.empty()) {
+        return get_tp_key_name(base_key, rank);
+    }
+    auto shard_axes = axes;
+    auto tp_axis_index = find_tp_axis_index(shard_axes);
+    if (!tp_axis_index.has_value()) {
+        return std::string();
+    }
+    shard_axes[*tp_axis_index].rank = rank;
+    return get_parallelism_key_name(base_key,
+                                    TensorParallelismSpec{shard_axes});
+}
+
+template <typename WriteShardFn>
+int execute_tp_tensor_write_impl(const std::string &key,
+                                 pybind11::object tensor,
+                                 const ReplicateConfig &config, int tp_size,
+                                 int split_dim,
+                                 const std::vector<ParallelAxisSpec> &axes,
+                                 const char *error_context,
+                                 WriteShardFn &&write_shard) {
+    try {
+        auto shard_infos = build_tp_shard_infos(
+            tensor, tp_size, split_dim,
+            [&](int rank) { return get_tp_write_shard_key(key, rank, axes); },
+            axes);
+        if (!shard_infos.has_value()) {
+            LOG(ERROR) << error_context;
+            return to_py_ret(ErrorCode::INVALID_PARAMS);
+        }
+
+        for (int rank = 0; rank < tp_size; ++rank) {
+            std::string shard_key = get_tp_write_shard_key(key, rank, axes);
+            if (shard_key.empty()) {
+                return to_py_ret(ErrorCode::INVALID_PARAMS);
+            }
+            int ret = write_shard(shard_key, (*shard_infos)[rank], config);
+            if (ret != 0) return ret;
+        }
+        return 0;
+    } catch (const std::exception &e) {
+        LOG(ERROR) << error_context << ": " << e.what();
+        return to_py_ret(ErrorCode::INVALID_PARAMS);
+    }
+}
+
+template <typename WritePartsFn>
+int execute_tp_tensor_write_from_impl(const std::string &key,
+                                      uintptr_t buffer_ptr, size_t size,
+                                      const ReplicateConfig &config,
+                                      int tp_size, int split_dim,
+                                      const std::vector<ParallelAxisSpec> &axes,
+                                      const char *error_context,
+                                      WritePartsFn &&write_parts) {
+    auto parsed =
+        parse_tensor_metadata_from_raw_buffer(buffer_ptr, size, error_context);
+    if (!parsed.has_value()) {
+        return to_py_ret(ErrorCode::INVALID_PARAMS);
+    }
+
+    for (int rank = 0; rank < tp_size; ++rank) {
+        std::string shard_key = get_tp_write_shard_key(key, rank, axes);
+        if (shard_key.empty()) {
+            return to_py_ret(ErrorCode::INVALID_PARAMS);
+        }
+
+        auto plan = build_raw_tensor_shard_write_plan(*parsed, split_dim, rank,
+                                                      tp_size, error_context);
+        if (!plan.has_value()) {
+            return to_py_ret(ErrorCode::INVALID_PARAMS);
+        }
+        const auto global_shape = TensorShapeToVector(
+            parsed->metadata.layout.global_shape, parsed->metadata.header.ndim);
+        const auto local_shape = TensorShapeToVector(
+            plan->metadata.layout.local_shape, plan->metadata.header.ndim);
+        if (axes.empty()) {
+            ParallelAxisSpec axis_spec{"tp",      rank,         tp_size,
+                                       split_dim, std::nullopt, std::nullopt};
+            auto metadata = build_shard_metadata_from_shapes(
+                parsed->metadata.header.dtype, global_shape, local_shape,
+                {axis_spec}, plan->metadata.header.data_bytes);
+            if (!metadata.has_value()) {
+                return to_py_ret(ErrorCode::INVALID_PARAMS);
+            }
+            plan->metadata = *metadata;
+        } else {
+            auto shard_axes = axes;
+            auto tp_axis_index = find_tp_axis_index(shard_axes);
+            if (!tp_axis_index.has_value()) {
+                return to_py_ret(ErrorCode::INVALID_PARAMS);
+            }
+            shard_axes[*tp_axis_index].rank = rank;
+            const auto local_shape = TensorShapeToVector(
+                plan->metadata.layout.local_shape, plan->metadata.header.ndim);
+            auto metadata = build_shard_metadata_from_shapes(
+                parsed->metadata.header.dtype, global_shape, local_shape,
+                shard_axes, plan->metadata.header.data_bytes);
+            if (!metadata.has_value()) {
+                return to_py_ret(ErrorCode::INVALID_PARAMS);
+            }
+            plan->metadata = *metadata;
+        }
+
+        std::vector<std::span<const char>> values;
+        values.emplace_back(reinterpret_cast<const char *>(&plan->metadata),
+                            plan->metadata.header.data_offset);
+        for (const auto &[src_offset, part_size] : plan->data_ranges) {
+            values.emplace_back(
+                reinterpret_cast<const char *>(buffer_ptr + src_offset),
+                part_size);
+        }
+        int ret = write_parts(shard_key, values, config);
+        if (ret != 0) {
+            return ret;
+        }
+    }
+    return 0;
+}
+
+template <typename BatchWriteFn>
+std::vector<int> batch_execute_tp_tensor_write_impl(
+    const std::vector<std::string> &base_keys,
+    const pybind11::list &tensors_list, const ReplicateConfig &config,
+    int tp_size, int split_dim, const char *error_context,
+    BatchWriteFn &&batch_write) {
+    std::vector<std::string> all_chunk_keys;
+    std::vector<PyTensorInfo> all_chunk_infos;
+    std::vector<size_t> processed_indices;
+    std::vector<int> final_results(base_keys.size(),
+                                   to_py_ret(ErrorCode::INVALID_PARAMS));
+    try {
+        for (size_t i = 0; i < base_keys.size(); ++i) {
+            py::object tensor = tensors_list[i];
+            if (tensor.is_none()) {
+                final_results[i] = to_py_ret(ErrorCode::INVALID_PARAMS);
+                continue;
+            }
+
+            auto shard_infos = build_tp_shard_infos(
+                tensor, tp_size, split_dim,
+                [&](int rank) { return get_tp_key_name(base_keys[i], rank); });
+            if (!shard_infos.has_value()) {
+                final_results[i] = to_py_ret(ErrorCode::INVALID_PARAMS);
+                continue;
+            }
+
+            processed_indices.push_back(i);
+            for (int rank = 0; rank < tp_size; ++rank) {
+                all_chunk_keys.push_back(get_tp_key_name(base_keys[i], rank));
+                all_chunk_infos.push_back((*shard_infos)[rank]);
+            }
+        }
+
+        if (all_chunk_keys.empty()) return final_results;
+
+        std::vector<int> chunk_results =
+            batch_write(all_chunk_keys, all_chunk_infos, config);
+        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;
+                    all_ok = false;
+                    break;
+                }
+            }
+            if (all_ok) {
+                final_results[original_idx] = 0;
+            }
+        }
+    } catch (const std::exception &e) {
+        LOG(ERROR) << error_context << ": " << e.what();
+    }
+
+    return final_results;
+}
+
+template <typename BatchTensorWriteFn>
+std::vector<int> batch_decode_tensor_buffers_and_write(
+    const std::vector<std::string> &keys,
+    const std::vector<uintptr_t> &buffer_ptrs, const std::vector<size_t> &sizes,
+    const char *decode_error_context,
+    BatchTensorWriteFn &&batch_tensor_write_fn) {
+    py::list tensors_list;
+    std::vector<size_t> processed_indices;
+    std::vector<int> final_results(keys.size(), 0);
+
+    for (size_t i = 0; i < keys.size(); ++i) {
+        if (!is_valid_tensor_object_buffer(
+                buffer_ptrs[i], sizes[i],
+                std::string("tensor object buffer at index ") +
+                    std::to_string(i))) {
+            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) << decode_error_context << " 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(keys[idx]);
+    }
+
+    std::vector<int> op_results =
+        batch_tensor_write_fn(valid_keys, tensors_list);
+    for (size_t i = 0; i < processed_indices.size(); ++i) {
+        final_results[processed_indices[i]] = op_results[i];
+    }
+    return final_results;
+}
+
+template <typename WriteShardFn>
+int execute_direct_parallelism_shard_write_from(
+    const std::string &key, uintptr_t buffer_ptr, size_t size,
+    const TensorParallelismSpec &parallelism, const ReplicateConfig &config,
+    const char *error_context, const TensorWriteStoreOps &ops,
+    WriteShardFn &&write_shard) {
+    auto parsed =
+        parse_tensor_metadata_from_raw_buffer(buffer_ptr, size, error_context);
+    if (!parsed.has_value()) {
+        return to_py_ret(ErrorCode::INVALID_PARAMS);
+    }
+
+    auto metadata = build_shard_metadata_from_shapes(
+        parsed->metadata.header.dtype,
+        TensorShapeToVector(parsed->metadata.layout.global_shape,
+                            parsed->metadata.header.ndim),
+        TensorShapeToVector(parsed->metadata.layout.local_shape,
+                            parsed->metadata.header.ndim),
+        parallelism.axes, parsed->data_bytes);
+    if (!metadata.has_value()) {
+        return to_py_ret(ErrorCode::INVALID_PARAMS);
+    }
+
+    std::string shard_key = get_parallelism_key_name(key, parallelism);
+    std::vector<std::span<const char>> values;
+    values.emplace_back(reinterpret_cast<const char *>(&*metadata),
+                        metadata->header.data_offset);
+    if (parsed->data_bytes > 0) {
+        values.emplace_back(
+            reinterpret_cast<const char *>(buffer_ptr + parsed->data_offset),
+            parsed->data_bytes);
+    }
+    return execute_tensor_parts_write(shard_key, values, config, ops,
+                                      std::forward<WriteShardFn>(write_shard));
+}
+
+template <typename WritePartsFn, typename WriteManifestFn>
+int execute_writer_partition_shard_write_from(
+    const std::string &key, uintptr_t buffer_ptr, size_t size,
+    const WriterPartitionSpec &writer, const ReplicateConfig &config,
+    const char *error_context, const TensorWriteStoreOps &ops,
+    WritePartsFn &&write_parts, WriteManifestFn &&write_manifest) {
+    auto parsed =
+        parse_tensor_metadata_from_raw_buffer(buffer_ptr, size, error_context);
+    if (!parsed.has_value()) {
+        return to_py_ret(ErrorCode::INVALID_PARAMS);
+    }
+
+    auto plan = build_raw_tensor_shard_write_plan(
+        *parsed, writer.split_dim, writer.rank, writer.size, error_context);
+    if (!plan.has_value()) {
+        return to_py_ret(ErrorCode::INVALID_PARAMS);
+    }
+
+    ParallelAxisSpec axis_spec{"tp",         writer.rank,
+                               writer.size,  writer.split_dim,
+                               std::nullopt, std::nullopt};
+    auto metadata = build_shard_metadata_from_shapes(
+        parsed->metadata.header.dtype,
+        TensorShapeToVector(parsed->metadata.layout.global_shape,
+                            parsed->metadata.header.ndim),
+        TensorShapeToVector(plan->metadata.layout.local_shape,
+                            plan->metadata.header.ndim),
+        {axis_spec}, plan->metadata.header.data_bytes);
+    if (!metadata.has_value()) {
+        return to_py_ret(ErrorCode::INVALID_PARAMS);
+    }
+    plan->metadata = *metadata;
+
+    std::vector<std::span<const char>> values;
+    values.emplace_back(reinterpret_cast<const char *>(&plan->metadata),
+                        plan->metadata.header.data_offset);
+    for (const auto &[src_offset, part_size] : plan->data_ranges) {
+        values.emplace_back(
+            reinterpret_cast<const char *>(buffer_ptr + src_offset), part_size);
+    }
+
+    int ret = execute_tensor_parts_write(
+        get_writer_shard_key_name(key, writer), values, config, ops,
+        std::forward<WritePartsFn>(write_parts));
+    if (ret != 0) {
+        return ret;
+    }
+    return write_manifest(
+        get_writer_manifest_key_name(key),
+        build_writer_shard_manifest_from_shape(
+            TensorShapeToVector(parsed->metadata.layout.global_shape,
+                                parsed->metadata.header.ndim),
+            parsed->metadata.header.dtype, writer),
+        config);
+}
+
+template <typename DirectFromWriteFn, typename WriterShardFromWriteFn,
+          typename TpShardFromWriteFn, typename GenericShardFromWriteFn,
+          typename TensorRouteWriteFn>
+int execute_parallelism_tensor_write_from_route(
+    const std::string &key, uintptr_t buffer_ptr, size_t size,
+    const ResolvedParallelismWriteRequest &request,
+    const ReplicateConfig &config, const char *error_context,
+    DirectFromWriteFn &&direct_from_write,
+    WriterShardFromWriteFn &&writer_shard_from_write,
+    TpShardFromWriteFn &&tp_shard_from_write,
+    GenericShardFromWriteFn &&generic_shard_from_write,
+    TensorRouteWriteFn &&tensor_route_write) {
+    switch (request.route) {
+        case ParallelismWriteStorageRoute::DIRECT_FULL_OBJECT:
+            return direct_from_write(key, buffer_ptr, size, config);
+        case ParallelismWriteStorageRoute::WRITER_PARTITION_SHARD:
+            if (!request.writer_partition.has_value()) {
+                return to_py_ret(ErrorCode::INVALID_PARAMS);
+            }
+            return writer_shard_from_write(key, buffer_ptr, size,
+                                           *request.writer_partition, config);
+        case ParallelismWriteStorageRoute::LEGACY_SINGLE_TP:
+        case ParallelismWriteStorageRoute::TP_SHARDED_PARALLELISM: {
+            if (!request.parallelism.has_value()) {
+                return to_py_ret(ErrorCode::INVALID_PARAMS);
+            }
+            auto tp_axis_index = find_tp_axis_index(request.parallelism->axes);
+            if (!tp_axis_index.has_value()) {
+                return to_py_ret(ErrorCode::INVALID_PARAMS);
+            }
+            const auto &axis = request.parallelism->axes[*tp_axis_index];
+            const auto axes =
+                request.route == ParallelismWriteStorageRoute::LEGACY_SINGLE_TP
+                    ? std::vector<ParallelAxisSpec>{}
+                    : request.parallelism->axes;
+            return tp_shard_from_write(key, buffer_ptr, size, axis.size,
+                                       axis.split_dim.value_or(0), axes,
+                                       config);
+        }
+        case ParallelismWriteStorageRoute::GENERIC_PARALLELISM_SHARD:
+            if (!request.parallelism.has_value()) {
+                return to_py_ret(ErrorCode::INVALID_PARAMS);
+            }
+            return generic_shard_from_write(key, buffer_ptr, size,
+                                            *request.parallelism, config);
+    }
+
+    pybind11::object tensor =
+        decode_tensor_object_buffer(buffer_ptr, size, error_context);
+    if (tensor.is_none()) {
+        return to_py_ret(ErrorCode::INVALID_PARAMS);
+    }
+    return tensor_route_write(key, tensor, request, config);
+}
+
+int put_tensor_with_writer_shards(
+    const std::string &key, pybind11::object tensor,
+    const WriterPartitionSpec &writer,
+    const ReplicateConfig &config = ReplicateConfig{}) {
+    if (!ensure_tensor_write_supported("put_tensor_with_writer_shards")) {
+        return to_py_ret(ErrorCode::INVALID_PARAMS);
+    }
+
+    auto writer_info = build_writer_shard_tensor_info(
+        key, tensor, writer, "put_tensor_with_writer_shards");
+    if (!writer_info.has_value()) {
+        return to_py_ret(ErrorCode::INVALID_PARAMS);
+    }
+
+    int ret =
+        put_tensor_info_impl(writer_info->shard_key, writer_info->info, config);
+    if (ret != 0) {
+        return ret;
+    }
+    return put_manifest_impl(get_writer_manifest_key_name(key),
+                             writer_info->manifest, config);
+}
+
+int put_tensor_parallelism_tp_impl(const std::string &key,
+                                   pybind11::object tensor,
+                                   const ReplicateConfig &config, int tp_rank,
+                                   int tp_size, int split_dim,
+                                   const std::vector<ParallelAxisSpec> &axes) {
+    std::string shard_key = get_tp_write_shard_key(key, tp_rank, axes);
+    if (shard_key.empty()) {
+        return to_py_ret(ErrorCode::INVALID_PARAMS);
+    }
+    auto shard_info = build_direct_parallelism_shard_info(
+        tensor, axes, shard_key, true /* infer_global_shape */);
+    if (!shard_info.has_value()) {
+        return to_py_ret(ErrorCode::INVALID_PARAMS);
+    }
+    int ret = put_tensor_info_impl(shard_key, shard_info->info, config);
+    if (ret != 0) {
+        return ret;
+    }
+    return put_manifest_impl(get_parallelism_manifest_key_name(key),
+                             shard_info->manifest, config);
+}
+
+int put_requested_parallelism_shard(const std::string &key,
+                                    pybind11::object tensor,
+                                    const TensorParallelismSpec &parallelism,
+                                    const ReplicateConfig &config) {
+    std::string shard_key = get_parallelism_key_name(key, parallelism);
+    auto shard_info = build_requested_parallelism_shard_info(
+        tensor, parallelism, shard_key, "put_tensor_with_parallelism");
+    if (!shard_info.has_value()) {
+        return to_py_ret(ErrorCode::INVALID_PARAMS);
+    }
+    int ret = put_tensor_info_impl(shard_key, shard_info->info, config);
+    if (ret != 0) {
+        return ret;
+    }
+    return put_manifest_impl(get_parallelism_manifest_key_name(key),
+                             shard_info->manifest, config);
+}
+
+int put_direct_parallelism_shard(const std::string &key,
+                                 pybind11::object tensor,
+                                 const TensorParallelismSpec &parallelism,
+                                 const ReplicateConfig &config) {
+    auto tp_axis_index = find_tp_axis_index(parallelism.axes);
+    if (tp_axis_index.has_value() && parallelism.axes.size() > 1) {
+        return put_requested_parallelism_shard(key, tensor, parallelism,
+                                               config);
+    }
+    auto info =
+        build_direct_parallelism_shard_info(tensor, parallelism.axes, key);
+    if (!info.has_value()) {
+        return to_py_ret(ErrorCode::INVALID_PARAMS);
+    }
+    return put_tensor_info_impl(get_parallelism_key_name(key, parallelism),
+                                *info, config);
+}
+
+int put_tensor_with_writer_shards_from(
+    const std::string &key, uintptr_t buffer_ptr, size_t size,
+    const WriterPartitionSpec &writer,
+    const ReplicateConfig &config = ReplicateConfig{}) {
+    return execute_writer_partition_shard_write_from(
+        key, buffer_ptr, size, writer, config,
+        "put_tensor_with_parallelism_from", TensorWriteStoreOps{"put_parts"},
+        [this](const std::string &shard_key,
+               const std::vector<std::span<const char>> &values,
+               const ReplicateConfig &write_config) {
+            return store_->put_parts(shard_key, values, write_config);
+        },
+        [this](const std::string &manifest_key,
+               const WriterShardManifest &manifest,
+               const ReplicateConfig &write_config) {
+            return put_manifest_impl(manifest_key, manifest, write_config);
+        });
+}
+
+int put_direct_parallelism_shard_from(const std::string &key,
+                                      uintptr_t buffer_ptr, size_t size,
+                                      const TensorParallelismSpec &parallelism,
+                                      const ReplicateConfig &config) {
+    return execute_direct_parallelism_shard_write_from(
+        key, buffer_ptr, size, parallelism, config,
+        "put_tensor_with_parallelism_from", TensorWriteStoreOps{"put_parts"},
+        [this](const std::string &shard_key,
+               const std::vector<std::span<const char>> &values,
+               const ReplicateConfig &write_config) {
+            return store_->put_parts(shard_key, values, write_config);
+        });
+}
+
+int execute_put_tensor_with_parallelism_route(
+    const std::string &key, pybind11::object tensor,
+    const ResolvedParallelismWriteRequest &request,
+    const ReplicateConfig &config) {
+    return execute_parallelism_tensor_write_route(
+        key, tensor, request, config, "put_tensor_with_parallelism",
+        [this](const std::string &write_key, pybind11::object write_tensor,
+               const ReplicateConfig &write_config) {
+            return put_tensor_impl(write_key, write_tensor, write_config);
+        },
+        [this](const std::string &write_key, pybind11::object write_tensor,
+               const WriterPartitionSpec &writer,
+               const ReplicateConfig &write_config) {
+            return put_tensor_with_writer_shards(write_key, write_tensor,
+                                                 writer, write_config);
+        },
+        [this](const std::string &write_key, pybind11::object write_tensor,
+               const ReplicateConfig &write_config, int rank, int size,
+               int split_dim) {
+            return put_tensor_with_tp_impl(write_key, write_tensor,
+                                           write_config, rank, size, split_dim);
+        },
+        [this](const std::string &write_key, pybind11::object write_tensor,
+               const ReplicateConfig &write_config, int rank, int size,
+               int split_dim, const std::vector<ParallelAxisSpec> &axes) {
+            if (axes.size() == 1) {
+                return put_tensor_parallelism_tp_impl(write_key, write_tensor,
+                                                      write_config, rank, size,
+                                                      split_dim, axes);
+            }
+            return put_requested_parallelism_shard(write_key, write_tensor,
+                                                   TensorParallelismSpec{axes},
+                                                   write_config);
+        },
+        [this](const std::string &write_key, pybind11::object write_tensor,
+               const TensorParallelismSpec &parallelism_spec,
+               const ReplicateConfig &write_config) {
+            return put_direct_parallelism_shard(write_key, write_tensor,
+                                                parallelism_spec, write_config);
+        });
+}
+
+int put_tensor_with_parallelism(
+    const std::string &key, pybind11::object tensor,
+    const py::object &parallelism = py::none(),
+    const ReplicateConfig &config = ReplicateConfig{},
+    const py::object &writer_partition = py::none()) {
+    auto request = resolve_parallelism_write_request(
+        parallelism, writer_partition, "put_tensor_with_parallelism");
+    if (!request.has_value()) {
+        return to_py_ret(ErrorCode::INVALID_PARAMS);
+    }
+    return execute_put_tensor_with_parallelism_route(key, tensor, *request,
+                                                     config);
+}
+
+std::vector<int> batch_put_tensor_with_parallelism(
+    const std::vector<std::string> &keys, const pybind11::list &tensors_list,
+    const py::object &parallelisms = py::none(),
+    const ReplicateConfig &config = ReplicateConfig{},
+    const py::object &writer_partitions = py::none()) {
+    auto group_ids_error = ValidateGroupIdsForBatchConfig(
+        config, keys.size(), "batch_put_tensor_with_parallelism");
+    if (!group_ids_error.empty()) {
+        return group_ids_error;
+    }
+
+    return execute_batch_parallelism_write_requests(
+        keys, tensors_list.size(), parallelisms, writer_partitions,
+        "batch_put_tensor_with_parallelism",
+        [this, &keys, &tensors_list, &config]() {
+            if (!is_client_initialized() || use_dummy_client_) {
+                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(), validate_result);
+            }
+            return batch_put_tensor_impl(keys, tensors_list, config);
+        },
+        [this, &keys, &tensors_list, &config](size_t i,
+                                              const py::handle &parallelism) {
+            ReplicateConfig key_config = config.ForSingleKey(i);
+            return put_tensor_with_parallelism(
+                keys[i], tensors_list[i],
+                py::reinterpret_borrow<py::object>(parallelism), key_config);
+        },
+        [this, &keys, &tensors_list, &config](
+            size_t i, const py::handle &writer_partition) {
+            ReplicateConfig key_config = config.ForSingleKey(i);
+            return put_tensor_with_parallelism(
+                keys[i], tensors_list[i], py::none(), key_config,
+                py::reinterpret_borrow<py::object>(writer_partition));
+        });
+}
+
+bool is_valid_tensor_object_buffer(uintptr_t buffer_ptr, size_t size,
+                                   const std::string &op_name) {
+    return parse_tensor_metadata_from_raw_buffer(buffer_ptr, size,
+                                                 op_name.c_str())
+        .has_value();
+}
+
+int execute_put_tensor_with_parallelism_from_route(
+    const std::string &key, uintptr_t buffer_ptr, size_t size,
+    const ResolvedParallelismWriteRequest &request,
+    const ReplicateConfig &config) {
+    return execute_parallelism_tensor_write_from_route(
+        key, buffer_ptr, size, request, config,
+        "put_tensor_with_parallelism_from",
+        [this](const std::string &write_key, uintptr_t write_buffer_ptr,
+               size_t write_size, const ReplicateConfig &write_config) {
+            if (!is_default_replicate_config(write_config)) {
+                if (!is_valid_tensor_object_buffer(
+                        write_buffer_ptr, write_size,
+                        "put_tensor_with_parallelism_from")) {
+                    return to_py_ret(ErrorCode::INVALID_PARAMS);
+                }
+                if (!ensure_tensor_write_supported(
+                        "put_tensor_with_parallelism_from")) {
+                    return to_py_ret(ErrorCode::INVALID_PARAMS);
+                }
+                int validate_result = validate_replicate_config(write_config);
+                if (validate_result) return validate_result;
+                py::gil_scoped_release release_gil;
+                return store_->put_from(
+                    write_key, reinterpret_cast<void *>(write_buffer_ptr),
+                    write_size, write_config);
+            }
+            return put_tensor_from(write_key, write_buffer_ptr, write_size);
+        },
+        [this](const std::string &write_key, uintptr_t write_buffer_ptr,
+               size_t write_size, const WriterPartitionSpec &writer,
+               const ReplicateConfig &write_config) {
+            return put_tensor_with_writer_shards_from(
+                write_key, write_buffer_ptr, write_size, writer, write_config);
+        },
+        [this](const std::string &write_key, uintptr_t write_buffer_ptr,
+               size_t write_size, int tp_size, int split_dim,
+               const std::vector<ParallelAxisSpec> &axes,
+               const ReplicateConfig &write_config) {
+            return execute_tp_tensor_write_from_impl(
+                write_key, write_buffer_ptr, write_size, write_config, tp_size,
+                split_dim, axes, "put_tensor_with_parallelism_from",
+                [this](const std::string &shard_key,
+                       const std::vector<std::span<const char>> &values,
+                       const ReplicateConfig &parts_config) {
+                    py::gil_scoped_release release_gil;
+                    int ret =
+                        store_->put_parts(shard_key, values, parts_config);
+                    if (ret != 0) {
+                        LOG(ERROR) << "put_parts failed for key " << shard_key
+                                   << " with code " << ret;
+                    }
+                    return ret;
+                });
+        },
+        [this](const std::string &write_key, uintptr_t write_buffer_ptr,
+               size_t write_size, const TensorParallelismSpec &parallelism_spec,
+               const ReplicateConfig &write_config) {
+            return put_direct_parallelism_shard_from(
+                write_key, write_buffer_ptr, write_size, parallelism_spec,
+                write_config);
+        },
+        [this](const std::string &write_key, pybind11::object write_tensor,
+               const ResolvedParallelismWriteRequest &write_request,
+               const ReplicateConfig &write_config) {
+            return execute_put_tensor_with_parallelism_route(
+                write_key, write_tensor, write_request, write_config);
+        });
+}
+
+int put_tensor_with_parallelism_from(
+    const std::string &key, uintptr_t buffer_ptr, size_t size,
+    const py::object &parallelism = py::none(),
+    const ReplicateConfig &config = ReplicateConfig{},
+    const py::object &writer_partition = py::none()) {
+    auto request = resolve_parallelism_write_request(
+        parallelism, writer_partition, "put_tensor_with_parallelism_from");
+    if (!request.has_value()) {
+        return to_py_ret(ErrorCode::INVALID_PARAMS);
+    }
+    return execute_put_tensor_with_parallelism_from_route(key, buffer_ptr, size,
+                                                          *request, config);
+}
+
+std::vector<int> batch_put_tensor_with_parallelism_from(
+    const std::vector<std::string> &keys,
+    const std::vector<uintptr_t> &buffer_ptrs, const std::vector<size_t> &sizes,
+    const py::object &parallelisms = py::none(),
+    const ReplicateConfig &config = ReplicateConfig{},
+    const py::object &writer_partitions = py::none()) {
+    auto group_ids_error = ValidateGroupIdsForBatchConfig(
+        config, keys.size(), "batch_put_tensor_with_parallelism_from");
+    if (!group_ids_error.empty()) {
+        return group_ids_error;
+    }
+
+    return execute_batch_parallelism_write_requests(
+        keys, buffer_ptrs.size(), parallelisms, writer_partitions,
+        "batch_put_tensor_with_parallelism_from",
+        [this, &keys, &buffer_ptrs, &sizes, &config]() {
+            if (!is_default_replicate_config(config)) {
+                if (!is_client_initialized() || use_dummy_client_) {
+                    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(), validate_result);
+                }
+                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 (!is_valid_tensor_object_buffer(
+                            buffer_ptrs[i], sizes[i],
+                            "tensor object buffer at index " +
+                                std::to_string(i))) {
+                        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, config);
+            }
+            return batch_put_tensor_from(keys, buffer_ptrs, sizes);
+        },
+        [this, &keys, &buffer_ptrs, &sizes, &config](
+            size_t i, const py::handle &parallelism) {
+            ReplicateConfig key_config = config.ForSingleKey(i);
+            return put_tensor_with_parallelism_from(
+                keys[i], buffer_ptrs[i], sizes[i],
+                py::reinterpret_borrow<py::object>(parallelism), key_config);
+        },
+        [this, &keys, &buffer_ptrs, &sizes, &config](
+            size_t i, const py::handle &writer_partition) {
+            ReplicateConfig key_config = config.ForSingleKey(i);
+            return put_tensor_with_parallelism_from(
+                keys[i], buffer_ptrs[i], sizes[i], py::none(), key_config,
+                py::reinterpret_borrow<py::object>(writer_partition));
+        });
+}
+
+int upsert_tensor_with_writer_shards(
+    const std::string &key, pybind11::object tensor,
+    const WriterPartitionSpec &writer,
+    const ReplicateConfig &config = ReplicateConfig{}) {
+    if (!ensure_tensor_write_supported("upsert_tensor_with_writer_shards")) {
+        return to_py_ret(ErrorCode::INVALID_PARAMS);
+    }
+
+    auto writer_info = build_writer_shard_tensor_info(
+        key, tensor, writer, "upsert_tensor_with_parallelism");
+    if (!writer_info.has_value()) {
+        return to_py_ret(ErrorCode::INVALID_PARAMS);
+    }
+
+    int ret = upsert_tensor_info_impl(writer_info->shard_key, writer_info->info,
+                                      config);
+    if (ret != 0) {
+        return ret;
+    }
+    return upsert_manifest_impl(get_writer_manifest_key_name(key),
+                                writer_info->manifest, config);
+}
+
+int upsert_tensor_parallelism_tp_impl(
+    const std::string &key, pybind11::object tensor,
+    const ReplicateConfig &config, int tp_rank, int tp_size, int split_dim,
+    const std::vector<ParallelAxisSpec> &axes) {
+    std::string shard_key = get_tp_write_shard_key(key, tp_rank, axes);
+    if (shard_key.empty()) {
+        return to_py_ret(ErrorCode::INVALID_PARAMS);
+    }
+    auto shard_info = build_direct_parallelism_shard_info(
+        tensor, axes, shard_key, true /* infer_global_shape */);
+    if (!shard_info.has_value()) {
+        return to_py_ret(ErrorCode::INVALID_PARAMS);
+    }
+    int ret = upsert_tensor_info_impl(shard_key, shard_info->info, config);
+    if (ret != 0) {
+        return ret;
+    }
+    return upsert_manifest_impl(get_parallelism_manifest_key_name(key),
+                                shard_info->manifest, config);
+}
+
+int upsert_requested_parallelism_shard(const std::string &key,
+                                       pybind11::object tensor,
+                                       const TensorParallelismSpec &parallelism,
+                                       const ReplicateConfig &config) {
+    std::string shard_key = get_parallelism_key_name(key, parallelism);
+    auto shard_info = build_requested_parallelism_shard_info(
+        tensor, parallelism, shard_key, "upsert_tensor_with_parallelism");
+    if (!shard_info.has_value()) {
+        return to_py_ret(ErrorCode::INVALID_PARAMS);
+    }
+    int ret = upsert_tensor_info_impl(shard_key, shard_info->info, config);
+    if (ret != 0) {
+        return ret;
+    }
+    return upsert_manifest_impl(get_parallelism_manifest_key_name(key),
+                                shard_info->manifest, config);
+}
+
+int upsert_direct_parallelism_shard(const std::string &key,
+                                    pybind11::object tensor,
+                                    const TensorParallelismSpec &parallelism,
+                                    const ReplicateConfig &config) {
+    auto tp_axis_index = find_tp_axis_index(parallelism.axes);
+    if (tp_axis_index.has_value() && parallelism.axes.size() > 1) {
+        return upsert_requested_parallelism_shard(key, tensor, parallelism,
+                                                  config);
+    }
+    auto info =
+        build_direct_parallelism_shard_info(tensor, parallelism.axes, key);
+    if (!info.has_value()) {
+        return to_py_ret(ErrorCode::INVALID_PARAMS);
+    }
+    return upsert_tensor_info_impl(get_parallelism_key_name(key, parallelism),
+                                   *info, config);
+}
+
+int upsert_tensor_with_writer_shards_from(
+    const std::string &key, uintptr_t buffer_ptr, size_t size,
+    const WriterPartitionSpec &writer,
+    const ReplicateConfig &config = ReplicateConfig{}) {
+    return execute_writer_partition_shard_write_from(
+        key, buffer_ptr, size, writer, config,
+        "upsert_tensor_with_parallelism_from",
+        TensorWriteStoreOps{"upsert_parts"},
+        [this](const std::string &shard_key,
+               const std::vector<std::span<const char>> &values,
+               const ReplicateConfig &write_config) {
+            return store_->upsert_parts(shard_key, values, write_config);
+        },
+        [this](const std::string &manifest_key,
+               const WriterShardManifest &manifest,
+               const ReplicateConfig &write_config) {
+            return upsert_manifest_impl(manifest_key, manifest, write_config);
+        });
+}
+
+int upsert_direct_parallelism_shard_from(
+    const std::string &key, uintptr_t buffer_ptr, size_t size,
+    const TensorParallelismSpec &parallelism, const ReplicateConfig &config) {
+    return execute_direct_parallelism_shard_write_from(
+        key, buffer_ptr, size, parallelism, config,
+        "upsert_tensor_with_parallelism_from",
+        TensorWriteStoreOps{"upsert_parts"},
+        [this](const std::string &shard_key,
+               const std::vector<std::span<const char>> &values,
+               const ReplicateConfig &write_config) {
+            return store_->upsert_parts(shard_key, values, write_config);
+        });
+}
+
+int execute_upsert_tensor_with_parallelism_route(
+    const std::string &key, pybind11::object tensor,
+    const ResolvedParallelismWriteRequest &request,
+    const ReplicateConfig &config) {
+    return execute_parallelism_tensor_write_route(
+        key, tensor, request, config, "upsert_tensor_with_parallelism",
+        [this](const std::string &write_key, pybind11::object write_tensor,
+               const ReplicateConfig &write_config) {
+            return upsert_tensor_impl(write_key, write_tensor, write_config);
+        },
+        [this](const std::string &write_key, pybind11::object write_tensor,
+               const WriterPartitionSpec &writer,
+               const ReplicateConfig &write_config) {
+            return upsert_tensor_with_writer_shards(write_key, write_tensor,
+                                                    writer, write_config);
+        },
+        [this](const std::string &write_key, pybind11::object write_tensor,
+               const ReplicateConfig &write_config, int rank, int size,
+               int split_dim) {
+            return upsert_tensor_with_tp_impl(write_key, write_tensor,
+                                              write_config, size, split_dim);
+        },
+        [this](const std::string &write_key, pybind11::object write_tensor,
+               const ReplicateConfig &write_config, int rank, int size,
+               int split_dim, const std::vector<ParallelAxisSpec> &axes) {
+            if (axes.size() == 1) {
+                return upsert_tensor_parallelism_tp_impl(
+                    write_key, write_tensor, write_config, rank, size,
+                    split_dim, axes);
+            }
+            return upsert_requested_parallelism_shard(
+                write_key, write_tensor, TensorParallelismSpec{axes},
+                write_config);
+        },
+        [this](const std::string &write_key, pybind11::object write_tensor,
+               const TensorParallelismSpec &parallelism_spec,
+               const ReplicateConfig &write_config) {
+            return upsert_direct_parallelism_shard(
+                write_key, write_tensor, parallelism_spec, write_config);
+        });
+}
+
+int upsert_tensor_with_parallelism(
+    const std::string &key, pybind11::object tensor,
+    const py::object &parallelism = py::none(),
+    const ReplicateConfig &config = ReplicateConfig{},
+    const py::object &writer_partition = py::none()) {
+    auto request = resolve_parallelism_write_request(
+        parallelism, writer_partition, "upsert_tensor_with_parallelism");
+    if (!request.has_value()) {
+        return to_py_ret(ErrorCode::INVALID_PARAMS);
+    }
+    return execute_upsert_tensor_with_parallelism_route(key, tensor, *request,
+                                                        config);
+}
+
+int execute_upsert_tensor_with_parallelism_from_route(
+    const std::string &key, uintptr_t buffer_ptr, size_t size,
+    const ResolvedParallelismWriteRequest &request,
+    const ReplicateConfig &config) {
+    return execute_parallelism_tensor_write_from_route(
+        key, buffer_ptr, size, request, config,
+        "upsert_tensor_with_parallelism_from",
+        [this](const std::string &write_key, uintptr_t write_buffer_ptr,
+               size_t write_size, const ReplicateConfig &write_config) {
+            if (!is_default_replicate_config(write_config)) {
+                if (!is_valid_tensor_object_buffer(
+                        write_buffer_ptr, write_size,
+                        "upsert_tensor_with_parallelism_from")) {
+                    return to_py_ret(ErrorCode::INVALID_PARAMS);
+                }
+                void *buffer = reinterpret_cast<void *>(write_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_with_parallelism_from is not "
+                                  "supported for dummy client";
+                    return to_py_ret(ErrorCode::INVALID_PARAMS);
+                }
+                int validate_result = validate_replicate_config(write_config);
+                if (validate_result) return validate_result;
+                py::gil_scoped_release release_gil;
+                return store_->upsert_from(write_key, buffer, write_size,
+                                           write_config);
+            }
+            return upsert_tensor_from(write_key, write_buffer_ptr, write_size);
+        },
+        [this](const std::string &write_key, uintptr_t write_buffer_ptr,
+               size_t write_size, const WriterPartitionSpec &writer,
+               const ReplicateConfig &write_config) {
+            return upsert_tensor_with_writer_shards_from(
+                write_key, write_buffer_ptr, write_size, writer, write_config);
+        },
+        [this](const std::string &write_key, uintptr_t write_buffer_ptr,
+               size_t write_size, int tp_size, int split_dim,
+               const std::vector<ParallelAxisSpec> &axes,
+               const ReplicateConfig &write_config) {
+            return execute_tp_tensor_write_from_impl(
+                write_key, write_buffer_ptr, write_size, write_config, tp_size,
+                split_dim, axes, "upsert_tensor_with_parallelism_from",
+                [this](const std::string &shard_key,
+                       const std::vector<std::span<const char>> &values,
+                       const ReplicateConfig &parts_config) {
+                    py::gil_scoped_release release_gil;
+                    int ret =
+                        store_->upsert_parts(shard_key, values, parts_config);
+                    if (ret != 0) {
+                        LOG(ERROR) << "upsert_parts failed for key "
+                                   << shard_key << " with code " << ret;
+                    }
+                    return ret;
+                });
+        },
+        [this](const std::string &write_key, uintptr_t write_buffer_ptr,
+               size_t write_size, const TensorParallelismSpec &parallelism_spec,
+               const ReplicateConfig &write_config) {
+            return upsert_direct_parallelism_shard_from(
+                write_key, write_buffer_ptr, write_size, parallelism_spec,
+                write_config);
+        },
+        [this](const std::string &write_key, pybind11::object write_tensor,
+               const ResolvedParallelismWriteRequest &write_request,
+               const ReplicateConfig &write_config) {
+            return execute_upsert_tensor_with_parallelism_route(
+                write_key, write_tensor, write_request, write_config);
+        });
+}
+
+int upsert_tensor_with_parallelism_from(
+    const std::string &key, uintptr_t buffer_ptr, size_t size,
+    const py::object &parallelism = py::none(),
+    const ReplicateConfig &config = ReplicateConfig{},
+    const py::object &writer_partition = py::none()) {
+    auto request = resolve_parallelism_write_request(
+        parallelism, writer_partition, "upsert_tensor_with_parallelism_from");
+    if (!request.has_value()) {
+        return to_py_ret(ErrorCode::INVALID_PARAMS);
+    }
+    return execute_upsert_tensor_with_parallelism_from_route(
+        key, buffer_ptr, size, *request, config);
+}
+
+std::vector<int> batch_upsert_tensor_with_parallelism(
+    const std::vector<std::string> &keys, const pybind11::list &tensors_list,
+    const py::object &parallelisms = py::none(),
+    const ReplicateConfig &config = ReplicateConfig{},
+    const py::object &writer_partitions = py::none()) {
+    auto group_ids_error = ValidateGroupIdsForBatchConfig(
+        config, keys.size(), "batch_upsert_tensor_with_parallelism");
+    if (!group_ids_error.empty()) {
+        return group_ids_error;
+    }
+
+    return execute_batch_parallelism_write_requests(
+        keys, tensors_list.size(), parallelisms, writer_partitions,
+        "batch_upsert_tensor_with_parallelism",
+        [this, &keys, &tensors_list, &config]() {
+            if (!is_client_initialized() || use_dummy_client_) {
+                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(), validate_result);
+            }
+            return batch_upsert_tensor_impl(keys, tensors_list, config);
+        },
+        [this, &keys, &tensors_list, &config](size_t i,
+                                              const py::handle &parallelism) {
+            ReplicateConfig key_config = config.ForSingleKey(i);
+            return upsert_tensor_with_parallelism(
+                keys[i], tensors_list[i],
+                py::reinterpret_borrow<py::object>(parallelism), key_config);
+        },
+        [this, &keys, &tensors_list, &config](
+            size_t i, const py::handle &writer_partition) {
+            ReplicateConfig key_config = config.ForSingleKey(i);
+            return upsert_tensor_with_parallelism(
+                keys[i], tensors_list[i], py::none(), key_config,
+                py::reinterpret_borrow<py::object>(writer_partition));
+        });
+}
+
+std::vector<int> batch_upsert_tensor_with_parallelism_from(
+    const std::vector<std::string> &keys,
+    const std::vector<uintptr_t> &buffer_ptrs, const std::vector<size_t> &sizes,
+    const py::object &parallelisms = py::none(),
+    const ReplicateConfig &config = ReplicateConfig{},
+    const py::object &writer_partitions = py::none()) {
+    auto group_ids_error = ValidateGroupIdsForBatchConfig(
+        config, keys.size(), "batch_upsert_tensor_with_parallelism_from");
+    if (!group_ids_error.empty()) {
+        return group_ids_error;
+    }
+
+    return execute_batch_parallelism_write_requests(
+        keys, buffer_ptrs.size(), parallelisms, writer_partitions,
+        "batch_upsert_tensor_with_parallelism_from",
+        [this, &keys, &buffer_ptrs, &sizes, &config]() {
+            if (!is_default_replicate_config(config)) {
+                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_with_parallelism_from "
+                                  "is not supported for dummy client";
+                    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(), validate_result);
+                }
+                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 (!is_valid_tensor_object_buffer(
+                            buffer_ptrs[i], sizes[i],
+                            "tensor object buffer at index " +
+                                std::to_string(i))) {
+                        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, config);
+            }
+            return batch_upsert_tensor_from(keys, buffer_ptrs, sizes);
+        },
+        [this, &keys, &buffer_ptrs, &sizes, &config](
+            size_t i, const py::handle &parallelism) {
+            ReplicateConfig key_config = config.ForSingleKey(i);
+            return upsert_tensor_with_parallelism_from(
+                keys[i], buffer_ptrs[i], sizes[i],
+                py::reinterpret_borrow<py::object>(parallelism), key_config);
+        },
+        [this, &keys, &buffer_ptrs, &sizes, &config](
+            size_t i, const py::handle &writer_partition) {
+            ReplicateConfig key_config = config.ForSingleKey(i);
+            return upsert_tensor_with_parallelism_from(
+                keys[i], buffer_ptrs[i], sizes[i], py::none(), key_config,
+                py::reinterpret_borrow<py::object>(writer_partition));
+        });
+}
diff --git a/mooncake-integration/transfer_engine/transfer_engine_py.cpp b/mooncake-integration/transfer_engine/transfer_engine_py.cpp
index 02bcab91..58eebe33 100644
--- a/mooncake-integration/transfer_engine/transfer_engine_py.cpp
+++ b/mooncake-integration/transfer_engine/transfer_engine_py.cpp
@@ -21,6 +21,14 @@
 #include <pybind11/stl.h>
 #include "transport/rpc_communicator/rpc_interface.h"
 
+#ifdef USE_TENT
+#include "tent/runtime/transport_selector.h"
+#endif
+
+#ifdef USE_EFA
+#include "transport/efa_transport/efa_transport.h"
+#endif
+
 #ifdef USE_HIP
 #include "transport/hip_transport/hip_transport.h"
 #endif
@@ -199,6 +207,16 @@ int TransferEnginePy::initializeExt(const char* local_hostname,
                   << engine_->getLocalTopology()->getHcaList().size()
                   << " devices.";
     }
+#elif defined(USE_CXI)
+
+    bool use_cxi = (proto == "cxi");
+    engine_ = std::make_unique<TransferEngine>(false, device_filter);
+    if (use_cxi) {
+        engine_->getLocalTopology()->discover(device_filter);
+        LOG(INFO) << "Topology discovery complete for CXI. Found "
+                  << engine_->getLocalTopology()->getHcaList().size()
+                  << " devices.";
+    }
 #else
     engine_ = std::make_unique<TransferEngine>(true, device_filter);
 #endif
@@ -239,6 +257,29 @@ int TransferEnginePy::initializeExt(const char* local_hostname,
         }
         LOG(INFO) << "TCP transport installed successfully";
     }
+#elif defined(USE_CXI)
+    if (use_cxi) {
+        LOG(INFO)
+            << "Installing CXI transport as requested by protocol parameter";
+        auto transport = engine_->installTransport("cxi", nullptr);
+        if (!transport) {
+            LOG(ERROR) << "Failed to install CXI transport";
+            return -1;
+        }
+        LOG(INFO) << "CXI transport installed successfully";
+    } else {
+        // For non-EFA protocols (e.g. TCP), manually install TCP transport
+        // since auto_discover is disabled to prevent RDMA installation
+        // (RDMA QP creation fails on EFA devices).
+        LOG(INFO)
+            << "Installing TCP transport (auto_discover disabled in CXI build)";
+        auto transport = engine_->installTransport("tcp", nullptr);
+        if (!transport) {
+            LOG(ERROR) << "Failed to install TCP transport";
+            return -1;
+        }
+        LOG(INFO) << "TCP transport installed successfully";
+    }
 #endif
 
     free_list_.resize(kSlabSizeKBTabLen);
@@ -268,6 +309,7 @@ int TransferEnginePy::findClassId(size_t size) {
 int TransferEnginePy::doBuddyAllocate(int class_id) {
     if (class_id == kMaxClassId) {
         auto buffer = allocateRawBuffer(kDefaultBufferCapacity);
+        if (!buffer) return -1;
         buffer_list_.push_back(buffer);
         for (size_t offset = 0; offset < kDefaultBufferCapacity;
              offset += 1024ull * kSlabSizeKB[kMaxClassId])
@@ -316,57 +358,80 @@ int TransferEnginePy::freeManagedBuffer(uintptr_t buffer_addr, size_t length) {
     return 0;
 }
 
+static int parseTransportHint(const std::string& name) {
+#ifdef USE_TENT
+    if (name.empty()) return mooncake::tent::UNSPEC;
+    auto type = mooncake::tent::TransportSelector::parseTransportType(name);
+    if (type == mooncake::tent::UNSPEC && name != "unspec") {
+        throw std::invalid_argument(
+            "Unknown transport_hint '" + name +
+            "' (valid: rdma, mnnvl, shm, nvlink, gds, io_uring, tcp, "
+            "ascend, sunrise_link, unspec, or empty string for no hint)");
+    }
+    return type;
+#else
+    return 0;
+#endif
+}
+
 int TransferEnginePy::transferSyncWrite(const char* target_hostname,
                                         uintptr_t buffer,
                                         uintptr_t peer_buffer_address,
-                                        size_t length) {
+                                        size_t length,
+                                        const std::string& transport_hint) {
     return transferSync(target_hostname, buffer, peer_buffer_address, length,
-                        TransferOpcode::WRITE);
+                        TransferOpcode::WRITE, nullptr, transport_hint);
 }
 
 int TransferEnginePy::transferSyncRead(const char* target_hostname,
                                        uintptr_t buffer,
                                        uintptr_t peer_buffer_address,
-                                       size_t length) {
+                                       size_t length,
+                                       const std::string& transport_hint) {
     return transferSync(target_hostname, buffer, peer_buffer_address, length,
-                        TransferOpcode::READ);
+                        TransferOpcode::READ, nullptr, transport_hint);
 }
 
 int TransferEnginePy::batchTransferSyncWrite(
     const char* target_hostname, std::vector<uintptr_t> buffers,
-    std::vector<uintptr_t> peer_buffer_addresses, std::vector<size_t> lengths) {
+    std::vector<uintptr_t> peer_buffer_addresses, std::vector<size_t> lengths,
+    const std::string& transport_hint) {
     return batchTransferSync(target_hostname, buffers, peer_buffer_addresses,
-                             lengths, TransferOpcode::WRITE);
+                             lengths, TransferOpcode::WRITE, nullptr,
+                             transport_hint);
 }
 
 int TransferEnginePy::batchTransferSyncRead(
     const char* target_hostname, std::vector<uintptr_t> buffers,
-    std::vector<uintptr_t> peer_buffer_addresses, std::vector<size_t> lengths) {
+    std::vector<uintptr_t> peer_buffer_addresses, std::vector<size_t> lengths,
+    const std::string& transport_hint) {
     return batchTransferSync(target_hostname, buffers, peer_buffer_addresses,
-                             lengths, TransferOpcode::READ);
+                             lengths, TransferOpcode::READ, nullptr,
+                             transport_hint);
 }
 
 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 std::vector<size_t>& lengths, const std::string& transport_hint) {
     return batchTransferAsync(target_hostname, buffers, peer_buffer_addresses,
-                              lengths, TransferOpcode::WRITE);
+                              lengths, TransferOpcode::WRITE, transport_hint);
 }
 
 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 std::vector<size_t>& lengths, const std::string& transport_hint) {
     return batchTransferAsync(target_hostname, buffers, peer_buffer_addresses,
-                              lengths, TransferOpcode::READ);
+                              lengths, TransferOpcode::READ, transport_hint);
 }
 
 int TransferEnginePy::transferSync(const char* target_hostname,
                                    uintptr_t buffer,
                                    uintptr_t peer_buffer_address, size_t length,
                                    TransferOpcode opcode,
-                                   TransferNotify* notify) {
+                                   TransferNotify* notify,
+                                   const std::string& transport_hint) {
     pybind11::gil_scoped_release release;
     Transport::SegmentHandle handle;
     {
@@ -404,6 +469,7 @@ int TransferEnginePy::transferSync(const char* target_hostname,
         entry.target_id = handle;
         entry.target_offset = peer_buffer_address;
         entry.advise_retry_cnt = retry;
+        entry.transport_hint = parseTransportHint(transport_hint);
 
         Status s =
             notify
@@ -458,7 +524,8 @@ int TransferEnginePy::transferSync(const char* target_hostname,
 int TransferEnginePy::batchTransferSync(
     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,
+    const std::string& transport_hint) {
     pybind11::gil_scoped_release release;
     Transport::SegmentHandle handle;
     {
@@ -496,6 +563,7 @@ int TransferEnginePy::batchTransferSync(
         entry.target_id = handle;
         entry.target_offset = peer_buffer_addresses[i];
         entry.advise_retry_cnt = 0;
+        entry.transport_hint = parseTransportHint(transport_hint);
         entries.push_back(entry);
     }
 
@@ -561,7 +629,8 @@ 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 std::vector<size_t>& lengths, TransferOpcode opcode,
+    const std::string& transport_hint) {
     pybind11::gil_scoped_release release;
     Transport::SegmentHandle handle;
     {
@@ -570,7 +639,9 @@ batch_id_t TransferEnginePy::batchTransferAsync(
             handle = handle_map_[target_hostname];
         } else {
             handle = engine_->openSegment(target_hostname);
-            if (handle == (Transport::SegmentHandle)-1) return -1;
+            // batch_id_t is unsigned; 0 is the documented failure sentinel
+            // (-1 would surface as 2^64 - 1 in Python).
+            if (handle == (Transport::SegmentHandle)-1) return 0;
             handle_map_[target_hostname] = handle;
         }
     }
@@ -598,6 +669,7 @@ batch_id_t TransferEnginePy::batchTransferAsync(
         entry.target_id = handle;
         entry.target_offset = peer_buffer_addresses[i];
         entry.advise_retry_cnt = 0;
+        entry.transport_hint = parseTransportHint(transport_hint);
         entries.push_back(entry);
     }
 
@@ -681,10 +753,10 @@ int TransferEnginePy::getBatchTransferStatus(
     return failed_or_timeout ? -1 : 0;
 }
 
-batch_id_t TransferEnginePy::transferSubmitWrite(const char* target_hostname,
-                                                 uintptr_t buffer,
-                                                 uintptr_t peer_buffer_address,
-                                                 size_t length) {
+batch_id_t TransferEnginePy::transferSubmitWrite(
+    const char* target_hostname, uintptr_t buffer,
+    uintptr_t peer_buffer_address, size_t length,
+    const std::string& transport_hint) {
     pybind11::gil_scoped_release release;
     Transport::SegmentHandle handle;
     {
@@ -693,7 +765,9 @@ batch_id_t TransferEnginePy::transferSubmitWrite(const char* target_hostname,
             handle = handle_map_[target_hostname];
         } else {
             handle = engine_->openSegment(target_hostname);
-            if (handle == (Transport::SegmentHandle)-1) return -1;
+            // batch_id_t is unsigned; 0 is the documented failure sentinel
+            // (-1 would surface as 2^64 - 1 in Python).
+            if (handle == (Transport::SegmentHandle)-1) return 0;
             handle_map_[target_hostname] = handle;
         }
     }
@@ -705,9 +779,13 @@ batch_id_t TransferEnginePy::transferSubmitWrite(const char* target_hostname,
     entry.source = (void*)buffer;
     entry.target_id = handle;
     entry.target_offset = peer_buffer_address;
+    entry.transport_hint = parseTransportHint(transport_hint);
 
     Status s = engine_->submitTransfer(batch_id, {entry});
-    if (!s.ok()) return -1;
+    if (!s.ok()) {
+        engine_->freeBatchID(batch_id);
+        return 0;
+    }
 
     return batch_id;
 }
@@ -731,7 +809,8 @@ int TransferEnginePy::transferCheckStatus(batch_id_t batch_id) {
 }
 
 int TransferEnginePy::batchRegisterMemory(
-    std::vector<uintptr_t> buffer_addresses, std::vector<size_t> capacities) {
+    std::vector<uintptr_t> buffer_addresses, std::vector<size_t> capacities,
+    const std::string& location) {
     pybind11::gil_scoped_release release;
     auto batch_size = buffer_addresses.size();
     std::vector<BufferEntry> buffers;
@@ -739,7 +818,7 @@ int TransferEnginePy::batchRegisterMemory(
         buffers.push_back(
             BufferEntry{(void*)buffer_addresses[i], capacities[i]});
     }
-    return engine_->registerLocalMemoryBatch(buffers, kWildcardLocation);
+    return engine_->registerLocalMemoryBatch(buffers, location);
 }
 
 int TransferEnginePy::batchUnregisterMemory(
@@ -753,9 +832,10 @@ int TransferEnginePy::batchUnregisterMemory(
     return engine_->unregisterLocalMemoryBatch(buffers);
 }
 
-int TransferEnginePy::registerMemory(uintptr_t buffer_addr, size_t capacity) {
+int TransferEnginePy::registerMemory(uintptr_t buffer_addr, size_t capacity,
+                                     const std::string& location) {
     char* buffer = reinterpret_cast<char*>(buffer_addr);
-    return engine_->registerLocalMemory(buffer, capacity);
+    return engine_->registerLocalMemory(buffer, capacity, location);
 }
 
 int TransferEnginePy::unregisterMemory(uintptr_t buffer_addr) {
@@ -851,7 +931,7 @@ 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,
-    uintptr_t stream_ptr) {
+    uintptr_t stream_ptr, const std::string& transport_hint) {
     pybind11::gil_scoped_release release;
     Transport::SegmentHandle handle;
     {
@@ -886,6 +966,7 @@ void TransferEnginePy::batchTransferOnCuda(
         entry.source = (void*)buffers[i];
         entry.target_id = handle;
         entry.target_offset = peer_buffer_addresses[i];
+        entry.transport_hint = parseTransportHint(transport_hint);
         entries.push_back(entry);
         total_bytes += lengths[i];
     }
@@ -911,10 +992,11 @@ void TransferEnginePy::batchTransferOnCuda(
 void TransferEnginePy::transferWriteOnCuda(const char* target_hostname,
                                            uintptr_t buffer,
                                            uintptr_t peer_buffer_address,
-                                           size_t length,
-                                           uintptr_t stream_ptr) {
+                                           size_t length, uintptr_t stream_ptr,
+                                           const std::string& transport_hint) {
     batchTransferOnCuda(target_hostname, {buffer}, {peer_buffer_address},
-                        {length}, TransferOpcode::WRITE, stream_ptr);
+                        {length}, TransferOpcode::WRITE, stream_ptr,
+                        transport_hint);
 }
 
 /**
@@ -923,9 +1005,11 @@ void TransferEnginePy::transferWriteOnCuda(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) {
+                                          size_t length, uintptr_t stream_ptr,
+                                          const std::string& transport_hint) {
     batchTransferOnCuda(target_hostname, {buffer}, {peer_buffer_address},
-                        {length}, TransferOpcode::READ, stream_ptr);
+                        {length}, TransferOpcode::READ, stream_ptr,
+                        transport_hint);
 }
 
 /**
@@ -934,9 +1018,11 @@ void TransferEnginePy::transferReadOnCuda(const char* target_hostname,
 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 std::vector<size_t>& lengths, uintptr_t stream_ptr,
+    const std::string& transport_hint) {
     batchTransferOnCuda(target_hostname, buffers, peer_buffer_addresses,
-                        lengths, TransferOpcode::WRITE, stream_ptr);
+                        lengths, TransferOpcode::WRITE, stream_ptr,
+                        transport_hint);
 }
 
 /**
@@ -945,12 +1031,28 @@ void TransferEnginePy::batchTransferWriteOnCuda(
 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 std::vector<size_t>& lengths, uintptr_t stream_ptr,
+    const std::string& transport_hint) {
     batchTransferOnCuda(target_hostname, buffers, peer_buffer_addresses,
-                        lengths, TransferOpcode::READ, stream_ptr);
+                        lengths, TransferOpcode::READ, stream_ptr,
+                        transport_hint);
 }
 #endif
 
+int TransferEnginePy::warmupEfaSegment(const std::string& segment_name) {
+#ifdef USE_EFA
+    pybind11::gil_scoped_release release;
+    auto* t = engine_->getTransport("efa");
+    if (!t) return 0;  // EFA transport not installed; nothing to do.
+    auto* efa = dynamic_cast<EfaTransport*>(t);
+    if (!efa) return 0;
+    return efa->warmupSegment(segment_name);
+#else
+    (void)segment_name;
+    return 0;
+#endif
+}
+
 uintptr_t TransferEnginePy::getFirstBufferAddress(
     const std::string& segment_name) {
     Transport::SegmentHandle segment_id =
@@ -993,6 +1095,12 @@ std::vector<TransferEnginePy::TransferNotify> TransferEnginePy::getNotifies() {
     return result;
 }
 
+int TransferEnginePy::sendProbe(const std::string& peer_server_name) {
+    if (!engine_) return -1;
+    pybind11::gil_scoped_release release;
+    return engine_->getMetadata()->sendProbe(peer_server_name);
+}
+
 namespace py = pybind11;
 
 // Implementation of coro_rpc_interface binding function
@@ -1008,6 +1116,36 @@ void bind_coro_rpc_interface(py::module_& m) {
 }
 
 PYBIND11_MODULE(engine, m) {
+#ifdef USE_EFA
+    m.attr("SUPPORT_EFA") = true;
+#else
+    m.attr("SUPPORT_EFA") = false;
+#endif
+
+#ifdef USE_HIP
+    m.attr("SUPPORT_HIP") = true;
+#else
+    m.attr("SUPPORT_HIP") = false;
+#endif
+
+#ifdef USE_MNNVL
+    m.attr("SUPPORT_MNNVL") = true;
+#else
+    m.attr("SUPPORT_MNNVL") = false;
+#endif
+
+#ifdef USE_INTRA_NVLINK
+    m.attr("SUPPORT_INTRA_NVLINK") = true;
+#else
+    m.attr("SUPPORT_INTRA_NVLINK") = false;
+#endif
+
+#ifdef USE_CUDA
+    m.attr("SUPPORT_CUDA") = true;
+#else
+    m.attr("SUPPORT_CUDA") = false;
+#endif
+
     py::enum_<TransferEnginePy::TransferOpcode> transfer_opcode(
         m, "TransferOpcode", py::arithmetic());
     transfer_opcode.value("Read", TransferEnginePy::TransferOpcode::READ)
@@ -1030,63 +1168,103 @@ PYBIND11_MODULE(engine, m) {
             .def("allocate_managed_buffer",
                  &TransferEnginePy::allocateManagedBuffer)
             .def("free_managed_buffer", &TransferEnginePy::freeManagedBuffer)
-            .def("transfer_sync_write", &TransferEnginePy::transferSyncWrite)
-            .def("transfer_sync_read", &TransferEnginePy::transferSyncRead)
+            .def("transfer_sync_write", &TransferEnginePy::transferSyncWrite,
+                 py::arg("target_hostname"), py::arg("buffer"),
+                 py::arg("peer_buffer_address"), py::arg("length"),
+                 py::arg("transport_hint") = "")
+            .def("transfer_sync_read", &TransferEnginePy::transferSyncRead,
+                 py::arg("target_hostname"), py::arg("buffer"),
+                 py::arg("peer_buffer_address"), py::arg("length"),
+                 py::arg("transport_hint") = "")
             .def("batch_transfer_sync_write",
-                 &TransferEnginePy::batchTransferSyncWrite)
+                 &TransferEnginePy::batchTransferSyncWrite,
+                 py::arg("target_hostname"), py::arg("buffers"),
+                 py::arg("peer_buffer_addresses"), py::arg("lengths"),
+                 py::arg("transport_hint") = "")
             .def("batch_transfer_sync_read",
-                 &TransferEnginePy::batchTransferSyncRead)
+                 &TransferEnginePy::batchTransferSyncRead,
+                 py::arg("target_hostname"), py::arg("buffers"),
+                 py::arg("peer_buffer_addresses"), py::arg("lengths"),
+                 py::arg("transport_hint") = "")
             .def("batch_transfer_async_write",
-                 &TransferEnginePy::batchTransferAsyncWrite)
+                 &TransferEnginePy::batchTransferAsyncWrite,
+                 py::arg("target_hostname"), py::arg("buffers"),
+                 py::arg("peer_buffer_addresses"), py::arg("lengths"),
+                 py::arg("transport_hint") = "")
             .def("batch_transfer_async_read",
-                 &TransferEnginePy::batchTransferAsyncRead)
+                 &TransferEnginePy::batchTransferAsyncRead,
+                 py::arg("target_hostname"), py::arg("buffers"),
+                 py::arg("peer_buffer_addresses"), py::arg("lengths"),
+                 py::arg("transport_hint") = "")
             .def("transfer_sync", &TransferEnginePy::transferSync,
                  py::arg("target_hostname"), py::arg("buffer"),
                  py::arg("peer_buffer_address"), py::arg("length"),
-                 py::arg("opcode"), py::arg("notify") = nullptr)
-            .def("batch_transfer_sync", &TransferEnginePy::batchTransferSync)
-            .def("batch_transfer_async", &TransferEnginePy::batchTransferAsync)
+                 py::arg("opcode"), py::arg("notify") = nullptr,
+                 py::arg("transport_hint") = "")
+            .def("batch_transfer_sync", &TransferEnginePy::batchTransferSync,
+                 py::arg("target_hostname"), py::arg("buffers"),
+                 py::arg("peer_buffer_addresses"), py::arg("lengths"),
+                 py::arg("opcode"), py::arg("notify") = nullptr,
+                 py::arg("transport_hint") = "")
+            .def("batch_transfer_async", &TransferEnginePy::batchTransferAsync,
+                 py::arg("target_hostname"), py::arg("buffers"),
+                 py::arg("peer_buffer_addresses"), py::arg("lengths"),
+                 py::arg("opcode"), py::arg("transport_hint") = "")
 #ifdef USE_CUDA
             .def("transfer_write_on_cuda",
                  &TransferEnginePy::transferWriteOnCuda,
                  py::arg("target_hostname"), py::arg("buffer"),
                  py::arg("peer_buffer_address"), py::arg("length"),
-                 py::arg("stream_ptr") = 0)
+                 py::arg("stream_ptr") = 0, py::arg("transport_hint") = "")
             .def("transfer_read_on_cuda", &TransferEnginePy::transferReadOnCuda,
                  py::arg("target_hostname"), py::arg("buffer"),
                  py::arg("peer_buffer_address"), py::arg("length"),
-                 py::arg("stream_ptr") = 0)
+                 py::arg("stream_ptr") = 0, py::arg("transport_hint") = "")
             .def("batch_transfer_write_on_cuda",
                  &TransferEnginePy::batchTransferWriteOnCuda,
                  py::arg("target_hostname"), py::arg("buffers"),
                  py::arg("peer_buffer_addresses"), py::arg("lengths"),
-                 py::arg("stream_ptr") = 0)
+                 py::arg("stream_ptr") = 0, py::arg("transport_hint") = "")
             .def("batch_transfer_read_on_cuda",
                  &TransferEnginePy::batchTransferReadOnCuda,
                  py::arg("target_hostname"), py::arg("buffers"),
                  py::arg("peer_buffer_addresses"), py::arg("lengths"),
-                 py::arg("stream_ptr") = 0)
+                 py::arg("stream_ptr") = 0, py::arg("transport_hint") = "")
 #endif
             .def("get_batch_transfer_status",
                  &TransferEnginePy::getBatchTransferStatus)
             .def("transfer_submit_write",
-                 &TransferEnginePy::transferSubmitWrite)
+                 &TransferEnginePy::transferSubmitWrite,
+                 py::arg("target_hostname"), py::arg("buffer"),
+                 py::arg("peer_buffer_address"), py::arg("length"),
+                 py::arg("transport_hint") = "")
             .def("transfer_check_status",
                  &TransferEnginePy::transferCheckStatus)
             .def("write_bytes_to_buffer", &TransferEnginePy::writeBytesToBuffer)
             .def("read_bytes_from_buffer",
                  &TransferEnginePy::readBytesFromBuffer)
-            .def("register_memory", &TransferEnginePy::registerMemory)
+            .def("register_memory", &TransferEnginePy::registerMemory,
+                 py::arg("buffer_addr"), py::arg("capacity"),
+                 py::arg("location") = kWildcardLocation)
             .def("unregister_memory", &TransferEnginePy::unregisterMemory)
             .def("batch_register_memory",
-                 &TransferEnginePy::batchRegisterMemory)
+                 &TransferEnginePy::batchRegisterMemory,
+                 py::arg("buffer_addresses"), py::arg("capacities"),
+                 py::arg("location") = kWildcardLocation)
             .def("batch_unregister_memory",
                  &TransferEnginePy::batchUnregisterMemory)
             .def("get_local_topology", &TransferEnginePy::getLocalTopology,
                  py::arg("device_name") = nullptr)
             .def("get_first_buffer_address",
                  &TransferEnginePy::getFirstBufferAddress)
+            .def("warmup_efa_segment", &TransferEnginePy::warmupEfaSegment,
+                 py::arg("segment_name"))
             .def("get_notifies", &TransferEnginePy::getNotifies)
+            .def("send_probe", &TransferEnginePy::sendProbe,
+                 py::arg("peer_server_name"),
+                 "Send a JSON-RPC probe to peer to verify reachability. "
+                 "Returns 0 on success, non-zero on failure. Used by "
+                 "SGLang's failed-session blacklist recovery.")
             .def("get_engine", &TransferEnginePy::getEngine)
             .def("get_engine_ptr", &TransferEnginePy::getEnginePtr);
 
diff --git a/mooncake-integration/transfer_engine/transfer_engine_py.h b/mooncake-integration/transfer_engine/transfer_engine_py.h
index 083ff357..d4af71af 100644
--- a/mooncake-integration/transfer_engine/transfer_engine_py.h
+++ b/mooncake-integration/transfer_engine/transfer_engine_py.h
@@ -70,41 +70,48 @@ class TransferEnginePy {
     int freeManagedBuffer(uintptr_t user_tensor, size_t length);
 
     int transferSyncWrite(const char *target_hostname, uintptr_t buffer,
-                          uintptr_t peer_buffer_address, size_t length);
+                          uintptr_t peer_buffer_address, size_t length,
+                          const std::string &transport_hint = "");
 
     batch_id_t transferSubmitWrite(const char *target_hostname,
                                    uintptr_t buffer,
-                                   uintptr_t peer_buffer_address,
-                                   size_t length);
+                                   uintptr_t peer_buffer_address, size_t length,
+                                   const std::string &transport_hint = "");
 
     int transferCheckStatus(batch_id_t batch_id);
 
     int transferSyncRead(const char *target_hostname, uintptr_t buffer,
-                         uintptr_t peer_buffer_address, size_t length);
+                         uintptr_t peer_buffer_address, size_t length,
+                         const std::string &transport_hint = "");
 
     int batchTransferSyncWrite(const char *target_hostname,
                                std::vector<uintptr_t> buffers,
                                std::vector<uintptr_t> peer_buffer_addresses,
-                               std::vector<size_t> lengths);
+                               std::vector<size_t> lengths,
+                               const std::string &transport_hint = "");
 
     int batchTransferSyncRead(const char *target_hostname,
                               std::vector<uintptr_t> buffers,
                               std::vector<uintptr_t> peer_buffer_addresses,
-                              std::vector<size_t> lengths);
+                              std::vector<size_t> lengths,
+                              const std::string &transport_hint = "");
 
     batch_id_t 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 std::vector<size_t> &lengths,
+        const std::string &transport_hint = "");
 
     batch_id_t 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 std::vector<size_t> &lengths,
+        const std::string &transport_hint = "");
 
     int transferSync(const char *target_hostname, uintptr_t buffer,
                      uintptr_t peer_buffer_address, size_t length,
-                     TransferOpcode opcode, TransferNotify *notify = nullptr);
+                     TransferOpcode opcode, TransferNotify *notify = nullptr,
+                     const std::string &transport_hint = "");
 
     // Known issue: in a few inference engines and benchmarks, accuracy
     // may be affected when using the batchTransferSync API. We currently
@@ -113,12 +120,14 @@ class TransferEnginePy {
                           std::vector<uintptr_t> buffers,
                           std::vector<uintptr_t> peer_buffer_addresses,
                           std::vector<size_t> lengths, TransferOpcode opcode,
-                          TransferNotify *notify = nullptr);
+                          TransferNotify *notify = nullptr,
+                          const std::string &transport_hint = "");
 
     batch_id_t 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 std::vector<size_t> &lengths, TransferOpcode opcode,
+        const std::string &transport_hint = "");
 
     int getBatchTransferStatus(const std::vector<batch_id_t> &batch_ids);
 
@@ -127,29 +136,38 @@ class TransferEnginePy {
         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 = 0);
+        uintptr_t stream_ptr = 0, const std::string &transport_hint = "");
 
     void transferWriteOnCuda(const char *target_hostname, uintptr_t buffer,
                              uintptr_t peer_buffer_address, size_t length,
-                             uintptr_t stream_ptr = 0);
+                             uintptr_t stream_ptr = 0,
+                             const std::string &transport_hint = "");
 
     void transferReadOnCuda(const char *target_hostname, uintptr_t buffer,
                             uintptr_t peer_buffer_address, size_t length,
-                            uintptr_t stream_ptr = 0);
+                            uintptr_t stream_ptr = 0,
+                            const std::string &transport_hint = "");
 
     void 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 = 0);
+        const std::vector<size_t> &lengths, uintptr_t stream_ptr = 0,
+        const std::string &transport_hint = "");
 
     void 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 = 0);
+        const std::vector<size_t> &lengths, uintptr_t stream_ptr = 0,
+        const std::string &transport_hint = "");
 #endif
 
     uintptr_t getFirstBufferAddress(const std::string &segment_name);
 
+    // Pre-connect every (local_ctx, peer_nic) pair for `segment_name` so the
+    // first submitTransfer does not stall on handshake RPC + fi_av_insert.
+    // No-op on non-EFA builds or when the EFA transport is not installed.
+    int warmupEfaSegment(const std::string &segment_name);
+
     int writeBytesToBuffer(uintptr_t dest_address, char *src_ptr,
                            size_t length) {
         memcpy((void *)dest_address, (void *)src_ptr, length);
@@ -164,13 +182,15 @@ class TransferEnginePy {
     }
 
     // FOR EXPERIMENT ONLY
-    int registerMemory(uintptr_t buffer_addr, size_t capacity);
+    int registerMemory(uintptr_t buffer_addr, size_t capacity,
+                       const std::string &location = kWildcardLocation);
 
     // must be called before TransferEnginePy::~TransferEnginePy()
     int unregisterMemory(uintptr_t buffer_addr);
 
     int batchRegisterMemory(std::vector<uintptr_t> buffer_addresses,
-                            std::vector<size_t> capacities);
+                            std::vector<size_t> capacities,
+                            const std::string &location = kWildcardLocation);
 
     int batchUnregisterMemory(std::vector<uintptr_t> buffer_addresses);
 
@@ -178,6 +198,8 @@ class TransferEnginePy {
 
     std::vector<TransferNotify> getNotifies();
 
+    int sendProbe(const std::string &peer_server_name);
+
     std::shared_ptr<TransferEngine> getEngine() const { return engine_; }
 
     uintptr_t getEnginePtr() const { return (uintptr_t)engine_.get(); }
diff --git a/mooncake-p2p-store/build.sh b/mooncake-p2p-store/build.sh
index 9b98a73e..774e6ba6 100644
--- a/mooncake-p2p-store/build.sh
+++ b/mooncake-p2p-store/build.sh
@@ -34,7 +34,8 @@ fi
 EXT_LDFLAGS="-L$BUILD_DIR/mooncake-transfer-engine/src"
 EXT_LDFLAGS+=" -L$BUILD_DIR/mooncake-transfer-engine/src/common/base"
 EXT_LDFLAGS+=" -L$BUILD_DIR/mooncake-common"
-EXT_LDFLAGS+=" -ltransfer_engine -lbase -lasio -lstdc++ -lnuma -lglog -libverbs -ljsoncpp"
+EXT_LDFLAGS+=" -L$BUILD_DIR/mooncake-common/src"
+EXT_LDFLAGS+=" -ltransfer_engine -lbase -lasio -lstdc++ -lnuma -lglog -libverbs -lmlx5 -ljsoncpp -lmooncake_common -lm"
 
 if [ -d "/usr/local/cuda/lib64/stubs" ]; then
     EXT_LDFLAGS+=" -L/usr/local/cuda/lib64/stubs"
@@ -52,6 +53,10 @@ if [ -d "/usr/local/musa/lib" ]; then
     EXT_LDFLAGS+=" -L/usr/local/musa/lib -lmusart"
 fi
 
+if [ -e "/usr/lib64/liburma.so" ]; then
+   EXT_LDFLAGS+=" -L/usr/lib64 -lurma"
+fi
+
 if [ "$USE_ETCD" = "ON" ]; then
     if [ "$USE_ETCD_LEGACY" = "ON" ]; then
         EXT_LDFLAGS+=" -letcd-cpp-api -lprotobuf -lgrpc++ -lgrpc"
diff --git a/mooncake-pg/BuildPgExt.cmake b/mooncake-pg/BuildPgExt.cmake
index 38bfadff..ead67c3e 100644
--- a/mooncake-pg/BuildPgExt.cmake
+++ b/mooncake-pg/BuildPgExt.cmake
@@ -11,6 +11,7 @@
 #   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
+#   EP_USE_MUSA         - set to "1" when building for MUSA (MTLink path)
 
 cmake_minimum_required(VERSION 3.16)
 
@@ -35,6 +36,11 @@ endif()
 set(ENV{MAKEFLAGS} "")
 set(ENV{MFLAGS} "")
 set(ENV{TORCH_CUDA_ARCH_LIST} "${TORCH_CUDA_ARCH_LIST}")
+if(EP_USE_MUSA)
+  set(ENV{MOONCAKE_EP_USE_MUSA} "1")
+else()
+  unset(ENV{MOONCAKE_EP_USE_MUSA})
+endif()
 
 # ---------------------------------------------------------------------------
 # 2. Ensure engine.so exists in mooncake-wheel/mooncake/ for setup.py linking.
diff --git a/mooncake-pg/include/connection_poller.h b/mooncake-pg/include/connection_poller.h
index a386439e..42417b1e 100644
--- a/mooncake-pg/include/connection_poller.h
+++ b/mooncake-pg/include/connection_poller.h
@@ -32,6 +32,12 @@ struct PeerConnection {
     std::optional<BatchID> warmupBatchId{std::nullopt};
     std::optional<TransferMetadata::SegmentID> segmentId{std::nullopt};
 
+    // Whether this peer has been counted in `totalConnectedPeers_`.
+    // Note that ConnectionPoller may establish connections for ranks beyond
+    // the current groupSize_ (when pollingLimit_ > groupSize_). In that case we
+    // delay counting until the rank officially enters the group.
+    bool countedInGroup{false};
+
     // Back off to avoid frequently checking store.
     std::chrono::steady_clock::time_point last_check_store;
     size_t check_store_backoff_ms{kCheckStoreInitialBackoffMs};
@@ -56,6 +62,11 @@ class ConnectionContext {
 
     std::atomic<int> groupSize_;
 
+    // Upper bound of peer polling range.
+    // This can be larger than groupSize_ so that existing ranks can observe
+    // joiners without calling extendGroupSizeTo().
+    std::atomic<int> pollingLimit_;
+
     bool isDummy_;
 
     // A mark tracking the group size for which all ranks
@@ -120,6 +131,9 @@ class ConnectionContext {
      */
     void extendGroupSizeTo(int newGroupSize);
 
+    // Allow polling ranks beyond groupSize_ without changing groupSize_.
+    void setPollingLimitTo(int pollingLimit);
+
     /**
      * @brief Checks whether all peers within the group have
      *        established connections.
@@ -184,14 +198,8 @@ class ConnectionContext {
         return "buffer_" + std::to_string(backendIndex) + "_" +
                std::to_string(rank);
     }
-    static std::string getExtensionTaskCountStoreKey(int backendIndex,
-                                                     int rank) {
-        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) + "_" +
+    static std::string getExtensionStateStoreKey(int backendIndex, int rank) {
+        return "extension_state_" + std::to_string(backendIndex) + "_" +
                std::to_string(rank);
     }
 
diff --git a/mooncake-pg/include/mooncake_backend.h b/mooncake-pg/include/mooncake_backend.h
index d6196ff4..81727175 100644
--- a/mooncake-pg/include/mooncake_backend.h
+++ b/mooncake-pg/include/mooncake_backend.h
@@ -11,24 +11,72 @@
 #include <sys/types.h>
 #include <torch/torch.h>
 #include <torch/csrc/distributed/c10d/Backend.hpp>
+#include <torch/csrc/distributed/c10d/ProcessGroup.hpp>
 #include <transfer_engine.h>
 
+#include <ATen/cuda/CUDAContext.h>
+
 namespace mooncake {
 
-class MooncakeBackend final : public ::c10d::Backend {
+// Forward declaration – MooncakeP2PShim holds a non-owning pointer to
+// MooncakeBackend, which is defined below.
+class MooncakeBackend;
+
+// Lightweight Backend shim that delegates P2P send/recv back to the owning
+// MooncakeBackend.  PyTorch's P2P dispatch (batch_isend_irecv, isend, irecv)
+// requires getBackend() to return a registered c10d::Backend instance.
+// Since MooncakeBackend inherits from ProcessGroup (not Backend), we register
+// this shim in the ProcessGroup's deviceTypeToBackend_ map so that the P2P
+// path can find it.  The shim holds a non-owning pointer to its owner and
+// delegates only the operations that the P2P dispatch path calls (send, recv,
+// getBackendName, supportsCoalescing).
+class MooncakeP2PShim final : public ::c10d::Backend {
    public:
-    struct MooncakeBackendOptions final : ::c10d::Backend::Options {
+    explicit MooncakeP2PShim(MooncakeBackend* owner);
+
+    const std::string getBackendName() const override;
+
+    bool supportsCoalescing() const override { return false; }
+
+    c10::intrusive_ptr<c10d::Work> send(std::vector<at::Tensor>& tensors,
+                                        int dstRank, int tag) override;
+
+    c10::intrusive_ptr<c10d::Work> recv(std::vector<at::Tensor>& tensors,
+                                        int srcRank, int tag) override;
+
+    c10::intrusive_ptr<c10d::Work> recvAnysource(
+        std::vector<at::Tensor>& tensors, int tag) override;
+
+    c10::intrusive_ptr<c10d::Work> barrier(
+        const c10d::BarrierOptions& opts) override;
+
+   private:
+    // Non-owning: the shim is stored in ProcessGroup's backend maps which are
+    // cleared on destruction, and MooncakeBackend always outlives the shim.
+    MooncakeBackend* owner_;
+};
+
+class MooncakeBackend final : public ::c10d::ProcessGroup {
+   public:
+    struct MooncakeBackendOptions final : torch::CustomClassHolder {
         explicit MooncakeBackendOptions(at::Tensor activeRanks)
-            : Options{"mooncake"}, activeRanks_{activeRanks} {}
+            : activeRanks_{activeRanks} {}
         MooncakeBackendOptions(at::Tensor activeRanks, bool isExtension)
-            : Options{"mooncake"},
-              activeRanks_{activeRanks},
-              isExtension_{isExtension} {}
+            : activeRanks_{activeRanks}, isExtension_{isExtension} {}
+        MooncakeBackendOptions(at::Tensor activeRanks, bool isExtension,
+                               int maxWorldSize)
+            : activeRanks_{activeRanks},
+              isExtension_{isExtension},
+              maxWorldSize_{maxWorldSize} {}
 
         ~MooncakeBackendOptions() override = default;
 
         at::Tensor activeRanks_;
         bool isExtension_ = false;
+        // Optional upper bound for connection polling / reserved rank slots.
+        // When > 0, the backend may pre-size internal rank metadata to this
+        // value (while PyTorch's group_size() remains unchanged).
+        int maxWorldSize_ = -1;
     };
 
     /**
@@ -50,20 +98,7 @@ 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_);
-    }
+    int getSize() const override { return meta_ ? meta_->activeSize : size_; }
 
     // Point-to-point send/recv for torch.distributed P2POp/batch_isend_irecv.
     // Only single-tensor ops are supported.
@@ -124,6 +159,13 @@ class MooncakeBackend final : public ::c10d::Backend {
         engine_->setWhitelistFilters(std::move(filters));
     }
 
+    /// Set an external TransferEngine to be used by MooncakeBackend
+    /// instead of creating its own. Must be called before
+    /// init_process_group(backend="mooncake"). The engine must already
+    /// be initialized. The caller is responsible for ensuring the engine
+    /// outlives all MooncakeBackend instances. Pass nullptr to reset.
+    static void setExternalEngine(TransferEngine* engine);
+
     std::string getPreferredHca(std::string location) {
         static std::once_flag topo_once;
         static std::shared_ptr<Topology> topology;
@@ -178,6 +220,11 @@ class MooncakeBackend final : public ::c10d::Backend {
     std::shared_ptr<MooncakeWorker> worker_;
     static bool engineInitialized_;
     static int backendIndex_;
+    // External engine injection: when set, MooncakeBackend uses this engine
+    // instead of the default self-created one. Non-owning pointer.
+    // The caller is responsible for ensuring the engine outlives all
+    // MooncakeBackend instances.
+    static TransferEngine* externalEngine_;
     const c10::intrusive_ptr<MooncakeBackendOptions> options_;
     bool isCpu_{false};
     static std::string hostIp_;
@@ -208,6 +255,14 @@ class MooncakeBackend final : public ::c10d::Backend {
     bool connectionPollerRegistered_{false};
 };
 
+struct ExtensionState {
+    std::vector<bool> activeRanks;
+    std::vector<uint32_t> p2pEpochs;
+    int taskCount = -1;
+};
+std::vector<uint8_t> serialize(const ExtensionState& state);
+ExtensionState deserialize(const std::vector<uint8_t>& buffer);
+
 }  // namespace mooncake
 
 #endif  // MOONCAKE_BACKEND_H
diff --git a/mooncake-pg/include/mooncake_worker.cuh b/mooncake-pg/include/mooncake_worker.cuh
index d7d41511..d8a6c1a5 100644
--- a/mooncake-pg/include/mooncake_worker.cuh
+++ b/mooncake-pg/include/mooncake_worker.cuh
@@ -1,13 +1,20 @@
 #ifndef MOONCAKE_WORKER_CUH
 #define MOONCAKE_WORKER_CUH
 
+#if !defined(__MUSA__)
+#include <ATen/ATen.h>
 #include <ATen/cuda/CUDAContext.h>
-#include <cuda_bf16.h>
-#include <cuda_runtime.h>
-#include <torch/torch.h>
+#include <c10/util/intrusive_ptr.h>
 #include <torch/csrc/distributed/c10d/Types.hpp>
 #include <torch/csrc/distributed/c10d/Work.hpp>
 #include <torch/csrc/distributed/c10d/Store.hpp>
+#else
+// MUSA device compilation: minimal includes to avoid mcc compiler crash
+#include <cstddef>
+#include <cstdint>
+#endif
+
+#include <cuda_alike.h>
 #include <transfer_engine.h>
 
 #include <memory>
@@ -25,31 +32,38 @@ static constexpr size_t kMaxNumRanks = 64;
 struct SegmentInfo {
     uint64_t send_buffer[2], recv_buffer[2], send_sync[2], recv_sync[2],
         warmup_buffer[2];
-    uint64_t p2p_send_buffer;
-    uint64_t p2p_recv_buffer;
-    uint64_t p2p_ctrl_send;
-    uint64_t p2p_ctrl_recv;
+    uint64_t p2p_credit_region;
+    uint64_t p2p_ack_region;
 };
 
 struct TransferGroupMeta {
     int rank;
-    int size;
+    int size;        // capacity: number of slots allocated (incl. inactive)
+    int activeSize;  // visible group size: number of ranks that participate
     int taskCount;
     bool* activeRanks;
     bool* activeRanksDevice;
+#if !defined(__MUSA__)
     at::Tensor activeRanksTensor;
+#endif
     bool peerConnected[kMaxNumRanks]{};
     TransferEngine* engine;
+#if !defined(__MUSA__)
     c10::intrusive_ptr<::c10d::Store> store;
+#endif
     int bufferBaseIndex;
     int backendIndex;
     TransferMetadata::SegmentID segmentIDs[kMaxNumRanks];
     SegmentInfo segmentInfos[kMaxNumRanks];
 };
 
-__global__ struct Task {
+#if defined(__CUDACC__) || defined(__MUSA__)
+__global__
+#endif
+    struct Task {
     volatile bool active = false;
-    c10d::OpType opType = c10d::OpType::UNKNOWN;
+    int opType =
+        0;  // c10d::OpType as int, for ABI compatibility with kernel code
     size_t tensorSize;  // In bytes
     int64_t broadcastRoot;
     int bufferOffset;
@@ -58,6 +72,7 @@ __global__ struct Task {
     void* transferGroupMeta;
 };
 
+#if !defined(__MUSA__)
 void launchReduceKernel(at::Tensor dst, size_t pos, size_t realSize, void* src,
                         size_t numRanks, c10d::ReduceOp op, bool* activeRanks,
                         cudaStream_t stream);
@@ -125,7 +140,7 @@ class MooncakeWorker {
     static constexpr size_t kPingTimeoutMicroseconds_ = 100;
     static constexpr size_t kDrainTasksTimeoutMs = 5000;  // 5s
 
-    bool running_ = false;
+    std::atomic<bool> running_{false};
     std::atomic<bool> started_{false};
     int cuda_device_index_;
 
@@ -160,6 +175,7 @@ class MooncakeWorkerManager {
     // detached threads must not outlive the MooncakeWorker object.
     std::unordered_map<int, std::shared_ptr<MooncakeWorker>> workers_;
 };
+#endif  // !defined(__MUSA__)
 
 }  // namespace mooncake
 
diff --git a/mooncake-pg/include/mooncake_worker_kernels.cuh b/mooncake-pg/include/mooncake_worker_kernels.cuh
new file mode 100644
index 00000000..67b56b9e
--- /dev/null
+++ b/mooncake-pg/include/mooncake_worker_kernels.cuh
@@ -0,0 +1,74 @@
+#ifndef MOONCAKE_WORKER_KERNELS_CUH
+#define MOONCAKE_WORKER_KERNELS_CUH
+
+// Include the main worker header for struct definitions (Task, SegmentInfo,
+// TransferGroupMeta). When compiled by mcc (__MUSA__ defined), the torch-
+// dependent parts are guarded out, making this safe for the MUSA compiler.
+#include <mooncake_worker.cuh>
+
+namespace mooncake {
+
+// Kernel function declarations — guarded so g++ doesn't see __global__
+// which it can't parse. Parameters use plain C++ types (int instead of
+// c10d::OpType / c10d::ReduceOp::RedOpType) so that mcc can compile them
+// without torch headers.
+#if defined(__CUDACC__) || defined(__MUSA__)
+__global__ void enqueueTaskKernel(int opType, size_t tensorSize,
+                                  int64_t broadcastRoot, int bufferOffset,
+                                  uint64_t submitSequence, void* meta,
+                                  Task* tasks, int numRanks,
+                                  const bool* activeRanks,
+                                  int* activeRanksTensor, size_t taskId);
+
+template <typename scalar_t>
+__global__ void reduceKernel(scalar_t* dst, const scalar_t* src,
+                             size_t numElements, size_t numRanks, int op,
+                             bool* activeRanks);
+#endif
+
+// Host-callable kernel launch wrappers (compiled by mcc/nvcc, callable from
+// g++) g++ cannot compile <<<>>> syntax, so these wrappers are compiled by the
+// GPU compiler and provide plain C++ functions that the host code can call.
+// Both CUDA and MUSA use cudaStream_t in the declaration: on MUSA,
+// cuda_alike.h typedefs cudaStream_t to musaStream_t.
+
+void launchEnqueueTaskKernel(int opType, size_t tensorSize,
+                             int64_t broadcastRoot, int bufferOffset,
+                             uint64_t submitSequence, void* meta, Task* tasks,
+                             int numRanks, const bool* activeRanks,
+                             int* activeRanksTensor, size_t taskId,
+                             cudaStream_t stream);
+
+void launchReduceKernel_uint8(uint8_t* dst, const uint8_t* src,
+                              size_t numElements, size_t numRanks, int op,
+                              bool* activeRanks, cudaStream_t stream);
+void launchReduceKernel_int8(int8_t* dst, const int8_t* src, size_t numElements,
+                             size_t numRanks, int op, bool* activeRanks,
+                             cudaStream_t stream);
+void launchReduceKernel_int16(int16_t* dst, const int16_t* src,
+                              size_t numElements, size_t numRanks, int op,
+                              bool* activeRanks, cudaStream_t stream);
+void launchReduceKernel_int32(int* dst, const int* src, size_t numElements,
+                              size_t numRanks, int op, bool* activeRanks,
+                              cudaStream_t stream);
+void launchReduceKernel_int64(int64_t* dst, const int64_t* src,
+                              size_t numElements, size_t numRanks, int op,
+                              bool* activeRanks, cudaStream_t stream);
+void launchReduceKernel_float(float* dst, const float* src, size_t numElements,
+                              size_t numRanks, int op, bool* activeRanks,
+                              cudaStream_t stream);
+void launchReduceKernel_double(double* dst, const double* src,
+                               size_t numElements, size_t numRanks, int op,
+                               bool* activeRanks, cudaStream_t stream);
+void launchReduceKernel_bool(bool* dst, const bool* src, size_t numElements,
+                             size_t numRanks, int op, bool* activeRanks,
+                             cudaStream_t stream);
+void launchReduceKernel_bf16(void* dst, const void* src, size_t numElements,
+                             size_t numRanks, int op, bool* activeRanks,
+                             cudaStream_t stream);
+
+void preloadReduceKernels();
+
+}  // namespace mooncake
+
+#endif  // MOONCAKE_WORKER_KERNELS_CUH
diff --git a/mooncake-pg/include/p2p_proxy.h b/mooncake-pg/include/p2p_proxy.h
index 168ebffb..69f508a7 100644
--- a/mooncake-pg/include/p2p_proxy.h
+++ b/mooncake-pg/include/p2p_proxy.h
@@ -1,5 +1,5 @@
-#ifndef MOONCAKE_P2P_PROXY_HH
-#define MOONCAKE_P2P_PROXY_HH
+#ifndef MOONCAKE_P2P_PROXY_H
+#define MOONCAKE_P2P_PROXY_H
 
 #include <mooncake_worker.cuh>
 #include <torch/torch.h>
@@ -9,51 +9,262 @@
 #include <cstddef>
 #include <cstdint>
 #include <deque>
+#include <limits>
 #include <memory>
 #include <mutex>
 #include <optional>
 #include <queue>
 #include <thread>
 #include <unordered_map>
+#include <vector>
 
 namespace mooncake {
 
-inline constexpr size_t kP2PBufferSize = 1u << 24;
-inline constexpr size_t kP2PNumSlots = 8;
-inline constexpr size_t kP2PSlotSize = kP2PBufferSize / kP2PNumSlots;
+// Memory layout of one P2PProxy (shown for chunk_size=8MiB, num_chunks=32):
+//
+//   +---------------------------------------------------------+
+//   |  SendPool  (32 x 8 MiB = 256 MiB)                       |
+//   |  +---------+  +---------+        +---------+            |
+//   |  | chunk 0 |  | chunk 1 |  ...   | chunk 31|            |
+//   |  | 8 MiB   |  | 8 MiB   |        | 8 MiB   |            |
+//   |  +---------+  +---------+        +---------+            |
+//   +---------------------------------------------------------+
+//   |  RecvPool (32 x 8 MiB = 256 MiB)                        |
+//   |  +---------+  +---------+        +---------+            |
+//   |  | chunk 0 |  | chunk 1 |  ...   | chunk 31|            |
+//   |  | 8 MiB   |  | 8 MiB   |        | 8 MiB   |            |
+//   |  +---------+  +---------+        +---------+            |
+//   +---------------------------------------------------------+
+//   |  Credit Control Region                                  |
+//   |  [peer 0 lane : 64 slots][peer 1 lane : 64 slots] ...   |
+//   +---------------------------------------------------------+
+//   |  Ack Control Region                                     |
+//   |  [peer 0 lane : 64 slots][peer 1 lane : 64 slots] ...   |
+//   +---------------------------------------------------------+
+//
+// Control lane addressing (isolated per sender-receiver pair):
+//
+//      Rank R wants to write a CreditSlot to peer P
+//      -----------------------------------------------------
+//      Target = P's credit base
+//               + (R * kP2PControlRingSize + seq % 64) * sizeof(CreditSlot)
+//
+//      Rank R wants to write an AckSlot to peer P
+//      -----------------------------------------------------
+//      Target = P's ack base
+//               + (R * kP2PControlRingSize + seq % 64) * sizeof(AckSlot)
+//
+//
+// P2PProxy implements a credit-based RDMA pull protocol.
+//
+// Protocol overview:
+//   - The RECEIVER drives the flow. It allocates chunks from its RecvPool and
+//     writes CreditSlots into the sender's CreditLane to grant credit to the
+//     sender to write data into the designated RecvPool offset.
+//   - The SENDER is passive. It polls its local CreditLane (written by the
+//     receiver via RDMA). Only after receiving a credit does it allocate
+//     a staging buffer from SendPool, copy user tensor data into it, and
+//     perform the RDMA Write to the remote RecvPool.
+//   - After the RDMA Write finishes, the sender writes an AckSlot back
+//     to the receiver's AckLane to acknowledge the transfer.
+//
+// Memory model:
+//   - SendPool / RecvPool are fixed-size chunk pools.
+//   - Control lanes are ring buffers.
+//
+// Per-chunk state machines:
+//   Sender:  Fetch Credit -> Copy-In (staging) -> RDMA Write -> Acknowledge
+//   Receiver: IssueCredit -> Poll Ack -> Copy-Out -> Free Chunk
+//
+// Concurrency rules:
+//   - All pool allocations are non-blocking. If a pool is exhausted the step
+//     function returns false immediately and retries on the next polling
+//     iteration. This prevents deadlocks.
+//   - SendPool and RecvPool are strictly separated so that a deadlock where
+//     "all chunks are reserved for receiving and none are left for sending"
+//     can never happen.
+//
+// Data flow example -- Rank 0 sends a 24 MiB tensor to Rank 1.
+// The five steps below are separated between Sender (Rank 0) and
+// Receiver (Rank 1) to make the protocol explicit.
+//
+//   Step 1 -- Rank 1 issues credits for free RecvPool chunks to Rank 0
+//   ------------------------------------------------------------------------
+//   Rank 1 (Receiver)                RDMA Write
+//   +-----------------------+              +-------------------------------+
+//   | RecvPool              |              | Rank 0's CreditLane[1]        |
+//   | chunk 7 : free        | -----------> | slot 0 : {seq=0,off=7,len=8M} |
+//   | chunk 2 : free        | -----------> | slot 1 : {seq=1,off=2,len=8M} |
+//   | chunk 1 : free        | -----------> | slot 2 : {seq=2,off=1,len=8M} |
+//   +-----------------------+              +-------------------------------+
+//
+//   Step 2 -- Rank 0 stages user tensor into SendPool
+//   ------------------------------------------------------------------------
+//   Rank 0 (Sender)
+//   +-----------------------+
+//   | User tensor           |
+//   |[0..8M)[8..16M)[16..24)|
+//   +-----------------------+
+//            |
+//            | cudaMemcpyAsync (device-to-device)
+//            v
+//   +----------------------------+
+//   | SendPool                   |
+//   | chunk 3 : bytes [0..8M)    |
+//   | chunk 0 : bytes [8..16M)   |
+//   | chunk 5 : bytes [16..24M)  |
+//   +----------------------------+
+//
+//   Step 3 -- Rank 0 RDMA-writes SendPool -> Rank 1's RecvPool
+//   ------------------------------------------------------------------------
+//   Rank 0's SendPool                    RDMA Write       Rank 1's RecvPool
+//   +---------------------+                           +---------------------+
+//   | chunk 3 : [0..8M)   | ------------------------> | chunk 7 : [0..8M)   |
+//   | chunk 0 : [8..16M)  | -------------------- ---> | chunk 2 : [8..16M)  |
+//   | chunk 5 : [16..24M) | ------------------------> | chunk 1 : [16..24M) |
+//   +---------------------+                           +---------------------+
+//
+//   Step 4 -- Rank 0 acknowledges with AckSlots
+//   ------------------------------------------------------------------------
+//   Rank 0 (Sender)                      RDMA Write       Rank 1 (Receiver)
+//   +-----------------------+             +---------------------------------+
+//   | AckLane[0]            |             | Rank 1's AckLane[0]             |
+//   | (local cache only)    | ----------> | slot 0 : {seq=0,len=8M,epoch=G} |
+//   +-----------------------+             | slot 1 : {seq=1,len=8M,epoch=G} |
+//                                         | slot 2 : {seq=2,len=8M,epoch=G} |
+//                                         +---------------------------------+
+//
+//   Step 5 -- Rank 1 copies out and recycles chunks
+//   ------------------------------------------------------------------------
+//   Rank 1 (Receiver)
+//   +-----------------------+
+//   | RecvPool              |
+//   | chunk 7 : [0..8M)     |
+//   | chunk 2 : [8..16M)    |
+//   | chunk 1 : [16..24M)   |
+//   +-----------------------+
+//            |
+//            | cudaMemcpyAsync (RecvPool -> user tensor)
+//            v
+//   +-----------------------+
+//   | User tensor (ready)   |
+//   |[0..8M)[8..16M)[16..24)|
+//   +-----------------------+
+//            |
+//            | Free chunks back to pool
+//            v
+//   +-----------------------+
+//   | RecvPool              |
+//   | chunk 7 : free        |
+//   | chunk 2 : free        |
+//   | chunk 1 : free        |
+//   +-----------------------+
+//
+// ---------------------------------------------------------------------------
+inline constexpr uint32_t kP2PControlRingSize = 64;
+// Pool size and chunk size can be overridden via environment variables:
+//   MOONCAKE_P2P_POOL_SIZE  -- total bytes per direction (default 128 MiB)
+//   MOONCAKE_P2P_CHUNK_SIZE -- granularity of each chunk (default 16 MiB)
+// The chunk count is pool_size / chunk_size (default 8).
+inline constexpr uint64_t kDefaultPoolSize = 128u * 1024 * 1024;  // 128 MiB
+inline constexpr uint64_t kDefaultChunkSize = 16u * 1024 * 1024;  // 16 MiB
 
-struct alignas(64) AtomicHeadTail {
-    uint32_t load(
-        std::memory_order order = std::memory_order_seq_cst) const noexcept {
-        return std::atomic_ref<uint32_t>(value).load(order);
-    }
+// Single-word atomic publication token.
+// Combines epoch and sequence to avoid torn reads.
+// kInvalidControlToken (all 1s) means "slot empty / not yet published".
+using ControlToken = uint64_t;
+inline constexpr ControlToken kInvalidControlToken =
+    std::numeric_limits<uint64_t>::max();
 
-    void store(uint32_t new_value,
-               std::memory_order order = std::memory_order_seq_cst) noexcept {
-        std::atomic_ref<uint32_t>(value).store(new_value, order);
-    }
+inline ControlToken makeControlToken(uint32_t epoch, uint32_t sequence) {
+    return (static_cast<ControlToken>(epoch) << 32) |
+           static_cast<ControlToken>(sequence & 0xFFFFFFFFu);
+}
 
-    mutable uint32_t value{0};
+// CreditSlot
+//
+// Header-Footer double-token guard:
+//   We store the same token at both ends of the 64-byte block.
+//     - Producer: write payload -> write footer -> write header (release).
+//     - Consumer: read header -> read payload -> read footer -> accept only
+//                 when header == footer.
+// If the NIC has only partially written the slot, header and footer will
+// mismatch (or one of them will still be kInvalidControlToken) and the
+// consumer safely retries on the next poll iteration.
+//
+// Layout:
+//   [0..7]   header_token  (epoch << 32 | sequence)
+//   [8..15]  recv_addr     (payload)
+//   [16..19] chunk_len     (payload)
+//   [20..55] padding       (kept zero, reserved for future use)
+//   [56..63] footer_token  (identical to header_token)
+struct alignas(64) CreditSlot {
+   private:
+    uint64_t header_token = kInvalidControlToken;
+    uint64_t recv_addr = 0;
+    uint32_t chunk_len = 0;
+    uint32_t _reserved = 0;
+    uint8_t _padding[32]{};
+    uint64_t footer_token = kInvalidControlToken;
+
+   public:
+    // Consumer-side reliable read.  Returns true when a consistent slot is
+    // observed (header == footer != kInvalidControlToken).
+    bool tryLoad(uint64_t& out_recv_addr, uint32_t& out_chunk_len,
+                 uint32_t& out_epoch, uint32_t& out_sequence) const;
+
+    // Producer-side reliable publish.  Writes payload and footer first,
+    // then releases header so consumers see a consistent slot.
+    void publish(uint32_t epoch, uint32_t seq, uint64_t addr, uint32_t len);
+
+    // Reset both tokens to kInvalidControlToken.
+    void reset();
 };
 
-// Ring-control metadata for one peer lane.
-//
-//   slot index:   0 -> 1 -> 2 -> ... -> N-1 -> 0   (N = kP2PNumSlots)
-//                   ^h (producer publish cursor)
-//                   ^t (consumer reclaim cursor)
-//
-// State by (head, tail), modulo N:
-// - empty: head == tail
-// - full : (head + 1) % N == tail   (one slot reserved to distinguish
-// full/empty)
-// - ready: head != tail
-//
-// Protocol:
-// - producer writes slot[head], then advances head (publish)
-// - consumer reads slot[tail], then advances tail (reclaim)
-struct P2PControlSlot {
-    AtomicHeadTail head;
-    AtomicHeadTail tail;
+// AckSlot Layout:
+//   [0..7]   header_token  (epoch << 32 | sequence)
+//   [8..11]  chunk_len     (payload)
+//   [12..55] padding       (kept zero, reserved for future use)
+//   [56..63] footer_token  (identical to header_token)
+struct alignas(64) AckSlot {
+   private:
+    uint64_t header_token = kInvalidControlToken;
+    uint32_t chunk_len = 0;
+    uint32_t _reserved = 0;
+    uint8_t _padding[40]{};
+    uint64_t footer_token = kInvalidControlToken;
+
+   public:
+    // See CreditSlot::tryLoad().
+    bool tryLoad(uint32_t& out_chunk_len, uint32_t& out_epoch,
+                 uint32_t& out_sequence) const;
+
+    // See CreditSlot::publish().
+    void publish(uint32_t epoch, uint32_t seq, uint32_t len);
+
+    // Reset both tokens to kInvalidControlToken.
+    void reset();
+};
+
+// P2PChunkPool
+// No lock needed because it is only accessed by one worker thread.
+class P2PChunkPool {
+   public:
+    P2PChunkPool() = default;
+
+    // Initialize pool with given base address, chunk size and number of chunks.
+    void init(void* base_addr, size_t chunk_size, uint32_t num_chunks);
+
+    // Acquire a free chunk. Returns nullptr if the pool is exhausted.
+    void* acquire();
+
+    // Release a chunk back to the pool.
+    void release(void* ptr);
+
+   private:
+    void* base_addr_ = nullptr;
+    size_t chunk_size_ = 0;
+    std::vector<void*> free_stack_;
 };
 
 class P2PDeviceWorker;
@@ -63,19 +274,21 @@ class P2PProxy {
    public:
     friend class P2PDeviceWorker;
 
+    enum class OpStatus : uint8_t { kPending = 0, kSuccess = 1, kFailed = 2 };
+
     struct Options {
         bool is_cpu = false;
         int rank = 0;
         int size = 0;
         int cuda_device_index = -1;
-        std::string location;
+        std::chrono::milliseconds transfer_timeout_ms{30000};
     };
 
     struct SendOp {
         at::Tensor tensor_;
         int peer_rank_ = -1;
         cudaStream_t cuda_stream_ = nullptr;
-        std::shared_ptr<std::atomic<bool>> completed_;
+        std::shared_ptr<std::atomic<OpStatus>> status_;
     };
 
     struct RecvOp {
@@ -83,26 +296,21 @@ class P2PProxy {
         at::Tensor original_tensor_;
         int peer_rank_ = -1;
         cudaStream_t cuda_stream_ = nullptr;
-        std::shared_ptr<std::atomic<bool>> completed_;
+        std::shared_ptr<std::atomic<OpStatus>> status_;
     };
 
     P2PProxy(TransferEngine* engine, const Options& options);
     ~P2PProxy();
 
-    void BindMeta(const std::shared_ptr<TransferGroupMeta>& meta);
-    void* send_buffer() const { return resources_.send_buffer_; }
-    void* recv_buffer() const { return resources_.recv_buffer_; }
-    P2PControlSlot* ctrl_send_region() const {
-        return resources_.ctrl_send_region_;
-    }
-    P2PControlSlot* ctrl_recv_region() const {
-        return resources_.ctrl_recv_region_;
-    }
+    CreditSlot* credit_region() const { return resources_.credit_region_; }
+    AckSlot* ack_region() const { return resources_.ack_region_; }
+    void bindMeta(const std::shared_ptr<TransferGroupMeta>& meta);
+    void extendGroupSizeTo(int new_size);
 
-    void EnqueueSend(SendOp op);
-    void EnqueueRecv(RecvOp op);
+    void enqueueSend(SendOp op);
+    void enqueueRecv(RecvOp op);
 
-    void ResetPeerState(int peer_rank);
+    void resetPeerState(int peer_rank);
 
     /**
      * @brief Waits for all active P2P send and receive tasks to complete.
@@ -114,7 +322,7 @@ class P2PProxy {
      * @return True if all tasks completed within the timeout; false if timed
      * out.
      */
-    bool DrainTasks() const;
+    bool drainTasks() const;
 
     /**
      * @brief Abandons resources instead of releasing them properly.
@@ -122,128 +330,225 @@ class P2PProxy {
      * When a hung operation prevents clean shutdown, this method marks
      * resources as abandoned to prevent crashes during destructor.
      */
-    void AbandonResources();
+    void abandonResources();
+
+    // Epoch for fault recovery.  All control slots carry this
+    // value so that stale messages from before a Reset can be detected.
+    uint32_t getEpoch(int peer_rank) const {
+        return peer_epoch_[peer_rank].load(std::memory_order_acquire);
+    }
+    void setEpoch(int peer_rank, uint32_t epoch) {
+        peer_epoch_[peer_rank].store(epoch, std::memory_order_release);
+    }
 
    private:
-    enum class TransferState {
-        kDataCopy,
-        kTransfer,
-        kDone,
+    // Sender-side per-chunk state machine.
+    enum class SendTaskState {
+        kCopyIn,  // Copying tensor slice into local SendPool staging buffer.
+                  // On CPU this is synchronous; on GPU we record a
+                  // cudaEvent on the op's stream and poll it.
+        kWriteRemote,  // Write from SendPool -> remote RecvPool.
+        kAck,          // Writing AckSlot back to the receiver.
+        kFinished,     // AckSlot write done, staging buffer freed.
+        kFailed,       // Transport error or timeout; staging buffer freed.
+    };
+
+    // Receiver-side per-chunk state machine.
+    enum class RecvTaskState {
+        kIssueCredit,  // CreditSlot Write is in flight.
+        kWaitAck,      // Waiting for sender's AckSlot.
+        kCopyOut,      // GPU: cudaMemcpyAsync RecvPool -> tensor in flight.
+        kFinished,     // Copy-out done, chunk returned to pool.
+        kFailed,       // Transport error or timeout; chunk returned to pool.
     };
 
     struct SendOpContext;
     struct RecvOpContext;
 
+    // One chunk of a SendOp.  The sender copies data from the user tensor
+    // into a SendPool staging buffer and then RDMA-writes it to the remote
+    // RecvPool offset that the receiver issued in the CreditSlot.
     struct SendTransferTask {
         SendTransferTask() = default;
-        SendTransferTask(uint64_t chunk_offset_in, uint64_t chunk_bytes_in,
-                         void* source_in, uint64_t target_offset_in);
+        SendTransferTask(uint64_t tensor_offset_in, uint32_t chunk_len_in,
+                         void* staging_addr_in, uint64_t remote_addr_in,
+                         uint32_t sequence_in, uint32_t epoch_in);
 
-        TransferState state_ = TransferState::kDataCopy;
-        uint64_t chunk_offset_ = 0;
-        uint64_t chunk_bytes_ = 0;
-        void* source_ = nullptr;
-        uint64_t target_offset_ = 0;
-        std::optional<BatchID> transfer_batch_id_;
-        cudaEvent_t copy_ready_event_ = nullptr;
+        SendTaskState state_ = SendTaskState::kCopyIn;
+        uint64_t tensor_offset_ = 0;  // Offset inside the user tensor.
+        uint32_t chunk_len_ = 0;      // Bytes in this chunk (<= kP2PChunkSize).
+        void* staging_addr_ = nullptr;  // Address inside SendPool.
+        uint64_t remote_addr_ = 0;      // Address inside REMOTE RecvPool.
+        uint32_t sequence_ = 0;         // Sequence number in the control ring.
+        uint32_t epoch_ = 0;            // Epoch for Reset detection.
+        std::optional<BatchID> transfer_batch_id_;  // RDMA Write batch id.
+        std::optional<BatchID> ack_batch_id_;       // AckSlot write batch id.
+        cudaEvent_t copy_ready_event_ = nullptr;  // Signals Copy-In done (GPU).
+        std::chrono::steady_clock::time_point last_update_time_;
     };
 
+    // Active send operation state.  Tasks are created in order and advance
+    // through the sender state machine (kCopyIn -> kWriteRemote -> kAck
+    // -> kFinished).
     struct SendOpContext {
         SendOpContext() = default;
         SendOpContext(SendOp&& op_in);
 
+        std::deque<SendTransferTask> tasks_;
+        std::shared_ptr<std::atomic<OpStatus>> status_;
+
         at::Tensor tensor_;
         int peer_rank_ = -1;
         cudaStream_t cuda_stream_ = nullptr;
-        std::shared_ptr<std::atomic<bool>> completed_;
         uint64_t total_bytes_ = 0;
-        uint64_t bytes_issued_ = 0;
-        std::optional<BatchID> head_update_batch_id_;
-        std::deque<SendTransferTask> tasks_;
+        // Number of bytes already pulled from the user tensor into SendPool
+        // staging buffers.  When bytes_staged_ == total_bytes_ every chunk
+        // has at least entered the Copy-In stage.
+        uint64_t bytes_staged_ = 0;
+
+        std::chrono::steady_clock::time_point last_update_time_;
     };
 
+    // One chunk of a RecvOp.  The receiver allocates a RecvPool chunk,
+    // issues it to the sender via a CreditSlot, waits for the sender
+    // to RDMA-write data into it, and finally copies the data into the user
+    // tensor before returning the chunk to RecvPool.
     struct RecvTransferTask {
         RecvTransferTask() = default;
-        RecvTransferTask(uint64_t chunk_offset_in, uint64_t chunk_bytes_in,
-                         void* source_in, void* target_in);
+        RecvTransferTask(uint64_t tensor_offset_in, uint32_t chunk_len_in,
+                         void* local_addr_in, uint32_t sequence_in,
+                         uint32_t epoch_in);
 
-        TransferState state_ = TransferState::kDataCopy;
-        uint64_t chunk_offset_ = 0;
-        uint64_t chunk_bytes_ = 0;
-        void* source_ = nullptr;
-        void* target_ = nullptr;
-        cudaEvent_t copy_ready_event_ = nullptr;
+        RecvTaskState state_ = RecvTaskState::kIssueCredit;
+        uint64_t tensor_offset_ = 0;  // Offset inside the user tensor.
+        uint32_t chunk_len_ = 0;      // Bytes in this chunk.
+        void* local_addr_ = nullptr;  // Address inside local RecvPool.
+        uint32_t sequence_ = 0;       // Sequence number in the control ring.
+        uint32_t epoch_ = 0;          // Epoch for Reset detection.
+        std::optional<BatchID>
+            credit_batch_id_;  // CreditSlot RDMA Write batch id.
+        cudaEvent_t copy_ready_event_ =
+            nullptr;  // Signals Copy-Out done (GPU).
+        std::chrono::steady_clock::time_point last_update_time_;
     };
 
+    // Active receive operation state.  Tasks are created in order and
+    // advance through the receiver state machine (kIssueCredit ->
+    // kWaitAck -> kCopyOut -> kFinished).
     struct RecvOpContext {
         RecvOpContext() = default;
         RecvOpContext(RecvOp&& op_in);
 
+        std::deque<RecvTransferTask> tasks_;
+        std::shared_ptr<std::atomic<OpStatus>> status_;
+
         at::Tensor tensor_;
         at::Tensor original_tensor_;
         int peer_rank_ = -1;
         cudaStream_t cuda_stream_ = nullptr;
-        std::shared_ptr<std::atomic<bool>> completed_;
         uint64_t total_bytes_ = 0;
-        uint64_t bytes_issued_ = 0;
-        std::optional<BatchID> tail_update_batch_id_;
-        std::deque<RecvTransferTask> tasks_;
+        // Number of bytes for which a RecvPool chunk has been reserved and a
+        // CreditSlot has been sent to the peer.  When bytes_credited_ ==
+        // total_bytes_ the entire tensor has been offered to the sender.
+        uint64_t bytes_credited_ = 0;
     };
 
+    // Per-peer sender state.  The sender consumes CreditSlots that the
+    // receiver writes into our local CreditLane for this peer.
     struct SendPeerLane {
         std::deque<SendOpContext> pending_send_ops_;
         std::optional<SendOpContext> active_send_op_;
-        uint32_t local_head_ = 0;
-        std::array<cudaEvent_t, kP2PNumSlots> copy_ready_events_;
+        // Sequence number of the next CreditSlot to consume from this peer.
+        // Monotonically increases; wraps around the ring via modulo.
+        uint64_t credit_consume_seq_ = 0;
+        std::array<cudaEvent_t, kP2PControlRingSize> copy_ready_events_;
     };
 
+    // Per-peer receiver state.  The receiver issues CreditSlots to grant credit
+    // to the peer to write data, and consumes AckSlots that the peer writes
+    // back to acknowledge finished transfers.
     struct RecvPeerLane {
         std::deque<RecvOp> pending_recv_ops_;
         std::optional<RecvOpContext> active_recv_op_;
-        uint32_t local_tail_ = 0;
-        std::array<cudaEvent_t, kP2PNumSlots> copy_ready_events_;
+        // Sequence number of the next CreditSlot to issue to this peer.
+        uint64_t credit_issue_seq_ = 0;
+        // Sequence number of the next AckSlot to consume from this peer.
+        uint64_t ack_consume_seq_ = 0;
+        std::array<cudaEvent_t, kP2PControlRingSize> copy_ready_events_;
     };
 
     // Resources are allocated and released by constructor/destructor
-    void AllocateResources();
-    void ReleaseResources();
+    void allocateResources();
+    void releaseResources();
 
     // For P2PDeviceWorker
-    bool StepSend();
-    bool StepRecv();
-    void SetDeviceWorker(P2PDeviceWorker*);
-    bool HasActiveSendWork() const;
-    bool HasActiveRecvWork() const;
+    bool stepSend();
+    bool stepRecv();
+    void attachToWorker(P2PDeviceWorker* worker, P2PChunkPool* send,
+                        P2PChunkPool* recv, size_t chunk_size);
+    bool hasActiveSendWork() const;
+    bool hasActiveRecvWork() const;
 
-    // Internal Steps
-    bool TryIssueSendTask(SendOpContext& op_ctx, uint32_t capacity);
-    bool StepSendTransferTask(SendOpContext& op_ctx, SendTransferTask& task);
-    bool StepSendDataCopy(SendTransferTask& task);
-    bool StepSendTransfer(SendOpContext& op_ctx, SendTransferTask& task);
-    bool StepSendHeadCommit(SendOpContext& op_ctx, uint32_t capacity);
-    bool IsSendDataPathCompleted(const SendOpContext& op_ctx) const;
-    bool IsSendOpCompleted(const SendOpContext& op_ctx) const;
-    void PerformSendReset(int peer_rank);
+    bool tryIssueRecvTask(RecvOpContext& op_ctx, RecvPeerLane& lane);
+    bool stepRecvTask(RecvTransferTask& task);
+    bool stepRecvCopyOut(RecvTransferTask& task);
+    bool stepRecvIssueCredit(RecvTransferTask& task);
+    bool pollRecvAckSlot(RecvOpContext& op_ctx, RecvPeerLane& lane,
+                         RecvTransferTask& head_task);
+    bool isRecvOpCompleted(const RecvOpContext& op_ctx) const;
+    void performRecvReset(int peer_rank);
 
-    bool TryIssueRecvTask(RecvOpContext& op_ctx, uint32_t capacity);
-    bool StepRecvTransferTask(RecvTransferTask& task);
-    bool StepRecvDataCopy(RecvTransferTask& task);
-    bool StepRecvTailCommit(RecvOpContext& op_ctx, uint32_t capacity);
-    bool IsRecvDataPathCompleted(const RecvOpContext& op_ctx) const;
-    void PerformRecvReset(int peer_rank);
+    bool tryIssueSendTask(SendOpContext& op_ctx, SendPeerLane& lane);
+    bool stepSendTask(SendOpContext& op_ctx, SendTransferTask& task);
+    bool stepSendCopyIn(SendTransferTask& task);
+    bool stepSendWriteRemote(SendOpContext& op_ctx, SendTransferTask& task);
+    bool stepSendAck(SendOpContext& op_ctx, SendTransferTask& task);
+    bool isSendOpCompleted(const SendOpContext& op_ctx) const;
+    void performSendReset(int peer_rank);
 
-    uint64_t GetLocalSendSlotAddress(int peer_rank, uint32_t slot_index) const;
-    uint64_t GetLocalRecvSlotAddress(int peer_rank, uint32_t slot_index) const;
-    uint64_t GetRemoteRecvSlotAddress(int peer_rank, uint32_t slot_index) const;
-    uint64_t GetRemoteCtrlRecvHeadOffset(int peer_rank) const;
-    uint64_t GetRemoteCtrlSendTailOffset(int peer_rank) const;
+    void reportBrokenPeer(int peer_rank);
+
+    // These helpers are used only during reset or shutdown.
+    // In the normal execution path, task and lane resources are released
+    // incrementally as work progresses.
+    void releaseSendTaskResources(SendTransferTask& task) const;
+    void releaseRecvTaskResources(RecvTransferTask& task) const;
+    void resetSendLane(SendPeerLane& lane);
+    void resetRecvLane(RecvPeerLane& lane);
+    void resetPeerControlLanes(int peer_rank);
+
+    // Control lane addressing.
+    //
+    // Each rank owns one contiguous Credit region and one Ack region.
+    // Within each region there are kMaxNumRanks lanes, one per peer.
+    // Lane index == peer rank.  Each lane has kP2PControlRingSize slots.
+    //
+    // Example: rank R wants to write a CreditSlot to peer P.
+    //   target = P's credit base + (R * kRingSize + seq % kRingSize) * sizeof
+    //
+    // This isolates control traffic per (sender, receiver) pair.
+    CreditSlot* getLocalCreditLane(int peer_rank) const;
+    AckSlot* getLocalAckLane(int peer_rank) const;
+    uint64_t getRemoteCreditSlot(int peer_rank, uint32_t sequence) const;
+    uint64_t getRemoteAckSlot(int peer_rank, uint32_t sequence) const;
+
+    CreditSlot* getLocalCreditStagingBuf(int peer_rank,
+                                         uint32_t sequence) const;
+    AckSlot* getLocalAckStagingBuf(int peer_rank, uint32_t sequence) const;
+
+    template <typename T>
+    bool isTimeout(const T& obj) const {
+        return std::chrono::steady_clock::now() - obj.last_update_time_ >
+               transfer_timeout_ms_;
+    }
 
    private:
     struct P2PResources {
-        void* send_buffer_ = nullptr;
-        void* recv_buffer_ = nullptr;
-        P2PControlSlot* ctrl_send_region_ = nullptr;
-        P2PControlSlot* ctrl_recv_region_ = nullptr;
+        CreditSlot* credit_region_ = nullptr;
+        AckSlot* ack_region_ = nullptr;
+        // Serve as RDMA write source address
+        CreditSlot* credit_staging_buf_ = nullptr;
+        AckSlot* ack_staging_buf_ = nullptr;
     };
 
     P2PDeviceWorker* device_worker_ = nullptr;
@@ -254,10 +559,15 @@ class P2PProxy {
     int rank_ = 0;
     int size_ = 0;
     int cuda_device_index_ = -1;
-    std::string location_;
+    std::chrono::milliseconds transfer_timeout_ms_{5000};  // 5s
     P2PResources resources_;
     bool resource_abandoned_{false};
 
+    size_t chunk_size_ = 0;
+
+    P2PChunkPool* send_pool_ = nullptr;
+    P2PChunkPool* recv_pool_ = nullptr;
+
     std::queue<SendOpContext> send_queue_;
     std::mutex send_queue_mutex_;
 
@@ -270,6 +580,11 @@ class P2PProxy {
     std::atomic<int> active_send_tasks_{0};
     std::atomic<int> active_recv_tasks_{0};
 
+    // Per-peer epoch for fault recovery.  Incremented in resetPeerState
+    // and performSend/RecvReset so that stale messages from a previous epoch
+    // can be detected on a per-peer basis.
+    std::array<std::atomic<uint32_t>, kMaxNumRanks> peer_epoch_;
+
     std::array<SendPeerLane, kMaxNumRanks> send_peer_lanes_;
     std::array<RecvPeerLane, kMaxNumRanks> recv_peer_lanes_;
 };
@@ -283,22 +598,19 @@ class P2PDeviceWorker {
     void registerProxy(const std::shared_ptr<P2PProxy>&);
     void removeProxy(const std::shared_ptr<P2PProxy>&);
 
-    P2PDeviceWorker(bool is_cpu, int cuda_device_index)
-        : is_cpu_(is_cpu), cuda_device_index_(cuda_device_index) {
-        Start();
-    }
+    P2PDeviceWorker(TransferEngine* engine, const std::string& location,
+                    bool is_cpu, int cuda_device_index);
 
-    ~P2PDeviceWorker() { Stop(); }
-
-    void WakeUpSend();
-    void WakeUpRecv();
+    ~P2PDeviceWorker();
+    void wakeUpSend();
+    void wakeUpRecv();
 
    private:
-    void Start();
-    void Stop();
+    void start();
+    void stop();
 
-    void SendWorkerThread();
-    void RecvWorkerThread();
+    void sendWorkerMainloop();
+    void recvWorkerMainloop();
 
     std::mutex send_wakeup_mutex_;
     std::condition_variable send_wakeup_cv_;
@@ -318,18 +630,36 @@ class P2PDeviceWorker {
 
     bool is_cpu_;
     int cuda_device_index_;
+
+    // Per-device shared chunk pools
+    //
+    // Memory footprint (with default env values):
+    //   - One direction (send or recv)    : 128 MiB
+    //   - Total per device (send + recv)  : 256 MiB
+    // This stays constant no matter how many backends or peer ranks are active.
+    TransferEngine* engine_ = nullptr;
+    P2PChunkPool send_pool_;
+    P2PChunkPool recv_pool_;
+    void* send_pool_base_ = nullptr;
+    void* recv_pool_base_ = nullptr;
+    size_t pool_bytes_ = 0;
+    size_t chunk_size_ = 0;
+    uint32_t num_chunks_ = 0;
+    void initPools(TransferEngine* engine, const std::string& location);
+    void releasePools();
 };
 
 class P2PDeviceWorkerManager {
    public:
-    static P2PDeviceWorkerManager& GetInstance() {
+    static P2PDeviceWorkerManager& getInstance() {
         // leaky singleton to avoid destructor fiasco problem
         static P2PDeviceWorkerManager* manager = new P2PDeviceWorkerManager;
         return *manager;
     }
 
-    std::shared_ptr<P2PDeviceWorker> GetCPUWorker();
-    std::shared_ptr<P2PDeviceWorker> GetCUDAWorker(int cuda_device_index);
+    std::shared_ptr<P2PDeviceWorker> getCPUWorker(TransferEngine* engine);
+    std::shared_ptr<P2PDeviceWorker> getCUDAWorker(int cuda_device_index,
+                                                   TransferEngine* engine);
 
    private:
     static constexpr int CPUWorkerID = -1;
diff --git a/mooncake-pg/setup.py b/mooncake-pg/setup.py
index 8a393773..e839339b 100644
--- a/mooncake-pg/setup.py
+++ b/mooncake-pg/setup.py
@@ -3,7 +3,23 @@ import re
 
 from setuptools import setup
 import torch
-from torch.utils.cpp_extension import BuildExtension, CUDAExtension, CUDA_HOME
+
+use_musa = os.getenv("MOONCAKE_EP_USE_MUSA", "").upper() in {"1", "ON", "TRUE", "YES"}
+if use_musa:
+    try:
+        import torchada  # noqa: F401
+    except ImportError as e:
+        raise ImportError(
+            "torchada is required to build the MUSA PG extension. "
+            "Please install it first using 'pip install torchada'."
+        ) from e
+
+
+from torch.utils.cpp_extension import (  # noqa: E402
+    BuildExtension,
+    CUDAExtension,
+    CUDA_HOME,
+)
 
 
 torch_version = re.match(r"\d+(?:\.\d+)*", torch.__version__).group()
@@ -13,18 +29,46 @@ 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.
+abi_define = f"-D_GLIBCXX_USE_CXX11_ABI={abi_flag}"
+cxx_args = [abi_define, "-std=c++20", "-O3", "-g0"]
+
 cuda_libraries = ["ibverbs", "mlx5"]
 cuda_library_dirs = []
+use_maca = hasattr(torch.version, "maca") and torch.version.maca is not None
 
-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)
-
+if use_musa:
+    musa_defines = ["-DUSE_MUSA", "-DMOONCAKE_EP_USE_MUSA=1"]
+    cxx_args += musa_defines
+    # torchada maps the "nvcc" key to "mcc".
+    device_args = [
+        abi_define,
+        *musa_defines,
+        "-std=c++20",
+        "--cuda-gpu-arch=mp_21",
+        "--cuda-gpu-arch=mp_31",
+        "-O3",
+    ]
+else:
+    if use_maca:
+        cxx_args.append("-DUSE_MACA")
+    device_args = [
+        abi_define,
+        "-std=c++20",
+        "-Xcompiler",
+        "-O3",
+        "-Xcompiler",
+        "-g0",
+    ]
+    if use_maca:
+        device_args.append("-DUSE_MACA")
+    # Link against the CUDA driver stub library if available.
+    # Same approach as mooncake-ep/setup.py.
+    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,
@@ -40,31 +84,19 @@ setup(
                 "src/mooncake_backend.cpp",
                 "src/p2p_proxy.cpp",
                 "src/mooncake_worker.cu",
+                "src/mooncake_worker_host.cpp",
                 "src/mooncake_worker_thread.cpp",
                 "src/connection_poller.cpp",
             ],
-            extra_compile_args={
-                "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",
-                ],
-            },
+            extra_compile_args={"cxx": cxx_args, "nvcc": device_args},
             libraries=cuda_libraries,
             library_dirs=cuda_library_dirs,
             extra_link_args=[
                 "-Wl,-rpath,$ORIGIN",
                 "-L" + os.path.join(current_dir, "../mooncake-wheel/mooncake"),
+                "-Wl,--push-state,--no-as-needed",
                 "-l:engine.so",
+                "-Wl,--pop-state",
             ],
         )
     ],
diff --git a/mooncake-pg/src/connection_poller.cpp b/mooncake-pg/src/connection_poller.cpp
index 9199806f..e516cc4e 100644
--- a/mooncake-pg/src/connection_poller.cpp
+++ b/mooncake-pg/src/connection_poller.cpp
@@ -1,9 +1,6 @@
 #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>
 #include <chrono>
@@ -23,6 +20,7 @@ namespace mooncake {
 // 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.
+#if !defined(MOONCAKE_EP_USE_MUSA) && !defined(USE_MACA)
 static bool supportFabricMem() {
     const char* nvlink_ipc = getenv("MC_USE_NVLINK_IPC");
 
@@ -41,6 +39,10 @@ static bool supportFabricMem() {
     }
     return true;
 }
+#else
+// MUSA and MACA do not use NVIDIA NVLink fabric memory handles here.
+static bool supportFabricMem() { return false; }
+#endif
 ConnectionContext::ConnectionContext(int backendIndex, int rank, int size,
                                      bool isDummy,
                                      uint64_t* local2global_rank_map,
@@ -51,6 +53,7 @@ ConnectionContext::ConnectionContext(int backendIndex, int rank, int size,
     : backendIndex_(backendIndex),
       rank_(rank),
       groupSize_(size),
+      pollingLimit_(size),
       isDummy_(isDummy),
       establishedGroupSize_(0),
       local2global_rank_map_(local2global_rank_map),
@@ -122,6 +125,18 @@ void ConnectionContext::extendGroupSizeTo(int newGroupSize) {
     }
 
     groupSize_.store(newGroupSize, std::memory_order_release);
+    // Keep polling range aligned with the largest known size.
+    pollingLimit_.store(newGroupSize, std::memory_order_release);
+}
+
+void ConnectionContext::setPollingLimitTo(int pollingLimit) {
+    const int groupSize = groupSize_.load(std::memory_order_acquire);
+    TORCH_CHECK(pollingLimit >= groupSize,
+                "pollingLimit must be >= current groupSize");
+    TORCH_CHECK(
+        pollingLimit >= 0 && static_cast<size_t>(pollingLimit) < kMaxNumRanks,
+        "Size out of range");
+    pollingLimit_.store(pollingLimit, std::memory_order_release);
 }
 
 bool ConnectionContext::isAllPeerConnected() const {
@@ -183,6 +198,7 @@ void ConnectionContext::bootstrapLocalPeer(const std::string& localServerName,
     ConnectionPoller::GetInstance()
         .global_peerConnected_[local2global_rank_map_[rank_]] = true;
     peerState.state = PeerConnectionState::CONNECTED;
+    peerState.countedInGroup = true;
 
     {
         std::lock_guard<std::mutex> lock(backend_wakeup_mutex_);
@@ -210,7 +226,8 @@ bool ConnectionContext::poll() {
     bool did_work = false;
 
     // Poll all peers sequentially.
-    for (int pollingRank = 0; pollingRank < groupSize_; ++pollingRank) {
+    const int pollingLimit = pollingLimit_.load(std::memory_order_acquire);
+    for (int pollingRank = 0; pollingRank < pollingLimit; ++pollingRank) {
         did_work |= pollPeer(pollingRank);
     }
 
@@ -279,9 +296,13 @@ bool ConnectionContext::pollPeer(int pollingRank) {
                 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();
+                    if (pollingRank <
+                        groupSize_.load(std::memory_order_acquire)) {
+                        totalConnectedPeers_.fetch_add(
+                            1, std::memory_order_release);
+                        peerState.countedInGroup = true;
+                        backend_wakeup_cv_.notify_all();
+                    }
                 }
             } else if (pollingRank <= rank_) {
                 // Send a warmup request to establish connections
@@ -322,9 +343,13 @@ bool ConnectionContext::pollPeer(int pollingRank) {
 
                 {
                     std::lock_guard<std::mutex> lock(backend_wakeup_mutex_);
-                    totalConnectedPeers_.fetch_add(1,
-                                                   std::memory_order_release);
-                    backend_wakeup_cv_.notify_all();
+                    if (pollingRank <
+                        groupSize_.load(std::memory_order_acquire)) {
+                        totalConnectedPeers_.fetch_add(
+                            1, std::memory_order_release);
+                        peerState.countedInGroup = true;
+                        backend_wakeup_cv_.notify_all();
+                    }
                 }
                 state_changed = true;
             } else if (status.s == TransferStatusEnum::FAILED) {
@@ -349,9 +374,13 @@ bool ConnectionContext::pollPeer(int pollingRank) {
                 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();
+                    if (pollingRank <
+                        groupSize_.load(std::memory_order_acquire)) {
+                        totalConnectedPeers_.fetch_add(
+                            1, std::memory_order_release);
+                        peerState.countedInGroup = true;
+                        backend_wakeup_cv_.notify_all();
+                    }
                 }
                 state_changed = true;
             }
@@ -359,6 +388,11 @@ bool ConnectionContext::pollPeer(int pollingRank) {
         }
 
         case PeerConnectionState::CONNECTED: {
+            // A peer may be warmed up (CONNECTED) while still outside of the
+            // current groupSize_. When it later enters the group, we need to
+            // update totalConnectedPeers_ lazily here; otherwise
+            // waitUntilAllConnected()/isAllPeerConnected() may hang.
+
             // ATTENTION: Ensure consistency of local (meta_->peerConnected)
             //            and global (global_peerConnected_).
             //
@@ -382,6 +416,14 @@ bool ConnectionContext::pollPeer(int pollingRank) {
 
             if (meta_->peerConnected[pollingRank] &&
                 global_peerConnected_[globalPollingRank]) {
+                if (!peerState.countedInGroup &&
+                    pollingRank < groupSize_.load(std::memory_order_acquire)) {
+                    std::lock_guard<std::mutex> lock(backend_wakeup_mutex_);
+                    totalConnectedPeers_.fetch_add(1,
+                                                   std::memory_order_release);
+                    peerState.countedInGroup = true;
+                    backend_wakeup_cv_.notify_all();
+                }
                 // happy path: both are connected.
                 break;
             }
@@ -391,7 +433,9 @@ bool ConnectionContext::pollPeer(int pollingRank) {
             global_peerConnected_[globalPollingRank] = false;
             meta_->peerConnected[pollingRank] = false;
             meta_->activeRanks[pollingRank] = false;
-            meta_->activeRanksTensor[pollingRank] = 0;
+            if (meta_->activeRanksTensor.device().is_cpu()) {
+                meta_->activeRanksTensor[pollingRank] = 0;
+            }
 
             // Reset store
             try {
@@ -400,9 +444,7 @@ bool ConnectionContext::pollPeer(int pollingRank) {
                 store_->deleteKey(
                     getBufferStoreKey(backendIndex_, pollingRank));
                 store_->deleteKey(
-                    getExtensionTaskCountStoreKey(backendIndex_, pollingRank));
-                store_->deleteKey(getExtensionActiveRanksStoreKey(backendIndex_,
-                                                                  pollingRank));
+                    getExtensionStateStoreKey(backendIndex_, pollingRank));
             } catch (const std::exception& e) {
                 LOG(WARNING) << "Rank " << rank_
                              << " got an exception when deleteKey for peer "
@@ -410,17 +452,22 @@ bool ConnectionContext::pollPeer(int pollingRank) {
             }
 
             // Reset warmup region
-            *reinterpret_cast<volatile int32_t*>(
-                &warmup_recv_region_[pollingRank]) = 0;
+            if (warmup_recv_region_) {
+                *reinterpret_cast<volatile int32_t*>(
+                    &warmup_recv_region_[pollingRank]) = 0;
+            }
 
             // Reset P2PProxy states
-            p2p_proxy_->ResetPeerState(pollingRank);
+            p2p_proxy_->resetPeerState(pollingRank);
 
             // Back to WAITING_STORE to reconnect it.
             peerState.state = PeerConnectionState::WAITING_STORE;
             engine_->closeSegment(peerState.segmentId.value());
             peerState.segmentId = std::nullopt;
-            totalConnectedPeers_.fetch_sub(1);
+            if (peerState.countedInGroup) {
+                totalConnectedPeers_.fetch_sub(1, std::memory_order_release);
+                peerState.countedInGroup = false;
+            }
             state_changed = true;
             break;
         }
diff --git a/mooncake-pg/src/mooncake_backend.cpp b/mooncake-pg/src/mooncake_backend.cpp
index 11c20658..ceb746b0 100644
--- a/mooncake-pg/src/mooncake_backend.cpp
+++ b/mooncake-pg/src/mooncake_backend.cpp
@@ -1,13 +1,13 @@
 #include <ATen/cuda/CUDAContext.h>
-#include <cuda_runtime.h>
+#include <cuda_alike.h>
 #include <torch/torch.h>
 #include <torch/csrc/distributed/c10d/Backend.hpp>
-#include <cuda_alike.h>
 #include <mooncake_backend.h>
 #include <p2p_proxy.h>
 #include <thread>
 #include <chrono>
 #include <atomic>
+#include <cstdlib>
 #include <memory>
 #include "connection_poller.h"
 #include "memory_location.h"
@@ -15,6 +15,18 @@
 #include "pg_utils.h"
 
 namespace mooncake {
+namespace {
+
+#ifdef USE_MACA
+static void requireMacaHostTransport() {
+    TORCH_CHECK(std::getenv("MC_MACA_HOST_TRANSPORT") != nullptr,
+                "MACA PG requires MC_MACA_HOST_TRANSPORT=1 so the transfer "
+                "engine uses a host transport.");
+}
+
+#endif
+
+}  // namespace
 
 constexpr const char* REGISTER_BUFFER_ERROR_MSG =
     "Failed to register local memory.";
@@ -32,60 +44,140 @@ TransferEngine* MooncakeBackend::engine_ = new TransferEngine(true);
 // worker_ is now owned per backend instance via MooncakeWorkerManager.
 bool MooncakeBackend::engineInitialized_ = false;
 int MooncakeBackend::backendIndex_ = 0;
+TransferEngine* MooncakeBackend::externalEngine_ = nullptr;
 
-namespace {
+std::vector<uint8_t> serialize(const ExtensionState& state) {
+    uint32_t rankCount = static_cast<uint32_t>(state.activeRanks.size());
 
-std::vector<uint8_t> serializeActiveRanks(const bool* activeRanks, int size) {
-    std::vector<uint8_t> bytes(size);
-    for (int i = 0; i < size; ++i) {
-        bytes[i] = activeRanks[i] ? 1 : 0;
+    // Calculate bytes needed for the bitmap: 1 bit per rank, rounded up to
+    // nearest byte
+    size_t bitmapSize = (rankCount + 7) / 8;
+
+    // Total size = count field + bitmap + p2pEpochs[] + taskCount
+    size_t totalSize = sizeof(uint32_t) + bitmapSize +
+                       sizeof(uint32_t) * rankCount + sizeof(int32_t);
+
+    std::vector<uint8_t> buffer(totalSize, 0);
+    uint8_t* ptr = buffer.data();
+
+    // 1. Store the number of ranks
+    std::memcpy(ptr, &rankCount, sizeof(uint32_t));
+    ptr += sizeof(uint32_t);
+
+    // 2. Store activeRanks as a bitset
+    for (size_t i = 0; i < rankCount; ++i) {
+        if (state.activeRanks[i]) {
+            // Set the i-th bit to 1 if the rank is active
+            ptr[i / 8] |= (1 << (i % 8));
+        }
     }
-    return bytes;
+    ptr += bitmapSize;
+
+    // 3. Store per-peer p2pEpochs
+    for (size_t i = 0; i < rankCount; ++i) {
+        std::memcpy(ptr + i * sizeof(uint32_t), &state.p2pEpochs[i],
+                    sizeof(uint32_t));
+    }
+    ptr += sizeof(uint32_t) * rankCount;
+
+    // 4. Store taskCount
+    int32_t taskCount = static_cast<int32_t>(state.taskCount);
+    std::memcpy(ptr, &taskCount, sizeof(int32_t));
+
+    return buffer;
 }
 
-void deserializeActiveRanks(const std::vector<uint8_t>& bytes,
-                            bool* activeRanks, int size) {
-    TORCH_CHECK(static_cast<int>(bytes.size()) == size,
-                "Unexpected active-ranks snapshot size.");
-    for (int i = 0; i < size; ++i) {
-        activeRanks[i] = (bytes[i] != 0);
-    }
-}
+ExtensionState deserialize(const std::vector<uint8_t>& buffer) {
+    ExtensionState state;
+    if (buffer.size() < sizeof(uint32_t)) return state;
 
-}  // namespace
+    const uint8_t* ptr = buffer.data();
+
+    // 1. Read the number of ranks
+    uint32_t rankCount = 0;
+    std::memcpy(&rankCount, ptr, sizeof(uint32_t));
+    ptr += sizeof(uint32_t);
+
+    // Calculate expected total size and verify buffer is sufficient before
+    // proceeding with further reads.
+    size_t bitmapSize = (rankCount + 7) / 8;
+    size_t expectedSize = sizeof(uint32_t) + bitmapSize +
+                          sizeof(uint32_t) * rankCount + sizeof(int32_t);
+    if (buffer.size() < expectedSize) return state;
+
+    // 2. Read the bitmap and reconstruct the activeRanks vector
+    state.activeRanks.resize(rankCount);
+    for (size_t i = 0; i < rankCount; ++i) {
+        // Check if the i-th bit is set
+        bool isActive = ptr[i / 8] & (1 << (i % 8));
+        state.activeRanks[i] = isActive;
+    }
+    ptr += bitmapSize;
+
+    // 3. Read per-peer p2pEpochs
+    state.p2pEpochs.resize(rankCount);
+    for (size_t i = 0; i < rankCount; ++i) {
+        std::memcpy(&state.p2pEpochs[i], ptr, sizeof(uint32_t));
+        ptr += sizeof(uint32_t);
+    }
+
+    // 4. Read taskCount
+    int32_t taskCount = 0;
+    std::memcpy(&taskCount, ptr, sizeof(int32_t));
+    state.taskCount = static_cast<int>(taskCount);
+
+    return state;
+}
 
 // Async Work implementation for P2P operations processed by worker threads.
 class MooncakeP2PWork : public ::c10d::Work {
    public:
-    explicit MooncakeP2PWork(std::shared_ptr<std::atomic<bool>> completed)
-        : Work(-1, c10d::OpType::UNKNOWN), completed_(completed) {}
+    explicit MooncakeP2PWork(
+        std::shared_ptr<std::atomic<P2PProxy::OpStatus>> status)
+        : Work(-1, c10d::OpType::UNKNOWN), status_(status) {}
 
     bool isCompleted() override {
-        return completed_->load(std::memory_order_acquire);
+        return status_->load(std::memory_order_acquire) !=
+               P2PProxy::OpStatus::kPending;
+    }
+
+    bool isSuccess() const override {
+        return status_->load(std::memory_order_acquire) ==
+               P2PProxy::OpStatus::kSuccess;
     }
 
     bool wait(std::chrono::milliseconds timeout) override {
-        if (completed_->load(std::memory_order_acquire)) {
-            return true;
-        }
-
         BackoffWaiterConfig cfg{};
         cfg.max_sleep = std::chrono::microseconds(10);
         BackoffWaiter waiter(cfg);
 
+        bool done = false;
         if (timeout.count() > 0) {
-            return waiter.wait_for(timeout, [this] {
-                return completed_->load(std::memory_order_acquire);
+            done = waiter.wait_for(timeout, [this] {
+                return status_->load(std::memory_order_acquire) !=
+                       P2PProxy::OpStatus::kPending;
             });
+        } else {
+            waiter.wait([this] {
+                return status_->load(std::memory_order_acquire) !=
+                       P2PProxy::OpStatus::kPending;
+            });
+            done = true;
         }
 
-        waiter.wait(
-            [this] { return completed_->load(std::memory_order_acquire); });
+        if (!done) {
+            return false;
+        }
+
+        if (status_->load(std::memory_order_acquire) ==
+            P2PProxy::OpStatus::kFailed) {
+            TORCH_CHECK(false, "Mooncake P2P operation failed.");
+        }
         return true;
     }
 
    private:
-    std::shared_ptr<std::atomic<bool>> completed_;
+    std::shared_ptr<std::atomic<P2PProxy::OpStatus>> status_;
 };
 
 /**
@@ -95,12 +187,21 @@ class MooncakeP2PWork : public ::c10d::Work {
 MooncakeBackend::MooncakeBackend(
     c10d::DistributedBackendOptions distBackendOpts,
     c10::intrusive_ptr<MooncakeBackendOptions> options, bool isCpu)
-    : Backend(distBackendOpts.group_rank, distBackendOpts.group_size),
+    : ProcessGroup(distBackendOpts.store, distBackendOpts.group_rank,
+                   distBackendOpts.group_size),
       options_(std::move(options)),
       isCpu_(isCpu) {
     auto store = std::move(distBackendOpts.store);
     const int rank = distBackendOpts.group_rank;
     const int size = distBackendOpts.group_size;
+    const int max_size = (options_ && options_->maxWorldSize_ > 0)
+                             ? options_->maxWorldSize_
+                             : size;
+
+    TORCH_CHECK(max_size >= 0 && static_cast<size_t>(max_size) <= kMaxNumRanks,
+                "max_world_size out of range");
+    TORCH_CHECK(max_size >= size,
+                "max_world_size must be >= process group size");
     const auto& globalRanks = distBackendOpts.global_ranks_in_group;
 
     // Memory location for device specific buffers
@@ -118,7 +219,14 @@ MooncakeBackend::MooncakeBackend(
     }
 
     // Initialize transfer engine
-    if (!engineInitialized_) {
+    if (externalEngine_) {
+        // Use externally-provided engine (already initialized), skip init.
+        engine_ = externalEngine_;
+        engineInitialized_ = true;
+    } else if (!engineInitialized_) {
+#ifdef USE_MACA
+        requireMacaHostTransport();
+#endif
         engine_->init(P2PHANDSHAKE, hostIp_);
         engineInitialized_ = true;
     }
@@ -134,6 +242,11 @@ MooncakeBackend::MooncakeBackend(
         }
     }
 
+    // Fill the remaining slots for polling / future joiners.
+    for (int i = size; i < max_size; ++i) {
+        local2global_rank_map_[i] = i;
+    }
+
     // Register buffers
     if (isCpu) {
         for (size_t i = 0; i < 2; i++) {
@@ -156,9 +269,31 @@ MooncakeBackend::MooncakeBackend(
             TORCH_CHECK(!rc, REGISTER_BUFFER_ERROR_MSG);
         }
 
+#ifdef USE_MACA
     } else {
         for (size_t i = 0; i < 2; i++) {
-            cudaError err = cudaMalloc(&send_buffer_[i], kBufferSize);
+            cudaError_t err = cudaMalloc(&send_buffer_[i], kBufferSize);
+            TORCH_CHECK(!err,
+                        c10::str("Failed to allocate MACA GPU send buffer"));
+
+            int rc = engine_->registerLocalMemory(send_buffer_[i], kBufferSize,
+                                                  location);
+            TORCH_CHECK(!rc, REGISTER_BUFFER_ERROR_MSG);
+        }
+
+        for (size_t i = 0; i < 2; i++) {
+            cudaError_t err = cudaMalloc(&recv_buffer_[i], kBufferSize);
+            TORCH_CHECK(!err,
+                        c10::str("Failed to allocate MACA GPU recv buffer"));
+
+            int rc = engine_->registerLocalMemory(recv_buffer_[i], kBufferSize,
+                                                  location);
+            TORCH_CHECK(!rc, REGISTER_BUFFER_ERROR_MSG);
+        }
+#else
+    } else {
+        for (size_t i = 0; i < 2; i++) {
+            cudaError_t err = cudaMalloc(&send_buffer_[i], kBufferSize);
             TORCH_CHECK(!err, c10::str("Failed to allocate CUDA send buffer"));
 
             int rc = engine_->registerLocalMemory(send_buffer_[i], kBufferSize,
@@ -167,13 +302,14 @@ MooncakeBackend::MooncakeBackend(
         }
 
         for (size_t i = 0; i < 2; i++) {
-            cudaError err = cudaMalloc(&recv_buffer_[i], kBufferSize);
+            cudaError_t err = cudaMalloc(&recv_buffer_[i], kBufferSize);
             TORCH_CHECK(!err, c10::str("Failed to allocate CUDA recv buffer"));
 
             int rc = engine_->registerLocalMemory(recv_buffer_[i], kBufferSize,
                                                   location);
             TORCH_CHECK(!rc, REGISTER_BUFFER_ERROR_MSG);
         }
+#endif
     }
 
     // Register CPU sync regions
@@ -195,13 +331,14 @@ MooncakeBackend::MooncakeBackend(
         TORCH_CHECK(!rc, REGISTER_BUFFER_ERROR_MSG);
     }
 
-    auto& dev_worker_mgr = P2PDeviceWorkerManager::GetInstance();
+    auto& dev_worker_mgr = P2PDeviceWorkerManager::getInstance();
     int cuda_device_index = isCpu_ ? -1 : at::cuda::current_device();
 
     if (isCpu_)
-        p2p_device_worker_ = dev_worker_mgr.GetCPUWorker();
+        p2p_device_worker_ = dev_worker_mgr.getCPUWorker(engine_);
     else
-        p2p_device_worker_ = dev_worker_mgr.GetCUDAWorker(cuda_device_index);
+        p2p_device_worker_ =
+            dev_worker_mgr.getCUDAWorker(cuda_device_index, engine_);
 
     auto& worker_mgr = MooncakeWorkerManager::GetInstance();
     if (isCpu_)
@@ -217,9 +354,8 @@ MooncakeBackend::MooncakeBackend(
         engine_, P2PProxy::Options{
                      .is_cpu = isCpu_,
                      .rank = rank_,
-                     .size = size_,
+                     .size = max_size,
                      .cuda_device_index = cuda_device_index,
-                     .location = location,
                  });
     p2p_device_worker_->registerProxy(p2p_proxy_);
 
@@ -228,6 +364,10 @@ MooncakeBackend::MooncakeBackend(
         backendIndex_, rank, size, options_ && options_->isExtension_,
         local2global_rank_map_, store, meta_, p2p_proxy_, engine_);
 
+    if (max_size != size) {
+        connection_ctx_->setPollingLimitTo(max_size);
+    }
+
     rank_info.send_buffer[0] = (uint64_t)send_buffer_[0];
     rank_info.send_buffer[1] = (uint64_t)send_buffer_[1];
     rank_info.recv_buffer[0] = (uint64_t)recv_buffer_[0];
@@ -240,16 +380,25 @@ MooncakeBackend::MooncakeBackend(
         (uint64_t)connection_ctx_->warmup_send_region();
     rank_info.warmup_buffer[1] =
         (uint64_t)connection_ctx_->warmup_recv_region();
-    rank_info.p2p_send_buffer = (uint64_t)p2p_proxy_->send_buffer();
-    rank_info.p2p_recv_buffer = (uint64_t)p2p_proxy_->recv_buffer();
-    rank_info.p2p_ctrl_send = (uint64_t)p2p_proxy_->ctrl_send_region();
-    rank_info.p2p_ctrl_recv = (uint64_t)p2p_proxy_->ctrl_recv_region();
+    rank_info.p2p_credit_region = (uint64_t)p2p_proxy_->credit_region();
+    rank_info.p2p_ack_region = (uint64_t)p2p_proxy_->ack_region();
 
     // Sync metadata
     std::vector<uint8_t> rank_info_bytes(sizeof(SegmentInfo));
     memcpy(rank_info_bytes.data(), &rank_info, sizeof(SegmentInfo));
     meta_->rank = rank;
-    meta_->size = size;
+    // NOTE: meta_->size is intentionally initialized to max_world_size (when
+    // provided) so that healthy ranks can activate joiners via recoverRanks()
+    // without calling extendGroupSizeTo(). Inactive slots are masked by
+    // meta_->activeRanks / meta_->activeRanksTensor.
+    meta_->size = max_size;
+    // activeSize tracks the visible group size (returned by getSize() /
+    // dist.get_world_size()). It starts at the actual member count and grows
+    // when extendGroupSizeTo() or recoverRanks() expands the group.
+    // For extension ranks, activeSize equals world_size (= max_world_size);
+    // the local-only behavior before joinGroup() is ensured by activeRanks
+    // masking, not by a smaller activeSize.
+    meta_->activeSize = size;
     meta_->taskCount = 0;
     if (isCpu) {
         meta_->activeRanks = new bool[kMaxNumRanks];
@@ -262,6 +411,11 @@ MooncakeBackend::MooncakeBackend(
     for (size_t i = 0; i < kMaxNumRanks; ++i) {
         meta_->activeRanks[i] = true;
     }
+
+    // Reserve extra slots as inactive so collectives won't wait on them.
+    for (int i = size; i < max_size; ++i) {
+        meta_->activeRanks[i] = false;
+    }
     if (options_ && options_->activeRanks_.defined()) {
         TORCH_CHECK(options_->activeRanks_.dtype() == at::kInt,
                     "activeRanks must be int.");
@@ -269,20 +423,29 @@ MooncakeBackend::MooncakeBackend(
             TORCH_CHECK(options_->activeRanks_.device().is_cpu(),
                         "activeRanks must be on CPU.");
         } else {
-            TORCH_CHECK(options_->activeRanks_.device().is_cuda(),
-                        "activeRanks must be on CUDA.");
+            TORCH_CHECK(
+                options_->activeRanks_.device().type() == c10::DeviceType::CUDA,
+                "activeRanks must be on GPU.");
+        }
+        if (max_size != size) {
+            TORCH_CHECK(options_->activeRanks_.numel() == max_size,
+                        "activeRanks must be sized to max_world_size when "
+                        "max_world_size is set");
         }
         meta_->activeRanksTensor = options_->activeRanks_;
     } else {
-        meta_->activeRanksTensor =
-            at::ones({size}, torch::dtype(torch::kInt32)
-                                 .device(isCpu ? torch::kCPU : torch::kCUDA));
+        meta_->activeRanksTensor = at::ones(
+            {max_size}, torch::dtype(torch::kInt32)
+                            .device(isCpu ? torch::kCPU : torch::kCUDA));
+        if (max_size != size) {
+            meta_->activeRanksTensor.slice(0, size, max_size).fill_(0);
+        }
     }
     meta_->engine = engine_;
     meta_->store = store;
     meta_->backendIndex = backendIndex_;
     meta_->bufferBaseIndex = backendIndex_ * 10;
-    p2p_proxy_->BindMeta(meta_);
+    p2p_proxy_->bindMeta(meta_);
 
     connection_ctx_->bootstrapLocalPeer(localServerName_, rank_info);
     if (options_ && options_->isExtension_) {
@@ -294,6 +457,16 @@ MooncakeBackend::MooncakeBackend(
         connection_ctx_->waitUntilAllConnected();
     }
 
+    // Register a lightweight Backend shim so that PyTorch's P2P dispatch path
+    // (batch_isend_irecv → _get_backend → getBackend) can find a registered
+    // Backend for this ProcessGroup.  The shim delegates send/recv back to us.
+    auto deviceType = isCpu ? c10::DeviceType::CPU : c10::DeviceType::CUDA;
+    auto shim = c10::make_intrusive<MooncakeP2PShim>(this);
+    setBackend(deviceType, BackendType::CUSTOM, shim);
+#ifndef MOONCAKE_EP_USE_MUSA
+    setDefaultBackend(BackendType::CUSTOM);
+#endif
+
     // Increment backend index
     ++backendIndex_;
 }
@@ -302,6 +475,35 @@ MooncakeBackend::~MooncakeBackend() { shutdown(); }
 
 const std::string MooncakeBackend::getBackendName() const { return "mooncake"; }
 
+// ---- MooncakeP2PShim implementation ----
+
+MooncakeP2PShim::MooncakeP2PShim(MooncakeBackend* owner)
+    : Backend(owner->getRank(), owner->getSize()), owner_(owner) {}
+
+const std::string MooncakeP2PShim::getBackendName() const { return "mooncake"; }
+
+c10::intrusive_ptr<c10d::Work> MooncakeP2PShim::send(
+    std::vector<at::Tensor>& tensors, int dstRank, int tag) {
+    return owner_->send(tensors, dstRank, tag);
+}
+
+c10::intrusive_ptr<c10d::Work> MooncakeP2PShim::recv(
+    std::vector<at::Tensor>& tensors, int srcRank, int tag) {
+    return owner_->recv(tensors, srcRank, tag);
+}
+
+c10::intrusive_ptr<c10d::Work> MooncakeP2PShim::recvAnysource(
+    std::vector<at::Tensor>& tensors, int tag) {
+    // MooncakeBackend doesn't implement recvAnysource; fall back to
+    // the base class which will raise a clear error.
+    return ::c10d::Backend::recvAnysource(tensors, tag);
+}
+
+c10::intrusive_ptr<c10d::Work> MooncakeP2PShim::barrier(
+    const c10d::BarrierOptions& opts) {
+    return owner_->barrier(opts);
+}
+
 c10::intrusive_ptr<c10d::Work> MooncakeBackend::send(
     std::vector<at::Tensor>& tensors, int dstRank, int tag) {
     connection_ctx_->waitUntilNewRanksConnected();
@@ -311,11 +513,12 @@ c10::intrusive_ptr<c10d::Work> MooncakeBackend::send(
     auto tensor = tensors.back();
 
     TORCH_CHECK(meta_->store, "P2P send requires a valid Store.");
-    TORCH_CHECK(dstRank >= 0 && dstRank < size_,
+    TORCH_CHECK(dstRank >= 0 && dstRank < meta_->size,
                 "P2P send: dstRank out of range.");
 
     auto contiguous = tensor.contiguous();
-    auto completed = std::make_shared<std::atomic<bool>>(false);
+    auto status = std::make_shared<std::atomic<P2PProxy::OpStatus>>(
+        P2PProxy::OpStatus::kPending);
     cudaStream_t stream = nullptr;
     if (!isCpu_) {
         auto current_stream =
@@ -324,14 +527,14 @@ c10::intrusive_ptr<c10d::Work> MooncakeBackend::send(
     }
 
     TORCH_CHECK(p2p_proxy_, "P2P send proxy is not initialized.");
-    p2p_proxy_->EnqueueSend(P2PProxy::SendOp{
+    p2p_proxy_->enqueueSend(P2PProxy::SendOp{
         .tensor_ = std::move(contiguous),
         .peer_rank_ = dstRank,
         .cuda_stream_ = stream,
-        .completed_ = completed,
+        .status_ = status,
     });
 
-    return c10::make_intrusive<MooncakeP2PWork>(completed);
+    return c10::make_intrusive<MooncakeP2PWork>(status);
 }
 
 c10::intrusive_ptr<c10d::Work> MooncakeBackend::recv(
@@ -343,11 +546,12 @@ c10::intrusive_ptr<c10d::Work> MooncakeBackend::recv(
     auto tensor = tensors.back();
 
     TORCH_CHECK(meta_->store, "P2P recv requires a valid Store.");
-    TORCH_CHECK(srcRank >= 0 && srcRank < size_,
+    TORCH_CHECK(srcRank >= 0 && srcRank < meta_->size,
                 "P2P recv: srcRank out of range.");
 
     auto target = tensor.is_contiguous() ? tensor : tensor.contiguous();
-    auto completed = std::make_shared<std::atomic<bool>>(false);
+    auto status = std::make_shared<std::atomic<P2PProxy::OpStatus>>(
+        P2PProxy::OpStatus::kPending);
     cudaStream_t stream = nullptr;
     if (!isCpu_) {
         auto current_stream =
@@ -356,15 +560,15 @@ c10::intrusive_ptr<c10d::Work> MooncakeBackend::recv(
     }
 
     TORCH_CHECK(p2p_proxy_, "P2P recv proxy is not initialized.");
-    p2p_proxy_->EnqueueRecv(P2PProxy::RecvOp{
+    p2p_proxy_->enqueueRecv(P2PProxy::RecvOp{
         .tensor_ = target,
         .original_tensor_ = tensor,
         .peer_rank_ = srcRank,
         .cuda_stream_ = stream,
-        .completed_ = completed,
+        .status_ = status,
     });
 
-    return c10::make_intrusive<MooncakeP2PWork>(completed);
+    return c10::make_intrusive<MooncakeP2PWork>(status);
 }
 
 c10::intrusive_ptr<c10d::Work> MooncakeBackend::broadcast(
@@ -493,15 +697,15 @@ c10::intrusive_ptr<c10d::Work> MooncakeBackend::_allgather_base(
     const c10d::AllgatherOptions& opts) {
     size_t tensorSize = inputBuffer.numel() * inputBuffer.element_size();
     if (isCpu_) {
-        auto numRanks = meta_->size;
         return worker_->putTaskCpu(
             c10d::OpType::_ALLGATHER_BASE, tensorSize, 0, meta_,
             connection_ctx_,
             [=](void* dst, size_t pos, size_t realSize) {
                 memcpy(dst, (char*)inputBuffer.data_ptr() + pos, realSize);
             },
-            [=](void* src, size_t pos, size_t realSize) {
-                for (const auto j : c10::irange(numRanks)) {
+            [=, this](void* src, size_t pos, size_t realSize) {
+                for (int j = 0; j < meta_->size; ++j) {
+                    if (!meta_->activeRanks[j]) continue;
                     memcpy(
                         (char*)outputBuffer.data_ptr() + j * tensorSize + pos,
                         (char*)src + j * realSize, realSize);
@@ -520,7 +724,8 @@ c10::intrusive_ptr<c10d::Work> MooncakeBackend::_allgather_base(
             },
             [=, this](void* src, size_t pos, size_t realSize,
                       const at::cuda::CUDAStream& enq_stream) {
-                for (const auto j : c10::irange(meta_->size)) {
+                for (int j = 0; j < meta_->size; ++j) {
+                    if (!meta_->activeRanks[j]) continue;
                     cudaMemcpyAsync(
                         (char*)outputBuffer.data_ptr() + j * tensorSize + pos,
                         (char*)src + j * realSize, realSize,
@@ -539,8 +744,9 @@ c10::intrusive_ptr<c10d::Work> MooncakeBackend::_reduce_scatter_base(
         return worker_->putTaskCpu(
             c10d::OpType::_REDUCE_SCATTER_BASE, tensorSize, 0, meta_,
             connection_ctx_,
-            [=](void* dst, size_t pos, size_t realSize) {
-                for (const auto j : c10::irange(numRanks)) {
+            [=, this](void* dst, size_t pos, size_t realSize) {
+                for (int j = 0; j < meta_->size; ++j) {
+                    if (!meta_->activeRanks[j]) continue;
                     memcpy((char*)dst + j * realSize,
                            (char*)inputBuffer.data_ptr() + j * tensorSize + pos,
                            realSize);
@@ -559,7 +765,8 @@ c10::intrusive_ptr<c10d::Work> MooncakeBackend::_reduce_scatter_base(
             connection_ctx_, stream,
             [=, this](void* dst, size_t pos, size_t realSize,
                       const at::cuda::CUDAStream& enq_stream) {
-                for (const auto j : c10::irange(meta_->size)) {
+                for (int j = 0; j < meta_->size; ++j) {
+                    if (!meta_->activeRanks[j]) continue;
                     cudaMemcpyAsync(
                         (char*)dst + j * realSize,
                         (char*)inputBuffer.data_ptr() + j * tensorSize + pos,
@@ -805,7 +1012,7 @@ void MooncakeBackend::shutdown() {
 
     // Phase 1: Drain P2P tasks
     p2p_device_worker_->removeProxy(p2p_proxy_);
-    has_hung_operation |= !p2p_proxy_->DrainTasks();
+    has_hung_operation |= !p2p_proxy_->drainTasks();
 
     // Phase 2: Drain collective tasks for this backend
     has_hung_operation |= !worker_->drainTasks(meta_.get());
@@ -825,7 +1032,7 @@ void MooncakeBackend::shutdown() {
 
     // Phase 5: Release resources if no hung operations
     if (has_hung_operation) {
-        p2p_proxy_->AbandonResources();
+        p2p_proxy_->abandonResources();
         connection_ctx_->abandonResources();
     }
 
@@ -845,6 +1052,13 @@ void MooncakeBackend::shutdown() {
                 cudaFree(recv_buffer_[i]);
             }
         }
+        if (isCpu_) {
+            delete[] meta_->activeRanks;
+        } else {
+            cudaFreeHost(meta_->activeRanks);
+        }
+        meta_->activeRanks = nullptr;
+        meta_->activeRanksDevice = nullptr;
     }
 }
 
@@ -896,25 +1110,43 @@ void MooncakeBackend::setLocalOnlyActiveRanks() {
 void MooncakeBackend::waitForExtensionState() {
     TORCH_CHECK(meta_->store, "Recovery join requires a valid Store.");
 
-    auto task_count_key = ConnectionContext::getExtensionTaskCountStoreKey(
-        meta_->backendIndex, rank_);
-    auto active_ranks_key = ConnectionContext::getExtensionActiveRanksStoreKey(
+    auto state_key = ConnectionContext::getExtensionStateStoreKey(
         meta_->backendIndex, rank_);
 
     BackoffWaiter waiter(
         BackoffWaiterConfig::constantSleep(std::chrono::milliseconds(50)));
 
-    waiter.wait([&] {
-        return meta_->store->check({task_count_key, active_ranks_key});
-    });
+    waiter.wait([&] { return meta_->store->check({state_key}); });
 
-    auto task_count_data = meta_->store->get(task_count_key);
-    std::string task_count(task_count_data.begin(), task_count_data.end());
-    meta_->taskCount = std::stoi(task_count);
+    auto state_data = meta_->store->get(state_key);
+    auto state = deserialize(state_data);
 
-    auto active_ranks = meta_->store->get(active_ranks_key);
-    deserializeActiveRanks(active_ranks, meta_->activeRanks, meta_->size);
+    // taskCount
+    meta_->taskCount = state.taskCount;
+
+    // p2pEpochs
+    TORCH_CHECK(static_cast<size_t>(meta_->size) == state.p2pEpochs.size(),
+                "Invalid p2pEpochs size");
+    for (int i = 0; i < meta_->size; ++i) {
+        p2p_proxy_->setEpoch(i, state.p2pEpochs[i]);
+    }
+
+    // activeRanks
+    TORCH_CHECK(static_cast<size_t>(meta_->size) == state.activeRanks.size(),
+                "Invalid activeRanks");
+    for (int i = 0; i < meta_->size; ++i) {
+        meta_->activeRanks[i] = state.activeRanks[i];
+    }
     syncActiveRanksTensor();
+
+    // activeSize: count the number of active ranks (contiguous from 0)
+    int newActiveSize = 0;
+    for (int i = 0; i < meta_->size; ++i) {
+        if (meta_->activeRanks[i]) {
+            newActiveSize = i + 1;
+        }
+    }
+    meta_->activeSize = newActiveSize;
 }
 
 int MooncakeBackend::getNumSyncedRanks() {
@@ -937,6 +1169,7 @@ int MooncakeBackend::getNumSyncedRanks() {
 
 void MooncakeBackend::extendGroupSizeTo(int newSize) {
     const int oldSize = meta_->size;
+    const int oldActiveSize = meta_->activeSize;
     if (newSize == oldSize) return;
 
     TORCH_CHECK(newSize >= 0 && static_cast<size_t>(newSize) < kMaxNumRanks,
@@ -947,19 +1180,29 @@ void MooncakeBackend::extendGroupSizeTo(int newSize) {
               << ": Group size extend to " << newSize;
 
     meta_->size = newSize;
+    meta_->activeSize = newSize;
     meta_->taskCount = 0;
 
     // Initialize new rank's metadata
     for (int i = oldSize; i < newSize; ++i) {
         local2global_rank_map_[i] = i;
-        meta_->activeRanks[i] = true;
+        // IMPORTANT: Newly-extended ranks must start as inactive.
+        // They will only participate in collectives after healthy ranks
+        // explicitly activate them via recoverRanks(). This enables a
+        // two-phase scale-up protocol (extend capacity -> poll readiness
+        // -> recover/activate) and avoids collectives including ranks that
+        // haven't joined yet.
+        meta_->activeRanks[i] = false;
     }
 
     auto& tensor = meta_->activeRanksTensor;
-    tensor.resize_({newSize});
-    tensor.slice(0, oldSize, newSize).fill_(1);
+    if (newSize > tensor.numel()) {
+        tensor.resize_({newSize});
+    }
+    tensor.slice(0, oldSize, newSize).fill_(0);
 
     connection_ctx_->extendGroupSizeTo(newSize);
+    p2p_proxy_->extendGroupSizeTo(newSize);
     // After extendGroupSizeTo, we don't `waitUntilNewRanksConnected` here
     // but do it in the first task. This enables client code to overlap
     // execution between `extendGroupSizeTo` and the first communication call.
@@ -1020,16 +1263,29 @@ void MooncakeBackend::recoverRanks(const std::vector<int>& ranks) {
         meta_->activeRanks[rank] = true;
     }
 
+    // Expand activeSize if any recovered rank is beyond the current boundary.
+    if (!ranks.empty()) {
+        const int max_rank = *std::max_element(ranks.begin(), ranks.end());
+        if (max_rank >= meta_->activeSize) {
+            meta_->activeSize = max_rank + 1;
+        }
+    }
+
     syncActiveRanksTensor();
-    auto active_ranks_snapshot =
-        serializeActiveRanks(meta_->activeRanks, meta_->size);
+    std::vector<uint32_t> epochs(meta_->size);
+    for (int i = 0; i < meta_->size; ++i) {
+        epochs[i] = p2p_proxy_->getEpoch(i);
+    }
+    ExtensionState state{
+        .activeRanks =
+            std::vector(meta_->activeRanks, meta_->activeRanks + meta_->size),
+        .p2pEpochs = std::move(epochs),
+        .taskCount = meta_->taskCount};
+    auto state_data = serialize(state);
     for (const int rank : ranks) {
-        meta_->store->set(ConnectionContext::getExtensionTaskCountStoreKey(
-                              meta_->backendIndex, rank),
-                          std::to_string(meta_->taskCount));
-        meta_->store->set(ConnectionContext::getExtensionActiveRanksStoreKey(
-                              meta_->backendIndex, rank),
-                          active_ranks_snapshot);
+        auto key = ConnectionContext::getExtensionStateStoreKey(
+            meta_->backendIndex, rank);
+        meta_->store->set(key, state_data);
     }
 }
 
@@ -1045,4 +1301,12 @@ void MooncakeBackend::joinGroup() {
     connection_ctx_->waitUntilAllConnected();
     waitForExtensionState();
 }
+
+void MooncakeBackend::setExternalEngine(TransferEngine* engine) {
+    externalEngine_ = engine;
+    if (engine) {
+        LOG(INFO) << "MooncakeBackend: external TransferEngine set (ptr="
+                  << engine << ")";
+    }
+}
 }  // namespace mooncake
diff --git a/mooncake-pg/src/mooncake_worker.cu b/mooncake-pg/src/mooncake_worker.cu
index 10e7aa56..127ecbcb 100644
--- a/mooncake-pg/src/mooncake_worker.cu
+++ b/mooncake-pg/src/mooncake_worker.cu
@@ -1,163 +1,21 @@
-#include <mooncake_backend.h>
-#include <cstdio>
-#include <memory>
-#include <thread>
-#include <mooncake_worker.cuh>
-#include <ATen/cuda/CUDAGraphsUtils.cuh>
+// mooncake_worker.cu — GPU kernel functions and launch wrappers.
+// Compiled by nvcc (CUDA) and mcc (MUSA, via mooncake_worker.mu symlink).
+// The __MUSA__ branch avoids torch headers to stay compatible with mcc.
 
-#include "pg_utils.h"
+#include <mooncake_worker_kernels.cuh>
+
+#ifdef __MUSA__
+#include <musa_bf16.h>
+#endif
 
 namespace mooncake {
 
-class MooncakeWorkCpu : public ::c10d::Work {
-   public:
-    MooncakeWorkCpu(c10d::OpType opType,
-                    c10::intrusive_ptr<c10::ivalue::Future> future,
-                    std::shared_ptr<TransferGroupMeta> meta)
-        : Work(-1, opType),
-          future_(std::move(future)),
-          meta_(std::move(meta)) {}
+// ── Kernel functions ──────────────────────────────────────────────
+// Both CUDA and MUSA share the same kernel bodies. Parameters use plain
+// C++ types (int instead of c10d::OpType / c10d::ReduceOp::RedOpType)
+// so that mcc can compile them without torch headers.
 
-    bool isCompleted() override { return future_->completed(); }
-
-    bool wait(std::chrono::milliseconds timeout) override {
-        future_->wait();
-        return future_->completed() && !future_->hasError();
-    }
-
-   private:
-    c10::intrusive_ptr<c10::ivalue::Future> future_;
-    std::shared_ptr<TransferGroupMeta> meta_;
-};
-
-class MooncakeWorkCuda : public ::c10d::Work {
-   public:
-    MooncakeWorkCuda(c10d::OpType opType, std::shared_ptr<torch::Event> event,
-                     std::shared_ptr<TransferGroupMeta> meta,
-                     const MooncakeWorker* worker,
-                     std::vector<CudaTaskSubmissionToken> submitted_tasks)
-        : Work(-1, opType),
-          event_(std::move(event)),
-          meta_(std::move(meta)),
-          worker_(worker),
-          submitted_tasks_(std::move(submitted_tasks)) {}
-
-    bool isCompleted() override { return event_->query(); }
-
-    bool wait(std::chrono::milliseconds timeout) override {
-        // Wait until the task has been submitted to TransferEngine:
-        // This tries to ensure that the CUDA kernels required for the transfer
-        // have been launched by the time `waitUntilTasksSubmitted` returns.
-        //
-        // Why is this needed? PyTorch documentation implies that collective
-        // operations should be enqueued when `wait()` returns. In practice, we
-        // found that violating this causes hangs.
-        //
-        // Our current hypothesis for the hang is: PyTorch assumes the kernels
-        // needed for the transfer are already launched when `wait` returns
-        // true. It may then launch subsequent operations after the collective
-        // (e.g., `.cpu()`). Such operations may acquire a process-wide lock in
-        // the CUDA runtime. Also, they may rely on the data produced by the
-        // collective, thus causing a synchronization on enq_stream. However,
-        // holding that runtime lock prevents cudaMemcpy(Async) in TE/TENT from
-        // launching. This means the transfer can't finish, and enq_stream won't
-        // complete. Thus, a deadlock occurs.
-        // (In practice, we found that replacing all cudaMemcpyAsync in TENT
-        // with cuMemcpyAsync actually alleviates this, which further suggests a
-        // deadlock in the CUDA runtime. However, that change is too invasive
-        // for TE/TENT, so we do not adopt it here.)
-        //
-        // Strictly speaking, the wait is needed for another reason: The current
-        // stream will be blocked on the event below. Any subsequent work on
-        // `current_stream` will wait on that event, which effectively waits for
-        // the task to be done. Therefore, we must ensure all kernels needed for
-        // the transfer task are launched BEFORE blocking the current stream, in
-        // case TE/TENT use `current_stream` to launch those kernels (though it
-        // is rare).
-        //
-        // Please note that this logic relies on the assumption that TE/TENT
-        // will launch all CUDA operations in `submitTransfer`.
-        // Unfortunately, TcpTransport in TE and TENT currently violates this
-        // assumption (cudaMemcpy(Async) may be called later from a callback),
-        // which can cause hangs in PG when a CUDA operation such as
-        // `x.cpu().item()` follows the collective. For TE's TcpTransport, the
-        // use of cudaMemcpy on the default stream may also contribute to the
-        // hang.
-        //
-        // Besides, for CPU-only transports (like RdmaTransport),
-        // waitUntilTasksSubmitted is totally unnecessary, but we keep it for
-        // uniform behavior to avoid invasive changes to TE/TENT.
-        bool submitted = true;
-        if (at::cuda::currentStreamCaptureStatus() ==
-            c10::cuda::CaptureStatus::None) {
-            // Normal execution: block until tasks are submitted.
-            submitted =
-                worker_->waitUntilTasksSubmitted(submitted_tasks_, timeout);
-        } else {
-            // During CUDA graph capture, kernels are recorded but not actually
-            // executed. The enqueueTaskKernel would never run, so
-            // waitUntilTasksSubmitted would hang because the CPU worker thread
-            // never sees task.active == true.
-            //
-            // Note that this also means NvlinkTransport (and TcpTransport too,
-            // of course) won't work with CUDA Graphs: Kernels launched inside
-            // TE/TENT can't be captured by the graph, and during replay they
-            // are not ordered with the graph execution. This may trigger the
-            // same deadlock described above.
-        }
-        if (!submitted) return false;
-
-        // Once all tasks have been submitted, use the event to synchronize
-        // the current stream and the enqueue stream, but do not wait on this
-        // event.
-        //
-        // See PyTorch docs for more details:
-        // https://docs.pytorch.org/docs/stable/distributed.html#synchronous-and-asynchronous-collective-operations
-        //   "wait() - in the case of CPU collectives, will block the process
-        //    until the operation is completed. In the case of CUDA collectives,
-        //    will block the currently active CUDA stream until the operation
-        //    is completed (but will not block the CPU)."
-        auto current_stream = at::cuda::getCurrentCUDAStream();
-        event_->block(current_stream);
-        return true;
-    }
-
-   protected:
-    std::shared_ptr<torch::Event> event_;
-    std::shared_ptr<TransferGroupMeta> meta_;
-    const MooncakeWorker* worker_;
-    std::vector<CudaTaskSubmissionToken> submitted_tasks_;
-};
-
-class MooncakeBarrierWorkCuda : public MooncakeWorkCuda {
-   public:
-    using MooncakeWorkCuda::MooncakeWorkCuda;
-
-    bool wait(std::chrono::milliseconds timeout) override {
-        // Skip host-side synchronization during CUDA graph capture.
-        // cudaEventSynchronize is not permitted while a stream is capturing.
-        if (at::cuda::currentStreamCaptureStatus() !=
-            c10::cuda::CaptureStatus::None) {
-            // We still need stream-level synchronization so that subsequent
-            // operations on the capture stream are ordered after the barrier
-            // task on the enqueue stream.
-            auto current_stream = at::cuda::getCurrentCUDAStream();
-            event_->block(current_stream);
-            return true;
-        }
-
-        if (timeout == kNoTimeout) {
-            event_->synchronize();
-            return true;
-        }
-
-        BackoffWaiter waiter(
-            BackoffWaiterConfig::constantSleep(std::chrono::microseconds(10)));
-        return waiter.wait_for(timeout, [this] { return event_->query(); });
-    }
-};
-
-__global__ void enqueueTaskKernel(c10d::OpType opType, size_t tensorSize,
+__global__ void enqueueTaskKernel(int opType, size_t tensorSize,
                                   int64_t broadcastRoot, int bufferOffset,
                                   uint64_t submitSequence, void* meta,
                                   Task* tasks, int numRanks,
@@ -186,371 +44,173 @@ __global__ void enqueueTaskKernel(c10d::OpType opType, size_t tensorSize,
 
 template <typename scalar_t>
 __global__ void reduceKernel(scalar_t* dst, const scalar_t* src,
-                             size_t numElements, size_t numRanks,
-                             c10d::ReduceOp::RedOpType op, bool* activeRanks) {
+                             size_t numElements, size_t numRanks, int op,
+                             bool* activeRanks) {
     size_t thread_idx = blockIdx.x * blockDim.x + threadIdx.x;
     size_t stride = blockDim.x * gridDim.x;
+
+#ifdef __MUSA__
+    // mt_bfloat16 lacks arithmetic and comparison operators; use float
+    // accumulator and fminf/fmaxf for that type only.  Other types
+    // (float, double, int64_t, etc.) use their native accumulator to
+    // preserve full precision.
+    constexpr bool kIsBf16 = std::is_same_v<scalar_t, mt_bfloat16>;
+#else
+    constexpr bool kIsBf16 = false;
+#endif
+    using acc_t = std::conditional_t<kIsBf16, float, scalar_t>;
+
     for (size_t elem_idx = thread_idx; elem_idx < numElements;
          elem_idx += stride) {
         bool valid = false;
-        scalar_t acc = 0;
+        acc_t acc = 0;
         for (size_t rank = 0; rank < numRanks; ++rank) {
             if (activeRanks[rank]) {
                 if (!valid) {
-                    acc = src[rank * numElements + elem_idx];
+                    if constexpr (kIsBf16) {
+                        acc = (float)src[rank * numElements + elem_idx];
+                    } else {
+                        acc = src[rank * numElements + elem_idx];
+                    }
                     valid = true;
                 } else {
-                    switch (op) {
-                        case c10d::ReduceOp::SUM:
-                            acc += src[rank * numElements + elem_idx];
-                            break;
-                        case c10d::ReduceOp::MIN:
-                            acc = std::min(src[rank * numElements + elem_idx],
-                                           acc);
-                            break;
-                        case c10d::ReduceOp::MAX:
-                            acc = std::max(src[rank * numElements + elem_idx],
-                                           acc);
-                            break;
-                        case c10d::ReduceOp::PRODUCT:
-                            acc *= src[rank * numElements + elem_idx];
-                            break;
-                        default:
-                            // never
-                    }
-                }
-            }
-        }
-        dst[elem_idx] = acc;
-    }
-}
-
-namespace {
-
-template <typename scalar_t>
-void preload_reduce_kernel(const char* name) {
-    cudaFuncAttributes attr{};
-    auto err = cudaFuncGetAttributes(
-        &attr, reinterpret_cast<const void*>(reduceKernel<scalar_t>));
-    TORCH_CHECK(err == cudaSuccess, "Failed to preload kernel ", name, ": ",
-                cudaGetErrorString(err));
-}
-
-}  // namespace
-
-void launchReduceKernel(at::Tensor dst, size_t pos, size_t realSize, void* src,
-                        size_t numRanks, c10d::ReduceOp op, bool* activeRanks,
-                        cudaStream_t stream) {
-    TORCH_CHECK(op == c10d::ReduceOp::SUM || op == c10d::ReduceOp::MIN ||
-                    op == c10d::ReduceOp::MAX || op == c10d::ReduceOp::PRODUCT,
-                "Only support SUM/MIN/MAX/PRODUCT for reduction.");
-    auto ptr = (char*)dst.data_ptr() + pos;
-    size_t num = realSize / dst.element_size();
-
-    switch (dst.scalar_type()) {
-        case c10::kByte:
-            reduceKernel<<<64, 256, 0, stream>>>((uint8_t*)ptr, (uint8_t*)src,
-                                                 num, numRanks, op.op_,
-                                                 activeRanks);
-            break;
-        case c10::kChar:
-            reduceKernel<<<64, 256, 0, stream>>>(
-                (int8_t*)ptr, (int8_t*)src, num, numRanks, op.op_, activeRanks);
-            break;
-        case c10::kShort:
-            reduceKernel<<<64, 256, 0, stream>>>((int16_t*)ptr, (int16_t*)src,
-                                                 num, numRanks, op.op_,
-                                                 activeRanks);
-            break;
-        case c10::kInt:
-            reduceKernel<<<64, 256, 0, stream>>>((int*)ptr, (int*)src, num,
-                                                 numRanks, op.op_, activeRanks);
-            break;
-        case c10::kLong:
-            reduceKernel<<<64, 256, 0, stream>>>((int64_t*)ptr, (int64_t*)src,
-                                                 num, numRanks, op.op_,
-                                                 activeRanks);
-            break;
-        case c10::kFloat:
-            reduceKernel<<<64, 256, 0, stream>>>((float*)ptr, (float*)src, num,
-                                                 numRanks, op.op_, activeRanks);
-            break;
-        case c10::kDouble:
-            reduceKernel<<<64, 256, 0, stream>>>(
-                (double*)ptr, (double*)src, num, numRanks, op.op_, activeRanks);
-            break;
-        case c10::kBool:
-            reduceKernel<<<64, 256, 0, stream>>>((bool*)ptr, (bool*)src, num,
-                                                 numRanks, op.op_, activeRanks);
-            break;
-        case c10::kBFloat16:
-            reduceKernel<<<64, 256, 0, stream>>>((at::BFloat16*)ptr,
-                                                 (at::BFloat16*)src, num,
-                                                 numRanks, op.op_, activeRanks);
-            break;
-        default:
-            TORCH_CHECK(false, c10::str("Unsupported reduce dtype: ",
-                                        dst.scalar_type()));
-    }
-}
-
-template <typename T>
-T applyReduceOp(const T& a, const T& b, c10d::ReduceOp op) {
-    switch (op) {
-        case c10d::ReduceOp::SUM:
-            return a + b;
-        case c10d::ReduceOp::PRODUCT:
-            return a * b;
-        case c10d::ReduceOp::MIN:
-            return std::min(a, b);
-        case c10d::ReduceOp::MAX:
-            return std::max(a, b);
-        default:
-            TORCH_CHECK(false, c10::str("Unsupported reduce op: ", op));
-    }
-}
-
-template <typename T>
-void reduceCpu(T* dst, const T* src, size_t numElements, size_t numRanks,
-               c10d::ReduceOp op, bool* activeRanks) {
-    at::parallel_for(0, numElements, 1024, [&](int64_t begin, int64_t end) {
-        for (int64_t i = begin; i < end; ++i) {
-            bool valid = false;
-            T acc{};
-            for (int64_t rank = 0; rank < numRanks; ++rank) {
-                if (activeRanks[rank]) {
-                    if (!valid) {
-                        acc = src[i + rank * numElements];
-                        valid = true;
+                    if constexpr (kIsBf16) {
+                        float val = (float)src[rank * numElements + elem_idx];
+                        switch (op) {
+                            case 0:  // SUM
+                                acc += val;
+                                break;
+                            case 2:  // PRODUCT
+                                acc *= val;
+                                break;
+                            case 3:  // MIN
+                                acc = fminf(acc, val);
+                                break;
+                            case 4:  // MAX
+                                acc = fmaxf(acc, val);
+                                break;
+                            default:
+                                break;
+                        }
                     } else {
-                        acc =
-                            applyReduceOp(acc, src[i + rank * numElements], op);
+                        switch (op) {
+                            case 0:  // SUM
+                                acc += src[rank * numElements + elem_idx];
+                                break;
+                            case 2:  // PRODUCT
+                                acc *= src[rank * numElements + elem_idx];
+                                break;
+                            case 3:  // MIN
+                                acc = std::min(
+                                    src[rank * numElements + elem_idx], acc);
+                                break;
+                            case 4:  // MAX
+                                acc = std::max(
+                                    src[rank * numElements + elem_idx], acc);
+                                break;
+                            default:
+                                break;
+                        }
                     }
                 }
             }
-            dst[i] = acc;
         }
-    });
-}
-
-void launchReduceCpu(at::Tensor dst, size_t pos, size_t realSize, void* src,
-                     size_t numRanks, c10d::ReduceOp op, bool* activeRanks) {
-    auto ptr = (char*)dst.data_ptr() + pos;
-    size_t num = realSize / dst.element_size();
-
-    switch (dst.scalar_type()) {
-        case c10::kByte:
-            reduceCpu((uint8_t*)ptr, (uint8_t*)src, num, numRanks, op,
-                      activeRanks);
-            break;
-        case c10::kChar:
-            reduceCpu((int8_t*)ptr, (int8_t*)src, num, numRanks, op,
-                      activeRanks);
-            break;
-        case c10::kShort:
-            reduceCpu((int16_t*)ptr, (int16_t*)src, num, numRanks, op,
-                      activeRanks);
-            break;
-        case c10::kInt:
-            reduceCpu((int*)ptr, (int*)src, num, numRanks, op, activeRanks);
-            break;
-        case c10::kLong:
-            reduceCpu((int64_t*)ptr, (int64_t*)src, num, numRanks, op,
-                      activeRanks);
-            break;
-        case c10::kFloat:
-            reduceCpu((float*)ptr, (float*)src, num, numRanks, op, activeRanks);
-            break;
-        case c10::kDouble:
-            reduceCpu((double*)ptr, (double*)src, num, numRanks, op,
-                      activeRanks);
-            break;
-        case c10::kBool:
-            reduceCpu((bool*)ptr, (bool*)src, num, numRanks, op, activeRanks);
-            break;
-        default:
-            TORCH_CHECK(false, c10::str("Unsupported reduce dtype: ",
-                                        dst.scalar_type()));
-    }
-}
-
-void preloadReduceKernels() {
-    preload_reduce_kernel<uint8_t>("reduceKernel<uint8_t>");
-    preload_reduce_kernel<int8_t>("reduceKernel<int8_t>");
-    preload_reduce_kernel<int16_t>("reduceKernel<int16_t>");
-    preload_reduce_kernel<int>("reduceKernel<int>");
-    preload_reduce_kernel<int64_t>("reduceKernel<int64_t>");
-    preload_reduce_kernel<float>("reduceKernel<float>");
-    preload_reduce_kernel<double>("reduceKernel<double>");
-    preload_reduce_kernel<bool>("reduceKernel<bool>");
-    preload_reduce_kernel<at::BFloat16>("reduceKernel<BFloat16>");
-}
-
-MooncakeWorker::MooncakeWorker(int cuda_device_index)
-    : cuda_device_index_(cuda_device_index) {
-    int deviceCount = 0;
-    cudaError err = cudaGetDeviceCount(&deviceCount);
-    if (!err && deviceCount > 0) {
-        // Pin memory for task array
-        cudaHostAlloc(&tasks_, kNumTasks_ * sizeof(Task), cudaHostAllocMapped);
-        cudaHostGetDevicePointer(&tasks_device_, tasks_, 0);
-    } else {
-        LOG(WARNING) << "No CUDA device found. Only the `mooncake-cpu` backend "
-                        "can be used.";
-        tasks_ = new Task[kNumTasks_];
-    }
-    for (size_t i = 0; i < kNumTasks_; ++i) {
-        tasks_[i].active = false;
-        tasks_[i].submitSequence = 0;
-        submitted_task_sequence_[i].store(0, std::memory_order_relaxed);
-    }
-}
-
-MooncakeWorker::~MooncakeWorker() {
-    running_ = false;
-    if (worker_thread_.joinable()) {
-        worker_thread_.join();
-    }
-}
-
-c10::intrusive_ptr<c10d::Work> MooncakeWorker::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)>&
-        bufferToTensor) {
-    connection_ctx->waitUntilNewRanksConnected();
-
-    size_t chunkSize = ((kBufferSize - 1) / meta->size) & ~(size_t)7;
-    auto future = c10::make_intrusive<c10::ivalue::Future>(
-        c10::ListType::create(c10::TensorType::get()));
-
-    struct IterState {
-        size_t currentPos = 0;
-    };
-    auto state = std::make_shared<IterState>();
-
-    auto processNextChunk = std::make_shared<std::function<void()>>();
-    std::weak_ptr<std::function<void()>> weakProcessNextChunk =
-        processNextChunk;
-
-    *processNextChunk = [this, weakProcessNextChunk, state, opType, tensorSize,
-                         chunkSize, broadcastRoot, meta, tensorToBuffer,
-                         bufferToTensor, future]() {
-        auto processNextChunk = weakProcessNextChunk.lock();
-
-        if (state->currentPos >= tensorSize) {
-            future->markCompleted(c10::IValue());
-            return;
+        if constexpr (kIsBf16) {
+            dst[elem_idx] = (scalar_t)acc;
+        } else {
+            dst[elem_idx] = acc;
         }
-
-        int taskId = cpuTaskCount % 2;
-        TORCH_CHECK(!tasks_[taskId].active);
-
-        size_t realSize = std::min(chunkSize, tensorSize - state->currentPos);
-        int bufferOffset = meta->taskCount % 2;
-
-        tasks_[taskId].opType = opType;
-        tasks_[taskId].tensorSize = realSize;
-        tasks_[taskId].broadcastRoot = broadcastRoot;
-        tasks_[taskId].bufferOffset = bufferOffset;
-        tasks_[taskId].transferGroupMeta = meta.get();
-        tensorToBuffer(
-            (void*)meta->segmentInfos[meta->rank].send_buffer[bufferOffset],
-            state->currentPos, realSize);
-
-        hasCallback_[taskId] = true;
-
-        callbacks_[taskId] = [this, processNextChunk, state, meta,
-                              bufferToTensor, bufferOffset, realSize,
-                              future]() {
-            for (int i = 0; i < meta->size; ++i) {
-                meta->activeRanksTensor[i] = meta->activeRanks[i] ? 1 : 0;
-            }
-            bufferToTensor(
-                (void*)meta->segmentInfos[meta->rank].recv_buffer[bufferOffset],
-                state->currentPos, realSize);
-
-            state->currentPos += realSize;
-
-            (*processNextChunk)();
-        };
-
-        tasks_[taskId].active = true;
-        ++cpuTaskCount;
-        ++meta->taskCount;
-    };
-
-    (*processNextChunk)();
-
-    return c10::make_intrusive<MooncakeWorkCpu>(opType, future, meta);
-}
-
-c10::intrusive_ptr<c10d::Work> MooncakeWorker::putTaskCuda(
-    c10d::OpType opType, size_t tensorSize, int64_t broadcastRoot,
-    const std::shared_ptr<TransferGroupMeta>& meta,
-    const std::shared_ptr<ConnectionContext>& connection_ctx,
-    const at::cuda::CUDAStream& issue_stream,
-    const std::function<void(void* dst, size_t pos, size_t realSize,
-                             const at::cuda::CUDAStream&)>& tensorToBuffer,
-    const std::function<void(void* src, size_t pos, size_t realSize,
-                             const at::cuda::CUDAStream&)>& bufferToTensor) {
-    connection_ctx->waitUntilNewRanksConnected();
-
-    // TORCH_CHECK(tensorSize * meta->size < kBufferSize, "Too large!");
-    //  Alternately use even-odd items to maintain tasks
-    size_t chunkSize = ((kBufferSize - 1) / meta->size) & ~(size_t)7;
-
-    // Get a non-blocking stream for enqueue:
-    //   The incoming `issue_stream` may be the Null Stream, which enforces
-    //   implicit synchronization semantics. Launching a spin-wait kernel
-    //   (enqueueTaskKernel) on such a stream can introduce potential deadlock.
-    at::cuda::CUDAStream enq_stream =
-        at::cuda::getStreamFromPool(false, issue_stream.device_index());
-
-    // Synchronize: enq_stream waits for issue_stream
-    auto event_start = std::make_shared<torch::Event>(torch::kCUDA);
-    event_start->record(issue_stream);
-    event_start->block(enq_stream);
-
-    std::vector<CudaTaskSubmissionToken> submitted_tasks;
-    submitted_tasks.reserve((tensorSize + chunkSize - 1) / chunkSize);
-    for (size_t pos = 0; pos < tensorSize; pos += chunkSize) {
-        size_t realSize = min(tensorSize, pos + chunkSize) - pos;
-        int taskId = cudaTaskCount % 2 + 2;
-        int bufferOffset = meta->taskCount % 2;
-        const uint64_t taskSequence =
-            next_cuda_task_sequence_.fetch_add(1, std::memory_order_relaxed);
-        submitted_tasks.push_back(
-            {.task_id = static_cast<size_t>(taskId), .sequence = taskSequence});
-        tensorToBuffer(
-            (void*)meta->segmentInfos[meta->rank].send_buffer[bufferOffset],
-            pos, realSize, enq_stream);
-
-        hasCallback_[taskId] = false;
-        enqueueTaskKernel<<<1, 1, 0, enq_stream>>>(
-            opType, realSize, broadcastRoot, bufferOffset, taskSequence,
-            meta.get(), tasks_device_, meta->size, meta->activeRanksDevice,
-            meta->activeRanksTensor.data_ptr<int>(), taskId);
-        bufferToTensor(
-            (void*)meta->segmentInfos[meta->rank].recv_buffer[bufferOffset],
-            pos, realSize, enq_stream);
-
-        ++cudaTaskCount;
-        ++meta->taskCount;
     }
-
-    auto event_end = std::make_shared<torch::Event>(torch::kCUDA);
-    event_end->record(enq_stream);
-
-    if (opType == c10d::OpType::BARRIER) {
-        return c10::make_intrusive<MooncakeBarrierWorkCuda>(
-            opType, event_end, meta, this, std::move(submitted_tasks));
-    }
-    return c10::make_intrusive<MooncakeWorkCuda>(opType, event_end, meta, this,
-                                                 std::move(submitted_tasks));
+}
+
+}  // namespace mooncake
+
+// ── Kernel launch wrappers ────────────────────────────────────────
+// g++ cannot compile <<<>>> syntax, so these wrappers are compiled by
+// the GPU compiler (nvcc/mcc) and provide plain C++ functions that the
+// host code can call.
+
+#ifdef __MUSA__
+#include <musa_runtime.h>
+#endif
+
+namespace mooncake {
+
+void launchEnqueueTaskKernel(int opType, size_t tensorSize,
+                             int64_t broadcastRoot, int bufferOffset,
+                             uint64_t submitSequence, void* meta, Task* tasks,
+                             int numRanks, const bool* activeRanks,
+                             int* activeRanksTensor, size_t taskId,
+                             cudaStream_t stream) {
+    enqueueTaskKernel<<<1, 1, 0, stream>>>(
+        opType, tensorSize, broadcastRoot, bufferOffset, submitSequence, meta,
+        tasks, numRanks, activeRanks, activeRanksTensor, taskId);
+}
+
+#define DEF_LAUNCH_REDUCE(scalar_t, suffix)                                   \
+    void launchReduceKernel_##suffix(                                         \
+        scalar_t* dst, const scalar_t* src, size_t numElements,               \
+        size_t numRanks, int op, bool* activeRanks, cudaStream_t stream) {    \
+        reduceKernel<<<64, 256, 0, stream>>>(dst, src, numElements, numRanks, \
+                                             op, activeRanks);                \
+    }
+
+DEF_LAUNCH_REDUCE(uint8_t, uint8)
+DEF_LAUNCH_REDUCE(int8_t, int8)
+DEF_LAUNCH_REDUCE(int16_t, int16)
+DEF_LAUNCH_REDUCE(int, int32)
+DEF_LAUNCH_REDUCE(int64_t, int64)
+DEF_LAUNCH_REDUCE(float, float)
+DEF_LAUNCH_REDUCE(double, double)
+DEF_LAUNCH_REDUCE(bool, bool)
+
+#undef DEF_LAUNCH_REDUCE
+
+void launchReduceKernel_bf16(void* dst, const void* src, size_t numElements,
+                             size_t numRanks, int op, bool* activeRanks,
+                             cudaStream_t stream) {
+#ifdef __MUSA__
+    reduceKernel<<<64, 256, 0, stream>>>((mt_bfloat16*)dst,
+                                         (const mt_bfloat16*)src, numElements,
+                                         numRanks, op, activeRanks);
+#else
+    reduceKernel<<<64, 256, 0, stream>>>((at::BFloat16*)dst,
+                                         (const at::BFloat16*)src, numElements,
+                                         numRanks, op, activeRanks);
+#endif
+}
+
+void preloadReduceKernels() {
+#ifdef __MUSA__
+    // MUSA: mcc has no cudaFuncGetAttributes, kernels are JIT-compiled on
+    // first use.
+#else
+    // CUDA: preload kernels to avoid JIT compilation overhead on first use.
+    auto preload = [](const char* name, auto kernel_ptr) {
+        cudaFuncAttributes attr{};
+        auto err = cudaFuncGetAttributes(&attr, kernel_ptr);
+        TORCH_CHECK(err == cudaSuccess, "Failed to preload kernel ", name, ": ",
+                    cudaGetErrorString(err));
+    };
+    preload("reduceKernel<uint8_t>",
+            reinterpret_cast<const void*>(reduceKernel<uint8_t>));
+    preload("reduceKernel<int8_t>",
+            reinterpret_cast<const void*>(reduceKernel<int8_t>));
+    preload("reduceKernel<int16_t>",
+            reinterpret_cast<const void*>(reduceKernel<int16_t>));
+    preload("reduceKernel<int>",
+            reinterpret_cast<const void*>(reduceKernel<int>));
+    preload("reduceKernel<int64_t>",
+            reinterpret_cast<const void*>(reduceKernel<int64_t>));
+    preload("reduceKernel<float>",
+            reinterpret_cast<const void*>(reduceKernel<float>));
+    preload("reduceKernel<double>",
+            reinterpret_cast<const void*>(reduceKernel<double>));
+    preload("reduceKernel<bool>",
+            reinterpret_cast<const void*>(reduceKernel<bool>));
+    preload("reduceKernel<BFloat16>",
+            reinterpret_cast<const void*>(reduceKernel<at::BFloat16>));
+#endif
 }
 
 }  // namespace mooncake
diff --git a/mooncake-pg/src/mooncake_worker_host.cpp b/mooncake-pg/src/mooncake_worker_host.cpp
new file mode 100644
index 00000000..98c6d994
--- /dev/null
+++ b/mooncake-pg/src/mooncake_worker_host.cpp
@@ -0,0 +1,459 @@
+// mooncake_worker_host.cpp — Host-side code for PG collectives.
+// Compiled by g++ for both CUDA and MUSA builds. Uses kernel launch wrappers
+// from mooncake_worker_kernels.cuh instead of <<<>>> syntax.
+
+#include <mooncake_backend.h>
+#include <cstdio>
+#include <memory>
+#include <thread>
+#include <mooncake_worker.cuh>
+#include <mooncake_worker_kernels.cuh>
+#include <ATen/cuda/CUDAGraphsUtils.cuh>
+
+#include "pg_utils.h"
+
+namespace mooncake {
+
+class MooncakeWorkCpu : public ::c10d::Work {
+   public:
+    MooncakeWorkCpu(c10d::OpType opType,
+                    c10::intrusive_ptr<c10::ivalue::Future> future,
+                    std::shared_ptr<TransferGroupMeta> meta)
+        : Work(-1, opType),
+          future_(std::move(future)),
+          meta_(std::move(meta)) {}
+
+    bool isCompleted() override { return future_->completed(); }
+
+    bool wait(std::chrono::milliseconds timeout) override {
+        future_->wait();
+        return future_->completed() && !future_->hasError();
+    }
+
+   private:
+    c10::intrusive_ptr<c10::ivalue::Future> future_;
+    std::shared_ptr<TransferGroupMeta> meta_;
+};
+
+class MooncakeWorkCuda : public ::c10d::Work {
+   public:
+    MooncakeWorkCuda(c10d::OpType opType, std::shared_ptr<torch::Event> event,
+                     std::shared_ptr<TransferGroupMeta> meta,
+                     const MooncakeWorker* worker,
+                     std::vector<CudaTaskSubmissionToken> submitted_tasks)
+        : Work(-1, opType),
+          event_(std::move(event)),
+          meta_(std::move(meta)),
+          worker_(worker),
+          submitted_tasks_(std::move(submitted_tasks)) {}
+
+    bool isCompleted() override { return event_->query(); }
+
+    bool wait(std::chrono::milliseconds timeout) override {
+        // Wait until the task has been submitted to TransferEngine:
+        // This tries to ensure that the CUDA kernels required for the transfer
+        // have been launched by the time `waitUntilTasksSubmitted` returns.
+        //
+        // Why is this needed? PyTorch documentation implies that collective
+        // operations should be enqueued when `wait()` returns. In practice, we
+        // found that violating this causes hangs.
+        //
+        // Our current hypothesis for the hang is: PyTorch assumes the kernels
+        // needed for the transfer are already launched when `wait` returns
+        // true. It may then launch subsequent operations after the collective
+        // (e.g., `.cpu()`). Such operations may acquire a process-wide lock in
+        // the CUDA runtime. Also, they may rely on the data produced by the
+        // collective, thus causing a synchronization on enq_stream. However,
+        // holding that runtime lock prevents cudaMemcpy(Async) in TE/TENT from
+        // launching. This means the transfer can't finish, and enq_stream won't
+        // complete. Thus, a deadlock occurs.
+        // (In practice, we found that replacing all cudaMemcpyAsync in TENT
+        // with cuMemcpyAsync actually alleviates this, which further suggests a
+        // deadlock in the CUDA runtime. However, that change is too invasive
+        // for TE/TENT, so we do not adopt it here.)
+        //
+        // Strictly speaking, the wait is needed for another reason: The current
+        // stream will be blocked on the event below. Any subsequent work on
+        // `current_stream` will wait on that event, which effectively waits for
+        // the task to be done. Therefore, we must ensure all kernels needed for
+        // the transfer task are launched BEFORE blocking the current stream, in
+        // case TE/TENT use `current_stream` to launch those kernels (though it
+        // is rare).
+        //
+        // Please note that this logic relies on the assumption that TE/TENT
+        // will launch all CUDA operations in `submitTransfer`.
+        // Unfortunately, TcpTransport in TE and TENT currently violates this
+        // assumption (cudaMemcpy(Async) may be called later from a callback),
+        // which can cause hangs in PG when a CUDA operation such as
+        // `x.cpu().item()` follows the collective. For TE's TcpTransport, the
+        // use of cudaMemcpy on the default stream may also contribute to the
+        // hang.
+        //
+        // Besides, for CPU-only transports (like RdmaTransport),
+        // waitUntilTasksSubmitted is totally unnecessary, but we keep it for
+        // uniform behavior to avoid invasive changes to TE/TENT.
+        bool submitted = true;
+        if (at::cuda::currentStreamCaptureStatus() ==
+            c10::cuda::CaptureStatus::None) {
+            // Normal execution: block until tasks are submitted.
+            submitted =
+                worker_->waitUntilTasksSubmitted(submitted_tasks_, timeout);
+        } else {
+            // During CUDA graph capture, kernels are recorded but not actually
+            // executed. The enqueueTaskKernel would never run, so
+            // waitUntilTasksSubmitted would hang because the CPU worker thread
+            // never sees task.active == true.
+            //
+            // Note that this also means NvlinkTransport (and TcpTransport too,
+            // of course) won't work with CUDA Graphs: Kernels launched inside
+            // TE/TENT can't be captured by the graph, and during replay they
+            // are not ordered with the graph execution. This may trigger the
+            // same deadlock described above.
+        }
+        if (!submitted) return false;
+
+        // Once all tasks have been submitted, use the event to synchronize
+        // the current stream and the enqueue stream, but do not wait on this
+        // event.
+        //
+        // See PyTorch docs for more details:
+        // https://docs.pytorch.org/docs/stable/distributed.html#synchronous-and-asynchronous-collective-operations
+        //   "wait() - in the case of CPU collectives, will block the process
+        //    until the operation is completed. In the case of CUDA collectives,
+        //    will block the currently active CUDA stream until the operation
+        //    is completed (but will not block the CPU)."
+        auto current_stream = at::cuda::getCurrentCUDAStream();
+        event_->block(current_stream);
+        return true;
+    }
+
+   protected:
+    std::shared_ptr<torch::Event> event_;
+    std::shared_ptr<TransferGroupMeta> meta_;
+    const MooncakeWorker* worker_;
+    std::vector<CudaTaskSubmissionToken> submitted_tasks_;
+};
+
+class MooncakeBarrierWorkCuda : public MooncakeWorkCuda {
+   public:
+    using MooncakeWorkCuda::MooncakeWorkCuda;
+
+    bool wait(std::chrono::milliseconds timeout) override {
+        // Skip host-side synchronization during CUDA graph capture.
+        // cudaEventSynchronize is not permitted while a stream is capturing.
+        if (at::cuda::currentStreamCaptureStatus() !=
+            c10::cuda::CaptureStatus::None) {
+            // We still need stream-level synchronization so that subsequent
+            // operations on the capture stream are ordered after the barrier
+            // task on the enqueue stream.
+            auto current_stream = at::cuda::getCurrentCUDAStream();
+            event_->block(current_stream);
+            return true;
+        }
+
+        if (timeout == kNoTimeout) {
+            event_->synchronize();
+            return true;
+        }
+
+        BackoffWaiter waiter(
+            BackoffWaiterConfig::constantSleep(std::chrono::microseconds(10)));
+        return waiter.wait_for(timeout, [this] { return event_->query(); });
+    }
+};
+
+void launchReduceKernel(at::Tensor dst, size_t pos, size_t realSize, void* src,
+                        size_t numRanks, c10d::ReduceOp op, bool* activeRanks,
+                        cudaStream_t stream) {
+    TORCH_CHECK(op == c10d::ReduceOp::SUM || op == c10d::ReduceOp::MIN ||
+                    op == c10d::ReduceOp::MAX || op == c10d::ReduceOp::PRODUCT,
+                "Only support SUM/MIN/MAX/PRODUCT for reduction.");
+    auto ptr = (char*)dst.data_ptr() + pos;
+    size_t num = realSize / dst.element_size();
+
+    switch (dst.scalar_type()) {
+        case c10::kByte:
+            launchReduceKernel_uint8((uint8_t*)ptr, (uint8_t*)src, num,
+                                     numRanks, (int)op, activeRanks, stream);
+            break;
+        case c10::kChar:
+            launchReduceKernel_int8((int8_t*)ptr, (int8_t*)src, num, numRanks,
+                                    (int)op, activeRanks, stream);
+            break;
+        case c10::kShort:
+            launchReduceKernel_int16((int16_t*)ptr, (int16_t*)src, num,
+                                     numRanks, (int)op, activeRanks, stream);
+            break;
+        case c10::kInt:
+            launchReduceKernel_int32((int*)ptr, (int*)src, num, numRanks,
+                                     (int)op, activeRanks, stream);
+            break;
+        case c10::kLong:
+            launchReduceKernel_int64((int64_t*)ptr, (int64_t*)src, num,
+                                     numRanks, (int)op, activeRanks, stream);
+            break;
+        case c10::kFloat:
+            launchReduceKernel_float((float*)ptr, (float*)src, num, numRanks,
+                                     (int)op, activeRanks, stream);
+            break;
+        case c10::kDouble:
+            launchReduceKernel_double((double*)ptr, (double*)src, num, numRanks,
+                                      (int)op, activeRanks, stream);
+            break;
+        case c10::kBool:
+            launchReduceKernel_bool((bool*)ptr, (bool*)src, num, numRanks,
+                                    (int)op, activeRanks, stream);
+            break;
+        case c10::kBFloat16:
+            launchReduceKernel_bf16(ptr, src, num, numRanks, (int)op,
+                                    activeRanks, stream);
+            break;
+        default:
+            TORCH_CHECK(false, c10::str("Unsupported reduce dtype: ",
+                                        dst.scalar_type()));
+    }
+}
+
+template <typename T>
+T applyReduceOp(const T& a, const T& b, c10d::ReduceOp op) {
+    switch (op) {
+        case c10d::ReduceOp::SUM:
+            return a + b;
+        case c10d::ReduceOp::PRODUCT:
+            return a * b;
+        case c10d::ReduceOp::MIN:
+            return std::min(a, b);
+        case c10d::ReduceOp::MAX:
+            return std::max(a, b);
+        default:
+            TORCH_CHECK(false, c10::str("Unsupported reduce op: ", op));
+    }
+}
+
+template <typename T>
+void reduceCpu(T* dst, const T* src, size_t numElements, size_t numRanks,
+               c10d::ReduceOp op, bool* activeRanks) {
+    at::parallel_for(0, numElements, 1024, [&](int64_t begin, int64_t end) {
+        for (int64_t i = begin; i < end; ++i) {
+            bool valid = false;
+            T acc{};
+            for (int64_t rank = 0; rank < numRanks; ++rank) {
+                if (activeRanks[rank]) {
+                    if (!valid) {
+                        acc = src[i + rank * numElements];
+                        valid = true;
+                    } else {
+                        acc =
+                            applyReduceOp(acc, src[i + rank * numElements], op);
+                    }
+                }
+            }
+            dst[i] = acc;
+        }
+    });
+}
+
+void launchReduceCpu(at::Tensor dst, size_t pos, size_t realSize, void* src,
+                     size_t numRanks, c10d::ReduceOp op, bool* activeRanks) {
+    auto ptr = (char*)dst.data_ptr() + pos;
+    size_t num = realSize / dst.element_size();
+
+    switch (dst.scalar_type()) {
+        case c10::kByte:
+            reduceCpu((uint8_t*)ptr, (uint8_t*)src, num, numRanks, op,
+                      activeRanks);
+            break;
+        case c10::kChar:
+            reduceCpu((int8_t*)ptr, (int8_t*)src, num, numRanks, op,
+                      activeRanks);
+            break;
+        case c10::kShort:
+            reduceCpu((int16_t*)ptr, (int16_t*)src, num, numRanks, op,
+                      activeRanks);
+            break;
+        case c10::kInt:
+            reduceCpu((int*)ptr, (int*)src, num, numRanks, op, activeRanks);
+            break;
+        case c10::kLong:
+            reduceCpu((int64_t*)ptr, (int64_t*)src, num, numRanks, op,
+                      activeRanks);
+            break;
+        case c10::kFloat:
+            reduceCpu((float*)ptr, (float*)src, num, numRanks, op, activeRanks);
+            break;
+        case c10::kDouble:
+            reduceCpu((double*)ptr, (double*)src, num, numRanks, op,
+                      activeRanks);
+            break;
+        case c10::kBool:
+            reduceCpu((bool*)ptr, (bool*)src, num, numRanks, op, activeRanks);
+            break;
+        default:
+            TORCH_CHECK(false, c10::str("Unsupported reduce dtype: ",
+                                        dst.scalar_type()));
+    }
+}
+
+MooncakeWorker::MooncakeWorker(int cuda_device_index)
+    : cuda_device_index_(cuda_device_index) {
+    int deviceCount = 0;
+    cudaError_t err = cudaGetDeviceCount(&deviceCount);
+    if (!err && deviceCount > 0) {
+        cudaHostAlloc(&tasks_, kNumTasks_ * sizeof(Task), cudaHostAllocMapped);
+        cudaHostGetDevicePointer(&tasks_device_, tasks_, 0);
+    } else {
+        LOG(WARNING) << "No GPU device found. Only the `mooncake-cpu` backend "
+                        "can be used.";
+        tasks_ = new Task[kNumTasks_];
+    }
+    for (size_t i = 0; i < kNumTasks_; ++i) {
+        tasks_[i].active = false;
+        tasks_[i].submitSequence = 0;
+        submitted_task_sequence_[i].store(0, std::memory_order_relaxed);
+    }
+}
+
+MooncakeWorker::~MooncakeWorker() {
+    running_ = false;
+    if (worker_thread_.joinable()) {
+        worker_thread_.join();
+    }
+}
+
+c10::intrusive_ptr<c10d::Work> MooncakeWorker::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)>&
+        bufferToTensor) {
+    connection_ctx->waitUntilNewRanksConnected();
+
+    size_t chunkSize = ((kBufferSize - 1) / meta->size) & ~(size_t)7;
+    auto future = c10::make_intrusive<c10::ivalue::Future>(
+        c10::ListType::create(c10::TensorType::get()));
+
+    struct IterState {
+        size_t currentPos = 0;
+    };
+    auto state = std::make_shared<IterState>();
+
+    auto processNextChunk = std::make_shared<std::function<void()>>();
+    std::weak_ptr<std::function<void()>> weakProcessNextChunk =
+        processNextChunk;
+
+    *processNextChunk = [this, weakProcessNextChunk, state, opType, tensorSize,
+                         chunkSize, broadcastRoot, meta, tensorToBuffer,
+                         bufferToTensor, future]() {
+        auto processNextChunk = weakProcessNextChunk.lock();
+
+        if (state->currentPos >= tensorSize) {
+            future->markCompleted(c10::IValue());
+            return;
+        }
+
+        int taskId = cpuTaskCount % 2;
+        TORCH_CHECK(!tasks_[taskId].active);
+
+        size_t realSize = std::min(chunkSize, tensorSize - state->currentPos);
+        int bufferOffset = meta->taskCount % 2;
+
+        tasks_[taskId].opType = (int)opType;
+        tasks_[taskId].tensorSize = realSize;
+        tasks_[taskId].broadcastRoot = broadcastRoot;
+        tasks_[taskId].bufferOffset = bufferOffset;
+        tasks_[taskId].transferGroupMeta = meta.get();
+        tensorToBuffer(
+            (void*)meta->segmentInfos[meta->rank].send_buffer[bufferOffset],
+            state->currentPos, realSize);
+
+        hasCallback_[taskId] = true;
+
+        callbacks_[taskId] = [this, processNextChunk, state, meta,
+                              bufferToTensor, bufferOffset, realSize,
+                              future]() {
+            if (meta->activeRanksTensor.device().is_cpu()) {
+                for (int i = 0; i < meta->size; ++i) {
+                    meta->activeRanksTensor[i] = meta->activeRanks[i] ? 1 : 0;
+                }
+            }
+            bufferToTensor(
+                (void*)meta->segmentInfos[meta->rank].recv_buffer[bufferOffset],
+                state->currentPos, realSize);
+
+            state->currentPos += realSize;
+
+            (*processNextChunk)();
+        };
+
+        tasks_[taskId].active = true;
+        ++cpuTaskCount;
+        ++meta->taskCount;
+    };
+
+    (*processNextChunk)();
+
+    return c10::make_intrusive<MooncakeWorkCpu>(opType, future, meta);
+}
+
+c10::intrusive_ptr<c10d::Work> MooncakeWorker::putTaskCuda(
+    c10d::OpType opType, size_t tensorSize, int64_t broadcastRoot,
+    const std::shared_ptr<TransferGroupMeta>& meta,
+    const std::shared_ptr<ConnectionContext>& connection_ctx,
+    const at::cuda::CUDAStream& issue_stream,
+    const std::function<void(void* dst, size_t pos, size_t realSize,
+                             const at::cuda::CUDAStream&)>& tensorToBuffer,
+    const std::function<void(void* src, size_t pos, size_t realSize,
+                             const at::cuda::CUDAStream&)>& bufferToTensor) {
+    connection_ctx->waitUntilNewRanksConnected();
+
+    size_t chunkSize = ((kBufferSize - 1) / meta->size) & ~(size_t)7;
+
+    at::cuda::CUDAStream enq_stream =
+        at::cuda::getStreamFromPool(false, issue_stream.device_index());
+
+    auto event_start = std::make_shared<torch::Event>(torch::kCUDA);
+    event_start->record(issue_stream);
+    event_start->block(enq_stream);
+
+    std::vector<CudaTaskSubmissionToken> submitted_tasks;
+    submitted_tasks.reserve((tensorSize + chunkSize - 1) / chunkSize);
+    for (size_t pos = 0; pos < tensorSize; pos += chunkSize) {
+        size_t realSize = std::min(tensorSize, pos + chunkSize) - pos;
+        int taskId = cudaTaskCount % 2 + 2;
+        int bufferOffset = meta->taskCount % 2;
+        const uint64_t taskSequence =
+            next_cuda_task_sequence_.fetch_add(1, std::memory_order_relaxed);
+        submitted_tasks.push_back(
+            {.task_id = static_cast<size_t>(taskId), .sequence = taskSequence});
+        tensorToBuffer(
+            (void*)meta->segmentInfos[meta->rank].send_buffer[bufferOffset],
+            pos, realSize, enq_stream);
+
+        hasCallback_[taskId] = false;
+        launchEnqueueTaskKernel(
+            (int)opType, realSize, broadcastRoot, bufferOffset, taskSequence,
+            meta.get(), tasks_device_, meta->size, meta->activeRanksDevice,
+            meta->activeRanksTensor.data_ptr<int>(), taskId,
+            enq_stream.stream());
+        bufferToTensor(
+            (void*)meta->segmentInfos[meta->rank].recv_buffer[bufferOffset],
+            pos, realSize, enq_stream);
+
+        ++cudaTaskCount;
+        ++meta->taskCount;
+    }
+
+    auto event_end = std::make_shared<torch::Event>(torch::kCUDA);
+    event_end->record(enq_stream);
+
+    if (opType == c10d::OpType::BARRIER) {
+        return c10::make_intrusive<MooncakeBarrierWorkCuda>(
+            opType, event_end, meta, this, std::move(submitted_tasks));
+    }
+    return c10::make_intrusive<MooncakeWorkCuda>(opType, event_end, meta, this,
+                                                 std::move(submitted_tasks));
+}
+
+}  // namespace mooncake
diff --git a/mooncake-pg/src/mooncake_worker_thread.cpp b/mooncake-pg/src/mooncake_worker_thread.cpp
index aeaba687..56e0c5c0 100644
--- a/mooncake-pg/src/mooncake_worker_thread.cpp
+++ b/mooncake-pg/src/mooncake_worker_thread.cpp
@@ -1,4 +1,4 @@
-#include <cuda_runtime.h>
+#include <cuda_alike.h>
 #include <thread>
 #include <mooncake_worker.cuh>
 #include <glog/logging.h>
@@ -16,6 +16,13 @@ enum WorkerTaskStatus {
 
 static constexpr size_t kInvalidTaskId = static_cast<size_t>(-1);
 
+static void setActiveRanksTensorValue(TransferGroupMeta* group, int rank,
+                                      int value) {
+    if (group->activeRanksTensor.device().is_cpu()) {
+        group->activeRanksTensor[rank] = value;
+    }
+}
+
 void MooncakeWorker::Start() {
     bool expected = false;
     if (started_.compare_exchange_strong(expected, true)) {
@@ -85,11 +92,12 @@ void MooncakeWorker::startWorker() {
                 }
 
                 auto group = (TransferGroupMeta*)task.transferGroupMeta;
-                bool skipTransfer = (task.opType == c10d::OpType::BROADCAST &&
-                                     group->rank != task.broadcastRoot) ||
-                                    (task.opType == c10d::OpType::SCATTER &&
-                                     group->rank != task.broadcastRoot) ||
-                                    task.opType == c10d::OpType::BARRIER;
+                bool skipTransfer =
+                    ((c10d::OpType)task.opType == c10d::OpType::BROADCAST &&
+                     group->rank != task.broadcastRoot) ||
+                    ((c10d::OpType)task.opType == c10d::OpType::SCATTER &&
+                     group->rank != task.broadcastRoot) ||
+                    (c10d::OpType)task.opType == c10d::OpType::BARRIER;
                 if (task_status[i].load(std::memory_order_acquire) == IDLE) {
                     const auto submit_sequence = task.submitSequence;
                     if (skipTransfer) {
@@ -107,15 +115,17 @@ void MooncakeWorker::startWorker() {
                         if (!group->activeRanks[j]) {
                             continue;
                         }
-                        if ((task.opType == c10d::OpType::GATHER ||
-                             task.opType == c10d::OpType::REDUCE) &&
+                        if (((c10d::OpType)task.opType ==
+                                 c10d::OpType::GATHER ||
+                             (c10d::OpType)task.opType ==
+                                 c10d::OpType::REDUCE) &&
                             j != task.broadcastRoot) {
                             continue;
                         }
                         uint64_t source = group->segmentInfos[group->rank]
                                               .send_buffer[task.bufferOffset];
 
-                        switch (task.opType) {
+                        switch ((c10d::OpType)task.opType) {
                             case c10d::OpType::BROADCAST:
                             case c10d::OpType::ALLREDUCE:
                             case c10d::OpType::ALLGATHER:
@@ -136,7 +146,7 @@ void MooncakeWorker::startWorker() {
                             group->segmentInfos[j]
                                 .recv_buffer[task.bufferOffset];
 
-                        switch (task.opType) {
+                        switch ((c10d::OpType)task.opType) {
                             case c10d::OpType::BROADCAST:
                             case c10d::OpType::SCATTER:
                                 break;
@@ -206,7 +216,7 @@ void MooncakeWorker::startWorker() {
                                     // connection poller to reconnect it.
                                     group->peerConnected[j] = false;
                                     group->activeRanks[j] = false;
-                                    group->activeRanksTensor[j] = 0;
+                                    setActiveRanksTensorValue(group, j, 0);
                                 } else {
                                     batch_done = false;
                                     break;
@@ -295,7 +305,7 @@ void MooncakeWorker::startWorker() {
                                 // connection poller to reconnect it.
                                 group->peerConnected[j] = false;
                                 group->activeRanks[j] = false;
-                                group->activeRanksTensor[j] = 0;
+                                setActiveRanksTensorValue(group, j, 0);
                             } else {
                                 task_done = false;
                                 break;
diff --git a/mooncake-pg/src/p2p_proxy.cpp b/mooncake-pg/src/p2p_proxy.cpp
index 7df056c7..e56dcdf0 100644
--- a/mooncake-pg/src/p2p_proxy.cpp
+++ b/mooncake-pg/src/p2p_proxy.cpp
@@ -1,8 +1,7 @@
 #include <memory>
 #include <mutex>
 #include <p2p_proxy.h>
-#include <ATen/cuda/CUDAContext.h>
-#include <cuda_runtime.h>
+#include <cuda_alike.h>
 #include <algorithm>
 #include <cstddef>
 #include <cstdlib>
@@ -16,14 +15,12 @@ namespace mooncake {
 
 namespace {
 
-constexpr size_t kP2PBytesPerRank = kP2PBufferSize;
-constexpr size_t kP2PNumSlotsPerRank = kP2PNumSlots;
+size_t getEnv_size_t(const char* name, size_t default_val) {
+    const char* val = std::getenv(name);
+    return val ? std::strtoull(val, nullptr, 10) : default_val;
+}
 
-static_assert(kP2PBufferSize % kP2PNumSlots == 0,
-              "kP2PBufferSize must be divisible by kP2PNumSlots");
-static_assert(kP2PNumSlots > 1, "P2P ring requires at least 2 slots per rank");
-
-void SetCudaDeviceIfNeeded(bool is_cpu, int cuda_device_index,
+void setCudaDeviceIfNeeded(bool is_cpu, int cuda_device_index,
                            const char* context) {
     if (is_cpu) {
         return;
@@ -37,13 +34,135 @@ void SetCudaDeviceIfNeeded(bool is_cpu, int cuda_device_index,
 
 }  // namespace
 
+void P2PChunkPool::init(void* base_addr, size_t chunk_size,
+                        uint32_t num_chunks) {
+    base_addr_ = base_addr;
+    chunk_size_ = chunk_size;
+    free_stack_.clear();
+    free_stack_.reserve(num_chunks);
+    for (uint32_t idx = 0; idx < num_chunks; ++idx) {
+        free_stack_.push_back(static_cast<uint8_t*>(base_addr_) +
+                              (num_chunks - 1 - idx) * chunk_size_);
+    }
+}
+
+void* P2PChunkPool::acquire() {
+    if (free_stack_.empty()) {
+        return nullptr;
+    }
+    void* ptr = free_stack_.back();
+    free_stack_.pop_back();
+    return ptr;
+}
+
+void P2PChunkPool::release(void* ptr) { free_stack_.push_back(ptr); }
+
+// Consumer-side reliable read.
+//
+// 1. Load header_token with acquire semantics.
+// 2. If it is kInvalidControlToken the slot is empty -> fail.
+// 3. Optimistically copy the payload.
+// 4. Issue an acquire fence so the payload reads cannot be reordered past
+//    the footer load that follows.
+// 5. Load footer_token.
+// 6. Accept the payload only when header == footer.
+bool CreditSlot::tryLoad(uint64_t& out_recv_addr, uint32_t& out_chunk_len,
+                         uint32_t& out_epoch, uint32_t& out_sequence) const {
+    uint64_t h = std::atomic_ref(header_token).load(std::memory_order_acquire);
+    if (h == kInvalidControlToken) return false;
+
+    // Optimistic payload copy – may be torn if the NIC is still writing.
+    uint64_t addr = std::atomic_ref(recv_addr).load(std::memory_order_relaxed);
+    uint32_t len = std::atomic_ref(chunk_len).load(std::memory_order_relaxed);
+
+    // Avoid moving payload reads below the footer check.
+    std::atomic_thread_fence(std::memory_order_acquire);
+
+    uint64_t f = std::atomic_ref(footer_token).load(std::memory_order_relaxed);
+
+    if (h == f) {
+        out_recv_addr = addr;
+        out_chunk_len = len;
+        out_epoch = static_cast<uint32_t>(h >> 32);
+        out_sequence = static_cast<uint32_t>(h);
+        return true;
+    }
+    // Torn write – caller retries next poll.
+    return false;
+}
+
+// Producer-side reliable publish.
+//
+// On the local CPU the store order matters:
+//   payload -> footer (relaxed) -> header (release).
+// The release store to header_token guarantees that every prior store is
+// visible to any consumer that sees the new header value.
+void CreditSlot::publish(uint32_t epoch, uint32_t seq, uint64_t addr,
+                         uint32_t len) {
+    uint64_t new_token = makeControlToken(epoch, seq);
+
+    recv_addr = addr;
+    chunk_len = len;
+
+    // Write footer first so it is never ahead of the header.
+    std::atomic_ref(footer_token).store(new_token, std::memory_order_relaxed);
+    // Release the header last
+    std::atomic_ref(header_token).store(new_token, std::memory_order_release);
+}
+
+void CreditSlot::reset() {
+    recv_addr = 0;
+    chunk_len = 0;
+    uint64_t inv = kInvalidControlToken;
+    std::atomic_ref(footer_token).store(inv, std::memory_order_relaxed);
+    std::atomic_ref(header_token).store(inv, std::memory_order_release);
+}
+
+// See CreditSlot::tryLoad() for the detailed description.
+bool AckSlot::tryLoad(uint32_t& out_chunk_len, uint32_t& out_epoch,
+                      uint32_t& out_sequence) const {
+    uint64_t h = std::atomic_ref(header_token).load(std::memory_order_acquire);
+    if (h == kInvalidControlToken) return false;
+
+    uint32_t len = std::atomic_ref(chunk_len).load(std::memory_order_relaxed);
+
+    std::atomic_thread_fence(std::memory_order_acquire);
+
+    uint64_t f = std::atomic_ref(footer_token).load(std::memory_order_relaxed);
+
+    if (h == f) {
+        out_chunk_len = len;
+        out_epoch = static_cast<uint32_t>(h >> 32);
+        out_sequence = static_cast<uint32_t>(h);
+        return true;
+    }
+    return false;
+}
+
+// See CreditSlot::publish() for the detailed description.
+void AckSlot::publish(uint32_t epoch, uint32_t seq, uint32_t len) {
+    uint64_t new_token = makeControlToken(epoch, seq);
+
+    chunk_len = len;
+
+    std::atomic_ref(footer_token).store(new_token, std::memory_order_relaxed);
+    std::atomic_ref(header_token).store(new_token, std::memory_order_release);
+}
+
+void AckSlot::reset() {
+    chunk_len = 0;
+    uint64_t inv = kInvalidControlToken;
+    std::atomic_ref(footer_token).store(inv, std::memory_order_relaxed);
+    std::atomic_ref(header_token).store(inv, std::memory_order_release);
+}
+
 P2PProxy::P2PProxy(TransferEngine* engine, const Options& options)
     : engine_(engine),
       is_cpu_(options.is_cpu),
       rank_(options.rank),
       size_(options.size),
       cuda_device_index_(options.cuda_device_index),
-      location_(options.location) {
+      transfer_timeout_ms_(options.transfer_timeout_ms) {
     if (!is_cpu_ && cuda_device_index_ < 0) {
         int current_device = -1;
         const cudaError_t get_device_error = cudaGetDevice(&current_device);
@@ -52,7 +171,7 @@ P2PProxy::P2PProxy(TransferEngine* engine, const Options& options)
                     cudaGetErrorString(get_device_error));
         cuda_device_index_ = current_device;
     }
-    AllocateResources();
+    allocateResources();
 }
 
 P2PProxy::~P2PProxy() {
@@ -62,85 +181,56 @@ P2PProxy::~P2PProxy() {
         return;
     }
 
-    ReleaseResources();
+    releaseResources();
 }
 
-void P2PProxy::BindMeta(const std::shared_ptr<TransferGroupMeta>& meta) {
+void P2PProxy::bindMeta(const std::shared_ptr<TransferGroupMeta>& meta) {
     meta_ = meta;
 }
 
-void P2PProxy::AllocateResources() {
+void P2PProxy::extendGroupSizeTo(int new_size) {
+    TORCH_CHECK(new_size >= size_, "extendGroupSizeTo: new_size < size_");
+    size_ = new_size;
+}
+
+// Allocate control regions only.
+// Send/Recv chunk pools are owned by P2PDeviceWorker and shared across
+// multiple P2PProxy instances on the same device.
+void P2PProxy::allocateResources() {
     TORCH_CHECK(engine_, "P2PProxy engine is null.");
-    if (resources_.send_buffer_ != nullptr ||
-        resources_.recv_buffer_ != nullptr ||
-        resources_.ctrl_send_region_ != nullptr ||
-        resources_.ctrl_recv_region_ != nullptr) {
+    if (resources_.credit_region_ != nullptr ||
+        resources_.ack_region_ != nullptr) {
         return;
     }
 
-    TORCH_CHECK(size_ >= 0, "P2PProxy invalid size_: ", size_);
+    TORCH_CHECK(size_ > 0, "P2PProxy invalid group size: ", size_);
     TORCH_CHECK(static_cast<size_t>(size_) <= kMaxNumRanks,
-                "P2PProxy size_ exceeds kMaxNumRanks: ", size_);
+                "P2PProxy group size exceeds kMaxNumRanks: ", size_);
 
-    if (size_ == 0) {
-        return;
+    const size_t ctrl_slots =
+        kMaxNumRanks * static_cast<size_t>(kP2PControlRingSize);
+    resources_.credit_region_ = new CreditSlot[ctrl_slots]{};
+    resources_.ack_region_ = new AckSlot[ctrl_slots]{};
+
+    for (size_t i = 0; i < ctrl_slots; ++i) {
+        resources_.credit_region_[i].reset();
+        resources_.ack_region_[i].reset();
+    }
+
+    for (size_t i = 0; i < kMaxNumRanks; ++i) {
+        peer_epoch_[i].store(1, std::memory_order_release);
     }
 
-    const size_t p2p_total_buffer_size =
-        kP2PBufferSize * static_cast<size_t>(size_);
-
-    if (is_cpu_) {
-        resources_.send_buffer_ = std::malloc(p2p_total_buffer_size);
-        TORCH_CHECK(resources_.send_buffer_ != nullptr,
-                    "Failed to allocate CPU P2P send buffer");
-        int rc = engine_->registerLocalMemory(resources_.send_buffer_,
-                                              p2p_total_buffer_size, location_);
-        TORCH_CHECK(rc == 0, "Failed to register CPU P2P send buffer");
-
-        resources_.recv_buffer_ = std::malloc(p2p_total_buffer_size);
-        TORCH_CHECK(resources_.recv_buffer_ != nullptr,
-                    "Failed to allocate CPU P2P recv buffer");
-        rc = engine_->registerLocalMemory(resources_.recv_buffer_,
-                                          p2p_total_buffer_size, location_);
-        TORCH_CHECK(rc == 0, "Failed to register CPU P2P recv buffer");
-    } else {
-        SetCudaDeviceIfNeeded(
-            is_cpu_, cuda_device_index_,
-            "P2PProxy AllocateResources cudaSetDevice failed");
-        cudaError_t err =
-            cudaMalloc(&resources_.send_buffer_, p2p_total_buffer_size);
-        TORCH_CHECK(err == cudaSuccess,
-                    "Failed to allocate CUDA P2P send buffer");
-        int rc = engine_->registerLocalMemory(resources_.send_buffer_,
-                                              p2p_total_buffer_size, location_);
-        TORCH_CHECK(rc == 0, "Failed to register CUDA P2P send buffer");
-
-        err = cudaMalloc(&resources_.recv_buffer_, p2p_total_buffer_size);
-        TORCH_CHECK(err == cudaSuccess,
-                    "Failed to allocate CUDA P2P recv buffer");
-        rc = engine_->registerLocalMemory(resources_.recv_buffer_,
-                                          p2p_total_buffer_size, location_);
-        TORCH_CHECK(rc == 0, "Failed to register CUDA P2P recv buffer");
-    }
-
-    resources_.ctrl_send_region_ = new P2PControlSlot[kMaxNumRanks]{};
-    resources_.ctrl_recv_region_ = new P2PControlSlot[kMaxNumRanks]{};
     for (size_t i = 0; i < kMaxNumRanks; ++i) {
-        resources_.ctrl_send_region_[i].head.store(0,
-                                                   std::memory_order_relaxed);
-        resources_.ctrl_send_region_[i].tail.store(0,
-                                                   std::memory_order_relaxed);
-        resources_.ctrl_recv_region_[i].head.store(0,
-                                                   std::memory_order_relaxed);
-        resources_.ctrl_recv_region_[i].tail.store(0,
-                                                   std::memory_order_relaxed);
-        send_peer_lanes_[i].local_head_ = 0;
         send_peer_lanes_[i].pending_send_ops_.clear();
         send_peer_lanes_[i].active_send_op_.reset();
+        send_peer_lanes_[i].credit_consume_seq_ = 0;
         send_peer_lanes_[i].copy_ready_events_.fill(nullptr);
-        recv_peer_lanes_[i].local_tail_ = 0;
+
         recv_peer_lanes_[i].pending_recv_ops_.clear();
         recv_peer_lanes_[i].active_recv_op_.reset();
+        recv_peer_lanes_[i].credit_issue_seq_ = 0;
+        recv_peer_lanes_[i].ack_consume_seq_ = 0;
         recv_peer_lanes_[i].copy_ready_events_.fill(nullptr);
     }
 
@@ -148,9 +238,6 @@ void P2PProxy::AllocateResources() {
         for (int peer_rank = 0; peer_rank < size_; ++peer_rank) {
             for (auto& copy_ready_event :
                  send_peer_lanes_[peer_rank].copy_ready_events_) {
-                if (copy_ready_event != nullptr) {
-                    continue;
-                }
                 const cudaError_t create_error = cudaEventCreateWithFlags(
                     &copy_ready_event, cudaEventDisableTiming);
                 TORCH_CHECK(create_error == cudaSuccess,
@@ -160,9 +247,6 @@ void P2PProxy::AllocateResources() {
 
             for (auto& copy_ready_event :
                  recv_peer_lanes_[peer_rank].copy_ready_events_) {
-                if (copy_ready_event != nullptr) {
-                    continue;
-                }
                 const cudaError_t create_error = cudaEventCreateWithFlags(
                     &copy_ready_event, cudaEventDisableTiming);
                 TORCH_CHECK(create_error == cudaSuccess,
@@ -171,137 +255,213 @@ void P2PProxy::AllocateResources() {
             }
         }
     }
-    int rc = engine_->registerLocalMemory(resources_.ctrl_send_region_,
-                                          kMaxNumRanks * sizeof(P2PControlSlot),
-                                          kWildcardLocation);
-    TORCH_CHECK(rc == 0, "Failed to register P2P ctrl send region");
 
-    rc = engine_->registerLocalMemory(resources_.ctrl_recv_region_,
-                                      kMaxNumRanks * sizeof(P2PControlSlot),
+    int rc = engine_->registerLocalMemory(resources_.credit_region_,
+                                          ctrl_slots * sizeof(CreditSlot),
+                                          kWildcardLocation);
+    TORCH_CHECK(rc == 0, "Failed to register P2P credit region");
+
+    rc = engine_->registerLocalMemory(resources_.ack_region_,
+                                      ctrl_slots * sizeof(AckSlot),
                                       kWildcardLocation);
-    TORCH_CHECK(rc == 0, "Failed to register P2P ctrl recv region");
+    TORCH_CHECK(rc == 0, "Failed to register P2P ack region");
+
+    // Staging buffers for control messages.  RDMA requires every
+    // transfer source to live in a registered MR.
+    resources_.credit_staging_buf_ = new CreditSlot[ctrl_slots]{};
+    rc = engine_->registerLocalMemory(resources_.credit_staging_buf_,
+                                      ctrl_slots * sizeof(CreditSlot),
+                                      kWildcardLocation);
+    TORCH_CHECK(rc == 0, "Failed to register P2P credit staging region");
+
+    resources_.ack_staging_buf_ = new AckSlot[ctrl_slots]{};
+    rc = engine_->registerLocalMemory(resources_.ack_staging_buf_,
+                                      ctrl_slots * sizeof(AckSlot),
+                                      kWildcardLocation);
+    TORCH_CHECK(rc == 0, "Failed to register P2P ack staging region");
 }
 
-void P2PProxy::ResetPeerState(int peer_rank) {
+void P2PProxy::resetPeerState(int peer_rank) {
     TORCH_CHECK(peer_rank >= 0 && peer_rank < size_,
                 "ResetPeerState: peer_rank out of range: ", peer_rank,
                 " size: ", size_);
+
+    // Epoch update:
+    //   We bump epoch_ so that any slots still in flight from the
+    //   old session (written by the sender/receiver before it learned about
+    //   the Reset) are recognized as stale and skipped.
+    peer_epoch_[peer_rank].fetch_add(1, std::memory_order_acq_rel);
+
+    // Request reset
     reset_send_req_[peer_rank].store(true, std::memory_order_release);
     reset_recv_req_[peer_rank].store(true, std::memory_order_release);
+
+    // Wake up worker
     if (device_worker_) {
-        device_worker_->WakeUpSend();
-        device_worker_->WakeUpRecv();
+        device_worker_->wakeUpSend();
+        device_worker_->wakeUpRecv();
     }
 }
 
-void P2PProxy::PerformSendReset(int peer_rank) {
-    auto& lane = send_peer_lanes_[peer_rank];
+// Reset sender state for peer_rank.
+void P2PProxy::performSendReset(int peer_rank) {
+    resetSendLane(send_peer_lanes_[peer_rank]);
+    resetPeerControlLanes(peer_rank);
+}
+
+// Reset receiver state for peer_rank.
+void P2PProxy::performRecvReset(int peer_rank) {
+    resetRecvLane(recv_peer_lanes_[peer_rank]);
+    resetPeerControlLanes(peer_rank);
+}
+
+void P2PProxy::reportBrokenPeer(int peer_rank) {
+    resetPeerState(peer_rank);
+    // Set peerConnected to notify the connection poller to reconnect it.
+    meta_->peerConnected[peer_rank] = false;
+    meta_->activeRanks[peer_rank] = false;
+    if (meta_->activeRanksTensor.device().is_cpu()) {
+        meta_->activeRanksTensor[peer_rank] = 0;
+    }
+    LOG(ERROR) << "Rank " << meta_->rank << " marking peer " << peer_rank
+               << " as broken during P2P transfer.";
+}
+
+void P2PProxy::releaseSendTaskResources(SendTransferTask& task) const {
+    if (task.staging_addr_ != nullptr) {
+        if (send_pool_) send_pool_->release(task.staging_addr_);
+        task.staging_addr_ = nullptr;
+    }
+    if (task.transfer_batch_id_.has_value()) {
+        if (engine_) engine_->freeBatchID(task.transfer_batch_id_.value());
+        task.transfer_batch_id_.reset();
+    }
+    if (task.ack_batch_id_.has_value()) {
+        if (engine_) engine_->freeBatchID(task.ack_batch_id_.value());
+        task.ack_batch_id_.reset();
+    }
+}
+
+void P2PProxy::releaseRecvTaskResources(RecvTransferTask& task) const {
+    if (task.local_addr_ != nullptr) {
+        if (recv_pool_) recv_pool_->release(task.local_addr_);
+        task.local_addr_ = nullptr;
+    }
+    if (task.credit_batch_id_.has_value()) {
+        if (engine_) engine_->freeBatchID(task.credit_batch_id_.value());
+        task.credit_batch_id_.reset();
+    }
+}
+
+void P2PProxy::resetSendLane(SendPeerLane& lane) {
+    for (auto& pending : lane.pending_send_ops_) {
+        pending.status_->store(OpStatus::kFailed, std::memory_order_release);
+        active_send_tasks_.fetch_sub(1, std::memory_order_release);
+    }
     lane.pending_send_ops_.clear();
-    lane.active_send_op_.reset();
-    lane.local_head_ = 0;
 
-    if (resources_.ctrl_send_region_) {
-        resources_.ctrl_send_region_[peer_rank].head.store(
-            0, std::memory_order_relaxed);
-        resources_.ctrl_send_region_[peer_rank].tail.store(
-            0, std::memory_order_relaxed);
+    if (lane.active_send_op_.has_value()) {
+        auto& op_ctx = lane.active_send_op_.value();
+        for (auto& task : op_ctx.tasks_) {
+            releaseSendTaskResources(task);
+        }
+        op_ctx.tasks_.clear();
+        op_ctx.status_->store(OpStatus::kFailed, std::memory_order_release);
+        active_send_tasks_.fetch_sub(1, std::memory_order_release);
+        lane.active_send_op_.reset();
     }
+    lane.credit_consume_seq_ = 0;
 }
 
-void P2PProxy::PerformRecvReset(int peer_rank) {
-    auto& lane = recv_peer_lanes_[peer_rank];
+void P2PProxy::resetRecvLane(RecvPeerLane& lane) {
+    for (auto& pending : lane.pending_recv_ops_) {
+        pending.status_->store(OpStatus::kFailed, std::memory_order_release);
+        active_recv_tasks_.fetch_sub(1, std::memory_order_release);
+    }
     lane.pending_recv_ops_.clear();
-    lane.active_recv_op_.reset();
-    lane.local_tail_ = 0;
 
-    if (resources_.ctrl_recv_region_) {
-        resources_.ctrl_recv_region_[peer_rank].head.store(
-            0, std::memory_order_relaxed);
-        resources_.ctrl_recv_region_[peer_rank].tail.store(
-            0, std::memory_order_relaxed);
+    if (lane.active_recv_op_.has_value()) {
+        auto& op_ctx = lane.active_recv_op_.value();
+        for (auto& task : op_ctx.tasks_) {
+            releaseRecvTaskResources(task);
+        }
+        op_ctx.tasks_.clear();
+        op_ctx.status_->store(OpStatus::kFailed, std::memory_order_release);
+        active_recv_tasks_.fetch_sub(1, std::memory_order_release);
+        lane.active_recv_op_.reset();
+    }
+    lane.credit_issue_seq_ = 0;
+    lane.ack_consume_seq_ = 0;
+}
+
+void P2PProxy::resetPeerControlLanes(int peer_rank) {
+    auto* credit_lane = getLocalCreditLane(peer_rank);
+    auto* ack_lane = getLocalAckLane(peer_rank);
+    for (uint32_t i = 0; i < kP2PControlRingSize; ++i) {
+        credit_lane[i].reset();
+        ack_lane[i].reset();
     }
 }
 
-void P2PProxy::ReleaseResources() {
+void P2PProxy::releaseResources() {
     TORCH_CHECK(!resource_abandoned_,
                 "Should not release abandoned resources.");
 
-    SetCudaDeviceIfNeeded(is_cpu_, cuda_device_index_,
+    setCudaDeviceIfNeeded(is_cpu_, cuda_device_index_,
                           "P2PProxy ReleaseResources cudaSetDevice failed");
 
     for (size_t i = 0; i < kMaxNumRanks; ++i) {
+        // Reset lane
         auto& send_lane = send_peer_lanes_[i];
-        for (auto& copy_ready_event : send_lane.copy_ready_events_) {
-            if (copy_ready_event == nullptr) {
-                continue;
-            }
-            const cudaError_t destroy_error =
-                cudaEventDestroy(copy_ready_event);
-            TORCH_CHECK(destroy_error == cudaSuccess,
-                        "Failed to destroy pooled send copy-ready event: ",
-                        cudaGetErrorString(destroy_error));
-            copy_ready_event = nullptr;
-        }
-        send_lane.pending_send_ops_.clear();
-        send_lane.active_send_op_.reset();
-        send_lane.local_head_ = 0;
-
         auto& recv_lane = recv_peer_lanes_[i];
-        for (auto& copy_ready_event : recv_lane.copy_ready_events_) {
-            if (copy_ready_event == nullptr) {
-                continue;
+        resetSendLane(send_lane);
+        resetRecvLane(recv_lane);
+
+        // Destroy cuda events
+        auto destroyEvents = [](auto& events) {
+            for (auto& ev : events) {
+                if (ev == nullptr) continue;
+                const cudaError_t err = cudaEventDestroy(ev);
+                TORCH_CHECK(err == cudaSuccess,
+                            "Failed to destroy pooled copy-ready event: ",
+                            cudaGetErrorString(err));
+                ev = nullptr;
             }
-            const cudaError_t destroy_error =
-                cudaEventDestroy(copy_ready_event);
-            TORCH_CHECK(destroy_error == cudaSuccess,
-                        "Failed to destroy pooled recv copy-ready event: ",
-                        cudaGetErrorString(destroy_error));
-            copy_ready_event = nullptr;
-        }
-        recv_lane.pending_recv_ops_.clear();
-        recv_lane.active_recv_op_.reset();
+        };
+        destroyEvents(send_lane.copy_ready_events_);
+        destroyEvents(recv_lane.copy_ready_events_);
     }
 
-    if (!engine_) {
-        return;
+    if (resources_.credit_staging_buf_ != nullptr) {
+        if (engine_)
+            engine_->unregisterLocalMemory(resources_.credit_staging_buf_);
+        delete[] resources_.credit_staging_buf_;
+        resources_.credit_staging_buf_ = nullptr;
     }
 
-    if (resources_.ctrl_send_region_ != nullptr) {
-        engine_->unregisterLocalMemory(resources_.ctrl_send_region_);
-        delete[] resources_.ctrl_send_region_;
-        resources_.ctrl_send_region_ = nullptr;
+    if (resources_.ack_staging_buf_ != nullptr) {
+        if (engine_)
+            engine_->unregisterLocalMemory(resources_.ack_staging_buf_);
+        delete[] resources_.ack_staging_buf_;
+        resources_.ack_staging_buf_ = nullptr;
     }
 
-    if (resources_.ctrl_recv_region_ != nullptr) {
-        engine_->unregisterLocalMemory(resources_.ctrl_recv_region_);
-        delete[] resources_.ctrl_recv_region_;
-        resources_.ctrl_recv_region_ = nullptr;
+    if (resources_.credit_region_ != nullptr) {
+        if (engine_) engine_->unregisterLocalMemory(resources_.credit_region_);
+        delete[] resources_.credit_region_;
+        resources_.credit_region_ = nullptr;
     }
 
-    if (resources_.send_buffer_ != nullptr) {
-        engine_->unregisterLocalMemory(resources_.send_buffer_);
-        if (is_cpu_) {
-            std::free(resources_.send_buffer_);
-        } else {
-            cudaFree(resources_.send_buffer_);
-        }
-        resources_.send_buffer_ = nullptr;
-    }
-
-    if (resources_.recv_buffer_ != nullptr) {
-        engine_->unregisterLocalMemory(resources_.recv_buffer_);
-        if (is_cpu_) {
-            std::free(resources_.recv_buffer_);
-        } else {
-            cudaFree(resources_.recv_buffer_);
-        }
-        resources_.recv_buffer_ = nullptr;
+    if (resources_.ack_region_ != nullptr) {
+        if (engine_) engine_->unregisterLocalMemory(resources_.ack_region_);
+        delete[] resources_.ack_region_;
+        resources_.ack_region_ = nullptr;
     }
 }
 
-void P2PProxy::AbandonResources() { resource_abandoned_ = true; }
+void P2PProxy::abandonResources() { resource_abandoned_ = true; }
 
-void P2PProxy::EnqueueSend(SendOp op) {
+void P2PProxy::enqueueSend(SendOp op) {
     op.tensor_ =
         op.tensor_.is_contiguous() ? op.tensor_ : op.tensor_.contiguous();
 
@@ -310,364 +470,221 @@ void P2PProxy::EnqueueSend(SendOp op) {
         send_queue_.emplace(std::move(op));
     }
     active_send_tasks_.fetch_add(1, std::memory_order_release);
-    if (device_worker_) device_worker_->WakeUpSend();
+    if (device_worker_) device_worker_->wakeUpSend();
 }
 
-void P2PProxy::EnqueueRecv(RecvOp op) {
+void P2PProxy::enqueueRecv(RecvOp op) {
     {
         std::lock_guard<std::mutex> lock(recv_queue_mutex_);
         recv_queue_.push(std::move(op));
     }
     active_recv_tasks_.fetch_add(1, std::memory_order_release);
-    if (device_worker_) device_worker_->WakeUpRecv();
+    if (device_worker_) device_worker_->wakeUpRecv();
 }
 
-P2PProxy::SendTransferTask::SendTransferTask(uint64_t chunk_offset_in,
-                                             uint64_t chunk_bytes_in,
-                                             void* source_in,
-                                             uint64_t target_offset_in)
-    : chunk_offset_(chunk_offset_in),
-      chunk_bytes_(chunk_bytes_in),
-      source_(source_in),
-      target_offset_(target_offset_in) {}
+P2PProxy::SendTransferTask::SendTransferTask(
+    uint64_t tensor_offset_in, uint32_t chunk_len_in, void* staging_addr_in,
+    uint64_t remote_addr_in, uint32_t sequence_in, uint32_t epoch_in)
+    : tensor_offset_(tensor_offset_in),
+      chunk_len_(chunk_len_in),
+      staging_addr_(staging_addr_in),
+      remote_addr_(remote_addr_in),
+      sequence_(sequence_in),
+      epoch_(epoch_in) {
+    last_update_time_ = std::chrono::steady_clock::now();
+}
 
 P2PProxy::SendOpContext::SendOpContext(SendOp&& op_in)
-    : tensor_(std::move(op_in.tensor_)),
+    : status_(std::move(op_in.status_)),
+      tensor_(std::move(op_in.tensor_)),
       peer_rank_(op_in.peer_rank_),
-      cuda_stream_(op_in.cuda_stream_),
-      completed_(std::move(op_in.completed_)) {
+      cuda_stream_(op_in.cuda_stream_) {
     total_bytes_ =
         tensor_.numel() * static_cast<uint64_t>(tensor_.element_size());
+    last_update_time_ = std::chrono::steady_clock::now();
 }
 
-P2PProxy::RecvTransferTask::RecvTransferTask(uint64_t chunk_offset_in,
-                                             uint64_t chunk_bytes_in,
-                                             void* source_in, void* target_in)
-    : chunk_offset_(chunk_offset_in),
-      chunk_bytes_(chunk_bytes_in),
-      source_(source_in),
-      target_(target_in) {}
+P2PProxy::RecvTransferTask::RecvTransferTask(uint64_t tensor_offset_in,
+                                             uint32_t chunk_len_in,
+                                             void* local_addr_in,
+                                             uint32_t sequence_in,
+                                             uint32_t epoch_in)
+    : tensor_offset_(tensor_offset_in),
+      chunk_len_(chunk_len_in),
+      local_addr_(local_addr_in),
+      sequence_(sequence_in),
+      epoch_(epoch_in) {
+    last_update_time_ = std::chrono::steady_clock::now();
+}
 
 P2PProxy::RecvOpContext::RecvOpContext(RecvOp&& op_in)
-    : tensor_(std::move(op_in.tensor_)),
+    : status_(std::move(op_in.status_)),
+      tensor_(std::move(op_in.tensor_)),
       original_tensor_(std::move(op_in.original_tensor_)),
       peer_rank_(op_in.peer_rank_),
-      cuda_stream_(op_in.cuda_stream_),
-      completed_(std::move(op_in.completed_)) {
+      cuda_stream_(op_in.cuda_stream_) {
     total_bytes_ =
         tensor_.numel() * static_cast<uint64_t>(tensor_.element_size());
 }
 
-uint64_t P2PProxy::GetLocalSendSlotAddress(int peer_rank,
-                                           uint32_t slot_index) const {
-    const uint64_t send_addr_base =
-        meta_->segmentInfos[rank_].p2p_send_buffer +
-        static_cast<uint64_t>(peer_rank) * kP2PBytesPerRank;
-    return send_addr_base + static_cast<uint64_t>(slot_index) * kP2PSlotSize;
+CreditSlot* P2PProxy::getLocalCreditLane(int peer_rank) const {
+    return resources_.credit_region_ +
+           static_cast<size_t>(peer_rank) * kP2PControlRingSize;
 }
 
-uint64_t P2PProxy::GetLocalRecvSlotAddress(int peer_rank,
-                                           uint32_t slot_index) const {
-    const uint64_t recv_addr_base =
-        meta_->segmentInfos[rank_].p2p_recv_buffer +
-        static_cast<uint64_t>(peer_rank) * kP2PBytesPerRank;
-    return recv_addr_base + static_cast<uint64_t>(slot_index) * kP2PSlotSize;
+AckSlot* P2PProxy::getLocalAckLane(int peer_rank) const {
+    return resources_.ack_region_ +
+           static_cast<size_t>(peer_rank) * kP2PControlRingSize;
 }
 
-uint64_t P2PProxy::GetRemoteRecvSlotAddress(int peer_rank,
-                                            uint32_t slot_index) const {
-    const uint64_t remote_recv_addr_base =
-        meta_->segmentInfos[peer_rank].p2p_recv_buffer +
-        static_cast<uint64_t>(rank_) * kP2PBytesPerRank;
-    return remote_recv_addr_base +
-           static_cast<uint64_t>(slot_index) * kP2PSlotSize;
+uint64_t P2PProxy::getRemoteCreditSlot(int peer_rank, uint32_t sequence) const {
+    const uint64_t slot_index = sequence % kP2PControlRingSize;
+    return meta_->segmentInfos[peer_rank].p2p_credit_region +
+           (static_cast<uint64_t>(rank_) * kP2PControlRingSize + slot_index) *
+               sizeof(CreditSlot);
 }
 
-uint64_t P2PProxy::GetRemoteCtrlRecvHeadOffset(int peer_rank) const {
-    return meta_->segmentInfos[peer_rank].p2p_ctrl_recv +
-           rank_ * sizeof(P2PControlSlot) + offsetof(P2PControlSlot, head);
+uint64_t P2PProxy::getRemoteAckSlot(int peer_rank, uint32_t sequence) const {
+    const uint64_t slot_index = sequence % kP2PControlRingSize;
+    return meta_->segmentInfos[peer_rank].p2p_ack_region +
+           (static_cast<uint64_t>(rank_) * kP2PControlRingSize + slot_index) *
+               sizeof(AckSlot);
 }
 
-uint64_t P2PProxy::GetRemoteCtrlSendTailOffset(int peer_rank) const {
-    return meta_->segmentInfos[peer_rank].p2p_ctrl_send +
-           rank_ * sizeof(P2PControlSlot) + offsetof(P2PControlSlot, tail);
+CreditSlot* P2PProxy::getLocalCreditStagingBuf(int peer_rank,
+                                               uint32_t sequence) const {
+    const size_t staging_idx =
+        static_cast<size_t>(peer_rank) * kP2PControlRingSize +
+        sequence % kP2PControlRingSize;
+    return resources_.credit_staging_buf_ + staging_idx;
+}
+AckSlot* P2PProxy::getLocalAckStagingBuf(int peer_rank,
+                                         uint32_t sequence) const {
+    const size_t staging_idx =
+        static_cast<size_t>(peer_rank) * kP2PControlRingSize +
+        static_cast<size_t>(sequence % kP2PControlRingSize);
+    return resources_.ack_staging_buf_ + staging_idx;
 }
 
-bool P2PProxy::TryIssueSendTask(SendOpContext& op_ctx, uint32_t capacity) {
-    if (op_ctx.bytes_issued_ >= op_ctx.total_bytes_) {
+// Receiver issues credit.
+//
+// Grab a free chunk from RecvPool and grant the sender credit to write into it.
+// We write a CreditSlot into the sender's CreditLane so the sender knows:
+//   (a) which sequence this is,
+//   (b) the RecvPool offset to write to,
+//   (c) how many bytes we expect.
+// If the pool is exhausted we return false and retry on the next polling
+// iteration (non-blocking).
+bool P2PProxy::tryIssueRecvTask(RecvOpContext& op_ctx, RecvPeerLane& lane) {
+    if (op_ctx.bytes_credited_ >= op_ctx.total_bytes_) {
         return false;
     }
 
-    const int peer_rank = op_ctx.peer_rank_;
-    auto& lane = send_peer_lanes_[peer_rank];
-    const uint32_t head = lane.local_head_;
-    const uint32_t slot_index = head;
-    const uint32_t remote_tail =
-        resources_.ctrl_send_region_[peer_rank].tail.load(
-            std::memory_order_acquire);
-    if (((head + 1) % capacity) == remote_tail) {
+    if (op_ctx.tasks_.size() >= kP2PControlRingSize) {
+        // No free control slot, retry on the next iteration
         return false;
     }
 
-    const uint64_t chunk_bytes =
-        std::min(static_cast<uint64_t>(kP2PSlotSize),
-                 op_ctx.total_bytes_ - op_ctx.bytes_issued_);
-    const uint64_t send_addr = GetLocalSendSlotAddress(peer_rank, head);
-    const uint64_t target_offset = GetRemoteRecvSlotAddress(peer_rank, head);
-    op_ctx.tasks_.emplace_back(op_ctx.bytes_issued_, chunk_bytes,
-                               reinterpret_cast<void*>(send_addr),
-                               target_offset);
+    void* local_addr = recv_pool_->acquire();
+    if (local_addr == nullptr) {
+        // No free recv chunk, retry on the next iteration
+        return false;
+    }
+
+    const uint32_t chunk_len = static_cast<uint32_t>(std::min<uint64_t>(
+        chunk_size_, op_ctx.total_bytes_ - op_ctx.bytes_credited_));
+    const uint64_t seq = lane.credit_issue_seq_;
+    const uint64_t remote_credit_offset =
+        getRemoteCreditSlot(op_ctx.peer_rank_, seq);
+
+    const uint32_t curr_epoch =
+        peer_epoch_[op_ctx.peer_rank_].load(std::memory_order_acquire);
+    op_ctx.tasks_.emplace_back(op_ctx.bytes_credited_, chunk_len, local_addr,
+                               seq, curr_epoch);
     auto& task = op_ctx.tasks_.back();
-    const auto* tensor_ptr =
-        static_cast<const uint8_t*>(op_ctx.tensor_.data_ptr());
-
-    if (is_cpu_) {
-        std::memcpy(task.source_, tensor_ptr + task.chunk_offset_,
-                    task.chunk_bytes_);
-        task.state_ = TransferState::kTransfer;
-    } else {
-        cudaError_t copy_error = cudaMemcpyAsync(
-            task.source_, tensor_ptr + task.chunk_offset_, task.chunk_bytes_,
-            cudaMemcpyDeviceToDevice, op_ctx.cuda_stream_);
-        TORCH_CHECK(!copy_error, "P2P send cudaMemcpyAsync failed: ",
-                    cudaGetErrorString(copy_error));
-        const cudaEvent_t pooled_copy_ready_event =
-            lane.copy_ready_events_[slot_index];
-        TORCH_CHECK(pooled_copy_ready_event != nullptr,
-                    "P2P send pooled copy-ready event is not initialized.");
-        task.copy_ready_event_ = pooled_copy_ready_event;
-        copy_error =
-            cudaEventRecord(task.copy_ready_event_, op_ctx.cuda_stream_);
-        if (copy_error != cudaSuccess) {
-            task.copy_ready_event_ = nullptr;
-            TORCH_CHECK(false, "P2P send cudaEventRecord failed: ",
-                        cudaGetErrorString(copy_error));
-        }
-    }
-
-    op_ctx.bytes_issued_ += task.chunk_bytes_;
-    lane.local_head_ = (head + 1) % capacity;
-    return true;
-}
-
-bool P2PProxy::StepSendTransferTask(SendOpContext& op_ctx,
-                                    SendTransferTask& task) {
-    bool did_work = false;
-
-    if (task.state_ == TransferState::kDataCopy && StepSendDataCopy(task)) {
-        did_work = true;
-    }
-    if (task.state_ == TransferState::kTransfer &&
-        StepSendTransfer(op_ctx, task)) {
-        did_work = true;
-    }
-
-    return did_work;
-}
-
-bool P2PProxy::StepSendDataCopy(SendTransferTask& task) {
-    if (task.copy_ready_event_ == nullptr) {
-        task.state_ = TransferState::kTransfer;
-        return true;
-    }
-
-    cudaError_t query_error = cudaSuccess;
-    query_error = cudaEventQuery(task.copy_ready_event_);
-
-    if (query_error == cudaSuccess) {
-        task.copy_ready_event_ = nullptr;
-        task.state_ = TransferState::kTransfer;
-        return true;
-    }
-    if (query_error == cudaErrorNotReady) {
-        return false;
-    }
-
-    task.copy_ready_event_ = nullptr;
-    TORCH_CHECK(false, "P2P send cudaEventQuery failed: ",
-                cudaGetErrorString(query_error));
-    return false;
-}
-
-bool P2PProxy::StepSendTransfer(SendOpContext& op_ctx, SendTransferTask& task) {
-    bool did_work = false;
-    if (!task.transfer_batch_id_.has_value()) {
-        const BatchID batch_id = engine_->allocateBatchID(1);
-        engine_->submitTransfer(
-            batch_id, {TransferRequest{
-                          .opcode = TransferRequest::WRITE,
-                          .source = task.source_,
-                          .target_id = meta_->segmentIDs[op_ctx.peer_rank_],
-                          .target_offset = task.target_offset_,
-                          .length = task.chunk_bytes_,
-                      }});
-        task.transfer_batch_id_ = batch_id;
-        did_work = true;
-    }
-
-    TransferStatus transfer_status;
-    engine_->getTransferStatus(task.transfer_batch_id_.value(), 0,
-                               transfer_status);
-    if (transfer_status.s == TransferStatusEnum::COMPLETED) {
-        engine_->freeBatchID(task.transfer_batch_id_.value());
-        task.transfer_batch_id_.reset();
-        task.state_ = TransferState::kDone;
-        return true;
-    }
-    if (transfer_status.s == TransferStatusEnum::FAILED) {
-        engine_->freeBatchID(task.transfer_batch_id_.value());
-        task.transfer_batch_id_.reset();
-        TORCH_CHECK(false, "P2P send transfer failed.");
-        return false;
-    }
-
-    return did_work;
-}
-
-bool P2PProxy::StepSendHeadCommit(SendOpContext& op_ctx, uint32_t capacity) {
-    bool did_work = false;
-    if (op_ctx.head_update_batch_id_.has_value()) {
-        TransferStatus head_status;
-        engine_->getTransferStatus(op_ctx.head_update_batch_id_.value(), 0,
-                                   head_status);
-        if (head_status.s == TransferStatusEnum::COMPLETED) {
-            engine_->freeBatchID(op_ctx.head_update_batch_id_.value());
-            op_ctx.head_update_batch_id_.reset();
-            did_work = true;
-        } else if (head_status.s == TransferStatusEnum::FAILED) {
-            engine_->freeBatchID(op_ctx.head_update_batch_id_.value());
-            op_ctx.head_update_batch_id_.reset();
-            TORCH_CHECK(false, "P2P ctrl head update failed.");
-            return false;
-        }
-    }
-
-    if (op_ctx.head_update_batch_id_.has_value()) {
-        return did_work;
-    }
-
-    uint32_t committed_tasks = 0;
-    while (!op_ctx.tasks_.empty() &&
-           op_ctx.tasks_.front().state_ == TransferState::kDone) {
-        op_ctx.tasks_.pop_front();
-        ++committed_tasks;
-        did_work = true;
-    }
-
-    if (committed_tasks == 0) {
-        return did_work;
-    }
-
-    const uint32_t current_head =
-        resources_.ctrl_send_region_[op_ctx.peer_rank_].head.load(
-            std::memory_order_relaxed);
-    const uint32_t next_head = (current_head + committed_tasks) % capacity;
-    resources_.ctrl_send_region_[op_ctx.peer_rank_].head.store(
-        next_head, std::memory_order_release);
-    void* head_source = static_cast<void*>(
-        &resources_.ctrl_send_region_[op_ctx.peer_rank_].head.value);
-    const uint64_t remote_head_offset =
-        GetRemoteCtrlRecvHeadOffset(op_ctx.peer_rank_);
 
+    auto* credit_staging_buf = getLocalCreditStagingBuf(op_ctx.peer_rank_, seq);
+    credit_staging_buf->publish(
+        curr_epoch, seq, reinterpret_cast<uint64_t>(local_addr), chunk_len);
     const BatchID batch_id = engine_->allocateBatchID(1);
     engine_->submitTransfer(
         batch_id, {TransferRequest{
                       .opcode = TransferRequest::WRITE,
-                      .source = head_source,
+                      .source = static_cast<void*>(credit_staging_buf),
                       .target_id = meta_->segmentIDs[op_ctx.peer_rank_],
-                      .target_offset = remote_head_offset,
-                      .length = sizeof(uint32_t),
+                      .target_offset = remote_credit_offset,
+                      .length = sizeof(CreditSlot),
                   }});
-    op_ctx.head_update_batch_id_ = batch_id;
-    did_work = true;
+    task.credit_batch_id_ = batch_id;
+    task.last_update_time_ = std::chrono::steady_clock::now();
 
-    return did_work;
-}
-
-bool P2PProxy::IsSendDataPathCompleted(const SendOpContext& op_ctx) const {
-    return op_ctx.bytes_issued_ == op_ctx.total_bytes_ && op_ctx.tasks_.empty();
-}
-
-bool P2PProxy::IsSendOpCompleted(const SendOpContext& op_ctx) const {
-    return IsSendDataPathCompleted(op_ctx) &&
-           !op_ctx.head_update_batch_id_.has_value();
-}
-
-bool P2PProxy::TryIssueRecvTask(RecvOpContext& op_ctx, uint32_t capacity) {
-    if (op_ctx.bytes_issued_ >= op_ctx.total_bytes_) {
-        return false;
-    }
-
-    auto* local_ctrl = &resources_.ctrl_recv_region_[op_ctx.peer_rank_];
-    auto& lane = recv_peer_lanes_[op_ctx.peer_rank_];
-    const uint32_t head = local_ctrl->head.load(std::memory_order_acquire);
-    const uint32_t tail = lane.local_tail_;
-    const uint32_t slot_index = tail;
-    if (head == tail) {
-        return false;
-    }
-
-    const uint64_t chunk_bytes =
-        std::min(static_cast<uint64_t>(kP2PSlotSize),
-                 op_ctx.total_bytes_ - op_ctx.bytes_issued_);
-    const uint64_t recv_addr = GetLocalRecvSlotAddress(op_ctx.peer_rank_, tail);
-    auto* tensor_ptr = static_cast<uint8_t*>(op_ctx.tensor_.data_ptr());
-    void* target_ptr = static_cast<void*>(tensor_ptr + op_ctx.bytes_issued_);
-
-    op_ctx.tasks_.emplace_back(op_ctx.bytes_issued_, chunk_bytes,
-                               reinterpret_cast<void*>(recv_addr), target_ptr);
-    auto& task = op_ctx.tasks_.back();
-
-    if (is_cpu_) {
-        std::memcpy(task.target_, task.source_, task.chunk_bytes_);
-        task.state_ = TransferState::kDone;
-    } else {
-        cudaError_t copy_error =
-            cudaMemcpyAsync(task.target_, task.source_, task.chunk_bytes_,
-                            cudaMemcpyDeviceToDevice, op_ctx.cuda_stream_);
-        TORCH_CHECK(!copy_error, "P2P recv cudaMemcpyAsync failed: ",
-                    cudaGetErrorString(copy_error));
-        const cudaEvent_t pooled_copy_ready_event =
-            lane.copy_ready_events_[slot_index];
-        TORCH_CHECK(pooled_copy_ready_event != nullptr,
-                    "P2P recv pooled copy-ready event is not initialized.");
-        task.copy_ready_event_ = pooled_copy_ready_event;
-        copy_error =
-            cudaEventRecord(task.copy_ready_event_, op_ctx.cuda_stream_);
-        if (copy_error != cudaSuccess) {
-            task.copy_ready_event_ = nullptr;
-            TORCH_CHECK(false, "P2P recv cudaEventRecord failed: ",
-                        cudaGetErrorString(copy_error));
-        }
-    }
-
-    op_ctx.bytes_issued_ += task.chunk_bytes_;
-    lane.local_tail_ = (tail + 1) % capacity;
+    ++lane.credit_issue_seq_;
+    op_ctx.bytes_credited_ += chunk_len;
     return true;
 }
 
-bool P2PProxy::StepRecvTransferTask(RecvTransferTask& task) {
-    bool did_work = false;
+// Drive a single receiver chunk through its state machine.
+//
+// kIssueCredit -> kWaitAck -> kCopyOut -> kFinished -> erase.
+bool P2PProxy::stepRecvTask(RecvTransferTask& task) {
+    switch (task.state_) {
+        case RecvTaskState::kIssueCredit:
+            return stepRecvIssueCredit(task);
 
-    if (task.state_ == TransferState::kDataCopy && StepRecvDataCopy(task)) {
-        did_work = true;
+        case RecvTaskState::kWaitAck:
+            // kWaitAck is polled in stepRecv
+            return false;
+
+        case RecvTaskState::kCopyOut:
+            return stepRecvCopyOut(task);
+
+        case RecvTaskState::kFinished:
+        case RecvTaskState::kFailed:
+            return false;
     }
-
-    return did_work;
+    return false;
 }
 
-bool P2PProxy::StepRecvDataCopy(RecvTransferTask& task) {
-    if (task.copy_ready_event_ == nullptr) {
-        task.state_ = TransferState::kDone;
+bool P2PProxy::stepRecvIssueCredit(RecvTransferTask& task) {
+    TORCH_CHECK(task.credit_batch_id_.has_value(),
+                "Expected a credit_batch_id in tryIssueRecvTask");
+
+    TransferStatus credit_status;
+    engine_->getTransferStatus(task.credit_batch_id_.value(), 0, credit_status);
+
+    if (credit_status.s == TransferStatusEnum::COMPLETED) {
+        engine_->freeBatchID(task.credit_batch_id_.value());
+        task.credit_batch_id_.reset();
+        task.state_ = RecvTaskState::kWaitAck;
+        task.last_update_time_ = std::chrono::steady_clock::now();
         return true;
     }
 
-    cudaError_t query_error = cudaSuccess;
-    query_error = cudaEventQuery(task.copy_ready_event_);
+    if (credit_status.s == TransferStatusEnum::FAILED || isTimeout(task)) {
+        LOG(ERROR) << "P2P credit transfer failed/timeout, seq="
+                   << task.sequence_;
+        engine_->freeBatchID(task.credit_batch_id_.value());
+        task.credit_batch_id_.reset();
+        task.state_ = RecvTaskState::kFailed;
+        return true;
+    }
+    return false;
+}
 
+// Poll the GPU Copy-Out event.
+//
+// After the AckSlot arrives we initiate cudaMemcpyAsync from the
+// RecvPool chunk to the user tensor and record an event.  This function
+// waits for that event and returns the chunk to RecvPool.
+bool P2PProxy::stepRecvCopyOut(RecvTransferTask& task) {
+    TORCH_CHECK(task.copy_ready_event_ != nullptr,
+                "Expected a copy_ready_event");
+    cudaError_t query_error = cudaEventQuery(task.copy_ready_event_);
     if (query_error == cudaSuccess) {
         task.copy_ready_event_ = nullptr;
-        task.state_ = TransferState::kDone;
+        recv_pool_->release(task.local_addr_);
+        task.local_addr_ = nullptr;
+        task.state_ = RecvTaskState::kFinished;
         return true;
     }
     if (query_error == cudaErrorNotReady) {
@@ -680,82 +697,286 @@ bool P2PProxy::StepRecvDataCopy(RecvTransferTask& task) {
     return false;
 }
 
-bool P2PProxy::StepRecvTailCommit(RecvOpContext& op_ctx, uint32_t capacity) {
-    bool did_work = false;
-    if (op_ctx.tail_update_batch_id_.has_value()) {
-        TransferStatus tail_status;
-        engine_->getTransferStatus(op_ctx.tail_update_batch_id_.value(), 0,
-                                   tail_status);
-        if (tail_status.s == TransferStatusEnum::COMPLETED) {
-            engine_->freeBatchID(op_ctx.tail_update_batch_id_.value());
-            op_ctx.tail_update_batch_id_.reset();
-            did_work = true;
-        } else if (tail_status.s == TransferStatusEnum::FAILED) {
-            engine_->freeBatchID(op_ctx.tail_update_batch_id_.value());
-            op_ctx.tail_update_batch_id_.reset();
-            TORCH_CHECK(false, "P2P ctrl tail update failed.");
-            return false;
+// Sender fetches credits.
+//
+// Poll the local CreditLane for the next expected sequence.  If the slot
+// matches our consume cursor we accept the credit: allocate a staging
+// chunk from SendPool, copy the corresponding slice of the user tensor into
+// it, and advance the cursor.  If the pool is full or no credit has
+// arrived we return false immediately (non-blocking).
+bool P2PProxy::tryIssueSendTask(SendOpContext& op_ctx, SendPeerLane& lane) {
+    // Check for timeout while waiting for the peer's CreditSlot.
+    if (isTimeout(op_ctx)) {
+        LOG(ERROR) << "P2P wait-for-credit timeout, peer=" << op_ctx.peer_rank_;
+        op_ctx.status_->store(OpStatus::kFailed, std::memory_order_release);
+        return false;
+    }
+
+    if (op_ctx.bytes_staged_ >= op_ctx.total_bytes_) {
+        return false;
+    }
+
+    if (op_ctx.tasks_.size() >= kP2PControlRingSize) {
+        return false;
+    }
+
+    CreditSlot* local_credit = getLocalCreditLane(op_ctx.peer_rank_);
+    const uint64_t seq = lane.credit_consume_seq_;
+    auto& slot = local_credit[static_cast<size_t>(seq % kP2PControlRingSize)];
+
+    // Step 1 -- Try load the slot
+    uint64_t recv_addr = 0;
+    uint32_t chunk_len = 0;
+    uint32_t slot_epoch = 0;
+    uint32_t slot_seq = 0;
+    if (!slot.tryLoad(recv_addr, chunk_len, slot_epoch, slot_seq)) {
+        // Slot is either empty or torn (partial RDMA write). Retry next poll.
+        return false;
+    }
+
+    // Step 2 -- Stale packet: the slot carries data from a previous epoch
+    // (before a Reset).  Clear it so the fresh credit can land safely.
+    const uint32_t curr_epoch =
+        peer_epoch_[op_ctx.peer_rank_].load(std::memory_order_acquire);
+    if (slot_epoch != curr_epoch) {
+        LOG(WARNING) << "[P2PProxy][Send] tryIssueSendTask peer="
+                     << op_ctx.peer_rank_ << " STALE_EPOCH seq=" << seq
+                     << " slot.epoch=" << slot_epoch
+                     << " curr_epoch=" << curr_epoch;
+        slot.reset();
+        return true;
+    }
+
+    // Step 3 -- Sequence check: make sure this is exactly the slot we expect.
+    if (slot_seq != static_cast<uint32_t>(seq)) {
+        return false;
+    }
+
+    void* staging_addr = send_pool_->acquire();
+    if (staging_addr == nullptr) {
+        return false;
+    }
+
+    const uint32_t expected_chunk_len =
+        static_cast<uint32_t>(std::min<uint64_t>(
+            chunk_size_, op_ctx.total_bytes_ - op_ctx.bytes_staged_));
+    TORCH_CHECK(chunk_len <= expected_chunk_len,
+                "P2P send got invalid chunk_len in credit slot.");
+
+    op_ctx.tasks_.emplace_back(op_ctx.bytes_staged_, chunk_len, staging_addr,
+                               recv_addr, seq, slot_epoch);
+    auto& task = op_ctx.tasks_.back();
+
+    slot.reset();
+
+    const auto* tensor_ptr =
+        static_cast<const uint8_t*>(op_ctx.tensor_.data_ptr());
+    if (is_cpu_) {
+        std::memcpy(staging_addr, tensor_ptr + task.tensor_offset_,
+                    task.chunk_len_);
+        task.state_ = SendTaskState::kWriteRemote;
+    } else {
+        cudaError_t copy_error = cudaMemcpyAsync(
+            staging_addr, tensor_ptr + task.tensor_offset_, task.chunk_len_,
+            cudaMemcpyDeviceToDevice, op_ctx.cuda_stream_);
+        TORCH_CHECK(!copy_error, "P2P send cudaMemcpyAsync failed: ",
+                    cudaGetErrorString(copy_error));
+
+        const cudaEvent_t pooled_copy_ready_event =
+            lane.copy_ready_events_[static_cast<size_t>(seq %
+                                                        kP2PControlRingSize)];
+        TORCH_CHECK(pooled_copy_ready_event != nullptr,
+                    "P2P send pooled copy-ready event is not initialized.");
+        task.copy_ready_event_ = pooled_copy_ready_event;
+        copy_error =
+            cudaEventRecord(task.copy_ready_event_, op_ctx.cuda_stream_);
+        if (copy_error != cudaSuccess) {
+            task.copy_ready_event_ = nullptr;
+            TORCH_CHECK(false, "P2P send cudaEventRecord failed: ",
+                        cudaGetErrorString(copy_error));
         }
     }
 
-    if (op_ctx.tail_update_batch_id_.has_value()) {
-        return did_work;
+    ++lane.credit_consume_seq_;
+    op_ctx.bytes_staged_ += task.chunk_len_;
+    task.last_update_time_ = std::chrono::steady_clock::now();
+    return true;
+}
+
+// Drive a single sender chunk through its state machine.
+//
+// kCopyIn -> kWriteRemote -> kAck -> kFinished -> erase.
+bool P2PProxy::stepSendTask(SendOpContext& op_ctx, SendTransferTask& task) {
+    switch (task.state_) {
+        case SendTaskState::kCopyIn:
+            return stepSendCopyIn(task);
+
+        case SendTaskState::kWriteRemote:
+            return stepSendWriteRemote(op_ctx, task);
+
+        case SendTaskState::kAck:
+            return stepSendAck(op_ctx, task);
+
+        case SendTaskState::kFinished:
+        case SendTaskState::kFailed:
+            return false;
+    }
+    return false;
+}
+
+// Poll the GPU Copy-In event (or skip on CPU).
+//
+// When the event signals we transition from kCopyIn to kWriteRemote so the
+// (RDMA) Write can be submitted.
+bool P2PProxy::stepSendCopyIn(SendTransferTask& task) {
+    if (task.copy_ready_event_ == nullptr) {
+        // CPU case
+        task.state_ = SendTaskState::kWriteRemote;
+        task.last_update_time_ = std::chrono::steady_clock::now();
+        return true;
     }
 
-    uint32_t committed_tasks = 0;
-    while (!op_ctx.tasks_.empty() &&
-           op_ctx.tasks_.front().state_ == TransferState::kDone) {
-        op_ctx.tasks_.pop_front();
-        ++committed_tasks;
+    cudaError_t query_error = cudaEventQuery(task.copy_ready_event_);
+    if (query_error == cudaErrorNotReady) {
+        return false;
+    }
+
+    task.copy_ready_event_ = nullptr;
+
+    if (query_error == cudaSuccess) {
+        task.state_ = SendTaskState::kWriteRemote;
+        task.last_update_time_ = std::chrono::steady_clock::now();
+        return true;
+    }
+
+    TORCH_CHECK(false, "P2P send cudaEventQuery failed: ",
+                cudaGetErrorString(query_error));
+    return false;
+}
+
+// Submit the Write from SendPool staging to remote RecvPool, then
+// poll for its completion.
+bool P2PProxy::stepSendWriteRemote(SendOpContext& op_ctx,
+                                   SendTransferTask& task) {
+    bool did_work = false;
+    if (!task.transfer_batch_id_.has_value()) {
+        const BatchID batch_id = engine_->allocateBatchID(1);
+        engine_->submitTransfer(
+            batch_id, {TransferRequest{
+                          .opcode = TransferRequest::WRITE,
+                          .source = task.staging_addr_,
+                          .target_id = meta_->segmentIDs[op_ctx.peer_rank_],
+                          .target_offset = task.remote_addr_,
+                          .length = task.chunk_len_,
+                      }});
+        task.transfer_batch_id_ = batch_id;
+        task.last_update_time_ = std::chrono::steady_clock::now();
         did_work = true;
     }
 
-    if (committed_tasks == 0) {
-        return did_work;
+    TransferStatus transfer_status;
+    engine_->getTransferStatus(task.transfer_batch_id_.value(), 0,
+                               transfer_status);
+    if (transfer_status.s == TransferStatusEnum::COMPLETED) {
+        engine_->freeBatchID(task.transfer_batch_id_.value());
+        task.transfer_batch_id_.reset();
+        task.state_ = SendTaskState::kAck;
+        task.last_update_time_ = std::chrono::steady_clock::now();
+        did_work = true;
+    } else if (transfer_status.s == TransferStatusEnum::FAILED ||
+               isTimeout(task)) {
+        LOG(ERROR) << "P2P send transfer failed/timeout, peer="
+                   << op_ctx.peer_rank_ << ", seq=" << task.sequence_;
+        engine_->freeBatchID(task.transfer_batch_id_.value());
+        task.transfer_batch_id_.reset();
+        task.state_ = SendTaskState::kFailed;
+        did_work = true;
     }
 
-    auto* local_ctrl = &resources_.ctrl_recv_region_[op_ctx.peer_rank_];
-    const uint32_t current_tail =
-        local_ctrl->tail.load(std::memory_order_relaxed);
-    const uint32_t next_tail = (current_tail + committed_tasks) % capacity;
-    local_ctrl->tail.store(next_tail, std::memory_order_release);
-
-    const uint64_t remote_tail_offset =
-        GetRemoteCtrlSendTailOffset(op_ctx.peer_rank_);
-    void* tail_source = static_cast<void*>(&local_ctrl->tail.value);
-
-    const BatchID batch_id = engine_->allocateBatchID(1);
-    engine_->submitTransfer(
-        batch_id, {TransferRequest{
-                      .opcode = TransferRequest::WRITE,
-                      .source = tail_source,
-                      .target_id = meta_->segmentIDs[op_ctx.peer_rank_],
-                      .target_offset = remote_tail_offset,
-                      .length = sizeof(uint32_t),
-                  }});
-    op_ctx.tail_update_batch_id_ = batch_id;
-    did_work = true;
-
     return did_work;
 }
 
-bool P2PProxy::IsRecvDataPathCompleted(const RecvOpContext& op_ctx) const {
-    return op_ctx.bytes_issued_ == op_ctx.total_bytes_ &&
-           op_ctx.tasks_.empty() && !op_ctx.tail_update_batch_id_.has_value();
+// Write the AckSlot back to the receiver.
+//
+// This tells the receiver that the data is present in its RecvPool and it
+// may begin the Copy-Out.  When the AckSlot write finishes we free
+// the staging chunk and transition to kFinished.
+bool P2PProxy::stepSendAck(SendOpContext& op_ctx, SendTransferTask& task) {
+    bool did_work = false;
+    if (!task.ack_batch_id_.has_value()) {
+        auto* ack_staging_buf =
+            getLocalAckStagingBuf(op_ctx.peer_rank_, task.sequence_);
+        ack_staging_buf->publish(task.epoch_, task.sequence_, task.chunk_len_);
+
+        const BatchID batch_id = engine_->allocateBatchID(1);
+        engine_->submitTransfer(
+            batch_id, {TransferRequest{
+                          .opcode = TransferRequest::WRITE,
+                          .source = static_cast<void*>(ack_staging_buf),
+                          .target_id = meta_->segmentIDs[op_ctx.peer_rank_],
+                          .target_offset = getRemoteAckSlot(op_ctx.peer_rank_,
+                                                            task.sequence_),
+                          .length = sizeof(AckSlot),
+                      }});
+        task.ack_batch_id_ = batch_id;
+        task.last_update_time_ = std::chrono::steady_clock::now();
+        did_work = true;
+    }
+
+    TransferStatus ack_status;
+    engine_->getTransferStatus(task.ack_batch_id_.value(), 0, ack_status);
+    if (ack_status.s == TransferStatusEnum::COMPLETED) {
+        engine_->freeBatchID(task.ack_batch_id_.value());
+        task.ack_batch_id_.reset();
+        send_pool_->release(task.staging_addr_);
+        task.staging_addr_ = nullptr;
+        task.state_ = SendTaskState::kFinished;
+        did_work = true;
+    } else if (ack_status.s == TransferStatusEnum::FAILED || isTimeout(task)) {
+        LOG(ERROR) << "P2P ack transfer failed/timeout, peer="
+                   << op_ctx.peer_rank_ << ", seq=" << task.sequence_;
+        engine_->freeBatchID(task.ack_batch_id_.value());
+        task.ack_batch_id_.reset();
+        task.state_ = SendTaskState::kFailed;
+        did_work = true;
+    }
+
+    return did_work;
 }
 
-bool P2PProxy::StepSend() {
-    const uint32_t capacity = static_cast<uint32_t>(kP2PNumSlotsPerRank);
+bool P2PProxy::isSendOpCompleted(const SendOpContext& op_ctx) const {
+    return op_ctx.bytes_staged_ == op_ctx.total_bytes_ && op_ctx.tasks_.empty();
+}
+
+bool P2PProxy::isRecvOpCompleted(const RecvOpContext& op_ctx) const {
+    return op_ctx.bytes_credited_ == op_ctx.total_bytes_ &&
+           op_ctx.tasks_.empty();
+}
+
+// Sender state machine
+//
+// Pipeline per peer:
+//   1. Drain the shared send_queue_ into the peer's pending_send_ops_.
+//   2. Promote the first pending op to active_send_op_.
+//   3. While we have credits (CreditSlots) and staging buffers,
+//      pull more tensor slices into SendPool (TryIssueSendTask).
+//   4. Advance every active chunk through its state machine
+//      (Copy-In -> Write Remote -> Ack write).
+//   5. Erase fully-acknowledged chunks.  When the op is empty and all
+//      bytes have been staged, mark it complete.
+bool P2PProxy::stepSend() {
     bool did_work = false;
 
+    // Handle reset first
     for (int peer_rank = 0; peer_rank < size_; ++peer_rank) {
         if (reset_send_req_[peer_rank].exchange(false,
                                                 std::memory_order_acquire)) {
-            PerformSendReset(peer_rank);
+            performSendReset(peer_rank);
             did_work = true;
         }
     }
 
+    // Drain send_queue_
     {
         std::lock_guard<std::mutex> lock(send_queue_mutex_);
         while (!send_queue_.empty()) {
@@ -767,6 +988,7 @@ bool P2PProxy::StepSend() {
         }
     }
 
+    // Promote active op
     for (int peer_rank = 0; peer_rank < size_; ++peer_rank) {
         auto& lane = send_peer_lanes_[peer_rank];
         if (lane.active_send_op_.has_value() ||
@@ -776,7 +998,8 @@ bool P2PProxy::StepSend() {
         SendOpContext op_ctx = std::move(lane.pending_send_ops_.front());
         lane.pending_send_ops_.pop_front();
         if (op_ctx.total_bytes_ == 0) {
-            op_ctx.completed_->store(true, std::memory_order_release);
+            op_ctx.status_->store(OpStatus::kSuccess,
+                                  std::memory_order_release);
             active_send_tasks_.fetch_sub(1, std::memory_order_release);
             did_work = true;
             continue;
@@ -785,28 +1008,49 @@ bool P2PProxy::StepSend() {
         did_work = true;
     }
 
+    // Advance state machine
     for (int peer_rank = 0; peer_rank < size_; ++peer_rank) {
         auto& lane = send_peer_lanes_[peer_rank];
         if (!lane.active_send_op_.has_value()) {
             continue;
         }
         auto& op_ctx = lane.active_send_op_.value();
-        while (TryIssueSendTask(op_ctx, capacity)) {
+
+        // Pull as many credits as we have free chunks
+        while (tryIssueSendTask(op_ctx, lane)) {
             did_work = true;
         }
-        for (auto& task : op_ctx.tasks_) {
-            if (task.state_ == TransferState::kDone) {
-                continue;
-            }
-            if (StepSendTransferTask(op_ctx, task)) {
-                did_work = true;
+
+        // Advance every chunk through the sender state machine.
+        // send op may fail due to credit-wait timeout in tryIssueSendTask.
+        auto op_status = op_ctx.status_->load(std::memory_order_acquire);
+        bool op_failed = op_status == OpStatus::kFailed;
+        if (!op_failed) {
+            for (auto it = op_ctx.tasks_.begin(); it != op_ctx.tasks_.end();) {
+                if (stepSendTask(op_ctx, *it)) {
+                    did_work = true;
+                }
+                if (it->state_ == SendTaskState::kFinished) {
+                    it = op_ctx.tasks_.erase(it);
+                    did_work = true;
+                } else if (it->state_ == SendTaskState::kFailed) {
+                    op_failed = true;
+                    break;
+                } else {
+                    ++it;
+                }
             }
         }
-        if (StepSendHeadCommit(op_ctx, capacity)) {
+
+        if (op_failed) {
+            reportBrokenPeer(peer_rank);
             did_work = true;
+            continue;
         }
-        if (IsSendOpCompleted(op_ctx)) {
-            op_ctx.completed_->store(true, std::memory_order_release);
+
+        if (isSendOpCompleted(op_ctx)) {
+            op_ctx.status_->store(OpStatus::kSuccess,
+                                  std::memory_order_release);
             lane.active_send_op_.reset();
             active_send_tasks_.fetch_sub(1, std::memory_order_release);
             did_work = true;
@@ -816,18 +1060,114 @@ bool P2PProxy::StepSend() {
     return did_work;
 }
 
-bool P2PProxy::StepRecv() {
-    const uint32_t capacity = static_cast<uint32_t>(kP2PNumSlotsPerRank);
+// Poll the local AckLane for the head task and initiate Copy-Out.
+//
+// Returns true if an ack was processed (including stale-slot clearing).
+bool P2PProxy::pollRecvAckSlot(RecvOpContext& op_ctx, RecvPeerLane& lane,
+                               RecvTransferTask& head_task) {
+    // Check for timeout while waiting for the peer's AckSlot.
+    if (isTimeout(head_task)) {
+        LOG(ERROR) << "P2P wait-for-ack timeout, peer=" << op_ctx.peer_rank_
+                   << " seq=" << head_task.sequence_;
+        head_task.state_ = RecvTaskState::kFailed;
+        return true;
+    }
+
+    auto* ack_lane = getLocalAckLane(op_ctx.peer_rank_);
+    auto& slot = ack_lane[head_task.sequence_ % kP2PControlRingSize];
+
+    // Step 1 -- Try load the slot
+    uint32_t ack_len = 0;
+    uint32_t slot_epoch = 0;
+    uint32_t slot_seq = 0;
+    if (!slot.tryLoad(ack_len, slot_epoch, slot_seq)) {
+        // Slot is either empty or torn. Retry next poll.
+        return false;
+    }
+
+    const uint32_t curr_epoch =
+        peer_epoch_[op_ctx.peer_rank_].load(std::memory_order_acquire);
+
+    // Step 2 -- Stale packet: data from a previous epoch (before Reset).
+    // Clear the slot so the fresh ack can land safely.
+    if (slot_epoch != curr_epoch) {
+        LOG(WARNING) << "[P2PProxy][Recv] pollRecvAckSlot peer="
+                     << op_ctx.peer_rank_
+                     << " front-seq=" << head_task.sequence_
+                     << " EPOCH_MISMATCH slot.epoch=" << slot_epoch
+                     << " curr_epoch=" << curr_epoch;
+        slot.reset();
+        return true;
+    }
+
+    // Step 3 -- Sequence check: make sure this is the exact ack
+    // we are waiting for.
+    if (slot_seq != head_task.sequence_) {
+        return false;
+    }
+
+    void* src_ptr = head_task.local_addr_;
+    auto* tensor_ptr = static_cast<uint8_t*>(op_ctx.tensor_.data_ptr());
+    void* dst_ptr = tensor_ptr + head_task.tensor_offset_;
+
+    if (is_cpu_) {
+        std::memcpy(dst_ptr, src_ptr, head_task.chunk_len_);
+        recv_pool_->release(head_task.local_addr_);
+        head_task.local_addr_ = nullptr;
+        head_task.state_ = RecvTaskState::kFinished;
+    } else {
+        cudaError_t copy_error =
+            cudaMemcpyAsync(dst_ptr, src_ptr, head_task.chunk_len_,
+                            cudaMemcpyDeviceToDevice, op_ctx.cuda_stream_);
+        TORCH_CHECK(!copy_error, "P2P recv cudaMemcpyAsync failed: ",
+                    cudaGetErrorString(copy_error));
+        const cudaEvent_t pooled_copy_ready_event =
+            lane.copy_ready_events_[static_cast<size_t>(head_task.sequence_ %
+                                                        kP2PControlRingSize)];
+        TORCH_CHECK(pooled_copy_ready_event != nullptr,
+                    "P2P recv pooled copy-ready event is not initialized.");
+        head_task.copy_ready_event_ = pooled_copy_ready_event;
+        copy_error =
+            cudaEventRecord(head_task.copy_ready_event_, op_ctx.cuda_stream_);
+        if (copy_error != cudaSuccess) {
+            head_task.copy_ready_event_ = nullptr;
+            TORCH_CHECK(false, "P2P recv cudaEventRecord failed: ",
+                        cudaGetErrorString(copy_error));
+        }
+        head_task.state_ = RecvTaskState::kCopyOut;
+    }
+
+    head_task.last_update_time_ = std::chrono::steady_clock::now();
+    slot.reset();
+    ++lane.ack_consume_seq_;
+    return true;
+}
+
+// Receiver state machine
+//
+// Pipeline per peer:
+//   1. Drain the shared recv_queue_ into the peer's pending_recv_ops_.
+//   2. Promote the first pending op to active_recv_op_.
+//   3. While we have free RecvPool chunks, issue CreditSlots to the sender
+//      (TryIssueRecvTask).
+//   4. Poll the local AckLane in order.  When an AckSlot
+//      arrives, initiate Copy-Out from RecvPool to the user tensor.
+//   5. Advance every active chunk through its state machine
+//      (IssueCredit -> Copy-Out -> erase).
+//   6. When all chunks are copied out, mark the op complete.
+bool P2PProxy::stepRecv() {
     bool did_work = false;
 
+    // Handle reset first
     for (int peer_rank = 0; peer_rank < size_; ++peer_rank) {
         if (reset_recv_req_[peer_rank].exchange(false,
                                                 std::memory_order_acquire)) {
-            PerformRecvReset(peer_rank);
+            performRecvReset(peer_rank);
             did_work = true;
         }
     }
 
+    // Drain recv_queue_
     {
         std::lock_guard<std::mutex> lock(recv_queue_mutex_);
         while (!recv_queue_.empty()) {
@@ -839,6 +1179,7 @@ bool P2PProxy::StepRecv() {
         }
     }
 
+    // Promote active op
     for (int peer_rank = 0; peer_rank < size_; ++peer_rank) {
         auto& lane = recv_peer_lanes_[peer_rank];
         if (lane.active_recv_op_.has_value() ||
@@ -847,36 +1188,62 @@ bool P2PProxy::StepRecv() {
         }
         RecvOp recv_op = std::move(lane.pending_recv_ops_.front());
         lane.pending_recv_ops_.pop_front();
-        lane.local_tail_ = resources_.ctrl_recv_region_[peer_rank].tail.load(
-            std::memory_order_acquire);
         RecvOpContext op_ctx(std::move(recv_op));
         lane.active_recv_op_ = std::move(op_ctx);
         did_work = true;
     }
 
+    // Advance state machine
     for (int peer_rank = 0; peer_rank < size_; ++peer_rank) {
         auto& lane = recv_peer_lanes_[peer_rank];
         if (!lane.active_recv_op_.has_value()) {
             continue;
         }
+
         auto& op_ctx = lane.active_recv_op_.value();
-        while (TryIssueRecvTask(op_ctx, capacity)) {
+
+        // Offer as many RecvPool chunks as we have free buffers and tensor
+        // bytes remaining.
+        while (tryIssueRecvTask(op_ctx, lane)) {
             did_work = true;
         }
-        for (auto& task : op_ctx.tasks_) {
-            if (task.state_ == TransferState::kDone) {
-                continue;
+
+        // In-order ack polling: we only look at the head of the queue
+        // because the sender acknowledges chunks in the same order it
+        // consumes credits.
+        if (!op_ctx.tasks_.empty()) {
+            auto& head = op_ctx.tasks_.front();
+            if (head.state_ == RecvTaskState::kWaitAck) {
+                if (pollRecvAckSlot(op_ctx, lane, head)) {
+                    did_work = true;
+                }
             }
-            if (StepRecvTransferTask(task)) {
+        }
+
+        // Advance every chunk through the receiver state machine.
+        bool op_failed = false;
+        for (auto it = op_ctx.tasks_.begin(); it != op_ctx.tasks_.end();) {
+            did_work |= stepRecvTask(*it);
+            if (it->state_ == RecvTaskState::kFinished) {
+                it = op_ctx.tasks_.erase(it);
                 did_work = true;
+            } else if (it->state_ == RecvTaskState::kFailed) {
+                op_failed = true;
+                break;
+            } else {
+                ++it;
             }
         }
-        if (StepRecvTailCommit(op_ctx, capacity)) {
+
+        if (op_failed) {
+            reportBrokenPeer(peer_rank);
             did_work = true;
+            continue;
         }
-        if (IsRecvDataPathCompleted(op_ctx)) {
+
+        if (isRecvOpCompleted(op_ctx)) {
             if (!op_ctx.original_tensor_.is_contiguous()) {
-                op_ctx.original_tensor_.copy_(op_ctx.tensor_);
+                (void)op_ctx.original_tensor_.copy_(op_ctx.tensor_);
                 if (!is_cpu_) {
                     const cudaError_t sync_error = cudaDeviceSynchronize();
                     TORCH_CHECK(sync_error == cudaSuccess,
@@ -885,7 +1252,8 @@ bool P2PProxy::StepRecv() {
                                 cudaGetErrorString(sync_error));
                 }
             }
-            op_ctx.completed_->store(true, std::memory_order_release);
+            op_ctx.status_->store(OpStatus::kSuccess,
+                                  std::memory_order_release);
             lane.active_recv_op_.reset();
             active_recv_tasks_.fetch_sub(1, std::memory_order_release);
             did_work = true;
@@ -895,46 +1263,50 @@ bool P2PProxy::StepRecv() {
     return did_work;
 }
 
-void P2PProxy::SetDeviceWorker(P2PDeviceWorker* worker) {
+void P2PProxy::attachToWorker(P2PDeviceWorker* worker, P2PChunkPool* send,
+                              P2PChunkPool* recv, size_t chunk_size) {
     device_worker_ = worker;
+    send_pool_ = send;
+    recv_pool_ = recv;
+    chunk_size_ = chunk_size;
 }
 
-bool P2PProxy::HasActiveSendWork() const {
+bool P2PProxy::hasActiveSendWork() const {
     for (int i = 0; i < size_; ++i) {
         if (reset_send_req_[i].load(std::memory_order_acquire)) return true;
     }
     return active_send_tasks_.load(std::memory_order_acquire) > 0;
 }
 
-bool P2PProxy::HasActiveRecvWork() const {
+bool P2PProxy::hasActiveRecvWork() const {
     for (int i = 0; i < size_; ++i) {
         if (reset_recv_req_[i].load(std::memory_order_acquire)) return true;
     }
     return active_recv_tasks_.load(std::memory_order_acquire) > 0;
 }
 
-bool P2PProxy::DrainTasks() const {
+bool P2PProxy::drainTasks() const {
     BackoffWaiter waiter;
     return waiter.wait_for(
         std::chrono::milliseconds(kDrainTasksTimeoutMs),
-        [this] { return !HasActiveSendWork() && !HasActiveRecvWork(); });
+        [this] { return !hasActiveSendWork() && !hasActiveRecvWork(); });
 }
 
-void P2PDeviceWorker::Start() {
+void P2PDeviceWorker::start() {
     bool expected_send = false;
     if (send_worker_running_.compare_exchange_strong(expected_send, true)) {
         send_worker_thread_ =
-            std::thread(&P2PDeviceWorker::SendWorkerThread, this);
+            std::thread(&P2PDeviceWorker::sendWorkerMainloop, this);
     }
 
     bool expected_recv = false;
     if (recv_worker_running_.compare_exchange_strong(expected_recv, true)) {
         recv_worker_thread_ =
-            std::thread(&P2PDeviceWorker::RecvWorkerThread, this);
+            std::thread(&P2PDeviceWorker::recvWorkerMainloop, this);
     }
 }
 
-void P2PDeviceWorker::Stop() {
+void P2PDeviceWorker::stop() {
     bool expected_send = true;
     std::unique_lock<std::mutex> s_lock(send_wakeup_mutex_);
     if (send_worker_running_.compare_exchange_strong(expected_send, false)) {
@@ -956,8 +1328,21 @@ void P2PDeviceWorker::Stop() {
     }
 }
 
+P2PDeviceWorker::P2PDeviceWorker(TransferEngine* engine,
+                                 const std::string& location, bool is_cpu,
+                                 int cuda_device_index)
+    : is_cpu_(is_cpu), cuda_device_index_(cuda_device_index) {
+    initPools(engine, location);
+    start();
+}
+
+P2PDeviceWorker::~P2PDeviceWorker() {
+    stop();
+    releasePools();
+}
+
 void P2PDeviceWorker::registerProxy(const std::shared_ptr<P2PProxy>& proxy) {
-    proxy->SetDeviceWorker(this);
+    proxy->attachToWorker(this, &send_pool_, &recv_pool_, chunk_size_);
     {
         std::lock_guard<std::mutex> lock(proxies_mutex_);
         proxies_.emplace_back(proxy);
@@ -981,13 +1366,96 @@ void P2PDeviceWorker::removeProxy(const std::shared_ptr<P2PProxy>& proxy) {
         std::lock_guard<std::mutex> r_lock(recv_wakeup_mutex_);
         proxies_version_.fetch_add(1, std::memory_order_release);
     }
-    proxy->SetDeviceWorker(nullptr);
     send_wakeup_cv_.notify_one();
     recv_wakeup_cv_.notify_one();
 }
 
+void P2PDeviceWorker::initPools(TransferEngine* engine,
+                                const std::string& location) {
+    engine_ = engine;
+
+    // Parse from environment variables
+    pool_bytes_ = getEnv_size_t("MOONCAKE_P2P_POOL_SIZE", kDefaultPoolSize);
+    chunk_size_ = getEnv_size_t("MOONCAKE_P2P_CHUNK_SIZE", kDefaultChunkSize);
+
+    // chunk_len is uint32_t fields in control slots
+    TORCH_CHECK(
+        chunk_size_ > 0 && chunk_size_ <= std::numeric_limits<uint32_t>::max(),
+        "Invalid MOONCAKE_P2P_CHUNK_SIZE: must be > 0 and <= 4GB");
+
+    TORCH_CHECK(
+        pool_bytes_ > 0 && pool_bytes_ % chunk_size_ == 0,
+        "Invalid pool size and chunk size (must hold 'pool_bytes_ > 0 && "
+        "pool_bytes_ % chunk_size_ == 0')");
+
+    num_chunks_ = static_cast<uint32_t>(pool_bytes_ / chunk_size_);
+    TORCH_CHECK(num_chunks_ > 0, "P2PDeviceWorker: num_chunks_ must be > 0");
+
+    if (is_cpu_) {
+        send_pool_base_ = std::malloc(pool_bytes_);
+        TORCH_CHECK(send_pool_base_ != nullptr,
+                    "Failed to allocate CPU P2P send pool");
+        int rc =
+            engine->registerLocalMemory(send_pool_base_, pool_bytes_, location);
+        TORCH_CHECK(rc == 0, "Failed to register CPU P2P send pool");
+
+        recv_pool_base_ = std::malloc(pool_bytes_);
+        TORCH_CHECK(recv_pool_base_ != nullptr,
+                    "Failed to allocate CPU P2P recv pool");
+        rc =
+            engine->registerLocalMemory(recv_pool_base_, pool_bytes_, location);
+        TORCH_CHECK(rc == 0, "Failed to register CPU P2P recv pool");
+    } else {
+        setCudaDeviceIfNeeded(is_cpu_, cuda_device_index_,
+                              "P2PDeviceWorker initPools cudaSetDevice failed");
+        cudaError_t err = cudaMalloc(&send_pool_base_, pool_bytes_);
+        TORCH_CHECK(
+            err == cudaSuccess,
+            "Failed to allocate CUDA P2P send pool: ", cudaGetErrorString(err));
+        int rc =
+            engine->registerLocalMemory(send_pool_base_, pool_bytes_, location);
+        TORCH_CHECK(rc == 0, "Failed to register CUDA P2P send pool");
+
+        err = cudaMalloc(&recv_pool_base_, pool_bytes_);
+        TORCH_CHECK(
+            err == cudaSuccess,
+            "Failed to allocate CUDA P2P recv pool: ", cudaGetErrorString(err));
+        rc =
+            engine->registerLocalMemory(recv_pool_base_, pool_bytes_, location);
+        TORCH_CHECK(rc == 0, "Failed to register CUDA P2P recv pool");
+    }
+
+    send_pool_.init(send_pool_base_, chunk_size_, num_chunks_);
+    recv_pool_.init(recv_pool_base_, chunk_size_, num_chunks_);
+}
+
+void P2PDeviceWorker::releasePools() {
+    setCudaDeviceIfNeeded(is_cpu_, cuda_device_index_,
+                          "P2PDeviceWorker releasePools cudaSetDevice failed");
+
+    if (send_pool_base_ != nullptr) {
+        if (engine_) engine_->unregisterLocalMemory(send_pool_base_);
+        if (is_cpu_) {
+            std::free(send_pool_base_);
+        } else {
+            cudaFree(send_pool_base_);
+        }
+        send_pool_base_ = nullptr;
+    }
+
+    if (recv_pool_base_ != nullptr) {
+        if (engine_) engine_->unregisterLocalMemory(recv_pool_base_);
+        if (is_cpu_) {
+            std::free(recv_pool_base_);
+        } else {
+            cudaFree(recv_pool_base_);
+        }
+        recv_pool_base_ = nullptr;
+    }
+}
+
 template <typename HasWorkFn, typename StepWorkFn>
-void WorkerThreadLoop(bool is_cpu, int cuda_device_index,
+void workerThreadLoop(bool is_cpu, int cuda_device_index,
                       std::atomic<bool>& worker_running,
                       std::atomic<uint64_t>& proxies_version,
                       std::mutex& proxies_mutex,
@@ -995,7 +1463,7 @@ void WorkerThreadLoop(bool is_cpu, int cuda_device_index,
                       std::mutex& wakeup_mutex,
                       std::condition_variable& wakeup_cv, HasWorkFn has_work,
                       StepWorkFn step_work) {
-    SetCudaDeviceIfNeeded(is_cpu, cuda_device_index,
+    setCudaDeviceIfNeeded(is_cpu, cuda_device_index,
                           "P2PDeviceWorker::WorkerThread cudaSetDevice failed");
 
     // A thread-local cache for proxies to avoid locking too frequently.
@@ -1091,20 +1559,20 @@ void WorkerThreadLoop(bool is_cpu, int cuda_device_index,
     }
 }
 
-void P2PDeviceWorker::SendWorkerThread() {
-    WorkerThreadLoop(
+void P2PDeviceWorker::sendWorkerMainloop() {
+    workerThreadLoop(
         is_cpu_, cuda_device_index_, send_worker_running_, proxies_version_,
         proxies_mutex_, proxies_, send_wakeup_mutex_, send_wakeup_cv_,
-        [](P2PProxy& p) { return p.HasActiveSendWork(); },
-        [](P2PProxy& p) { return p.StepSend(); });
+        [](P2PProxy& p) { return p.hasActiveSendWork(); },
+        [](P2PProxy& p) { return p.stepSend(); });
 }
 
-void P2PDeviceWorker::RecvWorkerThread() {
-    WorkerThreadLoop(
+void P2PDeviceWorker::recvWorkerMainloop() {
+    workerThreadLoop(
         is_cpu_, cuda_device_index_, recv_worker_running_, proxies_version_,
         proxies_mutex_, proxies_, recv_wakeup_mutex_, recv_wakeup_cv_,
-        [](P2PProxy& p) { return p.HasActiveRecvWork(); },
-        [](P2PProxy& p) { return p.StepRecv(); });
+        [](P2PProxy& p) { return p.hasActiveRecvWork(); },
+        [](P2PProxy& p) { return p.stepRecv(); });
 }
 
 // Standard practice is to update states shared with condition_variables under
@@ -1114,16 +1582,17 @@ void P2PDeviceWorker::RecvWorkerThread() {
 //  acquired in WakeUpSend/Recv)
 //
 // Ref: https://stackoverflow.com/a/21439617
-void P2PDeviceWorker::WakeUpSend() {
+void P2PDeviceWorker::wakeUpSend() {
     std::lock_guard<std::mutex> lock(send_wakeup_mutex_);
     send_wakeup_cv_.notify_one();
 }
-void P2PDeviceWorker::WakeUpRecv() {
+void P2PDeviceWorker::wakeUpRecv() {
     std::lock_guard<std::mutex> lock(recv_wakeup_mutex_);
     recv_wakeup_cv_.notify_one();
 }
 
-std::shared_ptr<P2PDeviceWorker> P2PDeviceWorkerManager::GetCPUWorker() {
+std::shared_ptr<P2PDeviceWorker> P2PDeviceWorkerManager::getCPUWorker(
+    TransferEngine* engine) {
     std::lock_guard<std::mutex> lock(manager_mutex_);
 
     auto it = workers_.find(CPUWorkerID);
@@ -1131,14 +1600,14 @@ std::shared_ptr<P2PDeviceWorker> P2PDeviceWorkerManager::GetCPUWorker() {
         if (auto ptr = it->second.lock()) return ptr;
     }
 
-    auto worker =
-        std::make_shared<P2PDeviceWorker>(/* is_cpu */ true, CPUWorkerID);
+    auto worker = std::make_shared<P2PDeviceWorker>(
+        engine, kWildcardLocation, /* is_cpu */ true, CPUWorkerID);
     workers_[CPUWorkerID] = worker;
     return worker;
 }
 
-std::shared_ptr<P2PDeviceWorker> P2PDeviceWorkerManager::GetCUDAWorker(
-    int cuda_device_index) {
+std::shared_ptr<P2PDeviceWorker> P2PDeviceWorkerManager::getCUDAWorker(
+    int cuda_device_index, TransferEngine* engine) {
     std::lock_guard<std::mutex> lock(manager_mutex_);
 
     auto it = workers_.find(cuda_device_index);
@@ -1146,9 +1615,11 @@ std::shared_ptr<P2PDeviceWorker> P2PDeviceWorkerManager::GetCUDAWorker(
         if (auto ptr = it->second.lock()) return ptr;
     }
 
-    auto worker = std::make_shared<P2PDeviceWorker>(/* is_cpu */ false,
-                                                    cuda_device_index);
+    auto worker = std::make_shared<P2PDeviceWorker>(
+        engine, GPU_PREFIX + std::to_string(cuda_device_index),
+        /* is_cpu */ false, cuda_device_index);
     workers_[cuda_device_index] = worker;
     return worker;
 }
+
 }  // namespace mooncake
diff --git a/mooncake-pg/src/pg_py.cpp b/mooncake-pg/src/pg_py.cpp
index 78e3269a..e39daffe 100644
--- a/mooncake-pg/src/pg_py.cpp
+++ b/mooncake-pg/src/pg_py.cpp
@@ -9,7 +9,7 @@ namespace py = pybind11;
 
 namespace mooncake {
 
-c10::intrusive_ptr<c10d::Backend> createMooncakeBackend(
+c10::intrusive_ptr<c10d::ProcessGroup> createMooncakeBackend(
     c10d::DistributedBackendOptions distBackendOpts,
     c10::intrusive_ptr<MooncakeBackend::MooncakeBackendOptions>
         backendOptions) {
@@ -17,7 +17,7 @@ c10::intrusive_ptr<c10d::Backend> createMooncakeBackend(
                                                 std::move(backendOptions));
 }
 
-c10::intrusive_ptr<c10d::Backend> createMooncakeCpuBackend(
+c10::intrusive_ptr<c10d::ProcessGroup> createMooncakeCpuBackend(
     c10d::DistributedBackendOptions distBackendOpts,
     c10::intrusive_ptr<MooncakeBackend::MooncakeBackendOptions>
         backendOptions) {
@@ -33,62 +33,91 @@ __attribute__((constructor)) static void MooncakeBackendConstructor() {
     kwargsCpu["devices"] = py::make_tuple("cpu");
     register_backend("mooncake-cpu", py::cpp_function(createMooncakeCpuBackend),
                      /* extended_api */ true, **kwargsCpu);
+#ifndef MOONCAKE_EP_USE_MUSA
     py::dict kwargsCuda;
     kwargsCuda["devices"] = py::make_tuple("cuda");
     register_backend("mooncake", py::cpp_function(createMooncakeBackend),
                      /* extended_api */ true, **kwargsCuda);
+#else
+    py::dict kwargsMusa;
+    kwargsMusa["devices"] = py::make_tuple("musa");
+    register_backend("mooncake", py::cpp_function(createMooncakeBackend),
+                     /* extended_api */ true, **kwargsMusa);
+#endif
 }
 
-std::string getPreferredHca(c10::intrusive_ptr<c10d::Backend> backend,
+std::string getPreferredHca(c10::intrusive_ptr<c10d::ProcessGroup> backend,
                             std::string location) {
     auto mooncakeBackend =
         c10::static_intrusive_pointer_cast<MooncakeBackend>(backend);
     return mooncakeBackend->getPreferredHca(location);
 }
 
-at::Tensor getActiveRanks(c10::intrusive_ptr<c10d::Backend> backend) {
+at::Tensor getActiveRanks(c10::intrusive_ptr<c10d::ProcessGroup> backend) {
     auto mooncakeBackend =
         c10::static_intrusive_pointer_cast<MooncakeBackend>(backend);
     return mooncakeBackend->getActiveRanksTensor();
 }
 
-int getNumSyncedRanks(c10::intrusive_ptr<c10d::Backend> backend) {
+int getNumSyncedRanks(c10::intrusive_ptr<c10d::ProcessGroup> backend) {
     auto mooncakeBackend =
         c10::static_intrusive_pointer_cast<MooncakeBackend>(backend);
     return mooncakeBackend->getNumSyncedRanks();
 }
 
-void extendGroupSizeTo(c10::intrusive_ptr<c10d::Backend> backend, int size) {
+void extendGroupSizeTo(c10::intrusive_ptr<c10d::ProcessGroup> backend,
+                       int size) {
     auto mooncakeBackend =
         c10::static_intrusive_pointer_cast<MooncakeBackend>(backend);
     mooncakeBackend->extendGroupSizeTo(size);
 }
 
-std::vector<bool> getPeerState(c10::intrusive_ptr<c10d::Backend> backend,
+std::vector<bool> getPeerState(c10::intrusive_ptr<c10d::ProcessGroup> backend,
                                const std::vector<int>& ranks) {
     auto mooncakeBackend =
         c10::static_intrusive_pointer_cast<MooncakeBackend>(backend);
     return mooncakeBackend->getPeerState(ranks);
 }
 
-void recoverRanks(c10::intrusive_ptr<c10d::Backend> backend,
+void recoverRanks(c10::intrusive_ptr<c10d::ProcessGroup> backend,
                   const std::vector<int>& ranks) {
     auto mooncakeBackend =
         c10::static_intrusive_pointer_cast<MooncakeBackend>(backend);
     mooncakeBackend->recoverRanks(ranks);
 }
 
-void joinGroup(c10::intrusive_ptr<c10d::Backend> backend) {
+void joinGroup(c10::intrusive_ptr<c10d::ProcessGroup> backend) {
     auto mooncakeBackend =
         c10::static_intrusive_pointer_cast<MooncakeBackend>(backend);
     mooncakeBackend->joinGroup();
 }
 
+/// Python-facing wrapper that extracts the raw TransferEngine* from a
+/// mooncake.engine.TransferEngine Python object and passes it to
+/// MooncakeBackend::setExternalEngine().  The caller must ensure the
+/// TransferEnginePy object outlives all MooncakeBackend instances.
+void setTransferEnginePy(pybind11::object engine_obj) {
+    if (engine_obj.is_none()) {
+        MooncakeBackend::setExternalEngine(nullptr);
+        return;
+    }
+    auto get_engine_ptr = engine_obj.attr("get_engine_ptr");
+    uintptr_t ptr = get_engine_ptr().cast<uintptr_t>();
+    auto* engine = reinterpret_cast<TransferEngine*>(ptr);
+    MooncakeBackend::setExternalEngine(engine);
+}
+
 PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
     m.def("createMooncakeBackend", &createMooncakeBackend);
     m.def("createMooncakeCpuBackend", &createMooncakeCpuBackend);
     m.def("set_host_ip", &MooncakeBackend::setHostIp);
     m.def("set_device_filter", &MooncakeBackend::setDeviceFilter);
+    m.def("set_transfer_engine", &setTransferEnginePy, py::arg("engine"),
+          "Set an external TransferEngine to be used by MooncakeBackend. "
+          "Must be called before init_process_group(). The engine must already "
+          "be initialized. Pass None to reset to default behavior. "
+          "The caller must ensure the TransferEngine object outlives all "
+          "MooncakeBackend instances.");
     m.def("get_preferred_hca", &getPreferredHca);
     m.def("get_active_ranks", &getActiveRanks);
     m.def("get_num_synced_ranks", &getNumSyncedRanks);
@@ -102,7 +131,9 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
         m, "MooncakeBackendOptions")
         .def(py::init<at::Tensor>(), py::arg("active_ranks"))
         .def(py::init<at::Tensor, bool>(), py::arg("active_ranks"),
-             py::arg("is_extension"));
+             py::arg("is_extension"))
+        .def(py::init<at::Tensor, bool, int>(), py::arg("active_ranks"),
+             py::arg("is_extension"), py::arg("max_world_size"));
 }
 
 }  // namespace mooncake
diff --git a/mooncake-pg/tests/pg_test_utils.py b/mooncake-pg/tests/pg_test_utils.py
index cd700451..87950c80 100644
--- a/mooncake-pg/tests/pg_test_utils.py
+++ b/mooncake-pg/tests/pg_test_utils.py
@@ -12,7 +12,6 @@ import torch.distributed as dist
 import torch.multiprocessing as mp
 from mooncake import pg
 
-
 DEVICE_FILTER_ENV_VAR = "MOONCAKE_PGTEST_DEVICE_FILTERS"
 MASTER_ADDR_ENV_VAR = "MOONCAKE_PGTEST_MASTER_ADDR"
 MASTER_PORT_ENV_VAR = "MOONCAKE_PGTEST_MASTER_PORT"
@@ -74,6 +73,12 @@ def temporary_env(updates: dict[str, str]):
                 os.environ[key] = value
 
 
+def musa_runtime_available(min_devices: int = 1) -> bool:
+    if not hasattr(torch, "musa") or not torch.musa.is_available():
+        return False
+    return torch.musa.device_count() >= min_devices
+
+
 def cuda_runtime_available(min_devices: int = 1) -> bool:
     if not torch.cuda.is_available():
         return False
@@ -81,6 +86,18 @@ def cuda_runtime_available(min_devices: int = 1) -> bool:
 
 
 def require_test_device(rank: int, device_type: str) -> torch.device:
+    if device_type == "musa":
+        device_count = torch.musa.device_count()
+        if device_count <= 0:
+            raise RuntimeError(
+                "MUSA backend requested but no MUSA devices are available"
+            )
+        if rank >= device_count:
+            raise RuntimeError(
+                f"rank {rank} requires a dedicated MUSA device but only {device_count} are visible"
+            )
+        torch.musa.set_device(rank)
+        return torch.device("musa", rank)
     if device_type == "cuda":
         device_count = torch.cuda.device_count()
         if device_count <= 0:
@@ -104,17 +121,21 @@ def mooncake_backend_options(
     *,
     active_value: int = 0,
     is_extension: bool = False,
+    max_world_size: int | None = None,
 ) -> pg.MooncakeBackendOptions:
     device = torch.device(device_type)
+    tensor_size = world_size if max_world_size is None else int(max_world_size)
     active_ranks = torch.full(
-        (world_size,),
+        (tensor_size,),
         int(active_value),
         dtype=torch.int32,
         device=device,
     )
-    if is_extension:
-        return pg.MooncakeBackendOptions(active_ranks, True)
-    return pg.MooncakeBackendOptions(active_ranks)
+    if max_world_size is None:
+        if is_extension:
+            return pg.MooncakeBackendOptions(active_ranks, True)
+        return pg.MooncakeBackendOptions(active_ranks)
+    return pg.MooncakeBackendOptions(active_ranks, bool(is_extension), tensor_size)
 
 
 def mooncake_cpu_options(world_size: int) -> pg.MooncakeBackendOptions:
@@ -140,6 +161,7 @@ def init_mooncake_group(
     use_pg_options: bool = True,
     is_extension: bool = False,
     active_value: int | None = None,
+    max_world_size: int | None = None,
 ) -> torch.device:
     device = require_test_device(rank, device_type)
     configure_mooncake_device_filter(device_filters)
@@ -148,8 +170,6 @@ def init_mooncake_group(
         "rank": rank,
         "world_size": world_size,
     }
-    if device_type == "cuda":
-        kwargs["device_id"] = device
     if use_pg_options:
         resolved_active_value = (
             1 if is_extension else 0 if active_value is None else active_value
@@ -159,6 +179,7 @@ def init_mooncake_group(
             device_type,
             active_value=resolved_active_value,
             is_extension=is_extension,
+            max_world_size=max_world_size,
         )
     dist.init_process_group(**kwargs)
     return device
@@ -184,7 +205,7 @@ def init_mooncake_cpu_group(
 def get_mooncake_backend(group=None, device_type: str = "cpu"):
     if group is None:
         group = dist.group.WORLD
-    return group._get_backend(torch.device(device_type))
+    return group
 
 
 @dataclass(slots=True)
@@ -216,6 +237,7 @@ class MooncakePGWorkerContext:
         use_pg_options: bool = True,
         is_extension: bool = False,
         active_value: int | None = None,
+        max_world_size: int | None = None,
     ) -> torch.device:
         self._device = init_mooncake_group(
             self.proc_rank if rank is None else rank,
@@ -228,6 +250,7 @@ class MooncakePGWorkerContext:
             use_pg_options=use_pg_options,
             is_extension=is_extension,
             active_value=active_value,
+            max_world_size=max_world_size,
         )
         return self._device
 
@@ -503,9 +526,16 @@ class BackendMultiProcessTestCase(MultiProcessTestCase):
             raise RuntimeError(
                 f"{cls.__name__} must inherit a concrete Mooncake PG backend test base class"
             )
+        if cls.device_type == "musa":
+            device_count = torch.musa.device_count() if hasattr(torch, "musa") and torch.musa.is_available() else 0
+            cls.configure_for_cuda_device_count(device_count)
         if cls.device_type == "cuda":
             device_count = torch.cuda.device_count() if torch.cuda.is_available() else 0
             cls.configure_for_cuda_device_count(device_count)
+        if cls.device_type == "musa" and not musa_runtime_available(cls.world_size):
+            raise unittest.SkipTest(
+                f"{cls.__name__} requires {cls.world_size} visible MUSA devices"
+            )
         if cls.device_type == "cuda" and not cuda_runtime_available(cls.world_size):
             raise unittest.SkipTest(
                 f"{cls.__name__} requires {cls.world_size} visible CUDA devices"
@@ -577,3 +607,8 @@ class MooncakePGCPUBackendTestCase(BackendMultiProcessTestCase):
 class MooncakePGCUDABackendTestCase(BackendMultiProcessTestCase):
     backend_name = "mooncake"
     device_type = "cuda"
+
+
+class MooncakePGMUSABackendTestCase(BackendMultiProcessTestCase):
+    backend_name = "mooncake"
+    device_type = "musa"
diff --git a/mooncake-pg/tests/test_pg_collectives.py b/mooncake-pg/tests/test_pg_collectives.py
index 0d07af00..52fbee7a 100644
--- a/mooncake-pg/tests/test_pg_collectives.py
+++ b/mooncake-pg/tests/test_pg_collectives.py
@@ -6,6 +6,7 @@ import torch.distributed as dist
 from pg_test_utils import (
     MooncakePGCPUBackendTestCase,
     MooncakePGCUDABackendTestCase,
+    MooncakePGMUSABackendTestCase,
     MooncakePGWorkerContext,
     wait_until,
 )
@@ -226,5 +227,9 @@ class TestMooncakePGCollectivesCUDA(_CollectiveTestMixin, MooncakePGCUDABackendT
     world_size = 2
 
 
+class TestMooncakePGCollectivesMUSA(_CollectiveTestMixin, MooncakePGMUSABackendTestCase):
+    world_size = 2
+
+
 if __name__ == "__main__":
     unittest.main()
diff --git a/mooncake-pg/tests/test_pg_elastic.py b/mooncake-pg/tests/test_pg_elastic.py
index 4f40d96f..e44f7ad0 100644
--- a/mooncake-pg/tests/test_pg_elastic.py
+++ b/mooncake-pg/tests/test_pg_elastic.py
@@ -11,6 +11,9 @@ from pg_test_utils import (
     MooncakePGCPUBackendTestCase,
     MooncakePGCUDABackendTestCase,
     MooncakePGWorkerContext,
+    configure_mooncake_device_filter,
+    get_mooncake_backend,
+    require_test_device,
     wait_until,
 )
 
@@ -18,6 +21,29 @@ from pg_test_utils import (
 BROKEN_RANK = 1
 
 
+def _dynamic_world_size_worker(
+    ctx: MooncakePGWorkerContext,
+) -> None:
+    """Worker for testing that dist.get_world_size() reflects dynamic size after extend."""
+    initial_world_size = ctx.world_size
+    ctx.init_group()
+    backend = ctx.get_backend()
+
+    initial_ws = dist.get_world_size()
+    assert initial_ws == initial_world_size, (
+        f"rank {ctx.rank}: initial world_size={initial_ws}, expected {initial_world_size}"
+    )
+
+    pg.extend_group_size_to(backend, initial_world_size + 1)
+
+    new_ws = dist.get_world_size()
+    assert new_ws == initial_world_size + 1, (
+        f"rank {ctx.rank}: after extend world_size={new_ws}, expected {initial_world_size + 1}"
+    )
+
+    ctx.record_result({"initial_ws": initial_ws, "new_ws": new_ws})
+
+
 def _extension_worker(
     ctx: MooncakePGWorkerContext,
     extend_event: mp.Event,
@@ -27,11 +53,24 @@ def _extension_worker(
     initial_world_size = ctx.world_size - 1
     extension_rank = ctx.world_size - 1
 
+    join_ranks = [extension_rank]
+
     if ctx.proc_rank < initial_world_size:
         # Original ranks
-        device = ctx.init_group(world_size=initial_world_size)
+        device = ctx.init_group(
+            world_size=initial_world_size,
+            max_world_size=ctx.world_size,
+        )
         backend = ctx.get_backend()
 
+        # group_size should equal initial_world_size immediately after init
+        # (max_world_size only pre-allocates capacity, does not change visible size)
+        actual_ws = dist.get_world_size()
+        assert actual_ws == initial_world_size, (
+            f"rank {ctx.proc_rank}: initial world_size={actual_ws}, "
+            f"expected initial_world_size={initial_world_size}"
+        )
+
         # First collective
         tensor = torch.tensor([ctx.proc_rank + 1], dtype=torch.int32, device=device)
         dist.all_reduce(tensor, op=dist.ReduceOp.SUM)
@@ -40,11 +79,25 @@ def _extension_worker(
         # Signal ready and extend
         if ctx.proc_rank == 0:
             extend_event.set()
-        pg.extend_group_size_to(backend, ctx.world_size)
 
-        # Wait for extension rank to complete init before collective
-        if not init_done_event.wait(timeout=30.0):
-            raise TimeoutError("timed out waiting for extension init")
+        # Two-phase extension protocol:
+        #   1) joiner publishes metadata + establishes transport readiness
+        #   2) healthy ranks recover/activate it via recover_ranks()
+        # Note: get_peer_state() is collective among *healthy ranks*.
+        wait_until(
+            lambda: all(pg.get_peer_state(backend, join_ranks)),
+            timeout_s=30.0,
+            poll_interval_s=0.05,
+            description=f"rank {ctx.proc_rank} waiting for joiner ready",
+        )
+        pg.recover_ranks(backend, join_ranks)
+
+        # After recover_ranks, world_size should now reflect the expanded group
+        actual_ws_after = dist.get_world_size()
+        assert actual_ws_after == ctx.world_size, (
+            f"rank {ctx.proc_rank}: world_size after recover={actual_ws_after}, "
+            f"expected max_world_size={ctx.world_size}"
+        )
 
         # Final collective
         final_tensor = torch.tensor([ctx.proc_rank + 1], dtype=torch.int32, device=device)
@@ -63,10 +116,41 @@ def _extension_worker(
         device = ctx.init_group(
             rank=extension_rank,
             world_size=ctx.world_size,
+            is_extension=True,
+            max_world_size=ctx.world_size,
         )
 
-        # Signal init complete before collective
-        init_done_event.set()
+        backend = ctx.get_backend()
+
+        # group_size for extension rank equals world_size passed at init.
+        # Note: this is world_size (= max_world_size for joiners), not 1,
+        # because the joiner's activeSize is initialized to world_size.
+        # The local-only behavior is ensured by activeRanks masking, not
+        # by a smaller activeSize.
+        actual_ws = dist.get_world_size()
+        assert actual_ws == ctx.world_size, (
+            f"extension rank: initial world_size={actual_ws}, "
+            f"expected {ctx.world_size}"
+        )
+
+        # In extension mode, joiner starts in local-only collectives.
+        local_tensor = torch.tensor([extension_rank + 1], dtype=torch.int32, device=device)
+        dist.all_reduce(local_tensor, op=dist.ReduceOp.SUM)
+        if int(local_tensor.cpu().item()) != extension_rank + 1:
+            raise AssertionError(
+                f"extension rank expected local-only sum {extension_rank + 1}, got {int(local_tensor.cpu().item())}"
+            )
+
+        # join_group publishes metadata and then blocks until recover_ranks()
+        # publishes the extension state.
+        pg.join_group(backend)
+
+        # After joinGroup, world_size should reflect the full group
+        actual_ws_after = dist.get_world_size()
+        assert actual_ws_after == ctx.world_size, (
+            f"extension rank: world_size after joinGroup={actual_ws_after}, "
+            f"expected {ctx.world_size}"
+        )
 
         # Final collective
         final_tensor = torch.tensor([extension_rank + 1], dtype=torch.int32, device=device)
@@ -78,6 +162,533 @@ def _extension_worker(
         })
 
 
+def _extension_p2p_worker(
+    ctx: MooncakePGWorkerContext,
+    extend_event: mp.Event,
+    direction: str,
+) -> None:
+    """Worker for testing P2P after max_world_size/recover_ranks scale-up."""
+    initial_world_size = ctx.world_size - 1
+    extension_rank = ctx.world_size - 1
+    primary_peer = initial_world_size - 1
+    join_ranks = [extension_rank]
+
+    if initial_world_size < 1:
+        raise AssertionError("elastic P2P test expects at least one primary rank")
+
+    if ctx.proc_rank < initial_world_size:
+        device = ctx.init_group(
+            world_size=initial_world_size,
+            max_world_size=ctx.world_size,
+        )
+        backend = ctx.get_backend()
+
+        if ctx.proc_rank == 0:
+            extend_event.set()
+
+        wait_until(
+            lambda: all(pg.get_peer_state(backend, join_ranks)),
+            timeout_s=30.0,
+            poll_interval_s=0.05,
+            description=f"rank {ctx.proc_rank} waiting for joiner ready",
+        )
+        pg.recover_ranks(backend, join_ranks)
+    else:
+        if not extend_event.wait(timeout=30.0):
+            raise TimeoutError("timed out waiting for extend_event")
+
+        device = ctx.init_group(
+            rank=extension_rank,
+            world_size=ctx.world_size,
+            is_extension=True,
+            max_world_size=ctx.world_size,
+        )
+        backend = ctx.get_backend()
+        pg.join_group(backend)
+
+    actual_ws_after = dist.get_world_size()
+    assert actual_ws_after == ctx.world_size, (
+        f"rank {ctx.proc_rank}: world_size after recovery={actual_ws_after}, "
+        f"expected {ctx.world_size}"
+    )
+
+    # Prove elastic collectives are functional before isolating P2P.
+    collective = torch.tensor([ctx.proc_rank + 1], dtype=torch.int32, device=device)
+    dist.all_reduce(collective, op=dist.ReduceOp.SUM)
+    expected_sum = ctx.world_size * (ctx.world_size + 1) // 2
+    if int(collective.cpu().item()) != expected_sum:
+        raise AssertionError(
+            f"rank {ctx.proc_rank}: post-recovery all_reduce expected "
+            f"{expected_sum}, got {int(collective.cpu().item())}"
+        )
+
+    dist.barrier()
+
+    if direction == "joiner_to_primary":
+        src_rank = extension_rank
+        dst_rank = primary_peer
+    elif direction == "primary_to_joiner":
+        src_rank = primary_peer
+        dst_rank = extension_rank
+    else:
+        raise AssertionError(f"unknown P2P direction: {direction}")
+
+    numel = 1024
+    if ctx.proc_rank == src_rank:
+        send_tensor = torch.full(
+            (numel,), src_rank, dtype=torch.int32, device=device
+        )
+        works = dist.batch_isend_irecv(
+            [dist.P2POp(op=dist.isend, tensor=send_tensor, peer=dst_rank)]
+        )
+        for work in works:
+            work.wait()
+        ctx.synchronize()
+        value = "sent"
+    elif ctx.proc_rank == dst_rank:
+        recv_tensor = torch.empty((numel,), dtype=torch.int32, device=device)
+        works = dist.batch_isend_irecv(
+            [dist.P2POp(op=dist.irecv, tensor=recv_tensor, peer=src_rank)]
+        )
+        for work in works:
+            work.wait()
+        ctx.synchronize()
+        expected = torch.full_like(recv_tensor, src_rank)
+        if not torch.equal(recv_tensor.cpu(), expected.cpu()):
+            raise AssertionError(
+                f"rank {ctx.proc_rank}: received unexpected P2P payload "
+                f"from rank {src_rank}"
+            )
+        value = "received"
+    else:
+        value = "idle"
+
+    dist.barrier()
+    ctx.record_result(
+        {
+            "direction": direction,
+            "role": (
+                "src"
+                if ctx.proc_rank == src_rank
+                else "dst"
+                if ctx.proc_rank == dst_rank
+                else "idle"
+            ),
+            "value": value,
+        }
+    )
+
+
+def _extension_worker_with_subgroups(
+    ctx: MooncakePGWorkerContext,
+    extend_event: mp.Event,
+) -> None:
+    """Multi-subgroup elastic extension test using split-ranks pattern.
+
+    Layout (world_size=4, primary=[0,1], joiners=[2,3]):
+      group_a: primary ranks=[0],   joiner ranks=[0,2],     max_world_size=2
+      group_b: primary ranks=[1],   joiner ranks=[1,3],     max_world_size=2
+      group_c: primary ranks=[0,1], joiner ranks=[0,1,2,3], max_world_size=4
+
+    Primary ranks must initialize WORLD with world_size=initial_world_size and
+    create subgroups using only their current membership; joiners wait for the
+    extend signal, then init WORLD with the full world_size and create subgroups
+    using the full eventual membership. PyTorch's new_group uses a monotonic
+    call counter for the store prefix, so primary and joiner side land on the
+    same prefix as long as the call order matches. backendIndex_ is also
+    process-local and increments only when the rank is an actual member of the
+    new group, so it stays aligned across primaries and joiners that all call
+    new_group in the same order.
+    """
+    configure_mooncake_device_filter(ctx.device_filters)
+    device = require_test_device(ctx.proc_rank, ctx.device_type)
+
+    assert ctx.world_size == 4, "this test assumes world_size=4"
+    initial_world_size = 2
+    join_ranks = [2, 3]
+    is_joiner = ctx.proc_rank >= initial_world_size
+
+    a_active = torch.tensor([1, 0], dtype=torch.int32, device=device)
+    b_active = torch.tensor([1, 0], dtype=torch.int32, device=device)
+    c_active = torch.tensor([1, 1, 0, 0], dtype=torch.int32, device=device)
+
+    if not is_joiner:
+        # Primary ranks: init WORLD with world_size=2, max_world_size=4
+        world_active = torch.tensor([1, 1, 0, 0], dtype=torch.int32, device=device)
+        dist_kwargs = {
+            "backend": ctx.backend_name,
+            "rank": ctx.proc_rank,
+            "world_size": initial_world_size,
+            "pg_options": pg.MooncakeBackendOptions(world_active, False, ctx.world_size),
+        }
+        if ctx.device_type == "cuda":
+            dist_kwargs["device_id"] = device
+        dist.init_process_group(**dist_kwargs)
+        world_backend = get_mooncake_backend(device_type=ctx.device_type)
+
+        # Subgroups with split-ranks pattern. All ranks in WORLD must call
+        # new_group in the same order even for groups they are not members of.
+        group_a = dist.new_group(
+            ranks=[0],
+            backend=ctx.backend_name,
+            pg_options=pg.MooncakeBackendOptions(a_active, False, 2),
+        )
+        group_b = dist.new_group(
+            ranks=[1],
+            backend=ctx.backend_name,
+            pg_options=pg.MooncakeBackendOptions(b_active, False, 2),
+        )
+        group_c = dist.new_group(
+            ranks=[0, 1],
+            backend=ctx.backend_name,
+            pg_options=pg.MooncakeBackendOptions(c_active, False, 4),
+        )
+        a_backend = get_mooncake_backend(group_a, device_type=ctx.device_type) if ctx.proc_rank == 0 else None
+        b_backend = get_mooncake_backend(group_b, device_type=ctx.device_type) if ctx.proc_rank == 1 else None
+        c_backend = get_mooncake_backend(group_c, device_type=ctx.device_type)
+
+        # Pre-activation: WORLD primary ranks sum to 1+2=3
+        t = torch.tensor([ctx.proc_rank + 1], dtype=torch.int32, device=device)
+        dist.all_reduce(t, op=dist.ReduceOp.SUM)
+        if int(t.cpu().item()) != 3:
+            raise AssertionError(f"WORLD pre: expected 3, got {int(t.cpu().item())}")
+
+        # group_c primaries (ranks [0,1]): sum to 1+2=3
+        tc = torch.tensor([ctx.proc_rank + 1], dtype=torch.int32, device=device)
+        dist.all_reduce(tc, op=dist.ReduceOp.SUM, group=group_c)
+        if int(tc.cpu().item()) != 3:
+            raise AssertionError(f"group_c pre: expected 3, got {int(tc.cpu().item())}")
+
+        if ctx.proc_rank == 0:
+            extend_event.set()
+
+        # WORLD: wait for joiners then recover
+        wait_until(
+            lambda: all(pg.get_peer_state(world_backend, join_ranks)),
+            timeout_s=60.0,
+            poll_interval_s=0.05,
+            description=f"rank {ctx.proc_rank} waiting for WORLD joiners",
+        )
+        pg.recover_ranks(world_backend, join_ranks)
+
+        # group_a: rank 0 waits for joiner (local rank 1 = global rank 2)
+        if ctx.proc_rank == 0:
+            wait_until(
+                lambda: pg.get_peer_state(a_backend, [1])[0],
+                timeout_s=60.0,
+                poll_interval_s=0.05,
+                description="rank 0 waiting for group_a joiner",
+            )
+            pg.recover_ranks(a_backend, [1])
+
+        # group_b: rank 1 waits for joiner (local rank 1 = global rank 3)
+        if ctx.proc_rank == 1:
+            wait_until(
+                lambda: pg.get_peer_state(b_backend, [1])[0],
+                timeout_s=60.0,
+                poll_interval_s=0.05,
+                description="rank 1 waiting for group_b joiner",
+            )
+            pg.recover_ranks(b_backend, [1])
+
+        # group_c: both primaries wait for both joiners (local ranks 2,3)
+        wait_until(
+            lambda: all(pg.get_peer_state(c_backend, [2, 3])),
+            timeout_s=60.0,
+            poll_interval_s=0.05,
+            description=f"rank {ctx.proc_rank} waiting for group_c joiners",
+        )
+        pg.recover_ranks(c_backend, [2, 3])
+
+        # Post-activation: WORLD all 4 ranks → 1+2+3+4=10
+        t = torch.tensor([ctx.proc_rank + 1], dtype=torch.int32, device=device)
+        dist.all_reduce(t, op=dist.ReduceOp.SUM)
+        if int(t.cpu().item()) != 10:
+            raise AssertionError(f"WORLD post: expected 10, got {int(t.cpu().item())}")
+
+        # group_a: ranks [0,2], values [1,3], sum=4
+        if ctx.proc_rank == 0:
+            ta = torch.tensor([ctx.proc_rank + 1], dtype=torch.int32, device=device)
+            dist.all_reduce(ta, op=dist.ReduceOp.SUM, group=group_a)
+            if int(ta.cpu().item()) != 4:
+                raise AssertionError(f"group_a post: expected 4, got {int(ta.cpu().item())}")
+
+        # group_b: ranks [1,3], values [2,4], sum=6
+        if ctx.proc_rank == 1:
+            tb = torch.tensor([ctx.proc_rank + 1], dtype=torch.int32, device=device)
+            dist.all_reduce(tb, op=dist.ReduceOp.SUM, group=group_b)
+            if int(tb.cpu().item()) != 6:
+                raise AssertionError(f"group_b post: expected 6, got {int(tb.cpu().item())}")
+
+        # group_c: all 4 ranks, sum=10
+        tc = torch.tensor([ctx.proc_rank + 1], dtype=torch.int32, device=device)
+        dist.all_reduce(tc, op=dist.ReduceOp.SUM, group=group_c)
+        if int(tc.cpu().item()) != 10:
+            raise AssertionError(f"group_c post: expected 10, got {int(tc.cpu().item())}")
+
+        ctx.record_result({"role": "extension_subgroups", "rank": ctx.proc_rank})
+    else:
+        # Joiners: wait for extend_event, then init WORLD with world_size=4
+        if not extend_event.wait(timeout=60.0):
+            raise TimeoutError("timed out waiting for extend_event")
+
+        world_active = torch.tensor([1, 1, 0, 0], dtype=torch.int32, device=device)
+        dist_kwargs = {
+            "backend": ctx.backend_name,
+            "rank": ctx.proc_rank,
+            "world_size": ctx.world_size,
+            "pg_options": pg.MooncakeBackendOptions(world_active, True, ctx.world_size),
+        }
+        if ctx.device_type == "cuda":
+            dist_kwargs["device_id"] = device
+        dist.init_process_group(**dist_kwargs)
+        world_backend = get_mooncake_backend(device_type=ctx.device_type)
+
+        # Subgroups: full eventual membership; matching call order with primaries.
+        group_a = dist.new_group(
+            ranks=[0, 2],
+            backend=ctx.backend_name,
+            pg_options=pg.MooncakeBackendOptions(a_active, True, 2),
+        )
+        group_b = dist.new_group(
+            ranks=[1, 3],
+            backend=ctx.backend_name,
+            pg_options=pg.MooncakeBackendOptions(b_active, True, 2),
+        )
+        group_c = dist.new_group(
+            ranks=[0, 1, 2, 3],
+            backend=ctx.backend_name,
+            pg_options=pg.MooncakeBackendOptions(c_active, True, 4),
+        )
+        a_backend = get_mooncake_backend(group_a, device_type=ctx.device_type) if ctx.proc_rank == 2 else None
+        b_backend = get_mooncake_backend(group_b, device_type=ctx.device_type) if ctx.proc_rank == 3 else None
+        c_backend = get_mooncake_backend(group_c, device_type=ctx.device_type)
+
+        # Joiners are local-only until join_group is called
+        t = torch.tensor([ctx.proc_rank + 1], dtype=torch.int32, device=device)
+        dist.all_reduce(t, op=dist.ReduceOp.SUM)
+        if int(t.cpu().item()) != ctx.proc_rank + 1:
+            raise AssertionError(
+                f"WORLD local-only: expected {ctx.proc_rank + 1}, got {int(t.cpu().item())}"
+            )
+
+        # Join groups in same order primaries created them
+        pg.join_group(world_backend)
+        if ctx.proc_rank == 2:
+            pg.join_group(a_backend)
+        if ctx.proc_rank == 3:
+            pg.join_group(b_backend)
+        pg.join_group(c_backend)
+
+        # Post-activation: WORLD all 4 ranks → 1+2+3+4=10
+        t = torch.tensor([ctx.proc_rank + 1], dtype=torch.int32, device=device)
+        dist.all_reduce(t, op=dist.ReduceOp.SUM)
+        if int(t.cpu().item()) != 10:
+            raise AssertionError(f"WORLD post: expected 10, got {int(t.cpu().item())}")
+
+        # group_a: ranks [0,2], sum=4
+        if ctx.proc_rank == 2:
+            ta = torch.tensor([ctx.proc_rank + 1], dtype=torch.int32, device=device)
+            dist.all_reduce(ta, op=dist.ReduceOp.SUM, group=group_a)
+            if int(ta.cpu().item()) != 4:
+                raise AssertionError(f"group_a post: expected 4, got {int(ta.cpu().item())}")
+
+        # group_b: ranks [1,3], sum=6
+        if ctx.proc_rank == 3:
+            tb = torch.tensor([ctx.proc_rank + 1], dtype=torch.int32, device=device)
+            dist.all_reduce(tb, op=dist.ReduceOp.SUM, group=group_b)
+            if int(tb.cpu().item()) != 6:
+                raise AssertionError(f"group_b post: expected 6, got {int(tb.cpu().item())}")
+
+        # group_c: all 4 ranks, sum=10
+        tc = torch.tensor([ctx.proc_rank + 1], dtype=torch.int32, device=device)
+        dist.all_reduce(tc, op=dist.ReduceOp.SUM, group=group_c)
+        if int(tc.cpu().item()) != 10:
+            raise AssertionError(f"group_c post: expected 10, got {int(tc.cpu().item())}")
+
+        ctx.record_result({"role": "extension_subgroups", "rank": ctx.proc_rank})
+
+
+def _run_allgather_reduce_scatter(
+    device: str,
+    active_world_size: int,
+    rank: int,
+) -> None:
+    """Run _allgather_base and _reduce_scatter_base and assert correctness.
+
+    active_world_size: number of currently active ranks (buffer size).
+    rank: this rank's logical rank (0-indexed).
+    Each rank contributes value (rank + 1).
+    """
+
+    # --- _allgather_base ---
+    # input: scalar (rank+1); output: flat buffer of active_world_size elements
+    input_t = torch.tensor([rank + 1], dtype=torch.int32, device=device)
+    output_t = torch.zeros(active_world_size, dtype=torch.int32, device=device)
+    dist.all_gather_into_tensor(output_t, input_t)
+    for j in range(active_world_size):
+        expected = j + 1
+        got = int(output_t[j].item())
+        if got != expected:
+            raise AssertionError(
+                f"allgather slot {j}: expected {expected}, got {got} "
+                f"(rank={rank}, active_world_size={active_world_size})"
+            )
+
+    # --- _reduce_scatter_base ---
+    # input: flat buffer [1, 2, ..., active_world_size]; output: scalar chunk
+    # With SUM and equal chunks of size 1, each rank receives its own slot's sum
+    # across all ranks. Since all ranks send the same input buffer [1..N],
+    # rank j receives sum of input[j] from all ranks = (j+1) * active_world_size.
+    input_rs = torch.arange(1, active_world_size + 1, dtype=torch.int32, device=device)
+    output_rs = torch.zeros(1, dtype=torch.int32, device=device)
+    dist.reduce_scatter_tensor(output_rs, input_rs)
+    expected_rs = (rank + 1) * active_world_size
+    got_rs = int(output_rs[0].item())
+    if got_rs != expected_rs:
+        raise AssertionError(
+            f"reduce_scatter slot {rank}: expected {expected_rs}, got {got_rs} "
+            f"(rank={rank}, active_world_size={active_world_size})"
+        )
+
+
+def _allgather_reduce_scatter_extension_worker(
+    ctx: MooncakePGWorkerContext,
+    extend_event: mp.Event,
+) -> None:
+    """Test _allgather_base and _reduce_scatter_base across elastic extension.
+
+    Layout: world_size=4, initial=3, extension_rank=3, max_world_size=4.
+    Pre-activation: 3 active ranks, max_world_size=4 → exercises the overflow path.
+    Post-activation: 4 active ranks → exercises correctness after extension.
+    """
+    configure_mooncake_device_filter(ctx.device_filters)
+    device = require_test_device(ctx.proc_rank, ctx.device_type)
+
+    assert ctx.world_size == 4
+    initial_world_size = 3
+    extension_rank = 3
+    join_ranks = [extension_rank]
+    is_joiner = ctx.proc_rank == extension_rank
+
+    if not is_joiner:
+        active = torch.tensor([1, 1, 1, 0], dtype=torch.int32, device=device)
+        dist_kwargs = {
+            "backend": ctx.backend_name,
+            "rank": ctx.proc_rank,
+            "world_size": initial_world_size,
+            "pg_options": pg.MooncakeBackendOptions(active, False, ctx.world_size),
+        }
+        if ctx.device_type == "cuda":
+            dist_kwargs["device_id"] = device
+        dist.init_process_group(**dist_kwargs)
+        backend = get_mooncake_backend(device_type=ctx.device_type)
+
+        # Pre-activation: 3 active ranks, max_world_size=4.
+        # This is the overflow path: buggy code would iterate 4 times into a
+        # buffer sized for 3.
+        _run_allgather_reduce_scatter(device, initial_world_size, ctx.proc_rank)
+
+        if ctx.proc_rank == 0:
+            extend_event.set()
+
+        wait_until(
+            lambda: all(pg.get_peer_state(backend, join_ranks)),
+            timeout_s=30.0,
+            poll_interval_s=0.05,
+            description=f"rank {ctx.proc_rank} waiting for joiner",
+        )
+        pg.recover_ranks(backend, join_ranks)
+
+        # Post-activation: all 4 ranks active.
+        _run_allgather_reduce_scatter(device, ctx.world_size, ctx.proc_rank)
+
+        ctx.record_result({"role": "primary", "rank": ctx.proc_rank})
+    else:
+        if not extend_event.wait(timeout=30.0):
+            raise TimeoutError("timed out waiting for extend_event")
+
+        active = torch.tensor([1, 1, 1, 0], dtype=torch.int32, device=device)
+        dist_kwargs = {
+            "backend": ctx.backend_name,
+            "rank": extension_rank,
+            "world_size": ctx.world_size,
+            "pg_options": pg.MooncakeBackendOptions(active, True, ctx.world_size),
+        }
+        if ctx.device_type == "cuda":
+            dist_kwargs["device_id"] = device
+        dist.init_process_group(**dist_kwargs)
+        backend = get_mooncake_backend(device_type=ctx.device_type)
+
+        pg.join_group(backend)
+
+        # Post-activation: all 4 ranks active.
+        _run_allgather_reduce_scatter(device, ctx.world_size, extension_rank)
+
+        ctx.record_result({"role": "joiner", "rank": extension_rank})
+
+
+def _allgather_reduce_scatter_recovery_worker(
+    ctx: MooncakePGWorkerContext,
+    broken_exited: mp.Event,
+    replacement_ready: mp.Event,
+    start_recovery: mp.Event,
+) -> None:
+    """Test _allgather_base and _reduce_scatter_base across rank recovery.
+
+    Layout: world_size=4, broken_rank=3, replacement takes rank 3.
+    Pre-failure: 4 active ranks.
+    Post-failure (3 survivors): 3 active ranks, max_world_size=4 → overflow path.
+    Post-recovery: 4 active ranks again.
+    """
+    broken_rank = ctx.world_size - 1
+    logical_rank = ctx.rank if ctx.proc_rank < ctx.world_size else broken_rank
+
+    if ctx.proc_rank < ctx.world_size:
+        device = ctx.init_group(rank=logical_rank)
+        backend = ctx.get_backend()
+
+        # Pre-failure: all 4 ranks active.
+        _run_allgather_reduce_scatter(device, ctx.world_size, logical_rank)
+
+        if logical_rank == broken_rank:
+            ctx.record_result({"role": "broken"})
+            broken_exited.set()
+            os._exit(0)
+
+        # Survivors: 3 active ranks, max_world_size=4 → overflow path.
+        broken_exited.wait()
+        _run_allgather_reduce_scatter(device, ctx.world_size - 1, logical_rank)
+
+        if logical_rank == 0:
+            start_recovery.set()
+
+        wait_until(
+            lambda: pg.get_peer_state(backend, [broken_rank])[0],
+            timeout_s=30.0,
+            poll_interval_s=2.0,
+            description=f"rank {logical_rank} waiting for replacement",
+        )
+        replacement_ready.wait()
+        pg.recover_ranks(backend, [broken_rank])
+
+        # Post-recovery: all 4 ranks active again.
+        _run_allgather_reduce_scatter(device, ctx.world_size, logical_rank)
+
+        ctx.record_result({"role": "survivor"})
+    else:
+        start_recovery.wait()
+        device = ctx.init_group(rank=logical_rank, is_extension=True)
+        backend = ctx.get_backend()
+        replacement_ready.set()
+        pg.join_group(backend)
+
+        # Post-recovery: all 4 ranks active.
+        _run_allgather_reduce_scatter(device, ctx.world_size, logical_rank)
+
+        ctx.record_result({"role": "replacement"})
+
+
 def _fault_detection_worker(
     ctx: MooncakePGWorkerContext,
     broken_exited: mp.Event,
@@ -187,6 +798,17 @@ class _ElasticMixin:
     world_size = 4
     spawn_timeout_s = 30.0
 
+    def test_dynamic_world_size(self) -> None:
+        """Test that dist.get_world_size() returns updated value after extend_group_size_to."""
+        rows = self.spawn_backend_and_collect(
+            _dynamic_world_size_worker,
+            timeout_s=30.0,
+        )
+
+        self.assert_all_ok(rows)
+        for row in rows:
+            self.assertEqual(row["new_ws"], self.world_size + 1)
+
     def test_failed_rank(self) -> None:
         """Test that survivors can continue collective after a rank fails."""
         spawn_ctx = mp.get_context("spawn")
@@ -259,6 +881,103 @@ class _ElasticMixin:
         for row in original_rows:
             self.assertEqual(row.get("baseline"), expected_baseline)
 
+    def test_extension_p2p_joiner_to_primary(self) -> None:
+        """Test P2P from an extension rank to an original rank after recovery."""
+        spawn_ctx = mp.get_context("spawn")
+        extend_event = spawn_ctx.Event()
+
+        rows = self.spawn_backend_and_collect(
+            _extension_p2p_worker,
+            extend_event,
+            "joiner_to_primary",
+            nprocs=self.world_size,
+            timeout_s=60.0,
+        )
+
+        self.assert_all_ok(rows)
+        self.assertEqual(len([r for r in rows if r.get("role") == "src"]), 1)
+        self.assertEqual(len([r for r in rows if r.get("role") == "dst"]), 1)
+
+    def test_extension_p2p_primary_to_joiner(self) -> None:
+        """Test P2P from an original rank to an extension rank after recovery."""
+        spawn_ctx = mp.get_context("spawn")
+        extend_event = spawn_ctx.Event()
+
+        rows = self.spawn_backend_and_collect(
+            _extension_p2p_worker,
+            extend_event,
+            "primary_to_joiner",
+            nprocs=self.world_size,
+            timeout_s=60.0,
+        )
+
+        self.assert_all_ok(rows)
+        self.assertEqual(len([r for r in rows if r.get("role") == "src"]), 1)
+        self.assertEqual(len([r for r in rows if r.get("role") == "dst"]), 1)
+
+    def test_extension_with_subgroups(self) -> None:
+        """Test extension with multiple disjoint subgroups using split-ranks pattern."""
+        spawn_ctx = mp.get_context("spawn")
+        extend_event = spawn_ctx.Event()
+
+        rows = self.spawn_backend_and_collect(
+            _extension_worker_with_subgroups,
+            extend_event,
+            nprocs=self.world_size,
+            timeout_s=60.0,
+        )
+
+        result_rows = [r for r in rows if r.get("role") == "extension_subgroups"]
+        self.assertEqual(len(result_rows), self.world_size)
+
+    def test_allgather_reduce_scatter_extension(self) -> None:
+        """Test _allgather_base/_reduce_scatter_base correctness across elastic extension.
+
+        Exercises the overflow path: pre-activation uses max_world_size=4 with only
+        3 active ranks, so the buggy code would access slot 3 of a size-3 buffer.
+        """
+        spawn_ctx = mp.get_context("spawn")
+        extend_event = spawn_ctx.Event()
+
+        rows = self.spawn_backend_and_collect(
+            _allgather_reduce_scatter_extension_worker,
+            extend_event,
+            nprocs=self.world_size,
+            timeout_s=60.0,
+        )
+
+        primary_rows = [r for r in rows if r.get("role") == "primary"]
+        joiner_rows = [r for r in rows if r.get("role") == "joiner"]
+        self.assertEqual(len(primary_rows), self.world_size - 1)
+        self.assertEqual(len(joiner_rows), 1)
+
+    def test_allgather_reduce_scatter_recovery(self) -> None:
+        """Test _allgather_base/_reduce_scatter_base correctness across rank recovery.
+
+        Exercises the overflow path: post-failure survivors run with 3 active ranks
+        and max_world_size=4, so the buggy code would access slot 3 of a size-3 buffer.
+        """
+        spawn_ctx = mp.get_context("spawn")
+        broken_exited = spawn_ctx.Event()
+        replacement_ready = spawn_ctx.Event()
+        start_recovery = spawn_ctx.Event()
+
+        rows = self.spawn_backend_and_collect(
+            _allgather_reduce_scatter_recovery_worker,
+            broken_exited,
+            replacement_ready,
+            start_recovery,
+            nprocs=self.world_size + 1,
+            timeout_s=60.0,
+        )
+
+        survivor_rows = [r for r in rows if r.get("role") == "survivor"]
+        replacement_rows = [r for r in rows if r.get("role") == "replacement"]
+        broken_rows = [r for r in rows if r.get("role") == "broken"]
+        self.assertEqual(len(survivor_rows), self.world_size - 1)
+        self.assertEqual(len(replacement_rows), 1)
+        self.assertGreaterEqual(len(broken_rows), 1)
+
 
 class TestMooncakePGElasticCPU(
     _ElasticMixin, MooncakePGCPUBackendTestCase
diff --git a/mooncake-store/CMakeLists.txt b/mooncake-store/CMakeLists.txt
index 54f3b1cf..5cf0f8a4 100644
--- a/mooncake-store/CMakeLists.txt
+++ b/mooncake-store/CMakeLists.txt
@@ -3,6 +3,34 @@ project(MooncakeStore VERSION 2.0.0)
 # Extract version components for C++ usage
 set(MOONCAKE_STORE_VERSION ${PROJECT_VERSION})
 
+# Read display version from mooncake-wheel/pyproject.toml
+set(PYPROJECT_TOML_PATH "${CMAKE_CURRENT_SOURCE_DIR}/../mooncake-wheel/pyproject.toml")
+if(EXISTS "${PYPROJECT_TOML_PATH}")
+    file(READ "${PYPROJECT_TOML_PATH}" PYPROJECT_CONTENT)
+    string(REGEX MATCH "version = \"([^\"]+)\"" _ "${PYPROJECT_CONTENT}")
+    set(MOONCAKE_DISPLAY_VERSION_BASE "${CMAKE_MATCH_1}")
+endif()
+if(NOT MOONCAKE_DISPLAY_VERSION_BASE)
+    message(WARNING "Could not read display version from ${PYPROJECT_TOML_PATH}; falling back to MOONCAKE_STORE_VERSION")
+    set(MOONCAKE_DISPLAY_VERSION_BASE "${MOONCAKE_STORE_VERSION}")
+endif()
+
+# Get short git commit hash for display
+find_package(Git QUIET)
+if(GIT_FOUND)
+    execute_process(
+        COMMAND ${GIT_EXECUTABLE} rev-parse --short HEAD
+        WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}
+        OUTPUT_VARIABLE GIT_SHORT_HASH
+        OUTPUT_STRIP_TRAILING_WHITESPACE
+        ERROR_QUIET
+    )
+endif()
+if(NOT GIT_SHORT_HASH)
+    set(GIT_SHORT_HASH "unknown")
+endif()
+set(MOONCAKE_DISPLAY_VERSION "${MOONCAKE_DISPLAY_VERSION_BASE} (git: ${GIT_SHORT_HASH})")
+
 configure_file(include/version.h.in include/version.h @ONLY)
 
 if (STORE_USE_ETCD)
@@ -10,13 +38,18 @@ if (STORE_USE_ETCD)
     set(ETCD_WRAPPER_LIB ${CMAKE_CURRENT_BINARY_DIR}/../mooncake-common/etcd/libetcd_wrapper.so)
 endif()
 
+if (STORE_USE_K8S_LEASE)
+    set(K8S_LEASE_WRAPPER_INCLUDE ${CMAKE_CURRENT_BINARY_DIR}/../mooncake-common/k8s-lease/)
+    set(K8S_LEASE_WRAPPER_LIB ${CMAKE_CURRENT_BINARY_DIR}/../mooncake-common/k8s-lease/libk8s_lease_wrapper.so)
+endif()
+
 if (STORE_USE_REDIS)
     find_path(MOONCAKE_STORE_HIREDIS_INCLUDE_DIR hiredis/hiredis.h REQUIRED)
     find_library(MOONCAKE_STORE_HIREDIS_LIBRARY hiredis REQUIRED)
     message(STATUS "Redis HA backend: Enabled")
 endif()
 
-if (NOT STORE_USE_ETCD AND NOT STORE_USE_REDIS)
+if (NOT STORE_USE_ETCD AND NOT STORE_USE_REDIS AND NOT STORE_USE_K8S_LEASE)
     message(STATUS "Store HA backends are disabled")
 endif()
 
@@ -37,5 +70,8 @@ add_subdirectory(src)
 
 if (BUILD_UNIT_TESTS)
     add_subdirectory(tests)
+endif()
+
+if (BUILD_BENCHMARK)
     add_subdirectory(benchmarks)
 endif()
diff --git a/mooncake-store/benchmarks/CMakeLists.txt b/mooncake-store/benchmarks/CMakeLists.txt
index 8798af89..76d23cda 100644
--- a/mooncake-store/benchmarks/CMakeLists.txt
+++ b/mooncake-store/benchmarks/CMakeLists.txt
@@ -8,6 +8,11 @@ add_executable(master_bench master_bench.cpp)
 target_link_libraries(master_bench PRIVATE cachelib_memory_allocator
                                            mooncake_store)
 
+# Add BatchGetReplicaList latency benchmark executable
+add_executable(batch_get_replica_bench batch_get_replica_bench.cpp)
+target_link_libraries(batch_get_replica_bench PRIVATE cachelib_memory_allocator
+                                                       mooncake_store)
+
 # Add file interface benchmark executable
 add_executable(file_interface_bench file_interface_bench.cpp)
 target_link_libraries(file_interface_bench PRIVATE mooncake_store glog::glog)
@@ -21,9 +26,16 @@ target_link_libraries(
                                 gflags::gflags glog::glog pthread)
 
 # Add allocation strategy benchmark executable
-# This benchmark tests AllocationStrategy performance with configurable 
+# This benchmark tests AllocationStrategy performance with configurable
 # segment counts, allocation sizes, replica counts, and workload patterns
 add_executable(allocation_strategy_bench allocation_strategy_bench.cpp)
 target_link_libraries(
   allocation_strategy_bench PRIVATE mooncake_store cachelib_memory_allocator
                                     gflags::gflags glog::glog pthread)
+
+# Add NoF worker pool benchmark executable
+if(USE_NOF)
+  add_executable(nof_worker_pool_bench nof_worker_pool_bench.cpp)
+  target_link_libraries(nof_worker_pool_bench PRIVATE mooncake_store glog::glog
+                                                       gflags::gflags)
+endif()
diff --git a/mooncake-store/benchmarks/README.md b/mooncake-store/benchmarks/README.md
new file mode 100644
index 00000000..e461836d
--- /dev/null
+++ b/mooncake-store/benchmarks/README.md
@@ -0,0 +1,85 @@
+# Mooncake Store Benchmarks
+
+This directory contains benchmark tools for Mooncake Store internals.
+
+## Allocation Strategy Benchmark
+
+`allocation_strategy_bench` evaluates Store allocation behavior across segment
+counts, replica counts, allocation strategies, and workload patterns.
+
+Build the benchmark from an existing CMake build directory:
+
+```bash
+cmake --build build --target allocation_strategy_bench -j$(nproc)
+```
+
+### Size-Class Churn Fragmentation Benchmark
+
+The `size_class_churn` workload measures fragmentation under mixed-size
+KVCache-like allocation pressure. It pre-fills the simulated cluster when
+`--prefill_pct` is set, then repeatedly allocates objects from weighted size
+classes. On allocation failure it randomly evicts a fraction of live objects and
+retries.
+
+When prefill is enabled, the prefill attempt cap is auto-derived from target
+utilization, total cluster capacity, weighted average object size, and replica
+count, with a 5000-attempt minimum for small cases.
+
+This is an allocation-strategy-layer benchmark. It complements the existing
+`dsa` workload by adding explicit fragmentation sampling and configurable
+weighted size-class patterns. It is not a replacement for `allocator_bench`,
+which remains the low-level `OffsetAllocator` microbenchmark.
+
+Run a small local validation:
+
+```bash
+./build/mooncake-store/benchmarks/allocation_strategy_bench \
+  --workload=size_class_churn \
+  --segment_capacity=1024 \
+  --num_allocations=10000 \
+  --prefill_pct=70
+```
+
+Run a larger baseline:
+
+```bash
+./build/mooncake-store/benchmarks/allocation_strategy_bench \
+  --workload=size_class_churn \
+  --segment_capacity=1024 \
+  --num_allocations=100000 \
+  --prefill_pct=80
+```
+
+Supported size-class patterns:
+
+- `kv_mixed`: 4KB at 70%, 256KB at 20%, and 3.12MB at 10%.
+- `dsa_pair`: 3.12MB KV pages at 50% and 643KB indexer entries at 50%.
+- `all`: run both patterns.
+
+Key output columns:
+
+- `Throughput`, `Avg(ns)`, `P50(ns)`, `P90(ns)`, and `P99(ns)` measure
+  allocation performance.
+- `Frag_avg`, `Frag_p50`, `Frag_p90`, and `Frag_p99` summarize sampled
+  fragmentation ratios.
+- `LargestFreeMB` shows the final largest contiguous free region.
+- `Evictions` counts fail-triggered eviction rounds during measurement.
+- `Full/Partial/Fail/Total` reports allocation outcomes. Only results with
+  `result->size() == replica_num` count as full success; shorter replica
+  results are counted as partial allocations.
+
+Fragmentation is computed per `OffsetBufferAllocator` and then averaged by free
+space:
+
+```text
+1 - largest_free_region / total_free_space
+```
+
+The weighted average avoids treating free space in different Store segments as
+one mergeable region. `LargestFreeMB` still reports the final largest contiguous
+free region across all segments.
+
+The benchmark also prints a one-line `Prefill summary`, `Fragmentation summary`,
+and `Size-class breakdown` after each result row, so reviewers can read the
+actual prefill utilization, fragmentation, and per-size-class latency numbers
+without manually deriving them from the table.
diff --git a/mooncake-store/benchmarks/allocation_strategy_bench.cpp b/mooncake-store/benchmarks/allocation_strategy_bench.cpp
index 9db7d546..10e45fe7 100644
--- a/mooncake-store/benchmarks/allocation_strategy_bench.cpp
+++ b/mooncake-store/benchmarks/allocation_strategy_bench.cpp
@@ -4,7 +4,10 @@
 #include <fstream>
 #include <iomanip>
 #include <iostream>
+#include <limits>
 #include <numeric>
+#include <random>
+#include <sstream>
 #include <string>
 #include <vector>
 #include <sys/resource.h>
@@ -21,13 +24,14 @@ DEFINE_int64(segment_capacity, 1024,
              "Per-segment capacity in MB (base capacity for skewed mode)");
 DEFINE_int32(num_allocations, 10000, "Number of allocations to benchmark");
 DEFINE_int32(convergence_sample_interval, 100,
-             "Sample utilization stddev every N allocations (scaleout only)");
-DEFINE_bool(
-    run_all, false,
-    "Also run the Scale-Out matrix in addition to the default Fillup matrix");
+             "Sample utilization stddev or fragmentation every N allocations");
+DEFINE_bool(run_all, false,
+            "Run all workload matrices instead of only the selected workload");
 
 // Scale-Out workload flags
-DEFINE_string(workload, "fillup", "Workload type: fillup (default), scaleout");
+DEFINE_string(workload, "fillup",
+              "Workload type: fillup (default), scaleout, dsa, "
+              "size_class_churn");
 DEFINE_int32(
     scale_out_trigger_pct, 50,
     "Pre-fill cluster to this utilization % before injecting new nodes "
@@ -42,6 +46,39 @@ DEFINE_int32(
     "measured allocation loop (0 = disabled, original behavior). "
     "When >0, early exit on consecutive failures is suppressed so that "
     "all num_allocations are attempted under near-full conditions.");
+DEFINE_int64(dsa_segment_capacity, 1024ULL * 1024,
+             "Per-segment capacity in MB for DSA workload");
+DEFINE_double(dsa_capacity_multiplier, 3.0,
+              "DSA workload measures (multiplier * cluster_capacity) worth "
+              "of allocations under steady-state churn. Combined with "
+              "--dsa_warmup_multiplier, total per-case work is "
+              "(warmup + capacity) * cluster_capacity, clamped to "
+              "[kDsaMinAllocs, --dsa_max_allocs]; --num_allocations is "
+              "ignored in DSA mode.");
+DEFINE_double(dsa_warmup_multiplier, 1.0,
+              "DSA warmup-phase write volume in units of cluster_capacity, "
+              "executed BEFORE the measurement window so the cluster is in "
+              "post-fill eviction churn when stats start being collected. "
+              "All stats during warmup (latencies, util_ratios, evict "
+              "rounds, succ/total) are discarded. Set to 0.0 to disable "
+              "warmup entirely (legacy behavior).");
+DEFINE_int32(dsa_max_allocs, 50000000,
+             "Per-case upper bound on the auto-derived DSA allocation count. "
+             "Raise this if you use very large segments (e.g. 100GB+) and "
+             "want big clusters to actually reach steady-state eviction; "
+             "lower it to bound benchmark runtime. Each allocation costs ~1-3 "
+             "us, so the default upper bound can take minutes per saturated "
+             "case.");
+DEFINE_double(dsa_evict_ratio, 0.05,
+              "Fraction of live objects to evict on each Allocate failure "
+              "(default 0.05 = 5%). Lower values keep "
+              "the steady-state cluster fill closer to the fragmentation "
+              "ceiling but trigger evictions more frequently.");
+DEFINE_string(size_class_pattern, "kv_mixed",
+              "Size-class churn pattern: kv_mixed, dsa_pair, or all");
+DEFINE_double(size_class_evict_ratio, 0.02,
+              "Fraction of live objects to evict on each Allocate failure "
+              "in size_class_churn mode.");
 
 using namespace mooncake;
 
@@ -58,12 +95,28 @@ constexpr size_t kMaxExpectedAllocs = 600000;
 constexpr double kLargeClusterThresholdGB = 500.0;
 constexpr int kPreFillSampleInterval = 100;
 constexpr int kMinMeasurementAllocs = 100;
-constexpr int kMaxMeasurementAllocs = 200000;
+constexpr int kMaxMeasurementAllocs = 50000000;
 constexpr int kMinSamplesForConvergence = 10;
 
+// DSA workload sizes (DeepSeek defaults; adjust in source for other models).
+constexpr size_t kDsaKvSize = 3274752;      // ~3.12 MB per KV page
+constexpr size_t kDsaIndexerSize = 658432;  // 643 KB per indexer entry
+constexpr int kDsaMaxBatch = 128;           // per-round batch upper bound
+constexpr int kDsaMaxRetries = 5;
+// Lower bound on the auto-derived allocation count per DSA case.
+constexpr int kDsaMinAllocs = 100;
+constexpr int kSizeClassMaxRetries = 5;
+constexpr size_t kSizeClassMinPrefillAttempts = 5000;
+// The theoretical prefill budget assumes every allocation succeeds. Use a
+// small multiplier to absorb partial allocations, retries, and fragmentation
+// without turning unreachable targets into long-running cases.
+constexpr double kSizeClassPrefillAttemptMultiplier = 2.0;
+
 enum class WorkloadType {
     FILL_UP,    // Only allocate, measure throughput/latency
     SCALE_OUT,  // Inject new nodes mid-run, measure adoption speed
+    DSA,        // DSA paired KV+indexer with random fail-triggered eviction
+    SIZE_CLASS_CHURN  // Mixed size classes with steady-state churn
 };
 
 struct BenchConfig {
@@ -84,6 +137,80 @@ struct BenchConfig {
     WorkloadType workload_type;
     int scale_out_trigger_pct;
     int scale_out_new_segments;
+
+    // DSA workload knobs (only used when workload_type == DSA).
+    bool dsa_paired = false;  // false = KV-only, true = KV+indexer pair
+
+    // Size-class churn knobs (only used when workload_type ==
+    // SIZE_CLASS_CHURN).
+    std::string size_class_pattern = "kv_mixed";
+};
+
+struct UtilRatioStats {
+    double min = 0.0;
+    double p99 = 0.0;
+    double p90 = 0.0;
+    double p50 = 0.0;
+    double max = 0.0;
+    double avg = 0.0;
+    bool valid = false;
+};
+
+struct DistributionStats {
+    double min = 0.0;
+    double p50 = 0.0;
+    double p90 = 0.0;
+    double p99 = 0.0;
+    double max = 0.0;
+    double avg = 0.0;
+    bool valid = false;
+};
+
+struct FragmentationSnapshot {
+    uint64_t total_free_space = 0;
+    uint64_t largest_free_region = 0;
+    uint64_t capacity = 0;
+    double fragmentation_ratio = 0.0;
+    bool valid = false;
+};
+
+struct SizeClassSpec {
+    std::string name;
+    size_t size;
+    int weight;
+};
+
+struct SizeClassStat {
+    std::string name;
+    size_t size = 0;
+    int weight = 0;
+    int success_count = 0;
+    int partial_count = 0;
+    int failed_count = 0;
+    int total_count = 0;
+    DistributionStats latency_stats;
+};
+
+enum class SizeClassAllocationStatus {
+    FAILED,
+    PARTIAL,
+    FULL,
+};
+
+struct SizeClassAllocationResult {
+    SizeClassAllocationStatus status = SizeClassAllocationStatus::FAILED;
+    size_t replica_count = 0;
+};
+
+struct SizeClassPrefillStats {
+    size_t attempts = 0;
+    size_t max_attempts = 0;
+    int full_count = 0;
+    int partial_count = 0;
+    int failed_count = 0;
+    int requested_pct = 0;
+    double achieved_util_pct = 0.0;
+    bool reached_target = false;
 };
 
 /**
@@ -109,10 +236,23 @@ struct BenchResultBase {
 
     int success_count = 0;  // successful allocations
     int total_count = 0;    // total attempted allocations
+    int partial_count = 0;  // partial allocations, if reported
+    int failed_count = 0;   // failed allocations, if reported
 
     double final_util_stddev;  // utilization stddev at run end
     double final_avg_util;     // average utilization at run end
 
+    // Number of times fail-triggered eviction fired (DSA workload only;
+    // remains 0 for fill-up / scale-out).
+    int evict_count = 0;
+
+    // DSA paired flag controls how AllocSize is rendered in print output.
+    // Defaults to false so existing fill-up / scale-out output is unchanged.
+    bool dsa_paired = false;
+
+    // Temporal cluster utilization stats sampled during the measured loop.
+    UtilRatioStats util_ratio_stats;
+
     virtual ~BenchResultBase() = default;
 };
 
@@ -135,6 +275,15 @@ struct ScaleOutResult : BenchResultBase {
     std::vector<double> new_node_util_over_time;
 };
 
+struct SizeClassChurnResult : BenchResultBase {
+    std::string pattern_name;
+    SizeClassPrefillStats prefill_stats;
+    DistributionStats fragmentation_stats;
+    DistributionStats largest_free_mb_stats;
+    FragmentationSnapshot final_fragmentation;
+    std::vector<SizeClassStat> size_class_stats;
+};
+
 static double computeClusterCapacityGB(int num_segments, size_t base_capacity,
                                        bool skewed) {
     double total = 0.0;
@@ -151,9 +300,9 @@ static double computeClusterCapacityGB(int num_segments, size_t base_capacity,
 
 static void setupResourceLimits() {
     struct rlimit rl;
-    // Cap virtual address space (RLIMIT_AS) to 4TB
-    rl.rlim_cur = 4096ULL * 1024 * 1024 * 1024;
-    rl.rlim_max = 4096ULL * 1024 * 1024 * 1024;
+    // DSA cases simulate multi-TB clusters via virtual address space.
+    rl.rlim_cur = 200ULL * 1024 * 1024 * 1024 * 1024;
+    rl.rlim_max = 200ULL * 1024 * 1024 * 1024 * 1024;
     setrlimit(RLIMIT_AS, &rl);
 
     rl.rlim_cur = RLIM_INFINITY;
@@ -304,6 +453,174 @@ static double computeAverageUtil(
     return count > 0 ? sum / count : 0.0;
 }
 
+static UtilRatioStats computeUtilRatioStats(std::vector<double>& util_ratios) {
+    UtilRatioStats stats;
+    if (util_ratios.empty()) return stats;
+
+    std::sort(util_ratios.begin(), util_ratios.end());
+
+    auto percentile = [&](double p) -> double {
+        size_t idx = static_cast<size_t>(util_ratios.size() * p);
+        if (idx >= util_ratios.size()) idx = util_ratios.size() - 1;
+        return util_ratios[idx];
+    };
+
+    stats.min = util_ratios.front();
+    stats.p99 = percentile(0.01);
+    stats.p90 = percentile(0.10);
+    stats.p50 = percentile(0.50);
+    stats.max = util_ratios.back();
+    stats.avg = std::accumulate(util_ratios.begin(), util_ratios.end(), 0.0) /
+                util_ratios.size();
+    stats.valid = true;
+    return stats;
+}
+
+static DistributionStats computeDistributionStats(std::vector<double>& values) {
+    DistributionStats stats;
+    if (values.empty()) return stats;
+
+    std::sort(values.begin(), values.end());
+
+    auto percentile = [&](double p) -> double {
+        size_t idx = static_cast<size_t>(std::round(p * (values.size() - 1)));
+        return values[idx];
+    };
+
+    stats.min = values.front();
+    stats.p50 = percentile(0.50);
+    stats.p90 = percentile(0.90);
+    stats.p99 = percentile(0.99);
+    stats.max = values.back();
+    stats.avg =
+        std::accumulate(values.begin(), values.end(), 0.0) / values.size();
+    stats.valid = true;
+    return stats;
+}
+
+static FragmentationSnapshot computeFragmentationSnapshot(
+    const AllocatorManager& manager) {
+    FragmentationSnapshot snapshot;
+    double weighted_fragmentation = 0.0;
+
+    for (const auto& name : manager.getNames()) {
+        const auto* allocs = manager.getAllocators(name);
+        if (!allocs) continue;
+
+        for (const auto& alloc : *allocs) {
+            auto offset_alloc =
+                std::dynamic_pointer_cast<OffsetBufferAllocator>(alloc);
+            if (!offset_alloc) continue;
+
+            auto allocator = offset_alloc->getOffsetAllocator();
+            if (!allocator) continue;
+
+            auto metrics = allocator->get_metrics();
+            snapshot.total_free_space += metrics.total_free_space_;
+            snapshot.capacity += metrics.capacity;
+            snapshot.largest_free_region = std::max(
+                snapshot.largest_free_region, metrics.largest_free_region_);
+            if (metrics.total_free_space_ > 0) {
+                double local_fragmentation =
+                    1.0 - (static_cast<double>(metrics.largest_free_region_) /
+                           static_cast<double>(metrics.total_free_space_));
+                local_fragmentation = std::clamp(local_fragmentation, 0.0, 1.0);
+                weighted_fragmentation +=
+                    local_fragmentation * metrics.total_free_space_;
+            }
+            snapshot.valid = true;
+        }
+    }
+
+    if (snapshot.valid && snapshot.total_free_space > 0) {
+        snapshot.fragmentation_ratio =
+            weighted_fragmentation /
+            static_cast<double>(snapshot.total_free_space);
+    }
+
+    return snapshot;
+}
+
+static std::vector<SizeClassSpec> getSizeClassSpecs(
+    const std::string& pattern_name) {
+    if (pattern_name == "kv_mixed") {
+        return {
+            {"small", 4 * KiB, 70},
+            {"medium", 256 * KiB, 20},
+            {"large", kDsaKvSize, 10},
+        };
+    }
+
+    if (pattern_name == "dsa_pair") {
+        return {
+            {"kv", kDsaKvSize, 50},
+            {"indexer", kDsaIndexerSize, 50},
+        };
+    }
+
+    return {};
+}
+
+static size_t chooseSizeClassIndex(const std::vector<SizeClassSpec>& specs,
+                                   std::mt19937& rng) {
+    if (specs.empty()) return 0;
+
+    int total_weight = 0;
+    for (const auto& spec : specs) {
+        total_weight += spec.weight;
+    }
+
+    if (total_weight <= 0) return 0;
+
+    std::uniform_int_distribution<int> dist(1, total_weight);
+    int pick = dist(rng);
+    for (size_t i = 0; i < specs.size(); ++i) {
+        pick -= specs[i].weight;
+        if (pick <= 0) return i;
+    }
+
+    return specs.size() - 1;
+}
+
+static double computeWeightedAverageObjectSize(
+    const std::vector<SizeClassSpec>& specs) {
+    double weighted_size = 0.0;
+    int total_weight = 0;
+
+    for (const auto& spec : specs) {
+        if (spec.weight <= 0) continue;
+        weighted_size += static_cast<double>(spec.size) * spec.weight;
+        total_weight += spec.weight;
+    }
+
+    if (total_weight <= 0) return 0.0;
+    return static_cast<double>(weighted_size / total_weight);
+}
+
+static size_t deriveSizeClassPrefillMaxAttempts(
+    const AllocatorManager& manager, const BenchConfig& cfg,
+    const std::vector<SizeClassSpec>& specs) {
+    if (cfg.prefill_pct <= 0 || cfg.replica_num <= 0) return 0;
+
+    const double avg_object_size = computeWeightedAverageObjectSize(specs);
+    if (avg_object_size <= 0.0) return kSizeClassMinPrefillAttempts;
+
+    const double target_bytes =
+        static_cast<double>(computeTotalCapacity(manager)) * cfg.prefill_pct /
+        100.0;
+    const double bytes_per_attempt =
+        avg_object_size * static_cast<double>(cfg.replica_num);
+    if (target_bytes <= 0.0 || bytes_per_attempt <= 0.0) {
+        return kSizeClassMinPrefillAttempts;
+    }
+
+    const double derived_attempts =
+        std::ceil((target_bytes / bytes_per_attempt) *
+                  kSizeClassPrefillAttemptMultiplier);
+    return std::max(kSizeClassMinPrefillAttempts,
+                    static_cast<size_t>(derived_attempts));
+}
+
 static std::string strategyName(AllocationStrategyType type) {
     switch (type) {
         case AllocationStrategyType::RANDOM:
@@ -667,8 +984,378 @@ static ScaleOutResult runScaleOutBenchmark(const BenchConfig& cfg) {
     return res;
 }
 
+static void evictRandomFraction(std::vector<std::vector<Replica>>& live,
+                                double ratio, std::mt19937& rng) {
+    if (live.empty()) return;
+
+    size_t to_drop =
+        std::max<size_t>(1, static_cast<size_t>(live.size() * ratio));
+    if (to_drop > live.size()) to_drop = live.size();
+
+    for (size_t i = 0; i < to_drop; ++i) {
+        std::uniform_int_distribution<size_t> dist(0, live.size() - 1);
+        size_t idx = dist(rng);
+        std::swap(live[idx], live.back());
+        live.pop_back();  // Replica destructor returns memory to allocator.
+    }
+}
+
+// Try to allocate once; on failure, sample utilization before eviction.
+static bool dsaAllocateWithEvict(
+    const std::shared_ptr<AllocationStrategy>& strategy,
+    AllocatorManager& manager, size_t size, int replica_num,
+    std::vector<std::vector<Replica>>& live, std::mt19937& rng,
+    int& evict_count, double evict_ratio,
+    std::vector<double>* pre_evict_util = nullptr) {
+    for (int attempt = 0; attempt <= kDsaMaxRetries; ++attempt) {
+        auto result = strategy->Allocate(manager, size, replica_num);
+        if (result.has_value()) {
+            live.push_back(std::move(result.value()));
+            return true;
+        }
+
+        if (live.empty()) return false;
+        if (attempt == kDsaMaxRetries) return false;
+
+        if (attempt == 0 && pre_evict_util != nullptr) {
+            pre_evict_util->push_back(computeAverageUtilAll(manager));
+        }
+
+        evictRandomFraction(live, evict_ratio, rng);
+        ++evict_count;
+    }
+
+    return false;
+}
+
+static SizeClassAllocationResult sizeClassAllocateWithEvict(
+    const std::shared_ptr<AllocationStrategy>& strategy,
+    AllocatorManager& manager, size_t size, int replica_num,
+    std::vector<std::vector<Replica>>& live, std::mt19937& rng,
+    int& evict_count, double evict_ratio) {
+    for (int attempt = 0; attempt <= kSizeClassMaxRetries; ++attempt) {
+        auto result = strategy->Allocate(manager, size, replica_num);
+        if (result.has_value()) {
+            size_t replica_count = result->size();
+            if (replica_count == 0) {
+                return {};
+            }
+            live.push_back(std::move(result.value()));
+            return {
+                replica_count == static_cast<size_t>(replica_num)
+                    ? SizeClassAllocationStatus::FULL
+                    : SizeClassAllocationStatus::PARTIAL,
+                replica_count,
+            };
+        }
+
+        if (live.empty()) return {};
+        if (attempt == kSizeClassMaxRetries) return {};
+
+        evictRandomFraction(live, evict_ratio, rng);
+        ++evict_count;
+    }
+
+    return {};
+}
+
+static SizeClassPrefillStats prefillSizeClassChurn(
+    const std::shared_ptr<AllocationStrategy>& strategy,
+    AllocatorManager& manager, const BenchConfig& cfg,
+    const std::vector<SizeClassSpec>& specs,
+    std::vector<std::vector<Replica>>& live_allocations, std::mt19937& rng) {
+    SizeClassPrefillStats stats;
+    if (cfg.prefill_pct <= 0 || specs.empty()) return stats;
+    stats.requested_pct = cfg.prefill_pct;
+    stats.max_attempts = deriveSizeClassPrefillMaxAttempts(manager, cfg, specs);
+
+    int consec_failures = 0;
+    int evict_throwaway = 0;
+    const int kMaxConsecFailures = 10;
+
+    while (stats.attempts < stats.max_attempts) {
+        if (stats.attempts % kPreFillSampleInterval == 0) {
+            stats.achieved_util_pct = computeAverageUtilAll(manager) * 100.0;
+            if (stats.achieved_util_pct >= cfg.prefill_pct) {
+                stats.reached_target = true;
+                break;
+            }
+        }
+
+        size_t class_idx = chooseSizeClassIndex(specs, rng);
+        auto alloc_result = sizeClassAllocateWithEvict(
+            strategy, manager, specs[class_idx].size, cfg.replica_num,
+            live_allocations, rng, evict_throwaway,
+            FLAGS_size_class_evict_ratio);
+        ++stats.attempts;
+        if (alloc_result.status == SizeClassAllocationStatus::FULL) {
+            ++stats.full_count;
+            consec_failures = 0;
+        } else if (alloc_result.status == SizeClassAllocationStatus::PARTIAL) {
+            ++stats.partial_count;
+            consec_failures = 0;
+        } else {
+            ++stats.failed_count;
+            if (++consec_failures >= kMaxConsecFailures) {
+                break;
+            }
+        }
+    }
+
+    stats.achieved_util_pct = computeAverageUtilAll(manager) * 100.0;
+    if (stats.achieved_util_pct >= cfg.prefill_pct) {
+        stats.reached_target = true;
+    }
+    return stats;
+}
+
+// Run DSA workload; the allocation count is derived from cluster capacity.
+static FillUpResult runDsaBenchmark(const BenchConfig& cfg) {
+    AllocatorManager manager =
+        createCluster(cfg.num_segments, cfg.segment_capacity, cfg.skewed);
+    auto strategy = CreateAllocationStrategy(cfg.strategy_type);
+
+    const size_t total_capacity = computeTotalCapacity(manager);
+    const size_t avg_obj_size =
+        cfg.dsa_paired ? (kDsaKvSize + kDsaIndexerSize) / 2 : kDsaKvSize;
+    const size_t bytes_per_alloc = avg_obj_size * cfg.replica_num;
+    const double warmup_mult = std::max(0.0, FLAGS_dsa_warmup_multiplier);
+    const double measure_mult = std::max(0.0, FLAGS_dsa_capacity_multiplier);
+    const double total_mult = warmup_mult + measure_mult;
+
+    int total_allocs =
+        (bytes_per_alloc == 0)
+            ? kDsaMinAllocs
+            : static_cast<int>(total_mult * total_capacity / bytes_per_alloc);
+    total_allocs =
+        std::clamp(total_allocs, kDsaMinAllocs, FLAGS_dsa_max_allocs);
+    int warmup_allocs =
+        (total_mult > 0.0)
+            ? static_cast<int>(total_allocs * (warmup_mult / total_mult))
+            : 0;
+    int measure_allocs = std::max(0, total_allocs - warmup_allocs);
+
+    std::vector<double> latencies;
+    latencies.reserve(measure_allocs);
+
+    std::vector<double> util_ratios;
+    util_ratios.reserve(measure_allocs);
+
+    // Use one live set so eviction matches allocator_bench.
+    std::vector<std::vector<Replica>> live_allocations;
+    live_allocations.reserve(std::min(total_allocs, 1 << 20));
+
+    std::mt19937 rng(42);
+    std::uniform_int_distribution<int> batch_dist(1, kDsaMaxBatch);
+
+    if (warmup_allocs > 0) {
+        int warmup_count = 0;
+        int warmup_evict_throwaway = 0;
+        auto warmup_run = [&](size_t size,
+                              std::vector<std::vector<Replica>>& live) {
+            (void)dsaAllocateWithEvict(strategy, manager, size, cfg.replica_num,
+                                       live, rng, warmup_evict_throwaway,
+                                       FLAGS_dsa_evict_ratio);
+            ++warmup_count;
+        };
+
+        while (warmup_count < warmup_allocs) {
+            int batch = batch_dist(rng);
+            for (int i = 0; i < batch && warmup_count < warmup_allocs; ++i) {
+                warmup_run(kDsaKvSize, live_allocations);
+            }
+            if (cfg.dsa_paired) {
+                for (int i = 0; i < batch && warmup_count < warmup_allocs;
+                     ++i) {
+                    warmup_run(kDsaIndexerSize, live_allocations);
+                }
+            }
+        }
+    }
+
+    int success_count = 0;
+    int total_count = 0;
+    int evict_count = 0;
+
+    auto run_one = [&](size_t size,
+                       std::vector<std::vector<Replica>>& live) -> bool {
+        auto t0 = std::chrono::high_resolution_clock::now();
+        bool ok = dsaAllocateWithEvict(strategy, manager, size, cfg.replica_num,
+                                       live, rng, evict_count,
+                                       FLAGS_dsa_evict_ratio, &util_ratios);
+        auto t1 = std::chrono::high_resolution_clock::now();
+
+        latencies.push_back(
+            std::chrono::duration<double, std::nano>(t1 - t0).count());
+        ++total_count;
+        if (ok) ++success_count;
+
+        return ok;
+    };
+
+    auto total_start = std::chrono::high_resolution_clock::now();
+
+    while (total_count < measure_allocs) {
+        int batch = batch_dist(rng);
+
+        for (int i = 0; i < batch && total_count < measure_allocs; ++i) {
+            run_one(kDsaKvSize, live_allocations);
+        }
+
+        if (cfg.dsa_paired) {
+            for (int i = 0; i < batch && total_count < measure_allocs; ++i) {
+                run_one(kDsaIndexerSize, live_allocations);
+            }
+        }
+    }
+
+    auto total_end = std::chrono::high_resolution_clock::now();
+    double total_us =
+        std::chrono::duration<double, std::micro>(total_end - total_start)
+            .count();
+
+    FillUpResult res;
+    res.strategy_name = cfg.strategy_name;
+    res.num_segments = cfg.num_segments;
+    res.alloc_size = cfg.alloc_size;  // KV size; AllocSize column displays it.
+    res.replica_num = cfg.replica_num;
+    res.skewed = cfg.skewed;
+    res.cluster_capacity_gb = computeClusterCapacityGB(
+        cfg.num_segments, cfg.segment_capacity, cfg.skewed);
+    res.final_util_stddev = computeUtilizationStdDev(manager);
+    res.final_avg_util = computeAverageUtilAll(manager);
+    res.success_count = success_count;
+    res.total_count = total_count;
+    res.evict_count = evict_count;
+    res.dsa_paired = cfg.dsa_paired;
+    res.util_ratio_stats = computeUtilRatioStats(util_ratios);
+    computeLatencyStats(latencies, total_us, total_count, res);
+    return res;
+}
+
+static SizeClassChurnResult runSizeClassChurnBenchmark(const BenchConfig& cfg) {
+    AllocatorManager manager =
+        createCluster(cfg.num_segments, cfg.segment_capacity, cfg.skewed);
+    auto strategy = CreateAllocationStrategy(cfg.strategy_type);
+    auto specs = getSizeClassSpecs(cfg.size_class_pattern);
+
+    std::vector<SizeClassStat> per_class_stats;
+    per_class_stats.reserve(specs.size());
+    std::vector<std::vector<double>> per_class_latencies(specs.size());
+    for (const auto& spec : specs) {
+        SizeClassStat stat;
+        stat.name = spec.name;
+        stat.size = spec.size;
+        stat.weight = spec.weight;
+        per_class_stats.push_back(std::move(stat));
+    }
+
+    std::vector<double> latencies;
+    latencies.reserve(cfg.num_allocations);
+    int sample_interval = std::max(1, FLAGS_convergence_sample_interval);
+    std::vector<double> fragmentation_samples;
+    fragmentation_samples.reserve(cfg.num_allocations / sample_interval + 2);
+    std::vector<double> largest_free_mb_samples;
+    largest_free_mb_samples.reserve(fragmentation_samples.capacity());
+
+    std::vector<std::vector<Replica>> live_allocations;
+    live_allocations.reserve(std::min(cfg.num_allocations, 1 << 20));
+
+    std::mt19937 rng(42);
+    SizeClassPrefillStats prefill_stats = prefillSizeClassChurn(
+        strategy, manager, cfg, specs, live_allocations, rng);
+
+    int success_count = 0;
+    int partial_count = 0;
+    int failed_count = 0;
+    int total_count = 0;
+    int evict_count = 0;
+    double instrumentation_time_us = 0.0;
+
+    auto total_start = std::chrono::high_resolution_clock::now();
+
+    for (int i = 0; i < cfg.num_allocations; ++i) {
+        size_t class_idx = chooseSizeClassIndex(specs, rng);
+        const auto& spec = specs[class_idx];
+
+        auto t0 = std::chrono::high_resolution_clock::now();
+        auto alloc_result = sizeClassAllocateWithEvict(
+            strategy, manager, spec.size, cfg.replica_num, live_allocations,
+            rng, evict_count, FLAGS_size_class_evict_ratio);
+        auto t1 = std::chrono::high_resolution_clock::now();
+
+        double latency_ns =
+            std::chrono::duration<double, std::nano>(t1 - t0).count();
+        latencies.push_back(latency_ns);
+        per_class_latencies[class_idx].push_back(latency_ns);
+
+        ++total_count;
+        ++per_class_stats[class_idx].total_count;
+        if (alloc_result.status == SizeClassAllocationStatus::FULL) {
+            ++success_count;
+            ++per_class_stats[class_idx].success_count;
+        } else if (alloc_result.status == SizeClassAllocationStatus::PARTIAL) {
+            ++partial_count;
+            ++per_class_stats[class_idx].partial_count;
+        } else {
+            ++failed_count;
+            ++per_class_stats[class_idx].failed_count;
+        }
+
+        if ((i + 1) % sample_interval == 0 || i == cfg.num_allocations - 1) {
+            auto s0 = std::chrono::high_resolution_clock::now();
+            auto snapshot = computeFragmentationSnapshot(manager);
+            if (snapshot.valid) {
+                fragmentation_samples.push_back(snapshot.fragmentation_ratio);
+                largest_free_mb_samples.push_back(
+                    static_cast<double>(snapshot.largest_free_region) / MiB);
+            }
+            auto s1 = std::chrono::high_resolution_clock::now();
+            instrumentation_time_us +=
+                std::chrono::duration<double, std::micro>(s1 - s0).count();
+        }
+    }
+
+    auto total_end = std::chrono::high_resolution_clock::now();
+    double total_us =
+        std::chrono::duration<double, std::micro>(total_end - total_start)
+            .count();
+    total_us -= instrumentation_time_us;
+    total_us = std::max(total_us, 1.0);
+
+    for (size_t i = 0; i < per_class_stats.size(); ++i) {
+        per_class_stats[i].latency_stats =
+            computeDistributionStats(per_class_latencies[i]);
+    }
+
+    SizeClassChurnResult res;
+    res.strategy_name = cfg.strategy_name;
+    res.num_segments = cfg.num_segments;
+    res.alloc_size = 0;
+    res.replica_num = cfg.replica_num;
+    res.skewed = cfg.skewed;
+    res.cluster_capacity_gb = computeClusterCapacityGB(
+        cfg.num_segments, cfg.segment_capacity, cfg.skewed);
+    res.final_util_stddev = computeUtilizationStdDev(manager);
+    res.final_avg_util = computeAverageUtilAll(manager);
+    res.success_count = success_count;
+    res.partial_count = partial_count;
+    res.failed_count = failed_count;
+    res.total_count = total_count;
+    res.evict_count = evict_count;
+    res.pattern_name = cfg.size_class_pattern;
+    res.prefill_stats = prefill_stats;
+    res.fragmentation_stats = computeDistributionStats(fragmentation_samples);
+    res.largest_free_mb_stats =
+        computeDistributionStats(largest_free_mb_samples);
+    res.final_fragmentation = computeFragmentationSnapshot(manager);
+    res.size_class_stats = std::move(per_class_stats);
+    computeLatencyStats(latencies, total_us, total_count, res);
+    return res;
+}
+
 static void printFillUpHeader() {
-    std::cout << std::string(170, '-') << std::endl;
+    std::cout << std::string(184, '-') << std::endl;
     std::cout << std::left << std::setw(18) << "Strategy" << std::setw(9)
               << "Replica" << std::setw(10) << "Segments" << std::setw(12)
               << "AllocSize" << std::setw(12) << "Cluster(GB)" << std::setw(8)
@@ -676,8 +1363,9 @@ static void printFillUpHeader() {
               << std::setw(12) << "Avg(ns)" << std::setw(12) << "P50(ns)"
               << std::setw(12) << "P90(ns)" << std::setw(12) << "P99(ns)"
               << std::setw(12) << "UtilStdDev" << std::setw(10) << "AvgUtil%"
-              << std::setw(15) << "Succ/Total" << std::endl;
-    std::cout << std::string(170, '-') << std::endl;
+              << std::setw(15) << "Succ/Total" << std::setw(14) << "Evictions"
+              << std::endl;
+    std::cout << std::string(184, '-') << std::endl;
 }
 
 static void printFillUpResult(const FillUpResult& r) {
@@ -685,18 +1373,66 @@ static void printFillUpResult(const FillUpResult& r) {
         std::to_string(r.success_count) + "/" + std::to_string(r.total_count);
     std::ostringstream cap_ss;
     cap_ss << std::fixed << std::setprecision(1) << r.cluster_capacity_gb;
+    std::string alloc_size_str =
+        r.dsa_paired ? (std::to_string(r.alloc_size / KiB) + "+" +
+                        std::to_string(kDsaIndexerSize / KiB) + "KB")
+                     : (std::to_string(r.alloc_size / KiB) + "KB");
     std::cout << std::left << std::setw(18) << r.strategy_name << std::setw(9)
               << r.replica_num << std::setw(10) << r.num_segments
-              << std::setw(12) << (std::to_string(r.alloc_size / KiB) + "KB")
-              << std::setw(12) << cap_ss.str() << std::setw(8)
-              << (r.skewed ? "yes" : "no") << std::right << std::fixed
+              << std::setw(12) << alloc_size_str << std::setw(12)
+              << cap_ss.str() << std::setw(8) << (r.skewed ? "yes" : "no")
+              << std::right << std::fixed << std::setprecision(0)
+              << std::setw(14) << r.throughput << std::setw(12) << r.avg_ns
+              << std::setw(12) << r.p50_ns << std::setw(12) << r.p90_ns
+              << std::setw(12) << r.p99_ns << std::setprecision(4)
+              << std::setw(12) << r.final_util_stddev << std::setprecision(2)
+              << std::setw(9) << (r.final_avg_util * 100.0) << "%"
+              << std::setw(15) << alloc_ratio << std::setw(14) << r.evict_count
+              << std::endl;
+}
+
+static void printDsaHeader() {
+    std::cout << std::string(264, '-') << std::endl;
+    std::cout << std::left << std::setw(18) << "Strategy" << std::setw(9)
+              << "Replica" << std::setw(10) << "Segments" << std::setw(12)
+              << "AllocSize" << std::setw(12) << "Cluster(GB)" << std::setw(8)
+              << "Skewed" << std::right << std::setw(14) << "UtilRatio_min"
+              << std::setw(14) << "UtilRatio_p99" << std::setw(14)
+              << "UtilRatio_p90" << std::setw(14) << "UtilRatio_p50"
+              << std::setw(14) << "UtilRatio_max" << std::setw(14)
+              << "UtilRatio_avg" << std::setw(14) << "Throughput"
+              << std::setw(12) << "Avg(ns)" << std::setw(12) << "P50(ns)"
+              << std::setw(12) << "P90(ns)" << std::setw(12) << "P99(ns)"
+              << std::setw(12) << "UtilStdDev" << std::setw(17)
+              << "Evictions(round)" << std::setw(20) << "Succ/Total"
+              << std::endl;
+    std::cout << std::string(264, '-') << std::endl;
+}
+
+static void printDsaResult(const FillUpResult& r) {
+    std::string alloc_ratio =
+        std::to_string(r.success_count) + "/" + std::to_string(r.total_count);
+    std::ostringstream cap_ss;
+    cap_ss << std::fixed << std::setprecision(1) << r.cluster_capacity_gb;
+    std::string alloc_size_str =
+        r.dsa_paired ? (std::to_string(r.alloc_size / KiB) + "+" +
+                        std::to_string(kDsaIndexerSize / KiB) + "KB")
+                     : (std::to_string(r.alloc_size / KiB) + "KB");
+    const auto& urs = r.util_ratio_stats;
+    std::cout << std::left << std::setw(18) << r.strategy_name << std::setw(9)
+              << r.replica_num << std::setw(10) << r.num_segments
+              << std::setw(12) << alloc_size_str << std::setw(12)
+              << cap_ss.str() << std::setw(8) << (r.skewed ? "yes" : "no")
+              << std::right << std::fixed << std::setprecision(6)
+              << std::setw(14) << urs.min << std::setw(14) << urs.p99
+              << std::setw(14) << urs.p90 << std::setw(14) << urs.p50
+              << std::setw(14) << urs.max << std::setw(14) << urs.avg
               << std::setprecision(0) << std::setw(14) << r.throughput
               << std::setw(12) << r.avg_ns << std::setw(12) << r.p50_ns
               << std::setw(12) << r.p90_ns << std::setw(12) << r.p99_ns
               << std::setprecision(4) << std::setw(12) << r.final_util_stddev
-              << std::setprecision(2) << std::setw(9)
-              << (r.final_avg_util * 100.0) << "%" << std::setw(15)
-              << alloc_ratio << std::endl;
+              << std::setprecision(0) << std::setw(17) << r.evict_count
+              << std::setw(20) << alloc_ratio << std::endl;
 }
 
 static void printScaleOutHeader() {
@@ -746,6 +1482,93 @@ static void printScaleOutResult(const ScaleOutResult& r) {
               << std::endl;
 }
 
+static void printSizeClassChurnHeader() {
+    std::cout << std::string(260, '-') << std::endl;
+    std::cout << std::left << std::setw(18) << "Strategy" << std::setw(9)
+              << "Replica" << std::setw(10) << "Segments" << std::setw(14)
+              << "Pattern" << std::setw(12) << "Cluster(GB)" << std::setw(8)
+              << "Skewed" << std::right << std::setw(14) << "Throughput"
+              << std::setw(12) << "Avg(ns)" << std::setw(12) << "P50(ns)"
+              << std::setw(12) << "P90(ns)" << std::setw(12) << "P99(ns)"
+              << std::setw(12) << "Frag_avg" << std::setw(12) << "Frag_p50"
+              << std::setw(12) << "Frag_p90" << std::setw(12) << "Frag_p99"
+              << std::setw(15) << "LargestFreeMB" << std::setw(10) << "AvgUtil%"
+              << std::setw(24) << "Full/Partial/Fail/Total" << std::setw(14)
+              << "Evictions" << std::endl;
+    std::cout << std::string(260, '-') << std::endl;
+}
+
+static void printSizeClassChurnResult(const SizeClassChurnResult& r) {
+    std::string alloc_ratio = std::to_string(r.success_count) + "/" +
+                              std::to_string(r.partial_count) + "/" +
+                              std::to_string(r.failed_count) + "/" +
+                              std::to_string(r.total_count);
+    std::ostringstream cap_ss;
+    cap_ss << std::fixed << std::setprecision(1) << r.cluster_capacity_gb;
+
+    double final_largest_free_mb =
+        static_cast<double>(r.final_fragmentation.largest_free_region) / MiB;
+
+    std::cout << std::left << std::setw(18) << r.strategy_name << std::setw(9)
+              << r.replica_num << std::setw(10) << r.num_segments
+              << std::setw(14) << r.pattern_name << std::setw(12)
+              << cap_ss.str() << std::setw(8) << (r.skewed ? "yes" : "no")
+              << std::right << std::fixed << std::setprecision(0)
+              << std::setw(14) << r.throughput << std::setw(12) << r.avg_ns
+              << std::setw(12) << r.p50_ns << std::setw(12) << r.p90_ns
+              << std::setw(12) << r.p99_ns << std::setprecision(4)
+              << std::setw(12) << r.fragmentation_stats.avg << std::setw(12)
+              << r.fragmentation_stats.p50 << std::setw(12)
+              << r.fragmentation_stats.p90 << std::setw(12)
+              << r.fragmentation_stats.p99 << std::setprecision(1)
+              << std::setw(15) << final_largest_free_mb << std::setprecision(2)
+              << std::setw(9) << (r.final_avg_util * 100.0) << "%"
+              << std::setw(24) << alloc_ratio << std::setw(14) << r.evict_count
+              << std::endl;
+
+    std::cout << "Prefill summary [" << r.strategy_name
+              << ", pattern=" << r.pattern_name
+              << ", segments=" << r.num_segments
+              << ", replica=" << r.replica_num
+              << "]: requested_pct=" << std::fixed << std::setprecision(2)
+              << r.prefill_stats.requested_pct
+              << ", achieved_pct=" << r.prefill_stats.achieved_util_pct
+              << ", reached=" << (r.prefill_stats.reached_target ? "yes" : "no")
+              << ", attempts=" << r.prefill_stats.attempts << "/"
+              << r.prefill_stats.max_attempts
+              << ", full/partial/failed=" << r.prefill_stats.full_count << "/"
+              << r.prefill_stats.partial_count << "/"
+              << r.prefill_stats.failed_count << std::endl;
+
+    std::cout << "Fragmentation summary [" << r.strategy_name
+              << ", pattern=" << r.pattern_name
+              << ", segments=" << r.num_segments
+              << ", replica=" << r.replica_num
+              << ", skewed=" << (r.skewed ? "yes" : "no")
+              << "]: avg=" << std::fixed << std::setprecision(4)
+              << r.fragmentation_stats.avg
+              << ", p50=" << r.fragmentation_stats.p50
+              << ", p90=" << r.fragmentation_stats.p90
+              << ", p99=" << r.fragmentation_stats.p99
+              << ", max=" << r.fragmentation_stats.max
+              << ", final_largest_free=" << std::setprecision(1)
+              << final_largest_free_mb << " MB" << std::endl;
+
+    std::cout << "Size-class breakdown:";
+    for (const auto& stat : r.size_class_stats) {
+        std::string ratio = std::to_string(stat.success_count) + "/" +
+                            std::to_string(stat.partial_count) + "/" +
+                            std::to_string(stat.failed_count) + "/" +
+                            std::to_string(stat.total_count);
+        std::cout << " " << stat.name << "(" << (stat.size / KiB)
+                  << "KB,w=" << stat.weight
+                  << ",full/partial/failed/total=" << ratio
+                  << ",p99_ns=" << std::fixed << std::setprecision(0)
+                  << stat.latency_stats.p99 << ")";
+    }
+    std::cout << std::endl;
+}
+
 static void runFillupBenchmarks() {
     std::vector<bool> skewed_options = {false, true};
     std::vector<int> segment_counts = {1, 10, 100, 512, 1024};
@@ -874,6 +1697,207 @@ static void runScaleOutMatrix() {
     }
 }
 
+// DSA matrix runner
+static void runDsaMatrix() {
+    std::vector<bool> skewed_options = {false, true};
+    std::vector<int> segment_counts = {1, 2, 4, 8, 16};
+    std::vector<int> replica_nums = {1, 2, 3};
+    std::vector<bool> paired_modes = {false, true};
+    std::vector<AllocationStrategyType> strategies = {
+        AllocationStrategyType::RANDOM,
+        AllocationStrategyType::FREE_RATIO_FIRST,
+    };
+
+    size_t seg_cap_mb = static_cast<size_t>(FLAGS_dsa_segment_capacity);
+
+    const double total_mult = std::max(0.0, FLAGS_dsa_warmup_multiplier) +
+                              std::max(0.0, FLAGS_dsa_capacity_multiplier);
+    std::cout
+        << "\n=== DSA Paired KV+Indexer Benchmark Matrix ===\n"
+        << "Workload: per-round burst of N kvcache (" << kDsaKvSize
+        << " B) allocations followed (paired mode only) by N indexer ("
+        << kDsaIndexerSize << " B) allocations; N uniform in [1, "
+        << kDsaMaxBatch << "].\n"
+        << "Eviction: on Allocate failure, drop "
+        << (FLAGS_dsa_evict_ratio * 100.0)
+        << "% of live objects at random; retry up to " << kDsaMaxRetries
+        << " times. Latency includes evict+retry time. Evictions(round) "
+           "counts how many such drop cycles fire during measurement, "
+           "NOT the number of objects evicted.\n"
+        << "UtilRatio_*: cluster utilization sampled at each eviction "
+           "trigger (peak fill the strategy reached before being forced "
+           "to evict), NOT a time-weighted average over all allocations.\n"
+        << "Phases: warmup writes " << FLAGS_dsa_warmup_multiplier
+        << "x cluster_capacity worth of allocations to reach steady state "
+           "(stats discarded), then measurement writes "
+        << FLAGS_dsa_capacity_multiplier
+        << "x cluster_capacity (stats collected). "
+           "Set --dsa_warmup_multiplier=0 to disable warmup.\n"
+        << "Config: dsa_segment_capacity=" << seg_cap_mb << " MB ("
+        << (seg_cap_mb / 1024.0)
+        << " GB), dsa_max_allocs=" << FLAGS_dsa_max_allocs << "\n"
+        << "Total per-case allocations auto-derived as " << total_mult
+        << "x * cluster_capacity / (avg_obj_size * replica), clamped to ["
+        << kDsaMinAllocs << ", " << FLAGS_dsa_max_allocs
+        << "]; warmup/measurement budgets split proportionally when clamp "
+           "activates. --num_allocations is ignored in DSA mode.\n"
+        << "Skewed setup: half nodes are (base + 50%) capacity, half are "
+           "(base - 50%)\n"
+        << std::endl;
+
+    std::vector<BenchConfig> configs;
+    for (auto paired : paired_modes) {
+        for (auto skew : skewed_options) {
+            for (auto strategy : strategies) {
+                for (auto segs : segment_counts) {
+                    for (auto rep : replica_nums) {
+                        if (rep > segs) continue;
+                        BenchConfig cfg;
+                        cfg.num_segments = segs;
+                        cfg.segment_capacity = seg_cap_mb * MiB;
+                        cfg.alloc_size = kDsaKvSize;
+                        cfg.replica_num = rep;
+                        // num_allocations is auto-derived inside
+                        // runDsaBenchmark from cluster capacity; the flag
+                        // value is intentionally not propagated.
+                        cfg.num_allocations = 0;
+                        cfg.skewed = skew;
+                        cfg.strategy_type = strategy;
+                        cfg.strategy_name = strategyName(strategy);
+                        cfg.workload_type = WorkloadType::DSA;
+                        cfg.dsa_paired = paired;
+                        configs.push_back(cfg);
+                    }
+                }
+            }
+        }
+    }
+
+    bool first = true;
+    bool prev_paired = false;
+    AllocationStrategyType prev_strategy = AllocationStrategyType::RANDOM;
+
+    for (const auto& cfg : configs) {
+        if (first || cfg.dsa_paired != prev_paired) {
+            std::cout << "\n--- "
+                      << (cfg.dsa_paired ? "Paired (KV + Indexer)" : "KV-only")
+                      << " ---" << std::endl;
+            prev_paired = cfg.dsa_paired;
+            first = true;  // force header re-print at start of new section
+        }
+
+        if (first || cfg.strategy_type != prev_strategy) {
+            printDsaHeader();
+            prev_strategy = cfg.strategy_type;
+            first = false;
+        }
+
+        auto result = runDsaBenchmark(cfg);
+        printDsaResult(result);
+    }
+}
+
+static void runSizeClassChurnMatrix() {
+    std::vector<bool> skewed_options = {false, true};
+    std::vector<int> segment_counts = {1, 10, 100};
+    std::vector<int> replica_nums = {1, 2, 3};
+    std::vector<AllocationStrategyType> strategies = {
+        AllocationStrategyType::RANDOM,
+        AllocationStrategyType::FREE_RATIO_FIRST,
+    };
+
+    std::vector<std::string> patterns;
+    if (FLAGS_size_class_pattern == "all") {
+        patterns = {"kv_mixed", "dsa_pair"};
+    } else {
+        patterns = {FLAGS_size_class_pattern};
+    }
+
+    if (FLAGS_size_class_evict_ratio <= 0.0 ||
+        FLAGS_size_class_evict_ratio > 1.0) {
+        std::cout << "Invalid size_class_evict_ratio: "
+                  << FLAGS_size_class_evict_ratio
+                  << ". Use a value in the range (0.0, 1.0]." << std::endl;
+        return;
+    }
+
+    for (const auto& pattern : patterns) {
+        if (getSizeClassSpecs(pattern).empty()) {
+            std::cout << "Invalid size_class_pattern: " << pattern
+                      << ". Use --size_class_pattern=kv_mixed, dsa_pair, or "
+                         "all."
+                      << std::endl;
+            return;
+        }
+    }
+
+    std::cout << "\n=== Size-Class Churn Fragmentation Benchmark Matrix ===\n"
+              << "Workload: prefill to --prefill_pct if set, then run "
+              << FLAGS_num_allocations
+              << " mixed-size allocation attempts with fail-triggered random "
+                 "eviction and retry.\n"
+              << "Fragmentation: 1 - largest_free_region / total_free_space, "
+                 "sampled every --convergence_sample_interval allocations.\n"
+              << "Config: segment_capacity=" << FLAGS_segment_capacity
+              << " MB, prefill_pct=" << FLAGS_prefill_pct
+              << ", evict_ratio=" << FLAGS_size_class_evict_ratio
+              << ", size_class_pattern=" << FLAGS_size_class_pattern << "\n"
+              << "Patterns: kv_mixed = 4KB:70%, 256KB:20%, 3198KB:10%; "
+                 "dsa_pair = 3198KB:50%, 643KB:50%.\n"
+              << "Skewed setup: half nodes are (base + 50%) capacity, half are "
+                 "(base - 50%)\n"
+              << std::endl;
+
+    std::vector<BenchConfig> configs;
+    for (const auto& pattern : patterns) {
+        for (auto skew : skewed_options) {
+            for (auto strategy : strategies) {
+                for (auto segs : segment_counts) {
+                    for (auto rep : replica_nums) {
+                        if (rep > segs) continue;
+                        BenchConfig cfg;
+                        cfg.num_segments = segs;
+                        cfg.segment_capacity =
+                            static_cast<size_t>(FLAGS_segment_capacity) * MiB;
+                        cfg.alloc_size = 0;
+                        cfg.replica_num = rep;
+                        cfg.num_allocations = FLAGS_num_allocations;
+                        cfg.skewed = skew;
+                        cfg.strategy_type = strategy;
+                        cfg.strategy_name = strategyName(strategy);
+                        cfg.prefill_pct = FLAGS_prefill_pct;
+                        cfg.workload_type = WorkloadType::SIZE_CLASS_CHURN;
+                        cfg.size_class_pattern = pattern;
+                        configs.push_back(cfg);
+                    }
+                }
+            }
+        }
+    }
+
+    bool first = true;
+    std::string prev_pattern;
+    AllocationStrategyType prev_strategy = AllocationStrategyType::RANDOM;
+
+    for (const auto& cfg : configs) {
+        if (first || cfg.size_class_pattern != prev_pattern) {
+            std::cout << "\n--- Pattern: " << cfg.size_class_pattern << " ---"
+                      << std::endl;
+            prev_pattern = cfg.size_class_pattern;
+            first = true;
+        }
+
+        if (first || cfg.strategy_type != prev_strategy) {
+            printSizeClassChurnHeader();
+            prev_strategy = cfg.strategy_type;
+            first = false;
+        }
+
+        auto result = runSizeClassChurnBenchmark(cfg);
+        printSizeClassChurnResult(result);
+    }
+}
+
 int main(int argc, char* argv[]) {
     gflags::SetUsageMessage(
         "AllocationStrategy performance benchmark.\n"
@@ -884,13 +1908,20 @@ int main(int argc, char* argv[]) {
     if (FLAGS_run_all) {
         runFillupBenchmarks();
         runScaleOutMatrix();
+        runDsaMatrix();
+        runSizeClassChurnMatrix();
     } else if (FLAGS_workload == "fillup") {
         runFillupBenchmarks();
     } else if (FLAGS_workload == "scaleout") {
         runScaleOutMatrix();
+    } else if (FLAGS_workload == "dsa") {
+        runDsaMatrix();
+    } else if (FLAGS_workload == "size_class_churn") {
+        runSizeClassChurnMatrix();
     } else {
         std::cout << "Invalid workload type: " << FLAGS_workload
-                  << ". Use --workload=fillup or --workload=scaleout."
+                  << ". Use --workload=fillup, --workload=scaleout, "
+                     "--workload=dsa, or --workload=size_class_churn."
                   << std::endl;
     }
 
diff --git a/mooncake-store/benchmarks/allocator_bench.cpp b/mooncake-store/benchmarks/allocator_bench.cpp
index e330b4c4..8c4a8d90 100644
--- a/mooncake-store/benchmarks/allocator_bench.cpp
+++ b/mooncake-store/benchmarks/allocator_bench.cpp
@@ -12,7 +12,7 @@ using namespace mooncake::offset_allocator;
 
 class OffsetAllocatorBenchHelper {
    public:
-    OffsetAllocatorBenchHelper(uint64_t baseAddress, uint32_t poolSize,
+    OffsetAllocatorBenchHelper(uint64_t baseAddress, size_t poolSize,
                                uint32_t maxAllocs)
         : pool_size_(poolSize),
           allocated_size_(0),
@@ -20,7 +20,8 @@ class OffsetAllocatorBenchHelper {
           rd_(),
           gen_(rd_()) {}
 
-    void allocate(uint32_t size) {
+    void allocate(uint32_t size, double evict_ratio = 0.0,
+                  std::vector<double>* pre_evict_util_ratios = nullptr) {
         while (true) {
             auto handle = allocator_->allocate(size);
             if (handle.has_value()) {
@@ -32,14 +33,26 @@ class OffsetAllocatorBenchHelper {
             if (allocated_.size() == 0) {
                 break;
             }
-            std::uniform_int_distribution<uint32_t> dist(0,
-                                                         allocated_.size() - 1);
-            auto index = dist(gen_);
-            std::swap(allocated_[index], allocated_.back());
-            std::swap(allocated_sizes_[index], allocated_sizes_.back());
-            allocated_size_ -= allocated_sizes_.back();
-            allocated_.pop_back();
-            allocated_sizes_.pop_back();
+            if (pre_evict_util_ratios != nullptr) {
+                pre_evict_util_ratios->push_back(get_allocated_ratio());
+            }
+            size_t evict_count = 1;
+            if (evict_ratio > 0.0) {
+                evict_count = std::max<size_t>(
+                    1, static_cast<size_t>(allocated_.size() * evict_ratio));
+            }
+            evict_count = std::min(evict_count, allocated_.size());
+
+            for (size_t i = 0; i < evict_count; ++i) {
+                std::uniform_int_distribution<uint32_t> dist(
+                    0, allocated_.size() - 1);
+                auto index = dist(gen_);
+                std::swap(allocated_[index], allocated_.back());
+                std::swap(allocated_sizes_[index], allocated_sizes_.back());
+                allocated_size_ -= allocated_sizes_.back();
+                allocated_.pop_back();
+                allocated_sizes_.pop_back();
+            }
         }
     }
 
@@ -160,13 +173,11 @@ void random_size_allocation_benchmark() {
         static_cast<double>(benchmark_num);
 
     std::sort(util_ratios.begin(), util_ratios.end());
-
     const double min_util = util_ratios.front();
     const double max_util = util_ratios.back();
     const double p50 = util_ratios[util_ratios.size() * 0.50];
     const double p90 = util_ratios[util_ratios.size() * 0.10];
     const double p99 = util_ratios[util_ratios.size() * 0.01];
-
     const double mean_util =
         std::accumulate(util_ratios.begin(), util_ratios.end(), 0.0) /
         util_ratios.size();
@@ -178,8 +189,93 @@ void random_size_allocation_benchmark() {
     std::cout << "avg alloc time: " << avg_time_ns << " ns/op" << std::endl;
 }
 
+template <typename BenchHelper>
+void paired_kv_indexer_allocation_benchmark() {
+    std::cout << std::endl
+              << "=== Paired KV/Indexer Allocation Benchmark (DSA) ==="
+              << std::endl;
+
+    const uint32_t kvcache_size = 3274752;                  // 3,274,752 B
+    const uint32_t indexer_size = 643u * 1024;              // 643 KB
+    const size_t pool_size = 1024ull * 1024 * 1024 * 1024;  // 1TB
+    const int max_per_round = 128;
+    const int warmup_rounds = 5000;
+    const int num_rounds = 50000;
+
+    size_t max_allocs = pool_size / indexer_size + 1024;
+    BenchHelper bench_helper(0x1000, pool_size, max_allocs);
+
+    std::random_device rd;
+    std::mt19937 gen(rd());
+    std::uniform_int_distribution<int> per_round_dist(1, max_per_round);
+
+    // Warmup
+    for (int round = 0; round < warmup_rounds; round++) {
+        int per_round = per_round_dist(gen);
+        for (int i = 0; i < per_round; i++) {
+            bench_helper.allocate(kvcache_size, 0.05);
+        }
+        for (int i = 0; i < per_round; i++) {
+            bench_helper.allocate(indexer_size, 0.05);
+        }
+    }
+
+    std::vector<double> util_ratios;
+    util_ratios.reserve(num_rounds);
+
+    auto start_time = std::chrono::high_resolution_clock::now();
+    size_t total_alloc_count = 0;
+    for (int round = 0; round < num_rounds; round++) {
+        int per_round = per_round_dist(gen);
+        for (int i = 0; i < per_round; i++) {
+            bench_helper.allocate(kvcache_size, 0.05, &util_ratios);
+            ++total_alloc_count;
+        }
+        for (int i = 0; i < per_round; i++) {
+            bench_helper.allocate(indexer_size, 0.05, &util_ratios);
+            ++total_alloc_count;
+        }
+    }
+    auto end_time = std::chrono::high_resolution_clock::now();
+
+    const double avg_time_ns =
+        std::chrono::duration_cast<std::chrono::nanoseconds>(end_time -
+                                                             start_time)
+            .count() /
+        static_cast<double>(total_alloc_count);
+
+    std::sort(util_ratios.begin(), util_ratios.end());
+    const double min_util = util_ratios.empty() ? 0.0 : util_ratios.front();
+    const double max_util = util_ratios.empty() ? 0.0 : util_ratios.back();
+    const double p50 =
+        util_ratios.empty() ? 0.0 : util_ratios[util_ratios.size() * 0.50];
+    const double p90 =
+        util_ratios.empty() ? 0.0 : util_ratios[util_ratios.size() * 0.10];
+    const double p99 =
+        util_ratios.empty() ? 0.0 : util_ratios[util_ratios.size() * 0.01];
+    const double mean_util =
+        util_ratios.empty()
+            ? 0.0
+            : std::accumulate(util_ratios.begin(), util_ratios.end(), 0.0) /
+                  util_ratios.size();
+
+    std::cout << std::fixed << std::setprecision(6);
+    std::cout << "kvcache size: " << kvcache_size
+              << " B, indexer size: " << indexer_size << " B" << std::endl;
+    std::cout << "pool size: " << (pool_size / (1024.0 * 1024 * 1024))
+              << " GB, warmup rounds: " << warmup_rounds
+              << ", benchmark rounds: " << num_rounds << std::endl;
+    std::cout << "util ratio (min / p99 / p90 / p50 / max / avg): " << min_util
+              << " / " << p99 << " / " << p90 << " / " << p50 << " / "
+              << max_util << " / " << mean_util << std::endl;
+    std::cout << "eviction-trigger samples: " << util_ratios.size()
+              << std::endl;
+    std::cout << "avg alloc time: " << avg_time_ns << " ns/op" << std::endl;
+}
+
 int main() {
     std::cout << "=== OffsetAllocator Benchmark ===" << std::endl;
     uniform_size_allocation_benchmark<OffsetAllocatorBenchHelper>();
     random_size_allocation_benchmark<OffsetAllocatorBenchHelper>();
-}
\ No newline at end of file
+    paired_kv_indexer_allocation_benchmark<OffsetAllocatorBenchHelper>();
+}
diff --git a/mooncake-store/benchmarks/batch_get_replica_bench.cpp b/mooncake-store/benchmarks/batch_get_replica_bench.cpp
new file mode 100644
index 00000000..60988592
--- /dev/null
+++ b/mooncake-store/benchmarks/batch_get_replica_bench.cpp
@@ -0,0 +1,717 @@
+// Copyright 2026 Alibaba Cloud and its affiliates
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include <algorithm>
+#include <atomic>
+#include <chrono>
+#include <cmath>
+#include <cstdint>
+#include <cstring>
+#include <fstream>
+#include <future>
+#include <iomanip>
+#include <iostream>
+#include <latch>
+#include <limits>
+#include <numeric>
+#include <random>
+#include <sstream>
+#include <string>
+#include <thread>
+#include <vector>
+
+#include <pthread.h>
+
+#include "gflags/gflags.h"
+#include "glog/logging.h"
+
+#include "master_client.h"
+
+static constexpr uint64_t KiB = 1024;
+static constexpr uint64_t MiB = 1024 * KiB;
+static constexpr uint64_t GiB = 1024 * MiB;
+static constexpr uintptr_t kSegmentBase = 0x100000000ULL;
+
+DEFINE_string(master_server, "127.0.0.1:50051", "Master server address");
+DEFINE_string(label, "run", "Label written to text and CSV output");
+DEFINE_string(csv_path, "", "Optional CSV output path");
+DEFINE_string(key_prefix, "batchget_bench", "Object key prefix");
+DEFINE_uint64(num_segments, 1, "Number of segments to mount");
+DEFINE_uint64(segment_size, 48 * GiB, "Size of each mounted segment");
+DEFINE_uint64(num_objects, 20 * 1000 * 1000,
+              "Number of objects to prefill and sample from");
+DEFINE_uint64(prefill_threads, 8, "Number of concurrent prefill threads");
+DEFINE_uint64(prefill_batch_size, 30000, "Batch size used by prefill writes");
+DEFINE_uint64(lookup_threads, 56,
+              "Number of concurrent BatchGetReplicaList threads");
+DEFINE_uint64(batch_size, 30000, "Keys per BatchGetReplicaList request");
+DEFINE_uint64(value_size, 2048, "Object value size used for master allocation");
+DEFINE_uint64(duration, 120, "Measured duration per cycle in seconds");
+DEFINE_uint64(cycles, 2, "Number of measured lookup cycles");
+DEFINE_uint64(requests_per_thread, 0,
+              "Fixed BatchGetReplicaList requests per lookup thread per "
+              "cycle. 0 uses duration-based cycles");
+DEFINE_uint64(lookup_interval_ms, 500,
+              "Sleep interval after each lookup request per thread");
+DEFINE_uint64(refill_threads, 0,
+              "Optional writer threads running during measured lookup");
+DEFINE_uint64(refill_batch_size, 30000,
+              "Batch size used by optional refill writer threads");
+DEFINE_uint64(refill_interval_ms, 0,
+              "Sleep interval after each refill write batch");
+DEFINE_uint64(random_seed, 1, "Base RNG seed for lookup threads");
+DEFINE_bool(skip_prefill, false,
+            "Skip prefill and assume keys already exist on the master");
+DEFINE_bool(log_progress, true, "Log prefill/refill progress");
+
+static inline void unset_cpu_affinity() {
+    cpu_set_t cpuset;
+    memset(&cpuset, -1, sizeof(cpuset));
+    pthread_setaffinity_np(pthread_self(), sizeof(cpuset), &cpuset);
+}
+
+class SegmentClient {
+   public:
+    SegmentClient(const std::string& name, const std::string& master_server,
+                  uintptr_t segment_base, uint64_t segment_size)
+        : master_client_(mooncake::generate_uuid()) {
+        auto ec = master_client_.Connect(master_server);
+        if (ec != mooncake::ErrorCode::OK) {
+            throw std::runtime_error("Cannot connect to master server at " +
+                                     master_server + ", ec=" + toString(ec));
+        }
+
+        segment_.id = mooncake::generate_uuid();
+        segment_.name = name;
+        segment_.base = segment_base;
+        segment_.size = segment_size;
+        segment_.te_endpoint = name;
+        auto mount_ec = master_client_.MountSegment(segment_);
+        if (!mount_ec.has_value()) {
+            throw std::runtime_error("Failed to mount segment " + name +
+                                     ", ec=" + toString(mount_ec.error()));
+        }
+    }
+
+    ~SegmentClient() {
+        if (remount_future_.valid()) {
+            remount_future_.wait();
+        }
+
+        auto unmount_result = master_client_.UnmountSegment(segment_.id);
+        if (!unmount_result.has_value()) {
+            LOG(ERROR) << "Failed to unmount segment " << segment_.name;
+        }
+    }
+
+    void Ping() {
+        if (remount_future_.valid() &&
+            remount_future_.wait_for(std::chrono::seconds(0)) ==
+                std::future_status::ready) {
+            remount_future_.get();
+            remount_future_ = std::future<void>();
+        }
+
+        auto ping_result = master_client_.Ping();
+        if (!ping_result.has_value()) {
+            throw std::runtime_error("Failed to ping master server");
+        }
+
+        if (ping_result.value().client_status ==
+                mooncake::ClientStatus::NEED_REMOUNT &&
+            !remount_future_.valid()) {
+            remount_future_ = std::async(std::launch::async, [&]() {
+                auto remount_ec = master_client_.ReMountSegment({segment_});
+                if (!remount_ec.has_value()) {
+                    throw std::runtime_error("Failed to remount segment");
+                }
+            });
+        }
+    }
+
+   private:
+    mooncake::MasterClient master_client_;
+    mooncake::Segment segment_;
+    std::future<void> remount_future_;
+};
+
+struct WriteStats {
+    uint64_t started = 0;
+    uint64_t completed = 0;
+    uint64_t start_failed = 0;
+    uint64_t end_failed = 0;
+
+    void Add(const WriteStats& other) {
+        started += other.started;
+        completed += other.completed;
+        start_failed += other.start_failed;
+        end_failed += other.end_failed;
+    }
+};
+
+struct ThreadStats {
+    std::vector<double> latency_ms;
+    uint64_t requests = 0;
+    uint64_t keys = 0;
+    uint64_t success_keys = 0;
+    uint64_t object_not_found = 0;
+    uint64_t replica_not_ready = 0;
+    uint64_t other_failures = 0;
+    uint64_t bad_responses = 0;
+
+    void Add(ThreadStats&& other) {
+        requests += other.requests;
+        keys += other.keys;
+        success_keys += other.success_keys;
+        object_not_found += other.object_not_found;
+        replica_not_ready += other.replica_not_ready;
+        other_failures += other.other_failures;
+        bad_responses += other.bad_responses;
+        latency_ms.insert(latency_ms.end(),
+                          std::make_move_iterator(other.latency_ms.begin()),
+                          std::make_move_iterator(other.latency_ms.end()));
+    }
+};
+
+struct ResultRow {
+    std::string label;
+    std::string cycle;
+    uint64_t requests = 0;
+    uint64_t keys = 0;
+    double duration_sec = 0;
+    double batch_per_sec = 0;
+    double key_per_sec = 0;
+    double avg_ms = 0;
+    double p50_ms = 0;
+    double p90_ms = 0;
+    double p99_ms = 0;
+    double max_ms = 0;
+    double object_not_found_rate = 0;
+    uint64_t success_keys = 0;
+    uint64_t object_not_found = 0;
+    uint64_t replica_not_ready = 0;
+    uint64_t other_failures = 0;
+    uint64_t bad_responses = 0;
+};
+
+std::string MakeKey(uint64_t id) {
+    std::ostringstream ss;
+    ss << FLAGS_key_prefix << "_" << std::setw(16) << std::setfill('0') << id;
+    return ss.str();
+}
+
+std::vector<std::string> MakeSequentialKeys(uint64_t start_id, uint64_t count) {
+    std::vector<std::string> keys;
+    keys.reserve(count);
+    for (uint64_t i = 0; i < count; ++i) {
+        keys.push_back(MakeKey(start_id + i));
+    }
+    return keys;
+}
+
+std::vector<std::vector<uint64_t>> MakeSliceLengths(size_t count) {
+    return std::vector<std::vector<uint64_t>>(
+        count, std::vector<uint64_t>{FLAGS_value_size});
+}
+
+WriteStats PutCompletedBatch(mooncake::MasterClient& client,
+                             const std::vector<std::string>& keys) {
+    WriteStats stats;
+    const mooncake::ReplicateConfig config;
+    auto put_start_result =
+        client.BatchPutStart(keys, MakeSliceLengths(keys.size()), config);
+
+    std::vector<std::string> started_keys;
+    started_keys.reserve(keys.size());
+    for (size_t i = 0; i < keys.size(); ++i) {
+        if (i < put_start_result.size() && put_start_result[i].has_value()) {
+            started_keys.push_back(keys[i]);
+            ++stats.started;
+        } else {
+            ++stats.start_failed;
+        }
+    }
+
+    if (started_keys.empty()) {
+        return stats;
+    }
+
+    auto put_end_result =
+        client.BatchPutEnd(started_keys, mooncake::ReplicaType::MEMORY);
+    for (const auto& result : put_end_result) {
+        if (result.has_value()) {
+            ++stats.completed;
+        } else {
+            ++stats.end_failed;
+        }
+    }
+    if (put_end_result.size() < started_keys.size()) {
+        stats.end_failed += started_keys.size() - put_end_result.size();
+    }
+    return stats;
+}
+
+WriteStats PrefillObjects() {
+    std::atomic<uint64_t> next_id{0};
+    std::atomic<uint64_t> completed{0};
+    std::atomic<uint64_t> started{0};
+    std::atomic<uint64_t> start_failed{0};
+    std::atomic<uint64_t> end_failed{0};
+
+    const uint64_t thread_count = std::max<uint64_t>(1, FLAGS_prefill_threads);
+    std::vector<std::thread> threads;
+    threads.reserve(thread_count);
+
+    auto started_at = std::chrono::steady_clock::now();
+    for (uint64_t thread_idx = 0; thread_idx < thread_count; ++thread_idx) {
+        threads.emplace_back([&, thread_idx] {
+            unset_cpu_affinity();
+            mooncake::MasterClient client(mooncake::generate_uuid());
+            auto ec = client.Connect(FLAGS_master_server);
+            if (ec != mooncake::ErrorCode::OK) {
+                throw std::runtime_error("Prefill client cannot connect, ec=" +
+                                         toString(ec));
+            }
+
+            uint64_t next_log = 100000;
+            while (true) {
+                const uint64_t start_id =
+                    next_id.fetch_add(FLAGS_prefill_batch_size);
+                if (start_id >= FLAGS_num_objects) {
+                    break;
+                }
+                const uint64_t count = std::min<uint64_t>(
+                    FLAGS_prefill_batch_size, FLAGS_num_objects - start_id);
+                auto stats = PutCompletedBatch(
+                    client, MakeSequentialKeys(start_id, count));
+                started.fetch_add(stats.started);
+                completed.fetch_add(stats.completed);
+                start_failed.fetch_add(stats.start_failed);
+                end_failed.fetch_add(stats.end_failed);
+
+                if (FLAGS_log_progress && thread_idx == 0 &&
+                    completed.load() >= next_log) {
+                    const auto elapsed =
+                        std::chrono::duration<double>(
+                            std::chrono::steady_clock::now() - started_at)
+                            .count();
+                    LOG(INFO)
+                        << "Prefill progress: completed=" << completed.load()
+                        << "/" << FLAGS_num_objects << ", elapsed=" << elapsed
+                        << "s";
+                    next_log += 100000;
+                }
+            }
+        });
+    }
+
+    for (auto& thread : threads) {
+        thread.join();
+    }
+
+    return {.started = started.load(),
+            .completed = completed.load(),
+            .start_failed = start_failed.load(),
+            .end_failed = end_failed.load()};
+}
+
+void RunRefillWorker(std::atomic<bool>& running, std::atomic<uint64_t>& next_id,
+                     WriteStats& stats) {
+    unset_cpu_affinity();
+    mooncake::MasterClient client(mooncake::generate_uuid());
+    auto ec = client.Connect(FLAGS_master_server);
+    if (ec != mooncake::ErrorCode::OK) {
+        throw std::runtime_error("Refill client cannot connect, ec=" +
+                                 toString(ec));
+    }
+
+    while (running.load(std::memory_order_relaxed)) {
+        const uint64_t start_id = next_id.fetch_add(FLAGS_refill_batch_size);
+        auto batch_stats = PutCompletedBatch(
+            client, MakeSequentialKeys(start_id, FLAGS_refill_batch_size));
+        stats.Add(batch_stats);
+        if (FLAGS_refill_interval_ms > 0) {
+            std::this_thread::sleep_for(
+                std::chrono::milliseconds(FLAGS_refill_interval_ms));
+        }
+    }
+}
+
+ThreadStats RunLookupWorker(uint64_t thread_idx,
+                            std::chrono::steady_clock::time_point deadline,
+                            std::shared_ptr<std::latch> barrier) {
+    unset_cpu_affinity();
+    ThreadStats stats;
+    stats.latency_ms.reserve(
+        FLAGS_lookup_interval_ms > 0
+            ? std::max<uint64_t>(
+                  1, FLAGS_duration * 1000 / FLAGS_lookup_interval_ms + 1)
+            : 1024);
+
+    mooncake::MasterClient client(mooncake::generate_uuid());
+    auto ec = client.Connect(FLAGS_master_server);
+    if (ec != mooncake::ErrorCode::OK) {
+        throw std::runtime_error("Lookup client cannot connect, ec=" +
+                                 toString(ec));
+    }
+
+    std::mt19937_64 rng(FLAGS_random_seed + thread_idx);
+    std::uniform_int_distribution<uint64_t> key_dist(0, FLAGS_num_objects - 1);
+
+    barrier->arrive_and_wait();
+    uint64_t request_idx = 0;
+    while (true) {
+        const bool fixed_request_mode = FLAGS_requests_per_thread > 0;
+        if (fixed_request_mode) {
+            if (request_idx >= FLAGS_requests_per_thread) {
+                break;
+            }
+        } else if (std::chrono::steady_clock::now() >= deadline) {
+            break;
+        }
+
+        std::vector<std::string> keys;
+        keys.reserve(FLAGS_batch_size);
+        for (uint64_t i = 0; i < FLAGS_batch_size; ++i) {
+            keys.push_back(MakeKey(key_dist(rng)));
+        }
+
+        const auto started_at = std::chrono::steady_clock::now();
+        auto results = client.BatchGetReplicaList(keys);
+        const auto finished_at = std::chrono::steady_clock::now();
+        const double latency_ms =
+            std::chrono::duration<double, std::milli>(finished_at - started_at)
+                .count();
+
+        stats.latency_ms.push_back(latency_ms);
+        ++stats.requests;
+        ++request_idx;
+        stats.keys += keys.size();
+
+        if (results.size() != keys.size()) {
+            ++stats.bad_responses;
+            stats.other_failures += keys.size();
+        } else {
+            for (const auto& result : results) {
+                if (result.has_value()) {
+                    ++stats.success_keys;
+                    continue;
+                }
+                switch (result.error()) {
+                    case mooncake::ErrorCode::OBJECT_NOT_FOUND:
+                        ++stats.object_not_found;
+                        break;
+                    case mooncake::ErrorCode::REPLICA_IS_NOT_READY:
+                        ++stats.replica_not_ready;
+                        break;
+                    default:
+                        ++stats.other_failures;
+                        break;
+                }
+            }
+        }
+
+        if (FLAGS_lookup_interval_ms > 0) {
+            const auto interval =
+                std::chrono::milliseconds(FLAGS_lookup_interval_ms);
+            if (fixed_request_mode) {
+                std::this_thread::sleep_until(finished_at + interval);
+            } else {
+                std::this_thread::sleep_until(
+                    std::min(finished_at + interval, deadline));
+            }
+        }
+    }
+    return stats;
+}
+
+double Percentile(const std::vector<double>& sorted_values, double percentile) {
+    if (sorted_values.empty()) {
+        return 0;
+    }
+    if (sorted_values.size() == 1) {
+        return sorted_values.front();
+    }
+    const double rank = percentile * (sorted_values.size() - 1);
+    const auto low = static_cast<size_t>(std::floor(rank));
+    const auto high = static_cast<size_t>(std::ceil(rank));
+    if (low == high) {
+        return sorted_values[low];
+    }
+    return sorted_values[low] +
+           (sorted_values[high] - sorted_values[low]) * (rank - low);
+}
+
+ResultRow BuildResultRow(const std::string& cycle, ThreadStats&& stats,
+                         double duration_sec) {
+    std::sort(stats.latency_ms.begin(), stats.latency_ms.end());
+    const double sum_ms =
+        std::accumulate(stats.latency_ms.begin(), stats.latency_ms.end(), 0.0);
+    ResultRow row;
+    row.label = FLAGS_label;
+    row.cycle = cycle;
+    row.requests = stats.requests;
+    row.keys = stats.keys;
+    row.duration_sec = duration_sec;
+    row.batch_per_sec = duration_sec > 0
+                            ? static_cast<double>(stats.requests) / duration_sec
+                            : 0;
+    row.key_per_sec =
+        duration_sec > 0 ? static_cast<double>(stats.keys) / duration_sec : 0;
+    row.avg_ms =
+        stats.latency_ms.empty() ? 0 : sum_ms / stats.latency_ms.size();
+    row.p50_ms = Percentile(stats.latency_ms, 0.50);
+    row.p90_ms = Percentile(stats.latency_ms, 0.90);
+    row.p99_ms = Percentile(stats.latency_ms, 0.99);
+    row.max_ms = stats.latency_ms.empty() ? 0 : stats.latency_ms.back();
+    row.object_not_found_rate =
+        stats.keys > 0 ? static_cast<double>(stats.object_not_found) /
+                             static_cast<double>(stats.keys)
+                       : 0;
+    row.success_keys = stats.success_keys;
+    row.object_not_found = stats.object_not_found;
+    row.replica_not_ready = stats.replica_not_ready;
+    row.other_failures = stats.other_failures;
+    row.bad_responses = stats.bad_responses;
+    return row;
+}
+
+ResultRow RunLookupCycle(uint64_t cycle_idx) {
+    std::atomic<bool> refill_running{true};
+    std::atomic<uint64_t> next_refill_id{FLAGS_num_objects +
+                                         cycle_idx * 1000000000ULL};
+    std::vector<WriteStats> refill_stats(FLAGS_refill_threads);
+    std::vector<std::thread> refill_threads;
+    refill_threads.reserve(FLAGS_refill_threads);
+    for (uint64_t i = 0; i < FLAGS_refill_threads; ++i) {
+        refill_threads.emplace_back(RunRefillWorker, std::ref(refill_running),
+                                    std::ref(next_refill_id),
+                                    std::ref(refill_stats[i]));
+    }
+
+    const auto started_at = std::chrono::steady_clock::now();
+    const auto deadline = started_at + std::chrono::seconds(FLAGS_duration);
+    auto barrier = std::make_shared<std::latch>(FLAGS_lookup_threads + 1);
+    std::vector<ThreadStats> per_thread_stats(FLAGS_lookup_threads);
+    std::vector<std::thread> lookup_threads;
+    lookup_threads.reserve(FLAGS_lookup_threads);
+
+    for (uint64_t i = 0; i < FLAGS_lookup_threads; ++i) {
+        lookup_threads.emplace_back([&, i] {
+            per_thread_stats[i] = RunLookupWorker(i, deadline, barrier);
+        });
+    }
+
+    barrier->arrive_and_wait();
+    for (auto& thread : lookup_threads) {
+        thread.join();
+    }
+    const auto finished_at = std::chrono::steady_clock::now();
+
+    refill_running.store(false, std::memory_order_relaxed);
+    for (auto& thread : refill_threads) {
+        thread.join();
+    }
+
+    if (FLAGS_refill_threads > 0) {
+        WriteStats total_refill;
+        for (const auto& stats : refill_stats) {
+            total_refill.Add(stats);
+        }
+        LOG(INFO) << "cycle=" << cycle_idx
+                  << " refill completed=" << total_refill.completed
+                  << ", start_failed=" << total_refill.start_failed
+                  << ", end_failed=" << total_refill.end_failed;
+    }
+
+    ThreadStats merged;
+    for (auto& stats : per_thread_stats) {
+        merged.Add(std::move(stats));
+    }
+    const double duration_sec =
+        std::chrono::duration<double>(finished_at - started_at).count();
+    return BuildResultRow(std::to_string(cycle_idx), std::move(merged),
+                          duration_sec);
+}
+
+ThreadStats RowsAsStats(const std::vector<ResultRow>& rows) {
+    ThreadStats stats;
+    for (const auto& row : rows) {
+        stats.requests += row.requests;
+        stats.keys += row.keys;
+        stats.success_keys += row.success_keys;
+        stats.object_not_found += row.object_not_found;
+        stats.replica_not_ready += row.replica_not_ready;
+        stats.other_failures += row.other_failures;
+        stats.bad_responses += row.bad_responses;
+    }
+    return stats;
+}
+
+void PrintRows(const std::vector<ResultRow>& rows) {
+    std::cout << std::fixed << std::setprecision(3);
+    std::cout << "\nBatchGetReplicaList latency summary\n";
+    std::cout << "label=" << FLAGS_label << ", objects=" << FLAGS_num_objects
+              << ", lookup_threads=" << FLAGS_lookup_threads
+              << ", batch_size=" << FLAGS_batch_size
+              << ", lookup_interval_ms=" << FLAGS_lookup_interval_ms
+              << ", requests_per_thread=" << FLAGS_requests_per_thread
+              << ", refill_threads=" << FLAGS_refill_threads << "\n";
+    std::cout << std::left << std::setw(8) << "cycle" << std::right
+              << std::setw(10) << "requests" << std::setw(14) << "keys/s"
+              << std::setw(12) << "batch/s" << std::setw(12) << "avg(ms)"
+              << std::setw(12) << "p50(ms)" << std::setw(12) << "p90(ms)"
+              << std::setw(12) << "p99(ms)" << std::setw(12) << "max(ms)"
+              << std::setw(12) << "miss(%)" << std::setw(14) << "not_found"
+              << "\n";
+    for (const auto& row : rows) {
+        std::cout << std::left << std::setw(8) << row.cycle << std::right
+                  << std::setw(10) << row.requests << std::setw(14)
+                  << row.key_per_sec << std::setw(12) << row.batch_per_sec
+                  << std::setw(12) << row.avg_ms << std::setw(12) << row.p50_ms
+                  << std::setw(12) << row.p90_ms << std::setw(12) << row.p99_ms
+                  << std::setw(12) << row.max_ms << std::setw(14)
+                  << row.object_not_found_rate * 100.0 << std::setw(14)
+                  << row.object_not_found << "\n";
+    }
+}
+
+void WriteCsv(const std::vector<ResultRow>& rows) {
+    if (FLAGS_csv_path.empty()) {
+        return;
+    }
+
+    const bool append_header = !std::ifstream(FLAGS_csv_path).good() ||
+                               std::ifstream(FLAGS_csv_path).peek() ==
+                                   std::ifstream::traits_type::eof();
+    std::ofstream out(FLAGS_csv_path, std::ios::app);
+    if (!out) {
+        throw std::runtime_error("Cannot open CSV output path: " +
+                                 FLAGS_csv_path);
+    }
+
+    if (append_header) {
+        out << "label,cycle,requests,total_keys,duration_sec,batch_per_sec,"
+               "key_per_sec,avg_ms,p50_ms,p90_ms,p99_ms,max_ms,"
+               "object_not_found_rate,success_keys,object_not_found,"
+               "replica_not_ready,other_failures,bad_responses\n";
+    }
+
+    out << std::fixed << std::setprecision(6);
+    for (const auto& row : rows) {
+        out << row.label << "," << row.cycle << "," << row.requests << ","
+            << row.keys << "," << row.duration_sec << "," << row.batch_per_sec
+            << "," << row.key_per_sec << "," << row.avg_ms << "," << row.p50_ms
+            << "," << row.p90_ms << "," << row.p99_ms << "," << row.max_ms
+            << "," << row.object_not_found_rate << "," << row.success_keys
+            << "," << row.object_not_found << "," << row.replica_not_ready
+            << "," << row.other_failures << "," << row.bad_responses << "\n";
+    }
+}
+
+void ValidateFlags() {
+    if (FLAGS_num_objects == 0 || FLAGS_batch_size == 0 ||
+        FLAGS_lookup_threads == 0 || FLAGS_cycles == 0 ||
+        FLAGS_prefill_batch_size == 0 || FLAGS_refill_batch_size == 0) {
+        throw std::invalid_argument(
+            "num_objects, batch_size, lookup_threads, cycles, "
+            "prefill_batch_size, and refill_batch_size must be positive");
+    }
+    if (FLAGS_requests_per_thread == 0 && FLAGS_duration == 0) {
+        throw std::invalid_argument(
+            "duration must be positive when requests_per_thread is 0");
+    }
+}
+
+int main(int argc, char** argv) {
+    google::InitGoogleLogging("BatchGetReplicaBench");
+    FLAGS_logtostderr = true;
+    gflags::ParseCommandLineFlags(&argc, &argv, false);
+
+    try {
+        ValidateFlags();
+
+        std::vector<std::unique_ptr<SegmentClient>> segment_clients;
+        std::mutex segment_clients_mutex;
+        std::jthread ping_thread([&](std::stop_token stop_token) {
+            static const auto OneSecond = std::chrono::seconds(1);
+            unset_cpu_affinity();
+            while (!stop_token.stop_requested()) {
+                const auto started_at = std::chrono::steady_clock::now();
+                {
+                    std::lock_guard<std::mutex> guard(segment_clients_mutex);
+                    for (auto& segment_client : segment_clients) {
+                        segment_client->Ping();
+                    }
+                }
+                const auto elapsed =
+                    std::chrono::steady_clock::now() - started_at;
+                if (elapsed < OneSecond) {
+                    std::this_thread::sleep_for(OneSecond - elapsed);
+                }
+            }
+        });
+
+        LOG(INFO) << "Mounting " << FLAGS_num_segments << " segments";
+        for (uint64_t i = 0; i < FLAGS_num_segments; ++i) {
+            auto segment_client = std::make_unique<SegmentClient>(
+                "batch_get_bench_segment_" + std::to_string(i),
+                FLAGS_master_server, kSegmentBase + i * FLAGS_segment_size,
+                FLAGS_segment_size);
+            std::lock_guard<std::mutex> guard(segment_clients_mutex);
+            segment_clients.push_back(std::move(segment_client));
+        }
+
+        if (!FLAGS_skip_prefill) {
+            LOG(INFO) << "Prefilling objects: objects=" << FLAGS_num_objects
+                      << ", threads=" << FLAGS_prefill_threads
+                      << ", batch_size=" << FLAGS_prefill_batch_size
+                      << ", value_size=" << FLAGS_value_size;
+            const auto prefill_started_at = std::chrono::steady_clock::now();
+            const auto prefill_stats = PrefillObjects();
+            const auto prefill_sec =
+                std::chrono::duration<double>(std::chrono::steady_clock::now() -
+                                              prefill_started_at)
+                    .count();
+            LOG(INFO) << "Prefill done: completed=" << prefill_stats.completed
+                      << ", started=" << prefill_stats.started
+                      << ", start_failed=" << prefill_stats.start_failed
+                      << ", end_failed=" << prefill_stats.end_failed
+                      << ", elapsed=" << prefill_sec << "s";
+        }
+
+        std::vector<ResultRow> rows;
+        rows.reserve(FLAGS_cycles);
+        for (uint64_t cycle = 1; cycle <= FLAGS_cycles; ++cycle) {
+            LOG(INFO) << "Starting lookup cycle " << cycle << "/"
+                      << FLAGS_cycles;
+            rows.push_back(RunLookupCycle(cycle));
+        }
+
+        PrintRows(rows);
+        WriteCsv(rows);
+
+        if (ping_thread.joinable()) {
+            ping_thread.request_stop();
+            ping_thread.join();
+        }
+        segment_clients.clear();
+        google::ShutdownGoogleLogging();
+        return 0;
+    } catch (const std::exception& ex) {
+        LOG(ERROR) << "benchmark failed: " << ex.what();
+        google::ShutdownGoogleLogging();
+        return 1;
+    }
+}
diff --git a/mooncake-store/benchmarks/nof_worker_pool_bench.cpp b/mooncake-store/benchmarks/nof_worker_pool_bench.cpp
new file mode 100644
index 00000000..ed3d011e
--- /dev/null
+++ b/mooncake-store/benchmarks/nof_worker_pool_bench.cpp
@@ -0,0 +1,1019 @@
+#include <gflags/gflags.h>
+#include <glog/logging.h>
+
+#include <algorithm>
+#include <array>
+#include <atomic>
+#include <chrono>
+#include <cctype>
+#include <cmath>
+#include <cstdint>
+#include <cstdlib>
+#include <cstring>
+#include <iomanip>
+#include <iostream>
+#include <limits>
+#include <memory>
+#include <random>
+#include <sstream>
+#include <string>
+#include <thread>
+#include <unordered_set>
+#include <vector>
+
+#include "spdk/spdk_wrapper.h"
+#include "transfer_task.h"
+#include "utils.h"
+
+namespace {
+
+constexpr uint64_t KiB = 1024;
+constexpr uint64_t MiB = 1024 * KiB;
+constexpr uint64_t GiB = 1024 * MiB;
+constexpr uint64_t kDefaultIoSize = 128 * KiB;
+constexpr uint64_t kDefaultIodepth = 64;
+constexpr uint64_t kDefaultSubmitThreads = 1;
+constexpr uint64_t kDefaultDurationSec = 30;
+constexpr uint64_t kDefaultReportIntervalMs = 1000;
+constexpr uint64_t kDefaultRangeBytes = 1 * GiB;
+constexpr uint64_t kDefaultWarmupSec = 3;
+constexpr char kDefaultOp[] = "read";
+
+DEFINE_string(
+    endpoints, "",
+    "Comma-separated NoF transport endpoints. To use multiple "
+    "SpdkNofWorkerPool workers with the current implementation, provide "
+    "multiple distinct endpoints/namespaces/targets.");
+DEFINE_string(op, "read", "Benchmark operation: read, write, or mixed.");
+DEFINE_string(
+    rw, "",
+    "fio-style alias for --op/--random_lba. Supported values: read, write, "
+    "randread, randwrite, mixed, randrw.");
+DEFINE_uint64(io_size, 128 * KiB,
+              "I/O size in bytes. Must align to device block size.");
+DEFINE_uint64(bs, 0, "fio-style alias for --io_size.");
+DEFINE_uint64(
+    iodepth, 64,
+    "Per-endpoint in-flight I/Os maintained by the benchmark. Total in-flight "
+    "depth is endpoints * iodepth.");
+DEFINE_uint64(
+    submit_threads, 1,
+    "Number of benchmark submit/completion threads. Endpoints are evenly "
+    "partitioned across submit threads.");
+DEFINE_uint64(numjobs, 0, "fio-style alias for --submit_threads.");
+DEFINE_uint64(duration_sec, 30, "Measurement duration in seconds.");
+DEFINE_uint64(runtime, 0, "fio-style alias for --duration_sec.");
+DEFINE_uint64(report_interval_ms, 1000,
+              "Periodic throughput report interval in milliseconds.");
+DEFINE_uint64(status_interval, 0,
+              "fio-style alias for --report_interval_ms, in seconds.");
+DEFINE_uint64(start_lba, 0, "Starting LBA offset within each endpoint.");
+DEFINE_uint64(range_bytes, 1 * GiB,
+              "Per-endpoint logical range in bytes used by the benchmark. "
+              "Must align to block size.");
+DEFINE_uint64(size, 0, "fio-style alias for --range_bytes.");
+DEFINE_bool(random_lba, false,
+            "Whether to choose LBA randomly inside the per-endpoint range.");
+DEFINE_uint64(seed, 1, "Random seed used when --random_lba or --op=mixed.");
+DEFINE_uint64(warmup_sec, 3,
+              "Warmup duration in seconds before measurement starts.");
+DEFINE_uint64(ramp_time, 0, "fio-style alias for --warmup_sec.");
+DEFINE_uint64(
+    nof_workers, 0,
+    "Optional override for MC_NOF_WORKERS. Equivalent to exporting "
+    "MC_NOF_WORKERS before launch. 0 means keep environment/default.");
+DEFINE_uint64(nof_submit_chunk_bytes, 0,
+              "Optional override for MC_NOF_SUBMIT_CHUNK_BYTES. Equivalent to "
+              "exporting MC_NOF_SUBMIT_CHUNK_BYTES before launch. 0 means keep "
+              "environment/default.");
+DEFINE_uint64(
+    nof_inflight_bytes_limit, 0,
+    "Optional override for MC_NOF_INFLIGHT_BYTES_LIMIT. Equivalent to "
+    "exporting MC_NOF_INFLIGHT_BYTES_LIMIT before launch. 0 means keep "
+    "environment/default.");
+DEFINE_int32(socket_id, -1,
+             "NUMA socket for benchmark buffers. -1 lets SPDK choose.");
+DEFINE_bool(fill_on_write, true,
+            "Fill buffers with a deterministic pattern for write/mixed ops.");
+
+using Clock = std::chrono::steady_clock;
+using TimePoint = Clock::time_point;
+
+enum class BenchOp {
+    READ = 0,
+    WRITE = 1,
+    MIXED = 2,
+};
+
+BenchOp ParseBenchOp(const std::string &value) {
+    std::string normalized(value);
+    std::transform(
+        normalized.begin(), normalized.end(), normalized.begin(),
+        [](unsigned char c) { return static_cast<char>(std::tolower(c)); });
+    if (normalized == "read") {
+        return BenchOp::READ;
+    }
+    if (normalized == "write") {
+        return BenchOp::WRITE;
+    }
+    if (normalized == "mixed") {
+        return BenchOp::MIXED;
+    }
+    throw std::invalid_argument(
+        "Invalid --op. Supported values: read, write, mixed");
+}
+
+struct ParsedRw {
+    std::string op;
+    bool random_lba{false};
+};
+
+ParsedRw ParseRwMode(const std::string &value) {
+    std::string normalized(value);
+    std::transform(
+        normalized.begin(), normalized.end(), normalized.begin(),
+        [](unsigned char c) { return static_cast<char>(std::tolower(c)); });
+    if (normalized == "read") {
+        return {"read", false};
+    }
+    if (normalized == "write") {
+        return {"write", false};
+    }
+    if (normalized == "mixed" || normalized == "rw") {
+        return {"mixed", false};
+    }
+    if (normalized == "randread") {
+        return {"read", true};
+    }
+    if (normalized == "randwrite") {
+        return {"write", true};
+    }
+    if (normalized == "randrw") {
+        return {"mixed", true};
+    }
+    throw std::invalid_argument(
+        "Invalid --rw. Supported values: read, write, randread, randwrite, "
+        "mixed, rw, randrw");
+}
+
+void ResolveUint64Alias(const char *legacy_name, uint64_t *legacy_value,
+                        uint64_t legacy_default, const char *alias_name,
+                        uint64_t alias_value) {
+    if (alias_value == 0) {
+        return;
+    }
+    if (*legacy_value != legacy_default && *legacy_value != alias_value) {
+        throw std::invalid_argument(std::string("Conflicting --") +
+                                    legacy_name + " and --" + alias_name +
+                                    " values");
+    }
+    *legacy_value = alias_value;
+}
+
+void SetEnvU64IfRequested(const char *name, uint64_t value) {
+    if (value == 0) {
+        return;
+    }
+#ifdef _WIN32
+    _putenv_s(name, std::to_string(value).c_str());
+#else
+    setenv(name, std::to_string(value).c_str(), 1);
+#endif
+}
+
+std::string FormatBytesPerSecond(double bytes_per_sec) {
+    std::ostringstream oss;
+    oss << std::fixed << std::setprecision(2);
+    if (bytes_per_sec >= static_cast<double>(GiB)) {
+        oss << (bytes_per_sec / GiB) << " GiB/s";
+    } else if (bytes_per_sec >= static_cast<double>(MiB)) {
+        oss << (bytes_per_sec / MiB) << " MiB/s";
+    } else if (bytes_per_sec >= static_cast<double>(KiB)) {
+        oss << (bytes_per_sec / KiB) << " KiB/s";
+    } else {
+        oss.unsetf(std::ios::fixed);
+        oss << bytes_per_sec << " B/s";
+    }
+    return oss.str();
+}
+
+std::string FormatLatencyNs(uint64_t latency_ns) {
+    std::ostringstream oss;
+    oss << std::fixed << std::setprecision(2);
+    if (latency_ns >= 1000ULL * 1000ULL * 1000ULL) {
+        oss << (static_cast<double>(latency_ns) / 1000.0 / 1000.0 / 1000.0)
+            << " s";
+    } else if (latency_ns >= 1000ULL * 1000ULL) {
+        oss << (static_cast<double>(latency_ns) / 1000.0 / 1000.0) << " ms";
+    } else if (latency_ns >= 1000ULL) {
+        oss << (static_cast<double>(latency_ns) / 1000.0) << " us";
+    } else {
+        oss.unsetf(std::ios::fixed);
+        oss << latency_ns << " ns";
+    }
+    return oss.str();
+}
+
+struct CounterSnapshot {
+    uint64_t submitted_ops{0};
+    uint64_t completed_ops{0};
+    uint64_t failed_ops{0};
+    uint64_t bytes{0};
+};
+
+struct LiveCounters {
+    std::atomic<uint64_t> submitted_ops{0};
+    std::atomic<uint64_t> completed_ops{0};
+    std::atomic<uint64_t> failed_ops{0};
+    std::atomic<uint64_t> bytes{0};
+
+    CounterSnapshot Snapshot() const {
+        return {
+            submitted_ops.load(std::memory_order_relaxed),
+            completed_ops.load(std::memory_order_relaxed),
+            failed_ops.load(std::memory_order_relaxed),
+            bytes.load(std::memory_order_relaxed),
+        };
+    }
+};
+
+class LatencyHistogram {
+   public:
+    static constexpr size_t kBuckets0 = 200;
+    static constexpr size_t kBuckets1 = 180;
+    static constexpr size_t kBuckets2 = 180;
+    static constexpr size_t kBuckets3 = 180;
+    static constexpr size_t kBucketCount =
+        kBuckets0 + kBuckets1 + kBuckets2 + kBuckets3 + 1;
+
+    void Record(uint64_t latency_ns);
+    void MergeFrom(const LatencyHistogram &other);
+    uint64_t Count() const { return count_; }
+    uint64_t SumNs() const { return sum_ns_; }
+    uint64_t PercentileNs(double percentile) const;
+
+   private:
+    static size_t BucketIndex(uint64_t latency_ns);
+    static uint64_t BucketUpperBoundNs(size_t index);
+
+    uint64_t count_{0};
+    uint64_t sum_ns_{0};
+    std::array<uint64_t, kBucketCount> buckets_{};
+};
+
+struct BenchStats {
+    uint64_t submitted_ops{0};
+    uint64_t completed_ops{0};
+    uint64_t failed_ops{0};
+    uint64_t bytes{0};
+    LatencyHistogram latency_histogram;
+
+    void RecordSubmit();
+    void RecordSuccess(uint64_t io_bytes, uint64_t latency_ns);
+    void RecordFailure();
+    void MergeFrom(const BenchStats &other);
+};
+
+struct EndpointContext {
+    std::string endpoint;
+    mooncake::nof_seg_handle *seg_handle{nullptr};
+    uint32_t block_size{0};
+    uint64_t range_blocks{0};
+    uint64_t io_blocks{0};
+    uint64_t next_block{0};
+    std::mt19937_64 rng;
+
+    EndpointContext(std::string ep, uint64_t seed_value)
+        : endpoint(std::move(ep)), rng(seed_value) {}
+};
+
+struct EndpointStats {
+    BenchStats measured;
+    LiveCounters live;
+};
+
+struct Slot {
+    void *buffer{nullptr};
+    size_t buffer_size{0};
+    std::shared_ptr<mooncake::SpdkNofOperationState> state;
+    std::unique_ptr<mooncake::TransferFuture> future;
+    uint64_t io_bytes{0};
+    size_t endpoint_index{0};
+    bool active{false};
+    bool measure{false};
+    TimePoint submit_time{};
+};
+
+struct ThreadEndpointState {
+    size_t endpoint_index{0};
+    std::vector<Slot> slots;
+};
+
+struct BenchThreadContext {
+    size_t thread_index{0};
+    std::vector<ThreadEndpointState> endpoints;
+};
+
+struct ThroughputView {
+    double bandwidth_bytes_per_sec{0.0};
+    double iops{0.0};
+    double fail_iops{0.0};
+};
+
+struct LatencyView {
+    uint64_t avg_ns{0};
+    uint64_t p50_ns{0};
+    uint64_t p95_ns{0};
+    uint64_t p99_ns{0};
+};
+
+ThroughputView ComputeThroughput(const CounterSnapshot &counters,
+                                 double duration_sec);
+LatencyView ComputeLatency(const BenchStats &stats);
+CounterSnapshot DeltaCounters(const CounterSnapshot &current,
+                              const CounterSnapshot &previous);
+uint64_t NextLba(EndpointContext &endpoint, bool random_lba);
+int PickTaskOp(BenchOp op_mode, std::mt19937_64 &rng);
+void FillPattern(void *buffer, size_t size, uint64_t seq);
+void PublishSubmit(LiveCounters &live);
+void PublishSuccess(LiveCounters &live, uint64_t io_bytes);
+void PublishFailure(LiveCounters &live);
+void SubmitSlot(Slot &slot, EndpointContext &endpoint, size_t endpoint_index,
+                mooncake::SpdkNofWorkerPool &pool, BenchOp op_mode,
+                uint64_t submit_seq, bool measure_now,
+                EndpointStats &endpoint_stats);
+void WorkerLoop(BenchThreadContext *thread_context,
+                std::vector<EndpointContext> *endpoints,
+                std::vector<EndpointStats> *endpoint_stats,
+                mooncake::SpdkNofWorkerPool *pool, BenchOp op_mode,
+                TimePoint warmup_end, TimePoint end,
+                std::atomic<size_t> *active_threads);
+void FreeThreadSlots(mooncake::SpdkWrapper &wrapper,
+                     std::vector<BenchThreadContext> &thread_contexts);
+
+void LatencyHistogram::Record(uint64_t latency_ns) {
+    ++count_;
+    sum_ns_ += latency_ns;
+    ++buckets_[BucketIndex(latency_ns)];
+}
+
+void LatencyHistogram::MergeFrom(const LatencyHistogram &other) {
+    count_ += other.count_;
+    sum_ns_ += other.sum_ns_;
+    for (size_t i = 0; i < kBucketCount; ++i) {
+        buckets_[i] += other.buckets_[i];
+    }
+}
+
+uint64_t LatencyHistogram::PercentileNs(double percentile) const {
+    if (count_ == 0) {
+        return 0;
+    }
+    if (percentile <= 0.0) {
+        return BucketUpperBoundNs(0);
+    }
+    if (percentile >= 1.0) {
+        return BucketUpperBoundNs(kBucketCount - 1);
+    }
+
+    const uint64_t target =
+        static_cast<uint64_t>(std::ceil(percentile * count_));
+    uint64_t seen = 0;
+    for (size_t i = 0; i < kBucketCount; ++i) {
+        seen += buckets_[i];
+        if (seen >= target) {
+            return BucketUpperBoundNs(i);
+        }
+    }
+    return BucketUpperBoundNs(kBucketCount - 1);
+}
+
+size_t LatencyHistogram::BucketIndex(uint64_t latency_ns) {
+    if (latency_ns < 1'000'000ULL) {
+        return static_cast<size_t>(latency_ns / 5'000ULL);
+    }
+    if (latency_ns < 10'000'000ULL) {
+        return kBuckets0 +
+               static_cast<size_t>((latency_ns - 1'000'000ULL) / 50'000ULL);
+    }
+    if (latency_ns < 100'000'000ULL) {
+        return kBuckets0 + kBuckets1 +
+               static_cast<size_t>((latency_ns - 10'000'000ULL) / 500'000ULL);
+    }
+    if (latency_ns < 1'000'000'000ULL) {
+        return kBuckets0 + kBuckets1 + kBuckets2 +
+               static_cast<size_t>((latency_ns - 100'000'000ULL) /
+                                   5'000'000ULL);
+    }
+    return kBucketCount - 1;
+}
+
+uint64_t LatencyHistogram::BucketUpperBoundNs(size_t index) {
+    if (index < kBuckets0) {
+        return (index + 1) * 5'000ULL;
+    }
+    index -= kBuckets0;
+    if (index < kBuckets1) {
+        return 1'000'000ULL + (index + 1) * 50'000ULL;
+    }
+    index -= kBuckets1;
+    if (index < kBuckets2) {
+        return 10'000'000ULL + (index + 1) * 500'000ULL;
+    }
+    index -= kBuckets2;
+    if (index < kBuckets3) {
+        return 100'000'000ULL + (index + 1) * 5'000'000ULL;
+    }
+    return std::numeric_limits<uint64_t>::max();
+}
+
+void BenchStats::RecordSubmit() { ++submitted_ops; }
+
+void BenchStats::RecordSuccess(uint64_t io_bytes, uint64_t latency_ns) {
+    ++completed_ops;
+    bytes += io_bytes;
+    latency_histogram.Record(latency_ns);
+}
+
+void BenchStats::RecordFailure() { ++failed_ops; }
+
+void BenchStats::MergeFrom(const BenchStats &other) {
+    submitted_ops += other.submitted_ops;
+    completed_ops += other.completed_ops;
+    failed_ops += other.failed_ops;
+    bytes += other.bytes;
+    latency_histogram.MergeFrom(other.latency_histogram);
+}
+
+ThroughputView ComputeThroughput(const CounterSnapshot &counters,
+                                 double duration_sec) {
+    ThroughputView view;
+    if (duration_sec <= 0.0) {
+        return view;
+    }
+    view.bandwidth_bytes_per_sec = counters.bytes / duration_sec;
+    view.iops = counters.completed_ops / duration_sec;
+    view.fail_iops = counters.failed_ops / duration_sec;
+    return view;
+}
+
+LatencyView ComputeLatency(const BenchStats &stats) {
+    LatencyView view;
+    const uint64_t count = stats.latency_histogram.Count();
+    if (count == 0) {
+        return view;
+    }
+    view.avg_ns = stats.latency_histogram.SumNs() / count;
+    view.p50_ns = stats.latency_histogram.PercentileNs(0.50);
+    view.p95_ns = stats.latency_histogram.PercentileNs(0.95);
+    view.p99_ns = stats.latency_histogram.PercentileNs(0.99);
+    return view;
+}
+
+CounterSnapshot DeltaCounters(const CounterSnapshot &current,
+                              const CounterSnapshot &previous) {
+    return {
+        current.submitted_ops - previous.submitted_ops,
+        current.completed_ops - previous.completed_ops,
+        current.failed_ops - previous.failed_ops,
+        current.bytes - previous.bytes,
+    };
+}
+
+uint64_t NextLba(EndpointContext &endpoint, bool random_lba) {
+    if (endpoint.range_blocks <= endpoint.io_blocks) {
+        return FLAGS_start_lba;
+    }
+
+    if (random_lba) {
+        std::uniform_int_distribution<uint64_t> dist(
+            0, endpoint.range_blocks - endpoint.io_blocks);
+        return FLAGS_start_lba + dist(endpoint.rng);
+    }
+
+    uint64_t lba = FLAGS_start_lba + endpoint.next_block;
+    endpoint.next_block += endpoint.io_blocks;
+    if (endpoint.next_block + endpoint.io_blocks > endpoint.range_blocks) {
+        endpoint.next_block = 0;
+    }
+    return lba;
+}
+
+int PickTaskOp(BenchOp op_mode, std::mt19937_64 &rng) {
+    switch (op_mode) {
+        case BenchOp::READ:
+            return 0;
+        case BenchOp::WRITE:
+            return 1;
+        case BenchOp::MIXED: {
+            std::uniform_int_distribution<int> dist(0, 1);
+            return dist(rng);
+        }
+    }
+    return 0;
+}
+
+void FillPattern(void *buffer, size_t size, uint64_t seq) {
+    auto *ptr = reinterpret_cast<uint8_t *>(buffer);
+    for (size_t i = 0; i < size; ++i) {
+        ptr[i] = static_cast<uint8_t>((seq + i) & 0xff);
+    }
+}
+
+void PublishSubmit(LiveCounters &live) {
+    live.submitted_ops.fetch_add(1, std::memory_order_relaxed);
+}
+
+void PublishSuccess(LiveCounters &live, uint64_t io_bytes) {
+    live.completed_ops.fetch_add(1, std::memory_order_relaxed);
+    live.bytes.fetch_add(io_bytes, std::memory_order_relaxed);
+}
+
+void PublishFailure(LiveCounters &live) {
+    live.failed_ops.fetch_add(1, std::memory_order_relaxed);
+}
+
+void SubmitSlot(Slot &slot, EndpointContext &endpoint, size_t endpoint_index,
+                mooncake::SpdkNofWorkerPool &pool, BenchOp op_mode,
+                uint64_t submit_seq, bool measure_now,
+                EndpointStats &endpoint_stats) {
+    int op = PickTaskOp(op_mode, endpoint.rng);
+    if (op == 1 && FLAGS_fill_on_write) {
+        FillPattern(slot.buffer, slot.buffer_size, submit_seq);
+    }
+
+    uint64_t lba = NextLba(endpoint, FLAGS_random_lba);
+    slot.io_bytes = slot.buffer_size;
+    slot.endpoint_index = endpoint_index;
+    slot.state = std::make_shared<mooncake::SpdkNofOperationState>();
+    slot.future = std::make_unique<mooncake::TransferFuture>(slot.state);
+    slot.measure = measure_now;
+    slot.submit_time = Clock::now();
+
+    mooncake::SpdkNofTask task(endpoint.seg_handle, slot.buffer, lba,
+                               static_cast<uint32_t>(endpoint.io_blocks), op,
+                               slot.state);
+    pool.submitTask(std::move(task));
+    slot.active = true;
+
+    if (measure_now) {
+        endpoint_stats.measured.RecordSubmit();
+        PublishSubmit(endpoint_stats.live);
+    }
+}
+
+void WorkerLoop(BenchThreadContext *thread_context,
+                std::vector<EndpointContext> *endpoints,
+                std::vector<EndpointStats> *endpoint_stats,
+                mooncake::SpdkNofWorkerPool *pool, BenchOp op_mode,
+                TimePoint warmup_end, TimePoint end,
+                std::atomic<size_t> *active_threads) {
+    uint64_t submit_seq = static_cast<uint64_t>(thread_context->thread_index)
+                          << 48;
+
+    for (auto &thread_endpoint : thread_context->endpoints) {
+        auto &endpoint = (*endpoints)[thread_endpoint.endpoint_index];
+        auto &stats = (*endpoint_stats)[thread_endpoint.endpoint_index];
+        for (auto &slot : thread_endpoint.slots) {
+            SubmitSlot(slot, endpoint, thread_endpoint.endpoint_index, *pool,
+                       op_mode, submit_seq++, false, stats);
+        }
+    }
+
+    while (true) {
+        const TimePoint now_before_scan = Clock::now();
+        const bool before_end = now_before_scan < end;
+        bool any_active = false;
+
+        for (auto &thread_endpoint : thread_context->endpoints) {
+            auto &endpoint = (*endpoints)[thread_endpoint.endpoint_index];
+            auto &stats = (*endpoint_stats)[thread_endpoint.endpoint_index];
+            for (auto &slot : thread_endpoint.slots) {
+                if (!slot.active) {
+                    continue;
+                }
+                any_active = true;
+                if (!slot.future || !slot.future->isReady()) {
+                    continue;
+                }
+
+                const TimePoint completion_time = Clock::now();
+                mooncake::ErrorCode result = slot.future->get();
+                slot.active = false;
+                slot.future.reset();
+
+                if (slot.measure) {
+                    if (result == mooncake::ErrorCode::OK) {
+                        const uint64_t latency_ns = static_cast<uint64_t>(
+                            std::chrono::duration_cast<
+                                std::chrono::nanoseconds>(completion_time -
+                                                          slot.submit_time)
+                                .count());
+                        stats.measured.RecordSuccess(slot.io_bytes, latency_ns);
+                        PublishSuccess(stats.live, slot.io_bytes);
+                    } else {
+                        stats.measured.RecordFailure();
+                        PublishFailure(stats.live);
+                    }
+                }
+
+                if (before_end) {
+                    const bool measure_now = completion_time >= warmup_end;
+                    SubmitSlot(slot, endpoint, thread_endpoint.endpoint_index,
+                               *pool, op_mode, submit_seq++, measure_now,
+                               stats);
+                    any_active = true;
+                }
+            }
+        }
+
+        if (!before_end && !any_active) {
+            break;
+        }
+
+        std::this_thread::sleep_for(std::chrono::microseconds(50));
+    }
+
+    active_threads->fetch_sub(1, std::memory_order_relaxed);
+}
+
+void FreeThreadSlots(mooncake::SpdkWrapper &wrapper,
+                     std::vector<BenchThreadContext> &thread_contexts) {
+    for (auto &thread_context : thread_contexts) {
+        for (auto &thread_endpoint : thread_context.endpoints) {
+            for (auto &slot : thread_endpoint.slots) {
+                if (slot.buffer) {
+                    wrapper.Free(slot.buffer);
+                    slot.buffer = nullptr;
+                }
+            }
+        }
+    }
+}
+
+}  // namespace
+
+int main(int argc, char **argv) {
+    google::InitGoogleLogging(argv[0]);
+    FLAGS_logtostderr = true;
+    gflags::SetUsageMessage(
+        "NoF worker pool benchmark. Use --helpshort to list benchmark "
+        "parameters.");
+    gflags::ParseCommandLineFlags(&argc, &argv, true);
+
+    try {
+        ResolveUint64Alias("io_size", &FLAGS_io_size, kDefaultIoSize, "bs",
+                           FLAGS_bs);
+        ResolveUint64Alias("submit_threads", &FLAGS_submit_threads,
+                           kDefaultSubmitThreads, "numjobs", FLAGS_numjobs);
+        ResolveUint64Alias("duration_sec", &FLAGS_duration_sec,
+                           kDefaultDurationSec, "runtime", FLAGS_runtime);
+        ResolveUint64Alias("range_bytes", &FLAGS_range_bytes,
+                           kDefaultRangeBytes, "size", FLAGS_size);
+        ResolveUint64Alias("warmup_sec", &FLAGS_warmup_sec, kDefaultWarmupSec,
+                           "ramp_time", FLAGS_ramp_time);
+        if (FLAGS_status_interval != 0) {
+            ResolveUint64Alias("report_interval_ms", &FLAGS_report_interval_ms,
+                               kDefaultReportIntervalMs, "status_interval",
+                               FLAGS_status_interval * 1000);
+        }
+        if (!FLAGS_rw.empty()) {
+            ParsedRw parsed_rw = ParseRwMode(FLAGS_rw);
+            if (FLAGS_op != kDefaultOp && FLAGS_op != parsed_rw.op) {
+                throw std::invalid_argument("Conflicting --op and --rw values");
+            }
+            if (parsed_rw.random_lba == false && FLAGS_random_lba) {
+                throw std::invalid_argument(
+                    "Conflicting --random_lba and --rw values");
+            }
+            FLAGS_op = parsed_rw.op;
+            FLAGS_random_lba = parsed_rw.random_lba;
+        }
+
+        BenchOp op_mode = ParseBenchOp(FLAGS_op);
+        if (FLAGS_endpoints.empty()) {
+            LOG(ERROR) << "--endpoints is required";
+            return 1;
+        }
+        if (FLAGS_iodepth == 0) {
+            LOG(ERROR) << "--iodepth must be greater than 0";
+            return 1;
+        }
+        if (FLAGS_submit_threads == 0) {
+            LOG(ERROR) << "--submit_threads must be greater than 0";
+            return 1;
+        }
+        if (FLAGS_duration_sec == 0) {
+            LOG(ERROR) << "--duration_sec must be greater than 0";
+            return 1;
+        }
+
+        SetEnvU64IfRequested("MC_NOF_WORKERS", FLAGS_nof_workers);
+        SetEnvU64IfRequested("MC_NOF_SUBMIT_CHUNK_BYTES",
+                             FLAGS_nof_submit_chunk_bytes);
+        SetEnvU64IfRequested("MC_NOF_INFLIGHT_BYTES_LIMIT",
+                             FLAGS_nof_inflight_bytes_limit);
+
+        auto endpoint_strings = mooncake::splitString(FLAGS_endpoints);
+        if (endpoint_strings.empty()) {
+            LOG(ERROR) << "No valid endpoints parsed from --endpoints";
+            return 1;
+        }
+
+        auto &wrapper = mooncake::SpdkWrapper::GetInstance();
+        if (!wrapper.InitializeEnv()) {
+            LOG(ERROR) << "Failed to initialize SPDK environment";
+            return 1;
+        }
+
+        std::vector<EndpointContext> endpoints;
+        endpoints.reserve(endpoint_strings.size());
+        std::unordered_set<mooncake::nof_seg_handle *> unique_handles;
+        for (size_t i = 0; i < endpoint_strings.size(); ++i) {
+            EndpointContext endpoint(endpoint_strings[i], FLAGS_seed + i);
+            endpoint.seg_handle = wrapper.OpenNofSegment(endpoint.endpoint);
+            if (!endpoint.seg_handle) {
+                LOG(ERROR) << "Failed to open NoF endpoint: "
+                           << endpoint.endpoint;
+                return 1;
+            }
+            endpoint.block_size = wrapper.GetBlockSize(endpoint.seg_handle);
+            if (endpoint.block_size == INVALID_BLOCK_SIZE ||
+                endpoint.block_size == 0) {
+                LOG(ERROR) << "Invalid block size for endpoint: "
+                           << endpoint.endpoint;
+                return 1;
+            }
+            if (FLAGS_io_size % endpoint.block_size != 0) {
+                LOG(ERROR) << "--io_size=" << FLAGS_io_size
+                           << " is not aligned to block size "
+                           << endpoint.block_size << " for endpoint "
+                           << endpoint.endpoint;
+                return 1;
+            }
+            if (FLAGS_range_bytes % endpoint.block_size != 0) {
+                LOG(ERROR) << "--range_bytes=" << FLAGS_range_bytes
+                           << " is not aligned to block size "
+                           << endpoint.block_size;
+                return 1;
+            }
+            endpoint.io_blocks = FLAGS_io_size / endpoint.block_size;
+            endpoint.range_blocks = FLAGS_range_bytes / endpoint.block_size;
+            if (endpoint.range_blocks < endpoint.io_blocks) {
+                LOG(ERROR)
+                    << "--range_bytes is smaller than one I/O for endpoint "
+                    << endpoint.endpoint;
+                return 1;
+            }
+            unique_handles.insert(endpoint.seg_handle);
+            endpoints.push_back(std::move(endpoint));
+        }
+
+        const size_t effective_submit_threads =
+            std::min<size_t>(FLAGS_submit_threads, endpoints.size());
+        const uint64_t total_iodepth =
+            static_cast<uint64_t>(endpoints.size()) * FLAGS_iodepth;
+
+        LOG(INFO) << "Bench config: endpoints=" << endpoints.size()
+                  << ", unique_handles=" << unique_handles.size()
+                  << ", configured_nof_workers="
+                  << (FLAGS_nof_workers == 0 ? mooncake::kDefaultSpdkNofWorkers
+                                             : FLAGS_nof_workers)
+                  << ", effective_worker_bindings<="
+                  << std::min<uint64_t>(unique_handles.size(),
+                                        FLAGS_nof_workers == 0
+                                            ? mooncake::kDefaultSpdkNofWorkers
+                                            : FLAGS_nof_workers)
+                  << ", io_size=" << FLAGS_io_size
+                  << ", iodepth_per_endpoint=" << FLAGS_iodepth
+                  << ", total_iodepth=" << total_iodepth
+                  << ", submit_threads=" << effective_submit_threads
+                  << ", warmup_sec=" << FLAGS_warmup_sec
+                  << ", duration_sec=" << FLAGS_duration_sec;
+        if (unique_handles.size() < endpoints.size()) {
+            LOG(WARNING)
+                << "Some endpoints resolved to the same nof_seg_handle. In the "
+                   "current SpdkNofWorkerPool implementation, the same handle "
+                   "binds to a single worker thread, so duplicate endpoints "
+                   "will "
+                   "not increase worker parallelism.";
+        }
+
+        std::vector<BenchThreadContext> thread_contexts(
+            effective_submit_threads);
+        for (size_t thread_index = 0; thread_index < effective_submit_threads;
+             ++thread_index) {
+            thread_contexts[thread_index].thread_index = thread_index;
+        }
+        for (size_t endpoint_index = 0; endpoint_index < endpoints.size();
+             ++endpoint_index) {
+            size_t thread_index = endpoint_index % effective_submit_threads;
+            ThreadEndpointState thread_endpoint;
+            thread_endpoint.endpoint_index = endpoint_index;
+            thread_endpoint.slots.resize(FLAGS_iodepth);
+            thread_contexts[thread_index].endpoints.push_back(
+                std::move(thread_endpoint));
+        }
+
+        for (auto &thread_context : thread_contexts) {
+            for (auto &thread_endpoint : thread_context.endpoints) {
+                for (auto &slot : thread_endpoint.slots) {
+                    slot.buffer_size = FLAGS_io_size;
+                    slot.buffer = wrapper.Alloc(slot.buffer_size, 0x1000,
+                                                FLAGS_socket_id);
+                    if (!slot.buffer) {
+                        LOG(ERROR) << "Failed to allocate DMA buffer of size "
+                                   << slot.buffer_size;
+                        FreeThreadSlots(wrapper, thread_contexts);
+                        return 1;
+                    }
+                }
+            }
+        }
+
+        std::vector<EndpointStats> endpoint_stats(endpoints.size());
+        mooncake::SpdkNofWorkerPool pool;
+
+        TimePoint start = Clock::now();
+        TimePoint warmup_end = start + std::chrono::seconds(FLAGS_warmup_sec);
+        TimePoint end = warmup_end + std::chrono::seconds(FLAGS_duration_sec);
+
+        std::vector<std::thread> workers;
+        workers.reserve(thread_contexts.size());
+        std::atomic<size_t> active_threads(thread_contexts.size());
+        for (auto &thread_context : thread_contexts) {
+            workers.emplace_back(WorkerLoop, &thread_context, &endpoints,
+                                 &endpoint_stats, &pool, op_mode, warmup_end,
+                                 end, &active_threads);
+        }
+
+        TimePoint next_report =
+            Clock::now() + std::chrono::milliseconds(FLAGS_report_interval_ms);
+        std::vector<CounterSnapshot> last_endpoint_snapshots(endpoints.size());
+
+        while (active_threads.load(std::memory_order_relaxed) > 0) {
+            TimePoint now = Clock::now();
+            if (now < next_report) {
+                std::this_thread::sleep_for(std::chrono::milliseconds(10));
+                continue;
+            }
+
+            if (now >= warmup_end) {
+                CounterSnapshot total_current;
+                CounterSnapshot total_last;
+                for (size_t endpoint_index = 0;
+                     endpoint_index < endpoints.size(); ++endpoint_index) {
+                    CounterSnapshot current =
+                        endpoint_stats[endpoint_index].live.Snapshot();
+                    total_current.submitted_ops += current.submitted_ops;
+                    total_current.completed_ops += current.completed_ops;
+                    total_current.failed_ops += current.failed_ops;
+                    total_current.bytes += current.bytes;
+                    total_last.submitted_ops +=
+                        last_endpoint_snapshots[endpoint_index].submitted_ops;
+                    total_last.completed_ops +=
+                        last_endpoint_snapshots[endpoint_index].completed_ops;
+                    total_last.failed_ops +=
+                        last_endpoint_snapshots[endpoint_index].failed_ops;
+                    total_last.bytes +=
+                        last_endpoint_snapshots[endpoint_index].bytes;
+                }
+
+                double report_window_sec =
+                    static_cast<double>(FLAGS_report_interval_ms) / 1000.0;
+                CounterSnapshot delta_total =
+                    DeltaCounters(total_current, total_last);
+                ThroughputView total_view =
+                    ComputeThroughput(delta_total, report_window_sec);
+                LOG(INFO) << "interval bw="
+                          << FormatBytesPerSecond(
+                                 total_view.bandwidth_bytes_per_sec)
+                          << ", iops=" << std::fixed << std::setprecision(2)
+                          << total_view.iops
+                          << ", fail_iops=" << total_view.fail_iops
+                          << ", submitted=" << total_current.submitted_ops
+                          << ", completed=" << total_current.completed_ops
+                          << ", failed=" << total_current.failed_ops;
+
+                for (size_t endpoint_index = 0;
+                     endpoint_index < endpoints.size(); ++endpoint_index) {
+                    CounterSnapshot current =
+                        endpoint_stats[endpoint_index].live.Snapshot();
+                    CounterSnapshot delta = DeltaCounters(
+                        current, last_endpoint_snapshots[endpoint_index]);
+                    ThroughputView endpoint_view =
+                        ComputeThroughput(delta, report_window_sec);
+                    LOG(INFO)
+                        << "interval endpoint[" << endpoint_index << "] bw="
+                        << FormatBytesPerSecond(
+                               endpoint_view.bandwidth_bytes_per_sec)
+                        << ", iops=" << std::fixed << std::setprecision(2)
+                        << endpoint_view.iops
+                        << ", fail_iops=" << endpoint_view.fail_iops
+                        << ", endpoint=" << endpoints[endpoint_index].endpoint;
+                    last_endpoint_snapshots[endpoint_index] = current;
+                }
+            }
+
+            next_report =
+                now + std::chrono::milliseconds(FLAGS_report_interval_ms);
+        }
+
+        for (auto &worker : workers) {
+            worker.join();
+        }
+
+        FreeThreadSlots(wrapper, thread_contexts);
+
+        BenchStats total_measured;
+        for (const auto &stats : endpoint_stats) {
+            total_measured.MergeFrom(stats.measured);
+        }
+
+        double duration = static_cast<double>(FLAGS_duration_sec);
+        ThroughputView total_view = ComputeThroughput(
+            {total_measured.submitted_ops, total_measured.completed_ops,
+             total_measured.failed_ops, total_measured.bytes},
+            duration);
+        LatencyView total_latency = ComputeLatency(total_measured);
+
+        std::cout << std::fixed << std::setprecision(2);
+        std::cout << "========== NoF Worker Pool Bench ==========\n";
+        std::cout << "endpoints=" << endpoints.size() << "\n";
+        std::cout << "unique_handles=" << unique_handles.size() << "\n";
+        std::cout << "configured_nof_workers="
+                  << (FLAGS_nof_workers == 0 ? mooncake::kDefaultSpdkNofWorkers
+                                             : FLAGS_nof_workers)
+                  << "\n";
+        std::cout << "effective_worker_bindings<="
+                  << std::min<uint64_t>(unique_handles.size(),
+                                        FLAGS_nof_workers == 0
+                                            ? mooncake::kDefaultSpdkNofWorkers
+                                            : FLAGS_nof_workers)
+                  << "\n";
+        std::cout << "numjobs=" << effective_submit_threads << "\n";
+        std::cout << "bs=" << FLAGS_io_size << "\n";
+        std::cout << "iodepth=" << FLAGS_iodepth << "\n";
+        std::cout << "total_iodepth=" << total_iodepth << "\n";
+        std::cout << "runtime=" << FLAGS_duration_sec << "\n";
+        std::cout << "completed_ops=" << total_measured.completed_ops << "\n";
+        std::cout << "failed_ops=" << total_measured.failed_ops << "\n";
+        std::cout << "bw="
+                  << FormatBytesPerSecond(total_view.bandwidth_bytes_per_sec)
+                  << "\n";
+        std::cout << "iops=" << total_view.iops << "\n";
+        std::cout << "fail_iops=" << total_view.fail_iops << "\n";
+        std::cout << "clat_mean=" << FormatLatencyNs(total_latency.avg_ns)
+                  << "\n";
+        std::cout << "clat_p50=" << FormatLatencyNs(total_latency.p50_ns)
+                  << "\n";
+        std::cout << "clat_p95=" << FormatLatencyNs(total_latency.p95_ns)
+                  << "\n";
+        std::cout << "clat_p99=" << FormatLatencyNs(total_latency.p99_ns)
+                  << "\n";
+
+        for (size_t endpoint_index = 0; endpoint_index < endpoints.size();
+             ++endpoint_index) {
+            ThroughputView endpoint_view = ComputeThroughput(
+                {endpoint_stats[endpoint_index].measured.submitted_ops,
+                 endpoint_stats[endpoint_index].measured.completed_ops,
+                 endpoint_stats[endpoint_index].measured.failed_ops,
+                 endpoint_stats[endpoint_index].measured.bytes},
+                duration);
+            LatencyView endpoint_latency =
+                ComputeLatency(endpoint_stats[endpoint_index].measured);
+            std::cout << "endpoint[" << endpoint_index
+                      << "]=" << endpoints[endpoint_index].endpoint << "\n";
+            std::cout << "endpoint[" << endpoint_index << "].completed_ops="
+                      << endpoint_stats[endpoint_index].measured.completed_ops
+                      << "\n";
+            std::cout << "endpoint[" << endpoint_index << "].failed_ops="
+                      << endpoint_stats[endpoint_index].measured.failed_ops
+                      << "\n";
+            std::cout << "endpoint[" << endpoint_index << "].bw="
+                      << FormatBytesPerSecond(
+                             endpoint_view.bandwidth_bytes_per_sec)
+                      << "\n";
+            std::cout << "endpoint[" << endpoint_index
+                      << "].iops=" << endpoint_view.iops << "\n";
+            std::cout << "endpoint[" << endpoint_index
+                      << "].fail_iops=" << endpoint_view.fail_iops << "\n";
+            std::cout << "endpoint[" << endpoint_index << "].clat_mean="
+                      << FormatLatencyNs(endpoint_latency.avg_ns) << "\n";
+            std::cout << "endpoint[" << endpoint_index << "].clat_p50="
+                      << FormatLatencyNs(endpoint_latency.p50_ns) << "\n";
+            std::cout << "endpoint[" << endpoint_index << "].clat_p95="
+                      << FormatLatencyNs(endpoint_latency.p95_ns) << "\n";
+            std::cout << "endpoint[" << endpoint_index << "].clat_p99="
+                      << FormatLatencyNs(endpoint_latency.p99_ns) << "\n";
+        }
+        std::cout << "==========================================\n";
+
+        google::ShutdownGoogleLogging();
+        return 0;
+    } catch (const std::exception &e) {
+        LOG(ERROR) << "Benchmark failed: " << e.what();
+        google::ShutdownGoogleLogging();
+        return 1;
+    }
+}
diff --git a/mooncake-store/benchmarks/store_kv_bench.md b/mooncake-store/benchmarks/store_kv_bench.md
new file mode 100644
index 00000000..4d8f85a5
--- /dev/null
+++ b/mooncake-store/benchmarks/store_kv_bench.md
@@ -0,0 +1,204 @@
+# `store_kv_bench.py`
+
+`store_kv_bench.py` is a Mooncake Store end-to-end KV benchmark tool. It talks
+to a real Mooncake cluster through the Python `store` binding and can exercise
+`put/get` as well as zero-copy `put_from/get_into` style APIs.
+
+## Scope
+
+This tool focuses on object-semantic benchmark scenarios:
+
+- Functional verification with read-after-write validation
+- Dataset fill for eviction / capacity tests
+- Pure write performance
+- Pure read performance
+- Read/write mixed mode with "existing-object read + new-object write"
+
+Fault injection, NoF register / unregister, heartbeat trigger, memory segment
+unmount, and target-side operations are intentionally out of scope. The tool
+supports phase gaps so external tools can finish those operations before the
+next phase continues.
+
+## Supported Scenarios
+
+- `verify_write`
+  - Fixed-count write followed by full readback verification
+- `fill`
+  - Fixed-count write used for filling a dataset / eviction watermark
+- `write_perf`
+  - Time-based or fixed-count write benchmark
+- `read_perf`
+  - Optional prepare-write phase, then read performance benchmark
+- `mixed_rw`
+  - Optional prepare-write phase, then mixed "read prepared objects + write new objects"
+
+## APIs
+
+- `--io-api=plain`
+  - Single object:
+    - `put`
+    - `get`
+  - Batch:
+    - `put_batch`
+    - `get_batch`
+- `--io-api=zcopy`
+  - Single object:
+    - `put_from`
+    - `get_into`
+  - Batch:
+    - `batch_put_from`
+    - `batch_get_into`
+
+`zcopy` mode automatically allocates temporary user buffers and registers them
+with `register_buffer`.
+
+## Key Rules
+
+- Keys are generated deterministically:
+  - `{prefix padded/truncated to fit}{16-digit object id}`
+- The same `key-prefix`, `key-size`, and `object-id-start` produce the same key sequence
+- `verify` currently requires `pattern`
+- Any write-involved scenario requires `value-size` to be 512-byte aligned
+- `memory-replica-num` and `nof-replica-num` cannot both be `0`
+- `prepare-objects`
+  - Controls how many objects are written by the prepare phase
+  - `0` means reuse `nr-objects`
+
+## Phase Gap
+
+Phase gaps are used when an external tool needs time to inject a fault or do an
+unmount / remount operation.
+
+- `--phase-gap-mode=none`
+  - Continue immediately
+- `--phase-gap-mode=sleep --phase-gap-sec=N`
+  - Sleep before the next phase
+- `--phase-gap-mode=manual`
+  - Wait for Enter
+- `--phase-gap-mode=file --phase-gap-file=/tmp/bench.ready`
+  - Wait until the file exists
+
+If `file` mode is used, make sure the marker file does not already exist before
+starting the benchmark.
+
+## Common Examples
+
+### 1. Functional verification (`1+0`)
+
+```bash
+python3 mooncake-store/benchmarks/store_kv_bench.py \
+  --scenario verify_write \
+  --io-api plain \
+  --local-hostname 127.0.0.1:50071 \
+  --metadata-server http://127.0.0.1:8080/metadata \
+  --master-server 127.0.0.1:50051 \
+  --protocol tcp \
+  --global-segment-size $((64*1024*1024)) \
+  --local-buffer-size $((32*1024*1024)) \
+  --nr-objects 16 \
+  --batch-size 4 \
+  --key-prefix verify \
+  --key-size 20 \
+  --value-size 4096 \
+  --memory-replica-num 1 \
+  --nof-replica-num 0 \
+  --verify \
+  --pattern 0xab
+```
+
+### 2. NoF-only functional verification (`0+1`)
+
+```bash
+python3 mooncake-store/benchmarks/store_kv_bench.py \
+  --scenario verify_write \
+  --io-api plain \
+  --local-hostname 127.0.0.1:50071 \
+  --metadata-server http://127.0.0.1:8080/metadata \
+  --master-server 127.0.0.1:50051 \
+  --protocol tcp \
+  --global-segment-size 0 \
+  --local-buffer-size $((8*1024*1024)) \
+  --nr-objects 8 \
+  --batch-size 2 \
+  --key-prefix nofonly \
+  --key-size 20 \
+  --value-size 4096 \
+  --memory-replica-num 0 \
+  --nof-replica-num 1 \
+  --verify \
+  --pattern 0xcd
+```
+
+### 3. Read performance with automatic prepare phase
+
+```bash
+python3 mooncake-store/benchmarks/store_kv_bench.py \
+  --scenario read_perf \
+  --prepare-mode auto \
+  --phase-gap-mode sleep \
+  --phase-gap-sec 1 \
+  --io-api plain \
+  --local-hostname 127.0.0.1:50071 \
+  --metadata-server http://127.0.0.1:8080/metadata \
+  --master-server 127.0.0.1:50051 \
+  --protocol tcp \
+  --nr-objects 32 \
+  --batch-size 4 \
+  --runtime 5 \
+  --key-prefix readperf \
+  --key-size 20 \
+  --value-size 4096 \
+  --memory-replica-num 1 \
+  --nof-replica-num 0 \
+  --verify \
+  --pattern 0xee
+```
+
+### 4. Mixed read/write with initial dataset
+
+```bash
+python3 mooncake-store/benchmarks/store_kv_bench.py \
+  --scenario mixed_rw \
+  --prepare-mode auto \
+  --io-api zcopy \
+  --local-hostname 127.0.0.1:50071 \
+  --metadata-server http://127.0.0.1:8080/metadata \
+  --master-server 127.0.0.1:50051 \
+  --protocol tcp \
+  --nr-objects 64 \
+  --write-objects 4096 \
+  --batch-size 4 \
+  --runtime 10 \
+  --rwmixread 70 \
+  --key-prefix mixed \
+  --key-size 20 \
+  --value-size 4096 \
+  --memory-replica-num 1 \
+  --nof-replica-num 1 \
+  --verify \
+  --pattern 0x5a
+```
+
+In `mixed_rw`, reads are served from the prepared object set, while writes
+always use fresh object ids. This keeps the workload as "existing-object read +
+new-object write" and avoids key overlap between the read and write streams.
+
+## Output
+
+Each phase prints:
+
+- request counts
+- KV counts
+- miss / verify-failure counts
+- bytes processed
+- duration
+- `req/s`
+- `kv/s`
+- `MiB/s`
+- `lat_mean`
+- `lat_p50`
+- `lat_p95`
+- `lat_p99`
+- aggregated error counts
+
+An overall summary is printed after all phases complete.
diff --git a/mooncake-store/benchmarks/store_kv_bench.py b/mooncake-store/benchmarks/store_kv_bench.py
new file mode 100644
index 00000000..68fa639b
--- /dev/null
+++ b/mooncake-store/benchmarks/store_kv_bench.py
@@ -0,0 +1,1028 @@
+#!/usr/bin/env python3
+"""Mooncake Store end-to-end KV benchmark."""
+
+from __future__ import annotations
+
+import argparse
+import ctypes
+import logging
+import math
+import os
+import random
+import statistics
+import threading
+import time
+from collections import Counter
+from dataclasses import dataclass, field
+from typing import Callable, Iterable, List, Optional
+from mooncake.store import MooncakeDistributedStore, ReplicateConfig, get_alloc_func_addr, get_free_func_addr
+
+
+LOG = logging.getLogger("store_kv_bench")
+
+
+def build_parser() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(
+        description="Mooncake Store end-to-end KV benchmark",
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
+    )
+    parser.add_argument(
+        "--scenario",
+        required=True,
+        choices=["verify_write", "fill", "write_perf", "read_perf", "mixed_rw"],
+        help="Benchmark scenario to execute.",
+    )
+
+    parser.add_argument("--local-hostname", default="127.0.0.1:50071")
+    parser.add_argument("--metadata-server", default="http://127.0.0.1:8080/metadata")
+    parser.add_argument("--master-server", default="127.0.0.1:50051")
+    parser.add_argument("--protocol", default="tcp")
+    parser.add_argument("--device-name", default="")
+    parser.add_argument("--global-segment-size", type=int, default=64 * 1024 * 1024)
+    parser.add_argument("--local-buffer-size", type=int, default=32 * 1024 * 1024)
+    parser.add_argument(
+        "--io-api",
+        choices=["plain", "zcopy"],
+        default="plain",
+        help="plain uses put/get/put_batch/get_batch, zcopy uses put_from/get_into/batch_put_from/batch_get_into",
+    )
+
+    parser.add_argument("--numjobs", type=int, default=1)
+    parser.add_argument("--iodepth", type=int, default=1)
+    parser.add_argument("--batch-size", type=int, default=1)
+    parser.add_argument("--runtime", type=int, default=0, help="Seconds. 0 means object-count based.")
+    parser.add_argument("--nr-objects", type=int, default=128)
+    parser.add_argument("--write-objects", type=int, default=0)
+    parser.add_argument(
+        "--prepare-objects",
+        type=int,
+        default=0,
+        help="Object count used by the prepare phase. 0 means reuse nr-objects.",
+    )
+    parser.add_argument("--object-id-start", type=int, default=0)
+    parser.add_argument("--key-prefix", default="kvbench")
+    parser.add_argument("--key-size", type=int, default=20)
+    parser.add_argument("--value-size", type=int, default=4096)
+    parser.add_argument("--rand-seed", type=int, default=1)
+
+    parser.add_argument("--memory-replica-num", type=int, default=1)
+    parser.add_argument("--nof-replica-num", type=int, default=0)
+
+    parser.add_argument("--verify", action="store_true")
+    parser.add_argument("--pattern", default="")
+    parser.add_argument("--prepare-mode", choices=["auto", "none", "write"], default="auto")
+    parser.add_argument("--rwmixread", type=int, default=70)
+
+    parser.add_argument(
+        "--phase-gap-mode",
+        choices=["none", "sleep", "manual", "file"],
+        default="none",
+    )
+    parser.add_argument("--phase-gap-sec", type=int, default=0)
+    parser.add_argument("--phase-gap-file", default="")
+    parser.add_argument("--phase-gap-timeout-sec", type=int, default=600)
+    parser.add_argument("--log-level", default="INFO")
+    return parser
+
+
+def setup_logging(level: str) -> None:
+    logging.basicConfig(
+        level=getattr(logging, level.upper(), logging.INFO),
+        format="%(asctime)s %(levelname)s %(name)s %(message)s",
+        datefmt="%Y-%m-%d %H:%M:%S",
+    )
+
+
+@dataclass
+class PhaseStats:
+    name: str
+    request_latencies: List[float] = field(default_factory=list)
+    requests: int = 0
+    successful_requests: int = 0
+    failed_requests: int = 0
+    kvs: int = 0
+    successful_kvs: int = 0
+    failed_kvs: int = 0
+    misses: int = 0
+    verify_failures: int = 0
+    bytes_processed: int = 0
+    error_counts: Counter = field(default_factory=Counter)
+    start_time: float = 0.0
+    end_time: float = 0.0
+    dataset_exhausted: bool = False
+
+
+@dataclass
+class RequestResult:
+    request_ok: bool
+    kv_successes: int
+    kv_failures: int
+    bytes_processed: int
+    successful_object_ids: List[int] = field(default_factory=list)
+    misses: int = 0
+    verify_failures: int = 0
+    error_counts: Counter = field(default_factory=Counter)
+
+
+class PayloadFactory:
+    def __init__(self, value_size: int, pattern: bytes):
+        self.value_size = value_size
+        self.pattern = pattern
+        self._default_cache: dict[int, bytes] = {}
+        self._pattern_payload = self._repeat(pattern) if pattern else b""
+
+    def _repeat(self, token: bytes) -> bytes:
+        repeat = (self.value_size + len(token) - 1) // len(token)
+        return (token * repeat)[: self.value_size]
+
+    def build(self, object_id: int) -> bytes:
+        if self.pattern:
+            return self._pattern_payload
+        fill_byte = object_id & 0xFF
+        payload = self._default_cache.get(fill_byte)
+        if payload is None:
+            payload = bytes([fill_byte]) * self.value_size
+            self._default_cache[fill_byte] = payload
+        return payload
+
+    def verify_payload(self, object_id: int, payload: bytes) -> bool:
+        return payload == self.build(object_id)
+
+
+class DatasetState:
+    def __init__(self, object_id_start: int):
+        self.write_lock = threading.Lock()
+        self.ids_lock = threading.Lock()
+        self.cursor_lock = threading.Lock()
+        self.next_write_id = object_id_start
+        self.prepared_ids: tuple[int, ...] = ()
+        self.written_ids: tuple[int, ...] = ()
+        self.read_cursor = 0
+
+    def reserve_write_ids(self, count: int, upper_bound: int) -> List[int]:
+        with self.write_lock:
+            if self.next_write_id >= upper_bound:
+                return []
+            end = min(self.next_write_id + count, upper_bound)
+            ids = list(range(self.next_write_id, end))
+            self.next_write_id = end
+            return ids
+
+    def mark_prepared(self, ids: Iterable[int]) -> None:
+        ids = tuple(ids)
+        if not ids:
+            return
+        with self.ids_lock:
+            self.prepared_ids = self.prepared_ids + ids
+            self.written_ids = self.written_ids + ids
+
+    def mark_runtime_written(self, ids: Iterable[int]) -> None:
+        ids = tuple(ids)
+        if not ids:
+            return
+        with self.ids_lock:
+            self.written_ids = self.written_ids + ids
+
+    def written_count(self) -> int:
+        with self.ids_lock:
+            return len(self.written_ids)
+
+    def snapshot_written_ids(self) -> List[int]:
+        with self.ids_lock:
+            return list(self.written_ids)
+
+    def next_read_ids(
+        self,
+        count: int,
+        *,
+        loop: bool,
+        sequential: bool,
+        rng,
+        source: str = "written",
+    ) -> List[int]:
+        with self.ids_lock:
+            readable_ids = self.prepared_ids if source == "prepared" else self.written_ids
+        if not readable_ids:
+            return []
+        if sequential:
+            with self.cursor_lock:
+                result = []
+                for _ in range(count):
+                    if self.read_cursor >= len(readable_ids):
+                        if not loop:
+                            break
+                        self.read_cursor = 0
+                    result.append(readable_ids[self.read_cursor])
+                    self.read_cursor += 1
+                return result
+        return [readable_ids[rng.randrange(len(readable_ids))] for _ in range(count)]
+
+
+def parse_pattern(pattern_text: str) -> bytes:
+    if not pattern_text:
+        return b""
+    if pattern_text.startswith("0x"):
+        hex_text = pattern_text[2:]
+        if len(hex_text) % 2 != 0:
+            raise ValueError("hex pattern length must be even")
+        return bytes.fromhex(hex_text)
+    return pattern_text.encode("utf-8")
+
+
+def make_key(prefix: str, key_size: int, object_id: int) -> str:
+    suffix = f"{object_id:016d}"
+    if key_size < len(suffix):
+        raise ValueError(f"key_size={key_size} is smaller than suffix length {len(suffix)}")
+    prefix_space = key_size - len(suffix)
+    prefix_part = prefix[:prefix_space].ljust(prefix_space, "_")
+    return f"{prefix_part}{suffix}"
+
+
+class StoreSession:
+    def __init__(
+        self,
+        args: argparse.Namespace,
+        lane_id: int,
+        payload_factory: PayloadFactory,
+        store_obj,
+        zcopy: Optional["ZcopyBufferView"] = None,
+    ):
+        self.args = args
+        self.lane_id = lane_id
+        self.payload_factory = payload_factory
+        self.store = store_obj
+        self.config = ReplicateConfig()
+        self.config.replica_num = args.memory_replica_num
+        self.config.nof_replica_num = args.nof_replica_num
+        self._zcopy = zcopy
+
+    def close(self) -> None:
+        self._zcopy = None
+
+    def put_ids(self, object_ids: List[int]) -> RequestResult:
+        keys = [make_key(self.args.key_prefix, self.args.key_size, object_id) for object_id in object_ids]
+        ret_codes = self._put_keys(keys, object_ids)
+
+        errors: Counter = Counter()
+        success_ids: List[int] = []
+        if self.args.io_api == "plain" and not (len(object_ids) == 1 and self.args.batch_size == 1):
+            ret = ret_codes[0]
+            if ret == 0:
+                success_ids.extend(object_ids)
+            else:
+                errors[ret] += 1
+        else:
+            for object_id, ret in zip(object_ids, ret_codes):
+                if ret == 0:
+                    success_ids.append(object_id)
+                else:
+                    errors[ret] += 1
+        request_ok = len(success_ids) == len(object_ids)
+        return RequestResult(
+            request_ok=request_ok,
+            kv_successes=len(success_ids),
+            kv_failures=len(object_ids) - len(success_ids),
+            bytes_processed=len(success_ids) * self.args.value_size,
+            successful_object_ids=success_ids,
+            error_counts=errors,
+        )
+
+    def get_ids(self, object_ids: List[int], verify: bool) -> RequestResult:
+        keys = [make_key(self.args.key_prefix, self.args.key_size, object_id) for object_id in object_ids]
+        errors: Counter = Counter()
+        kv_successes = 0
+        misses = 0
+        verify_failures = 0
+        if self.args.io_api == "plain":
+            payloads = self._get_payloads_plain(keys)
+            for object_id, payload in zip(object_ids, payloads):
+                if payload in (None, b""):
+                    misses += 1
+                    errors["MISS"] += 1
+                    continue
+                if verify and not self.payload_factory.verify_payload(object_id, payload):
+                    verify_failures += 1
+                    errors["VERIFY_FAIL"] += 1
+                    continue
+                kv_successes += 1
+        else:
+            assert self._zcopy is not None
+            lengths = self._get_lengths_zcopy(keys, len(object_ids))
+            for slot, (object_id, length) in enumerate(zip(object_ids, lengths)):
+                if length < 0:
+                    if self.store.isExist(keys[slot]) == 0:
+                        misses += 1
+                        errors["MISS"] += 1
+                    else:
+                        errors[length] += 1
+                    continue
+                payload = self._zcopy.read_bytes(slot, length)
+                if verify:
+                    if length != self.args.value_size:
+                        verify_failures += 1
+                        errors["VERIFY_SIZE_MISMATCH"] += 1
+                        continue
+                    if not self.payload_factory.verify_payload(object_id, payload):
+                        verify_failures += 1
+                        errors["VERIFY_FAIL"] += 1
+                        continue
+                kv_successes += 1
+
+        kv_failures = len(object_ids) - kv_successes
+        return RequestResult(
+            request_ok=(kv_failures == 0),
+            kv_successes=kv_successes,
+            kv_failures=kv_failures,
+            bytes_processed=kv_successes * self.args.value_size,
+            misses=misses,
+            verify_failures=verify_failures,
+            error_counts=errors,
+        )
+
+    def _put_keys(self, keys: List[str], object_ids: List[int]) -> List[int]:
+        values = [self.payload_factory.build(object_id) for object_id in object_ids]
+        if self.args.io_api == "plain":
+            if len(object_ids) == 1 and self.args.batch_size == 1:
+                return [self.store.put(keys[0], values[0], self.config)]
+            return [self.store.put_batch(keys, values, self.config)]
+
+        assert self._zcopy is not None
+        ptrs = self._zcopy.fill_write_buffers(values)
+        sizes = [len(value) for value in values]
+        if len(object_ids) == 1 and self.args.batch_size == 1:
+            return [self.store.put_from(keys[0], ptrs[0], sizes[0], self.config)]
+        return list(self.store.batch_put_from(keys, ptrs, sizes, self.config))
+
+    def _get_payloads_plain(self, keys: List[str]) -> List[bytes]:
+        if len(keys) == 1 and self.args.batch_size == 1:
+            return [self.store.get(keys[0])]
+        return list(self.store.get_batch(keys))
+
+    def _get_lengths_zcopy(self, keys: List[str], slot_count: int) -> List[int]:
+        assert self._zcopy is not None
+        ptrs = self._zcopy.prepare_read_buffers(slot_count)
+        sizes = [self.args.value_size] * slot_count
+        if slot_count == 1 and self.args.batch_size == 1:
+            return [self.store.get_into(keys[0], ptrs[0], sizes[0])]
+        return list(self.store.batch_get_into(keys, ptrs, sizes))
+
+
+class ZcopyBufferPool:
+    def __init__(self, store_obj, value_size: int, slots: int):
+        self.store = store_obj
+        self.value_size = value_size
+        self.slots = slots
+        self.total_size = self.value_size * self.slots
+        self._alloc_fn = None
+        self._free_fn = None
+        self._registered = False
+        self.base_ptr = 0
+
+        alloc_addr = get_alloc_func_addr()
+        free_addr = get_free_func_addr()
+        if alloc_addr is None or free_addr is None:
+            raise RuntimeError("store module does not expose hugepage alloc/free helpers")
+
+        self._alloc_fn = ctypes.CFUNCTYPE(ctypes.c_void_p, ctypes.c_size_t)(
+            get_alloc_func_addr()
+        )
+        self._free_fn = ctypes.CFUNCTYPE(None, ctypes.c_void_p)(
+            get_free_func_addr()
+        )
+
+        raw_ptr = self._alloc_fn(self.total_size)
+        self.base_ptr = ctypes.cast(raw_ptr, ctypes.c_void_p).value or 0
+        if self.base_ptr == 0:
+            raise RuntimeError(
+                f"direct hugepage alloc failed for zcopy pool: size={self.total_size}"
+            )
+        ret = self.store.register_buffer(self.base_ptr, self.total_size)
+        if ret != 0:
+            failed_ptr = self.base_ptr
+            self._free_fn(ctypes.c_void_p(self.base_ptr))
+            self.base_ptr = 0
+            raise RuntimeError(
+                f"register_buffer failed for direct zcopy pool ptr={failed_ptr}: {ret}"
+            )
+        self._registered = True
+        self._buffer = (ctypes.c_ubyte * self.total_size).from_address(self.base_ptr)
+
+    def close(self) -> None:
+        self._buffer = None
+        if self.base_ptr:
+            if self._registered:
+                try:
+                    self.store.unregister_buffer(self.base_ptr)
+                except Exception:
+                    LOG.debug("unregister_buffer failed for direct zcopy pool ptr=%s", self.base_ptr, exc_info=True)
+                self._registered = False
+            if self._free_fn is not None:
+                self._free_fn(ctypes.c_void_p(self.base_ptr))
+            self.base_ptr = 0
+
+    def slot_ptr(self, slot: int) -> int:
+        if slot < 0 or slot >= self.slots:
+            raise IndexError(f"zcopy slot {slot} is out of range [0, {self.slots})")
+        return self.base_ptr + slot * self.value_size
+
+
+class ZcopyBufferView:
+    def __init__(self, pool: ZcopyBufferPool, slot_offset: int, slots: int):
+        self.pool = pool
+        self.slot_offset = slot_offset
+        self.slots = slots
+
+    def _slot_ptr(self, slot: int) -> int:
+        if slot < 0 or slot >= self.slots:
+            raise IndexError(f"zcopy view slot {slot} is out of range [0, {self.slots})")
+        return self.pool.slot_ptr(self.slot_offset + slot)
+
+    def fill_write_buffers(self, payloads: List[bytes]) -> List[int]:
+        ptrs: List[int] = []
+        for slot, payload in enumerate(payloads):
+            if len(payload) > self.pool.value_size:
+                raise ValueError(
+                    f"Payload size {len(payload)} exceeds slot size {self.pool.value_size}"
+                )
+            ptr = self._slot_ptr(slot)
+            ctypes.memmove(ptr, payload, len(payload))
+            ptrs.append(ptr)
+        return ptrs
+
+    def prepare_read_buffers(self, slot_count: int) -> List[int]:
+        ptrs: List[int] = []
+        for slot in range(slot_count):
+            ptr = self._slot_ptr(slot)
+            ctypes.memset(ptr, 0, self.pool.value_size)
+            ptrs.append(ptr)
+        return ptrs
+
+    def read_bytes(self, slot: int, size: int) -> bytes:
+        if size > self.pool.value_size:
+            raise ValueError(
+                f"Read size {size} exceeds slot size {self.pool.value_size}"
+            )
+        return ctypes.string_at(self._slot_ptr(slot), size)
+
+
+class StoreRuntime:
+    def __init__(self, args: argparse.Namespace, lane_count: int):
+
+        self.lane_count = lane_count
+        self.store = MooncakeDistributedStore()
+        setup_ret = self.store.setup(
+            args.local_hostname,
+            args.metadata_server,
+            args.global_segment_size,
+            args.local_buffer_size,
+            args.protocol,
+            args.device_name,
+            args.master_server,
+        )
+        if setup_ret != 0:
+            raise RuntimeError(f"setup failed: {setup_ret}")
+
+        self.zcopy_pool: Optional[ZcopyBufferPool] = None
+        if args.io_api == "zcopy":
+            slots = max(1, args.batch_size) * lane_count
+            self.zcopy_pool = ZcopyBufferPool(
+                self.store, args.value_size, slots
+            )
+
+    def make_session(
+        self,
+        args: argparse.Namespace,
+        lane_id: int,
+        payload_factory: PayloadFactory,
+    ) -> StoreSession:
+        zcopy_view: Optional[ZcopyBufferView] = None
+        if self.zcopy_pool is not None:
+            slots_per_lane = max(1, args.batch_size)
+            zcopy_view = ZcopyBufferView(
+                self.zcopy_pool, lane_id * slots_per_lane, slots_per_lane
+            )
+        return StoreSession(
+            args,
+            lane_id,
+            payload_factory,
+            self.store,
+            zcopy_view,
+        )
+
+    def close(self) -> None:
+        if self.zcopy_pool is not None:
+            self.zcopy_pool.close()
+            self.zcopy_pool = None
+        if hasattr(self.store, "close"):
+            try:
+                self.store.close()
+            except Exception:
+                LOG.debug("shared store close failed", exc_info=True)
+        elif hasattr(self.store, "tearDownAll"):
+            try:
+                self.store.tearDownAll()
+            except Exception:
+                LOG.debug("shared tearDownAll failed", exc_info=True)
+
+
+def merge_stats(name: str, stats_list: List[PhaseStats]) -> PhaseStats:
+    merged = PhaseStats(name=name)
+    if not stats_list:
+        return merged
+    merged.start_time = min((s.start_time for s in stats_list if s.start_time), default=0.0)
+    merged.end_time = max((s.end_time for s in stats_list if s.end_time), default=0.0)
+    for stats in stats_list:
+        merged.request_latencies.extend(stats.request_latencies)
+        merged.requests += stats.requests
+        merged.successful_requests += stats.successful_requests
+        merged.failed_requests += stats.failed_requests
+        merged.kvs += stats.kvs
+        merged.successful_kvs += stats.successful_kvs
+        merged.failed_kvs += stats.failed_kvs
+        merged.misses += stats.misses
+        merged.verify_failures += stats.verify_failures
+        merged.bytes_processed += stats.bytes_processed
+        merged.error_counts.update(stats.error_counts)
+        merged.dataset_exhausted = merged.dataset_exhausted or stats.dataset_exhausted
+    return merged
+
+
+def percentile(values: List[float], p: float) -> float:
+    if not values:
+        return 0.0
+    ordered = sorted(values)
+    if len(ordered) == 1:
+        return ordered[0]
+    rank = (len(ordered) - 1) * p
+    low = math.floor(rank)
+    high = math.ceil(rank)
+    if low == high:
+        return ordered[low]
+    return ordered[low] + (ordered[high] - ordered[low]) * (rank - low)
+
+
+def summarize_stats(stats: PhaseStats) -> dict:
+    duration = max(stats.end_time - stats.start_time, 0.0)
+    return {
+        "requests": stats.requests,
+        "successful_requests": stats.successful_requests,
+        "failed_requests": stats.failed_requests,
+        "kvs": stats.kvs,
+        "successful_kvs": stats.successful_kvs,
+        "failed_kvs": stats.failed_kvs,
+        "misses": stats.misses,
+        "verify_failures": stats.verify_failures,
+        "bytes": stats.bytes_processed,
+        "duration_sec": duration,
+        "req_per_sec": (stats.requests / duration) if duration > 0 else 0.0,
+        "kv_per_sec": (stats.kvs / duration) if duration > 0 else 0.0,
+        "MiB_per_sec": (stats.bytes_processed / duration / (1024 * 1024)) if duration > 0 else 0.0,
+        "lat_mean_ms": statistics.mean(stats.request_latencies) * 1000 if stats.request_latencies else 0.0,
+        "lat_p50_ms": percentile(stats.request_latencies, 0.50) * 1000,
+        "lat_p95_ms": percentile(stats.request_latencies, 0.95) * 1000,
+        "lat_p99_ms": percentile(stats.request_latencies, 0.99) * 1000,
+        "dataset_exhausted": stats.dataset_exhausted,
+        "error_counts": dict(stats.error_counts),
+    }
+
+
+def log_phase_stats(stats: PhaseStats) -> None:
+    summary = summarize_stats(stats)
+    LOG.info("=== phase %s ===", stats.name)
+    LOG.info(
+        "requests=%d successful_requests=%d failed_requests=%d kvs=%d successful_kvs=%d failed_kvs=%d",
+        summary["requests"],
+        summary["successful_requests"],
+        summary["failed_requests"],
+        summary["kvs"],
+        summary["successful_kvs"],
+        summary["failed_kvs"],
+    )
+    LOG.info(
+        "misses=%d verify_failures=%d bytes=%d duration=%.3fs req/s=%.2f kv/s=%.2f MiB/s=%.2f",
+        summary["misses"],
+        summary["verify_failures"],
+        summary["bytes"],
+        summary["duration_sec"],
+        summary["req_per_sec"],
+        summary["kv_per_sec"],
+        summary["MiB_per_sec"],
+    )
+    LOG.info(
+        "lat_mean=%.3fms lat_p50=%.3fms lat_p95=%.3fms lat_p99=%.3fms dataset_exhausted=%s",
+        summary["lat_mean_ms"],
+        summary["lat_p50_ms"],
+        summary["lat_p95_ms"],
+        summary["lat_p99_ms"],
+        summary["dataset_exhausted"],
+    )
+    if summary["error_counts"]:
+        LOG.info("errors=%s", summary["error_counts"])
+
+
+class BenchmarkRunner:
+    def __init__(self, args: argparse.Namespace):
+        self.args = args
+        self.pattern = parse_pattern(args.pattern)
+        self.payload_factory = PayloadFactory(args.value_size, self.pattern)
+        self.dataset = DatasetState(args.object_id_start)
+        self.lane_count = args.numjobs * args.iodepth
+        self._sessions: Optional[List[StoreSession]] = None
+        self._runtime: Optional[StoreRuntime] = None
+        self._validate_args()
+
+    def _validate_args(self) -> None:
+        if self.args.numjobs <= 0 or self.args.iodepth <= 0:
+            raise ValueError("numjobs and iodepth must be > 0")
+        if self.args.batch_size <= 0:
+            raise ValueError("batch-size must be > 0")
+        if self.args.value_size <= 0:
+            raise ValueError("value-size must be > 0")
+        if self.args.key_size <= 0:
+            raise ValueError("key-size must be > 0")
+        if self.args.nr_objects <= 0:
+            raise ValueError("nr-objects must be > 0")
+        if self.args.write_objects < 0:
+            raise ValueError("write-objects must be >= 0")
+        if self.args.prepare_objects < 0:
+            raise ValueError("prepare-objects must be >= 0")
+        if self.args.rwmixread < 0 or self.args.rwmixread > 100:
+            raise ValueError("rwmixread must be within [0, 100]")
+        if self.args.verify and not self.pattern:
+            raise ValueError("verify mode currently requires --pattern")
+        if self.args.memory_replica_num == 0 and self.args.nof_replica_num == 0:
+            raise ValueError("memory_replica_num and nof_replica_num cannot both be 0")
+        if self.args.phase_gap_mode == "sleep" and self.args.phase_gap_sec <= 0:
+            raise ValueError("phase-gap-sec must be > 0 when phase-gap-mode=sleep")
+        if self.args.phase_gap_mode == "file" and not self.args.phase_gap_file:
+            raise ValueError("phase-gap-file must be set when phase-gap-mode=file")
+        if self.args.scenario == "mixed_rw" and self.args.runtime <= 0:
+            raise ValueError("mixed_rw requires --runtime > 0")
+        if self._scenario_has_write() and self.args.value_size % 512 != 0:
+            raise ValueError("write-involved scenarios require value-size to be 512B aligned")
+        make_key(self.args.key_prefix, self.args.key_size, self.args.object_id_start)
+
+    def _scenario_has_write(self) -> bool:
+        return self.args.scenario in {"verify_write", "fill", "write_perf", "mixed_rw"}
+
+    def _write_budget(self) -> int:
+        return self.args.write_objects if self.args.write_objects > 0 else self.args.nr_objects
+
+    def _prepare_budget(self) -> int:
+        return self.args.prepare_objects if self.args.prepare_objects > 0 else self.args.nr_objects
+
+    def _make_sessions(self) -> List[StoreSession]:
+        if self._sessions is None:
+            self._runtime = StoreRuntime(self.args, self.lane_count)
+            self._sessions = [
+                self._runtime.make_session(self.args, lane_id, self.payload_factory)
+                for lane_id in range(self.lane_count)
+            ]
+        return self._sessions
+
+    def close(self) -> None:
+        if self._sessions is not None:
+            for session in self._sessions:
+                session.close()
+            self._sessions = None
+        if self._runtime is not None:
+            self._runtime.close()
+            self._runtime = None
+
+    def _phase_gap(self, label: str) -> None:
+        mode = self.args.phase_gap_mode
+        if mode == "none":
+            return
+        LOG.info("phase gap before %s, mode=%s", label, mode)
+        if mode == "sleep":
+            LOG.info("sleeping %d seconds before %s", self.args.phase_gap_sec, label)
+            time.sleep(self.args.phase_gap_sec)
+            return
+        if mode == "manual":
+            input(f"phase '{label}' is waiting, finish external operations then press Enter to continue...")
+            return
+        deadline = time.time() + self.args.phase_gap_timeout_sec
+        while time.time() < deadline:
+            if os.path.exists(self.args.phase_gap_file):
+                LOG.info("detected phase gap file %s, continuing to %s", self.args.phase_gap_file, label)
+                return
+            time.sleep(1.0)
+        raise TimeoutError(f"timed out waiting for phase gap file {self.args.phase_gap_file}")
+
+    def _run_threads(self, phase_name: str, worker_builder: Callable[[StoreSession, int], Callable[[PhaseStats], None]]) -> PhaseStats:
+        sessions = self._make_sessions()
+        per_lane_stats: List[Optional[PhaseStats]] = [None] * self.lane_count
+        threads: List[threading.Thread] = []
+
+        def runner(index: int, session: StoreSession) -> None:
+            stats = PhaseStats(name=f"{phase_name}/lane{index}")
+            stats.start_time = time.perf_counter()
+            worker_builder(session, index)(stats)
+            stats.end_time = time.perf_counter()
+            per_lane_stats[index] = stats
+
+        for lane_id, session in enumerate(sessions):
+            thread = threading.Thread(target=runner, args=(lane_id, session), name=f"{phase_name}-lane{lane_id}")
+            threads.append(thread)
+            thread.start()
+
+        for thread in threads:
+            thread.join()
+
+        merged = merge_stats(phase_name, [s for s in per_lane_stats if s is not None])
+        log_phase_stats(merged)
+        return merged
+
+    def _record(self, stats: PhaseStats, latency: float, request: RequestResult, kv_count: int) -> None:
+        stats.request_latencies.append(latency)
+        stats.requests += 1
+        stats.kvs += kv_count
+        if request.request_ok:
+            stats.successful_requests += 1
+        else:
+            stats.failed_requests += 1
+        stats.successful_kvs += request.kv_successes
+        stats.failed_kvs += request.kv_failures
+        stats.misses += request.misses
+        stats.verify_failures += request.verify_failures
+        stats.bytes_processed += request.bytes_processed
+        stats.error_counts.update(request.error_counts)
+
+    def _run_fixed_write(
+        self,
+        phase_name: str,
+        total_objects: int,
+        *,
+        strict: bool,
+        write_scope: str = "runtime",
+    ) -> PhaseStats:
+        write_upper = self.dataset.next_write_id + total_objects
+
+        def worker(session: StoreSession, _lane_id: int) -> Callable[[PhaseStats], None]:
+            def run(stats: PhaseStats) -> None:
+                while True:
+                    object_ids = self.dataset.reserve_write_ids(self.args.batch_size, write_upper)
+                    if not object_ids:
+                        break
+                    start = time.perf_counter()
+                    result = session.put_ids(object_ids)
+                    latency = time.perf_counter() - start
+                    self._record(stats, latency, result, len(object_ids))
+                    if result.successful_object_ids:
+                        if write_scope == "prepared":
+                            self.dataset.mark_prepared(result.successful_object_ids)
+                        else:
+                            self.dataset.mark_runtime_written(result.successful_object_ids)
+            return run
+
+        stats = self._run_threads(phase_name, worker)
+        expected = total_objects
+        if stats.successful_kvs < expected:
+            stats.dataset_exhausted = True
+        if strict and (stats.failed_kvs > 0 or stats.successful_kvs != expected):
+            raise RuntimeError(
+                f"{phase_name} strict write failed: expected {expected} objects, "
+                f"got success={stats.successful_kvs}, failed={stats.failed_kvs}"
+            )
+        return stats
+
+    def _run_time_based_write(self, phase_name: str, total_objects: int) -> PhaseStats:
+        deadline = time.time() + self.args.runtime
+        write_upper = self.dataset.next_write_id + total_objects
+        stop_event = threading.Event()
+
+        def worker(session: StoreSession, _lane_id: int) -> Callable[[PhaseStats], None]:
+            def run(stats: PhaseStats) -> None:
+                while time.time() < deadline and not stop_event.is_set():
+                    object_ids = self.dataset.reserve_write_ids(self.args.batch_size, write_upper)
+                    if not object_ids:
+                        stats.dataset_exhausted = True
+                        stop_event.set()
+                        break
+                    start = time.perf_counter()
+                    result = session.put_ids(object_ids)
+                    latency = time.perf_counter() - start
+                    self._record(stats, latency, result, len(object_ids))
+                    if result.successful_object_ids:
+                        self.dataset.mark_runtime_written(result.successful_object_ids)
+            return run
+
+        return self._run_threads(phase_name, worker)
+
+    def _run_read_phase(
+        self,
+        phase_name: str,
+        *,
+        verify: bool,
+        sequential: bool,
+        loop: bool,
+        runtime_sec: int = 0,
+    ) -> PhaseStats:
+        seed_base = self.args.rand_seed
+        if runtime_sec > 0:
+            deadline = time.time() + runtime_sec
+
+            def worker(session: StoreSession, lane_id: int) -> Callable[[PhaseStats], None]:
+                rng = random.Random(seed_base + lane_id)
+
+                def run(stats: PhaseStats) -> None:
+                    while time.time() < deadline:
+                        object_ids = self.dataset.next_read_ids(
+                            self.args.batch_size,
+                            loop=True,
+                            sequential=sequential,
+                            rng=rng,
+                            source="prepared",
+                        )
+                        if not object_ids:
+                            stats.dataset_exhausted = True
+                            break
+                        start = time.perf_counter()
+                        result = session.get_ids(object_ids, verify)
+                        latency = time.perf_counter() - start
+                        self._record(stats, latency, result, len(object_ids))
+
+                return run
+
+            return self._run_threads(phase_name, worker)
+
+        def worker(session: StoreSession, lane_id: int) -> Callable[[PhaseStats], None]:
+            rng = random.Random(seed_base + lane_id)
+
+            def run(stats: PhaseStats) -> None:
+                while True:
+                    object_ids = self.dataset.next_read_ids(
+                        self.args.batch_size,
+                        loop=loop,
+                        sequential=sequential,
+                        rng=rng,
+                        source="prepared",
+                    )
+                    if not object_ids:
+                        break
+                    start = time.perf_counter()
+                    result = session.get_ids(object_ids, verify)
+                    latency = time.perf_counter() - start
+                    self._record(stats, latency, result, len(object_ids))
+
+            return run
+
+        return self._run_threads(phase_name, worker)
+
+    def _run_mixed_phase(self, phase_name: str, extra_write_budget: int) -> PhaseStats:
+        deadline = time.time() + self.args.runtime
+        write_upper = self.dataset.next_write_id + extra_write_budget
+        stop_event = threading.Event()
+        seed_base = self.args.rand_seed
+
+        def worker(session: StoreSession, lane_id: int) -> Callable[[PhaseStats], None]:
+            rng = random.Random(seed_base + lane_id)
+
+            def run(stats: PhaseStats) -> None:
+                while time.time() < deadline and not stop_event.is_set():
+                    do_read = rng.randrange(100) < self.args.rwmixread
+                    if do_read:
+                        object_ids = self.dataset.next_read_ids(
+                            self.args.batch_size,
+                            loop=True,
+                            sequential=False,
+                            rng=rng,
+                            source="prepared",
+                        )
+                        if not object_ids:
+                            continue
+                        start = time.perf_counter()
+                        result = session.get_ids(object_ids, verify=self.args.verify)
+                        latency = time.perf_counter() - start
+                        self._record(stats, latency, result, len(object_ids))
+                        continue
+
+                    object_ids = self.dataset.reserve_write_ids(self.args.batch_size, write_upper)
+                    if not object_ids:
+                        stats.dataset_exhausted = True
+                        stop_event.set()
+                        break
+                    start = time.perf_counter()
+                    result = session.put_ids(object_ids)
+                    latency = time.perf_counter() - start
+                    self._record(stats, latency, result, len(object_ids))
+                    if result.successful_object_ids:
+                        self.dataset.mark_runtime_written(result.successful_object_ids)
+
+            return run
+
+        return self._run_threads(phase_name, worker)
+
+    def _maybe_prepare_dataset(self) -> Optional[PhaseStats]:
+        if self.args.prepare_mode == "none":
+            return None
+        if self.args.prepare_mode == "write" or self.args.scenario in {"read_perf", "mixed_rw"}:
+            stats = self._run_fixed_write(
+                "prepare_write",
+                self._prepare_budget(),
+                strict=True,
+                write_scope="prepared",
+            )
+            self._phase_gap("main_run")
+            return stats
+        return None
+
+    def run(self) -> List[PhaseStats]:
+        LOG.info(
+            "scenario=%s io_api=%s numjobs=%d iodepth=%d lanes=%d batch_size=%d value_size=%d nr_objects=%d prepare_objects=%d write_objects=%d memory_replica_num=%d nof_replica_num=%d verify=%s",
+            self.args.scenario,
+            self.args.io_api,
+            self.args.numjobs,
+            self.args.iodepth,
+            self.lane_count,
+            self.args.batch_size,
+            self.args.value_size,
+            self.args.nr_objects,
+            self._prepare_budget(),
+            self.args.write_objects,
+            self.args.memory_replica_num,
+            self.args.nof_replica_num,
+            self.args.verify,
+        )
+
+        phases: List[PhaseStats] = []
+        if self.args.scenario == "verify_write":
+            phases.append(
+                self._run_fixed_write(
+                    "write_verify",
+                    self._write_budget(),
+                    strict=True,
+                    write_scope="prepared",
+                )
+            )
+            self._phase_gap("verify_read")
+            phases.append(self._run_read_phase("verify_read", verify=True, sequential=True, loop=False))
+            return phases
+
+        if self.args.scenario == "fill":
+            phases.append(self._run_fixed_write("fill_write", self._write_budget(), strict=False))
+            return phases
+
+        if self.args.scenario == "write_perf":
+            total_objects = self._write_budget()
+            if self.args.runtime > 0:
+                phases.append(self._run_time_based_write("write_perf", total_objects))
+            else:
+                phases.append(self._run_fixed_write("write_perf", total_objects, strict=False))
+            return phases
+
+        if self.args.scenario == "read_perf":
+            prepared = self._maybe_prepare_dataset()
+            if prepared is not None:
+                phases.append(prepared)
+            phases.append(
+                self._run_read_phase(
+                    "read_perf",
+                    verify=self.args.verify,
+                    sequential=True,
+                    loop=(self.args.runtime > 0),
+                    runtime_sec=self.args.runtime,
+                )
+            )
+            return phases
+
+        prepared = self._maybe_prepare_dataset()
+        if prepared is not None:
+            phases.append(prepared)
+        phases.append(self._run_mixed_phase("mixed_rw", self._write_budget()))
+        return phases
+
+
+def log_overall_summary(phases: List[PhaseStats]) -> None:
+    overall = merge_stats("overall", phases)
+    LOG.info("=== overall summary ===")
+    log_phase_stats(overall)
+
+
+def main() -> int:
+    parser = build_parser()
+    args = parser.parse_args()
+    setup_logging(args.log_level)
+    runner: Optional[BenchmarkRunner] = None
+    try:
+        runner = BenchmarkRunner(args)
+        phases = runner.run()
+        log_overall_summary(phases)
+        if any(phase.verify_failures > 0 for phase in phases):
+            return 20
+        if args.verify and any(phase.misses > 0 for phase in phases if "read" in phase.name):
+            return 21
+        return 0
+    except KeyboardInterrupt:
+        LOG.warning("benchmark interrupted")
+        return 130
+    except Exception as exc:  # pragma: no cover - CLI entry path
+        LOG.exception("benchmark failed: %s", exc)
+        return 1
+    finally:
+        if runner is not None:
+            runner.close()
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())
diff --git a/mooncake-store/conf/master.yaml b/mooncake-store/conf/master.yaml
index 02d06b9e..2b772e3a 100644
--- a/mooncake-store/conf/master.yaml
+++ b/mooncake-store/conf/master.yaml
@@ -13,6 +13,10 @@ allow_evict_soft_pinned_objects: true
 eviction_ratio: 0.1
 eviction_high_watermark_ratio: 1.0
 
+enable_multi_tenants: false
+tenant_quota_connector_type: "file"
+tenant_quota_connector_uri: ""
+
 enable_ha: false
 etcd_endpoints: "http://localhost:2379"
 root_fs_dir: ""
diff --git a/mooncake-store/go/build.sh b/mooncake-store/go/build.sh
index c1cbf812..ac190e8b 100755
--- a/mooncake-store/go/build.sh
+++ b/mooncake-store/go/build.sh
@@ -40,8 +40,9 @@ CGO_LDFLAGS+=" -L${BUILD_DIR}/mooncake-store/src/cachelib_memory_allocator"
 CGO_LDFLAGS+=" -L${BUILD_DIR}/mooncake-transfer-engine/src"
 CGO_LDFLAGS+=" -L${BUILD_DIR}/mooncake-transfer-engine/src/common/base"
 CGO_LDFLAGS+=" -L${BUILD_DIR}/mooncake-common"
-CGO_LDFLAGS+=" -lmooncake_store -lcachelib_memory_allocator -ltransfer_engine -lbase -lasio"
-CGO_LDFLAGS+=" -lstdc++ -lnuma -lglog -lgflags -libverbs -ljsoncpp -lzstd -lcurl"
+CGO_LDFLAGS+=" -L${BUILD_DIR}/mooncake-common/src"
+CGO_LDFLAGS+=" -lmooncake_store -lcachelib_memory_allocator -ltransfer_engine -lbase -lasio -lmooncake_common -lxxhash -lyaml-cpp"
+CGO_LDFLAGS+=" -lstdc++ -lnuma -lglog -lgflags -libverbs -lmlx5 -ljsoncpp -lzstd -lcurl -lm"
 
 if [ -d "/usr/local/cuda/lib64" ]; then
     CGO_LDFLAGS+=" -L/usr/local/cuda/lib64 -lcudart"
diff --git a/mooncake-store/include/aligned_client_buffer.hpp b/mooncake-store/include/aligned_client_buffer.hpp
index 57f81e92..436fc9bb 100644
--- a/mooncake-store/include/aligned_client_buffer.hpp
+++ b/mooncake-store/include/aligned_client_buffer.hpp
@@ -56,6 +56,7 @@ class AlignedClientBufferAllocator : public ClientBufferAllocator {
     // Store whether we own the memory (for cleanup)
     bool owns_memory_;
     size_t allocated_size_;  // Store the actual allocated size for cleanup
+    std::string protocol_;
 };
 
 }  // namespace mooncake
diff --git a/mooncake-store/include/allocation_strategy.h b/mooncake-store/include/allocation_strategy.h
index 18d7f465..008c8722 100644
--- a/mooncake-store/include/allocation_strategy.h
+++ b/mooncake-store/include/allocation_strategy.h
@@ -119,6 +119,14 @@ class AllocatorManager {
     friend class SegmentSerializer;  // for fork serialize
 };
 
+class SsdMetricsProvider {
+   public:
+    virtual ~SsdMetricsProvider() = default;
+    virtual int64_t getSsdTotalCapacity(
+        const std::string& segment_name) const = 0;
+    virtual int64_t getSsdUsedBytes(const std::string& segment_name) const = 0;
+};
+
 /**
  * @brief Abstract interface for allocation strategy, responsible for
  *        allocating a slice (with one or more replicas) using available
@@ -163,7 +171,20 @@ class AllocationStrategy {
         const std::vector<std::string>& preferred_segments =
             std::vector<std::string>(),
         const std::set<std::string>& excluded_segments =
-            std::set<std::string>()) = 0;
+            std::set<std::string>(),
+        const ReplicaType replica_type = ReplicaType::MEMORY) = 0;
+
+    virtual tl::expected<std::vector<Replica>, ErrorCode> Allocate(
+        const AllocatorManager& allocator_manager, const size_t slice_length,
+        const size_t replica_num,
+        const std::vector<std::string>& preferred_segments,
+        const std::set<std::string>& excluded_segments,
+        const ReplicaType replica_type,
+        const SsdMetricsProvider* ssd_provider) {
+        (void)ssd_provider;
+        return Allocate(allocator_manager, slice_length, replica_num,
+                        preferred_segments, excluded_segments, replica_type);
+    }
 
     /**
      * @brief Allocate one replica from the specified segment.
@@ -208,7 +229,8 @@ class RandomAllocationStrategy : public AllocationStrategy {
         const std::vector<std::string>& preferred_segments =
             std::vector<std::string>(),
         const std::set<std::string>& excluded_segments =
-            std::set<std::string>()) {
+            std::set<std::string>(),
+        const ReplicaType replica_type = ReplicaType::MEMORY) override {
         // Validate input parameters
         if (slice_length == 0 || replica_num == 0) {
             return tl::make_unexpected(ErrorCode::INVALID_PARAMS);
@@ -236,7 +258,7 @@ class RandomAllocationStrategy : public AllocationStrategy {
                                          slice_length, generator);
             if (buffer) {
                 replicas.emplace_back(std::move(buffer),
-                                      ReplicaStatus::PROCESSING);
+                                      ReplicaStatus::PROCESSING, replica_type);
                 return replicas;
             }
             return tl::make_unexpected(ErrorCode::NO_AVAILABLE_HANDLE);
@@ -256,7 +278,7 @@ class RandomAllocationStrategy : public AllocationStrategy {
                                          slice_length, generator);
             if (buffer) {
                 replicas.emplace_back(std::move(buffer),
-                                      ReplicaStatus::PROCESSING);
+                                      ReplicaStatus::PROCESSING, replica_type);
                 if (replicas.size() == replica_num) {
                     return replicas;
                 }
@@ -289,7 +311,7 @@ class RandomAllocationStrategy : public AllocationStrategy {
                                          slice_length, generator);
             if (buffer) {
                 replicas.emplace_back(std::move(buffer),
-                                      ReplicaStatus::PROCESSING);
+                                      ReplicaStatus::PROCESSING, replica_type);
                 // Nit: no need to insert names[index] into used_segments here
                 // because we only traverse all names once, thus there is no
                 // chance to try allocating from a segment for the second time.
@@ -345,8 +367,9 @@ class RandomAllocationStrategy : public AllocationStrategy {
         // Randomly select a start point to distribute
         // allocations across all segments
         std::uniform_int_distribution<size_t> dist(0, num_segs - 1);
-        size_t seg_offset = dist(generator);
-        for (size_t i = 0; i < num_segs; i++) {
+        size_t seg_offset =
+            dist(generator);  // select a start segment to place replica
+        for (size_t i = 0; i < num_segs; i++) {  // only allocate one replica
             auto& allocator = (*allocators)[(i + seg_offset) % num_segs];
             if (auto buffer = allocator->allocate(slice_length)) {
                 return buffer;
@@ -386,7 +409,8 @@ class FreeRatioFirstAllocationStrategy : public RandomAllocationStrategy {
         const std::vector<std::string>& preferred_segments =
             std::vector<std::string>(),
         const std::set<std::string>& excluded_segments =
-            std::set<std::string>()) override {
+            std::set<std::string>(),
+        const ReplicaType replica_type = ReplicaType::MEMORY) override {
         if (slice_length == 0 || replica_num == 0) {
             return tl::make_unexpected(ErrorCode::INVALID_PARAMS);
         }
@@ -413,7 +437,7 @@ class FreeRatioFirstAllocationStrategy : public RandomAllocationStrategy {
                                          slice_length, generator);
             if (buffer) {
                 replicas.emplace_back(std::move(buffer),
-                                      ReplicaStatus::PROCESSING);
+                                      ReplicaStatus::PROCESSING, replica_type);
                 used_segments.insert(preferred_segment);
                 if (replicas.size() == replica_num) {
                     return replicas;
@@ -469,7 +493,7 @@ class FreeRatioFirstAllocationStrategy : public RandomAllocationStrategy {
                                          generator);
             if (buffer) {
                 replicas.emplace_back(std::move(buffer),
-                                      ReplicaStatus::PROCESSING);
+                                      ReplicaStatus::PROCESSING, replica_type);
                 used_segments.insert(name);
             }
         }
@@ -499,7 +523,7 @@ class FreeRatioFirstAllocationStrategy : public RandomAllocationStrategy {
                                          slice_length, generator);
             if (buffer) {
                 replicas.emplace_back(std::move(buffer),
-                                      ReplicaStatus::PROCESSING);
+                                      ReplicaStatus::PROCESSING, replica_type);
                 used_segments.insert(names[index]);
             }
         }
@@ -534,6 +558,155 @@ class FreeRatioFirstAllocationStrategy : public RandomAllocationStrategy {
     }
 };
 
+class SsdFreeRatioFirstAllocationStrategy : public RandomAllocationStrategy {
+   public:
+    SsdFreeRatioFirstAllocationStrategy() = default;
+
+    tl::expected<std::vector<Replica>, ErrorCode> Allocate(
+        const AllocatorManager& allocator_manager, const size_t slice_length,
+        const size_t replica_num,
+        const std::vector<std::string>& preferred_segments,
+        const std::set<std::string>& excluded_segments,
+        const ReplicaType replica_type,
+        const SsdMetricsProvider* ssd_provider) override {
+        if (slice_length == 0 || replica_num == 0) {
+            return tl::make_unexpected(ErrorCode::INVALID_PARAMS);
+        }
+
+        const auto& names = allocator_manager.getNames();
+        if (names.empty()) {
+            return tl::make_unexpected(ErrorCode::NO_AVAILABLE_HANDLE);
+        }
+
+        static thread_local std::mt19937 generator(std::random_device{}());
+
+        std::vector<Replica> replicas;
+        replicas.reserve(replica_num);
+        std::set<std::string> used_segments;
+
+        // Handle preferred segments first
+        for (const auto& preferred_segment : preferred_segments) {
+            if (excluded_segments.contains(preferred_segment) ||
+                used_segments.contains(preferred_segment)) {
+                continue;
+            }
+
+            auto buffer = allocateSingle(allocator_manager, preferred_segment,
+                                         slice_length, generator);
+            if (buffer) {
+                replicas.emplace_back(std::move(buffer),
+                                      ReplicaStatus::PROCESSING, replica_type);
+                used_segments.insert(preferred_segment);
+                if (replicas.size() == replica_num) {
+                    return replicas;
+                }
+            }
+        }
+
+        const size_t remaining = replica_num - replicas.size();
+
+        // Sample candidates and sort by SSD free ratio
+        size_t sample_count =
+            std::min(kCandidateMultiplier * remaining, names.size());
+
+        std::uniform_int_distribution<size_t> start_dist(0, names.size() - 1);
+        size_t start_idx = start_dist(generator);
+
+        struct Candidate {
+            size_t name_idx;
+            double ssd_free_ratio;
+        };
+        std::vector<Candidate> candidates;
+        candidates.reserve(sample_count);
+
+        for (size_t i = 0; i < sample_count; ++i) {
+            size_t idx = (start_idx + i) % names.size();
+            const auto& name = names[idx];
+
+            if (excluded_segments.contains(name) ||
+                used_segments.contains(name)) {
+                continue;
+            }
+
+            double ssd_free_ratio = getSegmentSsdFreeRatio(name, ssd_provider);
+            candidates.push_back({idx, ssd_free_ratio});
+        }
+
+        std::sort(candidates.begin(), candidates.end(),
+                  [](const Candidate& a, const Candidate& b) {
+                      return a.ssd_free_ratio > b.ssd_free_ratio;
+                  });
+
+        for (const auto& candidate : candidates) {
+            if (replicas.size() >= replica_num) {
+                break;
+            }
+
+            const auto& name = names[candidate.name_idx];
+            auto buffer = allocateSingle(allocator_manager, name, slice_length,
+                                         generator);
+            if (buffer) {
+                replicas.emplace_back(std::move(buffer),
+                                      ReplicaStatus::PROCESSING, replica_type);
+                used_segments.insert(name);
+            }
+        }
+
+        if (replicas.size() >= replica_num) {
+            return replicas;
+        }
+
+        // Fallback: Random allocation for remaining replicas
+        std::uniform_int_distribution<size_t> distribution(0, names.size() - 1);
+        size_t fallback_idx = distribution(generator);
+        const size_t max_retry = std::min(kMaxRetryLimit, names.size());
+        size_t try_count = 0;
+
+        while (replicas.size() < replica_num && try_count < max_retry) {
+            auto index = fallback_idx % names.size();
+            fallback_idx++;
+            try_count++;
+
+            const auto& name = names[index];
+
+            if (excluded_segments.contains(name) ||
+                used_segments.contains(name)) {
+                continue;
+            }
+
+            auto buffer = allocateSingle(allocator_manager, name, slice_length,
+                                         generator);
+            if (buffer) {
+                replicas.emplace_back(std::move(buffer),
+                                      ReplicaStatus::PROCESSING, replica_type);
+                used_segments.insert(name);
+            }
+        }
+
+        if (replicas.empty()) {
+            return tl::make_unexpected(ErrorCode::NO_AVAILABLE_HANDLE);
+        }
+        return replicas;
+    }
+
+    using AllocationStrategy::Allocate;
+
+   private:
+    static constexpr size_t kMaxRetryLimit = 100;
+    static constexpr size_t kCandidateMultiplier = 6;
+
+    double getSegmentSsdFreeRatio(
+        const std::string& name, const SsdMetricsProvider* ssd_provider) const {
+        if (!ssd_provider) return 1.0;
+        int64_t total = ssd_provider->getSsdTotalCapacity(name);
+        if (total <= 0) return 1.0;
+        int64_t used = ssd_provider->getSsdUsedBytes(name);
+        used = std::clamp<int64_t>(used, 0, total);
+        int64_t free_bytes = total - used;
+        return static_cast<double>(free_bytes) / static_cast<double>(total);
+    }
+};
+
 class CxlAllocationStrategy : public AllocationStrategy {
    public:
     CxlAllocationStrategy() = default;
@@ -543,7 +716,8 @@ class CxlAllocationStrategy : public AllocationStrategy {
         const std::vector<std::string>& preferred_segments =
             std::vector<std::string>(),
         const std::set<std::string>& excluded_segments =
-            std::set<std::string>()) {
+            std::set<std::string>(),
+        const ReplicaType replica_type = ReplicaType::MEMORY) override {
         if (slice_length == 0 || replica_num == 0) {
             return tl::make_unexpected(ErrorCode::INVALID_PARAMS);
         }
@@ -579,7 +753,8 @@ class CxlAllocationStrategy : public AllocationStrategy {
         }
 
         buffer->change_to_cxl(cxl_segment_name);
-        replicas.emplace_back(std::move(buffer), ReplicaStatus::PROCESSING);
+        replicas.emplace_back(std::move(buffer), ReplicaStatus::PROCESSING,
+                              replica_type);
 
         VLOG(1) << "Successfully allocated " << replicas.size()
                 << " CXL replica.";
@@ -605,6 +780,8 @@ inline std::shared_ptr<AllocationStrategy> CreateAllocationStrategy(
             return std::make_shared<FreeRatioFirstAllocationStrategy>();
         case AllocationStrategyType::CXL:
             return std::make_shared<CxlAllocationStrategy>();
+        case AllocationStrategyType::SSD_FREE_RATIO_FIRST:
+            return std::make_shared<SsdFreeRatioFirstAllocationStrategy>();
         default:
             return std::make_shared<RandomAllocationStrategy>();
     }
diff --git a/mooncake-store/include/allocator.h b/mooncake-store/include/allocator.h
index 629f53ae..e28b7d5c 100644
--- a/mooncake-store/include/allocator.h
+++ b/mooncake-store/include/allocator.h
@@ -15,6 +15,17 @@ using facebook::cachelib::PoolId;
 
 namespace mooncake {
 
+/**
+ * @brief Type of buffer allocator used in the system
+ */
+enum class ReplicaType {
+    MEMORY,      // Memory replica
+    DISK,        // Disk replica
+    LOCAL_DISK,  // Local disk replica
+    NOF_SSD,     // Nvme-oF SSD replica
+    ALL,         // All memory and NoF replicas in put finalize path
+};
+
 // Constant for unknown free space in allocators that don't track it precisely
 static constexpr size_t kAllocatorUnknownFreeSpace =
     std::numeric_limits<size_t>::max();
@@ -142,7 +153,8 @@ class CachelibBufferAllocator
       public std::enable_shared_from_this<CachelibBufferAllocator> {
    public:
     CachelibBufferAllocator(std::string segment_name, size_t base, size_t size,
-                            std::string transport_endpoint);
+                            std::string transport_endpoint,
+                            ReplicaType replica_type = ReplicaType::MEMORY);
 
     ~CachelibBufferAllocator() override;
 
@@ -173,6 +185,7 @@ class CachelibBufferAllocator
     const size_t total_size_;
     std::atomic_size_t cur_size_;
     const std::string transport_endpoint_;
+    const ReplicaType replica_type_;
 
     // metrics - removed allocated_bytes_ member
     // ylt::metric::gauge_t* allocated_bytes_{nullptr};
@@ -193,7 +206,8 @@ class OffsetBufferAllocator
       public std::enable_shared_from_this<OffsetBufferAllocator> {
    public:
     OffsetBufferAllocator(std::string segment_name, size_t base, size_t size,
-                          std::string transport_endpoint);
+                          std::string transport_endpoint,
+                          ReplicaType replica_type = ReplicaType::MEMORY);
 
     ~OffsetBufferAllocator() override;
 
@@ -226,6 +240,7 @@ class OffsetBufferAllocator
     const size_t total_size_;
     std::atomic_size_t cur_size_;
     const std::string transport_endpoint_;
+    const ReplicaType replica_type_;
 
     // offset allocator implementation
     std::shared_ptr<offset_allocator::OffsetAllocator> offset_allocator_;
diff --git a/mooncake-store/include/client_buffer.hpp b/mooncake-store/include/client_buffer.hpp
index 44bb6f06..b9f84089 100644
--- a/mooncake-store/include/client_buffer.hpp
+++ b/mooncake-store/include/client_buffer.hpp
@@ -30,7 +30,7 @@ class ClientBufferAllocator
     // Create for heap-allocated memory
     static std::shared_ptr<ClientBufferAllocator> create(
         size_t size, const std::string& protocol = "",
-        bool use_hugepage = false);
+        bool use_hugepage = false, bool use_spdk_dma = false);
 
     // Create for shared memory
     static std::shared_ptr<ClientBufferAllocator> create(
@@ -57,17 +57,18 @@ class ClientBufferAllocator
     ClientBufferAllocator(void* addr, size_t size, const std::string& protocol);
 
     void* buffer_;
+    size_t buffer_size_;
     bool use_hugepage_ = false;
 
    private:
     ClientBufferAllocator(size_t size, const std::string& protocol,
-                          bool use_hugepage);
+                          bool use_hugepage, bool use_spdk_dma);
 
     std::shared_ptr<offset_allocator::OffsetAllocator> allocator_;
 
     std::string protocol;
-    size_t buffer_size_;
     bool is_external_memory_ = false;
+    bool use_spdk_dma_ = false;
 };
 
 /**
diff --git a/mooncake-store/include/client_service.h b/mooncake-store/include/client_service.h
index 445d75cd..0094f4f1 100644
--- a/mooncake-store/include/client_service.h
+++ b/mooncake-store/include/client_service.h
@@ -2,6 +2,8 @@
 
 #include <atomic>
 #include <boost/functional/hash.hpp>
+#include <condition_variable>
+#include <functional>
 #include <memory>
 #include <mutex>
 #include <optional>
@@ -60,9 +62,10 @@ class QueryResult {
  */
 class Client {
    public:
-    ~Client();
+    virtual ~Client();
 
     const UUID& getClientId() const { return client_id_; }
+    const std::string& tenant_id() const { return master_client_.tenant_id(); }
 
     /**
      * @brief Creates and initializes a new Client instance
@@ -84,7 +87,8 @@ class Client {
         const std::optional<std::string>& device_names = std::nullopt,
         const std::string& master_server_entry = kDefaultMasterAddress,
         const std::shared_ptr<TransferEngine>& transfer_engine = nullptr,
-        std::map<std::string, std::string> labels = {});
+        std::map<std::string, std::string> labels = {},
+        const std::string& tenant_id = "default");
 
     /**
      * @brief Retrieves data for a given key
@@ -142,6 +146,9 @@ class Client {
      */
     std::vector<tl::expected<QueryResult, ErrorCode>> BatchQuery(
         const std::vector<std::string>& object_keys);
+    std::vector<tl::expected<QueryResult, ErrorCode>> BatchQuery(
+        const std::vector<std::string>& object_keys,
+        const std::string& tenant_id);
 
     /**
      * @brief Batch clear KV cache for specified object keys on a specific
@@ -272,9 +279,15 @@ class Client {
      */
     tl::expected<void, ErrorCode> EvictDiskReplica(const std::string& key,
                                                    ReplicaType replica_type);
+    tl::expected<void, ErrorCode> EvictDiskReplica(const std::string& key,
+                                                   const std::string& tenant_id,
+                                                   ReplicaType replica_type);
 
     std::vector<tl::expected<void, ErrorCode>> BatchEvictDiskReplica(
         const std::vector<std::string>& keys, ReplicaType replica_type);
+    std::vector<tl::expected<void, ErrorCode>> BatchEvictDiskReplica(
+        const std::vector<std::string>& keys, const std::string& tenant_id,
+        ReplicaType replica_type);
 
     /**
      * @brief Registers a memory segment to master for allocation
@@ -295,6 +308,23 @@ class Client {
     tl::expected<void, ErrorCode> UnmountSegment(const void* buffer,
                                                  size_t size);
 
+    /**
+     * @brief Mounts a memory segment and returns its generated Segment UUID.
+     *        Logic is identical to MountSegment, but returns the segment id.
+     */
+    tl::expected<UUID, ErrorCode> MountSegmentAndGetId(
+        const void* buffer, size_t size, const std::string& protocol = "tcp",
+        const std::string& location = kWildcardLocation);
+
+    /**
+     * @brief Unmounts a segment by its UUID.
+     *        Logic is identical to UnmountSegment, but looks up by id.
+     * @param grace_period_ms 0 = immediate unmount (legacy behavior).
+     */
+    tl::expected<void, ErrorCode> UnmountSegmentById(
+        const UUID& segment_id, uint64_t grace_period_ms = 0,
+        std::function<void(const UUID&)> cleanup_callback = {});
+
     /**
      * @brief Registers memory buffer with TransferEngine for data transfer
      * @param addr Memory address to register
@@ -342,6 +372,9 @@ class Client {
      */
     tl::expected<UUID, ErrorCode> CreateCopyTask(
         const std::string& key, const std::vector<std::string>& targets);
+    tl::expected<UUID, ErrorCode> CreateCopyTask(
+        const std::string& key, const std::string& tenant_id,
+        const std::vector<std::string>& targets);
 
     /**
      * @brief Create a move task to move an object's replica from source segment
@@ -355,6 +388,10 @@ class Client {
     tl::expected<UUID, ErrorCode> CreateMoveTask(const std::string& key,
                                                  const std::string& source,
                                                  const std::string& target);
+    tl::expected<UUID, ErrorCode> CreateMoveTask(const std::string& key,
+                                                 const std::string& tenant_id,
+                                                 const std::string& source,
+                                                 const std::string& target);
 
     /**
      * @brief Query a task by task id
@@ -381,12 +418,66 @@ class Client {
      * set of non-offloaded objects.
      * @param enable_offloading Indicates whether offloading is enabled for this
      * segment.
-     * @param offloading_objects On return, contains a map from object key to
-     * size (in bytes) for all objects that require offload.
+     * @param offloading_objects On return, contains the tenant-scoped object
+     * tasks that require offload.
      */
     tl::expected<void, ErrorCode> OffloadObjectHeartbeat(
         bool enable_offloading,
-        std::unordered_map<std::string, int64_t>& offloading_objects);
+        std::vector<OffloadTaskItem>& offloading_objects);
+
+    tl::expected<void, ErrorCode> ReportSsdCapacity(
+        int64_t ssd_total_capacity_bytes);
+
+    /**
+     * @brief Heartbeat-driven pull of pending L2->L1 promotion work for this
+     * client. Mirror of OffloadObjectHeartbeat. Returns tenant-scoped tasks the
+     * caller (FileStorage) must read from local SSD and stage as MEMORY
+     * replicas via PromotionAllocStart + NotifyPromotionSuccess.
+     */
+    // Virtual to enable subclassing in unit tests.
+    virtual tl::expected<void, ErrorCode> PromotionObjectHeartbeat(
+        std::vector<PromotionTaskItem>& promotion_objects);
+
+    /**
+     * @brief Stage a PROCESSING MEMORY replica for an existing key during
+     * L2->L1 promotion. Returns the new replica's descriptor that the caller
+     * writes via Transfer Engine before calling NotifyPromotionSuccess.
+     */
+    virtual tl::expected<PromotionAllocStartResponse, ErrorCode>
+    PromotionAllocStart(const std::string& key, uint64_t size,
+                        const std::vector<std::string>& preferred_segments);
+    virtual tl::expected<PromotionAllocStartResponse, ErrorCode>
+    PromotionAllocStart(const std::string& key, const std::string& tenant_id,
+                        uint64_t size,
+                        const std::vector<std::string>& preferred_segments);
+
+    /**
+     * @brief Commit a staged MEMORY replica to COMPLETE; called after the
+     * client has written the bytes via Transfer Engine.
+     */
+    virtual tl::expected<void, ErrorCode> NotifyPromotionSuccess(
+        const std::string& key);
+    virtual tl::expected<void, ErrorCode> NotifyPromotionSuccess(
+        const std::string& key, const std::string& tenant_id);
+
+    /**
+     * @brief Release master-side promotion task after a client-side failure
+     * between PromotionAllocStart and the transfer's completion. Idempotent.
+     */
+    virtual tl::expected<void, ErrorCode> NotifyPromotionFailure(
+        const std::string& key);
+    virtual tl::expected<void, ErrorCode> NotifyPromotionFailure(
+        const std::string& key, const std::string& tenant_id);
+
+    /**
+     * @brief Write `slices` into the memory replica described by
+     * `memory_descriptor` via Transfer Engine. Used by FileStorage to fill a
+     * PROCESSING memory replica staged by PromotionAllocStart before calling
+     * NotifyPromotionSuccess.
+     */
+    virtual ErrorCode PromotionWrite(
+        const Replica::Descriptor& memory_descriptor,
+        std::vector<Slice>& slices);
 
     /**
      * @brief Performs a batched read of multiple objects using a
@@ -403,7 +494,8 @@ class Client {
         const std::string& transfer_engine_addr,
         const std::vector<std::string>& keys,
         const std::vector<uintptr_t>& pointers,
-        const std::unordered_map<std::string, Slice>& batch_slices);
+        const std::unordered_map<std::string, std::vector<Slice>>&
+            batch_slices);
 
     /**
      * @brief Notifies the master that offloading of specified objects has
@@ -416,6 +508,9 @@ class Client {
     tl::expected<void, ErrorCode> NotifyOffloadSuccess(
         const std::vector<std::string>& keys,
         const std::vector<StorageObjectMetadata>& metadatas);
+    tl::expected<void, ErrorCode> NotifyOffloadSuccess(
+        const std::vector<OffloadTaskItem>& tasks,
+        const std::vector<StorageObjectMetadata>& metadatas);
 
     /**
      * @brief Fetch tasks assigned to a client
@@ -474,11 +569,34 @@ class Client {
         return transfer_engine_->getLocalIpAndPort();
     }
 
+    [[nodiscard]] const std::string& GetProtocol() const { return protocol_; }
+
+    /**
+     * @brief Get the endpoint address for segment operations.
+     * @return For P2PHANDSHAKE mode, returns the actual RPC endpoint (IP:Port).
+     *         For other modes, returns the logical local hostname used for
+     * segment registration.
+     */
+    [[nodiscard]] std::string GetSegmentEndpoint() {
+        return (metadata_connstring_ == P2PHANDSHAKE) ? GetTransportEndpoint()
+                                                      : local_hostname_;
+    }
+
     // Return sorted NUMA node IDs that have at least one RDMA NIC.
     [[nodiscard]] std::vector<int> GetNicNumaNodes() const;
 
     tl::expected<Replica::Descriptor, ErrorCode> GetPreferredReplica(
         const std::vector<Replica::Descriptor>& replica_list);
+
+    std::unordered_set<std::string> GetLocalEndpoints() const {
+        std::lock_guard<std::mutex> lock(mounted_segments_mutex_);
+        std::unordered_set<std::string> endpoints;
+        for (const auto& [segment_id, segment] : mounted_segments_) {
+            endpoints.insert(segment.te_endpoint);
+        }
+        return endpoints;
+    }
+
     /**
      * @brief Check if local hot cache is enabled
      * @return true if hot cache is enabled, false otherwise
@@ -531,14 +649,17 @@ class Client {
 
     bool IsReplicaOnLocalMemory(const Replica::Descriptor& replica);
 
-   private:
+   protected:
     /**
-     * @brief Private constructor to enforce creation through Create() method
+     * @brief Constructor exposed to subclasses for testing only; production
+     * code must go through Create().
      */
     Client(const std::string& local_hostname,
            const std::string& metadata_connstring, const std::string& protocol,
-           const std::map<std::string, std::string>& labels = {});
+           const std::map<std::string, std::string>& labels = {},
+           const std::string& tenant_id = "default");
 
+   private:
     /**
      * @brief Internal helper functions for initialization and data transfer
      */
@@ -580,6 +701,11 @@ class Client {
      */
     ErrorCode InitLocalHotCache();
 
+    /**
+     * @brief Unregister local hot cache backing memory from TransferEngine.
+     */
+    void UnregisterLocalHotCacheMemory();
+
     /**
      * @brief Read MC_STORE_LOCAL_HOT_CACHE_SIZE from environment variable
      * @return Cache size in bytes, or 0 if not set or invalid
@@ -663,9 +789,32 @@ class Client {
     std::unique_ptr<TransferSubmitter> transfer_submitter_;
 
     // Mutex to protect mounted_segments_
-    std::mutex mounted_segments_mutex_;
+    mutable std::mutex mounted_segments_mutex_;
     std::unordered_map<UUID, Segment, boost::hash<UUID>> mounted_segments_;
 
+    // Segments in graceful unmount: readable by remote peers, not allocatable
+    // locally. TE MR remains registered until master confirms removal.
+    std::unordered_map<UUID, Segment, boost::hash<UUID>>
+        gracefully_unmounting_segments_;
+    std::unordered_map<UUID, std::function<void(const UUID&)>,
+                       boost::hash<UUID>>
+        graceful_unmount_cleanup_callbacks_;
+
+    /**
+     * @brief Internal helper to unmount a segment by iterator.
+     *        Caller must hold mounted_segments_mutex_.
+     */
+    tl::expected<void, ErrorCode> UnmountSegmentImpl(
+        std::unordered_map<UUID, Segment, boost::hash<UUID>>::iterator it);
+
+    void StartGracefulUnmountTimer(const UUID& segment_id,
+                                   uint64_t grace_period_ms);
+    void OnGracefulUnmountTimer(const UUID& segment_id, int retry_left);
+    bool WaitForGracefulUnmountDelay(std::chrono::milliseconds delay);
+    std::mutex graceful_unmount_timer_mutex_;
+    std::condition_variable graceful_unmount_timer_cv_;
+    bool graceful_unmount_timer_stopping_{false};
+
     // Configuration
     const std::string local_hostname_;
     const std::string metadata_connstring_;
@@ -690,6 +839,8 @@ class Client {
     std::thread task_poll_thread_;
     std::atomic<bool> task_poll_running_{false};
     std::atomic<bool> last_ping_success_{false};
+    std::atomic<bool> segment_desc_publish_pending_{false};
+    std::atomic<bool> rpc_meta_publish_pending_{false};
     ErrorCode SwitchLeader(const ha::MasterView& target_view);
     void LeaderMonitorThreadMain();
     void StorageHeartbeatThreadMain();
@@ -726,6 +877,10 @@ class Client {
     tl::expected<void, ErrorCode> Copy(const std::string& key,
                                        const std::string& source,
                                        const std::vector<std::string>& targets);
+    tl::expected<void, ErrorCode> Copy(const std::string& key,
+                                       const std::string& tenant_id,
+                                       const std::string& source,
+                                       const std::vector<std::string>& targets);
 
     /**
      * @brief Move an object's replica from source segment to target segment
@@ -737,6 +892,10 @@ class Client {
     tl::expected<void, ErrorCode> Move(const std::string& key,
                                        const std::string& source,
                                        const std::string& target);
+    tl::expected<void, ErrorCode> Move(const std::string& key,
+                                       const std::string& tenant_id,
+                                       const std::string& source,
+                                       const std::string& target);
 
     // Task thread pool for async task execution
     ThreadPool task_thread_pool_;
@@ -751,6 +910,7 @@ class Client {
     // Local hot cache and async handler
     std::shared_ptr<LocalHotCache> hot_cache_;
     std::unique_ptr<LocalHotCacheHandler> hot_cache_handler_;
+    bool hot_cache_memory_registered_{false};
 
     // Frequency admission: only cache keys whose CMS count >= threshold
     std::unique_ptr<CountMinSketch> admission_sketch_;
diff --git a/mooncake-store/include/deadline_scheduler.h b/mooncake-store/include/deadline_scheduler.h
new file mode 100644
index 00000000..0d428cde
--- /dev/null
+++ b/mooncake-store/include/deadline_scheduler.h
@@ -0,0 +1,140 @@
+#pragma once
+
+#include <algorithm>
+#include <atomic>
+#include <chrono>
+#include <condition_variable>
+#include <functional>
+#include <mutex>
+#include <thread>
+#include <utility>
+#include <vector>
+
+namespace mooncake {
+
+template <typename Id>
+class DeadlineScheduler {
+   public:
+    using Clock = std::chrono::steady_clock;
+    using Callback = std::function<void(const Id& id)>;
+
+    explicit DeadlineScheduler(Callback callback)
+        : callback_(std::move(callback)) {}
+
+    ~DeadlineScheduler() { Stop(); }
+
+    void Schedule(Id id, Clock::time_point deadline) {
+        {
+            std::lock_guard<std::mutex> lock(mutex_);
+            if (stopping_) {
+                return;
+            }
+            PushRecord(Record{std::move(id), deadline});
+            if (!timer_running_.load()) {
+                timer_running_.store(true);
+                timer_thread_ = std::thread([this]() { this->TimerLoop(); });
+            }
+        }
+        timer_cv_.notify_one();
+    }
+
+    template <typename Predicate>
+    void RemoveIf(Predicate predicate) {
+        std::lock_guard<std::mutex> lock(mutex_);
+        auto erase_begin = std::remove_if(
+            records_.begin(), records_.end(),
+            [&](const Record& rec) { return predicate(rec.id); });
+        if (erase_begin != records_.end()) {
+            records_.erase(erase_begin, records_.end());
+            std::make_heap(records_.begin(), records_.end(),
+                           std::greater<Record>{});
+        }
+        timer_cv_.notify_one();
+    }
+
+    void Stop() {
+        {
+            std::lock_guard<std::mutex> lock(mutex_);
+            stopping_ = true;
+            timer_running_.store(false);
+        }
+        timer_cv_.notify_all();
+        if (timer_thread_.joinable()) {
+            timer_thread_.join();
+        }
+    }
+
+   private:
+    void TimerLoop() {
+        while (timer_running_.load()) {
+            std::vector<Record> expired;
+            {
+                std::unique_lock<std::mutex> lock(mutex_);
+                if (records_.empty()) {
+                    timer_cv_.wait(lock, [this]() {
+                        return !timer_running_.load() || !records_.empty();
+                    });
+                    if (!timer_running_.load()) break;
+                    continue;
+                }
+
+                auto now = Clock::now();
+                auto next_deadline = NextRecord().deadline;
+                if (next_deadline > now) {
+                    timer_cv_.wait_until(
+                        lock, next_deadline, [this, next_deadline]() {
+                            return !timer_running_.load() || records_.empty() ||
+                                   NextRecord().deadline < next_deadline;
+                        });
+                    if (!timer_running_.load()) break;
+                    continue;
+                }
+
+                while (!records_.empty() && NextRecord().deadline <= now) {
+                    expired.push_back(PopNextRecord());
+                }
+            }
+
+            // Run callbacks with the scheduler mutex released: callbacks may
+            // acquire other locks, so holding mutex_ here risks lock-order
+            // inversion deadlocks.
+            for (auto& rec : expired) {
+                callback_(rec.id);
+            }
+        }
+    }
+
+    struct Record {
+        Id id;
+        Clock::time_point deadline;
+
+        bool operator>(const Record& other) const {
+            return deadline > other.deadline;
+        }
+    };
+
+    void PushRecord(Record record) {
+        records_.push_back(std::move(record));
+        std::push_heap(records_.begin(), records_.end(),
+                       std::greater<Record>{});
+    }
+
+    Record PopNextRecord() {
+        std::pop_heap(records_.begin(), records_.end(), std::greater<Record>{});
+        auto record = std::move(records_.back());
+        records_.pop_back();
+        return record;
+    }
+
+    const Record& NextRecord() const { return records_.front(); }
+
+    Callback callback_;
+    std::mutex mutex_;
+    std::vector<Record> records_;
+    std::thread timer_thread_;
+    std::atomic<bool> timer_running_{false};
+    bool stopping_{false};
+    std::condition_variable timer_cv_;
+};
+
+}  // namespace mooncake
diff --git a/mooncake-store/include/dummy_client.h b/mooncake-store/include/dummy_client.h
index 13128c43..d5387088 100644
--- a/mooncake-store/include/dummy_client.h
+++ b/mooncake-store/include/dummy_client.h
@@ -2,12 +2,14 @@
 
 #include <atomic>
 #include <csignal>
+#include <mutex>
+#include <shared_mutex>
+#include <unordered_map>
 #include <ylt/coro_rpc/coro_rpc_client.hpp>
 
-#include "pyclient.h"
-#include "real_client.h"
-#include "shm_helper.h"
 #include "client_metric.h"
+#include "pyclient.h"
+#include "shm_helper.h"
 #include <memory>
 
 namespace mooncake {
@@ -27,7 +29,8 @@ class DummyClient : public PyClient {
                    const std::shared_ptr<TransferEngine> &transfer_engine,
                    const std::string &ipc_socket_path,
                    bool enable_ssd_offload = false,
-                   const std::string &ssd_offload_path = "") {
+                   const std::string &ssd_offload_path = "",
+                   const std::string &tenant_id = "default") {
         // Dummy client does not support real setup
         return -1;
     };
@@ -58,8 +61,11 @@ class DummyClient : public PyClient {
         const std::vector<std::vector<std::string>> &all_keys,
         const std::vector<std::vector<std::vector<size_t>>> &all_dst_offsets,
         const std::vector<std::vector<std::vector<size_t>>> &all_src_offsets,
-        const std::vector<std::vector<std::vector<size_t>>> &all_sizes)
-        override;
+        const std::vector<std::vector<std::vector<size_t>>> &all_sizes,
+        const QueryResultCache *query_result_cache = nullptr) override;
+
+    std::vector<tl::expected<QueryResult, ErrorCode>> batch_query(
+        const std::vector<std::string> &keys) override;
 
     std::vector<int64_t> batch_get_into(const std::vector<std::string> &keys,
                                         const std::vector<void *> &buffers,
@@ -179,6 +185,15 @@ class DummyClient : public PyClient {
     int register_shm_via_ipc(const ShmHelper::ShmSegment *shm,
                              bool is_local = false);
 
+#if defined(USE_ASCEND_DIRECT)
+    int register_device_buffer_for_reconnect(void *buffer, size_t size);
+
+    int unregister_device_buffer_for_reconnect(void *buffer);
+
+    [[nodiscard]] std::vector<ShmHelper::ShmSegment>
+    get_registered_device_buffers() const;
+#endif
+
     /**
      * @brief Generic RPC invocation helper for single-result operations
      * @tparam ServiceMethod Pointer to WrappedMasterService member function
@@ -276,6 +291,11 @@ class DummyClient : public PyClient {
     void ping_thread_main();
     std::atomic<bool> connected_{false};
 
+#if defined(USE_ASCEND_DIRECT)
+    mutable std::mutex registered_device_buffers_mutex_;
+    std::unordered_map<uint64_t, size_t> registered_device_buffers_;
+#endif
+
     // Ascend physical device id for dummy-real RPC to real, set in setup_dummy
     int32_t device_id_ = 0;
 
diff --git a/mooncake-store/include/engram/engram_store.h b/mooncake-store/include/engram/engram_store.h
new file mode 100644
index 00000000..b655bd60
--- /dev/null
+++ b/mooncake-store/include/engram/engram_store.h
@@ -0,0 +1,83 @@
+#pragma once
+
+#include <cstdint>
+#include <memory>
+#include <string>
+#include <vector>
+
+#include "engram/engram_store_config.h"
+
+namespace mooncake {
+
+class PyClient;
+
+namespace engram {
+
+/**
+ * Mooncake backend for Engram embedding tables.
+ *
+ * This class intentionally owns only storage-side concerns:
+ * - per-head table naming/layout in Mooncake Store
+ * - batch populate / remove
+ * - row-id based embedding lookup
+ *
+ * It does not implement tokenizer compression, N-gram hashing, routing,
+ * gating, convolution, or any other model-side Engram logic.
+ */
+class EngramStore {
+   public:
+    EngramStore(int layer_id, const EngramStoreConfig& config,
+                std::shared_ptr<PyClient> store = nullptr);
+
+    ~EngramStore() = default;
+
+    /**
+     * Lookup embedding rows for a batch of precomputed row IDs.
+     * @param row_ids [B, L, H] precomputed row IDs, where H == num_heads
+     * @param output [B, L, H, D] output buffer
+     * @param output_size Size of output buffer in bytes
+     * @return 0 on success, negative on error
+     */
+    int lookup_rows(
+        const std::vector<std::vector<std::vector<int64_t>>>& row_ids,
+        void* output, size_t output_size) const;
+
+    /**
+     * Fast path for contiguous row-id buffers with shape [B, L, H].
+     */
+    int lookup_rows_contiguous(const int64_t* row_ids, int B, int L,
+                               void* output, size_t output_size) const;
+
+    std::vector<int64_t> get_table_vocab_sizes() const;
+    std::vector<std::string> get_store_keys() const;
+    int get_num_heads() const;
+    int get_embedding_dim() const;
+
+    /**
+     * Remove all head tables owned by this EngramStore layer from Mooncake
+     * Store. Missing keys are ignored. Returns the number of removed tables on
+     * success, or a negative error code on failure.
+     */
+    int remove_from_store(bool force = false);
+
+    /**
+     * Populate Store with per-head embedding tensors.
+     * @param embedding_buffers Buffers for each head [N_h, D]
+     * @param buffer_sizes Size in bytes for each buffer
+     * @return 0 on success, negative on error
+     */
+    int populate(const std::vector<void*>& embedding_buffers,
+                 const std::vector<size_t>& buffer_sizes);
+
+   private:
+    int lookup_rows_flat(const int64_t* row_ids, int B, int L, void* output,
+                         size_t output_size) const;
+
+    std::shared_ptr<PyClient> store_;
+    std::vector<int64_t> table_vocab_sizes_;
+    int embedding_dim_;
+    std::vector<std::string> embed_keys_;
+};
+
+}  // namespace engram
+}  // namespace mooncake
diff --git a/mooncake-store/include/engram/engram_store_config.h b/mooncake-store/include/engram/engram_store_config.h
new file mode 100644
index 00000000..09578d4d
--- /dev/null
+++ b/mooncake-store/include/engram/engram_store_config.h
@@ -0,0 +1,22 @@
+#pragma once
+
+#include <cstdint>
+#include <vector>
+
+namespace mooncake {
+namespace engram {
+
+/**
+ * Physical table layout for Mooncake's EngramStore backend.
+ *
+ * Mooncake does not derive table sizes from tokenizer/hash/model config. The
+ * caller must provide the final per-head table sizes and the per-row embedding
+ * width it wants Mooncake to store and query.
+ */
+struct EngramStoreConfig {
+    std::vector<int64_t> table_vocab_sizes = {1024};
+    int embedding_dim = 64;
+};
+
+}  // namespace engram
+}  // namespace mooncake
diff --git a/mooncake-store/include/etcd_helper.h b/mooncake-store/include/etcd_helper.h
index bb0f5625..677bef9a 100644
--- a/mooncake-store/include/etcd_helper.h
+++ b/mooncake-store/include/etcd_helper.h
@@ -26,6 +26,15 @@ class EtcdHelper {
     static ErrorCode ConnectToEtcdStoreClient(
         const std::string& etcd_endpoints);
 
+    /*
+     * @brief Reset the global store etcd client and reconnect to the given
+     *        endpoints. This cancels active store watches/keepalives in the Go
+     *        wrapper and creates a fresh clientv3.Client.
+     * @param etcd_endpoints: The endpoints of the etcd store client.
+     * @return: Error code.
+     */
+    static ErrorCode ResetEtcdStoreClient(const std::string& etcd_endpoints);
+
     /*
      * @brief Get the value of a key from the etcd.
      * @param key: The key to get the value of.
diff --git a/mooncake-store/include/file_storage.h b/mooncake-store/include/file_storage.h
index eef6741d..6aa2ac3d 100644
--- a/mooncake-store/include/file_storage.h
+++ b/mooncake-store/include/file_storage.h
@@ -51,6 +51,7 @@ class FileStorage {
 
    private:
     friend class FileStorageTest;
+    friend class FileStoragePromotionTest;
     struct AllocatedBatch {
         uint64_t batch_id;
         std::vector<BufferHandle> handles;
@@ -74,7 +75,7 @@ class FileStorage {
      * @return tl::expected<void, ErrorCode> indicating operation status.
      */
     tl::expected<void, ErrorCode> OffloadObjects(
-        const std::unordered_map<std::string, int64_t>& offloading_objects);
+        const std::vector<OffloadTaskItem>& offloading_objects);
 
     /**
      * @brief Performs a heartbeat operation for the FileStorage component.
@@ -82,17 +83,31 @@ class FileStorage {
      * client.
      * 2. Receives feedback on which objects should be offloaded.
      * 3. Triggers asynchronous offloading of pending objects.
+     * 4. Pulls and processes any pending L2->L1 promotion tasks queued by the
+     *    master (mirror of step 1+2 in the reverse direction).
      * @return tl::expected<void, ErrorCode> indicating operation status.
      */
     tl::expected<void, ErrorCode> Heartbeat();
 
+    /**
+     * @brief Drives the L2->L1 promotion pipeline for one heartbeat tick.
+     * Pulls promotion work from the master, stages a MEMORY replica for each
+     * key, copies the bytes from local SSD into that replica, and notifies the
+     * master on success. A failure on any single key is logged and skipped;
+     * the master-side reaper decrements the source replica's refcnt and
+     * erases the task entry on TTL expiry, and any orphaned PROCESSING
+     * MEMORY replica is reaped via the standard discarded-replicas path.
+     * @return tl::expected<void, ErrorCode> indicating operation status.
+     */
+    tl::expected<void, ErrorCode> ProcessPromotionTasks();
+
     tl::expected<bool, ErrorCode> IsEnableOffloading();
 
     tl::expected<void, ErrorCode> BatchLoad(
         std::unordered_map<std::string, Slice>& batch_object);
 
     tl::expected<void, ErrorCode> BatchQuerySegmentSlices(
-        const std::vector<std::string>& keys,
+        const std::vector<std::string>& keys, const std::string& tenant_id,
         std::unordered_map<std::string, std::vector<Slice>>& batched_slices);
 
     tl::expected<void, ErrorCode> RegisterLocalMemory();
@@ -103,6 +118,13 @@ class FileStorage {
 
     void ClientBufferGCThreadFunc();
 
+    /**
+     * @brief Re-registers all offloaded objects with the master.
+     * Called after master restart recovery to sync SSD object metadata.
+     * This is the same logic as the ScanMeta step in Init().
+     */
+    tl::expected<void, ErrorCode> ReRegisterOffloadedObjects();
+
     std::shared_ptr<Client> client_;
     SsdMetric* ssd_metric_{nullptr};
     std::string local_rpc_addr_;
@@ -121,6 +143,8 @@ class FileStorage {
     std::thread heartbeat_thread_;
     std::atomic<bool> client_buffer_gc_running_;
     std::thread client_buffer_gc_thread_;
+    std::future<void> rescan_future_;
+    std::atomic<bool> metadata_resync_pending_{false};
 };
 
-}  // namespace mooncake
\ No newline at end of file
+}  // namespace mooncake
diff --git a/mooncake-store/include/gpu_staging_utils.h b/mooncake-store/include/gpu_staging_utils.h
index d99bd477..07982aa0 100644
--- a/mooncake-store/include/gpu_staging_utils.h
+++ b/mooncake-store/include/gpu_staging_utils.h
@@ -15,7 +15,8 @@ namespace gpu_staging {
 // Detect whether ptr resides in accelerator device memory.
 // If so, writes the device ID to *out_device_id for subsequent SetDevice.
 inline bool IsDevicePointer(const void* ptr, int* out_device_id) {
-#if defined(USE_CUDA) || defined(USE_MUSA) || defined(USE_MACA)
+#if defined(USE_CUDA) || defined(USE_MUSA) || defined(USE_MACA) || \
+    defined(USE_HYGON) || defined(USE_COREX)
     cudaPointerAttributes attr{};
     if (cudaPointerGetAttributes(&attr, ptr) == cudaSuccess &&
         attr.type == cudaMemoryTypeDevice) {
@@ -45,7 +46,8 @@ inline bool IsDevicePointer(const void* ptr, int* out_device_id) {
 
 // Copy device memory to host. Caller must have called SetDevice first.
 inline bool CopyDeviceToHost(void* dst, const void* src, size_t size) {
-#if defined(USE_CUDA) || defined(USE_MUSA) || defined(USE_MACA)
+#if defined(USE_CUDA) || defined(USE_MUSA) || defined(USE_MACA) || \
+    defined(USE_HYGON) || defined(USE_COREX)
     return cudaMemcpy(dst, src, size, cudaMemcpyDeviceToHost) == cudaSuccess;
 #elif defined(USE_HIP)
     return hipMemcpy(dst, src, size, hipMemcpyDeviceToHost) == hipSuccess;
@@ -60,10 +62,43 @@ inline bool CopyDeviceToHost(void* dst, const void* src, size_t size) {
 #endif
 }
 
+// Auto-direction copy: runtime determines the transfer direction from pointer
+// attributes (cudaMemcpyDefault). Works for H2H, H2D, D2H, and D2D.
+// Caller must have called SetDevice first when device memory is involved.
+inline bool CopyAuto(void* dst, const void* src, size_t size) {
+#if defined(USE_CUDA) || defined(USE_MUSA) || defined(USE_MACA) || \
+    defined(USE_HYGON) || defined(USE_COREX)
+    return cudaMemcpy(dst, src, size, cudaMemcpyDefault) == cudaSuccess;
+#elif defined(USE_HIP)
+    return hipMemcpy(dst, src, size, hipMemcpyDefault) == hipSuccess;
+#elif defined(USE_ASCEND) || defined(USE_ASCEND_DIRECT) || defined(USE_UBSHMEM)
+    aclrtPtrAttributes src_attr{}, dst_attr{};
+    bool src_dev = aclrtPointerGetAttributes(const_cast<void*>(src),
+                                             &src_attr) == ACL_SUCCESS &&
+                   src_attr.location.type == ACL_MEM_LOCATION_TYPE_DEVICE;
+    bool dst_dev = aclrtPointerGetAttributes(dst, &dst_attr) == ACL_SUCCESS &&
+                   dst_attr.location.type == ACL_MEM_LOCATION_TYPE_DEVICE;
+    aclrtMemcpyKind kind = ACL_MEMCPY_HOST_TO_HOST;
+    if (src_dev && dst_dev)
+        kind = ACL_MEMCPY_DEVICE_TO_DEVICE;
+    else if (src_dev)
+        kind = ACL_MEMCPY_DEVICE_TO_HOST;
+    else if (dst_dev)
+        kind = ACL_MEMCPY_HOST_TO_DEVICE;
+    return aclrtMemcpy(dst, size, src, size, kind) == ACL_SUCCESS;
+#else
+    (void)dst;
+    (void)src;
+    (void)size;
+    return false;
+#endif
+}
+
 // Bind the calling thread to the given device context.
 inline void SetDevice(int device_id) {
     if (device_id < 0) return;
-#if defined(USE_CUDA) || defined(USE_MUSA) || defined(USE_MACA)
+#if defined(USE_CUDA) || defined(USE_MUSA) || defined(USE_MACA) || \
+    defined(USE_HYGON) || defined(USE_COREX)
     cudaSetDevice(device_id);
 #elif defined(USE_HIP)
     hipSetDevice(device_id);
@@ -72,5 +107,60 @@ inline void SetDevice(int device_id) {
 #endif
 }
 
+// Copy host memory to device. Caller must have called SetDevice first.
+inline bool CopyHostToDevice(void* dst, const void* src, size_t size) {
+#if defined(USE_CUDA) || defined(USE_MUSA) || defined(USE_MACA) || \
+    defined(USE_HYGON) || defined(USE_COREX)
+    return cudaMemcpy(dst, src, size, cudaMemcpyHostToDevice) == cudaSuccess;
+#elif defined(USE_HIP)
+    return hipMemcpy(dst, src, size, hipMemcpyHostToDevice) == hipSuccess;
+#elif defined(USE_ASCEND) || defined(USE_ASCEND_DIRECT) || defined(USE_UBSHMEM)
+    return aclrtMemcpy(dst, size, src, size, ACL_MEMCPY_HOST_TO_DEVICE) ==
+           ACL_SUCCESS;
+#else
+    (void)dst;
+    (void)src;
+    (void)size;
+    return false;
+#endif
+}
+
+// Detect whether ptr resides in host (CPU) memory.
+// Used together with IsDevicePointer for safe pointer-type dispatching:
+//   if IsDevicePointer  -> CopyHostToDevice / CopyDeviceToHost
+//   else if IsHostPointer -> memcpy
+//   else                  -> reject (unknown type, e.g. non-standard allocator)
+//
+// Pageable host memory (not tracked by CUDA runtime) is treated as host.
+inline bool IsHostPointer(const void* ptr) {
+#if defined(USE_CUDA) || defined(USE_MUSA) || defined(USE_MACA) || \
+    defined(USE_HYGON) || defined(USE_COREX)
+    cudaPointerAttributes attr{};
+    if (cudaPointerGetAttributes(&attr, ptr) != cudaSuccess) {
+        // Query failed: pageable host memory not tracked by the runtime.
+        cudaGetLastError();  // clear sticky error
+        return true;
+    }
+    return attr.type != cudaMemoryTypeDevice;
+#elif defined(USE_HIP)
+    hipPointerAttribute_t attr{};
+    if (hipPointerGetAttributes(&attr, ptr) != hipSuccess) {
+        hipGetLastError();  // clear sticky error
+        return true;
+    }
+    return attr.type != hipMemoryTypeDevice;
+#elif defined(USE_ASCEND) || defined(USE_ASCEND_DIRECT) || defined(USE_UBSHMEM)
+    aclrtPtrAttributes attr{};
+    if (aclrtPointerGetAttributes(const_cast<void*>(ptr), &attr) !=
+        ACL_SUCCESS) {
+        // Query failed: likely pageable host memory not tracked by the runtime.
+        return true;
+    }
+    return attr.location.type != ACL_MEM_LOCATION_TYPE_DEVICE;
+#else
+    (void)ptr;
+    return true;  // CPU-only build: all pointers are host
+#endif
+}
 }  // namespace gpu_staging
 }  // namespace mooncake
diff --git a/mooncake-store/include/ha/ha_types.h b/mooncake-store/include/ha/ha_types.h
index e0cedf2e..82a9181c 100644
--- a/mooncake-store/include/ha/ha_types.h
+++ b/mooncake-store/include/ha/ha_types.h
@@ -52,6 +52,33 @@ inline std::optional<HABackendType> ParseHABackendType(std::string_view type) {
     return std::nullopt;
 }
 
+inline ErrorCode ValidateHABackendAvailability(HABackendType type) {
+    switch (type) {
+        case HABackendType::UNKNOWN:
+            return ErrorCode::INVALID_PARAMS;
+        case HABackendType::ETCD:
+#ifdef STORE_USE_ETCD
+            return ErrorCode::OK;
+#else
+            return ErrorCode::UNAVAILABLE_IN_CURRENT_MODE;
+#endif
+        case HABackendType::REDIS:
+#ifdef STORE_USE_REDIS
+            return ErrorCode::OK;
+#else
+            return ErrorCode::UNAVAILABLE_IN_CURRENT_MODE;
+#endif
+        case HABackendType::K8S:
+#ifdef STORE_USE_K8S_LEASE
+            return ErrorCode::OK;
+#else
+            return ErrorCode::UNAVAILABLE_IN_CURRENT_MODE;
+#endif
+    }
+
+    return ErrorCode::INVALID_PARAMS;
+}
+
 struct HABackendSpec {
     HABackendType type = HABackendType::UNKNOWN;
     std::string connstring;
diff --git a/mooncake-store/include/ha/leadership/backends/etcd/etcd_leader_coordinator.h b/mooncake-store/include/ha/leadership/backends/etcd/etcd_leader_coordinator.h
index 432d4da7..30e5c8e0 100644
--- a/mooncake-store/include/ha/leadership/backends/etcd/etcd_leader_coordinator.h
+++ b/mooncake-store/include/ha/leadership/backends/etcd/etcd_leader_coordinator.h
@@ -50,6 +50,7 @@ class EtcdLeaderCoordinator final : public LeaderCoordinator {
     static LeadershipLossReason ClassifyLeadershipLossReason(ErrorCode err);
 
     ErrorCode EnsureConnected();
+    ErrorCode ResetConnection();
     ErrorCode ShutdownKeepAliveThread();
     void ClearLeadershipMonitorStateLocked();
     bool IsSameViewVersion(const std::optional<MasterView>& current_view,
diff --git a/mooncake-store/include/ha/leadership/backends/k8s/k8s_leader_coordinator.h b/mooncake-store/include/ha/leadership/backends/k8s/k8s_leader_coordinator.h
new file mode 100644
index 00000000..9621c26f
--- /dev/null
+++ b/mooncake-store/include/ha/leadership/backends/k8s/k8s_leader_coordinator.h
@@ -0,0 +1,71 @@
+#pragma once
+
+#include <atomic>
+#include <memory>
+#include <mutex>
+#include <thread>
+
+#include "k8s_lease_helper.h"
+#include "ha/leadership/leader_coordinator.h"
+
+namespace mooncake {
+namespace ha {
+namespace backends {
+namespace k8s {
+
+class K8sLeaderCoordinator final : public LeaderCoordinator {
+   public:
+    explicit K8sLeaderCoordinator(const HABackendSpec& spec);
+    ~K8sLeaderCoordinator() override;
+
+    ErrorCode Connect();
+
+    tl::expected<std::optional<MasterView>, ErrorCode> ReadCurrentView()
+        override;
+
+    tl::expected<AcquireLeadershipResult, ErrorCode> TryAcquireLeadership(
+        const std::string& leader_address) override;
+
+    tl::expected<bool, ErrorCode> RenewLeadership(
+        const LeadershipSession& session) override;
+
+    tl::expected<ViewChangeResult, ErrorCode> WaitForViewChange(
+        std::optional<ViewVersionId> known_version,
+        std::chrono::milliseconds timeout) override;
+
+    tl::expected<std::unique_ptr<LeadershipMonitorHandle>, ErrorCode>
+    StartLeadershipMonitor(const LeadershipSession& session,
+                           LeadershipLostCallback on_leadership_lost) override;
+
+    ErrorCode ReleaseLeadership(const LeadershipSession& session) override;
+
+   private:
+    ErrorCode EnsureConnected();
+    ErrorCode ShutdownElection();
+    void ClearLeadershipMonitorStateLocked();
+    bool IsSameViewVersion(const std::optional<MasterView>& current_view,
+                           std::optional<ViewVersionId> known_version) const;
+    static tl::expected<std::pair<std::string, std::string>, ErrorCode>
+    ParseConnstring(const std::string& connstring);
+
+    HABackendSpec spec_;
+    std::string namespace_;
+    std::string lease_name_;
+    bool connected_ = false;
+
+    std::mutex election_mutex_;
+    std::thread election_monitor_thread_;
+    std::string election_identity_;
+    OwnerToken election_owner_token_;
+    bool election_active_ = false;
+    bool election_monitor_stopped_ = false;
+    bool election_shutdown_requested_ = false;
+    LeadershipLostCallback leadership_monitor_callback_;
+    std::shared_ptr<std::atomic<bool>> leadership_monitor_armed_;
+    OwnerToken leadership_monitor_owner_token_;
+};
+
+}  // namespace k8s
+}  // namespace backends
+}  // namespace ha
+}  // namespace mooncake
diff --git a/mooncake-store/include/ha/leadership/leader_label_reconciler.h b/mooncake-store/include/ha/leadership/leader_label_reconciler.h
new file mode 100644
index 00000000..dc0c37cc
--- /dev/null
+++ b/mooncake-store/include/ha/leadership/leader_label_reconciler.h
@@ -0,0 +1,99 @@
+#pragma once
+
+#include <chrono>
+#include <condition_variable>
+#include <functional>
+#include <mutex>
+#include <optional>
+#include <thread>
+#include <utility>
+
+#include <glog/logging.h>
+
+#include "types.h"
+
+namespace mooncake {
+namespace ha {
+
+// Drives a single boolean "leader label" toward a desired state on a background
+// thread. Callers only flip the desired flag (non-blocking); the worker runs
+// the synchronous apply and retries until the actual state converges, so a
+// transient failure cannot leave a stale leader label on a pod that is no
+// longer serving.
+class LeaderLabelReconciler {
+   public:
+    // Applies the label mutation: true sets the leader label, false clears it.
+    // Must return ErrorCode::OK only when the cluster state reflects `desired`.
+    using ApplyFn = std::function<ErrorCode(bool desired)>;
+
+    LeaderLabelReconciler(bool enabled, ApplyFn apply,
+                          std::chrono::milliseconds retry_interval)
+        : enabled_(enabled),
+          apply_(std::move(apply)),
+          retry_interval_(retry_interval) {
+        if (enabled_) {
+            worker_ = std::thread([this] { Run(); });
+        }
+    }
+
+    ~LeaderLabelReconciler() {
+        if (!enabled_) return;
+        {
+            std::lock_guard<std::mutex> lock(mutex_);
+            stop_ = true;
+        }
+        cv_.notify_all();
+        if (worker_.joinable()) worker_.join();
+    }
+
+    LeaderLabelReconciler(const LeaderLabelReconciler&) = delete;
+    LeaderLabelReconciler& operator=(const LeaderLabelReconciler&) = delete;
+
+    void SetLeader(bool desired) {
+        if (!enabled_) return;
+        {
+            std::lock_guard<std::mutex> lock(mutex_);
+            desired_leader_ = desired;
+        }
+        cv_.notify_all();
+    }
+
+   private:
+    void Run() {
+        std::optional<bool> applied;
+        std::unique_lock<std::mutex> lock(mutex_);
+        while (true) {
+            cv_.wait(lock, [&] { return stop_ || applied != desired_leader_; });
+            if (stop_) return;
+            bool target = desired_leader_;
+            lock.unlock();
+
+            ErrorCode err = apply_(target);
+
+            lock.lock();
+            if (err == ErrorCode::OK) {
+                applied = target;
+                continue;
+            }
+            applied.reset();
+            LOG(WARNING) << "Failed to " << (target ? "set" : "clear")
+                         << " leader label: " << toString(err);
+            cv_.wait_for(lock, retry_interval_,
+                         [&] { return stop_ || desired_leader_ != target; });
+            if (stop_) return;
+        }
+    }
+
+    const bool enabled_;
+    const ApplyFn apply_;
+    const std::chrono::milliseconds retry_interval_;
+
+    std::mutex mutex_;
+    std::condition_variable cv_;
+    bool desired_leader_ = false;
+    bool stop_ = false;
+    std::thread worker_;
+};
+
+}  // namespace ha
+}  // namespace mooncake
diff --git a/mooncake-store/include/ha/oplog/oplog_store_factory.h b/mooncake-store/include/ha/oplog/oplog_store_factory.h
index 2075e408..7b358ca5 100644
--- a/mooncake-store/include/ha/oplog/oplog_store_factory.h
+++ b/mooncake-store/include/ha/oplog/oplog_store_factory.h
@@ -2,6 +2,7 @@
 #pragma once
 
 #include <algorithm>
+#include <cctype>
 #include <memory>
 #include <string>
 
@@ -31,7 +32,8 @@ static constexpr OpLogStoreType kDefaultOpLogStoreType =
 // Returns kDefaultOpLogStoreType for unrecognized strings.
 inline OpLogStoreType ParseOpLogStoreType(const std::string& type_str) {
     std::string lower = type_str;
-    std::transform(lower.begin(), lower.end(), lower.begin(), ::tolower);
+    std::transform(lower.begin(), lower.end(), lower.begin(),
+                   [](unsigned char c) -> char { return std::tolower(c); });
     if (lower == "localfs" || lower == "local_fs") {
         return OpLogStoreType::LOCAL_FS;
     }
diff --git a/mooncake-store/include/ha/snapshot/catalog/backends/embedded/embedded_snapshot_catalog_store.h b/mooncake-store/include/ha/snapshot/catalog/backends/embedded/embedded_snapshot_catalog_store.h
index a73b79ff..33a8926b 100644
--- a/mooncake-store/include/ha/snapshot/catalog/backends/embedded/embedded_snapshot_catalog_store.h
+++ b/mooncake-store/include/ha/snapshot/catalog/backends/embedded/embedded_snapshot_catalog_store.h
@@ -10,7 +10,8 @@ namespace embedded {
 
 class EmbeddedSnapshotCatalogStore : public SnapshotCatalogStore {
    public:
-    explicit EmbeddedSnapshotCatalogStore(SnapshotObjectStore* object_store);
+    EmbeddedSnapshotCatalogStore(SnapshotObjectStore* object_store,
+                                 const std::string& cluster_id);
 
     ErrorCode Publish(const SnapshotDescriptor& snapshot) override;
 
@@ -22,8 +23,13 @@ class EmbeddedSnapshotCatalogStore : public SnapshotCatalogStore {
 
     ErrorCode Delete(const SnapshotId& snapshot_id) override;
 
+    const std::string& GetSnapshotRoot() const override {
+        return snapshot_root_;
+    }
+
    private:
     SnapshotObjectStore* object_store_;
+    std::string snapshot_root_;
 };
 
 }  // namespace embedded
diff --git a/mooncake-store/include/ha/snapshot/catalog/backends/redis/redis_snapshot_catalog_store.h b/mooncake-store/include/ha/snapshot/catalog/backends/redis/redis_snapshot_catalog_store.h
index c0855541..8d7223f6 100644
--- a/mooncake-store/include/ha/snapshot/catalog/backends/redis/redis_snapshot_catalog_store.h
+++ b/mooncake-store/include/ha/snapshot/catalog/backends/redis/redis_snapshot_catalog_store.h
@@ -26,6 +26,10 @@ class RedisSnapshotCatalogStore final : public SnapshotCatalogStore {
 
     ErrorCode Delete(const SnapshotId& snapshot_id) override;
 
+    const std::string& GetSnapshotRoot() const override {
+        return snapshot_root_;
+    }
+
    private:
     static ClusterNamespace ResolveClusterNamespace(
         const ClusterNamespace& cluster_namespace);
@@ -36,6 +40,7 @@ class RedisSnapshotCatalogStore final : public SnapshotCatalogStore {
     SnapshotObjectStore* object_store_;
     std::string connstring_;
     ClusterNamespace cluster_namespace_;
+    std::string snapshot_root_;
     std::string latest_key_;
     std::string index_key_;
 };
diff --git a/mooncake-store/include/ha/snapshot/catalog/snapshot_catalog_store.h b/mooncake-store/include/ha/snapshot/catalog/snapshot_catalog_store.h
index fe7baff8..763e242c 100644
--- a/mooncake-store/include/ha/snapshot/catalog/snapshot_catalog_store.h
+++ b/mooncake-store/include/ha/snapshot/catalog/snapshot_catalog_store.h
@@ -18,11 +18,24 @@ namespace ha {
 
 namespace snapshot_catalog_store_detail {
 
-constexpr std::string_view kSnapshotRoot = "mooncake_master_snapshot/";
+constexpr std::string_view kSnapshotRootBase = "mooncake_master_snapshot";
 constexpr std::string_view kSnapshotLatest = "latest.txt";
 constexpr std::string_view kSnapshotManifest = "manifest.txt";
 constexpr std::string_view kSnapshotDescriptor = "descriptor.txt";
 
+/// Build the per-cluster snapshot root path.
+/// When cluster_id is non-empty: "mooncake_master_snapshot/{cluster_id}/"
+/// When cluster_id is empty:     "mooncake_master_snapshot/"
+inline std::string BuildSnapshotRoot(const std::string& cluster_id) {
+    std::string root(kSnapshotRootBase);
+    root += '/';
+    if (!cluster_id.empty()) {
+        root += cluster_id;
+        root += '/';
+    }
+    return root;
+}
+
 inline bool IsAsciiDigit(char ch) {
     return std::isdigit(static_cast<unsigned char>(ch)) != 0;
 }
@@ -59,28 +72,33 @@ inline std::string TrimAsciiWhitespace(std::string value) {
     return value.substr(first, last - first + 1);
 }
 
-inline std::string BuildSnapshotPrefix(const SnapshotId& snapshot_id) {
-    return std::string(kSnapshotRoot) + snapshot_id + "/";
+inline std::string BuildSnapshotPrefix(const std::string& snapshot_root,
+                                       const SnapshotId& snapshot_id) {
+    return snapshot_root + snapshot_id + "/";
 }
 
-inline std::string BuildManifestKey(const SnapshotId& snapshot_id) {
-    return BuildSnapshotPrefix(snapshot_id) + std::string(kSnapshotManifest);
+inline std::string BuildManifestKey(const std::string& snapshot_root,
+                                    const SnapshotId& snapshot_id) {
+    return BuildSnapshotPrefix(snapshot_root, snapshot_id) +
+           std::string(kSnapshotManifest);
 }
 
-inline std::string BuildDescriptorKey(const SnapshotId& snapshot_id) {
-    return BuildSnapshotPrefix(snapshot_id) + std::string(kSnapshotDescriptor);
+inline std::string BuildDescriptorKey(const std::string& snapshot_root,
+                                      const SnapshotId& snapshot_id) {
+    return BuildSnapshotPrefix(snapshot_root, snapshot_id) +
+           std::string(kSnapshotDescriptor);
 }
 
-inline std::string BuildLatestKey() {
-    return std::string(kSnapshotRoot) + std::string(kSnapshotLatest);
+inline std::string BuildLatestKey(const std::string& snapshot_root) {
+    return snapshot_root + std::string(kSnapshotLatest);
 }
 
 inline SnapshotDescriptor MakeSnapshotDescriptor(
-    const SnapshotId& snapshot_id) {
+    const std::string& snapshot_root, const SnapshotId& snapshot_id) {
     SnapshotDescriptor descriptor;
     descriptor.snapshot_id = snapshot_id;
-    descriptor.manifest_key = BuildManifestKey(snapshot_id);
-    descriptor.object_prefix = BuildSnapshotPrefix(snapshot_id);
+    descriptor.manifest_key = BuildManifestKey(snapshot_root, snapshot_id);
+    descriptor.object_prefix = BuildSnapshotPrefix(snapshot_root, snapshot_id);
     return descriptor;
 }
 
@@ -100,7 +118,8 @@ inline std::string SerializeSnapshotDescriptor(
 }
 
 inline tl::expected<SnapshotDescriptor, ErrorCode>
-DeserializeSnapshotDescriptor(const SnapshotId& snapshot_id,
+DeserializeSnapshotDescriptor(const std::string& snapshot_root,
+                              const SnapshotId& snapshot_id,
                               std::string_view payload) {
     const auto first = payload.find('|');
     const auto second = first == std::string_view::npos
@@ -110,7 +129,8 @@ DeserializeSnapshotDescriptor(const SnapshotId& snapshot_id,
         return tl::make_unexpected(ErrorCode::DESERIALIZE_FAIL);
     }
 
-    SnapshotDescriptor descriptor = MakeSnapshotDescriptor(snapshot_id);
+    SnapshotDescriptor descriptor =
+        MakeSnapshotDescriptor(snapshot_root, snapshot_id);
     if (!ParseDecimal(payload.substr(0, first), descriptor.last_included_seq) ||
         !ParseDecimal(payload.substr(first + 1, second - first - 1),
                       descriptor.producer_view_version) ||
@@ -135,6 +155,9 @@ class SnapshotCatalogStore {
         size_t limit) = 0;
 
     virtual ErrorCode Delete(const SnapshotId& snapshot_id) = 0;
+
+    /// Returns the per-cluster snapshot root path used by this catalog store.
+    virtual const std::string& GetSnapshotRoot() const = 0;
 };
 
 }  // namespace ha
diff --git a/mooncake-store/include/ha/snapshot/snapshot_logger.h b/mooncake-store/include/ha/snapshot/snapshot_logger.h
index 00f9d7f1..9d54f8b7 100644
--- a/mooncake-store/include/ha/snapshot/snapshot_logger.h
+++ b/mooncake-store/include/ha/snapshot/snapshot_logger.h
@@ -16,7 +16,7 @@ inline void SnapshotLogWrite(const char* level, const std::string& msg) {
     if (g_snapshot_log_pipe_fd < 0) return;
     std::string line = std::string("[") + level + "] " + msg + "\n";
     // write() is async-signal-safe, won't deadlock after fork
-    ::write(g_snapshot_log_pipe_fd, line.c_str(), line.size());
+    std::ignore = ::write(g_snapshot_log_pipe_fd, line.c_str(), line.size());
 }
 
 // Macros for child process logging - writes to pipe instead of glog
diff --git a/mooncake-store/include/ha_metric_manager.h b/mooncake-store/include/ha_metric_manager.h
index e3f1cda3..d10d040c 100644
--- a/mooncake-store/include/ha_metric_manager.h
+++ b/mooncake-store/include/ha_metric_manager.h
@@ -27,13 +27,6 @@ class HAMetricManager {
     // --- Singleton Access ---
     static HAMetricManager& instance();
 
-    /**
-     * @brief Explicitly initialize the singleton instance.
-     * Use this at startup (e.g. main) to ensure thread-safe initialization
-     * of the underlying metric library components.
-     */
-    static void Init() { instance(); }
-
     HAMetricManager(const HAMetricManager&) = delete;
     HAMetricManager& operator=(const HAMetricManager&) = delete;
     HAMetricManager(HAMetricManager&&) = delete;
diff --git a/mooncake-store/include/http_metadata_server.h b/mooncake-store/include/http_metadata_server.h
index 411070c0..366ce4d7 100644
--- a/mooncake-store/include/http_metadata_server.h
+++ b/mooncake-store/include/http_metadata_server.h
@@ -5,6 +5,7 @@
 #include <string>
 #include <unordered_map>
 #include <mutex>
+#include <vector>
 
 #include <ylt/coro_http/coro_http_server.hpp>
 
@@ -35,6 +36,14 @@ class HttpMetadataServer {
     // Check if the server is running
     bool is_running() const { return running_; }
 
+    // Remove a key from the metadata store (for internal use by MasterService)
+    // Returns true if key was found and removed, false if key did not exist
+    bool removeKey(const std::string& key);
+
+    // Remove multiple keys from the metadata store
+    // Returns the number of keys that were successfully removed
+    size_t removeKeys(const std::vector<std::string>& keys);
+
     // Non-copyable
     HttpMetadataServer(const HttpMetadataServer&) = delete;
     HttpMetadataServer& operator=(const HttpMetadataServer&) = delete;
diff --git a/mooncake-store/include/k8s_lease_helper.h b/mooncake-store/include/k8s_lease_helper.h
new file mode 100644
index 00000000..4312f693
--- /dev/null
+++ b/mooncake-store/include/k8s_lease_helper.h
@@ -0,0 +1,52 @@
+#pragma once
+
+#include <cstdint>
+#include <mutex>
+#include <string>
+
+#include "types.h"
+
+namespace mooncake {
+
+class K8sLeaseHelper {
+   public:
+    static ErrorCode Init();
+
+    static ErrorCode RunElection(const std::string& ns,
+                                 const std::string& lease,
+                                 const std::string& identity, int lease_dur,
+                                 int renew_deadline, int retry_period);
+
+    static ErrorCode WaitElected(const std::string& ns,
+                                 const std::string& lease, int timeout_sec,
+                                 int64_t& lease_transitions);
+
+    static ErrorCode WaitLost(const std::string& ns, const std::string& lease);
+
+    static ErrorCode CancelElection(const std::string& ns,
+                                    const std::string& lease);
+
+    static ErrorCode GetHolder(const std::string& ns, const std::string& lease,
+                               std::string& holder, int64_t& lease_transitions);
+
+    static ErrorCode WatchHolder(
+        const std::string& ns, const std::string& lease, void* callback_context,
+        void (*callback_func)(void*, const char*, size_t, int64_t));
+
+    static ErrorCode CancelWatch(const std::string& ns,
+                                 const std::string& lease);
+
+    static ErrorCode SetPodLabel(const std::string& ns, const std::string& pod,
+                                 const std::string& key,
+                                 const std::string& value);
+
+    static ErrorCode ClearPodLabel(const std::string& ns,
+                                   const std::string& pod,
+                                   const std::string& key);
+
+   private:
+    static std::mutex init_mutex_;
+    static bool initialized_;
+};
+
+}  // namespace mooncake
diff --git a/mooncake-store/include/local_hot_cache.h b/mooncake-store/include/local_hot_cache.h
index 42f0c7b3..3fa827f3 100644
--- a/mooncake-store/include/local_hot_cache.h
+++ b/mooncake-store/include/local_hot_cache.h
@@ -19,6 +19,16 @@
 
 namespace mooncake {
 
+/**
+ * @brief Token captured at async hot cache fill submission time.
+ * Invalidated when RemoveHotKey, BumpKeyGeneration, or Clear bumps
+ * generation/epoch.
+ */
+struct HotCachePutToken {
+    uint64_t cache_epoch = 0;
+    uint64_t key_generation = 0;
+};
+
 /**
  * @brief Memory block metadata for hot cache.
  */
@@ -61,6 +71,15 @@ class LocalHotCache {
      */
     bool PutHotKey(HotMemBlock* block);
 
+    /**
+     * @brief Insert a populated block only if its async fill token is still
+     * valid, checking the token and publishing atomically under one lock.
+     * If the token is stale (the key was removed/overwritten since the fill
+     * started), the block is returned to the pool instead of being published.
+     * @return true if the block was published, false if cancelled or on error.
+     */
+    bool PutHotKey(HotMemBlock* block, const HotCachePutToken& token);
+
     /**
      * @brief Check if the key exists in cache.
      * @param key Cache key: {request key}
@@ -91,6 +110,65 @@ class LocalHotCache {
      */
     bool TouchHotKey(const std::string& key);
 
+    /**
+     * @brief Remove a key from the hot cache immediately.
+     * Bumps the key generation so in-flight async fills are invalidated.
+     * @param key The key to remove from cache.
+     * @return true if a published cache entry was removed, false otherwise.
+     */
+    bool RemoveHotKey(const std::string& key);
+
+    /**
+     * @brief Remove a batch of keys from the hot cache under one lock.
+     * Bumps each key generation so in-flight async fills are invalidated.
+     * @return number of published cache entries removed.
+     */
+    size_t RemoveHotKeys(const std::vector<std::string>& keys);
+
+    /**
+     * @brief Remove cached keys matching a regex from the hot cache.
+     * Bumps key generation for matching published entries.
+     * @return number of published cache entries removed.
+     */
+    size_t RemoveHotKeysByRegex(const std::string& regex_pattern);
+
+    /**
+     * @brief Remove every published hot cache entry and invalidate async fills.
+     * @return number of published cache entries removed.
+     */
+    size_t RemoveAllHotKeys();
+
+    /**
+     * @brief Invalidate in-flight async fills for a key without evicting it.
+     */
+    void BumpKeyGeneration(const std::string& key);
+
+    /**
+     * @brief Invalidate in-flight async fills for multiple keys.
+     */
+    void BumpKeyGenerations(const std::vector<std::string>& keys);
+
+    /**
+     * @brief Invalidate all in-flight async fills without evicting entries.
+     */
+    void BumpCacheEpoch();
+
+    /**
+     * @brief Clear all hot cache entries and invalidate in-flight async fills.
+     */
+    void Clear();
+
+    /**
+     * @brief Capture the current put token for async hot cache fill validation.
+     */
+    HotCachePutToken AcquirePutToken(const std::string& key);
+
+    /**
+     * @brief Check whether an async put token is still valid.
+     */
+    bool IsPutTokenValid(const std::string& key,
+                         const HotCachePutToken& token) const;
+
     /**
      * @brief Get a free block for writing.
      * Detaches a block from the LRU tail (evicting if necessary) and returns
@@ -112,6 +190,16 @@ class LocalHotCache {
      */
     size_t GetBlockSize() const { return block_size_; }
 
+    /**
+     * @brief Get the base address of the bulk allocation.
+     */
+    void* GetBaseAddress() const { return bulk_memory_standard_; }
+
+    /**
+     * @brief Get the total size of the bulk allocation in bytes.
+     */
+    size_t GetTotalSize() const { return bulk_memory_size_; }
+
     /**
      * @brief Get the shm segment backing this cache.
      * Only valid when constructed with use_shm=true.
@@ -136,6 +224,11 @@ class LocalHotCache {
    private:
     // Drain deferred LRU touches: splice accessed blocks to front
     void drainDeferredTouches();
+    bool putHotKeyLocked(HotMemBlock* block);
+    bool removeHotKeyLocked(const std::string& key);
+    bool hasActiveBlockForKeyLocked(const std::string& key) const;
+    bool isPutTokenValidLocked(const std::string& key,
+                               const HotCachePutToken& token) const;
 
     size_t block_size_;  // Actual block size used by this cache
 
@@ -155,6 +248,9 @@ class LocalHotCache {
     // key -> iterator of lru_queue_
     std::unordered_map<std::string, std::list<HotMemBlock*>::iterator>
         key_to_lru_it_ GUARDED_BY(lru_mutex_);
+    std::unordered_map<std::string, uint64_t> key_generation_
+        GUARDED_BY(lru_mutex_);
+    std::atomic<uint64_t> cache_epoch_{0};
 };
 
 /**
@@ -165,13 +261,19 @@ struct HotCachePutTask {
     HotMemBlock* block;  // Pointer to the allocated block
     size_t size;
     std::shared_ptr<LocalHotCache> hot_cache;
+    HotCachePutToken token;
 
     // Default constructor for empty task
     HotCachePutTask() : block(nullptr), size(0), hot_cache(nullptr) {}
 
     HotCachePutTask(const std::string& k, const Slice& slice, HotMemBlock* blk,
-                    std::shared_ptr<LocalHotCache> cache)
-        : key(k), block(blk), size(slice.size), hot_cache(std::move(cache)) {
+                    std::shared_ptr<LocalHotCache> cache,
+                    HotCachePutToken put_token)
+        : key(k),
+          block(blk),
+          size(slice.size),
+          hot_cache(std::move(cache)),
+          token(put_token) {
         // No data copy here; memcpy is done by SubmitPutTask into block->addr.
     }
 };
diff --git a/mooncake-store/include/master_admin_service.h b/mooncake-store/include/master_admin_service.h
new file mode 100644
index 00000000..1898cb43
--- /dev/null
+++ b/mooncake-store/include/master_admin_service.h
@@ -0,0 +1,119 @@
+#pragma once
+
+#include <csignal>
+
+#include <atomic>
+#include <cstdint>
+#include <memory>
+#include <mutex>
+#include <optional>
+#include <semaphore>
+#include <string>
+#include <thread>
+
+#include <ylt/coro_http/coro_http_server.hpp>
+
+#include "ha/ha_types.h"
+
+namespace mooncake {
+
+extern const uint64_t kMetricReportIntervalSeconds;
+
+class WrappedMasterService;
+
+class MasterAdminServer {
+   public:
+    MasterAdminServer(uint16_t http_port, bool enable_metric_reporting);
+
+    ~MasterAdminServer();
+
+    bool Start();
+
+    void Stop();
+
+    void SetRuntimeState(ha::MasterRuntimeState state);
+
+    void SetObservedLeader(const std::optional<ha::MasterView>& leader_view);
+
+    void SetServiceDelegate(std::shared_ptr<WrappedMasterService> service);
+
+    void SetServiceAvailable(bool available);
+
+   private:
+    struct RuntimeSnapshot {
+        ha::MasterRuntimeState state = ha::MasterRuntimeState::kStarting;
+        std::optional<ha::MasterView> leader_view;
+        std::shared_ptr<WrappedMasterService> service;
+        bool service_available = false;
+    };
+
+    RuntimeSnapshot SnapshotState() const;
+
+    std::string BuildMetricsText() const;
+
+    std::string BuildTenantQuotaMetricsText() const;
+
+    std::string BuildMetricsSummaryText() const;
+
+    std::shared_ptr<WrappedMasterService> GetActiveService() const;
+
+    template <typename Handler>
+    void WithActiveService(coro_http::coro_http_response& resp,
+                           Handler&& handler) const;
+
+    void HandleMetrics(coro_http::coro_http_request& req,
+                       coro_http::coro_http_response& resp);
+    void HandleMetricsSummary(coro_http::coro_http_request& req,
+                              coro_http::coro_http_response& resp);
+    void HandleHealth(coro_http::coro_http_request& req,
+                      coro_http::coro_http_response& resp);
+    void HandleRole(coro_http::coro_http_request& req,
+                    coro_http::coro_http_response& resp);
+    void HandleHaStatus(coro_http::coro_http_request& req,
+                        coro_http::coro_http_response& resp);
+    void HandleLeader(coro_http::coro_http_request& req,
+                      coro_http::coro_http_response& resp);
+    void HandleQueryKey(coro_http::coro_http_request& req,
+                        coro_http::coro_http_response& resp);
+    void HandleGetAllKeys(coro_http::coro_http_request& req,
+                          coro_http::coro_http_response& resp);
+    void HandleGetAllSegments(coro_http::coro_http_request& req,
+                              coro_http::coro_http_response& resp);
+    void HandleGetSegmentsDetail(coro_http::coro_http_request& req,
+                                 coro_http::coro_http_response& resp);
+    void HandleQuerySegment(coro_http::coro_http_request& req,
+                            coro_http::coro_http_response& resp);
+    void HandleCreateDrainJob(coro_http::coro_http_request& req,
+                              coro_http::coro_http_response& resp);
+    void HandleQueryDrainJob(coro_http::coro_http_request& req,
+                             coro_http::coro_http_response& resp);
+    void HandleCancelDrainJob(coro_http::coro_http_request& req,
+                              coro_http::coro_http_response& resp);
+    void HandleSegmentStatus(coro_http::coro_http_request& req,
+                             coro_http::coro_http_response& resp);
+    void HandleBatchQueryKeys(coro_http::coro_http_request& req,
+                              coro_http::coro_http_response& resp);
+    void HandleGetTenantQuotas(coro_http::coro_http_request& req,
+                               coro_http::coro_http_response& resp);
+    void HandleUpsertTenantQuota(coro_http::coro_http_request& req,
+                                 coro_http::coro_http_response& resp);
+    void HandleDeleteTenantQuota(coro_http::coro_http_request& req,
+                                 coro_http::coro_http_response& resp);
+
+    void RegisterHandler();
+
+    uint16_t http_port_;
+    bool enable_metric_reporting_ = false;
+    coro_http::coro_http_server http_server_;
+    std::thread metric_report_thread_;
+    std::atomic<bool> metric_report_running_{false};
+    std::binary_semaphore metric_report_stop_sem_{0};
+    std::atomic<bool> started_{false};
+    mutable std::mutex state_mutex_;
+    ha::MasterRuntimeState state_{ha::MasterRuntimeState::kStarting};
+    std::optional<ha::MasterView> leader_view_;
+    std::shared_ptr<WrappedMasterService> service_;
+    bool service_available_ = false;
+};
+
+}  // namespace mooncake
diff --git a/mooncake-store/include/master_client.h b/mooncake-store/include/master_client.h
index 40ebb55e..7b541230 100644
--- a/mooncake-store/include/master_client.h
+++ b/mooncake-store/include/master_client.h
@@ -14,6 +14,7 @@
 
 #include "client_metric.h"
 #include "replica.h"
+#include "segment.h"
 #include "types.h"
 #include "rpc_types.h"
 #include "master_metric_manager.h"
@@ -51,8 +52,11 @@ inline void MaybeEnableRdmaSocketConfig(SocketConfigVariant& socket_config) {
  */
 class MasterClient {
    public:
-    MasterClient(const UUID& client_id, MasterClientMetric* metrics = nullptr)
-        : client_id_(client_id), metrics_(metrics) {
+    MasterClient(const UUID& client_id, MasterClientMetric* metrics = nullptr,
+                 std::string tenant_id = "default")
+        : client_id_(client_id),
+          tenant_id_(NormalizeTenantId(std::move(tenant_id))),
+          metrics_(metrics) {
         coro_io::client_pool<coro_rpc::coro_rpc_client>::pool_config
             pool_conf{};
 
@@ -66,12 +70,29 @@ class MasterClient {
             detail::MaybeEnableRdmaSocketConfig(
                 pool_conf.client_config.socket_config);
         }
+
+        // Per-request timeout for all client->master RPCs. coro_rpc's
+        // send_request falls back to this config value when no per-call
+        // timeout is given, so setting it here covers every RPC method
+        // uniformly. Default stays at coro_rpc's built-in 30s if unset;
+        // a negative value disables the timeout (no timer is armed).
+        if (const char* timeout_ms = std::getenv("MC_RPC_TIMEOUT_MS")) {
+            pool_conf.client_config.request_timeout_duration =
+                std::chrono::milliseconds(std::atoll(timeout_ms));
+        }
+        // Optional override for the TCP/RDMA connect timeout (default 30s).
+        if (const char* connect_ms = std::getenv("MC_RPC_CONNECT_TIMEOUT_MS")) {
+            pool_conf.client_config.connect_timeout_duration =
+                std::chrono::milliseconds(std::atoll(connect_ms));
+        }
         client_pools_ =
             std::make_shared<coro_io::client_pools<coro_rpc::coro_rpc_client>>(
                 pool_conf);
     }
     ~MasterClient();
 
+    const std::string& tenant_id() const { return tenant_id_; }
+
     MasterClient(const MasterClient&) = delete;
     MasterClient& operator=(const MasterClient&) = delete;
 
@@ -100,9 +121,11 @@ class MasterClient {
         const std::vector<std::string>& object_keys);
 
     /**
-     * @brief Calculate cache hit rate metrics
+     * @brief Calculate Store-observed cache reuse metrics
      * @param object_keys None
-     * @return Map containing metrics
+     * @return Map containing metrics. Legacy hit-rate keys describe cumulative
+     * Store-side hits normalized by current cached object counts, not
+     * end-to-end request/token hit ratios.
      */
     [[nodiscard]] tl::expected<MasterMetricManager::CacheHitStatDict, ErrorCode>
     CalcCacheStats();
@@ -140,6 +163,8 @@ class MasterClient {
      */
     [[nodiscard]] tl::expected<GetReplicaListResponse, ErrorCode>
     GetReplicaList(const std::string& object_key);
+    [[nodiscard]] tl::expected<GetReplicaListResponse, ErrorCode>
+    GetReplicaList(const std::string& object_key, const std::string& tenant_id);
 
     /**
      * @brief Retrieves replica lists for object keys that match a regex
@@ -161,6 +186,9 @@ class MasterClient {
      */
     [[nodiscard]] std::vector<tl::expected<GetReplicaListResponse, ErrorCode>>
     BatchGetReplicaList(const std::vector<std::string>& object_keys);
+    [[nodiscard]] std::vector<tl::expected<GetReplicaListResponse, ErrorCode>>
+    BatchGetReplicaList(const std::vector<std::string>& object_keys,
+                        const std::string& tenant_id);
 
     /**
      * @brief Starts a put operation
@@ -204,7 +232,8 @@ class MasterClient {
      * @return ErrorCode indicating success/failure
      */
     [[nodiscard]] std::vector<tl::expected<void, ErrorCode>> BatchPutEnd(
-        const std::vector<std::string>& keys);
+        const std::vector<std::string>& keys,
+        ReplicaType replica_type = ReplicaType::ALL);
 
     /**
      * @brief Revokes a put operation
@@ -221,7 +250,8 @@ class MasterClient {
      * @return ErrorCode indicating success/failure
      */
     [[nodiscard]] std::vector<tl::expected<void, ErrorCode>> BatchPutRevoke(
-        const std::vector<std::string>& keys);
+        const std::vector<std::string>& keys,
+        ReplicaType replica_type = ReplicaType::ALL);
 
     /**
      * @brief Starts an upsert operation (insert or update)
@@ -296,6 +326,17 @@ class MasterClient {
     [[nodiscard]] tl::expected<void, ErrorCode> MountSegment(
         const Segment& segment);
 
+    [[nodiscard]] tl::expected<void, ErrorCode> MountSSDSegment(
+        const Segment& segment);
+
+    /**
+     * @brief Registers a NoF ssd segment to master for allocation
+     * @param segment Segment to register
+     * @return tl::expected<void, ErrorCode> indicating success/failure
+     */
+    [[nodiscard]] tl::expected<void, ErrorCode> MountNoFSegment(
+        const NoFSegment& segment);
+
     /**
      * @brief Re-mount segments, invoked when the client is the first time to
      * connect to the master or the client Ping TTL is expired and need
@@ -307,6 +348,17 @@ class MasterClient {
     [[nodiscard]] tl::expected<void, ErrorCode> ReMountSegment(
         const std::vector<Segment>& segments);
 
+    /**
+     * @brief Re-mount NoF ssd segments, invoked when the client is the first
+     * time to connect to the master or the client Ping TTL is expired and need
+     * to remount. This function is idempotent. Client should retry if the
+     * return code is not ErrorCode::OK.
+     * @param segments Segments to remount
+     * @return tl::expected<void, ErrorCode> indicating success/failure
+     */
+    [[nodiscard]] tl::expected<void, ErrorCode> ReMountNoFSegment(
+        const std::vector<NoFSegment>& segments);
+
     /**
      * @brief Unregisters a memory segment from master
      * @param segment_id ID of the segment to unmount
@@ -315,6 +367,34 @@ class MasterClient {
     [[nodiscard]] tl::expected<void, ErrorCode> UnmountSegment(
         const UUID& segment_id);
 
+    [[nodiscard]] tl::expected<void, ErrorCode> GracefulUnmountSegment(
+        const UUID& segment_id, uint64_t grace_period_ms);
+
+    /**
+     * @brief Unregisters a NoF ssd segment from master
+     * @param segment_id ID of the segment to unmount
+     * @return tl::expected<void, ErrorCode> indicating success/failure
+     */
+    [[nodiscard]] tl::expected<void, ErrorCode> UnmountNoFSegment(
+        const UUID& segment_id);
+
+    /**
+     * @brief Gets all mounted NoF ssd segments from master
+     * @return tl::expected<std::vector<MountedNoFSegmentSnapshot>, ErrorCode>
+     * containing all mounted segments
+     */
+    [[nodiscard]] tl::expected<std::vector<NoFSegment>, ErrorCode>
+    GetAllNoFSegments();
+
+    /**
+     * @brief Gets all mounted NoF segments that match a segment name together
+     * with their owner client ids.
+     * @param segment_name Mounted NoF segment name
+     * @return Matching segment owner info list
+     */
+    [[nodiscard]] tl::expected<std::vector<NoFSegmentOwnerInfo>, ErrorCode>
+    GetNoFSegmentsByName(const std::string& segment_name);
+
     /**
      * @brief Gets the cluster ID for the current client to use as subdirectory
      * name
@@ -322,6 +402,9 @@ class MasterClient {
      */
     [[nodiscard]] tl::expected<std::string, ErrorCode> GetFsdir();
 
+    [[nodiscard]] tl::expected<SegmentStatus, ErrorCode> QuerySegmentStatusById(
+        const UUID& segment_id);
+
     [[nodiscard]] tl::expected<GetStorageConfigResponse, ErrorCode>
     GetStorageConfig();
 
@@ -345,10 +428,12 @@ class MasterClient {
      * @param enable_offloading Indicates whether persistence is enabled for
      * this segment.
      */
-    [[nodiscard]] tl::expected<std::unordered_map<std::string, int64_t>,
-                               ErrorCode>
+    [[nodiscard]] tl::expected<std::vector<OffloadTaskItem>, ErrorCode>
     OffloadObjectHeartbeat(const UUID& client_id, bool enable_offloading);
 
+    [[nodiscard]] tl::expected<void, ErrorCode> ReportSsdCapacity(
+        const UUID& client_id, int64_t ssd_total_capacity_bytes);
+
     /**
      * @brief Adds multiple new objects to a specified client in batch.
      * @param keys         A list of object keys (names) that were successfully
@@ -359,6 +444,53 @@ class MasterClient {
     [[nodiscard]] tl::expected<void, ErrorCode> NotifyOffloadSuccess(
         const UUID& client_id, const std::vector<std::string>& keys,
         const std::vector<StorageObjectMetadata>& metadatas);
+    [[nodiscard]] tl::expected<void, ErrorCode> NotifyOffloadSuccess(
+        const UUID& client_id, const std::vector<OffloadTaskItem>& tasks,
+        const std::vector<StorageObjectMetadata>& metadatas);
+
+    /**
+     * @brief Heartbeat-driven pull of pending L2->L1 promotion work for a
+     * client. Returns tenant-scoped tasks the caller should read from local
+     * SSD and stage as MEMORY replicas via PromotionAllocStart +
+     * NotifyPromotionSuccess.
+     */
+    [[nodiscard]] tl::expected<std::vector<PromotionTaskItem>, ErrorCode>
+    PromotionObjectHeartbeat(const UUID& client_id);
+
+    /**
+     * @brief Stage a PROCESSING MEMORY replica for an existing key during
+     * promotion. Returns the new replica's descriptor that the caller writes
+     * via Transfer Engine.
+     */
+    [[nodiscard]] tl::expected<PromotionAllocStartResponse, ErrorCode>
+    PromotionAllocStart(const UUID& client_id, const std::string& key,
+                        uint64_t size,
+                        const std::vector<std::string>& preferred_segments);
+    [[nodiscard]] tl::expected<PromotionAllocStartResponse, ErrorCode>
+    PromotionAllocStart(const UUID& client_id, const std::string& key,
+                        const std::string& tenant_id, uint64_t size,
+                        const std::vector<std::string>& preferred_segments);
+
+    /**
+     * @brief Release master-side promotion task state after a client-side
+     * failure that prevents the holder from calling NotifyPromotionSuccess.
+     * Idempotent; returns OK if the task was already swept by the reaper.
+     */
+    [[nodiscard]] tl::expected<void, ErrorCode> NotifyPromotionFailure(
+        const UUID& client_id, const std::string& key);
+    [[nodiscard]] tl::expected<void, ErrorCode> NotifyPromotionFailure(
+        const UUID& client_id, const std::string& key,
+        const std::string& tenant_id);
+
+    /**
+     * @brief Commit a staged MEMORY replica to COMPLETE; called after the
+     * client has written the bytes via Transfer Engine.
+     */
+    [[nodiscard]] tl::expected<void, ErrorCode> NotifyPromotionSuccess(
+        const UUID& client_id, const std::string& key);
+    [[nodiscard]] tl::expected<void, ErrorCode> NotifyPromotionSuccess(
+        const UUID& client_id, const std::string& key,
+        const std::string& tenant_id);
 
     /**
      * @brief Start a copy operation
@@ -371,6 +503,10 @@ class MasterClient {
     [[nodiscard]] tl::expected<CopyStartResponse, ErrorCode> CopyStart(
         const std::string& key, const std::string& src_segment,
         const std::vector<std::string>& tgt_segments);
+    [[nodiscard]] tl::expected<CopyStartResponse, ErrorCode> CopyStart(
+        const std::string& key, const std::string& tenant_id,
+        const std::string& src_segment,
+        const std::vector<std::string>& tgt_segments);
 
     /**
      * @brief End a copy operation
@@ -378,6 +514,8 @@ class MasterClient {
      * @return tl::expected<void, ErrorCode> indicating success/failure
      */
     [[nodiscard]] tl::expected<void, ErrorCode> CopyEnd(const std::string& key);
+    [[nodiscard]] tl::expected<void, ErrorCode> CopyEnd(
+        const std::string& key, const std::string& tenant_id);
 
     /**
      * @brief Revoke a copy operation
@@ -386,6 +524,8 @@ class MasterClient {
      */
     [[nodiscard]] tl::expected<void, ErrorCode> CopyRevoke(
         const std::string& key);
+    [[nodiscard]] tl::expected<void, ErrorCode> CopyRevoke(
+        const std::string& key, const std::string& tenant_id);
 
     /**
      * @brief Start a move operation
@@ -398,6 +538,9 @@ class MasterClient {
     [[nodiscard]] tl::expected<MoveStartResponse, ErrorCode> MoveStart(
         const std::string& key, const std::string& src_segment,
         const std::string& tgt_segment);
+    [[nodiscard]] tl::expected<MoveStartResponse, ErrorCode> MoveStart(
+        const std::string& key, const std::string& tenant_id,
+        const std::string& src_segment, const std::string& tgt_segment);
 
     /**
      * @brief End a move operation
@@ -405,6 +548,8 @@ class MasterClient {
      * @return tl::expected<void, ErrorCode> indicating success/failure
      */
     [[nodiscard]] tl::expected<void, ErrorCode> MoveEnd(const std::string& key);
+    [[nodiscard]] tl::expected<void, ErrorCode> MoveEnd(
+        const std::string& key, const std::string& tenant_id);
 
     /**
      * @brief Revoke a move operation
@@ -413,6 +558,8 @@ class MasterClient {
      */
     [[nodiscard]] tl::expected<void, ErrorCode> MoveRevoke(
         const std::string& key);
+    [[nodiscard]] tl::expected<void, ErrorCode> MoveRevoke(
+        const std::string& key, const std::string& tenant_id);
 
     /**
      * @brief Create a task to copy an object's replica to target segments
@@ -423,6 +570,9 @@ class MasterClient {
      */
     [[nodiscard]] tl::expected<UUID, ErrorCode> CreateCopyTask(
         const std::string& key, const std::vector<std::string>& targets);
+    [[nodiscard]] tl::expected<UUID, ErrorCode> CreateCopyTask(
+        const std::string& key, const std::string& tenant_id,
+        const std::vector<std::string>& targets);
 
     /**
      * @brief Create a task to move an object's replica from source segment to
@@ -436,6 +586,9 @@ class MasterClient {
     [[nodiscard]] tl::expected<UUID, ErrorCode> CreateMoveTask(
         const std::string& key, const std::string& source,
         const std::string& target);
+    [[nodiscard]] tl::expected<UUID, ErrorCode> CreateMoveTask(
+        const std::string& key, const std::string& tenant_id,
+        const std::string& source, const std::string& target);
 
     /**
      * @brief Query a task by task id
@@ -471,6 +624,9 @@ class MasterClient {
      */
     [[nodiscard]] tl::expected<void, ErrorCode> EvictDiskReplica(
         const std::string& key, ReplicaType replica_type);
+    [[nodiscard]] tl::expected<void, ErrorCode> EvictDiskReplica(
+        const std::string& key, const std::string& tenant_id,
+        ReplicaType replica_type);
 
     /**
      * @brief Batch notify master that disk replicas were evicted locally.
@@ -482,6 +638,10 @@ class MasterClient {
     [[nodiscard]] std::vector<tl::expected<void, ErrorCode>>
     BatchEvictDiskReplica(const std::vector<std::string>& keys,
                           ReplicaType replica_type);
+    [[nodiscard]] std::vector<tl::expected<void, ErrorCode>>
+    BatchEvictDiskReplica(const std::vector<std::string>& keys,
+                          const std::string& tenant_id,
+                          ReplicaType replica_type);
 
    private:
     /**
@@ -539,6 +699,9 @@ class MasterClient {
     // The client identification.
     const UUID client_id_;
 
+    // Tenant identity for this client instance.
+    const std::string tenant_id_;
+
     // Metrics for tracking RPC operations
     MasterClientMetric* metrics_;
     std::shared_ptr<coro_io::client_pools<coro_rpc::coro_rpc_client>>
diff --git a/mooncake-store/include/master_config.h b/mooncake-store/include/master_config.h
index ea3bf4d8..b5929a85 100644
--- a/mooncake-store/include/master_config.h
+++ b/mooncake-store/include/master_config.h
@@ -1,7 +1,10 @@
 #pragma once
 
+#include <cmath>
 #include <optional>
 #include <stdexcept>
+#include <string_view>
+#include <unordered_map>
 
 #include <glog/logging.h>
 
@@ -10,6 +13,53 @@
 
 namespace mooncake {
 
+// Forwarded to the HA serve phase via MasterServiceSupervisorConfig.
+class HttpMetadataServer;
+
+struct ObjectTypeEvictionScorePolicy {
+    double reuse_scale{1.0};
+    double soft_pin_weight{1.0};
+    int64_t eviction_grace{0};
+};
+
+struct ObjectTypeEvictionPolicy {
+    double budget_ratio{1.0};
+};
+
+inline void ValidateObjectTypeEvictionPolicy(
+    const ObjectTypeEvictionPolicy& policy) {
+    if (!std::isfinite(policy.budget_ratio) || policy.budget_ratio < 0.0 ||
+        policy.budget_ratio > 1.0) {
+        throw std::invalid_argument(
+            "budget_ratio must be finite and between 0.0 and 1.0");
+    }
+}
+
+inline void ValidateObjectTypeEvictionScorePolicy(
+    const ObjectTypeEvictionScorePolicy& policy) {
+    if (!std::isfinite(policy.reuse_scale) || policy.reuse_scale <= 0.0) {
+        throw std::invalid_argument(
+            "reuse_scale must be finite and greater than 0");
+    }
+    if (!std::isfinite(policy.soft_pin_weight) ||
+        policy.soft_pin_weight < 0.0) {
+        throw std::invalid_argument(
+            "soft_pin_weight must be finite and non-negative");
+    }
+}
+
+inline std::string ResolveConfiguredHABackendConnstring(
+    std::string_view ha_backend_type, std::string_view ha_backend_connstring,
+    std::string_view etcd_endpoints) {
+    if (!ha_backend_connstring.empty()) {
+        return std::string(ha_backend_connstring);
+    }
+    if (ha_backend_type == "etcd") {
+        return std::string(etcd_endpoints);
+    }
+    return {};
+}
+
 // The configuration for the master server
 struct MasterConfig {
     bool enable_metric_reporting;
@@ -23,10 +73,19 @@ struct MasterConfig {
 
     uint64_t default_kv_lease_ttl;
     uint64_t default_kv_soft_pin_ttl;
+    std::unordered_map<ObjectDataType, ObjectTypeEvictionScorePolicy>
+        object_type_eviction_score_policies;
+    std::unordered_map<ObjectDataType, ObjectTypeEvictionPolicy>
+        object_type_eviction_policies;
     bool allow_evict_soft_pinned_objects;
     double eviction_ratio;
     double eviction_high_watermark_ratio;
+    double nof_eviction_ratio;
+    double nof_eviction_high_watermark_ratio;
     int64_t client_live_ttl_sec;
+    int64_t nof_heartbeat_interval_sec;
+    uint32_t nof_heartbeat_probe_timeout_ms;
+    uint32_t nof_heartbeat_failures_threshold;
 
     bool enable_ha;
     bool enable_offload;
@@ -44,6 +103,15 @@ struct MasterConfig {
     bool enable_http_metadata_server;
     uint32_t http_metadata_server_port;
     std::string http_metadata_server_host;
+    // Enable cleanup of HTTP metadata (mooncake/ram/*, mooncake/rpc_meta/*)
+    // when client heartbeat times out. Works in two modes: (1) co-located
+    // (enable_http_metadata_server=true) via in-process removal, or
+    // (2) separately-deployed metadata server via async HTTP DELETE.
+    bool enable_metadata_cleanup_on_timeout;
+
+    // Pod identity for K8s label-based routing
+    std::string pod_name;
+    std::string pod_namespace;
 
     uint64_t put_start_discard_timeout_sec;
     uint64_t put_start_release_timeout_sec;
@@ -51,6 +119,9 @@ struct MasterConfig {
     // Storage backend eviction configuration
     bool enable_disk_eviction;
     uint64_t quota_bytes;
+    bool enable_multi_tenants = false;
+    std::string tenant_quota_connector_type = "file";
+    std::string tenant_quota_connector_uri;
 
     bool enable_snapshot_restore;
     bool enable_snapshot;
@@ -85,6 +156,20 @@ struct MasterConfig {
     // Offload-on-evict: defer LOCAL_DISK offload to eviction time
     bool offload_on_evict = false;
     bool offload_force_evict = false;
+    size_t offloading_queue_limit = 50000;
+    double offload_cap_ratio = 0.5;
+
+    // Promotion-on-hit: when Get observes a LOCAL_DISK-only key, queue an
+    // async copy back to MEMORY so the next Get is fast.
+    bool promotion_on_hit = false;
+    uint32_t promotion_admission_threshold = 2;
+    uint32_t promotion_queue_limit = 50000;
+    // Max promotion tasks PromotionObjectHeartbeat returns to a single
+    // client per call. Each task is a synchronous SSD-read + RDMA-write
+    // on the client; serializing them avoids blocking past the client-
+    // liveness window. Default 1 is conservative; small-object or RDMA-
+    // rich clusters may safely raise it.
+    uint32_t promotion_max_per_heartbeat = 1;
 };
 
 class MasterServiceSupervisorConfig {
@@ -94,12 +179,25 @@ class MasterServiceSupervisorConfig {
     RequiredParam<int> metrics_port{"metrics_port"};
     RequiredParam<int64_t> default_kv_lease_ttl{"default_kv_lease_ttl"};
     RequiredParam<int64_t> default_kv_soft_pin_ttl{"default_kv_soft_pin_ttl"};
+    std::unordered_map<ObjectDataType, ObjectTypeEvictionScorePolicy>
+        object_type_eviction_score_policies;
+    std::unordered_map<ObjectDataType, ObjectTypeEvictionPolicy>
+        object_type_eviction_policies;
     RequiredParam<bool> allow_evict_soft_pinned_objects{
         "allow_evict_soft_pinned_objects"};
     RequiredParam<double> eviction_ratio{"eviction_ratio"};
     RequiredParam<double> eviction_high_watermark_ratio{
         "eviction_high_watermark_ratio"};
+    RequiredParam<double> nof_eviction_ratio{"nof_eviction_ratio"};
+    RequiredParam<double> nof_eviction_high_watermark_ratio{
+        "nof_eviction_high_watermark_ratio"};
     RequiredParam<int64_t> client_live_ttl_sec{"client_live_ttl_sec"};
+    RequiredParam<int64_t> nof_heartbeat_interval_sec{
+        "nof_heartbeat_interval_sec"};
+    RequiredParam<uint32_t> nof_heartbeat_probe_timeout_ms{
+        "nof_heartbeat_probe_timeout_ms"};
+    RequiredParam<uint32_t> nof_heartbeat_failures_threshold{
+        "nof_heartbeat_failures_threshold"};
     RequiredParam<bool> enable_offload{"enable_offload"};
     RequiredParam<int> rpc_port{"rpc_port"};
     RequiredParam<size_t> rpc_thread_num{"rpc_thread_num"};
@@ -121,6 +219,9 @@ class MasterServiceSupervisorConfig {
     uint64_t put_start_release_timeout_sec = DEFAULT_PUT_START_RELEASE_TIMEOUT;
     bool enable_disk_eviction = true;
     uint64_t quota_bytes = 0;
+    bool enable_multi_tenants = false;
+    std::string tenant_quota_connector_type = "file";
+    std::string tenant_quota_connector_uri;
     uint32_t max_total_finished_tasks = DEFAULT_MAX_TOTAL_FINISHED_TASKS;
     uint32_t max_total_pending_tasks = DEFAULT_MAX_TOTAL_PENDING_TASKS;
     uint32_t max_total_processing_tasks = DEFAULT_MAX_TOTAL_PROCESSING_TASKS;
@@ -146,6 +247,24 @@ class MasterServiceSupervisorConfig {
     bool enable_cxl = false;
     bool offload_on_evict = false;
     bool offload_force_evict = false;
+    size_t offloading_queue_limit = 50000;
+    double offload_cap_ratio = 0.5;
+    bool promotion_on_hit = false;
+    uint32_t promotion_admission_threshold = 2;
+    uint32_t promotion_queue_limit = 50000;
+    uint32_t promotion_max_per_heartbeat = 1;
+
+    // Pod identity for K8s label-based routing
+    std::string pod_name;
+    std::string pod_namespace;
+
+    // Metadata cleanup on client timeout. Resolved in main() (not from
+    // MasterConfig) and forwarded to the serving primary's
+    // WrappedMasterService. Co-located: in-process server pointer; separate:
+    // derived http(s) URL.
+    HttpMetadataServer* http_metadata_server = nullptr;
+    std::string http_metadata_remote_url;
+
     MasterServiceSupervisorConfig() = default;
 
     // From MasterConfig
@@ -155,14 +274,30 @@ class MasterServiceSupervisorConfig {
         metrics_port = static_cast<int>(config.metrics_port);
         default_kv_lease_ttl = config.default_kv_lease_ttl;
         default_kv_soft_pin_ttl = config.default_kv_soft_pin_ttl;
+        object_type_eviction_score_policies =
+            config.object_type_eviction_score_policies;
+        object_type_eviction_policies = config.object_type_eviction_policies;
         allow_evict_soft_pinned_objects =
             config.allow_evict_soft_pinned_objects;
         eviction_ratio = config.eviction_ratio;
         eviction_high_watermark_ratio = config.eviction_high_watermark_ratio;
+        nof_eviction_ratio = config.nof_eviction_ratio;
+        nof_eviction_high_watermark_ratio =
+            config.nof_eviction_high_watermark_ratio;
         client_live_ttl_sec = config.client_live_ttl_sec;
+        nof_heartbeat_interval_sec = config.nof_heartbeat_interval_sec;
+        nof_heartbeat_probe_timeout_ms = config.nof_heartbeat_probe_timeout_ms;
+        nof_heartbeat_failures_threshold =
+            config.nof_heartbeat_failures_threshold;
         enable_offload = config.enable_offload;
         offload_on_evict = config.offload_on_evict;
         offload_force_evict = config.offload_force_evict;
+        offloading_queue_limit = config.offloading_queue_limit;
+        offload_cap_ratio = config.offload_cap_ratio;
+        promotion_on_hit = config.promotion_on_hit;
+        promotion_admission_threshold = config.promotion_admission_threshold;
+        promotion_queue_limit = config.promotion_queue_limit;
+        promotion_max_per_heartbeat = config.promotion_max_per_heartbeat;
         rpc_port = static_cast<int>(config.rpc_port);
         rpc_thread_num = static_cast<size_t>(config.rpc_thread_num);
 
@@ -172,11 +307,9 @@ class MasterServiceSupervisorConfig {
             std::chrono::seconds(config.rpc_conn_timeout_seconds);
         rpc_enable_tcp_no_delay = config.rpc_enable_tcp_no_delay;
         ha_backend_type = config.ha_backend_type;
-        ha_backend_connstring = config.ha_backend_connstring;
         etcd_endpoints = config.etcd_endpoints;
-        if (ha_backend_connstring.empty()) {
-            ha_backend_connstring = etcd_endpoints;
-        }
+        ha_backend_connstring = ResolveConfiguredHABackendConnstring(
+            ha_backend_type, config.ha_backend_connstring, etcd_endpoints);
         local_hostname = rpc_address + ":" + std::to_string(rpc_port);
         cluster_id = config.cluster_id;
         root_fs_dir = config.root_fs_dir;
@@ -193,6 +326,9 @@ class MasterServiceSupervisorConfig {
         put_start_release_timeout_sec = config.put_start_release_timeout_sec;
         enable_disk_eviction = config.enable_disk_eviction;
         quota_bytes = config.quota_bytes;
+        enable_multi_tenants = config.enable_multi_tenants;
+        tenant_quota_connector_type = config.tenant_quota_connector_type;
+        tenant_quota_connector_uri = config.tenant_quota_connector_uri;
 
         enable_snapshot_restore = config.enable_snapshot_restore;
         enable_snapshot = config.enable_snapshot;
@@ -214,6 +350,9 @@ class MasterServiceSupervisorConfig {
         cxl_path = config.cxl_path;
         cxl_size = config.cxl_size;
         enable_cxl = config.enable_cxl;
+
+        pod_name = config.pod_name;
+        pod_namespace = config.pod_namespace;
         validate();
     }
 
@@ -245,9 +384,27 @@ class MasterServiceSupervisorConfig {
             throw std::runtime_error(
                 "eviction_high_watermark_ratio is not set");
         }
+        if (!nof_eviction_ratio.IsSet()) {
+            throw std::runtime_error("nof_eviction_ratio is not set");
+        }
+        if (!nof_eviction_high_watermark_ratio.IsSet()) {
+            throw std::runtime_error(
+                "nof_eviction_high_watermark_ratio is not set");
+        }
         if (!client_live_ttl_sec.IsSet()) {
             throw std::runtime_error("client_live_ttl_sec is not set");
         }
+        if (!nof_heartbeat_interval_sec.IsSet()) {
+            throw std::runtime_error("nof_heartbeat_interval_sec is not set");
+        }
+        if (!nof_heartbeat_probe_timeout_ms.IsSet()) {
+            throw std::runtime_error(
+                "nof_heartbeat_probe_timeout_ms is not set");
+        }
+        if (!nof_heartbeat_failures_threshold.IsSet()) {
+            throw std::runtime_error(
+                "nof_heartbeat_failures_threshold is not set");
+        }
         if (!rpc_port.IsSet()) {
             throw std::runtime_error("rpc_port is not set");
         }
@@ -264,6 +421,10 @@ class WrappedMasterServiceConfig {
 
     // Optional parameters (with default values)
     uint64_t default_kv_soft_pin_ttl = DEFAULT_KV_SOFT_PIN_TTL_MS;
+    std::unordered_map<ObjectDataType, ObjectTypeEvictionScorePolicy>
+        object_type_eviction_score_policies;
+    std::unordered_map<ObjectDataType, ObjectTypeEvictionPolicy>
+        object_type_eviction_policies;
     bool allow_evict_soft_pinned_objects =
         DEFAULT_ALLOW_EVICT_SOFT_PINNED_OBJECTS;
     bool enable_metric_reporting = true;
@@ -271,12 +432,26 @@ class WrappedMasterServiceConfig {
     double eviction_ratio = DEFAULT_EVICTION_RATIO;
     double eviction_high_watermark_ratio =
         DEFAULT_EVICTION_HIGH_WATERMARK_RATIO;
+    double nof_eviction_ratio = DEFAULT_NOF_EVICTION_RATIO;
+    double nof_eviction_high_watermark_ratio =
+        DEFAULT_NOF_EVICTION_HIGH_WATERMARK_RATIO;
     ViewVersionId view_version = 0;
     int64_t client_live_ttl_sec = DEFAULT_CLIENT_LIVE_TTL_SEC;
+    int64_t nof_heartbeat_interval_sec = DEFAULT_NOF_HEARTBEAT_INTERVAL_SEC;
+    uint32_t nof_heartbeat_probe_timeout_ms =
+        DEFAULT_NOF_HEARTBEAT_PROBE_TIMEOUT_MS;
+    uint32_t nof_heartbeat_failures_threshold =
+        DEFAULT_NOF_HEARTBEAT_FAILURES_THRESHOLD;
     bool enable_ha = false;
     bool enable_offload = false;
     bool offload_on_evict = false;
     bool offload_force_evict = false;
+    size_t offloading_queue_limit = 50000;
+    double offload_cap_ratio = 0.5;
+    bool promotion_on_hit = false;
+    uint32_t promotion_admission_threshold = 2;
+    uint32_t promotion_queue_limit = 50000;
+    uint32_t promotion_max_per_heartbeat = 1;
     std::string ha_backend_type = "etcd";
     std::string ha_backend_connstring;
     std::string cluster_id = DEFAULT_CLUSTER_ID;
@@ -289,6 +464,9 @@ class WrappedMasterServiceConfig {
     uint64_t put_start_release_timeout_sec = DEFAULT_PUT_START_RELEASE_TIMEOUT;
     bool enable_disk_eviction = true;
     uint64_t quota_bytes = 0;
+    bool enable_multi_tenants = false;
+    std::string tenant_quota_connector_type = "file";
+    std::string tenant_quota_connector_uri;
 
     bool enable_snapshot_restore = false;
     bool enable_snapshot = false;
@@ -322,28 +500,46 @@ class WrappedMasterServiceConfig {
 
         // Set optional parameters (these have default values)
         default_kv_soft_pin_ttl = config.default_kv_soft_pin_ttl;
+        object_type_eviction_score_policies =
+            config.object_type_eviction_score_policies;
+        object_type_eviction_policies = config.object_type_eviction_policies;
         allow_evict_soft_pinned_objects =
             config.allow_evict_soft_pinned_objects;
         enable_metric_reporting = config.enable_metric_reporting;
         http_port = static_cast<uint16_t>(config.metrics_port);
         eviction_ratio = config.eviction_ratio;
         eviction_high_watermark_ratio = config.eviction_high_watermark_ratio;
+        nof_eviction_ratio = config.nof_eviction_ratio;
+        nof_eviction_high_watermark_ratio =
+            config.nof_eviction_high_watermark_ratio;
         view_version = view_version_param;
         client_live_ttl_sec = config.client_live_ttl_sec;
+        nof_heartbeat_interval_sec = config.nof_heartbeat_interval_sec;
+        nof_heartbeat_probe_timeout_ms = config.nof_heartbeat_probe_timeout_ms;
+        nof_heartbeat_failures_threshold =
+            config.nof_heartbeat_failures_threshold;
         enable_ha = config.enable_ha;
         enable_offload = config.enable_offload;
         offload_on_evict = config.offload_on_evict;
         offload_force_evict = config.offload_force_evict;
+        offloading_queue_limit = config.offloading_queue_limit;
+        offload_cap_ratio = config.offload_cap_ratio;
+        promotion_on_hit = config.promotion_on_hit;
+        promotion_admission_threshold = config.promotion_admission_threshold;
+        promotion_queue_limit = config.promotion_queue_limit;
+        promotion_max_per_heartbeat = config.promotion_max_per_heartbeat;
         ha_backend_type = config.ha_backend_type;
-        ha_backend_connstring = config.ha_backend_connstring;
-        if (ha_backend_connstring.empty()) {
-            ha_backend_connstring = config.etcd_endpoints;
-        }
+        ha_backend_connstring = ResolveConfiguredHABackendConnstring(
+            ha_backend_type, config.ha_backend_connstring,
+            config.etcd_endpoints);
         cluster_id = config.cluster_id;
         root_fs_dir = config.root_fs_dir;
         global_file_segment_size = config.global_file_segment_size;
         enable_disk_eviction = config.enable_disk_eviction;
         quota_bytes = config.quota_bytes;
+        enable_multi_tenants = config.enable_multi_tenants;
+        tenant_quota_connector_type = config.tenant_quota_connector_type;
+        tenant_quota_connector_uri = config.tenant_quota_connector_uri;
 
         // Convert string memory_allocator to BufferAllocatorType enum
         if (config.memory_allocator == "cachelib") {
@@ -359,12 +555,15 @@ class WrappedMasterServiceConfig {
             allocation_strategy_type = AllocationStrategyType::CXL;
         } else if (config.allocation_strategy == "random") {
             allocation_strategy_type = AllocationStrategyType::RANDOM;
+        } else if (config.allocation_strategy == "ssd_free_ratio_first") {
+            allocation_strategy_type =
+                AllocationStrategyType::SSD_FREE_RATIO_FIRST;
         } else {
             LOG(WARNING) << "Unrecognized allocation_strategy value: '"
                          << config.allocation_strategy
                          << "'. Defaulting to 'random'. "
-                         << "Valid options are: random, free_ratio_first, cxl "
-                            "(case-sensitive)";
+                         << "Valid options are: random, free_ratio_first, cxl, "
+                            "ssd_free_ratio_first (case-sensitive)";
             allocation_strategy_type = AllocationStrategyType::RANDOM;
         }
 
@@ -401,12 +600,18 @@ class WrappedMasterServiceConfig {
 
         // Set optional parameters (these have default values)
         default_kv_soft_pin_ttl = config.default_kv_soft_pin_ttl;
+        object_type_eviction_score_policies =
+            config.object_type_eviction_score_policies;
+        object_type_eviction_policies = config.object_type_eviction_policies;
         allow_evict_soft_pinned_objects =
             config.allow_evict_soft_pinned_objects;
         enable_metric_reporting = config.enable_metric_reporting;
         http_port = static_cast<uint16_t>(config.metrics_port);
         eviction_ratio = config.eviction_ratio;
         eviction_high_watermark_ratio = config.eviction_high_watermark_ratio;
+        nof_eviction_ratio = config.nof_eviction_ratio;
+        nof_eviction_high_watermark_ratio =
+            config.nof_eviction_high_watermark_ratio;
         view_version = view_version_param;
         client_live_ttl_sec = config.client_live_ttl_sec;
         enable_ha =
@@ -414,17 +619,25 @@ class WrappedMasterServiceConfig {
         enable_offload = config.enable_offload;
         offload_on_evict = config.offload_on_evict;
         offload_force_evict = config.offload_force_evict;
+        offloading_queue_limit = config.offloading_queue_limit;
+        offload_cap_ratio = config.offload_cap_ratio;
+        promotion_on_hit = config.promotion_on_hit;
+        promotion_admission_threshold = config.promotion_admission_threshold;
+        promotion_queue_limit = config.promotion_queue_limit;
+        promotion_max_per_heartbeat = config.promotion_max_per_heartbeat;
         ha_backend_type = config.ha_backend_type;
-        ha_backend_connstring = config.ha_backend_connstring;
-        if (ha_backend_connstring.empty()) {
-            ha_backend_connstring = config.etcd_endpoints;
-        }
+        ha_backend_connstring = ResolveConfiguredHABackendConnstring(
+            ha_backend_type, config.ha_backend_connstring,
+            config.etcd_endpoints);
         cluster_id = config.cluster_id;
         root_fs_dir = config.root_fs_dir;
         global_file_segment_size = config.global_file_segment_size;
         memory_allocator = config.memory_allocator;
         enable_disk_eviction = config.enable_disk_eviction;
         quota_bytes = config.quota_bytes;
+        enable_multi_tenants = config.enable_multi_tenants;
+        tenant_quota_connector_type = config.tenant_quota_connector_type;
+        tenant_quota_connector_uri = config.tenant_quota_connector_uri;
         put_start_discard_timeout_sec = config.put_start_discard_timeout_sec;
         put_start_release_timeout_sec = config.put_start_release_timeout_sec;
 
@@ -459,13 +672,25 @@ class MasterServiceConfigBuilder {
    private:
     uint64_t default_kv_lease_ttl_ = DEFAULT_DEFAULT_KV_LEASE_TTL;
     uint64_t default_kv_soft_pin_ttl_ = DEFAULT_KV_SOFT_PIN_TTL_MS;
+    std::unordered_map<ObjectDataType, ObjectTypeEvictionScorePolicy>
+        object_type_eviction_score_policies_;
+    std::unordered_map<ObjectDataType, ObjectTypeEvictionPolicy>
+        object_type_eviction_policies_;
     bool allow_evict_soft_pinned_objects_ =
         DEFAULT_ALLOW_EVICT_SOFT_PINNED_OBJECTS;
     double eviction_ratio_ = DEFAULT_EVICTION_RATIO;
     double eviction_high_watermark_ratio_ =
         DEFAULT_EVICTION_HIGH_WATERMARK_RATIO;
+    double nof_eviction_ratio_ = DEFAULT_NOF_EVICTION_RATIO;
+    double nof_eviction_high_watermark_ratio_ =
+        DEFAULT_NOF_EVICTION_HIGH_WATERMARK_RATIO;
     ViewVersionId view_version_ = 0;
     int64_t client_live_ttl_sec_ = DEFAULT_CLIENT_LIVE_TTL_SEC;
+    int64_t nof_heartbeat_interval_sec_ = DEFAULT_NOF_HEARTBEAT_INTERVAL_SEC;
+    uint32_t nof_heartbeat_probe_timeout_ms_ =
+        DEFAULT_NOF_HEARTBEAT_PROBE_TIMEOUT_MS;
+    uint32_t nof_heartbeat_failures_threshold_ =
+        DEFAULT_NOF_HEARTBEAT_FAILURES_THRESHOLD;
     bool enable_ha_ = false;
     bool enable_offload_ = false;
     std::string ha_backend_type_ = "etcd";
@@ -478,6 +703,9 @@ class MasterServiceConfigBuilder {
         AllocationStrategyType::RANDOM;
     bool enable_disk_eviction_ = true;
     uint64_t quota_bytes_ = 0;
+    bool enable_multi_tenants_ = false;
+    std::string tenant_quota_connector_type_ = "file";
+    std::string tenant_quota_connector_uri_;
     uint64_t put_start_discard_timeout_sec_ = DEFAULT_PUT_START_DISCARD_TIMEOUT;
     uint64_t put_start_release_timeout_sec_ = DEFAULT_PUT_START_RELEASE_TIMEOUT;
     bool enable_snapshot_restore_ = false;
@@ -514,6 +742,20 @@ class MasterServiceConfigBuilder {
         return *this;
     }
 
+    MasterServiceConfigBuilder& set_object_type_eviction_score_policy(
+        ObjectDataType data_type, ObjectTypeEvictionScorePolicy policy) {
+        ValidateObjectTypeEvictionScorePolicy(policy);
+        object_type_eviction_score_policies_[data_type] = policy;
+        return *this;
+    }
+
+    MasterServiceConfigBuilder& set_object_type_eviction_policy(
+        ObjectDataType data_type, ObjectTypeEvictionPolicy policy) {
+        ValidateObjectTypeEvictionPolicy(policy);
+        object_type_eviction_policies_[data_type] = policy;
+        return *this;
+    }
+
     MasterServiceConfigBuilder& set_allow_evict_soft_pinned_objects(
         bool allow) {
         allow_evict_soft_pinned_objects_ = allow;
@@ -531,6 +773,17 @@ class MasterServiceConfigBuilder {
         return *this;
     }
 
+    MasterServiceConfigBuilder& set_nof_eviction_ratio(double ratio) {
+        nof_eviction_ratio_ = ratio;
+        return *this;
+    }
+
+    MasterServiceConfigBuilder& set_nof_eviction_high_watermark_ratio(
+        double ratio) {
+        nof_eviction_high_watermark_ratio_ = ratio;
+        return *this;
+    }
+
     MasterServiceConfigBuilder& set_view_version(ViewVersionId version) {
         view_version_ = version;
         return *this;
@@ -541,6 +794,23 @@ class MasterServiceConfigBuilder {
         return *this;
     }
 
+    MasterServiceConfigBuilder& set_nof_heartbeat_interval_sec(int64_t ttl) {
+        nof_heartbeat_interval_sec_ = ttl;
+        return *this;
+    }
+
+    MasterServiceConfigBuilder& set_nof_heartbeat_probe_timeout_ms(
+        uint32_t timeout_ms) {
+        nof_heartbeat_probe_timeout_ms_ = timeout_ms;
+        return *this;
+    }
+
+    MasterServiceConfigBuilder& set_nof_heartbeat_failures_threshold(
+        uint32_t threshold) {
+        nof_heartbeat_failures_threshold_ = threshold;
+        return *this;
+    }
+
     MasterServiceConfigBuilder& set_enable_ha(bool enable) {
         enable_ha_ = enable;
         return *this;
@@ -591,6 +861,23 @@ class MasterServiceConfigBuilder {
         return *this;
     }
 
+    MasterServiceConfigBuilder& set_enable_multi_tenants(bool enable) {
+        enable_multi_tenants_ = enable;
+        return *this;
+    }
+
+    MasterServiceConfigBuilder& set_tenant_quota_connector_type(
+        const std::string& type) {
+        tenant_quota_connector_type_ = type;
+        return *this;
+    }
+
+    MasterServiceConfigBuilder& set_tenant_quota_connector_uri(
+        const std::string& uri) {
+        tenant_quota_connector_uri_ = uri;
+        return *this;
+    }
+
     MasterServiceConfigBuilder& set_put_start_discard_timeout_sec(
         uint64_t put_start_discard_timeout_sec) {
         put_start_discard_timeout_sec_ = put_start_discard_timeout_sec;
@@ -741,17 +1028,35 @@ class MasterServiceConfig {
    public:
     uint64_t default_kv_lease_ttl = DEFAULT_DEFAULT_KV_LEASE_TTL;
     uint64_t default_kv_soft_pin_ttl = DEFAULT_KV_SOFT_PIN_TTL_MS;
+    std::unordered_map<ObjectDataType, ObjectTypeEvictionScorePolicy>
+        object_type_eviction_score_policies;
+    std::unordered_map<ObjectDataType, ObjectTypeEvictionPolicy>
+        object_type_eviction_policies;
     bool allow_evict_soft_pinned_objects =
         DEFAULT_ALLOW_EVICT_SOFT_PINNED_OBJECTS;
     double eviction_ratio = DEFAULT_EVICTION_RATIO;
     double eviction_high_watermark_ratio =
         DEFAULT_EVICTION_HIGH_WATERMARK_RATIO;
+    double nof_eviction_ratio = DEFAULT_NOF_EVICTION_RATIO;
+    double nof_eviction_high_watermark_ratio =
+        DEFAULT_NOF_EVICTION_HIGH_WATERMARK_RATIO;
     ViewVersionId view_version = 0;
     int64_t client_live_ttl_sec = DEFAULT_CLIENT_LIVE_TTL_SEC;
+    int64_t nof_heartbeat_interval_sec = DEFAULT_NOF_HEARTBEAT_INTERVAL_SEC;
+    uint32_t nof_heartbeat_probe_timeout_ms =
+        DEFAULT_NOF_HEARTBEAT_PROBE_TIMEOUT_MS;
+    uint32_t nof_heartbeat_failures_threshold =
+        DEFAULT_NOF_HEARTBEAT_FAILURES_THRESHOLD;
     bool enable_ha = false;
     bool enable_offload = false;
     bool offload_on_evict = false;
     bool offload_force_evict = false;
+    size_t offloading_queue_limit = 50000;
+    double offload_cap_ratio = 0.5;
+    bool promotion_on_hit = false;
+    uint32_t promotion_admission_threshold = 2;
+    uint32_t promotion_queue_limit = 50000;
+    uint32_t promotion_max_per_heartbeat = 1;
     std::string ha_backend_type = "etcd";
     std::string ha_backend_connstring;
     std::string cluster_id = DEFAULT_CLUSTER_ID;
@@ -764,6 +1069,9 @@ class MasterServiceConfig {
     uint64_t put_start_release_timeout_sec = DEFAULT_PUT_START_RELEASE_TIMEOUT;
     bool enable_disk_eviction = true;
     uint64_t quota_bytes = 0;
+    bool enable_multi_tenants = false;
+    std::string tenant_quota_connector_type = "file";
+    std::string tenant_quota_connector_uri;
 
     bool enable_snapshot_restore = false;
     bool enable_snapshot = false;
@@ -795,16 +1103,32 @@ class MasterServiceConfig {
 
         default_kv_lease_ttl = config.default_kv_lease_ttl;
         default_kv_soft_pin_ttl = config.default_kv_soft_pin_ttl;
+        object_type_eviction_score_policies =
+            config.object_type_eviction_score_policies;
+        object_type_eviction_policies = config.object_type_eviction_policies;
         allow_evict_soft_pinned_objects =
             config.allow_evict_soft_pinned_objects;
         eviction_ratio = config.eviction_ratio;
         eviction_high_watermark_ratio = config.eviction_high_watermark_ratio;
+        nof_eviction_ratio = config.nof_eviction_ratio;
+        nof_eviction_high_watermark_ratio =
+            config.nof_eviction_high_watermark_ratio;
         view_version = config.view_version;
         client_live_ttl_sec = config.client_live_ttl_sec;
+        nof_heartbeat_interval_sec = config.nof_heartbeat_interval_sec;
+        nof_heartbeat_probe_timeout_ms = config.nof_heartbeat_probe_timeout_ms;
+        nof_heartbeat_failures_threshold =
+            config.nof_heartbeat_failures_threshold;
         enable_ha = config.enable_ha;
         enable_offload = config.enable_offload;
         offload_on_evict = config.offload_on_evict;
         offload_force_evict = config.offload_force_evict;
+        offloading_queue_limit = config.offloading_queue_limit;
+        offload_cap_ratio = config.offload_cap_ratio;
+        promotion_on_hit = config.promotion_on_hit;
+        promotion_admission_threshold = config.promotion_admission_threshold;
+        promotion_queue_limit = config.promotion_queue_limit;
+        promotion_max_per_heartbeat = config.promotion_max_per_heartbeat;
         ha_backend_type = config.ha_backend_type;
         ha_backend_connstring = config.ha_backend_connstring;
         cluster_id = config.cluster_id;
@@ -815,6 +1139,9 @@ class MasterServiceConfig {
         allocation_strategy_type = config.allocation_strategy_type;
         enable_disk_eviction = config.enable_disk_eviction;
         quota_bytes = config.quota_bytes;
+        enable_multi_tenants = config.enable_multi_tenants;
+        tenant_quota_connector_type = config.tenant_quota_connector_type;
+        tenant_quota_connector_uri = config.tenant_quota_connector_uri;
         put_start_discard_timeout_sec = config.put_start_discard_timeout_sec;
         put_start_release_timeout_sec = config.put_start_release_timeout_sec;
 
@@ -854,11 +1181,20 @@ inline MasterServiceConfig MasterServiceConfigBuilder::build() const {
     MasterServiceConfig config;
     config.default_kv_lease_ttl = default_kv_lease_ttl_;
     config.default_kv_soft_pin_ttl = default_kv_soft_pin_ttl_;
+    config.object_type_eviction_score_policies =
+        object_type_eviction_score_policies_;
+    config.object_type_eviction_policies = object_type_eviction_policies_;
     config.allow_evict_soft_pinned_objects = allow_evict_soft_pinned_objects_;
     config.eviction_ratio = eviction_ratio_;
     config.eviction_high_watermark_ratio = eviction_high_watermark_ratio_;
+    config.nof_eviction_ratio = nof_eviction_ratio_;
+    config.nof_eviction_high_watermark_ratio =
+        nof_eviction_high_watermark_ratio_;
     config.view_version = view_version_;
     config.client_live_ttl_sec = client_live_ttl_sec_;
+    config.nof_heartbeat_interval_sec = nof_heartbeat_interval_sec_;
+    config.nof_heartbeat_probe_timeout_ms = nof_heartbeat_probe_timeout_ms_;
+    config.nof_heartbeat_failures_threshold = nof_heartbeat_failures_threshold_;
     config.enable_ha = enable_ha_;
     config.enable_offload = enable_offload_;
     config.ha_backend_type = ha_backend_type_;
@@ -872,6 +1208,9 @@ inline MasterServiceConfig MasterServiceConfigBuilder::build() const {
     config.put_start_release_timeout_sec = put_start_release_timeout_sec_;
     config.enable_disk_eviction = enable_disk_eviction_;
     config.quota_bytes = quota_bytes_;
+    config.enable_multi_tenants = enable_multi_tenants_;
+    config.tenant_quota_connector_type = tenant_quota_connector_type_;
+    config.tenant_quota_connector_uri = tenant_quota_connector_uri_;
     config.enable_snapshot_restore = enable_snapshot_restore_;
     config.enable_snapshot = enable_snapshot_;
     config.snapshot_backup_dir = snapshot_backup_dir_;
diff --git a/mooncake-store/include/master_metric_manager.h b/mooncake-store/include/master_metric_manager.h
index 78c1cb27..5eb96f14 100644
--- a/mooncake-store/include/master_metric_manager.h
+++ b/mooncake-store/include/master_metric_manager.h
@@ -1,7 +1,10 @@
 #pragma once
 
+#include <atomic>
+#include <chrono>
 #include <mutex>
 #include <string>
+#include <unordered_map>
 
 #include "ylt/metric/counter.hpp"
 #include "ylt/metric/gauge.hpp"
@@ -34,10 +37,20 @@ class MasterMetricManager {
     void inc_file_cache_nums(int64_t val = 1);
     void dec_mem_cache_nums(int64_t val = 1);
     void dec_file_cache_nums(int64_t val = 1);
+    void reset_cache_total_nums();
 
     void inc_valid_get_nums(int64_t val = 1);
     void inc_total_get_nums(int64_t val = 1);
 
+    // NoF segment Metrics
+    void inc_allocated_nof_size(const std::string& segment, int64_t val = 1);
+    void dec_allocated_nof_size(const std::string& segment, int64_t val = 1);
+    void reset_allocated_nof_size();
+    void inc_total_nof_capacity(const std::string& segment, int64_t val = 1);
+    void dec_total_nof_capacity(const std::string& segment, int64_t val = 1);
+    void reset_total_nof_capacity();
+    double get_global_nof_used_ratio(void);
+
     enum class CacheHitStat {
         MEMORY_HITS,
         SSD_HITS,
@@ -46,7 +59,16 @@ class MasterMetricManager {
         MEMORY_HIT_RATE,
         SSD_HIT_RATE,
         OVERALL_HIT_RATE,
-        VALID_GET_RATE
+        VALID_GET_RATE,
+        // Clearer aliases for Store-observed counters. The legacy names above
+        // are kept to preserve RPC/API enum values.
+        MEMORY_CURRENT_CACHED_OBJECTS = MEMORY_TOTAL,
+        SSD_CURRENT_CACHED_OBJECTS = SSD_TOTAL,
+        // These values may exceed 1.0 because the numerator is cumulative
+        // while the denominator is the current cached object count.
+        MEMORY_HITS_PER_CURRENT_CACHED_OBJECT = MEMORY_HIT_RATE,
+        SSD_HITS_PER_CURRENT_CACHED_OBJECT = SSD_HIT_RATE,
+        OVERALL_HITS_PER_CURRENT_CACHED_OBJECT = OVERALL_HIT_RATE
     };
     using CacheHitStatDict = std::unordered_map<CacheHitStat, double>;
     void add_stat_to_dict(CacheHitStatDict&, CacheHitStat, double);
@@ -65,6 +87,17 @@ class MasterMetricManager {
     int64_t get_segment_allocated_mem_size(const std::string& segment);
     int64_t get_segment_total_mem_capacity(const std::string& segment);
 
+    // NoF segment Metrics
+    void inc_allocated_nof_size(int64_t val = 1);
+    void dec_allocated_nof_size(int64_t val = 1);
+    void inc_total_nof_capacity(int64_t val = 1);
+    void dec_total_nof_capacity(int64_t val = 1);
+    int64_t get_allocated_nof_size();
+    int64_t get_total_nof_capacity();
+    double get_segment_nof_used_ratio(const std::string& segment);
+    int64_t get_segment_allocated_nof_size(const std::string& segment);
+    int64_t get_segment_total_nof_capacity(const std::string& segment);
+
     // File Storage Metrics
     void inc_allocated_file_size(int64_t val = 1);
     void dec_allocated_file_size(int64_t val = 1);
@@ -73,6 +106,10 @@ class MasterMetricManager {
     int64_t get_allocated_file_size();
     int64_t get_total_file_capacity();
     double get_global_file_used_ratio(void);
+    // When true, get_total_file_capacity() returns INT64_MAX to indicate
+    // unlimited DFS capacity instead of relying on the gauge value.
+    void set_dfs_capacity_unlimited(bool unlimited);
+    bool is_dfs_capacity_unlimited() const;
 
     // Key/Value Metrics
     void inc_key_count(int64_t val = 1);
@@ -96,6 +133,7 @@ class MasterMetricManager {
     // Operation Statistics (Counters)
     void inc_put_start_requests(int64_t val = 1);
     void inc_put_start_failures(int64_t val = 1);
+    void inc_put_start_alloc_failures(int64_t val = 1);
     void inc_put_end_requests(int64_t val = 1);
     void inc_put_end_failures(int64_t val = 1);
     void inc_put_revoke_requests(int64_t val = 1);
@@ -114,12 +152,23 @@ class MasterMetricManager {
     void inc_remove_all_failures(int64_t val = 1);
     void inc_mount_segment_requests(int64_t val = 1);
     void inc_mount_segment_failures(int64_t val = 1);
+    void inc_mount_nof_segment_requests(int64_t val = 1);
+    void inc_mount_nof_segment_failures(int64_t val = 1);
     void inc_unmount_segment_requests(int64_t val = 1);
     void inc_unmount_segment_failures(int64_t val = 1);
+    void inc_unmount_nof_segment_requests(int64_t val = 1);
+    void inc_unmount_nof_segment_failures(int64_t val = 1);
     void inc_remount_segment_requests(int64_t val = 1);
     void inc_remount_segment_failures(int64_t val = 1);
+    void inc_remount_nof_segment_requests(int64_t val = 1);
+    void inc_remount_nof_segment_failures(int64_t val = 1);
     void inc_ping_requests(int64_t val = 1);
     void inc_ping_failures(int64_t val = 1);
+    void inc_nof_heartbeat_success_total(int64_t val = 1);
+    void inc_nof_heartbeat_failure_total(int64_t val = 1);
+    void inc_nof_heartbeat_timeout_total(int64_t val = 1);
+    void inc_nof_segments_unmounted_by_heartbeat_total(int64_t val = 1);
+    void observe_nof_heartbeat_probe_latency_ms(int64_t latency_ms);
 
     // Batch Operation Statistics (Counters)
     void inc_batch_exist_key_requests(int64_t items);
@@ -147,6 +196,7 @@ class MasterMetricManager {
     // Operation Statistics Getters
     int64_t get_put_start_requests();
     int64_t get_put_start_failures();
+    int64_t get_put_start_alloc_failures();
     int64_t get_put_end_requests();
     int64_t get_put_end_failures();
     int64_t get_put_revoke_requests();
@@ -210,14 +260,32 @@ class MasterMetricManager {
     int64_t get_batch_put_revoke_failed_items();
 
     // Eviction Metrics
+    // total eviction metrics
     void inc_eviction_success(int64_t key_count, int64_t size);
     void inc_eviction_fail();  // not a single object is evicted
+    // mem eviction metrics
+    void inc_mem_eviction_success(int64_t key_count, int64_t size);
+    void inc_mem_eviction_fail();  // not a single object is evicted
+    // nof eviction metrics
+    void inc_nof_eviction_success(int64_t key_count, int64_t size);
+    void inc_nof_eviction_fail();  // not a single object is evicted
 
     // Eviction Metrics Getters
+    // total eviction metrics
     int64_t get_eviction_success();
     int64_t get_eviction_attempts();
     int64_t get_evicted_key_count();
     int64_t get_evicted_size();
+    // mem eviction metrics
+    int64_t get_mem_eviction_success();
+    int64_t get_mem_eviction_attempts();
+    int64_t get_mem_evicted_key_count();
+    int64_t get_mem_evicted_size();
+    // nof eviction metrics
+    int64_t get_nof_eviction_success();
+    int64_t get_nof_eviction_attempts();
+    int64_t get_nof_evicted_key_count();
+    int64_t get_nof_evicted_size();
 
     // PutStart Discard Metrics
     void inc_put_start_discard_cnt(int64_t count, int64_t size);
@@ -228,6 +296,36 @@ class MasterMetricManager {
     int64_t get_put_start_release_cnt();
     int64_t get_put_start_discarded_staging_size();
 
+    // Promotion-on-hit Metrics
+    void inc_promotion_in_flight(int64_t val = 1);
+    void dec_promotion_in_flight(int64_t val = 1);
+    void inc_promotion_admitted(int64_t val = 1);
+    void inc_promotion_completed(int64_t val = 1);
+    void inc_promotion_completed_bytes(int64_t bytes);
+    void inc_promotion_expired(int64_t val = 1);
+    void inc_promotion_failed(int64_t val = 1);
+    void inc_promotion_cancelled(int64_t val = 1);
+    void inc_promotion_rejected_frequency(int64_t val = 1);
+    void inc_promotion_rejected_watermark(int64_t val = 1);
+    void inc_promotion_rejected_cap(int64_t val = 1);
+
+    // Tenant quota metrics
+    void inc_tenant_quota_reject(const std::string& tenant_id,
+                                 const std::string& reason, int64_t val = 1);
+    void inc_tenant_evict_bytes(const std::string& tenant_id, int64_t bytes);
+
+    // Promotion-on-hit Metrics Getters
+    int64_t get_promotion_in_flight();
+    int64_t get_promotion_admitted();
+    int64_t get_promotion_completed();
+    int64_t get_promotion_completed_bytes();
+    int64_t get_promotion_expired();
+    int64_t get_promotion_failed();
+    int64_t get_promotion_cancelled();
+    int64_t get_promotion_rejected_frequency();
+    int64_t get_promotion_rejected_watermark();
+    int64_t get_promotion_rejected_cap();
+
     // CopyStart, CopyEnd, CopyRevoke, MoveStart, MoveEnd, MoveRevoke Metrics
     void inc_copy_start_requests(int64_t val = 1);
     void inc_copy_start_failures(int64_t val = 1);
@@ -297,6 +395,7 @@ class MasterMetricManager {
      * @return A string containing the formatted summary.
      */
     std::string get_summary_string();
+    std::string get_summary_string_and_update_snapshot();
 
    private:
     // --- Private Constructor & Destructor ---
@@ -305,8 +404,108 @@ class MasterMetricManager {
 
     // Update all metrics once to ensure zero values are serialized
     void update_metrics_for_zero_output();
+    std::string get_summary_string(bool update_summary_snapshot);
+
+    struct SummaryCounters {
+        int64_t exist_keys = 0;
+        int64_t exist_key_fails = 0;
+        int64_t put_starts = 0;
+        int64_t put_start_fails = 0;
+        int64_t put_start_alloc_fails = 0;
+        int64_t put_ends = 0;
+        int64_t put_end_fails = 0;
+        int64_t put_revoke_requests = 0;
+        int64_t put_revoke_fails = 0;
+        int64_t get_replicas = 0;
+        int64_t get_replica_fails = 0;
+        int64_t removes = 0;
+        int64_t remove_fails = 0;
+        int64_t remove_all = 0;
+        int64_t remove_all_fails = 0;
+        int64_t create_move_tasks = 0;
+        int64_t create_move_task_fails = 0;
+        int64_t create_copy_tasks = 0;
+        int64_t create_copy_task_fails = 0;
+        int64_t query_tasks = 0;
+        int64_t query_task_fails = 0;
+        int64_t fetch_tasks = 0;
+        int64_t fetch_task_fails = 0;
+        int64_t copy_starts = 0;
+        int64_t copy_start_fails = 0;
+        int64_t copy_ends = 0;
+        int64_t copy_end_fails = 0;
+        int64_t copy_revokes = 0;
+        int64_t copy_revoke_fails = 0;
+        int64_t move_starts = 0;
+        int64_t move_start_fails = 0;
+        int64_t move_ends = 0;
+        int64_t move_end_fails = 0;
+        int64_t move_revokes = 0;
+        int64_t move_revoke_fails = 0;
+        int64_t evict_disk_replicas = 0;
+        int64_t evict_disk_replica_fails = 0;
+        int64_t batch_put_start_requests = 0;
+        int64_t batch_put_start_fails = 0;
+        int64_t batch_put_start_partial_successes = 0;
+        int64_t batch_put_start_items = 0;
+        int64_t batch_put_start_failed_items = 0;
+        int64_t batch_put_end_requests = 0;
+        int64_t batch_put_end_fails = 0;
+        int64_t batch_put_end_partial_successes = 0;
+        int64_t batch_put_end_items = 0;
+        int64_t batch_put_end_failed_items = 0;
+        int64_t batch_put_revoke_requests = 0;
+        int64_t batch_put_revoke_fails = 0;
+        int64_t batch_put_revoke_partial_successes = 0;
+        int64_t batch_put_revoke_items = 0;
+        int64_t batch_put_revoke_failed_items = 0;
+        int64_t batch_get_replica_list_requests = 0;
+        int64_t batch_get_replica_list_fails = 0;
+        int64_t batch_get_replica_list_partial_successes = 0;
+        int64_t batch_get_replica_list_items = 0;
+        int64_t batch_get_replica_list_failed_items = 0;
+        int64_t batch_exist_key_requests = 0;
+        int64_t batch_exist_key_fails = 0;
+        int64_t batch_exist_key_partial_successes = 0;
+        int64_t batch_exist_key_items = 0;
+        int64_t batch_exist_key_failed_items = 0;
+        int64_t batch_query_ip_requests = 0;
+        int64_t batch_query_ip_fails = 0;
+        int64_t batch_query_ip_partial_successes = 0;
+        int64_t batch_query_ip_items = 0;
+        int64_t batch_query_ip_failed_items = 0;
+        int64_t batch_replica_clear_requests = 0;
+        int64_t batch_replica_clear_fails = 0;
+        int64_t batch_replica_clear_partial_successes = 0;
+        int64_t batch_replica_clear_items = 0;
+        int64_t batch_replica_clear_failed_items = 0;
+        int64_t eviction_success = 0;
+        int64_t eviction_attempts = 0;
+        int64_t evicted_key_count = 0;
+        int64_t evicted_size = 0;
+        int64_t mem_eviction_success = 0;
+        int64_t mem_eviction_attempts = 0;
+        int64_t mem_evicted_key_count = 0;
+        int64_t mem_evicted_size = 0;
+        int64_t nof_eviction_success = 0;
+        int64_t nof_eviction_attempts = 0;
+        int64_t nof_evicted_key_count = 0;
+        int64_t nof_evicted_size = 0;
+        int64_t ping = 0;
+        int64_t ping_fails = 0;
+        int64_t mark_task_to_complete_requests = 0;
+        int64_t mark_task_to_complete_fails = 0;
+    };
+
+    struct SummarySnapshot {
+        bool initialized = false;
+        std::chrono::steady_clock::time_point timestamp;
+        SummaryCounters counters;
+    };
 
     // --- Metric Members ---
+    std::mutex summary_snapshot_mutex_;
+    SummarySnapshot summary_snapshot_;
 
     // Memory Storage Metrics
     ylt::metric::gauge_t
@@ -320,9 +519,22 @@ class MasterMetricManager {
         mem_total_capacity_per_segment_;  // Segment memory capacity update for
                                           // gauge
 
+    // NoF Segment Metrics
+    ylt::metric::gauge_t
+        nof_allocated_size_;  // Overall NoF SSD usage update for gauge
+    ylt::metric::gauge_t
+        nof_total_capacity_;  // Overall NoF SSD capacity update for gauge
+    ylt::metric::dynamic_gauge_1t
+        nof_allocated_size_per_segment_;  // NoF segment usage update for
+                                          // gauge
+    ylt::metric::dynamic_gauge_1t
+        nof_total_capacity_per_segment_;  // NoF segment capacity update for
+                                          // gauge
+
     // File Storage Metrics
     ylt::metric::gauge_t file_allocated_size_;
     ylt::metric::gauge_t file_total_capacity_;
+    std::atomic<bool> dfs_capacity_unlimited_{false};
 
     // Key/Value Metrics
     ylt::metric::gauge_t key_count_;
@@ -335,6 +547,7 @@ class MasterMetricManager {
     // Operation Statistics
     ylt::metric::counter_t put_start_requests_;
     ylt::metric::counter_t put_start_failures_;
+    ylt::metric::counter_t put_start_alloc_failures_;
     ylt::metric::counter_t put_end_requests_;
     ylt::metric::counter_t put_end_failures_;
     ylt::metric::counter_t put_revoke_requests_;
@@ -357,8 +570,19 @@ class MasterMetricManager {
     ylt::metric::counter_t unmount_segment_failures_;
     ylt::metric::counter_t remount_segment_requests_;
     ylt::metric::counter_t remount_segment_failures_;
+    ylt::metric::counter_t mount_nof_segment_requests_;
+    ylt::metric::counter_t mount_nof_segment_failures_;
+    ylt::metric::counter_t unmount_nof_segment_requests_;
+    ylt::metric::counter_t unmount_nof_segment_failures_;
+    ylt::metric::counter_t remount_nof_segment_requests_;
+    ylt::metric::counter_t remount_nof_segment_failures_;
     ylt::metric::counter_t ping_requests_;
     ylt::metric::counter_t ping_failures_;
+    ylt::metric::counter_t nof_heartbeat_success_total_;
+    ylt::metric::counter_t nof_heartbeat_failure_total_;
+    ylt::metric::counter_t nof_heartbeat_timeout_total_;
+    ylt::metric::counter_t nof_segments_unmounted_by_heartbeat_total_;
+    ylt::metric::histogram_t nof_heartbeat_probe_latency_ms_;
 
     // Batch Operation Statistics
     ylt::metric::counter_t batch_exist_key_requests_;
@@ -397,7 +621,8 @@ class MasterMetricManager {
     ylt::metric::counter_t batch_put_revoke_items_;
     ylt::metric::counter_t batch_put_revoke_failed_items_;
 
-    // cache hit Statistics
+    // Store-observed cache reuse statistics. These counters do not represent
+    // end-to-end request/token-level cache hit ratio.
     ylt::metric::counter_t mem_cache_hit_nums_;
     ylt::metric::counter_t file_cache_hit_nums_;
     ylt::metric::gauge_t mem_cache_nums_;
@@ -417,16 +642,42 @@ class MasterMetricManager {
                        {CacheHitStat::VALID_GET_RATE, "valid_get_rate"}};
 
     // Eviction Metrics
+    // total eviction metrics
     ylt::metric::counter_t eviction_success_;
     ylt::metric::counter_t eviction_attempts_;
     ylt::metric::counter_t evicted_key_count_;
     ylt::metric::counter_t evicted_size_;
+    // mem eviction metrics
+    ylt::metric::counter_t mem_eviction_success_;
+    ylt::metric::counter_t mem_eviction_attempts_;
+    ylt::metric::counter_t mem_evicted_key_count_;
+    ylt::metric::counter_t mem_evicted_size_;
+    // nof eviction metrics
+    ylt::metric::counter_t nof_eviction_success_;
+    ylt::metric::counter_t nof_eviction_attempts_;
+    ylt::metric::counter_t nof_evicted_key_count_;
+    ylt::metric::counter_t nof_evicted_size_;
 
     // PutStart Discard Metrics
     ylt::metric::counter_t put_start_discard_cnt_;
     ylt::metric::counter_t put_start_release_cnt_;
     ylt::metric::gauge_t put_start_discarded_staging_size_;
 
+    // Promotion-on-hit Metrics
+    ylt::metric::gauge_t promotion_in_flight_metric_;
+    ylt::metric::counter_t promotion_admitted_;
+    ylt::metric::counter_t promotion_completed_;
+    ylt::metric::counter_t promotion_completed_bytes_;
+    ylt::metric::counter_t promotion_expired_;
+    ylt::metric::counter_t promotion_failed_;
+    ylt::metric::counter_t promotion_cancelled_;
+    ylt::metric::counter_t promotion_rejected_frequency_;
+    ylt::metric::counter_t promotion_rejected_watermark_;
+    ylt::metric::counter_t promotion_rejected_cap_;
+
+    ylt::metric::dynamic_counter_2t tenant_quota_reject_total_;
+    ylt::metric::dynamic_counter_1t tenant_evict_bytes_total_;
+
     // Snapshot Metrics
     ylt::metric::histogram_t snapshot_duration_ms_;
     ylt::metric::counter_t snapshot_success_;
diff --git a/mooncake-store/include/master_service.h b/mooncake-store/include/master_service.h
index fc4f2c32..ca7d7c35 100644
--- a/mooncake-store/include/master_service.h
+++ b/mooncake-store/include/master_service.h
@@ -1,13 +1,16 @@
 #pragma once
 
+#include <array>
 #include <atomic>
 #include <boost/functional/hash.hpp>
 #include <boost/lockfree/queue.hpp>
 #include <chrono>
 #include <condition_variable>
 #include <cstdint>
+#include <functional>
 #include <list>
 #include <memory>
+#include <mutex>
 #include <optional>
 #include <shared_mutex>
 #include <string>
@@ -19,9 +22,13 @@
 #include <ylt/util/tl/expected.hpp>
 
 #include "allocation_strategy.h"
+#include "count_min_sketch.h"
+#include "deadline_scheduler.h"
 #include "master_metric_manager.h"
 #include "mutex.h"
 #include "segment.h"
+#include "tenant_quota.h"
+#include "tenant_quota_policy_store.h"
 #include "types.h"
 #include "master_config.h"
 #include "rpc_types.h"
@@ -40,30 +47,68 @@ class EtcdOpLogStore;
 // Forward declarations
 class AllocationStrategy;
 class EvictionStrategy;
+class HttpMetadataServer;
+struct MetadataStoragePlugin;
 
 // Forward declarations for test classes
 namespace test {
 class MasterServiceSnapshotTestBase;
 class SnapshotChildProcessTest;
+// Friended so the promotion-on-hit tests can drive a serialize/reset/
+// deserialize cycle directly via the otherwise-private
+// MetadataSerializer, and inspect private clamp fields. This avoids
+// standing up a full snapshot catalog + child-process harness, and
+// exposing test-only accessors on MasterService itself.
+class PromotionOnHitTest;
+class MasterServiceTenantQuotaTest;
 }  // namespace test
 
 /*
  * @brief MasterService is the main class for the master server.
  * Lock order: To avoid deadlocks, the following lock order should be followed:
  * 1. client_mutex_
- * 2. metadata_shards_[shard_idx_].mutex
- * 3. segment_mutex_
+ * 2. tenant_quota_policy_mutex_
+ * 3. snapshot_mutex_
+ * 4. metadata_shards_[shard_idx_].mutex
+ * 5. tenant_quota_shards_[shard_idx_].mutex
+ * 6. segment_mutex_
+ *
+ * Strict tenant admission and policy mutation paths that need both
+ * tenant_quota_policy_mutex_ and snapshot_mutex_ must acquire the tenant
+ * policy mutex first, then snapshot_mutex_.
  */
 class MasterService {
     // Test friend class for snapshot/restore testing
     friend class test::MasterServiceSnapshotTestBase;
     friend class test::SnapshotChildProcessTest;
+    friend class test::PromotionOnHitTest;
+    friend class test::MasterServiceTenantQuotaTest;
 
    public:
+    using NoFProbeFn =
+        std::function<bool(const std::string&, uint32_t, std::string*)>;
+
     MasterService();
     MasterService(const MasterServiceConfig& config);
     ~MasterService();
 
+    void SetNoFProbeFnForTesting(NoFProbeFn fn);
+    size_t GetMountedNoFSegmentCountForTesting();
+    bool IsNoFSegmentMountedForTesting(const UUID& segment_id);
+    std::optional<uint32_t> GetNoFHeartbeatFailureCountForTesting(
+        const UUID& segment_id);
+    std::optional<TenantQuotaSnapshot> GetTenantQuotaSnapshotForTesting(
+        const std::string& tenant_id) const;
+    bool IsTenantQuotaEnabled() const;
+    std::vector<TenantQuotaSnapshot> ListTenantQuotaSnapshots() const;
+    std::optional<TenantQuotaSnapshot> GetTenantQuotaSnapshot(
+        const std::string& tenant_id) const;
+    tl::expected<TenantQuotaSnapshot, ErrorCode> UpsertTenantQuotaPolicy(
+        const std::string& tenant_id, uint64_t requested_quota_bytes);
+    tl::expected<std::optional<TenantQuotaSnapshot>, ErrorCode>
+    DeleteTenantQuotaPolicy(const std::string& tenant_id);
+    uint64_t GetTenantQuotaAllocatableCapacityBytes();
+
     /**
      * @brief Mount a memory segment for buffer allocation. This function is
      * idempotent.
@@ -76,6 +121,18 @@ class MasterService {
     auto MountSegment(const Segment& segment, const UUID& client_id)
         -> tl::expected<void, ErrorCode>;
 
+    /**
+     * @brief Mount a NoF SSD segment for buffer allocation. This function is
+     * idempotent.
+     * @return ErrorCode::OK on success,
+     *         ErrorCode::INVALID_PARAMS on invalid parameters,
+     *         ErrorCode::UNAVAILABLE_IN_CURRENT_STATUS if the segment cannot
+     *         be mounted temporarily,
+     *         ErrorCode::INTERNAL_ERROR on internal errors.
+     */
+    auto MountNoFSegment(const NoFSegment& segment, const UUID& client_id)
+        -> tl::expected<void, ErrorCode>;
+
     /**
      * @brief Re-mount segments, invoked when the client is the first time to
      * connect to the master or the client Ping TTL is expired and need
@@ -90,6 +147,21 @@ class MasterService {
     auto ReMountSegment(const std::vector<Segment>& segments,
                         const UUID& client_id) -> tl::expected<void, ErrorCode>;
 
+    /**
+     * @brief Re-mount NoF SSD segments, invoked when the client is the first
+     * time to connect to the master or the client Ping TTL is expired and need
+     * to remount. This function is idempotent. Client should retry if the
+     * return code is not ErrorCode::OK.
+     * @return ErrorCode::OK means either all segments are remounted
+     * successfully or the fail is not solvable by a new remount request.
+     *         ErrorCode::UNAVAILABLE_IN_CURRENT_STATUS if the segment cannot
+     *         be mounted temporarily.
+     *         ErrorCode::INTERNAL_ERROR if something temporary error happens.
+     */
+    auto ReMountNoFSegment(const std::vector<NoFSegment>& segments,
+                           const UUID& client_id)
+        -> tl::expected<void, ErrorCode>;
+
     /**
      * @brief Unmount a memory segment. This function is idempotent.
      * @return ErrorCode::OK on success,
@@ -99,20 +171,35 @@ class MasterService {
     auto UnmountSegment(const UUID& segment_id, const UUID& client_id)
         -> tl::expected<void, ErrorCode>;
 
+    auto GracefulUnmountSegment(const UUID& segment_id, const UUID& client_id,
+                                uint64_t grace_period_ms)
+        -> tl::expected<void, ErrorCode>;
+
+    /**
+     * @brief Unmount a NoF ssd segment. This function is idempotent.
+     * @return ErrorCode::OK on success,
+     *         ErrorCode::UNAVAILABLE_IN_CURRENT_STATUS if the segment is
+     *         currently unmounting.
+     */
+    auto UnmountNoFSegment(const UUID& segment_id, const UUID& client_id)
+        -> tl::expected<void, ErrorCode>;
+
     /**
      * @brief Check if an object exists
      * @return ErrorCode::OK if exists, otherwise return other ErrorCode
      */
-    auto ExistKey(const std::string& key) -> tl::expected<bool, ErrorCode>;
+    auto ExistKey(const std::string& key, const std::string& tenant_id)
+        -> tl::expected<bool, ErrorCode>;
 
     std::vector<tl::expected<bool, ErrorCode>> BatchExistKey(
-        const std::vector<std::string>& keys);
+        const std::vector<std::string>& keys, const std::string& tenant_id);
 
     /**
-     * @brief Fetch all keys
+     * @brief Fetch all keys for a single tenant.
      * @return ErrorCode::OK if exists
      */
-    auto GetAllKeys() -> tl::expected<std::vector<std::string>, ErrorCode>;
+    auto GetAllKeys(const std::string& tenant_id)
+        -> tl::expected<std::vector<std::string>, ErrorCode>;
 
     /**
      * @brief Fetch all segments, each node has a unique real client with fixed
@@ -122,6 +209,51 @@ class MasterService {
      */
     auto GetAllSegments() -> tl::expected<std::vector<std::string>, ErrorCode>;
 
+    /**
+     * @brief Fetch all mounted NoF segments.
+     * @return std::vector<MountedNoFSegmentSnapshot> on success, error code
+     * otherwise.
+     */
+    auto GetAllNoFSegments()
+        -> tl::expected<std::vector<NoFSegment>, ErrorCode>;
+
+    /**
+     * @brief Query mounted NoF segments by segment name and return their
+     * segment ids together with owner client ids.
+     * @param segment_name Mounted NoF segment name.
+     * @return Matching segment owner info list on success, error code
+     * otherwise.
+     */
+    auto GetNoFSegmentsByName(const std::string& segment_name)
+        -> tl::expected<std::vector<NoFSegmentOwnerInfo>, ErrorCode>;
+
+    /**
+     * @brief Detailed information about a single segment.
+     * Keeps original types so callers can use values directly without
+     * needing to parse strings back to uuid/address/enum.
+     */
+    struct SegmentDetailInfo {
+        std::string segment_name;
+        UUID segment_id{0, 0};
+        UUID client_id{0, 0};
+        uintptr_t base_address{0};
+        uint64_t size_bytes{0};
+        std::string te_endpoint;
+        std::string protocol;
+        SegmentStatus status{SegmentStatus::UNDEFINED};
+        uint64_t allocator_used_bytes{0};
+        uint64_t allocator_capacity_bytes{0};
+    };
+
+    /**
+     * @brief Get detailed information of all segments, including the
+     * relationships between segment_id, client_id, segment_name, status,
+     * allocator used/capacity, etc.
+     * @return A vector of SegmentDetailInfo on success, error code otherwise.
+     */
+    auto GetSegmentsDetail()
+        -> tl::expected<std::vector<SegmentDetailInfo>, ErrorCode>;
+
     /**
      * @brief Query a segment's capacity and used size in bytes.
      * Conductor should use these information to schedule new requests.
@@ -175,7 +307,8 @@ class MasterService {
      * @return An expected object containing a map from object keys to their
      * replica descriptors on success, or an ErrorCode on failure.
      */
-    auto GetReplicaListByRegex(const std::string& regex_pattern)
+    auto GetReplicaListByRegex(const std::string& regex_pattern,
+                               const std::string& tenant_id)
         -> tl::expected<
             std::unordered_map<std::string, std::vector<Replica::Descriptor>>,
             ErrorCode>;
@@ -186,9 +319,16 @@ class MasterService {
      * @return ErrorCode::OK on success, ErrorCode::REPLICA_IS_NOT_READY if not
      * ready
      */
-    auto GetReplicaList(const std::string& key)
+    auto GetReplicaList(const std::string& key, const std::string& tenant_id)
         -> tl::expected<GetReplicaListResponse, ErrorCode>;
 
+    /**
+     * @brief Get replica lists for a batch of objects.
+     */
+    std::vector<tl::expected<GetReplicaListResponse, ErrorCode>>
+    BatchGetReplicaList(const std::vector<std::string>& keys,
+                        const std::string& tenant_id);
+
     /**
      * @brief Start a put operation for an object
      * @param[out] replica_list Vector to store replica information for the
@@ -198,7 +338,8 @@ class MasterService {
      *         ErrorCode::INVALID_PARAMS if slice size is invalid
      */
     auto PutStart(const UUID& client_id, const std::string& key,
-                  const uint64_t slice_length, const ReplicateConfig& config)
+                  const std::string& tenant_id, const uint64_t slice_length,
+                  const ReplicateConfig& config)
         -> tl::expected<std::vector<Replica::Descriptor>, ErrorCode>;
 
     /**
@@ -208,13 +349,15 @@ class MasterService {
      * found, ErrorCode::INVALID_WRITE if replica status is invalid
      */
     auto PutEnd(const UUID& client_id, const std::string& key,
-                ReplicaType replica_type) -> tl::expected<void, ErrorCode>;
+                const std::string& tenant_id, ReplicaType replica_type)
+        -> tl::expected<void, ErrorCode>;
 
     /**
      * @brief Adds a replica instance associated with the given client and key.
      */
     auto AddReplica(const UUID& client_id, const std::string& key,
-                    Replica& replica) -> tl::expected<void, ErrorCode>;
+                    const std::string& tenant_id, Replica& replica)
+        -> tl::expected<bool, ErrorCode>;
 
     /**
      * @brief Revoke a put operation, replica_type indicates the type of
@@ -223,7 +366,8 @@ class MasterService {
      * found, ErrorCode::INVALID_WRITE if replica status is invalid
      */
     auto PutRevoke(const UUID& client_id, const std::string& key,
-                   ReplicaType replica_type) -> tl::expected<void, ErrorCode>;
+                   const std::string& tenant_id, ReplicaType replica_type)
+        -> tl::expected<void, ErrorCode>;
 
     /**
      * @brief Complete a batch of put operations
@@ -231,7 +375,9 @@ class MasterService {
      * found, ErrorCode::INVALID_WRITE if replica status is invalid
      */
     std::vector<tl::expected<void, ErrorCode>> BatchPutEnd(
-        const UUID& client_id, const std::vector<std::string>& keys);
+        const UUID& client_id, const std::vector<std::string>& keys,
+        const std::string& tenant_id,
+        ReplicaType replica_type = ReplicaType::ALL);
 
     /**
      * @brief Revoke a batch of put operations
@@ -239,7 +385,9 @@ class MasterService {
      * found, ErrorCode::INVALID_WRITE if replica status is invalid
      */
     std::vector<tl::expected<void, ErrorCode>> BatchPutRevoke(
-        const UUID& client_id, const std::vector<std::string>& keys);
+        const UUID& client_id, const std::vector<std::string>& keys,
+        const std::string& tenant_id,
+        ReplicaType replica_type = ReplicaType::ALL);
 
     /**
      * @brief Start an upsert operation. If the key does not exist, behaves
@@ -251,20 +399,22 @@ class MasterService {
      * progress), OBJECT_REPLICA_BUSY (replicas have non-zero refcnt).
      */
     auto UpsertStart(const UUID& client_id, const std::string& key,
-                     const uint64_t slice_length, const ReplicateConfig& config)
+                     const std::string& tenant_id, const uint64_t slice_length,
+                     const ReplicateConfig& config)
         -> tl::expected<std::vector<Replica::Descriptor>, ErrorCode>;
 
     /**
      * @brief Complete an upsert operation. Delegates to PutEnd.
      */
     auto UpsertEnd(const UUID& client_id, const std::string& key,
-                   ReplicaType replica_type) -> tl::expected<void, ErrorCode>;
+                   const std::string& tenant_id, ReplicaType replica_type)
+        -> tl::expected<void, ErrorCode>;
 
     /**
      * @brief Revoke an upsert operation. Delegates to PutRevoke.
      */
     auto UpsertRevoke(const UUID& client_id, const std::string& key,
-                      ReplicaType replica_type)
+                      const std::string& tenant_id, ReplicaType replica_type)
         -> tl::expected<void, ErrorCode>;
 
     /**
@@ -273,6 +423,7 @@ class MasterService {
     std::vector<tl::expected<std::vector<Replica::Descriptor>, ErrorCode>>
     BatchUpsertStart(const UUID& client_id,
                      const std::vector<std::string>& keys,
+                     const std::string& tenant_id,
                      const std::vector<uint64_t>& slice_lengths,
                      const ReplicateConfig& config);
 
@@ -280,13 +431,15 @@ class MasterService {
      * @brief Complete a batch of upsert operations. Delegates to BatchPutEnd.
      */
     std::vector<tl::expected<void, ErrorCode>> BatchUpsertEnd(
-        const UUID& client_id, const std::vector<std::string>& keys);
+        const UUID& client_id, const std::vector<std::string>& keys,
+        const std::string& tenant_id);
 
     /**
      * @brief Revoke a batch of upsert operations. Delegates to BatchPutRevoke.
      */
     std::vector<tl::expected<void, ErrorCode>> BatchUpsertRevoke(
-        const UUID& client_id, const std::vector<std::string>& keys);
+        const UUID& client_id, const std::vector<std::string>& keys,
+        const std::string& tenant_id);
 
     /**
      * @brief Evict a disk replica for a key (triggered by client-side disk
@@ -297,6 +450,7 @@ class MasterService {
      * @return ErrorCode::OK on success, OBJECT_NOT_FOUND if key missing
      */
     auto EvictDiskReplica(const UUID& client_id, const std::string& key,
+                          const std::string& tenant_id,
                           ReplicaType replica_type)
         -> tl::expected<void, ErrorCode>;
 
@@ -309,7 +463,7 @@ class MasterService {
      */
     std::vector<tl::expected<void, ErrorCode>> BatchEvictDiskReplica(
         const UUID& client_id, const std::vector<std::string>& keys,
-        ReplicaType replica_type);
+        const std::string& tenant_id, ReplicaType replica_type);
 
     /**
      * @brief Start a copy operation
@@ -326,14 +480,16 @@ class MasterService {
      */
     tl::expected<CopyStartResponse, ErrorCode> CopyStart(
         const UUID& client_id, const std::string& key,
-        const std::string& src_segment,
+        const std::string& tenant_id, const std::string& src_segment,
         const std::vector<std::string>& tgt_segments);
 
     tl::expected<void, ErrorCode> CopyEnd(const UUID& client_id,
-                                          const std::string& key);
+                                          const std::string& key,
+                                          const std::string& tenant_id);
 
     tl::expected<void, ErrorCode> CopyRevoke(const UUID& client_id,
-                                             const std::string& key);
+                                             const std::string& key,
+                                             const std::string& tenant_id);
 
     /**
      * @brief Start a move operation
@@ -350,13 +506,16 @@ class MasterService {
      */
     tl::expected<MoveStartResponse, ErrorCode> MoveStart(
         const UUID& client_id, const std::string& key,
-        const std::string& src_segment, const std::string& tgt_segment);
+        const std::string& tenant_id, const std::string& src_segment,
+        const std::string& tgt_segment);
 
     tl::expected<void, ErrorCode> MoveEnd(const UUID& client_id,
-                                          const std::string& key);
+                                          const std::string& key,
+                                          const std::string& tenant_id);
 
     tl::expected<void, ErrorCode> MoveRevoke(const UUID& client_id,
-                                             const std::string& key);
+                                             const std::string& key,
+                                             const std::string& tenant_id);
 
     /**
      * @brief Remove an object and its replicas
@@ -365,8 +524,8 @@ class MasterService {
      * @return ErrorCode::OK on success, ErrorCode::OBJECT_NOT_FOUND if not
      * found
      */
-    auto Remove(const std::string& key, bool force = false)
-        -> tl::expected<void, ErrorCode>;
+    auto Remove(const std::string& key, const std::string& tenant_id,
+                bool force = false) -> tl::expected<void, ErrorCode>;
 
     /**
      * @brief Removes objects from the master whose keys match a regex pattern.
@@ -375,23 +534,32 @@ class MasterService {
      * @return An expected object containing the number of removed objects on
      * success, or an ErrorCode on failure.
      */
-    auto RemoveByRegex(const std::string& str, bool force = false)
-        -> tl::expected<long, ErrorCode>;
+    auto RemoveByRegex(const std::string& str, const std::string& tenant_id,
+                       bool force = false) -> tl::expected<long, ErrorCode>;
 
     /**
-     * @brief Remove all objects and their replicas
+     * @brief Remove all objects and their replicas across all tenants.
      * @param force If true, skip lease and replication task checks.
      * @return return the number of objects removed
      */
     long RemoveAll(bool force = false);
 
+    /**
+     * @brief Remove all objects and their replicas for a single tenant.
+     * @param tenant_id The tenant whose objects should be removed.
+     * @param force If true, skip lease and replication task checks.
+     * @return return the number of objects removed
+     */
+    long RemoveAll(const std::string& tenant_id, bool force = false);
+
     /**
      * @brief Batch remove objects and their replicas
      * @param keys The list of keys to remove.
      * @param force If true, skip lease and replication task checks.
      * @return Vector of expected results for each key.
      */
-    auto BatchRemove(const std::vector<std::string>& keys, bool force = false)
+    auto BatchRemove(const std::vector<std::string>& keys,
+                     const std::string& tenant_id, bool force = false)
         -> std::vector<tl::expected<void, ErrorCode>>;
 
     /**
@@ -437,27 +605,90 @@ class MasterService {
      * segment.
      */
     auto OffloadObjectHeartbeat(const UUID& client_id, bool enable_offloading)
-        -> tl::expected<std::unordered_map<std::string, int64_t>, ErrorCode>;
+        -> tl::expected<std::vector<OffloadTaskItem>, ErrorCode>;
+
+    auto ReportSsdCapacity(const UUID& client_id,
+                           int64_t ssd_total_capacity_bytes)
+        -> tl::expected<void, ErrorCode>;
 
     /**
      * @brief Notifies the master that offloading of specified objects has
      * succeeded.
-     * @param keys         A list of object keys (names) that were successfully
-     * offloaded.
+     * @param tasks        A list of tenant-scoped objects that were
+     * successfully offloaded.
      * @param metadatas    The corresponding metadata for each offloaded object,
      * including size, storage location, etc.
      */
     auto NotifyOffloadSuccess(
-        const UUID& client_id, const std::vector<std::string>& keys,
+        const UUID& client_id, const std::vector<OffloadTaskItem>& tasks,
         const std::vector<StorageObjectMetadata>& metadatas)
         -> tl::expected<void, ErrorCode>;
 
+    /**
+     * @brief Heartbeat-driven pull of pending promotion work for a client.
+     * Returns tenant-scoped promotion tasks for the holder client and clears
+     * its per-client promotion_objects queue. The per-shard promotion_tasks
+     * map remains populated as the source of truth until NotifyPromotionSuccess
+     * commits the new MEMORY replica.
+     */
+    auto PromotionObjectHeartbeat(const UUID& client_id)
+        -> tl::expected<std::vector<PromotionTaskItem>, ErrorCode>;
+
+    /**
+     * @brief Stage a PROCESSING MEMORY replica for an existing key. Allocates
+     * DRAM via the existing AllocationStrategy, optionally biased toward the
+     * caller's local memory segment via preferred_segments. The new replica is
+     * invisible to readers until NotifyPromotionSuccess flips it to COMPLETE.
+     *
+     * Only the holder client (the one owning the source LOCAL_DISK replica)
+     * is authorized to call this. Other clients receive INVALID_PARAMS.
+     * `size` must match the source replica's object_size captured at task
+     * admission; mismatch returns INVALID_PARAMS to avoid allocating an
+     * arbitrary buffer size from a buggy or malicious caller.
+     */
+    auto PromotionAllocStart(const UUID& client_id, const std::string& key,
+                             const std::string& tenant_id, uint64_t size,
+                             const std::vector<std::string>& preferred_segments)
+        -> tl::expected<PromotionAllocStartResponse, ErrorCode>;
+
+    /**
+     * @brief Commit a staged MEMORY replica to COMPLETE; decrement source
+     * refcnt; erase per-shard and per-client task entries. Mirror of
+     * NotifyOffloadSuccess.
+     */
+    auto NotifyPromotionSuccess(const UUID& client_id, const std::string& key,
+                                const std::string& tenant_id)
+        -> tl::expected<void, ErrorCode>;
+
+    /**
+     * @brief Holder-side failure notification: the client got past
+     * PromotionAllocStart but a downstream step (local SSD read, RDMA
+     * write, etc.) failed and it will not be calling
+     * NotifyPromotionSuccess. Releases the master-side task state
+     * immediately rather than waiting put_start_release_timeout_sec_
+     * for the reaper to do it. Without this call every transient
+     * client-side error (SSD throttling, RDMA flake, etc.) pins a
+     * task slot and a staged DRAM buffer for the full reaper TTL,
+     * which can saturate promotion_queue_limit_ on busy clusters.
+     *
+     * Authorization is the same as NotifyPromotionSuccess: only the
+     * holder client may release a task. Effects mirror the reaper's
+     * expiry path: drop source LOCAL_DISK refcnt, pop the staged
+     * PROCESSING MEMORY replica if alloc_id was recorded, erase the
+     * task, decrement the global in-flight counter, and clear the
+     * holder's promotion_objects entry.
+     */
+    auto NotifyPromotionFailure(const UUID& client_id, const std::string& key,
+                                const std::string& tenant_id)
+        -> tl::expected<void, ErrorCode>;
+
     /**
      * @brief Create a copy task to copy an object's replicas to target segments
      * @return Copy task ID on success, ErrorCode on failure
      */
     tl::expected<UUID, ErrorCode> CreateCopyTask(
-        const std::string& key, const std::vector<std::string>& targets);
+        const std::string& key, const std::string& tenant_id,
+        const std::vector<std::string>& targets);
 
     /**
      * @brief Create a move task to move an object's replica from source segment
@@ -465,6 +696,7 @@ class MasterService {
      * @return Move task ID on success, ErrorCode on failure
      */
     tl::expected<UUID, ErrorCode> CreateMoveTask(const std::string& key,
+                                                 const std::string& tenant_id,
                                                  const std::string& source,
                                                  const std::string& target);
 
@@ -490,6 +722,12 @@ class MasterService {
     tl::expected<SegmentStatus, ErrorCode> QuerySegmentStatus(
         const std::string& segment_name);
 
+    /**
+     * @brief Query current segment lifecycle state by segment id.
+     */
+    tl::expected<SegmentStatus, ErrorCode> QuerySegmentStatusById(
+        const UUID& segment_id);
+
     /**
      * @brief Query the status of a task
      * @return Task basic info
@@ -512,6 +750,24 @@ class MasterService {
     tl::expected<void, ErrorCode> MarkTaskToComplete(
         const UUID& client_id, const TaskCompleteRequest& request);
 
+    /**
+     * @brief Set the HttpMetadataServer pointer for cleanup on client timeout.
+     * @param server Pointer to HttpMetadataServer. If nullptr, cleanup is
+     * disabled.
+     */
+    void setHttpMetadataServer(HttpMetadataServer* server);
+
+    /**
+     * @brief Configure cleanup against a separately-deployed HTTP metadata
+     * server (not co-located in the master process). The master sends HTTP
+     * DELETE requests to this endpoint when a client times out. Only http://
+     * and https:// connection strings are supported; other schemes (etcd /
+     * redis / P2PHANDSHAKE) are ignored with a warning and leave cleanup
+     * disabled.
+     * @param metadata_connstring e.g. "http://host:8080/metadata".
+     */
+    void setHttpMetadataRemoteUrl(const std::string& metadata_connstring);
+
    private:
     void SnapshotThreadFunc();
 
@@ -561,6 +817,14 @@ class MasterService {
     // evict_ratio_lowerbound, the second pass will be triggered and try to
     // fulfill evict ratio lowerbound.
     void BatchEvict(double evict_ratio_target, double evict_ratio_lowerbound);
+    void NoFBatchEvict(double evict_ratio_target,
+                       double evict_ratio_lowerbound);
+    struct TenantQuotaEvictionResult {
+        uint64_t freed_bytes{0};
+        uint64_t evicted_objects{0};
+    };
+    TenantQuotaEvictionResult EvictTenantMemoryForQuota(
+        const std::string& tenant_id, uint64_t target_bytes);
 
     // Helper to get a snapshot of alive clients (under client_mutex_ shared
     // lock)
@@ -579,6 +843,11 @@ class MasterService {
     void JobDispatchThreadFunc();
 
     // Internal data structures
+    struct ObjectIdentity {
+        std::string tenant_id;
+        std::string user_key;
+    };
+
     struct ObjectMetadata {
         // RAII-style metric management
         ~ObjectMetadata() {
@@ -594,10 +863,17 @@ class MasterService {
             const UUID& client_id_,
             const std::chrono::system_clock::time_point put_start_time_,
             size_t value_length, std::vector<Replica>&& reps,
-            bool enable_soft_pin, bool enable_hard_pin = false)
+            bool enable_soft_pin, bool enable_hard_pin = false,
+            ObjectDataType data_type_ = ObjectDataType::UNKNOWN,
+            std::string group_id_ = "", std::string tenant_id_ = "default",
+            std::string user_key_ = {})
             : client_id(client_id_),
               put_start_time(put_start_time_),
               size(value_length),
+              data_type(data_type_),
+              group_id(std::move(group_id_)),
+              tenant_id(std::move(tenant_id_)),
+              user_key(std::move(user_key_)),
               lease_timeout(),
               soft_pin_timeout(std::nullopt),
               hard_pinned(enable_hard_pin),
@@ -620,6 +896,10 @@ class MasterService {
         // Updated by UpsertStart (Case B) to reset the discard timeout.
         std::chrono::system_clock::time_point put_start_time;
         const size_t size;
+        const ObjectDataType data_type{ObjectDataType::UNKNOWN};
+        const std::string group_id;
+        const std::string tenant_id;
+        const std::string user_key;
 
         mutable SpinLock lock;
         // Default constructor, creates a time_point representing
@@ -630,6 +910,11 @@ class MasterService {
             soft_pin_timeout GUARDED_BY(lock);  // optional soft pin, only
                                                 // set for vip objects
         const bool hard_pinned{false};          // immutable, set at creation
+        bool memory_cache_total_accounted{false};
+        bool disk_cache_total_accounted{false};
+        uint64_t reserved_quota_charge_bytes{0};
+        uint64_t committed_quota_charge_bytes{0};
+        uint64_t pending_replaced_quota_charge_bytes{0};
 
         void AddReplicas(std::vector<Replica>&& replicas) {
             replicas_.insert(replicas_.end(),
@@ -746,6 +1031,32 @@ class MasterService {
             return num_erased > 0;
         }
 
+        size_t EraseReplica(ReplicaType replica_type) {
+            return EraseReplicas([replica_type](const Replica& replica) {
+                if (replica_type == ReplicaType::ALL) {
+                    return replica.is_memory_replica() ||
+                           replica.is_nof_replica();
+                }
+                return replica.type() == replica_type;
+            });
+        }
+
+        bool HasMemReplica() const {
+            return HasReplica(&Replica::fn_is_memory_replica);
+        }
+
+        bool HasNoFReplica() const {
+            return HasReplica(&Replica::fn_is_nof_replica);
+        }
+
+        size_t GetMemReplicaCount() const {
+            return CountReplicas(&Replica::fn_is_memory_replica);
+        }
+
+        size_t GetNoFReplicaCount() const {
+            return CountReplicas(&Replica::fn_is_nof_replica);
+        }
+
         Replica* GetReplicaBySegmentName(const std::string& segment_name) {
             return GetFirstReplica([&segment_name](const Replica& replica) {
                 auto names = replica.get_segment_names();
@@ -773,6 +1084,18 @@ class MasterService {
             }
         }
 
+        bool NeedsLeaseRefresh(const uint64_t ttl,
+                               const uint64_t soft_ttl) const {
+            SpinLocker locker(&lock);
+            const auto now = std::chrono::system_clock::now();
+            if (lease_timeout <= now + std::chrono::milliseconds(ttl / 2)) {
+                return true;
+            }
+            return soft_pin_timeout &&
+                   *soft_pin_timeout <=
+                       now + std::chrono::milliseconds(soft_ttl / 2);
+        }
+
         // Check if the lease has expired
         bool IsLeaseExpired() const {
             SpinLocker locker(&lock);
@@ -800,6 +1123,8 @@ class MasterService {
 
         bool IsHardPinned() const { return hard_pinned; }
 
+        bool IsGrouped() const { return !group_id.empty(); }
+
         // Check if the metadata is valid
         // Valid means it has at least one valid replica and size is greater
         // than 0
@@ -837,6 +1162,7 @@ class MasterService {
         } type;
         ReplicaID source_id;
         std::vector<ReplicaID> replica_ids;
+        uint64_t reserved_quota_charge_bytes{0};
     };
 
     struct OffloadingTask {
@@ -844,21 +1170,107 @@ class MasterService {
         std::chrono::system_clock::time_point start_time;
     };
 
+    // Tracks an in-flight LOCAL_DISK -> MEMORY copy. The source
+    // LOCAL_DISK replica is refcnt-pinned for the duration of the task
+    // so it cannot be evicted.
+    //
+    // alloc_id pins down which staged PROCESSING MEMORY replica
+    // NotifyPromotionSuccess should commit, so a concurrent Put on the
+    // same key cannot be confused with ours. 0 until
+    // PromotionAllocStart records the new replica.
+    //
+    // start_time is the reaper deadline anchor. Set at task admission
+    // and reset at PromotionAllocStart so each phase (queue-wait and
+    // active-transfer) gets its own full put_start_release_timeout_sec_
+    // window. Without the reset a backlogged task could enter active
+    // transfer with little TTL left, and the reaper could free the
+    // staged replica via EraseReplicaByID mid-RDMA-write.
+    //
+    // holder_id is the client owning the source LOCAL_DISK segment and
+    // the only one authorized to commit (NotifyPromotionSuccess) or
+    // abort (NotifyPromotionFailure) the task. Without it, any client
+    // knowing the key could flip the staged PROCESSING replica to
+    // COMPLETE before the holder's RDMA write landed, exposing torn
+    // data to readers.
+    struct PromotionTask {
+        ReplicaID source_id;    // the LOCAL_DISK replica being promoted
+        ReplicaID alloc_id{0};  // the new MEMORY replica staged by AllocStart
+        uint64_t object_size;
+        uint64_t reserved_quota_charge_bytes{0};
+        std::chrono::system_clock::time_point start_time;
+        UUID holder_id;  // owner of source LOCAL_DISK; only Notifier allowed
+    };
+
     static constexpr size_t kNumShards = 1024;  // Number of metadata shards
 
+    struct TenantState {
+        std::unordered_map<std::string, ObjectMetadata> metadata;
+        std::unordered_set<std::string> processing_keys;
+        std::unordered_map<std::string, const ReplicationTask>
+            replication_tasks;
+        std::unordered_map<std::string, const OffloadingTask> offloading_tasks;
+        std::unordered_map<std::string, PromotionTask> promotion_tasks;
+
+        std::unordered_map<std::string, std::unordered_set<std::string>>
+            group_members;  // group_id → set of keys
+
+        bool Empty() const {
+            return metadata.empty() && processing_keys.empty() &&
+                   replication_tasks.empty() && offloading_tasks.empty() &&
+                   promotion_tasks.empty() && group_members.empty();
+        }
+    };
+
     // Sharded metadata maps and their mutexes
     struct MetadataShard {
         mutable SharedMutex mutex;
-        std::unordered_map<std::string, ObjectMetadata> metadata
-            GUARDED_BY(mutex);
-        std::unordered_set<std::string> processing_keys GUARDED_BY(mutex);
-        std::unordered_map<std::string, const ReplicationTask> replication_tasks
-            GUARDED_BY(mutex);
-        std::unordered_map<std::string, const OffloadingTask> offloading_tasks
-            GUARDED_BY(mutex);
+        std::unordered_map<std::string, TenantState> tenants GUARDED_BY(mutex);
+        // Count of objects that have at least one completed LOCAL_DISK replica.
+        // Used to compute eviction_base = metadata.size() - disk_object_count,
+        // excluding disk-only objects from the eviction denominator.
+        long disk_object_count GUARDED_BY(mutex) = 0;
     };
     std::array<MetadataShard, kNumShards> metadata_shards_;
 
+    static bool HasCompletedMemoryCacheReplica(const ObjectMetadata& metadata);
+    static bool HasCompletedDiskCacheReplica(const ObjectMetadata& metadata);
+    static void SyncCacheTotalAccounting(ObjectMetadata& metadata);
+    void RebuildCacheTotalAccounting();
+    static void AccountCacheTotalRemoval(ObjectMetadata& metadata);
+    std::vector<Replica> PopReplicasWithCacheTotalAccounting(
+        ObjectMetadata& metadata,
+        const std::function<bool(const Replica&)>& pred_fn);
+    std::vector<Replica> PopReplicasWithCacheTotalAccounting(
+        ObjectMetadata& metadata);
+    size_t EraseReplicasWithCacheTotalAccounting(
+        ObjectMetadata& metadata,
+        const std::function<bool(const Replica&)>& pred_fn);
+
+    static constexpr size_t kNumTenantQuotaShards = 1024;
+    struct TenantQuotaShard {
+        mutable std::mutex mutex;
+        std::unordered_map<std::string, TenantQuotaState> tenants
+            GUARDED_BY(mutex);
+    };
+    std::array<TenantQuotaShard, kNumTenantQuotaShards> tenant_quota_shards_;
+
+    std::unordered_map<std::string, std::string> object_group_ids_
+        GUARDED_BY(group_routing_mutex_);
+    mutable std::unordered_set<std::string> groups_needing_lease_refresh_
+        GUARDED_BY(group_routing_mutex_);
+    mutable std::shared_mutex group_routing_mutex_;
+
+    static constexpr size_t kObjectOperationLockStripes = 4096;
+
+    struct ObjectOperationLock {
+        std::unique_lock<std::mutex> lock;
+    };
+
+    ObjectOperationLock AcquireObjectOperationLock(const std::string& tenant_id,
+                                                   const std::string& key);
+
+    std::array<std::mutex, kObjectOperationLockStripes> object_operation_locks_;
+
     // For accessing a metadata shard with read-write permission
     class MetadataShardAccessorRW {
        public:
@@ -871,6 +1283,38 @@ class MasterService {
 
         const MetadataShard* operator->() const { return &shard_; }
 
+        MetadataShard& get() { return shard_; }
+
+        const MetadataShard& get() const { return shard_; }
+
+        // Called after adding a LOCAL_DISK replica. Increments
+        // disk_object_count if this is the first completed LOCAL_DISK
+        // replica for the object (i.e., exactly 1 completed disk replica now).
+        void OnDiskReplicaAdded(const ObjectMetadata& metadata) {
+            size_t disk_count = metadata.CountReplicas([](const Replica& r) {
+                return r.is_local_disk_replica() && r.is_completed();
+            });
+            if (disk_count == 1) shard_.disk_object_count++;
+        }
+
+        // Called after removing a LOCAL_DISK replica, or when erasing an
+        // object that had one. Pass had_completed_disk=true if the object
+        // had at least one completed LOCAL_DISK replica before the removal.
+        // When the entire object is being erased, call the one-arg overload.
+        void OnDiskReplicaRemoved(bool had_completed_disk,
+                                  const ObjectMetadata& metadata) {
+            if (!had_completed_disk) return;
+            bool still_has_disk = metadata.HasReplica([](const Replica& r) {
+                return r.is_local_disk_replica() && r.is_completed();
+            });
+            if (!still_has_disk) shard_.disk_object_count--;
+        }
+
+        // Overload for full object erasure — no metadata needed.
+        void OnDiskReplicaRemoved(bool had_completed_disk) {
+            if (had_completed_disk) shard_.disk_object_count--;
+        }
+
        private:
         MetadataShard& shard_;
         SharedMutexLocker lock_;
@@ -886,28 +1330,133 @@ class MasterService {
 
         const MetadataShard* operator->() const { return &shard_; }
 
+        const MetadataShard& get() const { return shard_; }
+
        private:
         const MetadataShard& shard_;
         SharedMutexLocker lock_;
     };
 
-    // Helper to get shard index from key
+    static ObjectIdentity MakeObjectIdentity(const std::string& user_key,
+                                             const std::string& tenant_id) {
+        return {NormalizeTenantId(tenant_id), user_key};
+    }
+    std::string NormalizeRequestTenantId(const std::string& tenant_id) const;
+    ObjectIdentity MakeObjectIdentityForRequest(
+        const std::string& user_key, const std::string& tenant_id) const;
+    tl::expected<std::string, ErrorCode> NormalizeTenantIdForWrite(
+        const std::string& tenant_id) const;
+    tl::expected<std::string, ErrorCode> NormalizeTenantIdForWriteLocked(
+        const std::string& tenant_id) const;
+    bool IsTenantRegistered(const std::string& tenant_id) const;
+    bool TenantHasObjects(const std::string& tenant_id) const;
+
+    static std::string MakeTenantScopedKey(const std::string& tenant_id,
+                                           const std::string& key) {
+        const auto normalized_tenant = NormalizeTenantId(tenant_id);
+        std::string scoped_key;
+        scoped_key.reserve(normalized_tenant.size() + key.size() + 1);
+        scoped_key.append(normalized_tenant);
+        scoped_key.push_back('\0');
+        scoped_key.append(key);
+        return scoped_key;
+    }
+
+    // Helper to get shard index from tenant-scoped object identity.
+    size_t getShardIndex(const std::string& tenant_id,
+                         const std::string& user_key) const {
+        const auto normalized_tenant = NormalizeTenantId(tenant_id);
+        if (normalized_tenant == "default") {
+            return std::hash<std::string>{}(user_key) % kNumShards;
+        }
+        size_t seed = std::hash<std::string>{}(normalized_tenant);
+        boost::hash_combine(seed, user_key);
+        return seed % kNumShards;
+    }
+
+    // Legacy helper routes plain keys to the default tenant.
     size_t getShardIndex(const std::string& key) const {
         return std::hash<std::string>{}(key) % kNumShards;
     }
 
+    size_t getMetadataShardIndex(const std::string& tenant_id,
+                                 const std::string& key) const;
+    size_t getTenantQuotaShardIndex(const std::string& tenant_id) const;
+    std::optional<std::string> GetGroupRoute(const std::string& tenant_id,
+                                             const std::string& key) const;
+    void RegisterGroupMember(TenantState& tenant_state,
+                             const std::string& tenant_id,
+                             const std::string& key,
+                             const std::string& group_id);
+    void UnregisterGroupMember(TenantState& tenant_state,
+                               const std::string& tenant_id,
+                               const std::string& key,
+                               const std::string& group_id);
+    std::unordered_map<std::string, ObjectMetadata>::iterator EraseMetadata(
+        TenantState& tenant_state,
+        std::unordered_map<std::string, ObjectMetadata>::iterator it,
+        const std::string& tenant_id);
+    void ReleaseLocalDiskUsage(const std::vector<Replica>& replicas);
+    enum class QuotaEraseMode {
+        kFull,
+        kPreserveOld,
+        kAbortOnly,
+    };
+    std::unordered_map<std::string, ObjectMetadata>::iterator EraseMetadata(
+        TenantState& tenant_state,
+        std::unordered_map<std::string, ObjectMetadata>::iterator it,
+        const std::string& tenant_id, QuotaEraseMode quota_mode);
+    uint64_t CompletedMemoryQuotaCharge(const ObjectMetadata& metadata) const;
+    uint64_t RequestedMemoryQuotaCharge(uint64_t value_length,
+                                        const ReplicateConfig& config) const;
+    bool ShouldProtectZeroChargeMetadataCreate(
+        uint64_t requested_quota_charge) const;
+    uint64_t ComputeTenantQuotaDeficit(const std::string& tenant_id,
+                                       uint64_t incoming_quota_charge);
+    tl::expected<void, ErrorCode> ReserveTenantQuota(
+        const std::string& tenant_id, uint64_t bytes);
+    void CommitTenantQuota(const std::string& tenant_id, uint64_t bytes);
+    void AbortTenantQuota(const std::string& tenant_id, uint64_t bytes);
+    void ReleaseTenantQuota(const std::string& tenant_id, uint64_t bytes);
+    void ReleaseTenantQuotaPartial(const std::string& tenant_id,
+                                   uint64_t bytes);
+    void CommitAdditionalTenantQuota(const std::string& tenant_id,
+                                     uint64_t bytes);
+    void AbortReplicationTaskQuota(const std::string& tenant_id,
+                                   const ReplicationTask& task);
+    void IncrementTenantMetadataObjectCount(const std::string& tenant_id);
+    void DecrementTenantMetadataObjectCount(const std::string& tenant_id);
+    void ReleaseCommittedQuotaCharge(ObjectMetadata& metadata, uint64_t bytes);
+    void RecomputeTenantEffectiveQuotas();
+    void RebuildTenantQuotaUsageFromMetadata();
+    void LoadTenantQuotaPoliciesFromStoreOrThrow();
+    void ApplyTenantQuotaPolicies(const TenantQuotaPolicySnapshot& snapshot);
+    TenantQuotaPolicySnapshot BuildTenantQuotaPolicySnapshot() const;
+    uint64_t GetTenantQuotaCapacityBytes();
+    std::unordered_map<std::string, ObjectMetadata>::iterator EraseMetadata(
+        TenantState& tenant_state,
+        std::unordered_map<std::string, ObjectMetadata>::iterator it,
+        const std::string& tenant_id, QuotaEraseMode quota_mode,
+        MetadataShardAccessorRW* shard);
+    void RebuildGroupRoutingIndex();
+    void GrantLeaseForGroup(const TenantState& tenant_state,
+                            const std::string& key,
+                            const ObjectMetadata& metadata) const;
+
     // Helper to clean up stale handles pointing to unmounted segments
     // or local_disk replicas whose owner client is no longer alive.
     bool CleanupStaleHandles(
         ObjectMetadata& metadata,
-        const std::unordered_set<UUID, boost::hash<UUID>>& alive_clients);
+        const std::unordered_set<UUID, boost::hash<UUID>>& alive_clients,
+        MetadataShardAccessorRW* shard = nullptr);
 
     // Helper: allocate replicas, create ObjectMetadata, insert into shard,
     // and return descriptor list.  Shared by PutStart and UpsertStart.
     auto AllocateAndInsertMetadata(
         MetadataShardAccessorRW& shard, const UUID& client_id,
         const std::string& key, uint64_t value_length,
-        const ReplicateConfig& config,
+        const ReplicateConfig& config, const std::string& group_id,
+        const std::string& tenant_id,
         const std::chrono::system_clock::time_point& now)
         -> tl::expected<std::vector<Replica::Descriptor>, ErrorCode>;
 
@@ -927,20 +1476,78 @@ class MasterService {
 
     // Eviction thread function
     void EvictionThreadFunc();
+    void NofHeartbeatThreadFunc();
+    bool TryUnmountNoFSegmentByHeartbeat(
+        const MountedNoFSegmentSnapshot& snapshot,
+        const std::string& error_reason);
+    bool ProbeNoFSegment(const std::string& te_endpoint,
+                         std::string* error_reason);
 
-    tl::expected<void, ErrorCode> PushOffloadingQueue(const std::string& key,
-                                                      Replica& replica);
+    tl::expected<void, ErrorCode> PushOffloadingQueue(
+        const ObjectIdentity& object_id, Replica& replica);
+
+    struct GracefulUnmountDeadlineRecord {
+        UUID segment_id;
+        UUID client_id;
+    };
+
+    DeadlineScheduler<GracefulUnmountDeadlineRecord>
+        graceful_unmount_scheduler_;
+
+    /**
+     * @brief Mirror of PushOffloadingQueue for promotion-on-hit. Inserts an
+     * entry into the holder client's LocalDiskSegment::promotion_objects map.
+     * Caller is responsible for refcnt-pinning the source replica and
+     * recording the task in the shard's promotion_tasks map.
+     */
+    tl::expected<void, ErrorCode> PushPromotionQueue(
+        const ObjectIdentity& object_id, Replica& source_replica);
+
+    /**
+     * @brief Helper invoked from GetReplicaList when an only-LOCAL_DISK key is
+     * observed. Applies the gating chain (frequency / watermark / dedup /
+     * cap), refcnt-pins the source LOCAL_DISK replica, records a
+     * PromotionTask, and pushes onto the holder client's promotion_objects
+     * map. Acquires its own RW shard accessor; safe to call after
+     * GetReplicaList's RO accessor has been released.
+     */
+    void TryPushPromotionQueue(const ObjectIdentity& object_id);
+
+    // Erase any in-flight PromotionTask for `key`, abort any staged promotion
+    // quota reservation, and decrement the cluster-wide in-flight counter. Safe
+    // no-op if no task exists.
+    void ErasePromotionTaskIfPresent(
+        TenantState& tenant_state, const std::string& key,
+        const std::string& tenant_id) NO_THREAD_SAFETY_ANALYSIS {
+        auto task_it = tenant_state.promotion_tasks.find(key);
+        if (task_it != tenant_state.promotion_tasks.end()) {
+            AbortTenantQuota(tenant_id,
+                             task_it->second.reserved_quota_charge_bytes);
+            tenant_state.promotion_tasks.erase(task_it);
+            promotion_in_flight_.fetch_sub(1, std::memory_order_relaxed);
+            MasterMetricManager::instance().dec_promotion_in_flight();
+            MasterMetricManager::instance().inc_promotion_cancelled();
+        }
+    }
 
     // Lease related members
     const uint64_t default_kv_lease_ttl_;     // in milliseconds
     const uint64_t default_kv_soft_pin_ttl_;  // in milliseconds
     const bool allow_evict_soft_pinned_objects_;
+    std::array<double, UINT8_MAX + 1> object_type_reuse_scales_;
+    std::array<double, UINT8_MAX + 1> object_type_soft_pin_weights_;
+    std::array<int64_t, UINT8_MAX + 1> object_type_eviction_graces_;
+    std::array<double, UINT8_MAX + 1> object_type_budget_ratios_;
 
     // Eviction related members
-    std::atomic<bool> need_eviction_{
-        false};  // Set to trigger eviction when not enough space left
-    const double eviction_ratio_;                 // in range [0.0, 1.0]
-    const double eviction_high_watermark_ratio_;  // in range [0.0, 1.0]
+    std::atomic<bool> need_mem_eviction_{
+        false};  // Set to trigger memory eviction when allocation fails
+    std::atomic<bool> need_nof_eviction_{
+        false};  // Set to trigger NoF eviction when allocation fails
+    const double eviction_ratio_;                     // in range [0.0, 1.0]
+    const double eviction_high_watermark_ratio_;      // in range [0.0, 1.0]
+    const double nof_eviction_ratio_;                 // in range [0.0, 1.0]
+    const double nof_eviction_high_watermark_ratio_;  // in range [0.0, 1.0]
 
     // Eviction thread related members
     std::thread eviction_thread_;
@@ -950,6 +1557,8 @@ class MasterService {
 
     std::thread snapshot_thread_;
     std::atomic<bool> snapshot_running_{false};
+    std::mutex snapshot_thread_mutex_;
+    std::condition_variable snapshot_thread_cv_;
     // Task cleanup thread related members
     std::thread task_cleanup_thread_;
     std::atomic<bool> task_cleanup_running_{false};
@@ -963,54 +1572,95 @@ class MasterService {
     // Helper class for accessing metadata with automatic locking and cleanup
     class MetadataAccessorRW {
        public:
-        MetadataAccessorRW(MasterService* service, const std::string& key)
+        MetadataAccessorRW(MasterService* service,
+                           const ObjectIdentity& object_id)
             : service_(service),
-              key_(key),
-              shard_idx_(service_->getShardIndex(key)),
+              object_id_(object_id),
+              shard_idx_(service_->getMetadataShardIndex(object_id_.tenant_id,
+                                                         object_id_.user_key)),
               shard_guard_(service_, shard_idx_),
-              it_(shard_guard_->metadata.find(key)),
-              processing_it_(shard_guard_->processing_keys.find(key)),
-              replication_task_it_(shard_guard_->replication_tasks.find(key)) {
+              tenant_it_(shard_guard_->tenants.find(object_id_.tenant_id)),
+              tenant_state_(tenant_it_ == shard_guard_->tenants.end()
+                                ? nullptr
+                                : &tenant_it_->second),
+              it_(tenant_state_ == nullptr
+                      ? ObjectMetadataIterator{}
+                      : tenant_state_->metadata.find(object_id_.user_key)),
+              processing_it_(tenant_state_ == nullptr
+                                 ? ProcessingIterator{}
+                                 : tenant_state_->processing_keys.find(
+                                       object_id_.user_key)),
+              replication_task_it_(tenant_state_ == nullptr
+                                       ? ReplicationTaskIterator{}
+                                       : tenant_state_->replication_tasks.find(
+                                             object_id_.user_key)) {
             // Automatically clean up invalid handles (memory replicas only).
             // Note: We only check memory replicas here to avoid lock order
             // violation (client_mutex_ must be acquired before metadata shard).
             // local_disk replicas are cleaned up by ClearInvalidHandles() in
             // ClientMonitorFunc.
-            if (it_ != shard_guard_->metadata.end()) {
+            if (tenant_state_ != nullptr &&
+                it_ != tenant_state_->metadata.end()) {
                 // Erase invalid memory replicas (those with unmounted
                 // segments). No client_mutex_ needed since we only check memory
                 // replicas.
-                it_->second.EraseReplicas([](const Replica& replica) {
-                    return replica.has_invalid_mem_handle();
-                });
+                const uint64_t before_charge =
+                    service_->CompletedMemoryQuotaCharge(it_->second);
+                service_->EraseReplicasWithCacheTotalAccounting(
+                    it_->second, [](const Replica& replica) {
+                        return replica.has_invalid_mem_handle();
+                    });
+                const uint64_t after_charge =
+                    service_->CompletedMemoryQuotaCharge(it_->second);
+                if (before_charge > after_charge) {
+                    service_->ReleaseCommittedQuotaCharge(
+                        it_->second, before_charge - after_charge);
+                }
                 // If no valid replicas remain, delete the whole object.
                 if (!it_->second.IsValid()) {
+                    const bool had_processing =
+                        processing_it_ != tenant_state_->processing_keys.end();
                     this->Erase();
-                    if (processing_it_ != shard_guard_->processing_keys.end()) {
+                    if (tenant_state_ != nullptr && had_processing) {
                         this->EraseFromProcessing();
                     }
+                    if (tenant_state_ != nullptr) {
+                        service_->ErasePromotionTaskIfPresent(
+                            *tenant_state_, object_id_.user_key,
+                            object_id_.tenant_id);
+                        MaybeEraseEmptyTenant();
+                    }
                 }
             }
         }
 
         // Check if metadata exists
         bool Exists() const NO_THREAD_SAFETY_ANALYSIS {
-            return it_ != shard_guard_->metadata.end() && it_->second.IsValid();
+            return tenant_state_ != nullptr &&
+                   it_ != tenant_state_->metadata.end() &&
+                   it_->second.IsValid();
         }
 
         bool InProcessing() const NO_THREAD_SAFETY_ANALYSIS {
-            return processing_it_ != shard_guard_->processing_keys.end();
+            return tenant_state_ != nullptr &&
+                   processing_it_ != tenant_state_->processing_keys.end();
         }
 
         bool HasReplicationTask() const NO_THREAD_SAFETY_ANALYSIS {
-            return replication_task_it_ !=
-                   shard_guard_->replication_tasks.end();
+            return tenant_state_ != nullptr &&
+                   replication_task_it_ !=
+                       tenant_state_->replication_tasks.end();
         }
 
         MetadataShardAccessorRW& GetShard() NO_THREAD_SAFETY_ANALYSIS {
             return shard_guard_;
         }
 
+        TenantState& GetTenantState() NO_THREAD_SAFETY_ANALYSIS {
+            EnsureTenantState();
+            return *tenant_state_;
+        }
+
         // Get metadata (only call when Exists() is true)
         ObjectMetadata& Get() NO_THREAD_SAFETY_ANALYSIS { return it_->second; }
 
@@ -1020,44 +1670,85 @@ class MasterService {
 
         // Delete current metadata (for PutRevoke or Remove operations)
         void Erase() NO_THREAD_SAFETY_ANALYSIS {
-            shard_guard_->metadata.erase(it_);
-            it_ = shard_guard_->metadata.end();
+            service_->EraseMetadata(*tenant_state_, it_, object_id_.tenant_id,
+                                    QuotaEraseMode::kFull, &shard_guard_);
+            it_ = tenant_state_->metadata.end();
+            MaybeEraseEmptyTenant();
         }
 
         void EraseFromProcessing() NO_THREAD_SAFETY_ANALYSIS {
-            shard_guard_->processing_keys.erase(processing_it_);
-            processing_it_ = shard_guard_->processing_keys.end();
+            tenant_state_->processing_keys.erase(processing_it_);
+            processing_it_ = tenant_state_->processing_keys.end();
+            MaybeEraseEmptyTenant();
         }
 
         void EraseReplicationTask() NO_THREAD_SAFETY_ANALYSIS {
-            shard_guard_->replication_tasks.erase(replication_task_it_);
-            replication_task_it_ = shard_guard_->replication_tasks.end();
+            tenant_state_->replication_tasks.erase(replication_task_it_);
+            replication_task_it_ = tenant_state_->replication_tasks.end();
+            MaybeEraseEmptyTenant();
         }
 
         void Create(const UUID& client_id, uint64_t total_length,
                     std::vector<Replica> replicas, bool enable_soft_pin,
-                    bool enable_hard_pin = false) {
+                    bool enable_hard_pin = false,
+                    ObjectDataType data_type = ObjectDataType::UNKNOWN,
+                    std::string group_id = "") {
             if (Exists()) {
                 throw std::logic_error("Already exists");
             }
             const auto now = std::chrono::system_clock::now();
-            auto result = shard_guard_->metadata.emplace(
-                std::piecewise_construct, std::forward_as_tuple(key_),
-                std::forward_as_tuple(client_id, now, total_length,
-                                      std::move(replicas), enable_soft_pin,
-                                      enable_hard_pin));
+            EnsureTenantState();
+            auto result = tenant_state_->metadata.emplace(
+                std::piecewise_construct,
+                std::forward_as_tuple(object_id_.user_key),
+                std::forward_as_tuple(
+                    client_id, now, total_length, std::move(replicas),
+                    enable_soft_pin, enable_hard_pin, data_type, group_id,
+                    object_id_.tenant_id, object_id_.user_key));
             it_ = result.first;
+            if (result.second) {
+                service_->IncrementTenantMetadataObjectCount(
+                    object_id_.tenant_id);
+            }
         }
 
        private:
+        using ObjectMetadataIterator =
+            std::unordered_map<std::string, ObjectMetadata>::iterator;
+        using ProcessingIterator = std::unordered_set<std::string>::iterator;
+        using ReplicationTaskIterator =
+            std::unordered_map<std::string, const ReplicationTask>::iterator;
+
+        void EnsureTenantState() NO_THREAD_SAFETY_ANALYSIS {
+            if (tenant_state_ != nullptr) {
+                return;
+            }
+            auto result =
+                shard_guard_->tenants.try_emplace(object_id_.tenant_id);
+            tenant_it_ = result.first;
+            tenant_state_ = &tenant_it_->second;
+            it_ = tenant_state_->metadata.end();
+            processing_it_ = tenant_state_->processing_keys.end();
+            replication_task_it_ = tenant_state_->replication_tasks.end();
+        }
+
+        void MaybeEraseEmptyTenant() NO_THREAD_SAFETY_ANALYSIS {
+            if (tenant_state_ == nullptr || !tenant_state_->Empty()) {
+                return;
+            }
+            shard_guard_->tenants.erase(tenant_it_);
+            tenant_state_ = nullptr;
+        }
+
         MasterService* service_;
-        std::string key_;
+        ObjectIdentity object_id_;
         size_t shard_idx_;
         MetadataShardAccessorRW shard_guard_;
-        std::unordered_map<std::string, ObjectMetadata>::iterator it_;
-        std::unordered_set<std::string>::iterator processing_it_;
-        std::unordered_map<std::string, const ReplicationTask>::iterator
-            replication_task_it_;
+        std::unordered_map<std::string, TenantState>::iterator tenant_it_;
+        TenantState* tenant_state_;
+        ObjectMetadataIterator it_;
+        ProcessingIterator processing_it_;
+        ReplicationTaskIterator replication_task_it_;
     };
 
     class MetadataSerializer {
@@ -1104,21 +1795,35 @@ class MasterService {
     friend class MetadataAccessor;
     class MetadataAccessorRO {
        public:
-        MetadataAccessorRO(const MasterService* service, const std::string& key)
+        MetadataAccessorRO(const MasterService* service,
+                           const ObjectIdentity& object_id)
             : service_(service),
-              key_(key),
-              shard_idx_(service_->getShardIndex(key)),
+              object_id_(object_id),
+              shard_idx_(service_->getMetadataShardIndex(object_id_.tenant_id,
+                                                         object_id_.user_key)),
               shard_guard_(service_, shard_idx_),
-              it_(shard_guard_->metadata.find(key)),
-              processing_it_(shard_guard_->processing_keys.find(key)) {}
+              tenant_it_(shard_guard_->tenants.find(object_id_.tenant_id)),
+              tenant_state_(tenant_it_ == shard_guard_->tenants.end()
+                                ? nullptr
+                                : &tenant_it_->second),
+              it_(tenant_state_ == nullptr
+                      ? ObjectMetadataConstIterator{}
+                      : tenant_state_->metadata.find(object_id_.user_key)),
+              processing_it_(tenant_state_ == nullptr
+                                 ? ProcessingConstIterator{}
+                                 : tenant_state_->processing_keys.find(
+                                       object_id_.user_key)) {}
 
         // Check if metadata exists
         bool Exists() const NO_THREAD_SAFETY_ANALYSIS {
-            return it_ != shard_guard_->metadata.end() && it_->second.IsValid();
+            return tenant_state_ != nullptr &&
+                   it_ != tenant_state_->metadata.end() &&
+                   it_->second.IsValid();
         }
 
         bool InProcessing() const NO_THREAD_SAFETY_ANALYSIS {
-            return processing_it_ != shard_guard_->processing_keys.end();
+            return tenant_state_ != nullptr &&
+                   processing_it_ != tenant_state_->processing_keys.end();
         }
 
         // Get metadata (only call when Exists() is true)
@@ -1130,13 +1835,24 @@ class MasterService {
             return shard_guard_;
         }
 
+        const TenantState* GetTenantState() const NO_THREAD_SAFETY_ANALYSIS {
+            return tenant_state_;
+        }
+
        private:
+        using ObjectMetadataConstIterator =
+            std::unordered_map<std::string, ObjectMetadata>::const_iterator;
+        using ProcessingConstIterator =
+            std::unordered_set<std::string>::const_iterator;
+
         const MasterService* service_;
-        const std::string key_;
+        const ObjectIdentity object_id_;
         const size_t shard_idx_;
         MetadataShardAccessorRO shard_guard_;
-        std::unordered_map<std::string, ObjectMetadata>::const_iterator it_;
-        std::unordered_set<std::string>::const_iterator processing_it_;
+        std::unordered_map<std::string, TenantState>::const_iterator tenant_it_;
+        const TenantState* tenant_state_;
+        ObjectMetadataConstIterator it_;
+        ProcessingConstIterator processing_it_;
     };
 
     friend class MetadataAccessorRW;
@@ -1162,6 +1878,27 @@ class MasterService {
         128 * 1024;  // Size of the client ping queue
     boost::lockfree::queue<PodUUID> client_ping_queue_{kClientPingQueueSize};
     const int64_t client_live_ttl_sec_;
+    const std::chrono::seconds nof_heartbeat_interval_sec_;
+    const std::chrono::milliseconds nof_heartbeat_probe_timeout_ms_;
+    const uint32_t nof_heartbeat_failures_threshold_;
+
+    struct NoFHeartbeatState {
+        UUID owner_client_id{0, 0};
+        std::string segment_name;
+        std::string te_endpoint;
+        std::chrono::steady_clock::time_point next_probe_at{};
+        std::chrono::steady_clock::time_point last_success_at{};
+        uint32_t consecutive_failures{0};
+        std::string last_error_reason;
+    };
+    std::mutex nof_heartbeat_mutex_;
+    std::unordered_map<UUID, NoFHeartbeatState, boost::hash<UUID>>
+        nof_heartbeat_states_;
+    std::thread nof_heartbeat_thread_;
+    std::atomic<bool> nof_heartbeat_running_{false};
+    static constexpr uint64_t kNoFHeartbeatThreadSleepMs = 100;
+    mutable std::mutex nof_probe_fn_mutex_;
+    NoFProbeFn nof_probe_fn_;
 
     // if high availability features enabled
     const bool enable_ha_;
@@ -1176,6 +1913,25 @@ class MasterService {
     // offload_on_evict_=true)
     bool offload_force_evict_{false};
 
+    // Promotion-on-hit: opt-in flag enabling LOCAL_DISK -> MEMORY promotion
+    // when a Get observes a key with only LOCAL_DISK replicas.
+    bool promotion_on_hit_{false};
+    uint32_t promotion_admission_threshold_{2};
+    uint32_t promotion_queue_limit_{50000};
+    uint32_t promotion_max_per_heartbeat_{1};
+    // Global in-flight task counter, checked against promotion_queue_limit_
+    // as the gate cap. Promotion specifically targets skewed
+    // access (hot keys re-accessed after eviction), so the global counter
+    // is the correct primitive. Incremented in TryPushPromotionQueue after
+    // successful enqueue; decremented in NotifyPromotionSuccess and in the
+    // promotion task reaper after the task entry is erased. Relaxed memory
+    // order is safe — the value is an advisory soft cap, not a barrier.
+    std::atomic<uint64_t> promotion_in_flight_{0};
+    // Master-side frequency sketch. Constructed only when promotion_on_hit_ is
+    // true. CountMinSketch is mutex-protected internally so we can call into it
+    // from any GetReplicaList caller without additional locking.
+    std::unique_ptr<CountMinSketch> promotion_sketch_;
+
     const std::string ha_backend_type_;
 
     const std::string ha_backend_connstring_;
@@ -1189,12 +1945,46 @@ class MasterService {
     // storage backend eviction configuration
     const bool enable_disk_eviction_;
     const uint64_t quota_bytes_;
+    const bool enable_multi_tenants_;
+    const std::string tenant_quota_connector_type_;
+    const std::string tenant_quota_connector_uri_;
+    std::unique_ptr<TenantQuotaPolicyStore> tenant_quota_policy_store_;
+    mutable std::mutex tenant_quota_policy_mutex_;
+    mutable std::mutex tenant_quota_recompute_mutex_;
+
+    // HTTP metadata server pointer for cleanup on client timeout
+    // nullptr means cleanup is disabled
+    HttpMetadataServer* http_metadata_server_{nullptr};
+
+    // Remote HTTP metadata client, used when the metadata server is deployed
+    // separately. nullptr = no remote cleanup (co-located prefers the pointer).
+    std::shared_ptr<MetadataStoragePlugin> http_metadata_remote_;
+
+    // Cached HTTP metadata key prefix (initialized once at startup)
+    std::string http_metadata_prefix_;
+
+    // Async worker for remote cleanup: segments are enqueued from the client
+    // monitor thread so a slow/unreachable server never blocks heartbeats.
+    std::thread http_metadata_cleanup_thread_;
+    std::atomic<bool> http_metadata_cleanup_running_{false};
+    std::mutex http_metadata_cleanup_mutex_;
+    std::condition_variable http_metadata_cleanup_cv_;
+    std::vector<std::string> http_metadata_cleanup_queue_;
+
+    void HttpMetadataCleanupThreadFunc();
+
+    // Clean up HTTP metadata (mooncake/ram/*, mooncake/rpc_meta/*) for a
+    // segment. For the co-located case this is synchronous (no network I/O);
+    // for the remote case it enqueues to the async cleanup worker.
+    void cleanupHttpMetadata(const std::string& segment_name);
 
     bool use_disk_replica_{false};
 
     // Segment management
     SegmentManager segment_manager_;
+    NoFSegmentManager nof_segment_manager_;
     BufferAllocatorType memory_allocator_type_;
+    const AllocationStrategyType allocation_strategy_type_;
     std::shared_ptr<AllocationStrategy> allocation_strategy_;
 
     bool enable_snapshot_restore_ = false;
@@ -1260,12 +2050,14 @@ class MasterService {
     std::list<DiscardedReplicas> discarded_replicas_
         GUARDED_BY(discarded_replicas_mutex_);
     size_t offloading_queue_limit_ = 50000;
+    double offload_cap_ratio_ = 0.5;
 
     // Task manager
     ClientTaskManager task_manager_;
 
     struct ActiveDrainTask {
         UUID task_id;
+        std::string tenant_id;
         std::string key;
         std::string source_segment;
         std::string target_segment;
@@ -1307,7 +2099,8 @@ class MasterService {
     std::optional<std::string> SelectDrainTargetForKey(
         const ObjectMetadata& metadata, const std::string& source_segment,
         const std::vector<std::string>& requested_targets);
-    std::string MakeDrainUnitKey(const std::string& key,
+    std::string MakeDrainUnitKey(const std::string& tenant_id,
+                                 const std::string& key,
                                  const std::string& source_segment) const;
 
     std::thread job_dispatch_thread_;
diff --git a/mooncake-store/include/memory_alloc.h b/mooncake-store/include/memory_alloc.h
new file mode 100644
index 00000000..55a6d192
--- /dev/null
+++ b/mooncake-store/include/memory_alloc.h
@@ -0,0 +1,7 @@
+#pragma once
+
+#include <cstdint>
+#include <cstddef>
+
+void *hugepage_memory_alloc(size_t size);
+void hugepage_memory_free(void *ptr);
diff --git a/mooncake-store/include/mmap_arena.h b/mooncake-store/include/mmap_arena.h
new file mode 100644
index 00000000..b731efb1
--- /dev/null
+++ b/mooncake-store/include/mmap_arena.h
@@ -0,0 +1,108 @@
+// Copyright 2026 KVCache.AI
+// Simple arena allocator for mmap-based memory allocation
+// Used by SGLang HiCache for fast buffer allocation
+
+#pragma once
+
+#include <atomic>
+#include <cstddef>
+#include <cstdint>
+#include <mutex>
+
+namespace mooncake {
+
+/**
+ * @brief Simple lock-free arena allocator for mmap'd memory
+ *
+ * Thread-safe: allocate() is lock-free (CAS), initialize() is mutex-guarded.
+ */
+class MmapArena {
+   public:
+    static constexpr size_t kMinAlignment = 64;
+
+    struct Stats {
+        size_t pool_size;
+        size_t reserved_bytes;
+        size_t peak_reserved_bytes;
+        size_t num_allocations;
+        size_t num_failed_allocs;
+    };
+
+    MmapArena();
+    ~MmapArena();
+
+    // Delete copy and move operations (class contains atomics)
+    MmapArena(const MmapArena&) = delete;
+    MmapArena& operator=(const MmapArena&) = delete;
+    MmapArena(MmapArena&&) = delete;
+    MmapArena& operator=(MmapArena&&) = delete;
+
+    /**
+     * Initialize arena with a large mmap'd pool
+     * @param pool_size Total size to pre-allocate (aligned to huge pages)
+     * @param alignment Allocation alignment (default 64 bytes)
+     * @param allow_regular_page_fallback When true, the arena may retry the
+     *        initial pool mmap on regular pages after a hugepage failure.
+     *        Callers that need strict hugepage semantics should pass false.
+     * @return true on success
+     */
+    bool initialize(size_t pool_size, size_t alignment = kMinAlignment,
+                    bool allow_regular_page_fallback = true);
+
+    /**
+     * Allocate memory from arena
+     * @param size Number of bytes to allocate
+     * @param alignment Per-call alignment override (0 = use arena default).
+     *        Effective alignment is max(arena default, this parameter).
+     * @return Pointer to allocated memory, or nullptr if OOM
+     */
+    void* allocate(size_t size, size_t alignment = 0);
+
+    /**
+     * Get current arena statistics
+     */
+    Stats getStats() const;
+
+    /**
+     * Check if arena is initialized (thread-safe)
+     */
+    bool isInitialized() const {
+        return pool_base_.load(std::memory_order_acquire) != nullptr;
+    }
+
+    /**
+     * Check if pointer was allocated from this arena
+     * Used by free_buffer_mmap_memory to determine allocation type
+     * @param ptr Pointer to check
+     * @return true if ptr is in arena's address range
+     */
+    bool owns(const void* ptr) const;
+
+    /**
+     * Get pool base and size (for debugging/testing)
+     */
+    void* getPoolBase() const {
+        return pool_base_.load(std::memory_order_acquire);
+    }
+
+    size_t getPoolSize() const {
+        return pool_size_.load(std::memory_order_acquire);
+    }
+
+   private:
+    std::atomic<void*>
+        pool_base_;  // Base address of mmap'd pool (atomic for thread-safety)
+    std::atomic<size_t>
+        pool_size_;  // Total pool size (atomic for thread-safety)
+    std::atomic<size_t>
+        alignment_;  // Default allocation alignment (atomic for thread-safety)
+
+    std::atomic<size_t> alloc_cursor_;       // Monotonic reservation cursor
+    std::atomic<size_t> peak_reserved_;      // Peak reservation cursor
+    std::atomic<size_t> num_allocations_;    // Total allocations
+    std::atomic<size_t> num_failed_allocs_;  // Failed allocations (OOM)
+
+    std::mutex init_mutex_;  // Guards initialize() against concurrent calls
+};
+
+}  // namespace mooncake
diff --git a/mooncake-store/include/mutex.h b/mooncake-store/include/mutex.h
index 336440aa..caeefca3 100644
--- a/mooncake-store/include/mutex.h
+++ b/mooncake-store/include/mutex.h
@@ -125,6 +125,10 @@ class CAPABILITY("mutex") SpinLock {
    public:
     // Acquire/lock the spinlock.
     void lock() ACQUIRE() {
+        // Fast path: try to acquire directly
+        if (!flag.test_and_set(std::memory_order_acquire)) return;
+
+        // Slow path: spin wait with relaxed loads, acquire on success
         do {
             while (flag.test(std::memory_order_relaxed)) {
                 PAUSE();
diff --git a/mooncake-store/include/pinned_buffer_pool.h b/mooncake-store/include/pinned_buffer_pool.h
index 63f27a5c..11cdb9d8 100644
--- a/mooncake-store/include/pinned_buffer_pool.h
+++ b/mooncake-store/include/pinned_buffer_pool.h
@@ -83,7 +83,8 @@ class PinnedBufferPool {
         buf.capacity = size;
         buf.is_pinned = false;
 
-#if defined(USE_CUDA) || defined(USE_MUSA) || defined(USE_MACA)
+#if defined(USE_CUDA) || defined(USE_MUSA) || defined(USE_MACA) || \
+    defined(USE_HYGON) || defined(USE_COREX)
         if (cudaMallocHost(reinterpret_cast<void**>(&buf.data), size) ==
             cudaSuccess) {
             buf.is_pinned = true;
@@ -119,7 +120,8 @@ class PinnedBufferPool {
             delete[] buf.data;
             return;
         }
-#if defined(USE_CUDA) || defined(USE_MUSA) || defined(USE_MACA)
+#if defined(USE_CUDA) || defined(USE_MUSA) || defined(USE_MACA) || \
+    defined(USE_HYGON) || defined(USE_COREX)
         cudaFreeHost(buf.data);
 #elif defined(USE_HIP)
         hipHostFree(buf.data);
diff --git a/mooncake-store/include/pyclient.h b/mooncake-store/include/pyclient.h
index c47bbfb9..c2e2201c 100644
--- a/mooncake-store/include/pyclient.h
+++ b/mooncake-store/include/pyclient.h
@@ -1,10 +1,13 @@
 #pragma once
 
-#include <csignal>
 #include <atomic>
-#include <thread>
-#include <string>
+#include <chrono>
+#include <csignal>
+#include <map>
 #include <memory>
+#include <string>
+#include <thread>
+#include <unordered_map>
 #include <vector>
 
 #include "client_service.h"
@@ -28,6 +31,8 @@ enum ShmSegmentType : uint32_t {
     SHM_SEG_HOT_CACHE = 0,  // local hot cache backing memory
 };
 
+constexpr int32_t kInvalidPhysicalDeviceId = -1;
+
 // Return codes for health_check()
 enum HealthCheckStatus : int {
     HC_HEALTHY = 0,          // Fully connected, all links up
@@ -116,6 +121,7 @@ struct ShmRegisterRequest {
     uint64_t client_id_second;
     uint64_t dummy_base_addr;
     uint64_t shm_size;
+    int32_t device_id = kInvalidPhysicalDeviceId;
     bool is_local_buffer;
 };
 
@@ -204,6 +210,12 @@ class ClientRequester {
 // Python-specific wrapper class for client interface
 class PyClient {
    public:
+    // Request-scoped cache of fresh batch_query() results that can be reused by
+    // subsequent ranged reads in the same operation. Callers should not retain
+    // entries across independent requests.
+    using QueryResultCache =
+        std::unordered_map<std::string, tl::expected<QueryResult, ErrorCode>>;
+
     virtual ~PyClient() = 0;
     virtual int setup_real(
         const std::string &local_hostname, const std::string &metadata_server,
@@ -212,7 +224,8 @@ class PyClient {
         const std::string &master_server_addr,
         const std::shared_ptr<TransferEngine> &transfer_engine,
         const std::string &ipc_socket_path, bool enable_ssd_offload = false,
-        const std::string &ssd_offload_path = "") = 0;
+        const std::string &ssd_offload_path = "",
+        const std::string &tenant_id = "default") = 0;
 
     virtual int setup_dummy(size_t mem_pool_size, size_t local_buffer_size,
                             const std::string &server_address,
@@ -234,12 +247,19 @@ class PyClient {
     virtual int64_t get_into(const std::string &key, void *buffer,
                              size_t size) = 0;
 
+    // query_result_cache is optional and is only used to reuse fresh
+    // batch_query() results for this read; callers still do not provide any
+    // replica-selection or ranged-read metadata directly.
     virtual std::vector<std::vector<std::vector<int64_t>>> get_into_ranges(
         const std::vector<void *> &buffers,
         const std::vector<std::vector<std::string>> &all_keys,
         const std::vector<std::vector<std::vector<size_t>>> &all_dst_offsets,
         const std::vector<std::vector<std::vector<size_t>>> &all_src_offsets,
-        const std::vector<std::vector<std::vector<size_t>>> &all_sizes) = 0;
+        const std::vector<std::vector<std::vector<size_t>>> &all_sizes,
+        const QueryResultCache *query_result_cache = nullptr) = 0;
+
+    virtual std::vector<tl::expected<QueryResult, ErrorCode>> batch_query(
+        const std::vector<std::string> &keys) = 0;
 
     virtual std::vector<int64_t> batch_get_into(
         const std::vector<std::string> &keys,
@@ -354,4 +374,79 @@ class PyClient {
     std::shared_ptr<ClientBufferAllocator> client_buffer_allocator_ = nullptr;
 };
 
+inline uint64_t remaining_lease_ttl_ms(
+    const QueryResult &query_result,
+    std::chrono::steady_clock::time_point now) {
+    if (query_result.IsLeaseExpired(now)) {
+        return 0;
+    }
+    return static_cast<uint64_t>(
+        std::chrono::duration_cast<std::chrono::milliseconds>(
+            query_result.lease_timeout - now)
+            .count());
+}
+
+inline CachedQueryResultResponse to_cached_query_result_response(
+    const tl::expected<QueryResult, ErrorCode> &query_result,
+    std::chrono::steady_clock::time_point now) {
+    if (!query_result) {
+        return CachedQueryResultResponse(query_result.error());
+    }
+    if (query_result->IsLeaseExpired(now)) {
+        return CachedQueryResultResponse(ErrorCode::OBJECT_NOT_FOUND);
+    }
+    return CachedQueryResultResponse(GetReplicaListResponse(
+        std::vector<Replica::Descriptor>(query_result->replicas.begin(),
+                                         query_result->replicas.end()),
+        remaining_lease_ttl_ms(*query_result, now)));
+}
+
+inline tl::expected<QueryResult, ErrorCode> from_cached_query_result_response(
+    const CachedQueryResultResponse &cached_result,
+    std::chrono::steady_clock::time_point now) {
+    if (!cached_result.success) {
+        return tl::make_unexpected(cached_result.error);
+    }
+    return QueryResult(
+        std::vector<Replica::Descriptor>(cached_result.value.replicas.begin(),
+                                         cached_result.value.replicas.end()),
+        now + std::chrono::milliseconds(cached_result.value.lease_ttl_ms));
+}
+
+inline PyClient::QueryResultCache build_query_result_cache_from_cached_results(
+    const std::map<std::string, CachedQueryResultResponse>
+        &cached_query_results) {
+    PyClient::QueryResultCache query_result_cache;
+    query_result_cache.reserve(cached_query_results.size());
+    auto now = std::chrono::steady_clock::now();
+    for (const auto &[key, cached_result] : cached_query_results) {
+        auto query_result =
+            from_cached_query_result_response(cached_result, now);
+        if (!query_result || query_result->IsLeaseExpired(now)) {
+            continue;
+        }
+        query_result_cache.emplace(key, std::move(query_result));
+    }
+    return query_result_cache;
+}
+
+inline std::map<std::string, CachedQueryResultResponse>
+build_cached_query_results_from_query_result_cache(
+    const PyClient::QueryResultCache *query_result_cache) {
+    std::map<std::string, CachedQueryResultResponse> cached_query_results;
+    if (query_result_cache == nullptr) {
+        return cached_query_results;
+    }
+
+    auto now = std::chrono::steady_clock::now();
+    for (const auto &[key, query_result] : *query_result_cache) {
+        if (query_result && query_result->IsLeaseExpired(now)) {
+            continue;
+        }
+        cached_query_results.emplace(
+            key, to_cached_query_result_response(query_result, now));
+    }
+    return cached_query_results;
+}
+
 }  // namespace mooncake
diff --git a/mooncake-store/include/real_client.h b/mooncake-store/include/real_client.h
index cb6b5452..06789644 100644
--- a/mooncake-store/include/real_client.h
+++ b/mooncake-store/include/real_client.h
@@ -3,6 +3,7 @@
 #include <atomic>
 #include <boost/lockfree/queue.hpp>
 #include <csignal>
+#include <map>
 #include <memory>
 #include <shared_mutex>
 #include <string>
@@ -19,6 +20,8 @@
 #include "rpc_types.h"
 #include <ylt/coro_http/coro_http_server.hpp>
 #include <ylt/coro_rpc/coro_rpc_server.hpp>
+#include <ylt/coro_io/coro_io.hpp>
+#include <async_simple/coro/Lazy.h>
 
 namespace mooncake {
 
@@ -80,7 +83,8 @@ class RealClient : public PyClient {
         const std::shared_ptr<TransferEngine> &transfer_engine = nullptr,
         const std::string &ipc_socket_path = "",
         bool enable_ssd_offload = false,
-        const std::string &ssd_offload_path = "");
+        const std::string &ssd_offload_path = "",
+        const std::string &tenant_id = "default");
 
     int setup_dummy(size_t mem_pool_size, size_t local_buffer_size,
                     const std::string &server_address,
@@ -101,15 +105,24 @@ class RealClient : public PyClient {
 
     int unregister_buffer(void *buffer);
 
+    struct WritableBufferRegion {
+        void *base{nullptr};
+        size_t size{0};
+        size_t offset{0};
+    };
+
+    std::optional<WritableBufferRegion> resolve_writable_buffer_region(
+        void *buffer) const;
+
     /**
      * @brief Get object data directly into a pre-allocated buffer
      * @param key Key of the object to get
-     * @param buffer Pointer to the pre-allocated buffer (must be registered
-     * with register_buffer)
+     * @param buffer Pointer to a writable Store buffer, either explicitly
+     * registered with register_buffer() or inside the setup-time local buffer
      * @param size Size of the buffer
      * @return Number of bytes read on success, negative value on error
-     * @note The buffer address must be previously registered with
-     * register_buffer() for zero-copy operations
+     * @note The buffer address must resolve to Store-managed registered memory
+     * for zero-copy operations
      */
     int64_t get_into(const std::string &key, void *buffer, size_t size);
 
@@ -118,20 +131,17 @@ class RealClient : public PyClient {
         const std::vector<std::vector<std::string>> &all_keys,
         const std::vector<std::vector<std::vector<size_t>>> &all_dst_offsets,
         const std::vector<std::vector<std::vector<size_t>>> &all_src_offsets,
-        const std::vector<std::vector<std::vector<size_t>>> &all_sizes)
-        override;
+        const std::vector<std::vector<std::vector<size_t>>> &all_sizes,
+        const QueryResultCache *query_result_cache = nullptr) override;
 
     /**
-     * @brief Get object data directly into pre-allocated buffers for multiple
-     * keys (batch version)
-     * @param keys Vector of keys of the objects to get
-     * @param buffers Vector of pointers to the pre-allocated buffers
-     * @param sizes Vector of sizes of the buffers
-     * @return Vector of 64-bit integers, where each element is the number of
-     * bytes read on success, or a negative value on error
-     * @note The buffer addresses must be previously registered with
-     * register_buffer() for zero-copy operations
+     * @brief Batch query object placement/lease metadata for later read reuse
+     * @param keys Vector of keys to query
+     * @return Vector of query results in the same order as keys
      */
+    std::vector<tl::expected<QueryResult, ErrorCode>> batch_query(
+        const std::vector<std::string> &keys) override;
+
     std::vector<int64_t> batch_get_into(const std::vector<std::string> &keys,
                                         const std::vector<void *> &buffers,
                                         const std::vector<size_t> &sizes);
@@ -145,8 +155,9 @@ class RealClient : public PyClient {
      * @param all_sizes Vector of vectors of sizes of the buffers
      * @return Vector of integers, where each element is the number of bytes
      * read on success, or a negative value on error
-     * @note The buffer addresses must be previously registered with
-     * register_buffer() for zero-copy operations
+     * @note The buffer addresses must resolve to Store-managed registered
+     * memory, either explicit register_buffer() regions or the setup-time local
+     * buffer
      */
     std::vector<int> batch_get_into_multi_buffers(
         const std::vector<std::string> &keys,
@@ -157,30 +168,28 @@ class RealClient : public PyClient {
     /**
      * @brief Put object data directly from a pre-allocated buffer
      * @param key Key of the object to put
-     * @param buffer Pointer to the buffer containing data (must be registered
-     * with register_buffer)
+     * @param buffer Pointer to Store-managed registered memory, either
+     * explicitly registered with register_buffer() or inside the setup-time
+     * local buffer
      * @param size Size of the data to put
      * @return 0 on success, negative value on error
-     * @note The buffer address must be previously registered with
-     * register_buffer() for zero-copy operations
+     * @note The buffer address must resolve to Store-managed registered memory,
+     * either an explicit register_buffer() region or the setup-time local
+     * buffer
      */
     int put_from(const std::string &key, void *buffer, size_t size,
                  const ReplicateConfig &config = ReplicateConfig{});
 
     /**
-     * @brief Put object data directly from pre-allocated buffers for multiple
-     * keys(metadata version, better not be directly used in Python)
-     * @param keys Vector of keys of the objects to put
-     * @param buffers Vector of pointers to the pre-allocated buffers
-     * @param metadata_buffers Vector of pointers to the pre-allocated metadata
-     * buffers
-     * @param size Number of sizes of the buffers
-     * @param metadata_size Number of sizes of the metadata buffers
+     * @brief Put one object directly from a registered data buffer plus a
+     * registered metadata buffer
+     * @param key Key of the object to put
+     * @param buffer Pointer to the registered data buffer
+     * @param metadata_buffer Pointer to the registered metadata buffer
+     * @param size Size of the data buffer in bytes
+     * @param metadata_size Size of the metadata buffer in bytes
      * @param config Replication configuration
-     * @return Vector of integers, where each element is 0 on success, or a
-     * negative value on error
-     * @note The buffer addresses must be previously registered with
-     * register_buffer() for zero-copy operations
+     * @return 0 on success, negative value on error
      */
     int put_from_with_metadata(
         const std::string &key, void *buffer, void *metadata_buffer,
@@ -196,8 +205,9 @@ class RealClient : public PyClient {
      * @param config Replication configuration
      * @return Vector of integers, where each element is 0 on success, or a
      * negative value on error
-     * @note The buffer addresses must be previously registered with
-     * register_buffer() for zero-copy operations
+     * @note The buffer addresses must resolve to Store-managed registered
+     * memory, either explicit register_buffer() regions or the setup-time local
+     * buffer
      */
 
     std::vector<int> batch_put_from(
@@ -215,8 +225,9 @@ class RealClient : public PyClient {
      * @param config Replication configuration
      * @return Vector of integers, where each element is 0 on success, or a
      * negative value on error
-     * @note The buffer addresses must be previously registered with
-     * register_buffer() for zero-copy operations
+     * @note The buffer addresses must resolve to Store-managed registered
+     * memory, either explicit register_buffer() regions or the setup-time local
+     * buffer
      */
     std::vector<int> batch_put_from_multi_buffers(
         const std::vector<std::string> &keys,
@@ -406,10 +417,11 @@ class RealClient : public PyClient {
         const std::string &key, std::vector<std::span<const char>> values,
         const ReplicateConfig &config, const UUID &client_id);
 
-    std::vector<tl::expected<int64_t, ErrorCode>> batch_get_into_dummy_helper(
-        const std::vector<std::string> &keys,
-        const std::vector<uint64_t> &buffers, const std::vector<size_t> &sizes,
-        int32_t device_id, const UUID &client_id);
+    async_simple::coro::Lazy<std::vector<tl::expected<int64_t, ErrorCode>>>
+    batch_get_into_dummy_helper(const std::vector<std::string> &keys,
+                                const std::vector<uint64_t> &buffers,
+                                const std::vector<size_t> &sizes,
+                                int32_t device_id, const UUID &client_id);
 
     std::vector<tl::expected<void, ErrorCode>> batch_put_from_dummy_helper(
         const std::vector<std::string> &keys,
@@ -444,6 +456,8 @@ class RealClient : public PyClient {
         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,
+        const std::map<std::string, CachedQueryResultResponse>
+            &cached_query_results,
         int32_t device_id, const UUID &client_id);
 
     // Share mem management for dummy client
@@ -453,6 +467,9 @@ class RealClient : public PyClient {
                                                    size_t shm_size,
                                                    bool is_local_buffer,
                                                    const UUID &client_id);
+    tl::expected<void, ErrorCode> map_shm_internal_with_device(
+        int fd, uint64_t shm_base_addr, size_t shm_size, bool is_local_buffer,
+        int32_t physical_device_id, const UUID &client_id);
 
     tl::expected<void, ErrorCode> unmap_shm_internal(const UUID &client_id);
 
@@ -469,6 +486,9 @@ class RealClient : public PyClient {
     tl::expected<void, ErrorCode> ascend_unmap_shm_internal(
         const UUID &client_id);
 
+    tl::expected<bool, ErrorCode> is_shm_mapped_internal(
+        uint64_t dummy_base_addr, const UUID &client_id);
+
     tl::expected<void, ErrorCode> unregister_shm_buffer_internal(
         uint64_t dummy_base_addr, const UUID &client_id);
 
@@ -485,7 +505,8 @@ class RealClient : public PyClient {
         const std::shared_ptr<TransferEngine> &transfer_engine = nullptr,
         const std::string &ipc_socket_path = "", int local_rpc_port = 50052,
         bool enable_ssd_offload = false, bool start_offload_rpc_server = false,
-        const std::string &ssd_offload_path = "");
+        const std::string &ssd_offload_path = "",
+        const std::string &tenant_id = "default");
 
     // Overload that accepts a configuration dictionary
     tl::expected<void, ErrorCode> setup_internal(const ConfigDict &config);
@@ -511,6 +532,11 @@ class RealClient : public PyClient {
         uint64_t total_size;
     };
 
+    tl::expected<RangedReadMetadata, ErrorCode>
+    build_ranged_read_metadata_from_query_result(
+        const std::string &key,
+        tl::expected<QueryResult, ErrorCode> query_result);
+
     tl::expected<RangedReadMetadata, ErrorCode> resolve_ranged_read_metadata(
         const std::string &key);
 
@@ -534,7 +560,8 @@ class RealClient : public PyClient {
         std::vector<std::vector<std::vector<tl::expected<int64_t, ErrorCode>>>>
             *prepared_results = nullptr,
         const std::vector<std::vector<std::vector<bool>>> *valid_fragments =
-            nullptr);
+            nullptr,
+        const QueryResultCache *query_result_cache = nullptr);
 
     std::vector<tl::expected<int64_t, ErrorCode>> batch_get_into_internal(
         const std::vector<std::string> &keys,
@@ -637,6 +664,8 @@ class RealClient : public PyClient {
 
     std::map<std::string, std::vector<Replica::Descriptor>>
     batch_get_replica_desc(const std::vector<std::string> &keys);
+    std::vector<CachedQueryResultResponse> batch_get_query_results(
+        const std::vector<std::string> &keys);
     std::vector<Replica::Descriptor> get_replica_desc(const std::string &key);
 
     std::vector<std::string> batch_replica_clear(
@@ -645,7 +674,8 @@ class RealClient : public PyClient {
 
     tl::expected<PingResponse, ErrorCode> ping(const UUID &client_id);
 
-    tl::expected<BatchGetOffloadObjectResponse, ErrorCode>
+    async_simple::coro::Lazy<
+        tl::expected<BatchGetOffloadObjectResponse, ErrorCode>>
     batch_get_offload_object(const std::vector<std::string> &keys,
                              const std::vector<int64_t> &sizes);
 
@@ -665,7 +695,59 @@ class RealClient : public PyClient {
      */
     tl::expected<void, ErrorCode> batch_get_into_offload_object_internal(
         const std::string &target_rpc_service_addr,
-        std::unordered_map<std::string, Slice> &objects);
+        std::unordered_map<std::string, std::vector<Slice>> &objects);
+
+    int64_t get_offload_rpc_read_count() const {
+        return offload_rpc_read_count_.load(std::memory_order_relaxed);
+    }
+
+    /**
+     * @brief Mount a shared memory file region and return segment ids.
+     *        If size > max_mr_size, it will be split into multiple chunks
+     *        and mounted separately. RealClient will open(path) + mmap
+     *        internally for each chunk.
+     */
+    int mountSegment(const std::string &path, size_t offset, size_t size,
+                     const std::string &protocol, const std::string &location,
+                     std::vector<std::string> &out_segment_ids);
+
+    /**
+     * @brief Unmount segments by their ids and clean up local mmap/fd.
+     * @param grace_period_seconds 0 = immediate unmount (legacy behavior).
+     */
+    int unmountSegment(const std::vector<std::string> &segment_ids,
+                       uint64_t grace_period_seconds = 0);
+
+    /**
+     * @brief Allocate memory internally and mount segments to master.
+     *        If size > max_mr_size, it will be split into multiple chunks.
+     *        Memory is allocated via allocate_buffer_allocator_memory.
+     *        The actual allocated size (aligned up to Slab::kSize) is written
+     *        to out_allocated_size if non-null.
+     */
+    int allocateAndMountSegment(size_t size, const std::string &protocol,
+                                const std::string &location,
+                                std::vector<std::string> &out_segment_ids,
+                                size_t *out_allocated_size = nullptr);
+
+    /**
+     * @brief Unmount segments by their ids and free locally allocated memory.
+     * @param grace_period_seconds 0 = immediate unmount (legacy behavior).
+     */
+    int unmountAndFreeSegment(const std::vector<std::string> &segment_ids,
+                              uint64_t grace_period_seconds = 0);
+
+    struct MountedSegmentRecord {
+        void *mmap_base = nullptr;
+        size_t size = 0;
+        std::string path;
+    };
+
+    struct AllocatedSegmentRecord {
+        void *base = nullptr;
+        size_t size = 0;
+        std::string protocol;
+    };
 
     std::unique_ptr<AutoPortBinder> port_binder_ = nullptr;
 
@@ -695,11 +777,22 @@ class RealClient : public PyClient {
         }
     };
 
+    struct UbSegmentDeleter {
+        size_t size = 0;
+        std::string protocol = "ub";
+        void operator()(void *ptr) const {
+            if (ptr && size > 0) {
+                free_memory(protocol.c_str(), ptr);
+            }
+        }
+    };
+
     std::vector<std::unique_ptr<void, HugepageSegmentDeleter>>
         hugepage_segment_ptrs_;
     std::vector<std::unique_ptr<void, SegmentDeleter>> segment_ptrs_;
     std::vector<std::unique_ptr<void, AscendSegmentDeleter>>
         ascend_segment_ptrs_;
+    std::vector<std::unique_ptr<void, UbSegmentDeleter>> ub_segment_ptrs_;
     std::string protocol;
     std::string device_name;
     std::string local_hostname;
@@ -718,7 +811,7 @@ class RealClient : public PyClient {
         bool is_ascend = false;
         bool is_ipc = false;
         // Ascend physical device id from dummy (dummy-real RPC).
-        int32_t device_id = -1;
+        int32_t device_id = kInvalidPhysicalDeviceId;
         uint64_t vmm_handle = 0;
         std::string ipc_key_data;
     };
@@ -737,6 +830,7 @@ class RealClient : public PyClient {
 
     mutable std::shared_mutex registered_buffer_mutex_;
     std::unordered_map<void *, size_t> registered_buffer_sizes_;
+    std::optional<WritableBufferRegion> local_buffer_region_;
 
     // Dummy VA -> real VA using mapped_shms; last_hit_shm caches locality.
     bool map_dummy_range_in_shm(const MappedShm &shm, uint64_t dummy_addr,
@@ -766,9 +860,13 @@ class RealClient : public PyClient {
     // Ensure cleanup executes at most once across multiple entry points
     std::atomic<bool> closed_{false};
 
+    // Counts every LOCAL_DISK read served via peer offload-RPC.
+    std::atomic<int64_t> offload_rpc_read_count_{0};
+
     // Dummy Client manage related members
     void dummy_client_monitor_func();
     int start_dummy_client_monitor();
+    void stop_dummy_client_monitor();
     std::thread dummy_client_monitor_thread_;
     std::atomic<bool> dummy_client_monitor_running_{false};
     static constexpr uint64_t kDummyClientMonitorSleepMs =
@@ -803,6 +901,20 @@ class RealClient : public PyClient {
     void teardown_ascend_shm_buffer(MappedShm &shm);
     tl::expected<void, ErrorCode> setup_ascend_internal(
         size_t local_buffer_size);
+
+   private:
+    std::unordered_map<std::string, MountedSegmentRecord>
+        mounted_segment_records_;
+    std::mutex mounted_segment_records_mutex_;
+
+    std::unordered_map<std::string, AllocatedSegmentRecord>
+        allocated_segment_records_;
+    std::mutex allocated_segment_records_mutex_;
+
+    void ReleaseMountedSegmentRecord(const std::string &segment_id);
+    void ReleaseAllMountedSegmentRecords();
+    void ReleaseAllocatedSegmentRecord(const std::string &segment_id);
+    void ReleaseAllAllocatedSegmentRecords();
 };
 
 }  // namespace mooncake
diff --git a/mooncake-store/include/replica.h b/mooncake-store/include/replica.h
index 0e776818..677e77b3 100644
--- a/mooncake-store/include/replica.h
+++ b/mooncake-store/include/replica.h
@@ -4,6 +4,7 @@
 
 #include <boost/functional/hash.hpp>
 
+#include <atomic>
 #include <cstdint>
 #include <memory>
 #include <string>
@@ -26,15 +27,6 @@ namespace mooncake {
  */
 using ReplicaID = uint64_t;
 
-/**
- * @brief Type of buffer allocator used in the system
- */
-enum class ReplicaType {
-    MEMORY,     // Memory replica
-    DISK,       // Disk replica
-    LOCAL_DISK  // Local disk replica
-};
-
 /**
  * @brief Stream operator for ReplicaType
  */
@@ -43,7 +35,9 @@ inline std::ostream& operator<<(std::ostream& os,
     static const std::unordered_map<ReplicaType, std::string_view>
         replica_type_strings{{ReplicaType::MEMORY, "MEMORY"},
                              {ReplicaType::DISK, "DISK"},
-                             {ReplicaType::LOCAL_DISK, "LOCAL_DISK"}};
+                             {ReplicaType::LOCAL_DISK, "LOCAL_DISK"},
+                             {ReplicaType::NOF_SSD, "NOF_SSD"},
+                             {ReplicaType::ALL, "ALL"}};
 
     os << (replica_type_strings.count(replicaType)
                ? replica_type_strings.at(replicaType)
@@ -86,17 +80,35 @@ inline std::ostream& operator<<(std::ostream& os,
  */
 struct ReplicateConfig {
     size_t replica_num{1};
+    size_t nof_replica_num{0};
     bool with_soft_pin{false};
     bool with_hard_pin{false};  // Hard pin: object cannot be evicted
     std::vector<std::string>
         preferred_segments{};         // Preferred segments for allocation
     std::string preferred_segment{};  // Deprecated: Single preferred segment
                                       // for backward compatibility
+    std::vector<std::string>
+        preferred_nof_segments{};  // Preferred NoF segments for allocation
     bool prefer_alloc_in_same_node{false};
+    ObjectDataType data_type{ObjectDataType::UNKNOWN};
+    // Optional per-key routing group IDs. Empty string keeps that key
+    // ungrouped. Grouped keys share metadata routing, coalesced lease refresh,
+    // and memory eviction behavior.
+    std::optional<std::vector<std::string>> group_ids{};
+
+    ReplicateConfig ForSingleKey(size_t key_index) const {
+        ReplicateConfig key_config = *this;
+        if (group_ids.has_value()) {
+            key_config.group_ids =
+                std::vector<std::string>{group_ids->at(key_index)};
+        }
+        return key_config;
+    }
 
     friend std::ostream& operator<<(std::ostream& os,
                                     const ReplicateConfig& config) noexcept {
         os << "ReplicateConfig: { replica_num: " << config.replica_num
+           << ", nof_replica_num: " << config.nof_replica_num
            << ", with_soft_pin: " << config.with_soft_pin
            << ", with_hard_pin: " << config.with_hard_pin
            << ", preferred_segments: [";
@@ -109,16 +121,53 @@ struct ReplicateConfig {
             os << ", preferred_segment (deprecated): "
                << config.preferred_segment;
         }
+        os << ", preferred_nof_segments: [";
+        for (size_t i = 0; i < config.preferred_nof_segments.size(); ++i) {
+            os << config.preferred_nof_segments[i];
+            if (i < config.preferred_nof_segments.size() - 1) os << ", ";
+        }
+        os << "]";
         os << ", prefer_alloc_in_same_node: "
-           << config.prefer_alloc_in_same_node << " }";
+           << config.prefer_alloc_in_same_node
+           << ", data_type: " << config.data_type;
+        if (config.group_ids.has_value()) {
+            os << ", group_ids: [";
+            for (size_t i = 0; i < config.group_ids->size(); ++i) {
+                os << config.group_ids->at(i);
+                if (i + 1 < config.group_ids->size()) os << ", ";
+            }
+            os << "]";
+        }
+        os << " }";
         return os;
     }
 };
 
+enum class ReplicaWriteMode {
+    SINGLE_REPLICA,
+    FLEXIBLE_DUAL_REPLICA,
+    RELIABLE_MULTI_REPLICA,
+};
+
+inline ReplicaWriteMode DetermineReplicaWriteMode(
+    const ReplicateConfig& config) {
+    if (config.replica_num == 1 && config.nof_replica_num == 1) {
+        return ReplicaWriteMode::FLEXIBLE_DUAL_REPLICA;
+    }
+    if (config.replica_num > 1 || config.nof_replica_num > 1) {
+        return ReplicaWriteMode::RELIABLE_MULTI_REPLICA;
+    }
+    return ReplicaWriteMode::SINGLE_REPLICA;
+}
+
 struct MemoryReplicaData {
     std::unique_ptr<AllocatedBuffer> buffer;
 };
 
+struct NoFReplicaData {
+    std::unique_ptr<AllocatedBuffer> buffer;
+};
+
 struct DiskReplicaData {
     std::string file_path;
     uint64_t object_size = 0;
@@ -135,6 +184,11 @@ struct MemoryDescriptor {
     YLT_REFL(MemoryDescriptor, buffer_descriptor);
 };
 
+struct NoFDescriptor {
+    AllocatedBuffer::Descriptor buffer_descriptor;
+    YLT_REFL(NoFDescriptor, buffer_descriptor);
+};
+
 struct DiskDescriptor {
     std::string file_path{};
     uint64_t object_size = 0;
@@ -159,6 +213,19 @@ class Replica {
           status_(status),
           refcnt_(0) {}
 
+    // nof ssd replica constructor
+    Replica(std::unique_ptr<AllocatedBuffer> buffer, ReplicaStatus status,
+            ReplicaType replica_type)
+        : id_(next_id_.fetch_add(1)), status_(status), refcnt_(0) {
+        if (replica_type == ReplicaType::MEMORY) {
+            data_ = MemoryReplicaData{std::move(buffer)};
+        } else if (replica_type == ReplicaType::NOF_SSD) {
+            data_ = NoFReplicaData{std::move(buffer)};
+        } else {
+            LOG(ERROR) << "Invalid buffered replica type: " << replica_type;
+        }
+    }
+
     // disk replica constructor
     Replica(std::string file_path, uint64_t object_size, ReplicaStatus status)
         : id_(next_id_.fetch_add(1)),
@@ -169,17 +236,25 @@ class Replica {
         MasterMetricManager::instance().inc_allocated_file_size(object_size);
     }
 
+    // local disk replica constructor
     Replica(UUID client_id, uint64_t object_size,
             std::string transport_endpoint, ReplicaStatus status)
-        : data_(LocalDiskReplicaData{client_id, object_size,
+        : id_(next_id_.fetch_add(1)),
+          data_(LocalDiskReplicaData{client_id, object_size,
                                      std::move(transport_endpoint)}),
-          status_(status) {}
+          status_(status),
+          refcnt_(0) {
+        MasterMetricManager::instance().inc_allocated_file_size(object_size);
+    }
 
     ~Replica() {
-        if (status_ != ReplicaStatus::UNDEFINED && is_disk_replica()) {
-            const auto& disk_data = std::get<DiskReplicaData>(data_);
+        if (status_ == ReplicaStatus::UNDEFINED) return;
+        if (is_disk_replica()) {
             MasterMetricManager::instance().dec_allocated_file_size(
-                disk_data.object_size);
+                std::get<DiskReplicaData>(data_).object_size);
+        } else if (is_local_disk_replica()) {
+            MasterMetricManager::instance().dec_allocated_file_size(
+                std::get<LocalDiskReplicaData>(data_).object_size);
         }
     }
 
@@ -205,10 +280,14 @@ class Replica {
         }
 
         // Decrement metric for the current object before overwriting.
-        if (status_ != ReplicaStatus::UNDEFINED && is_disk_replica()) {
-            const auto& disk_data = std::get<DiskReplicaData>(data_);
-            MasterMetricManager::instance().dec_allocated_file_size(
-                disk_data.object_size);
+        if (status_ != ReplicaStatus::UNDEFINED) {
+            if (is_disk_replica()) {
+                MasterMetricManager::instance().dec_allocated_file_size(
+                    std::get<DiskReplicaData>(data_).object_size);
+            } else if (is_local_disk_replica()) {
+                MasterMetricManager::instance().dec_allocated_file_size(
+                    std::get<LocalDiskReplicaData>(data_).object_size);
+            }
         }
 
         id_ = src.id_;
@@ -261,6 +340,14 @@ class Replica {
         return replica.is_memory_replica();
     }
 
+    [[nodiscard]] bool is_nof_replica() const {
+        return std::holds_alternative<NoFReplicaData>(data_);
+    }
+
+    [[nodiscard]] static bool fn_is_nof_replica(const Replica& replica) {
+        return replica.is_nof_replica();
+    }
+
     [[nodiscard]] bool is_disk_replica() const {
         return std::holds_alternative<DiskReplicaData>(data_);
     }
@@ -285,6 +372,14 @@ class Replica {
         return false;  // DiskReplicaData does not have handles
     }
 
+    [[nodiscard]] bool has_invalid_nof_handle() const {
+        if (is_nof_replica()) {
+            const auto& nof_data = std::get<NoFReplicaData>(data_);
+            return !nof_data.buffer->isAllocatorValid();
+        }
+        return false;
+    }
+
     /**
      * @brief Check if a local_disk replica's owner client is still alive.
      * Used by CleanupStaleHandles to remove replicas belonging to expired
@@ -358,6 +453,9 @@ class Replica {
         ReplicaType operator()(const MemoryReplicaData&) const {
             return ReplicaType::MEMORY;
         }
+        ReplicaType operator()(const NoFReplicaData&) const {
+            return ReplicaType::NOF_SSD;
+        }
         ReplicaType operator()(const DiskReplicaData&) const {
             return ReplicaType::DISK;
         }
@@ -368,7 +466,8 @@ class Replica {
 
     struct Descriptor {
         ReplicaID id;
-        std::variant<MemoryDescriptor, DiskDescriptor, LocalDiskDescriptor>
+        std::variant<MemoryDescriptor, NoFDescriptor, DiskDescriptor,
+                     LocalDiskDescriptor>
             descriptor_variant;
         ReplicaStatus status;
         YLT_REFL(Descriptor, id, descriptor_variant, status);
@@ -382,6 +481,14 @@ class Replica {
             return std::holds_alternative<MemoryDescriptor>(descriptor_variant);
         }
 
+        bool is_nof_replica() noexcept {
+            return std::holds_alternative<NoFDescriptor>(descriptor_variant);
+        }
+
+        bool is_nof_replica() const noexcept {
+            return std::holds_alternative<NoFDescriptor>(descriptor_variant);
+        }
+
         bool is_disk_replica() noexcept {
             return std::holds_alternative<DiskDescriptor>(descriptor_variant);
         }
@@ -408,6 +515,13 @@ class Replica {
             throw std::runtime_error("Expected MemoryDescriptor");
         }
 
+        NoFDescriptor& get_nof_descriptor() {
+            if (auto* desc = std::get_if<NoFDescriptor>(&descriptor_variant)) {
+                return *desc;
+            }
+            throw std::runtime_error("Expected NoFDescriptor");
+        }
+
         DiskDescriptor& get_disk_descriptor() {
             if (auto* desc = std::get_if<DiskDescriptor>(&descriptor_variant)) {
                 return *desc;
@@ -431,6 +545,13 @@ class Replica {
             throw std::runtime_error("Expected MemoryDescriptor");
         }
 
+        const NoFDescriptor& get_nof_descriptor() const {
+            if (auto* desc = std::get_if<NoFDescriptor>(&descriptor_variant)) {
+                return *desc;
+            }
+            throw std::runtime_error("Expected NoFDescriptor");
+        }
+
         const DiskDescriptor& get_disk_descriptor() const {
             if (auto* desc = std::get_if<DiskDescriptor>(&descriptor_variant)) {
                 return *desc;
@@ -451,7 +572,8 @@ class Replica {
     inline static std::atomic<ReplicaID> next_id_{1};
 
     ReplicaID id_;
-    std::variant<MemoryReplicaData, DiskReplicaData, LocalDiskReplicaData>
+    std::variant<MemoryReplicaData, NoFReplicaData, DiskReplicaData,
+                 LocalDiskReplicaData>
         data_;
     ReplicaStatus status_{ReplicaStatus::UNDEFINED};
 
@@ -477,6 +599,18 @@ inline Replica::Descriptor Replica::get_descriptor() const {
             LOG(ERROR) << "Trying to get invalid memory replica descriptor";
         }
         desc.descriptor_variant = std::move(mem_desc);
+    } else if (is_nof_replica()) {
+        const auto& nof_data = std::get<NoFReplicaData>(data_);
+        NoFDescriptor nof_desc;
+        if (nof_data.buffer) {
+            nof_desc.buffer_descriptor = nof_data.buffer->get_descriptor();
+        } else {
+            nof_desc.buffer_descriptor.size_ = 0;
+            nof_desc.buffer_descriptor.buffer_address_ = 0;
+            nof_desc.buffer_descriptor.transport_endpoint_ = "";
+            LOG(ERROR) << "Trying to get invalid nof replica descriptor";
+        }
+        desc.descriptor_variant = std::move(nof_desc);
     } else if (is_disk_replica()) {
         const auto& disk_data = std::get<DiskReplicaData>(data_);
         DiskDescriptor disk_desc;
@@ -506,6 +640,15 @@ inline std::vector<std::optional<std::string>> Replica::get_segment_names()
             segment_names.push_back(std::nullopt);
         }
         return segment_names;
+    } else if (is_nof_replica()) {
+        const auto& nof_data = std::get<NoFReplicaData>(data_);
+        std::vector<std::optional<std::string>> segment_names;
+        if (nof_data.buffer && nof_data.buffer->isAllocatorValid()) {
+            segment_names.push_back(nof_data.buffer->getSegmentName());
+        } else {
+            segment_names.push_back(std::nullopt);
+        }
+        return segment_names;
     }
     return std::vector<std::optional<std::string>>();
 }
@@ -521,6 +664,13 @@ inline std::ostream& operator<<(std::ostream& os, const Replica& replica) {
             os << *mem_data.buffer;
         }
         os << "]";
+    } else if (replica.is_nof_replica()) {
+        const auto& nof_data = std::get<NoFReplicaData>(replica.data_);
+        os << "type: NOF_SSD, buffers: [";
+        if (nof_data.buffer) {
+            os << *nof_data.buffer;
+        }
+        os << "]";
     } else if (replica.is_disk_replica()) {
         const auto& disk_data = std::get<DiskReplicaData>(replica.data_);
         os << "type: DISK, file_path: " << disk_data.file_path
diff --git a/mooncake-store/include/rpc_service.h b/mooncake-store/include/rpc_service.h
index f771774f..bf93f386 100644
--- a/mooncake-store/include/rpc_service.h
+++ b/mooncake-store/include/rpc_service.h
@@ -1,41 +1,47 @@
 #pragma once
 
 #include <csignal>
-#include <atomic>
-#include <memory>
-#include <mutex>
-#include <optional>
+
 #include <string>
 #include <boost/functional/hash.hpp>
 #include <cstdint>
-#include <thread>
-#include <ylt/coro_http/coro_http_server.hpp>
 #include <ylt/coro_rpc/coro_rpc_server.hpp>
 #include <ylt/util/tl/expected.hpp>
 
-#include "ha/ha_types.h"
 #include "master_service.h"
 #include "types.h"
 #include "rpc_types.h"
 #include "master_config.h"
+#include "segment.h"
 
 namespace mooncake {
 
-extern const uint64_t kMetricReportIntervalSeconds;
-
+// Forward declaration
+class HttpMetadataServer;
 class WrappedMasterService {
    public:
-    WrappedMasterService(const WrappedMasterServiceConfig& config);
+    // Constructor with optional metadata-cleanup-on-timeout configuration.
+    // - http_metadata_server: in-process pointer used when the HTTP metadata
+    //   server is co-located in the master process (nullptr = not co-located).
+    // - http_metadata_remote_url: http(s) connection string used when the
+    //   metadata server is deployed separately (empty = none). Only consulted
+    //   when http_metadata_server is nullptr. If both are unset, cleanup is
+    //   disabled.
+    WrappedMasterService(const WrappedMasterServiceConfig& config,
+                         HttpMetadataServer* http_metadata_server = nullptr,
+                         const std::string& http_metadata_remote_url = "");
 
     ~WrappedMasterService();
 
-    tl::expected<bool, ErrorCode> ExistKey(const std::string& key);
+    tl::expected<bool, ErrorCode> ExistKey(
+        const std::string& key, const std::string& tenant_id = "default");
 
     tl::expected<MasterMetricManager::CacheHitStatDict, ErrorCode>
     CalcCacheStats();
 
     std::vector<tl::expected<bool, ErrorCode>> BatchExistKey(
-        const std::vector<std::string>& keys);
+        const std::vector<std::string>& keys,
+        const std::string& tenant_id = "default");
 
     tl::expected<
         std::unordered_map<UUID, std::vector<std::string>, boost::hash<UUID>>,
@@ -49,81 +55,120 @@ class WrappedMasterService {
     tl::expected<
         std::unordered_map<std::string, std::vector<Replica::Descriptor>>,
         ErrorCode>
-    GetReplicaListByRegex(const std::string& str);
+    GetReplicaListByRegex(const std::string& str,
+                          const std::string& tenant_id = "default");
 
     tl::expected<GetReplicaListResponse, ErrorCode> GetReplicaList(
-        const std::string& key);
+        const std::string& key, const std::string& tenant_id = "default");
 
     std::vector<tl::expected<GetReplicaListResponse, ErrorCode>>
-    BatchGetReplicaList(const std::vector<std::string>& keys);
+    BatchGetReplicaList(const std::vector<std::string>& keys,
+                        const std::string& tenant_id = "default");
 
     tl::expected<std::vector<Replica::Descriptor>, ErrorCode> PutStart(
         const UUID& client_id, const std::string& key,
-        const uint64_t slice_length, const ReplicateConfig& config);
+        const uint64_t slice_length, const ReplicateConfig& config,
+        const std::string& tenant_id = "default");
 
-    tl::expected<void, ErrorCode> PutEnd(const UUID& client_id,
-                                         const std::string& key,
-                                         ReplicaType replica_type);
+    tl::expected<void, ErrorCode> PutEnd(
+        const UUID& client_id, const std::string& key,
+        ReplicaType replica_type = ReplicaType::ALL,
+        const std::string& tenant_id = "default");
 
-    tl::expected<void, ErrorCode> PutRevoke(const UUID& client_id,
-                                            const std::string& key,
-                                            ReplicaType replica_type);
+    tl::expected<void, ErrorCode> PutRevoke(
+        const UUID& client_id, const std::string& key,
+        ReplicaType replica_type = ReplicaType::ALL,
+        const std::string& tenant_id = "default");
 
     std::vector<tl::expected<std::vector<Replica::Descriptor>, ErrorCode>>
     BatchPutStart(const UUID& client_id, const std::vector<std::string>& keys,
                   const std::vector<uint64_t>& slice_lengths,
-                  const ReplicateConfig& config);
+                  const ReplicateConfig& config,
+                  const std::string& tenant_id = "default");
 
     std::vector<tl::expected<void, ErrorCode>> BatchPutEnd(
-        const UUID& client_id, const std::vector<std::string>& keys);
+        const UUID& client_id, const std::vector<std::string>& keys,
+        ReplicaType replica_type = ReplicaType::ALL,
+        const std::string& tenant_id = "default");
 
     std::vector<tl::expected<void, ErrorCode>> BatchPutRevoke(
-        const UUID& client_id, const std::vector<std::string>& keys);
+        const UUID& client_id, const std::vector<std::string>& keys,
+        ReplicaType replica_type = ReplicaType::ALL,
+        const std::string& tenant_id = "default");
 
     tl::expected<std::vector<Replica::Descriptor>, ErrorCode> UpsertStart(
         const UUID& client_id, const std::string& key,
-        const uint64_t slice_length, const ReplicateConfig& config);
+        const uint64_t slice_length, const ReplicateConfig& config,
+        const std::string& tenant_id = "default");
 
-    tl::expected<void, ErrorCode> UpsertEnd(const UUID& client_id,
-                                            const std::string& key,
-                                            ReplicaType replica_type);
+    tl::expected<void, ErrorCode> UpsertEnd(
+        const UUID& client_id, const std::string& key,
+        ReplicaType replica_type = ReplicaType::ALL,
+        const std::string& tenant_id = "default");
 
-    tl::expected<void, ErrorCode> UpsertRevoke(const UUID& client_id,
-                                               const std::string& key,
-                                               ReplicaType replica_type);
+    tl::expected<void, ErrorCode> UpsertRevoke(
+        const UUID& client_id, const std::string& key,
+        ReplicaType replica_type = ReplicaType::ALL,
+        const std::string& tenant_id = "default");
 
     std::vector<tl::expected<std::vector<Replica::Descriptor>, ErrorCode>>
     BatchUpsertStart(const UUID& client_id,
                      const std::vector<std::string>& keys,
                      const std::vector<uint64_t>& slice_lengths,
-                     const ReplicateConfig& config);
+                     const ReplicateConfig& config,
+                     const std::string& tenant_id = "default");
 
     std::vector<tl::expected<void, ErrorCode>> BatchUpsertEnd(
-        const UUID& client_id, const std::vector<std::string>& keys);
+        const UUID& client_id, const std::vector<std::string>& keys,
+        const std::string& tenant_id = "default");
 
     std::vector<tl::expected<void, ErrorCode>> BatchUpsertRevoke(
-        const UUID& client_id, const std::vector<std::string>& keys);
+        const UUID& client_id, const std::vector<std::string>& keys,
+        const std::string& tenant_id = "default");
 
-    tl::expected<void, ErrorCode> Remove(const std::string& key,
-                                         bool force = false);
+    tl::expected<void, ErrorCode> Remove(
+        const std::string& key, bool force = false,
+        const std::string& tenant_id = "default");
 
-    tl::expected<long, ErrorCode> RemoveByRegex(const std::string& str,
-                                                bool force = false);
+    tl::expected<long, ErrorCode> RemoveByRegex(
+        const std::string& str, bool force = false,
+        const std::string& tenant_id = "default");
 
-    long RemoveAll(bool force = false);
+    long RemoveAll(bool force = false,
+                   const std::string& tenant_id = "default");
 
     std::vector<tl::expected<void, ErrorCode>> BatchRemove(
-        const std::vector<std::string>& keys, bool force = false);
+        const std::vector<std::string>& keys, bool force = false,
+        const std::string& tenant_id = "default");
 
     tl::expected<void, ErrorCode> MountSegment(const Segment& segment,
                                                const UUID& client_id);
 
+    tl::expected<void, ErrorCode> MountNoFSegment(const NoFSegment& segment,
+                                                  const UUID& client_id);
+
     tl::expected<void, ErrorCode> ReMountSegment(
         const std::vector<Segment>& segments, const UUID& client_id);
 
+    tl::expected<void, ErrorCode> ReMountNoFSegment(
+        const std::vector<NoFSegment>& segments, const UUID& client_id);
+
     tl::expected<void, ErrorCode> UnmountSegment(const UUID& segment_id,
                                                  const UUID& client_id);
 
+    tl::expected<void, ErrorCode> GracefulUnmountSegment(
+        const UUID& segment_id, const UUID& client_id,
+        uint64_t grace_period_ms);
+
+    tl::expected<void, ErrorCode> UnmountNoFSegment(const UUID& segment_id,
+                                                    const UUID& client_id);
+
+    [[nodiscard]] tl::expected<std::vector<NoFSegment>, ErrorCode>
+    GetAllNoFSegments();
+
+    [[nodiscard]] tl::expected<std::vector<NoFSegmentOwnerInfo>, ErrorCode>
+    GetNoFSegmentsByName(const std::string& segment_name);
+
     tl::expected<std::string, ErrorCode> GetFsdir();
 
     tl::expected<GetStorageConfigResponse, ErrorCode> GetStorageConfig();
@@ -132,23 +177,56 @@ class WrappedMasterService {
 
     tl::expected<std::string, ErrorCode> ServiceReady();
 
+    tl::expected<std::vector<TenantQuotaSnapshot>, ErrorCode>
+    ListTenantQuotaSnapshots();
+    tl::expected<TenantQuotaSnapshot, ErrorCode> GetTenantQuotaSnapshot(
+        const std::string& tenant_id);
+    tl::expected<TenantQuotaSnapshot, ErrorCode> UpsertTenantQuotaPolicy(
+        const std::string& tenant_id, uint64_t requested_quota_bytes);
+    tl::expected<std::optional<TenantQuotaSnapshot>, ErrorCode>
+    DeleteTenantQuotaPolicy(const std::string& tenant_id);
+    tl::expected<uint64_t, ErrorCode> GetTenantQuotaAllocatableCapacityBytes();
+
     tl::expected<std::vector<std::string>, ErrorCode> GetAllKeysForAdmin();
 
     tl::expected<std::vector<std::string>, ErrorCode> GetAllSegmentsForAdmin();
 
+    tl::expected<std::vector<MasterService::SegmentDetailInfo>, ErrorCode>
+    GetSegmentsDetailForAdmin();
+
     tl::expected<std::pair<uint64_t, uint64_t>, ErrorCode> QuerySegmentForAdmin(
         const std::string& segment);
 
     tl::expected<void, ErrorCode> MountLocalDiskSegment(const UUID& client_id,
                                                         bool enable_offloading);
 
-    tl::expected<std::unordered_map<std::string, int64_t>, ErrorCode>
+    tl::expected<std::vector<OffloadTaskItem>, ErrorCode>
     OffloadObjectHeartbeat(const UUID& client_id, bool enable_offloading);
 
+    tl::expected<void, ErrorCode> ReportSsdCapacity(
+        const UUID& client_id, int64_t ssd_total_capacity_bytes);
+
     tl::expected<void, ErrorCode> NotifyOffloadSuccess(
-        const UUID& client_id, const std::vector<std::string>& keys,
+        const UUID& client_id, const std::vector<OffloadTaskItem>& tasks,
         const std::vector<StorageObjectMetadata>& metadatas);
 
+    // Promotion-on-hit RPCs.
+    tl::expected<std::vector<PromotionTaskItem>, ErrorCode>
+    PromotionObjectHeartbeat(const UUID& client_id);
+
+    tl::expected<PromotionAllocStartResponse, ErrorCode> PromotionAllocStart(
+        const UUID& client_id, const std::string& key,
+        const std::string& tenant_id, uint64_t size,
+        const std::vector<std::string>& preferred_segments);
+
+    tl::expected<void, ErrorCode> NotifyPromotionSuccess(
+        const UUID& client_id, const std::string& key,
+        const std::string& tenant_id);
+
+    tl::expected<void, ErrorCode> NotifyPromotionFailure(
+        const UUID& client_id, const std::string& key,
+        const std::string& tenant_id);
+
     tl::expected<UUID, ErrorCode> CreateDrainJob(
         const CreateDrainJobRequest& request);
 
@@ -158,10 +236,14 @@ class WrappedMasterService {
 
     tl::expected<SegmentStatus, ErrorCode> QuerySegmentStatus(
         const std::string& segment_name);
+    tl::expected<SegmentStatus, ErrorCode> QuerySegmentStatusById(
+        const UUID& segment_id);
     tl::expected<UUID, ErrorCode> CreateCopyTask(
-        const std::string& key, const std::vector<std::string>& targets);
+        const std::string& key, const std::string& tenant_id,
+        const std::vector<std::string>& targets);
 
     tl::expected<UUID, ErrorCode> CreateMoveTask(const std::string& key,
+                                                 const std::string& tenant_id,
                                                  const std::string& source,
                                                  const std::string& target);
 
@@ -175,90 +257,43 @@ class WrappedMasterService {
 
     tl::expected<CopyStartResponse, ErrorCode> CopyStart(
         const UUID& client_id, const std::string& key,
-        const std::string& src_segment,
+        const std::string& tenant_id, const std::string& src_segment,
         const std::vector<std::string>& tgt_segments);
 
     tl::expected<void, ErrorCode> CopyEnd(const UUID& client_id,
-                                          const std::string& key);
+                                          const std::string& key,
+                                          const std::string& tenant_id);
 
     tl::expected<void, ErrorCode> CopyRevoke(const UUID& client_id,
-                                             const std::string& key);
+                                             const std::string& key,
+                                             const std::string& tenant_id);
 
     tl::expected<MoveStartResponse, ErrorCode> MoveStart(
         const UUID& client_id, const std::string& key,
-        const std::string& src_segment, const std::string& tgt_segment);
+        const std::string& tenant_id, const std::string& src_segment,
+        const std::string& tgt_segment);
 
     tl::expected<void, ErrorCode> MoveEnd(const UUID& client_id,
-                                          const std::string& key);
+                                          const std::string& key,
+                                          const std::string& tenant_id);
 
     tl::expected<void, ErrorCode> MoveRevoke(const UUID& client_id,
-                                             const std::string& key);
+                                             const std::string& key,
+                                             const std::string& tenant_id);
 
     tl::expected<void, ErrorCode> EvictDiskReplica(const UUID& client_id,
                                                    const std::string& key,
+                                                   const std::string& tenant_id,
                                                    ReplicaType replica_type);
 
     std::vector<tl::expected<void, ErrorCode>> BatchEvictDiskReplica(
         const UUID& client_id, const std::vector<std::string>& keys,
-        ReplicaType replica_type);
+        const std::string& tenant_id, ReplicaType replica_type);
 
    private:
     MasterService master_service_;
 };
 
-class MasterAdminServer {
-   public:
-    MasterAdminServer(uint16_t http_port, bool enable_metric_reporting);
-
-    ~MasterAdminServer();
-
-    bool Start();
-
-    void Stop();
-
-    void SetRuntimeState(ha::MasterRuntimeState state);
-
-    void SetObservedLeader(const std::optional<ha::MasterView>& leader_view);
-
-    void SetServiceDelegate(std::shared_ptr<WrappedMasterService> service);
-
-    void SetServiceAvailable(bool available);
-
-   private:
-    struct RuntimeSnapshot {
-        ha::MasterRuntimeState state = ha::MasterRuntimeState::kStarting;
-        std::optional<ha::MasterView> leader_view;
-        std::shared_ptr<WrappedMasterService> service;
-        bool service_available = false;
-    };
-
-    RuntimeSnapshot SnapshotState() const;
-
-    std::string BuildMetricsText() const;
-
-    std::string BuildMetricsSummaryText() const;
-
-    std::string BuildHealthJson() const;
-
-    std::string BuildLeaderJson() const;
-
-    std::shared_ptr<WrappedMasterService> GetActiveService() const;
-
-    void InitHttpServer();
-
-    uint16_t http_port_;
-    bool enable_metric_reporting_ = false;
-    coro_http::coro_http_server http_server_;
-    std::thread metric_report_thread_;
-    std::atomic<bool> metric_report_running_{false};
-    std::atomic<bool> started_{false};
-    mutable std::mutex state_mutex_;
-    ha::MasterRuntimeState state_{ha::MasterRuntimeState::kStarting};
-    std::optional<ha::MasterView> leader_view_;
-    std::shared_ptr<WrappedMasterService> service_;
-    bool service_available_ = false;
-};
-
 void RegisterRpcService(coro_rpc::coro_rpc_server& server,
                         mooncake::WrappedMasterService& wrapped_master_service);
 
diff --git a/mooncake-store/include/rpc_types.h b/mooncake-store/include/rpc_types.h
index dcaeeacc..fbf4fe37 100644
--- a/mooncake-store/include/rpc_types.h
+++ b/mooncake-store/include/rpc_types.h
@@ -41,6 +41,20 @@ struct GetReplicaListResponse {
 };
 YLT_REFL(GetReplicaListResponse, replicas, lease_ttl_ms);
 
+struct CachedQueryResultResponse {
+    bool success;
+    GetReplicaListResponse value;
+    ErrorCode error;
+
+    CachedQueryResultResponse()
+        : success(false), value(), error(ErrorCode::INVALID_PARAMS) {}
+    CachedQueryResultResponse(GetReplicaListResponse&& value_param)
+        : success(true), value(std::move(value_param)), error(ErrorCode::OK) {}
+    CachedQueryResultResponse(ErrorCode error_param)
+        : success(false), value(), error(error_param) {}
+};
+YLT_REFL(CachedQueryResultResponse, success, value, error);
+
 /**
  * @brief Response structure for GetStorageConfig operation
  */
@@ -58,6 +72,17 @@ struct GetStorageConfigResponse {
 };
 YLT_REFL(GetStorageConfigResponse, fsdir, enable_disk_eviction, quota_bytes);
 
+struct NoFSegmentOwnerInfo {
+    UUID segment_id;
+    UUID client_id;
+
+    NoFSegmentOwnerInfo() = default;
+    NoFSegmentOwnerInfo(const UUID& segment_id_param,
+                        const UUID& client_id_param)
+        : segment_id(segment_id_param), client_id(client_id_param) {}
+};
+YLT_REFL(NoFSegmentOwnerInfo, segment_id, client_id);
+
 /**
  * @brief Response structure for CopyStart operation
  */
@@ -67,6 +92,15 @@ struct CopyStartResponse {
 };
 YLT_REFL(CopyStartResponse, source, targets);
 
+/**
+ * @brief Response structure for PromotionAllocStart (L2->L1 promotion-on-hit).
+ * Carries the staged PROCESSING MEMORY replica descriptor.
+ */
+struct PromotionAllocStartResponse {
+    Replica::Descriptor memory_descriptor;
+};
+YLT_REFL(PromotionAllocStartResponse, memory_descriptor);
+
 /**
  * @brief Response structure for MoveStart operation
  */
diff --git a/mooncake-store/include/segment.h b/mooncake-store/include/segment.h
index 0feb9633..a9acd6f6 100644
--- a/mooncake-store/include/segment.h
+++ b/mooncake-store/include/segment.h
@@ -1,6 +1,7 @@
 #pragma once
 
 #include <boost/functional/hash.hpp>
+#include <chrono>
 #include <ostream>
 #include <shared_mutex>
 #include <string>
@@ -10,6 +11,7 @@
 
 #include "allocation_strategy.h"
 #include "allocator.h"
+#include "rpc_types.h"
 #include "types.h"
 
 namespace mooncake {
@@ -21,7 +23,8 @@ enum class SegmentStatus {
     OK,             // Segment is mounted and available for allocation
     DRAINING,       // Segment remains readable but accepts no new allocations
     DRAINED,        // Segment has been drained and awaits unmount
-    UNMOUNTING,     // Segment is under unmounting
+    GRACEFULLY_UNMOUNTING,  // Readable, no new allocations, timer running
+    UNMOUNTING,             // Segment is under unmounting
 };
 
 /**
@@ -30,11 +33,13 @@ enum class SegmentStatus {
 inline std::ostream& operator<<(std::ostream& os,
                                 const SegmentStatus& status) noexcept {
     static const std::unordered_map<SegmentStatus, std::string_view>
-        status_strings{{SegmentStatus::UNDEFINED, "UNDEFINED"},
-                       {SegmentStatus::OK, "OK"},
-                       {SegmentStatus::DRAINING, "DRAINING"},
-                       {SegmentStatus::DRAINED, "DRAINED"},
-                       {SegmentStatus::UNMOUNTING, "UNMOUNTING"}};
+        status_strings{
+            {SegmentStatus::UNDEFINED, "UNDEFINED"},
+            {SegmentStatus::OK, "OK"},
+            {SegmentStatus::DRAINING, "DRAINING"},
+            {SegmentStatus::DRAINED, "DRAINED"},
+            {SegmentStatus::GRACEFULLY_UNMOUNTING, "GRACEFULLY_UNMOUNTING"},
+            {SegmentStatus::UNMOUNTING, "UNMOUNTING"}};
 
     os << (status_strings.count(status) ? status_strings.at(status)
                                         : "UNKNOWN");
@@ -47,11 +52,49 @@ struct MountedSegment {
     std::shared_ptr<BufferAllocatorBase> buf_allocator;
 };
 
+struct MountedNoFSegment {
+    NoFSegment segment;
+    UUID client_id;
+    SegmentStatus status;
+    std::shared_ptr<BufferAllocatorBase> buf_allocator;
+};
+
+struct MountedNoFSegmentSnapshot {
+    UUID segment_id;
+    UUID client_id;
+    NoFSegment segment;
+    SegmentStatus status;
+};
+
+/**
+ * @brief Stream operator for MountedNoFSegmentSnapshot
+ */
+inline std::ostream& operator<<(
+    std::ostream& os, const MountedNoFSegmentSnapshot& snapshot) noexcept {
+    os << "{segment_id=" << snapshot.segment_id
+       << ", client_id=" << snapshot.client_id
+       << ", segment.id=" << snapshot.segment.id
+       << ", segment.name=" << snapshot.segment.name
+       << ", segment.base=" << snapshot.segment.base
+       << ", segment.size=" << snapshot.segment.size
+       << ", segment.te_endpoint=" << snapshot.segment.te_endpoint
+       << ", status=" << snapshot.status << "}";
+    return os;
+}
+
 struct LocalDiskSegment {
     mutable Mutex offloading_mutex_;
     bool enable_offloading;
-    std::unordered_map<std::string, int64_t> GUARDED_BY(offloading_mutex_)
-        offloading_objects;
+    int64_t ssd_total_capacity_bytes = 0;  // last reported by client heartbeat
+    std::atomic<int64_t> ssd_used_bytes{0};
+    std::unordered_map<std::string, OffloadTaskItem> GUARDED_BY(
+        offloading_mutex_) offloading_objects;
+    // Promotion-on-hit pending work for this client. Populated by master's
+    // TryPushPromotionQueue when a Get hits a LOCAL_DISK-only key on this
+    // client. Drained by PromotionObjectHeartbeat. Same locking as
+    // offloading_objects (offloading_mutex_).
+    std::unordered_map<std::string, PromotionTaskItem> GUARDED_BY(
+        offloading_mutex_) promotion_objects;
     explicit LocalDiskSegment(bool enable_offloading)
         : enable_offloading(enable_offloading) {}
 
@@ -101,6 +144,12 @@ class ScopedSegmentAccess {
     ErrorCode PrepareUnmountSegment(const UUID& segment_id,
                                     size_t& metrics_dec_capacity);
 
+    /**
+     * @brief Prepare a segment for graceful unmount: remove allocator but keep
+     *        segment metadata. Status becomes GRACEFULLY_UNMOUNTING.
+     */
+    ErrorCode PrepareGracefulUnmountSegment(const UUID& segment_id);
+
     /**
      * @brief Deleting the segment to complete the unmounting operation
      */
@@ -164,6 +213,12 @@ class ScopedSegmentAccess {
     ErrorCode GetSegmentStatusByName(const std::string& segment_name,
                                      SegmentStatus& status) const;
 
+    /**
+     * @brief Query the lifecycle status of a segment by id.
+     */
+    ErrorCode GetSegmentStatusById(const UUID& segment_id,
+                                   SegmentStatus& status) const;
+
     /**
      * @brief Update the lifecycle status of a segment by name.
      */
@@ -181,6 +236,78 @@ class ScopedSegmentAccess {
     std::unique_lock<std::shared_mutex> lock_;
 };
 
+class NoFSegmentManager;
+
+/**
+ * @brief RAII-style access to NoF ssd segment mutex for thread-safe segment
+ * operations
+ */
+class ScopedNoFSegmentAccess {
+   public:
+    /**
+     * @brief Acquires a lock on the segment mutex
+     * @param mutex Reference to the segment mutex
+     */
+    explicit ScopedNoFSegmentAccess(NoFSegmentManager* segment_manager,
+                                    std::shared_mutex& mutex)
+        : nof_segment_manager_(segment_manager), lock_(mutex) {}
+
+    /**
+     * @brief Mount a segment
+     */
+    ErrorCode MountSegment(const NoFSegment& segment, const UUID& client_id);
+
+    /**
+     * @brief Re-mount a segment. To avoid infinite remount trying, only the
+     * errors that may be solved by subsequent remount tryings are considered as
+     * errors. When encounters unsolvable errors, the segment will not be
+     * mounted while the return value will be OK.
+     */
+    ErrorCode ReMountSegment(const std::vector<NoFSegment>& segments,
+                             const UUID& client_id);
+
+    /**
+     * @brief Prepare to unmount a segment by deleting its allocator
+     */
+    ErrorCode PrepareUnmountSegment(const UUID& segment_id,
+                                    size_t& metrics_dec_capacity);
+
+    /**
+     * @brief Deleting the segment to complete the unmounting operation
+     */
+    ErrorCode CommitUnmountSegment(const UUID& segment_id,
+                                   const UUID& client_id,
+                                   const size_t& metrics_dec_capacity);
+
+    /**
+     * @brief Get all the segments of a client
+     */
+    ErrorCode GetClientSegments(const UUID& client_id,
+                                std::vector<NoFSegment>& segments) const;
+
+    /**
+     * @brief Get all mounted NoF segments with owner information
+     */
+    ErrorCode GetMountedSegments(
+        std::vector<MountedNoFSegmentSnapshot>& segments) const;
+
+    /**
+     * @brief Get the names of all the segments
+     */
+    ErrorCode GetAllSegments(std::vector<std::string>& all_segments);
+
+    /**
+     * @brief Get the segment by name. If there are multiple segments with the
+     * same name, return the first one.
+     */
+    ErrorCode QuerySegments(const std::string& segment, size_t& used,
+                            size_t& capacity);
+
+   private:
+    NoFSegmentManager* nof_segment_manager_;
+    std::unique_lock<std::shared_mutex> lock_;
+};
+
 /**
  * @brief RAII-style access to allocators for thread-safe allocator usage
  */
@@ -201,7 +328,7 @@ class ScopedAllocatorAccess {
  * @brief RAII-style access to LocalDiskOffloadingQueues for thread-safe
  * LocalDiskOffloadingQueue usage
  */
-class ScopedLocalDiskSegmentAccess {
+class ScopedLocalDiskSegmentAccess : public SsdMetricsProvider {
    public:
     explicit ScopedLocalDiskSegmentAccess(
         std::unordered_map<std::string, UUID>& client_by_name,
@@ -222,6 +349,9 @@ class ScopedLocalDiskSegmentAccess {
         return client_local_disk_segment_;
     }
 
+    int64_t getSsdTotalCapacity(const std::string& segment_name) const override;
+    int64_t getSsdUsedBytes(const std::string& segment_name) const override;
+
    private:
     const std::unordered_map<std::string, UUID>&
         client_by_name_;  // segment name -> client_id
@@ -336,4 +466,76 @@ class SegmentManager {
     friend class SegmentSerializer;  // for fork serialize
 };
 
+class NoFSegmentManager {
+   public:
+    /**
+     * @brief Constructor for SegmentManager
+     * @param memory_allocator Type of buffer allocator to use for new segments
+     */
+    explicit NoFSegmentManager(
+        BufferAllocatorType memory_allocator = BufferAllocatorType::CACHELIB)
+        : memory_allocator_(memory_allocator) {}
+
+    /**
+     * @brief Get RAII-style access to segment management operations
+     * @return ScopedSegmentAccess object that holds the lock
+     */
+    ScopedNoFSegmentAccess getNoFSegmentAccess() {
+        return ScopedNoFSegmentAccess(this, segment_mutex_);
+    }
+
+    /**
+     * @brief Get RAII-style access to use allocators
+     * @return ScopedAllocatorAccess object that holds the lock
+     */
+    ScopedAllocatorAccess getAllocatorAccess() {
+        return ScopedAllocatorAccess(allocator_manager_, segment_mutex_);
+    }
+
+    /**
+     * @brief Get the number of mounted NoF segments
+     * @return the mounted NoF segment count
+     */
+    int getMountedSegmentCount() const {
+        std::shared_lock<std::shared_mutex> lock(segment_mutex_);
+        return mounted_segments_.size();
+    }
+
+    void GetMountedSegmentsSnapshot(
+        std::vector<MountedNoFSegmentSnapshot>& segments) const;
+
+    tl::expected<std::vector<NoFSegmentOwnerInfo>, ErrorCode> GetSegmentsByName(
+        const std::string& segment_name) const {
+        std::shared_lock<std::shared_mutex> lock(segment_mutex_);
+        std::vector<NoFSegmentOwnerInfo> result;
+        for (const auto& [segment_id, mounted_segment] : mounted_segments_) {
+            if (mounted_segment.segment.name == segment_name) {
+                result.emplace_back(segment_id, mounted_segment.client_id);
+            }
+        }
+        if (result.empty()) {
+            return tl::make_unexpected(ErrorCode::SEGMENT_NOT_FOUND);
+        }
+        return result;
+    }
+
+   private:
+    mutable std::shared_mutex segment_mutex_;
+    std::shared_ptr<AllocationStrategy> allocation_strategy_;
+    const BufferAllocatorType
+        memory_allocator_;  // Type of buffer allocator to use
+    // allocator_manager_ only contains allocators whose segment status is OK.
+    AllocatorManager allocator_manager_;
+    std::unordered_map<UUID, MountedNoFSegment, boost::hash<UUID>>
+        mounted_segments_;  // segment_id -> mounted segment
+    std::unordered_map<UUID, std::vector<UUID>, boost::hash<UUID>>
+        client_segments_;  // client_id -> segment_ids
+
+    std::unordered_map<std::string, UUID>
+        client_by_name_;  // segment name -> client_id
+
+    friend class ScopedNoFSegmentAccess;
+    friend class SegmentTest;
+};
+
 }  // namespace mooncake
diff --git a/mooncake-store/include/spdk/spdk_wrapper.h b/mooncake-store/include/spdk/spdk_wrapper.h
new file mode 100644
index 00000000..3dbb621f
--- /dev/null
+++ b/mooncake-store/include/spdk/spdk_wrapper.h
@@ -0,0 +1,109 @@
+#pragma once
+#include <atomic>
+#include <cstddef>
+#include <cstdint>
+#include <map>
+#include <memory>
+#include <mutex>
+#include <stack>
+#include <string>
+#include <vector>
+#include <spdk/env.h>
+#include <spdk/nvme.h>
+
+namespace mooncake {
+
+#define INVALID_BLOCK_SIZE 0xFFFFFFFF
+
+constexpr int kSpdkNofOpRead = 0;
+constexpr int kSpdkNofOpWrite = 1;
+constexpr int kSpdkNofOpNum = 2;
+
+struct nof_seg_handle;
+struct tr_info;
+struct ctrlr_info;
+
+class SpdkWrapper {
+   public:
+    SpdkWrapper(const SpdkWrapper &) = delete;
+    SpdkWrapper &operator=(const SpdkWrapper &) = delete;
+
+    static SpdkWrapper &GetInstance();
+
+    bool InitializeEnv();
+
+    void Cleanup();
+
+    void *Alloc(size_t size, size_t align, int socket_id = -1);
+
+    void Free(void *ptr);
+
+    int64_t NvmePollProcessCompletion(nof_seg_handle *seg,
+                                      uint32_t complete_per_seg);
+
+    /** @brief Open a NoF segment. */
+    nof_seg_handle *OpenNofSegment(const std::string &tr_str);
+
+    uint32_t GetBlockSize(const nof_seg_handle *seg_handle);
+
+    int SubmitRequest(const nof_seg_handle *seg_handle, void *ptr, uint64_t lba,
+                      uint32_t lba_count, int op, spdk_nvme_cmd_cb cb_fn,
+                      void *cb_ctx);
+
+    bool ProbeNofSegment(const std::string &tr_str, uint32_t timeout_ms,
+                         std::string *error_reason = nullptr);
+
+   private:
+    struct ProbeBuffer {
+        void *ptr{nullptr};
+        uint32_t size{0};
+
+        ProbeBuffer() = default;
+        ProbeBuffer(const ProbeBuffer &) = delete;
+        ProbeBuffer &operator=(const ProbeBuffer &) = delete;
+        ProbeBuffer(ProbeBuffer &&) = delete;
+        ProbeBuffer &operator=(ProbeBuffer &&) = delete;
+    };
+
+    struct ProbeRequestContext {
+        std::atomic<bool> done{false};
+        std::atomic<bool> success{false};
+        std::mutex error_mutex;
+        std::string error_reason;
+        SpdkWrapper *owner{nullptr};
+
+        void Reset(SpdkWrapper *wrapper) {
+            std::lock_guard<std::mutex> lock(error_mutex);
+            owner = wrapper;
+            done.store(false, std::memory_order_release);
+            success.store(false, std::memory_order_release);
+            error_reason.clear();
+        }
+    };
+
+    explicit SpdkWrapper();
+    ~SpdkWrapper();
+
+    int ParseTransPortStr(const std::string &tr_str, tr_info *info);
+    int ConnectController(const struct spdk_nvme_transport_id *trid,
+                          ctrlr_info *info);
+    ProbeBuffer *GetOrCreateProbeBuffer(const std::string &tr_str,
+                                        uint32_t block_size,
+                                        std::string *error_reason);
+    ProbeRequestContext *AcquireProbeRequestContext();
+    void RecycleProbeRequestContext(ProbeRequestContext *ctx);
+    void ReplenishProbeRequestContextPoolLocked(size_t count);
+    static void ProbeReadComplete(void *ctx, const struct spdk_nvme_cpl *cpl);
+
+    std::atomic<bool> initialized{false};
+    std::mutex init_mutex;
+    std::map<std::string, std::unique_ptr<ctrlr_info>> connected_ctrlrs;
+    std::mutex ctrlrs_mutex;
+    std::map<std::string, std::unique_ptr<ProbeBuffer>> probe_buffers_;
+    std::mutex probe_buffers_mutex_;
+    std::vector<std::unique_ptr<ProbeRequestContext>> probe_request_contexts_;
+    std::stack<ProbeRequestContext *> probe_request_context_pool_;
+    std::mutex probe_request_context_pool_mutex_;
+};
+
+}  // namespace mooncake
diff --git a/mooncake-store/include/ssd_register_client.h b/mooncake-store/include/ssd_register_client.h
new file mode 100644
index 00000000..e8ae8cd0
--- /dev/null
+++ b/mooncake-store/include/ssd_register_client.h
@@ -0,0 +1,39 @@
+#ifndef MOONCAKE_SSD_REGISTER_CLIENT_H
+#define MOONCAKE_SSD_REGISTER_CLIENT_H
+
+#include "master_client.h"
+#include "types.h"
+#include <string>
+
+namespace mooncake {
+#define OPERATION_OK 0
+#define OPERATION_FAILED -1
+
+class NoFRegisterClient {
+   public:
+    NoFRegisterClient();
+    ~NoFRegisterClient();
+
+    int set_register(const std::string &nqn, size_t nsid,
+                     const std::string &traddr, size_t trsvcid, uintptr_t base,
+                     size_t size, const std::string &master_server_addr);
+
+    /**
+     * @brief Unregister a NoF SSD segment by its te_endpoint
+     * @param nqn NQN of the SSD
+     * @param nsid Namespace ID
+     * @param traddr Transport address
+     * @param trsvcid Transport service ID
+     * @param master_server_addr Master server address
+     * @return int OPERATION_OK or OPERATION_FAILED
+     */
+    int set_unregister_by_endpoint(const std::string &nqn, size_t nsid,
+                                   const std::string &traddr, size_t trsvcid,
+                                   const std::string &master_server_addr);
+
+   private:
+    MasterClient master_client_;
+};
+}  // namespace mooncake
+
+#endif  // MOONCAKE_SSD_REGISTER_CLIENT_H
diff --git a/mooncake-store/include/storage/distributed/distributed_storage_backend.h b/mooncake-store/include/storage/distributed/distributed_storage_backend.h
new file mode 100644
index 00000000..d4c935bb
--- /dev/null
+++ b/mooncake-store/include/storage/distributed/distributed_storage_backend.h
@@ -0,0 +1,67 @@
+#pragma once
+
+#include <memory>
+
+#include "fs_adapter.h"
+#include "storage_backend.h"
+
+namespace mooncake {
+
+struct DistributedStorageConfig {
+    std::string fsdir = "distributed_dir";
+    std::string fs_adapter_type = "hf3fs";
+    bool enable_health_check = false;
+    int hash_bucket_count = 256;
+
+    bool Validate() const;
+    static DistributedStorageConfig FromEnvironment();
+};
+
+/**
+ * @brief Distributed filesystem storage backend.
+ *
+ * Implements StorageBackendInterface, delegating I/O to a FileSystemAdapter.
+ * Does not handle eviction (DFS manages its own space).
+ */
+class DistributedStorageBackend : public StorageBackendInterface {
+   public:
+    DistributedStorageBackend(
+        const FileStorageConfig& file_storage_config,
+        const DistributedStorageConfig& distributed_config,
+        std::unique_ptr<FileSystemAdapter> fs_adapter);
+
+    tl::expected<void, ErrorCode> Init() override;
+
+    tl::expected<int64_t, ErrorCode> BatchOffload(
+        const std::unordered_map<std::string, std::vector<Slice>>& batch_object,
+        std::function<ErrorCode(const std::vector<std::string>& keys,
+                                std::vector<StorageObjectMetadata>& metadatas)>
+            complete_handler,
+        std::function<void(const std::vector<std::string>& evicted_keys)>
+            eviction_handler = nullptr) override;
+
+    tl::expected<void, ErrorCode> BatchLoad(
+        std::unordered_map<std::string, Slice>& batched_slices) override;
+
+    tl::expected<bool, ErrorCode> IsExist(const std::string& key) override;
+
+    tl::expected<bool, ErrorCode> IsEnableOffloading() override;
+
+    tl::expected<void, ErrorCode> ScanMeta(
+        const std::function<ErrorCode(
+            const std::vector<std::string>& keys,
+            std::vector<StorageObjectMetadata>& metadatas)>& handler) override;
+
+   private:
+    std::string GetObjectPath(const std::string& key) const;
+    static std::string EscapeFilename(const std::string& key);
+    static std::string UnescapeFilename(const std::string& name);
+
+    std::unique_ptr<FileSystemAdapter> fs_adapter_;
+    DistributedStorageConfig distributed_config_;
+    std::string root_dir_;
+    int hash_bucket_count_;
+    bool initialized_ = false;
+};
+
+}  // namespace mooncake
diff --git a/mooncake-store/include/storage/distributed/fs_adapter.h b/mooncake-store/include/storage/distributed/fs_adapter.h
new file mode 100644
index 00000000..defda616
--- /dev/null
+++ b/mooncake-store/include/storage/distributed/fs_adapter.h
@@ -0,0 +1,112 @@
+#pragma once
+
+#include <sys/stat.h>
+#include <sys/uio.h>
+
+#include <span>
+#include <string>
+#include <vector>
+
+#include "types.h"
+
+namespace mooncake {
+
+struct FileInfo {
+    std::string name;
+    size_t size;
+};
+
+/**
+ * @brief Abstract interface for distributed filesystem adapters.
+ *
+ * Encapsulates file-level I/O differences across DFS implementations
+ * (3FS, CephFS, JuiceFS, etc.). DistributedStorageBackend depends only
+ * on this interface and is unaware of the concrete DFS.
+ */
+class FileSystemAdapter {
+   public:
+    virtual ~FileSystemAdapter() = default;
+
+    // === File I/O ===
+
+    virtual tl::expected<size_t, ErrorCode> WriteFile(
+        const std::string& path, std::span<const char> data) = 0;
+
+    // Read file into a pre-allocated buffer (zero-copy into Slice.ptr)
+    virtual tl::expected<size_t, ErrorCode> ReadFile(const std::string& path,
+                                                     void* buf, size_t len) = 0;
+
+    virtual tl::expected<size_t, ErrorCode> VectorWriteFile(
+        const std::string& path, const iovec* iov, int iovcnt,
+        off_t offset) = 0;
+
+    virtual tl::expected<size_t, ErrorCode> VectorReadFile(
+        const std::string& path, const iovec* iov, int iovcnt,
+        off_t offset) = 0;
+
+    // === File management ===
+
+    virtual tl::expected<void, ErrorCode> DeleteFile(
+        const std::string& path) = 0;
+
+    virtual tl::expected<bool, ErrorCode> FileExists(
+        const std::string& path) = 0;
+
+    virtual tl::expected<std::vector<std::string>, ErrorCode> ListFiles(
+        const std::string& dir) = 0;
+
+    virtual tl::expected<size_t, ErrorCode> GetFileSize(
+        const std::string& path) {
+        struct stat st;
+        if (::stat(path.c_str(), &st) != 0) {
+            if (errno == ENOENT) {
+                return tl::make_unexpected(ErrorCode::FILE_NOT_FOUND);
+            }
+            return tl::make_unexpected(ErrorCode::FILE_READ_FAIL);
+        }
+        return static_cast<size_t>(st.st_size);
+    }
+
+    // === Batch operations (default implementations, adapters may override) ===
+
+    virtual tl::expected<void, ErrorCode> DeleteFiles(
+        const std::vector<std::string>& paths) {
+        for (const auto& path : paths) {
+            auto result = DeleteFile(path);
+            if (!result) return result;
+        }
+        return {};
+    }
+
+    virtual tl::expected<std::vector<FileInfo>, ErrorCode> ListFilesWithInfo(
+        const std::string& dir) {
+        auto files = ListFiles(dir);
+        if (!files) return tl::make_unexpected(files.error());
+
+        std::vector<FileInfo> result;
+        result.reserve(files->size());
+        for (const auto& name : *files) {
+            std::string full_path = dir + "/" + name;
+            auto size = GetFileSize(full_path);
+            if (size) {
+                result.push_back({name, *size});
+            } else if (size.error() != ErrorCode::FILE_NOT_FOUND) {
+                return tl::make_unexpected(size.error());
+            }
+        }
+        return result;
+    }
+
+    // === Lifecycle ===
+
+    virtual tl::expected<void, ErrorCode> Init(
+        const std::string& mount_path) = 0;
+
+    virtual tl::expected<void, ErrorCode> Shutdown() = 0;
+
+    // === Identity ===
+
+    virtual const char* GetName() const = 0;
+};
+
+}  // namespace mooncake
diff --git a/mooncake-store/include/storage/distributed/hf3fs_adapter.h b/mooncake-store/include/storage/distributed/hf3fs_adapter.h
new file mode 100644
index 00000000..4b0fa05f
--- /dev/null
+++ b/mooncake-store/include/storage/distributed/hf3fs_adapter.h
@@ -0,0 +1,51 @@
+#pragma once
+
+#include <memory>
+
+#include "storage/distributed/fs_adapter.h"
+
+namespace mooncake {
+
+// Forward declaration: avoid including hf3fs/hf3fs.h in the header.
+// Full type is only used in hf3fs_adapter.cpp.
+class USRBIOResourceManager;
+
+class Hf3fsAdapter : public FileSystemAdapter {
+   public:
+    Hf3fsAdapter();
+    ~Hf3fsAdapter() override;  // defined in .cpp (needs complete type for
+                               // unique_ptr deleter)
+
+    tl::expected<size_t, ErrorCode> WriteFile(
+        const std::string& path, std::span<const char> data) override;
+
+    tl::expected<size_t, ErrorCode> ReadFile(const std::string& path, void* buf,
+                                             size_t len) override;
+
+    tl::expected<size_t, ErrorCode> VectorWriteFile(const std::string& path,
+                                                    const iovec* iov,
+                                                    int iovcnt,
+                                                    off_t offset) override;
+
+    tl::expected<size_t, ErrorCode> VectorReadFile(const std::string& path,
+                                                   const iovec* iov, int iovcnt,
+                                                   off_t offset) override;
+
+    tl::expected<void, ErrorCode> DeleteFile(const std::string& path) override;
+
+    tl::expected<bool, ErrorCode> FileExists(const std::string& path) override;
+
+    tl::expected<std::vector<std::string>, ErrorCode> ListFiles(
+        const std::string& dir) override;
+
+    tl::expected<void, ErrorCode> Init(const std::string& mount_path) override;
+
+    tl::expected<void, ErrorCode> Shutdown() override;
+
+    const char* GetName() const override { return "hf3fs"; }
+
+   private:
+    std::unique_ptr<USRBIOResourceManager> resource_manager_;
+};
+
+}  // namespace mooncake
diff --git a/mooncake-store/include/storage_backend.h b/mooncake-store/include/storage_backend.h
index a686fd97..5fd071b5 100644
--- a/mooncake-store/include/storage_backend.h
+++ b/mooncake-store/include/storage_backend.h
@@ -155,7 +155,12 @@ struct OffloadMetadata {
 
 enum class FileMode { Read, Write };
 
-enum class StorageBackendType { kFilePerKey, kBucket, kOffsetAllocator };
+enum class StorageBackendType {
+    kFilePerKey,
+    kBucket,
+    kOffsetAllocator,
+    kDistributed
+};
 
 static constexpr size_t kKB = 1024;
 static constexpr size_t kMB = kKB * 1024;
@@ -267,6 +272,11 @@ class StorageBackendInterface {
             const std::vector<std::string>& keys,
             std::vector<StorageObjectMetadata>& metadatas)>& handler) = 0;
 
+    // Reset internal scan iterator so that the next ScanMeta() call
+    // starts from the beginning.  Required for backends that use
+    // cursor-based iteration (e.g. BucketStorageBackend).
+    virtual void ResetScanIterator() {}
+
     // Test-only: Set predicate to force failures for specific keys in
     // BatchOffload. Default implementation does nothing (no failures injected).
     // Concrete backends can override to provide test failure injection.
@@ -289,40 +299,25 @@ class StorageBackendInterface {
  */
 class StorageBackend {
    public:
-/**
- * @brief Constructs a new StorageBackend instance
- * @param root_dir Root directory path for object storage
- * @param fsdir  subdirectory name
- * @note Directory existence is not checked in constructor
- */
-#ifdef USE_3FS
-    explicit StorageBackend(const std::string& root_dir,
-                            const std::string& fsdir, bool is_3fs_dir,
-                            bool enable_eviction = true)
-        : root_dir_(root_dir),
-          fsdir_(fsdir),
-          is_3fs_dir_(is_3fs_dir),
-          enable_eviction_(enable_eviction) {
-        resource_manager_ = std::make_unique<USRBIOResourceManager>();
-        Hf3fsConfig config;
-        config.mount_root = root_dir;
-        resource_manager_->setDefaultParams(config);
-    }
-#else
+    /**
+     * @brief Constructs a new StorageBackend instance
+     * @param root_dir Root directory path for object storage
+     * @param fsdir  subdirectory name
+     * @note Directory existence is not checked in constructor
+     */
     explicit StorageBackend(const std::string& root_dir,
                             const std::string& fsdir,
                             bool enable_eviction = true)
         : root_dir_(root_dir),
           fsdir_(fsdir),
           enable_eviction_(enable_eviction) {}
-#endif
 
     /**
      * @brief Factory method to create a StorageBackend instance
      * @param root_dir Root directory path for object storage
      * @param fsdir  subdirectory name
      * @param enable_eviction Whether to enable disk eviction feature (default:
-     * true) Note: Eviction is automatically disabled for 3FS mode
+     * true) Note: Eviction is controlled by the enable_eviction parameter
      * @return shared_ptr to new instance or nullptr if directory is invalid
      *
      * Performs validation of the root directory before creating the instance:
@@ -347,15 +342,8 @@ class StorageBackend {
         fs::path root_path(root_dir);
 
         std::string real_fsdir = "moon_" + fsdir;
-#ifdef USE_3FS
-        bool is_3fs_dir = fs::exists(root_path / "3fs-virt") &&
-                          fs::is_directory(root_path / "3fs-virt");
-        return std::make_shared<StorageBackend>(root_dir, real_fsdir,
-                                                is_3fs_dir, enable_eviction);
-#else
         return std::make_shared<StorageBackend>(root_dir, real_fsdir,
                                                 enable_eviction);
-#endif
     }
 
     /**
@@ -471,12 +459,6 @@ class StorageBackend {
         true};  // User-configurable flag to enable/disable eviction
     bool use_uring_{false};  // Use io_uring for file I/O
 
-#ifdef USE_3FS
-    bool is_3fs_dir_{false};  // Flag to indicate if the storage is using 3FS
-                              // directory structure
-    std::unique_ptr<USRBIOResourceManager> resource_manager_;
-#endif
-
    private:
     // File write queue for disk eviction - tracks files in FIFO order
     std::list<FileRecord> file_write_queue_;
@@ -578,8 +560,7 @@ class StorageBackend {
 
     /**
      * @brief Checks if disk eviction is enabled for this storage backend.
-     * @return true if eviction is enabled (local mode), false if disabled (3FS
-     * mode).
+     * @return true if eviction is enabled, false otherwise.
      */
     bool IsEvictionEnabled() const;
 
@@ -781,6 +762,11 @@ class BucketStorageBackend : public StorageBackendInterface {
             const std::vector<std::string>& keys,
             std::vector<StorageObjectMetadata>& metadatas)>& handler) override;
 
+    void ResetScanIterator() override {
+        MutexLocker locker(&iterator_mutex_);
+        next_bucket_ = -1;
+    }
+
     /**
      * @brief Checks whether the backend is allowed to continue offloading.
      * @return tl::expected<bool, ErrorCode>
diff --git a/mooncake-store/include/task_manager.h b/mooncake-store/include/task_manager.h
index d12152ed..8f4578ce 100644
--- a/mooncake-store/include/task_manager.h
+++ b/mooncake-store/include/task_manager.h
@@ -92,18 +92,20 @@ struct Task {
 };
 
 struct ReplicaCopyPayload {
+    std::string tenant_id = "default";
     std::string key;
     std::string source;
     std::vector<std::string> targets;
 };
-YLT_REFL(ReplicaCopyPayload, key, source, targets);
+YLT_REFL(ReplicaCopyPayload, tenant_id, key, source, targets);
 
 struct ReplicaMovePayload {
+    std::string tenant_id = "default";
     std::string key;
     std::string source;
     std::string target;
 };
-YLT_REFL(ReplicaMovePayload, key, source, target);
+YLT_REFL(ReplicaMovePayload, tenant_id, key, source, target);
 
 template <TaskType T>
 struct TaskPayloadTraits;
diff --git a/mooncake-store/include/tenant_quota.h b/mooncake-store/include/tenant_quota.h
new file mode 100644
index 00000000..bd189db2
--- /dev/null
+++ b/mooncake-store/include/tenant_quota.h
@@ -0,0 +1,80 @@
+#pragma once
+
+#include <cstdint>
+#include <map>
+#include <optional>
+#include <string>
+#include <vector>
+
+#include <ylt/util/tl/expected.hpp>
+
+namespace mooncake {
+
+struct TenantQuotaState {
+    uint64_t requested_quota_bytes = 0;
+    uint64_t effective_quota_bytes = 0;
+    uint64_t used_bytes = 0;
+    uint64_t reserved_bytes = 0;
+    uint64_t committed_count = 0;
+    uint64_t metadata_object_count = 0;
+    bool has_explicit_policy = false;
+    bool over_quota = false;
+};
+
+struct TenantQuotaSnapshot {
+    std::string tenant_id;
+    uint64_t requested_quota_bytes = 0;
+    uint64_t effective_quota_bytes = 0;
+    uint64_t used_bytes = 0;
+    uint64_t reserved_bytes = 0;
+    uint64_t committed_count = 0;
+    uint64_t metadata_object_count = 0;
+    bool has_explicit_policy = false;
+    bool over_quota = false;
+};
+
+struct TenantQuotaAssignment {
+    std::string tenant_id;
+    uint64_t effective_quota_bytes = 0;
+};
+
+enum class TenantQuotaError {
+    kQuotaExceeded,
+    kInvalidArgument,
+    kAccountingMismatch,
+};
+
+using TenantQuotaResult = tl::expected<void, TenantQuotaError>;
+
+std::vector<TenantQuotaAssignment> BuildEffectiveQuotaAssignments(
+    const std::map<std::string, TenantQuotaState>& tenants,
+    uint64_t allocatable_capacity_bytes);
+
+class TenantQuotaTable {
+   public:
+    TenantQuotaResult UpsertTenantPolicy(std::string tenant_id,
+                                         uint64_t requested_quota_bytes);
+    void EraseTenantPolicy(std::string tenant_id);
+
+    void RecomputeEffectiveQuotas(uint64_t allocatable_capacity_bytes);
+
+    std::optional<TenantQuotaSnapshot> GetTenantSnapshot(
+        std::string tenant_id) const;
+    std::vector<TenantQuotaSnapshot> ListTenantSnapshots() const;
+
+    TenantQuotaResult Reserve(std::string tenant_id, uint64_t bytes);
+    TenantQuotaResult Commit(std::string tenant_id, uint64_t bytes);
+    TenantQuotaResult Abort(std::string tenant_id, uint64_t bytes);
+    TenantQuotaResult Release(std::string tenant_id, uint64_t bytes);
+    TenantQuotaResult ReleasePartial(std::string tenant_id, uint64_t bytes);
+
+   private:
+    TenantQuotaState& GetOrCreateState(const std::string& tenant_id);
+    TenantQuotaSnapshot MakeSnapshot(const std::string& tenant_id,
+                                     const TenantQuotaState& state) const;
+    void RefreshOverQuota(TenantQuotaState* state) const;
+
+    std::map<std::string, TenantQuotaState> tenants_;
+};
+
+}  // namespace mooncake
diff --git a/mooncake-store/include/tenant_quota_policy_store.h b/mooncake-store/include/tenant_quota_policy_store.h
new file mode 100644
index 00000000..30f3d85d
--- /dev/null
+++ b/mooncake-store/include/tenant_quota_policy_store.h
@@ -0,0 +1,51 @@
+#pragma once
+
+#include <cstdint>
+#include <map>
+#include <memory>
+#include <mutex>
+#include <string>
+
+#include <ylt/util/tl/expected.hpp>
+
+namespace mooncake {
+
+struct TenantQuotaPolicySnapshot {
+    std::map<std::string, uint64_t> tenant_quotas;
+};
+
+tl::expected<uint64_t, std::string> ParseTenantQuotaBytes(
+    const std::string& value);
+
+tl::expected<TenantQuotaPolicySnapshot, std::string> ParseTenantQuotaPolicyYaml(
+    const std::string& yaml);
+
+std::string FormatTenantQuotaPolicyYaml(
+    const TenantQuotaPolicySnapshot& snapshot);
+
+class TenantQuotaPolicyStore {
+   public:
+    virtual ~TenantQuotaPolicyStore() = default;
+
+    virtual tl::expected<TenantQuotaPolicySnapshot, std::string> Load() = 0;
+    virtual tl::expected<void, std::string> Save(
+        const TenantQuotaPolicySnapshot& snapshot) = 0;
+};
+
+class YamlTenantQuotaPolicyStore final : public TenantQuotaPolicyStore {
+   public:
+    explicit YamlTenantQuotaPolicyStore(std::string path);
+
+    tl::expected<TenantQuotaPolicySnapshot, std::string> Load() override;
+    tl::expected<void, std::string> Save(
+        const TenantQuotaPolicySnapshot& snapshot) override;
+
+   private:
+    std::string path_;
+    std::mutex mutex_;
+};
+
+tl::expected<std::unique_ptr<TenantQuotaPolicyStore>, std::string>
+CreateTenantQuotaPolicyStore(const std::string& type, const std::string& uri);
+
+}  // namespace mooncake
diff --git a/mooncake-store/include/transfer_task.h b/mooncake-store/include/transfer_task.h
index 75a96ac0..63b06191 100644
--- a/mooncake-store/include/transfer_task.h
+++ b/mooncake-store/include/transfer_task.h
@@ -4,11 +4,13 @@
 #include <condition_variable>
 #include <cstdint>
 #include <cstring>
+#include <map>
 #include <memory>
 #include <mutex>
 #include <optional>
 #include <ostream>
 #include <queue>
+#include <stack>
 #include <string>
 #include <thread>
 #include <vector>
@@ -18,6 +20,9 @@
 #include "replica.h"
 #include "storage_backend.h"
 #include "client_metric.h"
+#ifdef USE_NOF
+#include "spdk/spdk_wrapper.h"
+#endif
 
 namespace mooncake {
 
@@ -28,7 +33,8 @@ enum class TransferStrategy {
     LOCAL_MEMCPY = 0,     // Local memory copy using memcpy
     TRANSFER_ENGINE = 1,  // Remote transfer using transfer engine
     FILE_READ = 2,        // File read operation
-    EMPTY = 3
+    EMPTY = 3,
+    SPDK_NVMF = 4  // Spdk nvmf operation
 };
 
 /**
@@ -41,8 +47,12 @@ inline std::ostream& operator<<(std::ostream& os,
             return os << "LOCAL_MEMCPY";
         case TransferStrategy::TRANSFER_ENGINE:
             return os << "TRANSFER_ENGINE";
+        case TransferStrategy::SPDK_NVMF:
+            return os << "SPDK_NVMF";
         case TransferStrategy::FILE_READ:
             return os << "FILE_READ";
+        case TransferStrategy::EMPTY:
+            return os << "EMPTY";
         default:
             return os << "UNKNOWN";
     }
@@ -141,6 +151,35 @@ class MemcpyOperationState : public OperationState {
     }
 };
 
+/**
+ * @brief Operation state for local memcpy transfers
+ */
+class SpdkNofOperationState : public OperationState {
+   public:
+    bool is_completed() override {
+        std::lock_guard<std::mutex> lock(mutex_);
+        return result_.has_value();
+    }
+
+    void set_completed(ErrorCode error_code) {
+        {
+            std::lock_guard<std::mutex> lock(mutex_);
+            assert(!result_.has_value());
+            result_.emplace(error_code);
+        }
+        cv_.notify_all();
+    }
+
+    void wait_for_completion() override {
+        std::unique_lock<std::mutex> lock(mutex_);
+        cv_.wait(lock, [this] { return result_.has_value(); });
+    }
+
+    TransferStrategy get_strategy() const override {
+        return TransferStrategy::SPDK_NVMF;
+    }
+};
+
 class FilereadOperationState : public OperationState {
    public:
     bool is_completed() override {
@@ -307,6 +346,127 @@ class MemcpyWorkerPool {
     std::atomic<bool> shutdown_;
 };
 
+#ifdef USE_NOF
+// struct SpdkNofSubTask;
+struct SpdkNofQos;
+
+/**
+ * @brief Spdk nvmf operation descriptor
+ */
+struct SpdkNofTask {
+    nof_seg_handle* seg_handle;
+    void* ptr;
+    uint64_t lba;
+    uint32_t lba_count;
+    int remaining_lba;
+    int outstanding_sub_io;
+    int op;   // kSpdkNofOpRead or kSpdkNofOpWrite
+    int idx;  // subop idx
+    bool failed;
+    bool on_chain;
+    std::shared_ptr<SpdkNofOperationState> state;
+    int64_t* io_count;
+    SpdkNofQos* nof_qos;
+    SpdkNofTask* nxt;
+
+    SpdkNofTask(nof_seg_handle* handle, void* buf, uint64_t off, uint32_t len,
+                int op_code, std::shared_ptr<SpdkNofOperationState> s)
+        : seg_handle(handle),
+          ptr(buf),
+          lba(off),
+          lba_count(len),
+          remaining_lba(lba_count),
+          outstanding_sub_io(0),
+          op(op_code),
+          idx(0),
+          failed(false),
+          on_chain(false),
+          state(std::move(s)),
+          io_count(nullptr),
+          nof_qos(nullptr),
+          nxt(nullptr) {}
+};
+
+struct SpdkNofSubTask {
+    SpdkNofTask* task;
+    int submit_lba_count;
+    std::stack<SpdkNofSubTask*>* sub_task_pool;
+};
+
+constexpr int kDefaultSpdkNofSubmitChunkBytes = (1 << 17);    // 128k
+constexpr int kDefaultSpdkNofInflightBytesLimit = (1 << 25);  // 32M
+struct SpdkNofQos {
+    int inflight_blocks[kSpdkNofOpNum];
+    int blocks_per_chunk;
+    int inflight_blocks_limit;
+    SpdkNofTask* head[kSpdkNofOpNum];
+    SpdkNofTask* tail[kSpdkNofOpNum];
+
+    explicit SpdkNofQos(uint32_t block_size);
+
+    bool Empty() const {
+        return (head[kSpdkNofOpRead] == nullptr &&
+                head[kSpdkNofOpWrite] == nullptr);
+    }
+
+    void PushTask(SpdkNofTask* task) {
+        int op = task->op;
+        if (head[op] == nullptr) {
+            head[op] = task;
+            tail[op] = task;
+        } else {
+            tail[op]->nxt = task;
+            tail[op] = task;
+        }
+    }
+
+    void PopTask(int op) {
+        if (head[op]) {
+            head[op] = head[op]->nxt;
+        }
+    }
+};
+
+/**
+ * @brief Thread pool for asynchronous spdk nvmf operations
+ *
+ * This class manages multiple worker thread that executes spdk nvmf operations
+ * asynchronously.
+ */
+constexpr int kDefaultSpdkNofWorkers = 4;
+class SpdkNofWorkerPool {
+   public:
+    explicit SpdkNofWorkerPool(int numa_socket_id = 0);
+    ~SpdkNofWorkerPool();
+
+    // Non-copyable, non-movable
+    SpdkNofWorkerPool(const SpdkNofWorkerPool&) = delete;
+    SpdkNofWorkerPool& operator=(const SpdkNofWorkerPool&) = delete;
+    SpdkNofWorkerPool(SpdkNofWorkerPool&&) = delete;
+    SpdkNofWorkerPool& operator=(SpdkNofWorkerPool&&) = delete;
+
+    /**
+     * @brief Submit a spdk nvmf task for async execution
+     * @param task The spdk nvmf task to execute
+     */
+    void submitTask(SpdkNofTask task);
+
+   private:
+    void workerThread(int work_idx);
+
+    int worker_count_;
+    int numa_socket_id_;
+    std::vector<std::thread> workers_;
+    std::unique_ptr<std::queue<SpdkNofTask>[]> task_queue_;
+    std::unique_ptr<std::mutex[]> queue_mutex_;
+    std::unique_ptr<std::condition_variable[]> queue_cv_;
+    std::atomic<bool> shutdown_;
+    std::mutex seg_mutex_;
+    int seg_num = 0;
+    std::map<nof_seg_handle*, int> seg_to_worker_;
+};
+#endif
+
 /**
  * @brief Fileread task for async execution
  */
@@ -370,7 +530,9 @@ class TransferSubmitter {
    public:
     explicit TransferSubmitter(TransferEngine& engine,
                                std::shared_ptr<StorageBackend>& backend,
-                               TransferMetric* transfer_metric = nullptr);
+                               const std::string& local_hostname,
+                               TransferMetric* transfer_metric = nullptr,
+                               int numa_socket_id = 0);
 
     /**
      * @brief Submit an asynchronous transfer operation
@@ -387,7 +549,8 @@ class TransferSubmitter {
      */
     std::optional<TransferFuture> submit(const Replica::Descriptor& replica,
                                          std::vector<Slice>& slices,
-                                         TransferRequest::OpCode op_code);
+                                         TransferRequest::OpCode op_code,
+                                         void* ptr = nullptr, size_t size = 0);
 
     /**
      * @brief Submit a range read: read [src_offset, src_offset+size) from
@@ -406,13 +569,34 @@ class TransferSubmitter {
         const std::string& transfer_engine_addr,
         const std::vector<std::string>& keys,
         const std::vector<uint64_t>& pointers,
-        const std::unordered_map<std::string, Slice>& batched_slices);
+        const std::unordered_map<std::string, std::vector<Slice>>&
+            batched_slices);
+
+    /**
+     * @brief Pure comparison helper: returns true iff both endpoints are
+     * non-empty and identical. Exposed for unit testing of the locality
+     * decision without instantiating a full TransferEngine.
+     *
+     * Two endpoints identify the same process only when their ip:port (or
+     * full hostname) match exactly; same-host different-process pairs share
+     * an IP but not a port and must NOT be treated as locally addressable.
+     */
+    static bool isSameProcessEndpoint(const std::string& handle_endpoint,
+                                      const std::string& local_endpoint);
 
    private:
     TransferEngine& engine_;
+    // Cached at construction: the local transport endpoint never changes for
+    // the lifetime of the TransferSubmitter, so we avoid calling
+    // engine_.getLocalIpAndPort() (which allocates a string) on every transfer.
+    const std::string local_endpoint_;
     std::unique_ptr<MemcpyWorkerPool> memcpy_pool_;
+#ifdef USE_NOF
+    std::unique_ptr<SpdkNofWorkerPool> spdk_nvmf_pool_;
+#endif
     std::unique_ptr<FilereadWorkerPool> fileread_pool_;
     bool memcpy_enabled_;
+    const std::string local_hostname_;
     TransferMetric* transfer_metric_;
 
     /**
@@ -440,6 +624,15 @@ class TransferSubmitter {
         const std::vector<Slice>& slices, const TransferRequest::OpCode op_code,
         uint64_t src_offset = 0);
 
+#ifdef USE_NOF
+    /**
+     * @brief Submit SPDK NVMe-oF operation asynchronously
+     */
+    std::optional<TransferFuture> submitSpdkNofOperation(
+        const AllocatedBuffer::Descriptor& handle, void* ptr, size_t size,
+        const TransferRequest::OpCode op_code);
+#endif
+
     /**
      * @brief Submit transfer engine operation asynchronously
      * @param src_offset Optional offset in source buffer (default 0)
diff --git a/mooncake-store/include/types.h b/mooncake-store/include/types.h
index 68dfb05c..b5836f8d 100644
--- a/mooncake-store/include/types.h
+++ b/mooncake-store/include/types.h
@@ -5,6 +5,7 @@
 #include <memory>
 #include <optional>
 #include <string>
+#include <string_view>
 #include <limits>
 #include <unordered_map>
 #include <utility>
@@ -88,8 +89,13 @@ static constexpr uint64_t DEFAULT_KV_SOFT_PIN_TTL_MS =
 static constexpr bool DEFAULT_ALLOW_EVICT_SOFT_PINNED_OBJECTS = true;
 static constexpr double DEFAULT_EVICTION_RATIO = 0.05;
 static constexpr double DEFAULT_EVICTION_HIGH_WATERMARK_RATIO = 0.95;
+static constexpr double DEFAULT_NOF_EVICTION_RATIO = 0.05;
+static constexpr double DEFAULT_NOF_EVICTION_HIGH_WATERMARK_RATIO = 0.95;
 static constexpr int64_t DEFAULT_MASTER_VIEW_LEASE_TTL_SEC = 5;  // in seconds
 static constexpr int64_t DEFAULT_CLIENT_LIVE_TTL_SEC = 10;       // in seconds
+static constexpr int64_t DEFAULT_NOF_HEARTBEAT_INTERVAL_SEC = 10;
+static constexpr uint32_t DEFAULT_NOF_HEARTBEAT_PROBE_TIMEOUT_MS = 1000;
+static constexpr uint32_t DEFAULT_NOF_HEARTBEAT_FAILURES_THRESHOLD = 3;
 static constexpr uint64_t DEFAULT_SNAPSHOT_INTERVAL_SEC =
     60 * 10;  // in seconds
 static constexpr uint64_t DEFAULT_SNAPSHOT_CHILD_TIMEOUT_SEC =
@@ -123,6 +129,48 @@ static constexpr uint64_t DEFAULT_PROCESSING_TASK_TIMEOUT_SEC =
     300;  // 0 to be no timeout
 static constexpr uint32_t DEFAULT_MAX_RETRY_ATTEMPTS = 10;
 
+/**
+ * @brief Data type classification for objects stored in Mooncake Store.
+ *
+ * This allows the store to track what kind of data each object holds,
+ * enabling future type-aware policies (eviction priority, replication
+ * strategies, etc.). Defaults to UNKNOWN for backward compatibility.
+ */
+enum class ObjectDataType : uint8_t {
+    UNKNOWN = 0,
+    KVCACHE = 1,
+    TENSOR = 2,
+    WEIGHT = 3,
+    SAMPLE = 4,
+    ACTIVATION = 5,
+    GRADIENT = 6,
+    OPTIMIZER_STATE = 7,
+    METADATA = 8,
+    GENERAL = 9,
+    HIDDEN_STATE = 10,
+    // 11-255 reserved for future types
+};
+
+inline std::ostream& operator<<(std::ostream& os,
+                                const ObjectDataType& type) noexcept {
+    static const std::unordered_map<ObjectDataType, std::string_view>
+        type_strings{{ObjectDataType::UNKNOWN, "UNKNOWN"},
+                     {ObjectDataType::KVCACHE, "KVCACHE"},
+                     {ObjectDataType::TENSOR, "TENSOR"},
+                     {ObjectDataType::WEIGHT, "WEIGHT"},
+                     {ObjectDataType::SAMPLE, "SAMPLE"},
+                     {ObjectDataType::ACTIVATION, "ACTIVATION"},
+                     {ObjectDataType::GRADIENT, "GRADIENT"},
+                     {ObjectDataType::OPTIMIZER_STATE, "OPTIMIZER_STATE"},
+                     {ObjectDataType::METADATA, "METADATA"},
+                     {ObjectDataType::GENERAL, "GENERAL"},
+                     {ObjectDataType::HIDDEN_STATE, "HIDDEN_STATE"}};
+
+    auto it = type_strings.find(type);
+    os << (it != type_strings.end() ? it->second : "UNKNOWN");
+    return os;
+}
+
 // Forward declarations
 class BufferAllocatorBase;
 class CachelibBufferAllocator;
@@ -169,6 +217,7 @@ constexpr const char* CONFIG_KEY_PROTOCOL = "protocol";
 constexpr const char* CONFIG_KEY_RDMA_DEVICES = "rdma_devices";
 constexpr const char* CONFIG_KEY_MASTER_SERVER_ADDR = "master_server_addr";
 constexpr const char* CONFIG_KEY_IPC_SOCKET_PATH = "ipc_socket_path";
+constexpr const char* CONFIG_KEY_TENANT_ID = "tenant_id";
 
 // Store client configuration defaults
 static constexpr size_t DEFAULT_GLOBAL_SEGMENT_SIZE = 1024 * 1024 * 16;  // 16MB
@@ -176,6 +225,67 @@ static constexpr size_t DEFAULT_LOCAL_BUFFER_SIZE = 1024 * 1024 * 16;    // 16MB
 constexpr const char* DEFAULT_PROTOCOL = "tcp";
 constexpr const char* DEFAULT_MASTER_SERVER_ADDR = "127.0.0.1:50051";
 
+inline std::string NormalizeTenantId(const std::string& tenant_id) {
+    return tenant_id.empty() ? "default" : tenant_id;
+}
+
+inline bool IsValidTenantId(const std::string& tenant_id) {
+    if (tenant_id.empty() || tenant_id.front() == '_') {
+        return false;
+    }
+    for (unsigned char c : tenant_id) {
+        if (c < 0x20 || c == 0x7f) {
+            return false;
+        }
+    }
+    return true;
+}
+
+inline std::string MakeTenantScopedStorageKey(const std::string& tenant_id,
+                                              const std::string& key) {
+    const auto normalized_tenant = NormalizeTenantId(tenant_id);
+    std::string scoped_key;
+    scoped_key.reserve(normalized_tenant.size() + key.size() + 1);
+    scoped_key.append(normalized_tenant);
+    scoped_key.push_back('\0');
+    scoped_key.append(key);
+    return scoped_key;
+}
+
+inline std::pair<std::string, std::string> ParseTenantScopedStorageKey(
+    const std::string& storage_key) {
+    const auto separator = storage_key.find('\0');
+    if (separator == std::string::npos) {
+        return {"default", storage_key};
+    }
+    return {NormalizeTenantId(storage_key.substr(0, separator)),
+            storage_key.substr(separator + 1)};
+}
+
+struct OffloadTaskItem {
+    std::string tenant_id;
+    std::string key;
+    int64_t size;
+
+    bool operator==(const OffloadTaskItem& other) const {
+        return tenant_id == other.tenant_id && key == other.key &&
+               size == other.size;
+    }
+};
+YLT_REFL(OffloadTaskItem, tenant_id, key, size);
+
+struct PromotionTaskItem {
+    std::string tenant_id;
+    std::string key;
+    int64_t size;
+
+    bool operator==(const PromotionTaskItem& other) const {
+        return tenant_id == other.tenant_id && key == other.key &&
+               size == other.size;
+    }
+};
+YLT_REFL(PromotionTaskItem, tenant_id, key, size);
+
 // Store client configuration validation limits
 static constexpr size_t MIN_SEGMENT_SIZE = 1024;                          // 1KB
 static constexpr size_t MAX_SEGMENT_SIZE = 1024ULL * 1024 * 1024 * 1024;  // 1TB
@@ -261,7 +371,8 @@ enum class ErrorCode : int32_t {
     TRANSFER_FAIL = -800,  ///< Transfer operation failed.
 
     // RPC errors (Range: -900 to -999)
-    RPC_FAIL = -900,  ///< RPC operation failed.
+    RPC_FAIL = -900,     ///< RPC operation failed.
+    RPC_TIMEOUT = -901,  ///< RPC call timed out (client-side deadline hit).
 
     // High availability errors (Range: -1000 to -1099)
     ETCD_OPERATION_ERROR = -1000,   ///< etcd operation failed.
@@ -270,6 +381,8 @@ enum class ErrorCode : int32_t {
     ETCD_CTX_CANCELLED = -1003,     ///< etcd context cancelled.
     OPLOG_ENTRY_NOT_FOUND =
         -1004,  ///< OpLog entry not found (backend-agnostic).
+    K8S_LEASE_OPERATION_ERROR = -1005,  ///< K8s Lease operation failed.
+    K8S_LEASE_NOT_FOUND = -1006,        ///< K8s Lease not found.
     UNAVAILABLE_IN_CURRENT_STATUS =
         -1010,  ///< Request cannot be done in current status.
     UNAVAILABLE_IN_CURRENT_MODE =
@@ -300,6 +413,17 @@ enum class ErrorCode : int32_t {
     TASK_PENDING_LIMIT_EXCEEDED =
         -1401,              ///< Total pending tasks exceed the limit.
     JOB_NOT_FOUND = -1402,  ///< Job not found.
+
+    // DFS errors (Range: -1600 to -1699)
+    DFS_NETWORK_TIMEOUT = -1600,      ///< DFS network timeout.
+    DFS_SERVICE_UNAVAILABLE = -1601,  ///< DFS service unavailable.
+    DFS_QUOTA_EXCEEDED = -1602,       ///< DFS quota exceeded.
+    DFS_PERMISSION_DENIED = -1603,    ///< DFS permission denied.
+    DFS_STALE_HANDLE = -1604,         ///< DFS file handle expired.
+    DFS_PARTIAL_WRITE = -1605,        ///< DFS partial write success.
+    TENANT_QUOTA_EXCEEDED = -1700,    ///< Tenant memory quota exceeded.
+    TENANT_NOT_REGISTERED = -1701,    ///< Tenant has no quota policy.
+    TENANT_NOT_EMPTY = -1702,         ///< Tenant still owns objects or quota.
 };
 
 int32_t toInt(ErrorCode errorCode) noexcept;
@@ -354,11 +478,26 @@ YLT_REFL(Segment, id, name, base, size, te_endpoint, protocol);
  * @brief Allocation strategy type for segment allocation
  */
 enum class AllocationStrategyType {
-    RANDOM = 0,        // Pure random allocation
-    FREE_RATIO_FIRST,  // Free-ratio-first allocation
-    CXL,               // CXL-specific allocation
+    RANDOM = 0,            // Pure random allocation
+    FREE_RATIO_FIRST,      // Free-ratio-first allocation
+    CXL,                   // CXL-specific allocation
+    SSD_FREE_RATIO_FIRST,  // SSD free-ratio-first allocation
 };
 
+/**
+ * @brief Represents a contiguous NoF ssd region
+ */
+struct NoFSegment {
+    UUID id{0, 0};
+    std::string name{};  // Logical segment name used for preferred allocation
+    uintptr_t base{0};
+    size_t size{0};
+    // TE p2p endpoint (ip:port) for transport-only addressing
+    std::string te_endpoint{};
+    NoFSegment() = default;
+};
+YLT_REFL(NoFSegment, id, name, base, size, te_endpoint);
+
 /**
  * @brief Client status from the master's perspective
  */
diff --git a/mooncake-store/include/utils.h b/mooncake-store/include/utils.h
index 8951f96c..160292c2 100644
--- a/mooncake-store/include/utils.h
+++ b/mooncake-store/include/utils.h
@@ -1,7 +1,10 @@
 #pragma once
 
+#include <algorithm>
+#include <cctype>
 #include <cstddef>
 #include <cstdlib>
+#include <optional>
 #include <linux/memfd.h>
 #include <linux/mman.h>
 #include <string>
@@ -168,14 +171,15 @@ std::string expected_to_str(const tl::expected<T, ErrorCode>& expected) {
 }
 
 /**
- * @brief Convert a string representation of size to bytes
+ * @brief Parse a string representation of size to bytes
  * @param str String representation of size (e.g., "1.5 GB", "1024 MB",
  * "1048576")
- * @return uint64_t Number of bytes, or 0 if parsing fails
+ * @return Parsed byte size, or std::nullopt if parsing fails
  */
-[[nodiscard]] inline uint64_t string_to_byte_size(const std::string& str) {
+[[nodiscard]] inline std::optional<uint64_t> try_string_to_byte_size(
+    const std::string& str) {
     if (str.empty()) {
-        return 0;
+        return std::nullopt;
     }
 
     // Create a copy for manipulation
@@ -186,7 +190,7 @@ std::string expected_to_str(const tl::expected<T, ErrorCode>& expected) {
     s.erase(s.find_last_not_of(" \t\r\n") + 1);
 
     if (s.empty()) {
-        return 0;
+        return std::nullopt;
     }
 
     // Handle special case for "infinite"
@@ -201,7 +205,10 @@ std::string expected_to_str(const tl::expected<T, ErrorCode>& expected) {
     try {
         value = std::stod(s, &pos);
     } catch (const std::exception&) {
-        return 0;  // Failed to parse number
+        return std::nullopt;  // Failed to parse number
+    }
+    if (value < 0) {
+        return std::nullopt;
     }
 
     if (pos >= s.length()) {
@@ -215,7 +222,8 @@ std::string expected_to_str(const tl::expected<T, ErrorCode>& expected) {
     unit.erase(0, unit.find_first_not_of(" \t\r\n"));
 
     // Convert to uppercase for comparison
-    std::transform(unit.begin(), unit.end(), unit.begin(), ::toupper);
+    std::transform(unit.begin(), unit.end(), unit.begin(),
+                   [](unsigned char c) -> char { return std::toupper(c); });
 
     // Apply unit multiplier
     const double KB = 1024.0;
@@ -235,10 +243,47 @@ std::string expected_to_str(const tl::expected<T, ErrorCode>& expected) {
         return static_cast<uint64_t>(value);
     } else {
         // Unknown unit
-        return 0;
+        return std::nullopt;
     }
 }
 
+/**
+ * @brief Convert a string representation of size to bytes
+ * @param str String representation of size (e.g., "1.5 GB", "1024 MB",
+ * "1048576")
+ * @return uint64_t Number of bytes, or 0 if parsing fails
+ */
+[[nodiscard]] inline uint64_t string_to_byte_size(const std::string& str) {
+    auto parsed = try_string_to_byte_size(str);
+    return parsed.value_or(0);
+}
+
+/**
+ * @brief Convert a boolean-like string to a bool
+ * @param str String representation ("1"/"true"/"yes"/"on" or
+ * "0"/"false"/"no"/"off")
+ * @return std::optional<bool> Parsed value, or std::nullopt if parsing fails
+ */
+[[nodiscard]] inline std::optional<bool> string_to_bool(std::string str) {
+    if (str.empty()) {
+        return std::nullopt;
+    }
+
+    str.erase(0, str.find_first_not_of(" \t\r\n"));
+    str.erase(str.find_last_not_of(" \t\r\n") + 1);
+    std::transform(str.begin(), str.end(), str.begin(),
+                   [](unsigned char c) { return std::tolower(c); });
+
+    if (str == "1" || str == "true" || str == "yes" || str == "on") {
+        return true;
+    }
+    if (str == "0" || str == "false" || str == "no" || str == "off") {
+        return false;
+    }
+
+    return std::nullopt;
+}
+
 /**
  * @brief Split a string by delimiter into a vector of strings
  * @param str The string to split
@@ -256,6 +301,7 @@ std::vector<std::string> splitString(const std::string& str,
 
 constexpr size_t SZ_2MB = 2 * 1024 * 1024;
 constexpr size_t SZ_1GB = 1024 * 1024 * 1024;
+constexpr double BYTES_PER_GIB = static_cast<double>(SZ_1GB);
 
 /**
  * @brief Allocates memory for the `BufferAllocator` class.
@@ -264,7 +310,8 @@ constexpr size_t SZ_1GB = 1024 * 1024 * 1024;
  */
 void* allocate_buffer_allocator_memory(
     size_t total_size, const std::string& protocol = "",
-    size_t alignment = facebook::cachelib::Slab::kSize);
+    size_t alignment = facebook::cachelib::Slab::kSize,
+    bool use_spdk_dma = false);
 
 inline size_t align_up(size_t size, size_t alignment) {
     if (alignment == 0) {
@@ -287,9 +334,6 @@ inline size_t align_up(size_t size, size_t alignment) {
         return 0;
     }
 
-    constexpr size_t SZ_2MB = 2 * 1024 * 1024;
-    constexpr size_t SZ_1GB = 1024 * 1024 * 1024;
-
     size_t size = SZ_2MB;  // Default to 2MB
 
     const char* size_env = std::getenv("MC_STORE_HUGEPAGE_SIZE");
@@ -333,8 +377,33 @@ inline size_t align_up(size_t size, size_t alignment) {
     return size;
 }
 
-// Hugepage-backed allocation helpers (MAP_HUGETLB + MADV_HUGEPAGE)
+/**
+ * Allocate mmap-backed buffer memory for host KV / transfer buffers.
+ *
+ * When the global mmap arena is enabled, this function serves allocations
+ * from the arena and still honors the caller's requested alignment.
+ * Arena-owned allocations remain owned by the arena until process shutdown.
+ *
+ * When the arena is disabled or unavailable, this falls back to a direct
+ * mmap() allocation and returns a pointer aligned to at least the system page
+ * size (or the configured hugepage size when available).
+ *
+ * @param total_size Total buffer size in bytes.
+ * @param alignment Minimum alignment requested by the caller.
+ * @return Pointer to the allocation, or nullptr on failure.
+ */
 void* allocate_buffer_mmap_memory(size_t total_size, size_t alignment);
+
+/**
+ * Release memory previously returned by allocate_buffer_mmap_memory().
+ *
+ * Direct-mmap allocations are unmapped immediately. Arena-owned pointers are
+ * intentionally not unmapped individually; in that case this function is a
+ * no-op and the arena releases the backing pool during process teardown.
+ *
+ * @param ptr Pointer previously returned by allocate_buffer_mmap_memory().
+ * @param total_size Original allocation size in bytes.
+ */
 void free_buffer_mmap_memory(void* ptr, size_t total_size);
 
 /**
diff --git a/mooncake-store/include/version.h.in b/mooncake-store/include/version.h.in
index 5fca0b53..8c27b0f6 100644
--- a/mooncake-store/include/version.h.in
+++ b/mooncake-store/include/version.h.in
@@ -4,7 +4,7 @@
 
 namespace mooncake {
 
-// Version information for the Mooncake system
+// Version information for the Mooncake system (used for RPC handshake)
 constexpr const char* MOONCAKE_STORE_VERSION = "@MOONCAKE_STORE_VERSION@";
 
 // Function to get the version string
@@ -13,4 +13,9 @@ inline const std::string& GetMooncakeStoreVersion() {
     return version;
 }
 
+// Human-readable display version: sourced from pyproject.toml + short git hash
+// e.g. "0.3.10.post2 (git: abc1234)"
+// Used for the --version flag only; does not affect RPC compatibility.
+constexpr const char* MOONCAKE_DISPLAY_VERSION = "@MOONCAKE_DISPLAY_VERSION@";
+
 }  // namespace mooncake
\ No newline at end of file
diff --git a/mooncake-store/rust/CMakeLists.txt b/mooncake-store/rust/CMakeLists.txt
index 54908527..49fa328f 100644
--- a/mooncake-store/rust/CMakeLists.txt
+++ b/mooncake-store/rust/CMakeLists.txt
@@ -1,39 +1,51 @@
-# CMake integration for the mooncake_store Rust crate.
+# CMake integration for the mooncake_store Rust package.
 #
 # This target is only added when the parent CMakeLists.txt sets
-# WITH_STORE_RUST=ON.  It depends on the mooncake_store C++ target so that
-# the static library is available before cargo tries to link against it.
+# WITH_STORE_RUST=ON.  It depends on the mooncake_store C++ target so that the
+# static library is available before cargo tries to link against it.
 
 add_custom_target(build_mooncake_store_rust DEPENDS mooncake_store)
 
 add_custom_command(
-    TARGET build_mooncake_store_rust
-    COMMAND
-        ${CMAKE_COMMAND} -E env
-            # Where cargo should place its build artefacts.
-            CARGO_TARGET_DIR=${CMAKE_CURRENT_BINARY_DIR}
-            # Directory containing libmooncake_store.a (or .so).
-            MOONCAKE_STORE_LIB_DIR=${CMAKE_CURRENT_BINARY_DIR}/../src
-            # Directory containing store_c.h.
-            MOONCAKE_STORE_INCLUDE_DIR=${CMAKE_CURRENT_SOURCE_DIR}/../include
-        cargo build --release
-    WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}
-    COMMENT "Building mooncake_store Rust crate"
-    VERBATIM
-)
+  TARGET build_mooncake_store_rust
+  COMMAND
+    ${CMAKE_COMMAND} -E env
+    # Where cargo should place its build artefacts.
+    CARGO_TARGET_DIR=${CMAKE_CURRENT_BINARY_DIR}
+    MOONCAKE_BUILD_DIR=${PROJECT_BINARY_DIR}
+    # Directory containing libmooncake_store.a (or .so).
+    MOONCAKE_STORE_LIB_DIR=${CMAKE_CURRENT_BINARY_DIR}/../src
+    # Directory containing store_c.h.
+    MOONCAKE_STORE_INCLUDE_DIR=${CMAKE_CURRENT_SOURCE_DIR}/../include cargo
+    build --release
+  WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}
+  COMMENT "Building mooncake_store Rust package"
+  VERBATIM)
 
-# Build the basic_usage example so that CI can verify it compiles.
 add_custom_target(build_mooncake_store_rust_example DEPENDS mooncake_store)
 
 add_custom_command(
-    TARGET build_mooncake_store_rust_example
-    COMMAND
-        ${CMAKE_COMMAND} -E env
-            CARGO_TARGET_DIR=${CMAKE_CURRENT_BINARY_DIR}
-            MOONCAKE_STORE_LIB_DIR=${CMAKE_CURRENT_BINARY_DIR}/../src
-            MOONCAKE_STORE_INCLUDE_DIR=${CMAKE_CURRENT_SOURCE_DIR}/../include
-        cargo build --example basic_usage --release
-    WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}
-    COMMENT "Building mooncake_store Rust basic_usage example"
-    VERBATIM
-)
+  TARGET build_mooncake_store_rust_example
+  COMMAND
+    ${CMAKE_COMMAND} -E env CARGO_TARGET_DIR=${CMAKE_CURRENT_BINARY_DIR}
+    MOONCAKE_BUILD_DIR=${PROJECT_BINARY_DIR}
+    MOONCAKE_STORE_LIB_DIR=${CMAKE_CURRENT_BINARY_DIR}/../src
+    MOONCAKE_STORE_INCLUDE_DIR=${CMAKE_CURRENT_SOURCE_DIR}/../include cargo
+    build --examples --release
+  WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}
+  COMMENT "Building mooncake_store Rust examples"
+  VERBATIM)
+
+add_custom_target(build_mooncake_store_rust_tests DEPENDS mooncake_store)
+
+add_custom_command(
+  TARGET build_mooncake_store_rust_tests
+  COMMAND
+    ${CMAKE_COMMAND} -E env CARGO_TARGET_DIR=${CMAKE_CURRENT_BINARY_DIR}
+    MOONCAKE_BUILD_DIR=${PROJECT_BINARY_DIR}
+    MOONCAKE_STORE_LIB_DIR=${CMAKE_CURRENT_BINARY_DIR}/../src
+    MOONCAKE_STORE_INCLUDE_DIR=${CMAKE_CURRENT_SOURCE_DIR}/../include cargo test
+    --tests --no-run --release
+  WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}
+  COMMENT "Building mooncake_store Rust tests"
+  VERBATIM)
diff --git a/mooncake-store/rust/README.md b/mooncake-store/rust/README.md
new file mode 100644
index 00000000..65d49125
--- /dev/null
+++ b/mooncake-store/rust/README.md
@@ -0,0 +1,91 @@
+# Mooncake Store Rust Bindings
+
+This directory contains the Rust bindings, examples, benchmarks, and smoke tests
+for Mooncake Store. The Rust crate links against the C++ Mooncake Store build, so
+run a CMake build before using standalone Cargo commands.
+
+## Prerequisites
+
+Install the Rust toolchain and libclang before building this crate. On Ubuntu:
+
+```bash
+sudo apt-get install libclang-dev clang
+```
+
+## Build With CMake
+
+Configure Mooncake with Store and Rust bindings enabled:
+
+```bash
+cmake -S . -B build -G Ninja \
+  -DWITH_STORE=ON \
+  -DWITH_STORE_RUST=ON
+```
+
+Build the Rust crate, examples, and tests through the CMake targets:
+
+```bash
+cmake --build build --target build_mooncake_store_rust
+cmake --build build --target build_mooncake_store_rust_example
+cmake --build build --target build_mooncake_store_rust_tests
+```
+
+## Build With Cargo
+
+Standalone Cargo builds need the CMake build directory and Store headers:
+
+```bash
+cd mooncake-store/rust
+
+export MOONCAKE_BUILD_DIR=/path/to/Mooncake/build
+export MOONCAKE_STORE_LIB_DIR=$MOONCAKE_BUILD_DIR/mooncake-store/src
+export MOONCAKE_STORE_INCLUDE_DIR=/path/to/Mooncake/mooncake-store/include
+
+cargo build --examples --release
+cargo test --tests --no-run --release
+```
+
+At runtime, the dynamic linker must be able to find the Mooncake shared
+libraries generated by CMake. If they are not installed in a system library
+directory, add the relevant build outputs to `LD_LIBRARY_PATH`:
+
+```bash
+export LD_LIBRARY_PATH=$MOONCAKE_BUILD_DIR/mooncake-common:\
+$MOONCAKE_BUILD_DIR/mooncake-common/etcd:\
+$MOONCAKE_BUILD_DIR/mooncake-store/src:\
+$MOONCAKE_BUILD_DIR/mooncake-store/src/cachelib_memory_allocator:\
+$MOONCAKE_BUILD_DIR/mooncake-transfer-engine/src:\
+$MOONCAKE_BUILD_DIR/mooncake-transfer-engine/src/common/base:\
+${LD_LIBRARY_PATH:-}
+```
+
+## Integration Smoke Test
+
+The integration smoke test needs a running metadata server and
+`mooncake_master`. The CI workflow is the best reference for the complete
+service setup. Once those services are running, execute:
+
+```bash
+cd mooncake-store/rust
+
+MC_RUST_STORE_RUN_INTEGRATION=true \
+MC_METADATA_SERVER=http://127.0.0.1:8080/metadata \
+MC_RUST_STORE_MASTER_ADDR=127.0.0.1:50051 \
+MC_RUST_STORE_LOCAL_HOSTNAME=127.0.0.1 \
+MC_RUST_STORE_PROTOCOL=tcp \
+MC_RUST_STORE_DEVICE_NAME= \
+cargo test --test minimal_smoke -- --nocapture
+```
+
+## Benchmark Smoke Test
+
+For a short benchmark run, reduce the iteration counts:
+
+```bash
+cd mooncake-store/rust
+
+MC_RUST_BENCH_ITERATIONS=4 \
+MC_RUST_BENCH_VALUE_SIZE=4096 \
+MC_RUST_BENCH_WARMUP=1 \
+cargo run --release --example store_benchmark
+```
diff --git a/mooncake-store/rust/build.rs b/mooncake-store/rust/build.rs
index 99d1714e..abcd95b0 100644
--- a/mooncake-store/rust/build.rs
+++ b/mooncake-store/rust/build.rs
@@ -13,7 +13,155 @@
 // limitations under the License.
 
 use std::env;
+use std::fs;
 use std::path::PathBuf;
+use std::process::Command;
+
+fn push_existing_dir(search_dirs: &mut Vec<PathBuf>, dir: PathBuf) {
+    if dir.is_dir() && !search_dirs.iter().any(|existing| existing == &dir) {
+        search_dirs.push(dir);
+    }
+}
+
+fn push_env_paths(search_dirs: &mut Vec<PathBuf>, name: &str) {
+    if let Some(value) = env::var_os(name) {
+        for dir in env::split_paths(&value) {
+            push_existing_dir(search_dirs, dir);
+        }
+    }
+}
+
+fn push_cmake_prefix_lib_dirs(search_dirs: &mut Vec<PathBuf>, prefix: PathBuf) {
+    push_existing_dir(search_dirs, prefix.join("lib"));
+    push_existing_dir(search_dirs, prefix.join("lib64"));
+}
+
+fn push_cmake_prefix_paths(search_dirs: &mut Vec<PathBuf>, value: &str) {
+    for prefix in value.split(';').filter(|prefix| !prefix.is_empty()) {
+        push_cmake_prefix_lib_dirs(search_dirs, PathBuf::from(prefix));
+    }
+}
+
+fn push_cmake_cache_library_dirs(search_dirs: &mut Vec<PathBuf>, build_dir: &PathBuf) {
+    let cache_path = build_dir.join("CMakeCache.txt");
+    let Ok(cache) = fs::read_to_string(cache_path) else {
+        return;
+    };
+
+    for line in cache.lines() {
+        let Some((key, value)) = line.split_once('=') else {
+            continue;
+        };
+
+        if key == "CMAKE_PREFIX_PATH:PATH" {
+            push_cmake_prefix_paths(search_dirs, value);
+            continue;
+        }
+
+        if key.ends_with("_DIR:PATH") {
+            let package_dir = PathBuf::from(value);
+            if package_dir.parent().and_then(|dir| dir.file_name())
+                == Some(std::ffi::OsStr::new("cmake"))
+            {
+                if let Some(lib_dir) = package_dir.ancestors().nth(2) {
+                    push_existing_dir(search_dirs, lib_dir.to_path_buf());
+                }
+            }
+            continue;
+        }
+
+        if key.ends_with("_LIBRARY:FILEPATH") || key.ends_with("_LIBRARIES:FILEPATH") {
+            let library_path = PathBuf::from(value);
+            if let Some(parent) = library_path.parent() {
+                push_existing_dir(search_dirs, parent.to_path_buf());
+            }
+        }
+    }
+}
+
+fn has_library(search_dirs: &[PathBuf], candidates: &[&str]) -> bool {
+    search_dirs.iter().any(|dir| {
+        candidates.iter().any(|candidate| {
+            ["a", "so", "dylib"]
+                .into_iter()
+                .map(|ext| dir.join(format!("lib{candidate}.{ext}")))
+                .any(|path| path.exists())
+        })
+    })
+}
+
+fn emit_link_searches(search_dirs: &[PathBuf]) {
+    for dir in search_dirs {
+        println!("cargo:rustc-link-search=native={}", dir.display());
+    }
+}
+
+fn emit_runtime_rpaths(search_dirs: &[PathBuf]) {
+    for dir in search_dirs {
+        let dir_str = dir.display();
+        // Use rustc-link-arg (not -tests) so that the rpath is also applied
+        // to the lib-test binary produced by `cargo test` for #[cfg(test)]
+        // modules inside src/lib.rs.
+        println!("cargo:rustc-link-arg=-Wl,-rpath,{dir_str}");
+    }
+}
+
+fn compiler_candidates() -> Vec<String> {
+    let mut tools = Vec::new();
+
+    for env_var in ["CC", "CXX"] {
+        if let Ok(value) = env::var(env_var) {
+            let tool = value.trim();
+            if !tool.is_empty() && !tools.iter().any(|existing| existing == tool) {
+                tools.push(tool.to_string());
+            }
+        }
+    }
+
+    for tool in ["gcc", "cc", "clang", "c++"] {
+        if !tools.iter().any(|existing| existing == tool) {
+            tools.push(tool.to_string());
+        }
+    }
+
+    tools
+}
+
+fn compiler_runtime_library(file_name: &str) -> Option<PathBuf> {
+    for tool in compiler_candidates() {
+        let output = Command::new(&tool)
+            .arg(format!("-print-file-name={file_name}"))
+            .output()
+            .ok()?;
+
+        if !output.status.success() {
+            continue;
+        }
+
+        let path = String::from_utf8(output.stdout).ok()?;
+        let path = PathBuf::from(path.trim());
+        if path.as_os_str().is_empty() || path == PathBuf::from(file_name) || !path.exists() {
+            continue;
+        }
+
+        return Some(path);
+    }
+
+    None
+}
+
+fn add_compiler_runtime_search_dir(search_dirs: &mut Vec<PathBuf>, file_name: &str) -> bool {
+    let Some(path) = compiler_runtime_library(file_name) else {
+        return false;
+    };
+
+    if let Some(parent) = path.parent() {
+        push_existing_dir(search_dirs, parent.to_path_buf());
+        return true;
+    }
+
+    false
+}
 
 fn main() {
     // -----------------------------------------------------------------------
@@ -29,33 +177,211 @@ fn main() {
         .unwrap_or_else(|_| "../../build/mooncake-store/src".to_string());
 
     println!("cargo:rustc-link-search=native={lib_dir}");
+
+    // mooncake_store depends on libasio.so (shared) built in mooncake-common.
+    let lib_path = PathBuf::from(&lib_dir);
+    let build_dir = lib_path.ancestors().nth(2).map(PathBuf::from).unwrap_or_else(|| {
+        println!("cargo:warning=MOONCAKE_STORE_LIB_DIR='{lib_dir}' does not have enough parent directories; using current directory");
+        PathBuf::from(".")
+    });
+    println!(
+        "cargo:rustc-link-search=native={}",
+        build_dir.join("mooncake-common").display()
+    );
+    // mooncake_common static library lives in the src/ subdirectory.
+    println!(
+        "cargo:rustc-link-search=native={}",
+        build_dir.join("mooncake-common/src").display()
+    );
+
+    // transfer_engine is built in a sibling directory.
+    println!(
+        "cargo:rustc-link-search=native={}",
+        build_dir.join("mooncake-transfer-engine/src").display()
+    );
+
+    // common/base library (contains mooncake::Status etc.)
+    println!(
+        "cargo:rustc-link-search=native={}",
+        build_dir.join("mooncake-transfer-engine/src/common/base").display()
+    );
+
+    // CUDA runtime libraries (needed by transfer_engine RDMA transport).
+    let cuda_home = env::var("CUDA_HOME")
+        .or_else(|_| env::var("CUDA_PATH"))
+        .unwrap_or_else(|_| "/usr/local/cuda".to_string());
+    println!("cargo:rustc-link-search=native={}/targets/x86_64-linux/lib", cuda_home);
+
+    // cachelib_memory_allocator is a static library built alongside mooncake_store.
+    println!(
+        "cargo:rustc-link-search=native={}",
+        build_dir.join("mooncake-store/src/cachelib_memory_allocator").display()
+    );
+
     println!("cargo:rustc-link-lib=mooncake_store");
 
     // Dependencies of mooncake_store that must be satisfied at link time.
     // The list mirrors what mooncake-store/src/CMakeLists.txt links against.
     println!("cargo:rustc-link-lib=transfer_engine");
+    println!("cargo:rustc-link-lib=mooncake_common"); // Environ::Get() and other common utilities
+    println!("cargo:rustc-link-lib=base"); // mooncake::Status etc.
+    println!("cargo:rustc-link-lib=asio"); // shared library built by mooncake-common
+    println!("cargo:rustc-link-lib=jsoncpp"); // transfer_engine dependency
+    println!("cargo:rustc-link-lib=cachelib_memory_allocator"); // static
     println!("cargo:rustc-link-lib=stdc++");
     println!("cargo:rustc-link-lib=glog");
     println!("cargo:rustc-link-lib=gflags");
+    println!("cargo:rustc-link-lib=numa");   // NUMA binding
+    println!("cargo:rustc-link-lib=curl");    // HTTP metadata plugin
+    println!("cargo:rustc-link-lib=ibverbs"); // RDMA transport
+    println!("cargo:rustc-link-lib=yaml-cpp"); // tenant quota policy connector
     println!("cargo:rustc-link-lib=pthread");
     println!("cargo:rustc-link-lib=xxhash");
 
     // -----------------------------------------------------------------------
     // Header path for bindgen
     // -----------------------------------------------------------------------
+    let manifest_dir = PathBuf::from(env::var("CARGO_MANIFEST_DIR").expect("missing CARGO_MANIFEST_DIR"));
+    let mut search_dirs = Vec::new();
+
+    let explicit_lib_dir = env::var("MOONCAKE_STORE_LIB_DIR").ok().map(PathBuf::from);
+    if let Some(dir) = explicit_lib_dir.clone() {
+        push_existing_dir(&mut search_dirs, dir);
+    }
+
+    if let Ok(build_dir) = env::var("MOONCAKE_BUILD_DIR") {
+        let build_dir = PathBuf::from(build_dir);
+        push_cmake_cache_library_dirs(&mut search_dirs, &build_dir);
+        for dir in [
+            build_dir.join("mooncake-store/src"),
+            build_dir.join("mooncake-store/src/cachelib_memory_allocator"),
+            build_dir.join("mooncake-transfer-engine/src"),
+            build_dir.join("mooncake-transfer-engine/src/common/base"),
+            build_dir.join("mooncake-asio"),
+            build_dir.join("mooncake-common/etcd"),
+        ] {
+            push_existing_dir(&mut search_dirs, dir);
+        }
+    }
+
+    let default_build_dir = manifest_dir.join("../../build");
+    push_cmake_cache_library_dirs(&mut search_dirs, &default_build_dir);
+    for dir in [
+        default_build_dir.join("mooncake-store/src"),
+        default_build_dir.join("mooncake-store/src/cachelib_memory_allocator"),
+        default_build_dir.join("mooncake-transfer-engine/src"),
+        default_build_dir.join("mooncake-transfer-engine/src/common/base"),
+        default_build_dir.join("mooncake-asio"),
+        default_build_dir.join("mooncake-common"),
+        default_build_dir.join("mooncake-common/etcd"),
+        PathBuf::from("/usr/local/lib"),
+        PathBuf::from("/usr/lib/x86_64-linux-gnu"),
+        PathBuf::from("/lib/x86_64-linux-gnu"),
+    ] {
+        push_existing_dir(&mut search_dirs, dir);
+    }
+
+    if let Ok(cmake_prefix_path) = env::var("CMAKE_PREFIX_PATH") {
+        push_cmake_prefix_paths(&mut search_dirs, &cmake_prefix_path);
+    }
+
+    push_env_paths(&mut search_dirs, "LD_LIBRARY_PATH");
+    push_env_paths(&mut search_dirs, "LIBRARY_PATH");
+
+    // Only link the AddressSanitizer runtime when explicitly requested via
+    // MOONCAKE_LINK_ASAN. Linking libasan whenever the toolchain merely *ships*
+    // it (gcc always does) pulls the ASan runtime into an otherwise
+    // non-sanitized build, even though a Release Mooncake has no `__asan_*`
+    // symbols. Any consumer that `dlopen()`s the resulting library then aborts
+    // at load with "ASan runtime does not come first in the initial library
+    // list", since a late-loaded plugin can never be first. Default to NOT
+    // linking it; sanitized builds opt in with `MOONCAKE_LINK_ASAN=1` (and run
+    // with `LD_PRELOAD=$(gcc -print-file-name=libasan.so)`).
+    let want_asan = env::var("MOONCAKE_LINK_ASAN")
+        .map(|value| {
+            !matches!(
+                value.trim().to_ascii_lowercase().as_str(),
+                "" | "0" | "false" | "off" | "no"
+            )
+        })
+        .unwrap_or(false);
+    let has_asan_runtime = if want_asan {
+        let asan_runtime_so = compiler_runtime_library("libasan.so");
+        if let Some(path) = asan_runtime_so.as_ref() {
+            if let Some(parent) = path.parent() {
+                push_existing_dir(&mut search_dirs, parent.to_path_buf());
+            }
+        }
+        asan_runtime_so.is_some() || add_compiler_runtime_search_dir(&mut search_dirs, "libasan.a")
+    } else {
+        false
+    };
+    let has_gcov_runtime = add_compiler_runtime_search_dir(&mut search_dirs, "libgcov.a")
+        || add_compiler_runtime_search_dir(&mut search_dirs, "libgcov.so");
+
+    emit_link_searches(&search_dirs);
+    emit_runtime_rpaths(&search_dirs);
+
+    if want_asan && (has_asan_runtime || has_library(&search_dirs, &["asan"])) {
+        println!("cargo:rustc-link-lib=asan");
+    }
+
+    for library in [
+        "mooncake_store",
+        "cachelib_memory_allocator",
+        "transfer_engine",
+        "base",
+        "asio",
+        "stdc++",
+        "glog",
+        "gflags",
+        "pthread",
+        "xxhash",
+        "numa",
+        "ibverbs",
+        "jsoncpp",
+        "yaml-cpp",
+        "zstd",
+        "m",
+    ] {
+        println!("cargo:rustc-link-lib={library}");
+    }
+
+    for (link_name, candidates) in [
+        ("etcd_wrapper", &["etcd_wrapper"] as &[&str]),
+        ("hiredis", &["hiredis"]),
+        ("curl", &["curl"]),
+        ("cuda", &["cuda"]),
+        ("cudart", &["cudart"]),
+        ("mlx5", &["mlx5"]), // IBGDA device transport (mlx5 DevX) pulled into transfer_engine, CUDA-only
+        ("uring", &["uring"]),
+    ] {
+        if has_library(&search_dirs, candidates) {
+            println!("cargo:rustc-link-lib={link_name}");
+        }
+    }
+
+    if has_gcov_runtime || has_library(&search_dirs, &["gcov"]) {
+        println!("cargo:rustc-link-lib=gcov");
+    }
+
     let include_dir = env::var("MOONCAKE_STORE_INCLUDE_DIR")
         .unwrap_or_else(|_| "../include".to_string());
 
     let header = format!("{include_dir}/store_c.h");
 
-    // Re-run this build script if the C header changes.
     println!("cargo:rerun-if-changed={header}");
+    println!("cargo:rerun-if-env-changed=MOONCAKE_BUILD_DIR");
     println!("cargo:rerun-if-env-changed=MOONCAKE_STORE_LIB_DIR");
     println!("cargo:rerun-if-env-changed=MOONCAKE_STORE_INCLUDE_DIR");
+    println!("cargo:rerun-if-env-changed=LD_LIBRARY_PATH");
+    println!("cargo:rerun-if-env-changed=LIBRARY_PATH");
+    println!("cargo:rerun-if-env-changed=CC");
+    println!("cargo:rerun-if-env-changed=CXX");
+    println!("cargo:rerun-if-env-changed=MOONCAKE_LINK_ASAN");
 
     let bindings = bindgen::Builder::default()
         .header(&header)
-        // Only pull in declarations from store_c.h (no transitive system headers).
         .allowlist_function("mooncake_store_.*")
         .allowlist_type("mooncake_.*")
         .generate()
diff --git a/mooncake-store/rust/examples/store_benchmark.rs b/mooncake-store/rust/examples/store_benchmark.rs
new file mode 100644
index 00000000..dc20ae03
--- /dev/null
+++ b/mooncake-store/rust/examples/store_benchmark.rs
@@ -0,0 +1,130 @@
+// Copyright 2024 KVCache.AI
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use std::time::{Instant, SystemTime, UNIX_EPOCH};
+
+use mooncake_store::MooncakeStore;
+
+fn env_or_default<T>(key: &str, default: T) -> T
+where
+    T: std::str::FromStr,
+{
+    std::env::var(key)
+        .ok()
+        .and_then(|value| value.parse::<T>().ok())
+        .unwrap_or(default)
+}
+
+fn unique_prefix() -> String {
+    let timestamp_ns = SystemTime::now()
+        .duration_since(UNIX_EPOCH)
+        .map(|duration| duration.as_nanos())
+        .unwrap_or_default();
+    format!("rust-bench-{}-{timestamp_ns}", std::process::id())
+}
+
+fn mib_per_second(total_bytes: usize, elapsed: std::time::Duration) -> f64 {
+    if elapsed.is_zero() {
+        return 0.0;
+    }
+    total_bytes as f64 / (1024.0 * 1024.0) / elapsed.as_secs_f64()
+}
+
+fn main() -> Result<(), Box<dyn std::error::Error>> {
+    let iterations = env_or_default("MC_RUST_BENCH_ITERATIONS", 64usize);
+    let payload_size = env_or_default("MC_RUST_BENCH_VALUE_SIZE", 64 * 1024usize);
+    let warmup_iterations = env_or_default("MC_RUST_BENCH_WARMUP", 4usize);
+    let metadata_server = std::env::var("MC_METADATA_SERVER")
+        .unwrap_or_else(|_| "http://127.0.0.1:8080/metadata".to_string());
+    let master_server_addr = std::env::var("MC_RUST_STORE_MASTER_ADDR")
+        .unwrap_or_else(|_| "127.0.0.1:50051".to_string());
+    let local_hostname =
+        std::env::var("MC_RUST_STORE_LOCAL_HOSTNAME").unwrap_or_else(|_| "localhost".to_string());
+    let protocol = std::env::var("MC_RUST_STORE_PROTOCOL").unwrap_or_else(|_| "tcp".to_string());
+    let device_name = std::env::var("MC_RUST_STORE_DEVICE_NAME").unwrap_or_default();
+    let global_segment_size = env_or_default("MC_RUST_STORE_GLOBAL_SEGMENT_SIZE", 512u64 << 20);
+    let local_buffer_size = env_or_default("MC_RUST_STORE_LOCAL_BUFFER_SIZE", 128u64 << 20);
+
+    let store = MooncakeStore::new()?;
+    store.setup(
+        &local_hostname,
+        &metadata_server,
+        global_segment_size,
+        local_buffer_size,
+        &protocol,
+        &device_name,
+        &master_server_addr,
+    )?;
+    store.health_check()?;
+
+    let prefix = unique_prefix();
+    let payload: Vec<u8> = (0..payload_size).map(|index| (index % 251) as u8).collect();
+
+    for iteration in 0..warmup_iterations {
+        let key = format!("{prefix}-warmup-{iteration}");
+        store.put(&key, &payload, None)?;
+        let fetched = store.get(&key)?;
+        assert_eq!(fetched, payload, "warmup round-trip mismatch");
+        store.remove(&key, true)?;
+    }
+
+    let keys: Vec<String> = (0..iterations)
+        .map(|iteration| format!("{prefix}-key-{iteration}"))
+        .collect();
+
+    let put_start = Instant::now();
+    for key in &keys {
+        store.put(key, &payload, None)?;
+    }
+    let put_elapsed = put_start.elapsed();
+
+    let get_start = Instant::now();
+    for key in &keys {
+        let fetched = store.get(key)?;
+        assert_eq!(fetched, payload, "benchmark round-trip mismatch for {key}");
+    }
+    let get_elapsed = get_start.elapsed();
+
+    let remove_start = Instant::now();
+    for key in &keys {
+        store.remove(key, true)?;
+    }
+    let remove_elapsed = remove_start.elapsed();
+
+    let total_bytes = iterations * payload.len();
+    println!("Mooncake Store Rust benchmark");
+    println!("iterations={iterations}");
+    println!("payload_size_bytes={payload_size}");
+    println!("put_seconds={:.6}", put_elapsed.as_secs_f64());
+    println!(
+        "put_mib_per_sec={:.2}",
+        mib_per_second(total_bytes, put_elapsed)
+    );
+    println!("get_seconds={:.6}", get_elapsed.as_secs_f64());
+    println!(
+        "get_mib_per_sec={:.2}",
+        mib_per_second(total_bytes, get_elapsed)
+    );
+    println!("remove_seconds={:.6}", remove_elapsed.as_secs_f64());
+    println!(
+        "remove_ops_per_sec={:.2}",
+        if remove_elapsed.is_zero() {
+            0.0
+        } else {
+            iterations as f64 / remove_elapsed.as_secs_f64()
+        }
+    );
+
+    Ok(())
+}
diff --git a/mooncake-store/rust/src/error.rs b/mooncake-store/rust/src/error.rs
index 8f04ef04..a854a3aa 100644
--- a/mooncake-store/rust/src/error.rs
+++ b/mooncake-store/rust/src/error.rs
@@ -40,4 +40,8 @@ pub enum StoreError {
     /// [`OperationFailed`](StoreError::OperationFailed) instead.
     #[error("key not found")]
     NotFound,
+
+    /// One or more arguments are invalid (e.g. mismatched array lengths).
+    #[error("invalid argument: {0}")]
+    InvalidArgument(String),
 }
diff --git a/mooncake-store/rust/src/store.rs b/mooncake-store/rust/src/store.rs
index 87888b39..cb59439c 100644
--- a/mooncake-store/rust/src/store.rs
+++ b/mooncake-store/rust/src/store.rs
@@ -418,6 +418,156 @@ impl MooncakeStore {
         Ok(())
     }
 
+    // -----------------------------------------------------------------------
+    // Batch operations
+    // -----------------------------------------------------------------------
+
+    /// Batch version of [`put_from`](MooncakeStore::put_from).
+    ///
+    /// Returns a `Vec<i32>` where each element is the result code for the
+    /// corresponding key (0 = success, non-zero = error).
+    ///
+    /// # Safety
+    ///
+    /// Same safety requirements as `put_from`: each buffer must be registered.
+    pub unsafe fn batch_put_from(
+        &self,
+        keys: &[&str],
+        buffers: &[*mut c_void],
+        sizes: &[usize],
+        config: Option<&ReplicateConfig>,
+    ) -> Result<Vec<i32>, StoreError> {
+        let count = keys.len();
+        if count == 0 {
+            return Ok(Vec::new());
+        }
+        if buffers.len() != count || sizes.len() != count {
+            return Err(StoreError::InvalidArgument(
+                "keys, buffers, and sizes must have the same length".to_string(),
+            ));
+        }
+
+        let key_strings: Vec<CString> = keys
+            .iter()
+            .map(|k| CString::new(*k))
+            .collect::<Result<_, _>>()?;
+        let key_ptrs: Vec<*const libc::c_char> =
+            key_strings.iter().map(|s| s.as_ptr()).collect();
+
+        let (_c_config, _strings, _ptrs) = Self::prepare_config(config)?;
+        let cfg_ptr = _c_config
+            .as_ref()
+            .map_or(std::ptr::null(), |c| c as *const _);
+
+        let mut results = vec![0i32; count];
+
+        let rc = ffi::mooncake_store_batch_put_from(
+            self.handle,
+            key_ptrs.as_ptr() as *mut *const libc::c_char,
+            buffers.as_ptr() as *mut *mut c_void,
+            sizes.as_ptr(),
+            count,
+            cfg_ptr,
+            results.as_mut_ptr(),
+        );
+
+        if rc != 0 {
+            return Err(StoreError::OperationFailed(rc));
+        }
+
+        Ok(results)
+    }
+
+    /// Batch version of [`get_into`](MooncakeStore::get_into).
+    ///
+    /// Returns a `Vec<i64>` where each element is the number of bytes written
+    /// (≥ 0) or an error code (< 0).
+    ///
+    /// # Safety
+    ///
+    /// Same safety requirements as `get_into`: each buffer must be writable.
+    pub unsafe fn batch_get_into(
+        &self,
+        keys: &[&str],
+        buffers: &[*mut c_void],
+        sizes: &[usize],
+    ) -> Result<Vec<i64>, StoreError> {
+        let count = keys.len();
+        if count == 0 {
+            return Ok(Vec::new());
+        }
+        if buffers.len() != count || sizes.len() != count {
+            return Err(StoreError::InvalidArgument(
+                "keys, buffers, and sizes must have the same length".to_string(),
+            ));
+        }
+
+        let key_strings: Vec<CString> = keys
+            .iter()
+            .map(|k| CString::new(*k))
+            .collect::<Result<_, _>>()?;
+        let key_ptrs: Vec<*const libc::c_char> =
+            key_strings.iter().map(|s| s.as_ptr()).collect();
+
+        let mut results = vec![0i64; count];
+
+        let rc = ffi::mooncake_store_batch_get_into(
+            self.handle,
+            key_ptrs.as_ptr() as *mut *const libc::c_char,
+            buffers.as_ptr() as *mut *mut c_void,
+            sizes.as_ptr(),
+            count,
+            results.as_mut_ptr(),
+        );
+
+        if rc != 0 {
+            return Err(StoreError::OperationFailed(rc));
+        }
+
+        Ok(results)
+    }
+
+    /// Batch check existence of multiple keys.
+    pub fn batch_is_exist(
+        &self,
+        keys: &[&str],
+    ) -> Result<Vec<bool>, StoreError> {
+        let count = keys.len();
+        if count == 0 {
+            return Ok(Vec::new());
+        }
+        let key_strings: Vec<CString> = keys
+            .iter()
+            .map(|k| CString::new(*k))
+            .collect::<Result<_, _>>()?;
+        let key_ptrs: Vec<*const libc::c_char> =
+            key_strings.iter().map(|s| s.as_ptr()).collect();
+
+        let mut results = vec![0i32; count];
+
+        let rc = unsafe {
+            ffi::mooncake_store_batch_is_exist(
+                self.handle,
+                key_ptrs.as_ptr() as *mut *const libc::c_char,
+                count,
+                results.as_mut_ptr(),
+            )
+        };
+
+        if rc != 0 {
+            return Err(StoreError::OperationFailed(rc));
+        }
+
+        results
+            .into_iter()
+            .map(|r| match r {
+                1 => Ok(true),
+                0 => Ok(false),
+                _ => Err(StoreError::OperationFailed(r)),
+            })
+            .collect()
+    }
+
     // -----------------------------------------------------------------------
     // Internal helpers
     // -----------------------------------------------------------------------
@@ -548,4 +698,43 @@ mod tests {
         assert_eq!(strings.len(), 2);
         assert_eq!(ptrs.len(), 2);
     }
+
+    // -----------------------------------------------------------------------
+    // Batch API tests
+    // -----------------------------------------------------------------------
+
+    #[test]
+    fn batch_is_exist_empty() {
+        let store = MooncakeStore::new().expect("new should succeed");
+        let results = store.batch_is_exist(&[]).expect("empty batch should succeed");
+        assert!(results.is_empty());
+    }
+
+    #[test]
+    fn batch_put_from_rejects_mismatched_lengths() {
+        let store = MooncakeStore::new().expect("new should succeed");
+        let keys = &["key1", "key2"];
+        let buffers: Vec<*mut c_void> = vec![std::ptr::null_mut(); 3];
+        let sizes = &[100usize, 200];
+
+        let result = unsafe { store.batch_put_from(keys, &buffers, sizes, None) };
+        assert!(matches!(
+            result,
+            Err(StoreError::InvalidArgument(_))
+        ));
+    }
+
+    #[test]
+    fn batch_get_into_rejects_mismatched_lengths() {
+        let store = MooncakeStore::new().expect("new should succeed");
+        let keys = &["key1"];
+        let buffers: Vec<*mut c_void> = vec![std::ptr::null_mut(); 2];
+        let sizes = &[100usize];
+
+        let result = unsafe { store.batch_get_into(keys, &buffers, sizes) };
+        assert!(matches!(
+            result,
+            Err(StoreError::InvalidArgument(_))
+        ));
+    }
 }
diff --git a/mooncake-store/rust/tests/minimal_smoke.rs b/mooncake-store/rust/tests/minimal_smoke.rs
new file mode 100644
index 00000000..1e95cff5
--- /dev/null
+++ b/mooncake-store/rust/tests/minimal_smoke.rs
@@ -0,0 +1,76 @@
+// Copyright 2024 KVCache.AI
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use std::time::{SystemTime, UNIX_EPOCH};
+
+use mooncake_store::MooncakeStore;
+
+fn env_or_default(key: &str, default: &str) -> String {
+    std::env::var(key).unwrap_or_else(|_| default.to_string())
+}
+
+fn should_run_integration() -> bool {
+    matches!(
+        std::env::var("MC_RUST_STORE_RUN_INTEGRATION").as_deref(),
+        Ok("1") | Ok("true") | Ok("TRUE")
+    )
+}
+
+fn unique_key() -> String {
+    let timestamp_ns = SystemTime::now()
+        .duration_since(UNIX_EPOCH)
+        .map(|duration| duration.as_nanos())
+        .unwrap_or_default();
+    format!("rust-smoke-{}-{timestamp_ns}", std::process::id())
+}
+
+#[test]
+fn smoke_round_trip() -> Result<(), Box<dyn std::error::Error>> {
+    if !should_run_integration() {
+        eprintln!("skipping Mooncake Store Rust smoke test; set MC_RUST_STORE_RUN_INTEGRATION=1 to enable");
+        return Ok(());
+    }
+
+    let metadata_server = env_or_default("MC_METADATA_SERVER", "http://127.0.0.1:8080/metadata");
+    let master_server_addr = env_or_default("MC_RUST_STORE_MASTER_ADDR", "127.0.0.1:50051");
+    let local_hostname = env_or_default("MC_RUST_STORE_LOCAL_HOSTNAME", "localhost");
+    let protocol = env_or_default("MC_RUST_STORE_PROTOCOL", "tcp");
+    let device_name = std::env::var("MC_RUST_STORE_DEVICE_NAME").unwrap_or_default();
+
+    let store = MooncakeStore::new()?;
+    store.setup(
+        &local_hostname,
+        &metadata_server,
+        512 << 20,
+        128 << 20,
+        &protocol,
+        &device_name,
+        &master_server_addr,
+    )?;
+    store.health_check()?;
+
+    let key = unique_key();
+    let value = format!("mooncake-rust-smoke-value-{key}").into_bytes();
+
+    store.put(&key, &value, None)?;
+    assert!(store.is_exist(&key)?);
+    assert_eq!(store.get_size(&key)? as usize, value.len());
+    assert_eq!(store.get(&key)?, value);
+    assert!(!store.get_hostname()?.is_empty());
+
+    store.remove(&key, true)?;
+    assert!(!store.is_exist(&key)?);
+
+    Ok(())
+}
diff --git a/mooncake-store/src/CMakeLists.txt b/mooncake-store/src/CMakeLists.txt
index 06fcd2bd..85ca4c5f 100644
--- a/mooncake-store/src/CMakeLists.txt
+++ b/mooncake-store/src/CMakeLists.txt
@@ -9,13 +9,17 @@ set(MOONCAKE_STORE_SOURCES
     types.cpp
     master_client.cpp
     utils.cpp
+    mmap_arena.cpp
     master_metric_manager.cpp
     storage_backend.cpp
     thread_pool.cpp
     etcd_helper.cpp
     segment.cpp
     transfer_task.cpp
+    tenant_quota.cpp
+    tenant_quota_policy_store.cpp
     rpc_service.cpp
+    master_admin_service.cpp
     offset_allocator.cpp
     posix_file.cpp
     client_buffer.cpp
@@ -26,6 +30,7 @@ set(MOONCAKE_STORE_SOURCES
     http_metadata_server.cpp
     file_storage.cpp
     serialize/serializer.cpp
+    storage/distributed/distributed_storage_backend.cpp
     ha/leadership/leader_coordinator_factory.cpp
     ha/leadership/backends/etcd/etcd_leader_coordinator.cpp
     ha/common/redis/redis_connection.cpp
@@ -54,9 +59,13 @@ set(MOONCAKE_STORE_SOURCES
     hot_standby_service.cpp
     standby_state_machine.cpp
     ha_metric_manager.cpp
-    store_c.cpp)
+    store_c.cpp
+    memory_alloc.cpp
+    ssd_register_client.cpp
+    engram/engram_store.cpp)
 
 set(EXTRA_LIBS "")
+set(SPDK_STATIC_LIBS "")
 
 # Find AWS SDK
 find_package(AWSSDK QUIET COMPONENTS s3)
@@ -111,6 +120,7 @@ endif()
 if(USE_3FS)
   add_subdirectory(hf3fs)
   list(APPEND MOONCAKE_STORE_SOURCES ${HF3FS_SOURCES})
+  list(APPEND MOONCAKE_STORE_SOURCES storage/distributed/hf3fs_adapter.cpp)
   find_library(
     HF3FS_API_LIB hf3fs_api_shared
     PATHS /usr/lib
@@ -118,7 +128,7 @@ if(USE_3FS)
   if(NOT HF3FS_API_LIB)
     message(FATAL_ERROR "hf3fs_api_shared library not found in /usr/lib")
   endif()
-  set(EXTRA_LIBS ${HF3FS_API_LIB})
+  list(APPEND EXTRA_LIBS ${HF3FS_API_LIB})
 endif()
 
 # io_uring support (auto-detected)
@@ -137,31 +147,126 @@ if(STORE_USE_REDIS)
   list(APPEND EXTRA_LIBS ${MOONCAKE_STORE_HIREDIS_LIBRARY})
 endif()
 
+list(APPEND MOONCAKE_STORE_SOURCES k8s_lease_helper.cpp)
+if(STORE_USE_K8S_LEASE)
+  list(APPEND MOONCAKE_STORE_SOURCES
+       ha/leadership/backends/k8s/k8s_leader_coordinator.cpp)
+  list(APPEND EXTRA_LIBS ${K8S_LEASE_WRAPPER_LIB})
+  set_source_files_properties(
+    k8s_lease_helper.cpp
+    PROPERTIES
+      OBJECT_DEPENDS
+      "${CMAKE_BINARY_DIR}/mooncake-common/k8s-lease/libk8s_lease_wrapper.h")
+endif()
+
+if(USE_NOF)
+  add_subdirectory(spdk)
+  list(APPEND MOONCAKE_STORE_SOURCES ${SPDK_WRAPPER_SOURCES})
+
+  set(SPDK_INCLUDE_DIR
+      "/usr/local/include"
+      CACHE PATH "SPDK include directory")
+  set(DPDK_INCLUDE_DIR
+      "/usr/local/include"
+      CACHE PATH "DPDK include directory")
+  set(SPDK_LIB_DIR
+      "/usr/local/lib"
+      CACHE PATH "SPDK static library directory")
+  set(DPDK_LIB_DIR
+      "/usr/local/lib"
+      CACHE PATH "DPDK static library directory")
+
+  set(SPDK_STATIC_LIBS
+      -Wl,--whole-archive
+      ${SPDK_LIB_DIR}/libspdk_event.a
+      ${SPDK_LIB_DIR}/libspdk_rdma.a
+      ${SPDK_LIB_DIR}/libspdk_nvme.a
+      ${SPDK_LIB_DIR}/libspdk_sock_posix.a
+      -Wl,--no-whole-archive
+      ${SPDK_LIB_DIR}/libspdk_vfio_user.a
+      ${SPDK_LIB_DIR}/libspdk_dma.a
+      ${SPDK_LIB_DIR}/libspdk_init.a
+      ${SPDK_LIB_DIR}/libspdk_env_dpdk.a
+      ${SPDK_LIB_DIR}/libspdk_log.a
+      ${SPDK_LIB_DIR}/libspdk_util.a
+      ${SPDK_LIB_DIR}/libspdk_blob.a
+      ${SPDK_LIB_DIR}/libspdk_blobfs.a
+      ${SPDK_LIB_DIR}/libspdk_bdev.a
+      ${SPDK_LIB_DIR}/libspdk_conf.a
+      ${SPDK_LIB_DIR}/libspdk_json.a
+      ${SPDK_LIB_DIR}/libspdk_jsonrpc.a
+      ${SPDK_LIB_DIR}/libspdk_sock.a
+      ${SPDK_LIB_DIR}/libspdk_thread.a
+      ${SPDK_LIB_DIR}/libspdk_trace.a
+      ${SPDK_LIB_DIR}/libspdk_vmd.a
+      ${SPDK_LIB_DIR}/libspdk_rpc.a
+      ${DPDK_LIB_DIR}/librte_eal.a
+      ${DPDK_LIB_DIR}/librte_mempool.a
+      ${DPDK_LIB_DIR}/librte_ring.a
+      ${DPDK_LIB_DIR}/librte_bus_pci.a
+      ${DPDK_LIB_DIR}/librte_pci.a
+      ${DPDK_LIB_DIR}/librte_timer.a
+      ${DPDK_LIB_DIR}/librte_telemetry.a
+      ${DPDK_LIB_DIR}/librte_kvargs.a
+      ${DPDK_LIB_DIR}/librte_mbuf.a
+      -Wl,--whole-archive
+      ${DPDK_LIB_DIR}/librte_mempool_ring.a
+      -Wl,--no-whole-archive
+      ${DPDK_LIB_DIR}/librte_power.a
+      ${DPDK_LIB_DIR}/librte_ethdev.a
+      ${DPDK_LIB_DIR}/librte_net.a
+      ${DPDK_LIB_DIR}/librte_vhost.a
+      ${DPDK_LIB_DIR}/librte_cryptodev.a
+      ${DPDK_LIB_DIR}/librte_hash.a
+      ${DPDK_LIB_DIR}/librte_rcu.a
+      pthread
+      dl
+      numa
+      rt
+      m
+      uuid
+      crypto
+      aio
+      z
+      elf
+      ibverbs
+      rdmacm
+      ssl
+      bsd)
+endif()
+
 # The cache_allocator library
 include_directories(${Python3_INCLUDE_DIRS})
 add_library(mooncake_store ${MOONCAKE_STORE_SOURCES})
 target_include_directories(mooncake_store PUBLIC ${XXHASH_INCLUDE_DIR})
+if(USE_NOF)
+  target_include_directories(mooncake_store PRIVATE ${SPDK_INCLUDE_DIR}
+                                                    ${DPDK_INCLUDE_DIR})
+endif()
 target_link_libraries(mooncake_store PUBLIC ${XXHASH_LIBRARY})
 if(STORE_USE_REDIS)
-    target_include_directories(mooncake_store PRIVATE
-        ${MOONCAKE_STORE_HIREDIS_INCLUDE_DIR})
+  target_include_directories(mooncake_store
+                             PRIVATE ${MOONCAKE_STORE_HIREDIS_INCLUDE_DIR})
 endif()
-# Note: transfer_engine is PRIVATE to avoid propagating its dependencies (e.g., vendor-specific hardware)
-# to targets that don't need it (e.g., mooncake_master).
-# Targets that need transfer_engine should link it explicitly.
-target_link_libraries(mooncake_store
-    PUBLIC
-        cachelib_memory_allocator
-        ${ETCD_WRAPPER_LIB}
-        glog::glog
-        gflags::gflags
-        ${EXTRA_LIBS}
-        asio_shared
-    PRIVATE
-        transfer_engine
-)
+# Note: transfer_engine is PRIVATE to avoid propagating its dependencies (e.g.,
+# vendor-specific hardware) to targets that don't need it (e.g.,
+# mooncake_master). Targets that need transfer_engine should link it explicitly.
+target_link_libraries(
+  mooncake_store
+  PUBLIC cachelib_memory_allocator
+         ${ETCD_WRAPPER_LIB}
+         glog::glog
+         gflags::gflags
+         ${EXTRA_LIBS}
+         ${SPDK_STATIC_LIBS}
+         yaml-cpp
+         asio_shared
+  PRIVATE transfer_engine)
 if(STORE_USE_ETCD)
-    add_dependencies(mooncake_store build_etcd_wrapper)
+  add_dependencies(mooncake_store build_etcd_wrapper)
+endif()
+if(STORE_USE_K8S_LEASE)
+  add_dependencies(mooncake_store build_k8s_lease_wrapper)
 endif()
 
 if(URING_LIB AND URING_INCLUDE)
@@ -194,6 +299,7 @@ target_link_libraries(
           pthread
           ibverbs
           mooncake_common
+          JsonCpp::JsonCpp
           ${ETCD_WRAPPER_LIB}
           ${MASTER_EXTRA_LIBS}
           asio_shared)
@@ -210,23 +316,22 @@ target_link_libraries(mooncake_client PRIVATE mooncake_store transfer_engine
 
 # Optimize binary sizes only in Release mode
 string(TOUPPER "${CMAKE_BUILD_TYPE}" CMAKE_BUILD_TYPE_UPPER)
-if (CMAKE_BUILD_TYPE_UPPER STREQUAL "RELEASE")
+if(CMAKE_BUILD_TYPE_UPPER STREQUAL "RELEASE")
   target_compile_options(mooncake_master PRIVATE -Os)
   target_link_options(mooncake_master PRIVATE -Os -s)
   target_compile_options(mooncake_client PRIVATE -Os)
   target_link_options(mooncake_client PRIVATE -Os -s)
 endif()
 
-
 # GPU runtime library for D2H staging in PutToLocalFile / OffloadObjects.
-# transfer_engine is PRIVATE-linked, so its CUDA/HIP/Ascend dependencies
-# are not propagated; we must detect and link them independently.
+# transfer_engine is PRIVATE-linked, so its CUDA/HIP/Ascend dependencies are not
+# propagated; we must detect and link them independently.
 #
-# Auto-detect each toolkit regardless of global USE_CUDA/USE_HIP flags,
-# because USE_CUDA may be OFF even when GPU pointers are present
-# (e.g. WITH_NVIDIA_PEERMEM=ON uses nvidia-peermem for RDMA without cudart).
-# Each detected toolkit gets both link libraries AND compile definitions,
-# so that gpu_staging_utils.h / pinned_buffer_pool.h pick the correct backend.
+# Auto-detect each toolkit regardless of global USE_CUDA/USE_HIP flags, because
+# USE_CUDA may be OFF even when GPU pointers are present (e.g.
+# WITH_NVIDIA_PEERMEM env var uses nvidia-peermem for RDMA without cudart). Each
+# detected toolkit gets both link libraries AND compile definitions, so that
+# gpu_staging_utils.h / pinned_buffer_pool.h pick the correct backend.
 #
 # NOTE: mooncake_store is a static library (.a). External consumers (Go CGo,
 # Python pybind) that link it must also link the GPU runtime (e.g. -lcudart).
@@ -234,30 +339,33 @@ endif()
 
 find_package(CUDAToolkit QUIET)
 if(CUDAToolkit_FOUND)
-    message(STATUS "mooncake_store: CUDAToolkit detected, enabling D2H staging")
-    target_compile_definitions(mooncake_store PRIVATE USE_CUDA)
-    target_compile_definitions(mooncake_client PRIVATE USE_CUDA)
-    target_include_directories(mooncake_store PRIVATE ${CUDAToolkit_INCLUDE_DIRS})
-    target_include_directories(mooncake_client PRIVATE ${CUDAToolkit_INCLUDE_DIRS})
-    target_link_libraries(mooncake_store PRIVATE CUDA::cudart)
-    target_link_libraries(mooncake_client PRIVATE CUDA::cudart)
+  message(STATUS "mooncake_store: CUDAToolkit detected, enabling D2H staging")
+  target_compile_definitions(mooncake_store PRIVATE USE_CUDA)
+  target_compile_definitions(mooncake_client PRIVATE USE_CUDA)
+  target_include_directories(mooncake_store PRIVATE ${CUDAToolkit_INCLUDE_DIRS})
+  target_include_directories(mooncake_client
+                             PRIVATE ${CUDAToolkit_INCLUDE_DIRS})
+  target_link_libraries(mooncake_store PRIVATE CUDA::cudart)
+  target_link_libraries(mooncake_client PRIVATE CUDA::cudart)
 endif()
 
 if(NOT CUDAToolkit_FOUND)
-    find_package(hip QUIET)
-    if(hip_FOUND)
-        message(STATUS "mooncake_store: HIP detected, enabling D2H staging")
-        target_compile_definitions(mooncake_store PRIVATE USE_HIP)
-        target_compile_definitions(mooncake_client PRIVATE USE_HIP)
-        target_link_libraries(mooncake_store PRIVATE hip::host)
-        target_link_libraries(mooncake_client PRIVATE hip::host)
-    endif()
+  find_package(hip QUIET)
+  if(hip_FOUND)
+    message(STATUS "mooncake_store: HIP detected, enabling D2H staging")
+    target_compile_definitions(mooncake_store PRIVATE USE_HIP)
+    target_compile_definitions(mooncake_client PRIVATE USE_HIP)
+    target_link_libraries(mooncake_store PRIVATE hip::host)
+    target_link_libraries(mooncake_client PRIVATE hip::host)
+  endif()
 endif()
 
-if(USE_ASCEND OR USE_ASCEND_DIRECT OR USE_UBSHMEM)
-    target_include_directories(mooncake_store PRIVATE $ENV{ASCEND_PATH}/include)
-    target_link_libraries(mooncake_store PRIVATE ascendcl)
-    target_include_directories(mooncake_client PRIVATE $ENV{ASCEND_PATH}/include)
-    target_link_libraries(mooncake_client PRIVATE ascendcl)
+if(USE_ASCEND
+   OR USE_ASCEND_DIRECT
+   OR USE_UBSHMEM)
+  target_include_directories(mooncake_store PRIVATE $ENV{ASCEND_PATH}/include)
+  target_link_libraries(mooncake_store PRIVATE ascendcl)
+  target_include_directories(mooncake_client PRIVATE $ENV{ASCEND_PATH}/include)
+  target_link_libraries(mooncake_client PRIVATE ascendcl)
 endif()
 install(TARGETS mooncake_master mooncake_client DESTINATION bin)
diff --git a/mooncake-store/src/aligned_client_buffer.cpp b/mooncake-store/src/aligned_client_buffer.cpp
index 6ac7ea90..232539d8 100644
--- a/mooncake-store/src/aligned_client_buffer.cpp
+++ b/mooncake-store/src/aligned_client_buffer.cpp
@@ -4,10 +4,78 @@
 #include <sys/mman.h>
 #include <cstdlib>
 #include <cstring>
+#include <string_view>
 
 #include "utils.h"
 
 namespace mooncake {
+namespace {
+constexpr std::string_view kAscendProtocol = "ascend";
+constexpr std::string_view kUbshmemProtocol = "ubshmem";
+
+bool UseProtocolAllocator(const std::string& protocol) {
+    return protocol == kAscendProtocol || protocol == kUbshmemProtocol;
+}
+
+void FreeAlignedBuffer(void* buffer, size_t size, const std::string& protocol,
+                       bool use_hugepage) {
+    if (use_hugepage) {
+        free_buffer_mmap_memory(buffer, size);
+    } else if (UseProtocolAllocator(protocol)) {
+        free_memory(protocol, buffer);
+    } else {
+        free(buffer);
+    }
+}
+
+void* AllocateProtocolAlignedBuffer(size_t aligned_size,
+                                    const std::string& protocol) {
+    void* aligned_buffer = allocate_buffer_allocator_memory(
+        aligned_size, protocol,
+        AlignedClientBufferAllocator::kDirectIOAlignment);
+    if (!aligned_buffer) {
+        LOG(ERROR) << "AlignedClientBufferAllocator: failed to allocate "
+                   << "protocol-aware memory of size " << aligned_size
+                   << " for protocol " << protocol;
+    }
+    return aligned_buffer;
+}
+
+void* AllocateHugepageAlignedBuffer(size_t aligned_size) {
+    void* aligned_buffer = allocate_buffer_mmap_memory(
+        aligned_size, AlignedClientBufferAllocator::kDirectIOAlignment);
+    if (!aligned_buffer) {
+        LOG(ERROR) << "AlignedClientBufferAllocator: failed to allocate "
+                   << "hugepage memory of size " << aligned_size;
+    }
+    return aligned_buffer;
+}
+
+void* AllocatePosixAlignedBuffer(size_t aligned_size) {
+    void* aligned_buffer = nullptr;
+    int ret = posix_memalign(&aligned_buffer,
+                             AlignedClientBufferAllocator::kDirectIOAlignment,
+                             aligned_size);
+    if (ret != 0) {
+        LOG(ERROR) << "AlignedClientBufferAllocator: posix_memalign failed "
+                   << "with error " << ret << " (" << strerror(ret) << ")";
+        return nullptr;
+    }
+    memset(aligned_buffer, 0, aligned_size);
+    return aligned_buffer;
+}
+
+void* AllocateAlignedBuffer(size_t aligned_size, const std::string& protocol,
+                            bool use_mmap_hugepage) {
+    if (UseProtocolAllocator(protocol)) {
+        return AllocateProtocolAlignedBuffer(aligned_size, protocol);
+    }
+    if (use_mmap_hugepage) {
+        return AllocateHugepageAlignedBuffer(aligned_size);
+    }
+    return AllocatePosixAlignedBuffer(aligned_size);
+}
+}  // namespace
 
 std::shared_ptr<AlignedClientBufferAllocator>
 AlignedClientBufferAllocator::create(size_t size, const std::string& protocol,
@@ -21,40 +89,20 @@ AlignedClientBufferAllocator::create(size_t size, const std::string& protocol,
     // Align size up to kDirectIOAlignment
     size_t aligned_size = align_up(size, kDirectIOAlignment);
 
-    void* aligned_buffer = nullptr;
-
-    if (use_hugepage) {
-        // Use hugepage allocation (already aligned)
-        aligned_buffer =
-            allocate_buffer_mmap_memory(aligned_size, kDirectIOAlignment);
-        if (!aligned_buffer) {
-            LOG(ERROR) << "AlignedClientBufferAllocator: failed to allocate "
-                       << "hugepage memory of size " << aligned_size;
-            return nullptr;
-        }
-    } else {
-        // Use posix_memalign for 4096-byte alignment
-        int ret =
-            posix_memalign(&aligned_buffer, kDirectIOAlignment, aligned_size);
-        if (ret != 0) {
-            LOG(ERROR) << "AlignedClientBufferAllocator: posix_memalign failed "
-                       << "with error " << ret << " (" << strerror(ret) << ")";
-            return nullptr;
-        }
-
-        // Zero-initialize the allocated memory
-        memset(aligned_buffer, 0, aligned_size);
+    const bool use_protocol_allocator = UseProtocolAllocator(protocol);
+    const bool use_mmap_hugepage = use_hugepage && !use_protocol_allocator;
+    void* aligned_buffer =
+        AllocateAlignedBuffer(aligned_size, protocol, use_mmap_hugepage);
+    if (!aligned_buffer) {
+        return nullptr;
     }
 
     // Verify alignment
     if (reinterpret_cast<uintptr_t>(aligned_buffer) % kDirectIOAlignment != 0) {
         LOG(ERROR) << "AlignedClientBufferAllocator: allocated buffer is not "
                    << "aligned to " << kDirectIOAlignment << " bytes";
-        if (use_hugepage) {
-            free_buffer_mmap_memory(aligned_buffer, aligned_size);
-        } else {
-            free(aligned_buffer);
-        }
+        FreeAlignedBuffer(aligned_buffer, aligned_size, protocol,
+                          use_mmap_hugepage);
         return nullptr;
     }
 
@@ -65,7 +113,7 @@ AlignedClientBufferAllocator::create(size_t size, const std::string& protocol,
     // Use custom deleter to properly free the aligned memory
     return std::shared_ptr<AlignedClientBufferAllocator>(
         new AlignedClientBufferAllocator(aligned_buffer, aligned_size, protocol,
-                                         use_hugepage));
+                                         use_mmap_hugepage));
 }
 
 AlignedClientBufferAllocator::AlignedClientBufferAllocator(
@@ -73,7 +121,8 @@ AlignedClientBufferAllocator::AlignedClientBufferAllocator(
     bool use_hugepage)
     : ClientBufferAllocator(aligned_buffer, size, protocol),
       owns_memory_(true),
-      allocated_size_(size) {
+      allocated_size_(size),
+      protocol_(protocol) {
     // Store hugepage flag in parent class's use_hugepage_ member
     // We need this for proper cleanup
     use_hugepage_ = use_hugepage;
@@ -88,12 +137,14 @@ AlignedClientBufferAllocator::~AlignedClientBufferAllocator() {
             LOG(INFO)
                 << "AlignedClientBufferAllocator: freeing hugepage memory "
                 << "at " << buffer_ << " (" << allocated_size_ << " bytes)";
-            free_buffer_mmap_memory(buffer_, allocated_size_);
+            FreeAlignedBuffer(buffer_, allocated_size_, protocol_,
+                              use_hugepage_);
         } else {
             LOG(INFO) << "AlignedClientBufferAllocator: freeing aligned memory "
                       << "at " << buffer_ << " (" << allocated_size_
                       << " bytes)";
-            free(buffer_);
+            FreeAlignedBuffer(buffer_, allocated_size_, protocol_,
+                              use_hugepage_);
         }
         buffer_ = nullptr;
     }
diff --git a/mooncake-store/src/allocator.cpp b/mooncake-store/src/allocator.cpp
index 72ef4e3a..23311b83 100644
--- a/mooncake-store/src/allocator.cpp
+++ b/mooncake-store/src/allocator.cpp
@@ -73,12 +73,14 @@ std::ostream& operator<<(std::ostream& os, const AllocatedBuffer& buffer) {
 // Removed allocated_bytes parameter and member initialization
 CachelibBufferAllocator::CachelibBufferAllocator(std::string segment_name,
                                                  size_t base, size_t size,
-                                                 std::string transport_endpoint)
+                                                 std::string transport_endpoint,
+                                                 ReplicaType replica_type)
     : segment_name_(segment_name),
       base_(base),
       total_size_(size),
       cur_size_(0),
-      transport_endpoint_(std::move(transport_endpoint)) {
+      transport_endpoint_(std::move(transport_endpoint)),
+      replica_type_(replica_type) {
     VLOG(1) << "initializing_buffer_allocator segment_name=" << segment_name
             << " base_address=" << reinterpret_cast<void*>(base)
             << " size=" << size;
@@ -110,8 +112,13 @@ CachelibBufferAllocator::CachelibBufferAllocator(std::string segment_name,
 }
 
 CachelibBufferAllocator::~CachelibBufferAllocator() {
-    MasterMetricManager::instance().dec_allocated_mem_size(segment_name_,
-                                                           cur_size_);
+    if (replica_type_ == ReplicaType::MEMORY) {
+        MasterMetricManager::instance().dec_allocated_mem_size(segment_name_,
+                                                               cur_size_);
+    } else if (replica_type_ == ReplicaType::NOF_SSD) {
+        MasterMetricManager::instance().dec_allocated_nof_size(segment_name_,
+                                                               cur_size_);
+    }
 };
 
 std::unique_ptr<AllocatedBuffer> CachelibBufferAllocator::allocate(
@@ -137,7 +144,13 @@ std::unique_ptr<AllocatedBuffer> CachelibBufferAllocator::allocate(
     VLOG(1) << "allocation_succeeded size=" << size
             << " segment=" << segment_name_ << " address=" << buffer;
     cur_size_.fetch_add(size);
-    MasterMetricManager::instance().inc_allocated_mem_size(segment_name_, size);
+    if (replica_type_ == ReplicaType::MEMORY) {
+        MasterMetricManager::instance().inc_allocated_mem_size(segment_name_,
+                                                               size);
+    } else if (replica_type_ == ReplicaType::NOF_SSD) {
+        MasterMetricManager::instance().inc_allocated_nof_size(segment_name_,
+                                                               size);
+    }
     return std::make_unique<AllocatedBuffer>(shared_from_this(), buffer, size);
 }
 
@@ -151,8 +164,13 @@ void CachelibBufferAllocator::deallocate(AllocatedBuffer* handle) {
         size_t freed_size =
             handle->size_;  // Store size before handle might become invalid
         cur_size_.fetch_sub(freed_size);
-        MasterMetricManager::instance().dec_allocated_mem_size(segment_name_,
-                                                               freed_size);
+        if (replica_type_ == ReplicaType::MEMORY) {
+            MasterMetricManager::instance().dec_allocated_mem_size(
+                segment_name_, freed_size);
+        } else if (replica_type_ == ReplicaType::NOF_SSD) {
+            MasterMetricManager::instance().dec_allocated_nof_size(
+                segment_name_, freed_size);
+        }
         VLOG(1) << "deallocation_succeeded address=" << handle->buffer_ptr_
                 << " size=" << freed_size << " segment=" << segment_name_;
     } catch (const std::exception& e) {
@@ -165,12 +183,14 @@ void CachelibBufferAllocator::deallocate(AllocatedBuffer* handle) {
 // OffsetBufferAllocator implementation
 OffsetBufferAllocator::OffsetBufferAllocator(std::string segment_name,
                                              size_t base, size_t size,
-                                             std::string transport_endpoint)
+                                             std::string transport_endpoint,
+                                             ReplicaType replica_type)
     : segment_name_(segment_name),
       base_(base),
       total_size_(size),
       cur_size_(0),
-      transport_endpoint_(std::move(transport_endpoint)) {
+      transport_endpoint_(std::move(transport_endpoint)),
+      replica_type_(replica_type) {
     VLOG(1) << "initializing_offset_buffer_allocator segment_name="
             << segment_name << " base_address=" << reinterpret_cast<void*>(base)
             << " size=" << size;
@@ -205,8 +225,13 @@ OffsetBufferAllocator::OffsetBufferAllocator(std::string segment_name,
 }
 
 OffsetBufferAllocator::~OffsetBufferAllocator() {
-    MasterMetricManager::instance().dec_allocated_mem_size(segment_name_,
-                                                           cur_size_);
+    if (replica_type_ == ReplicaType::MEMORY) {
+        MasterMetricManager::instance().dec_allocated_mem_size(segment_name_,
+                                                               cur_size_);
+    } else if (replica_type_ == ReplicaType::NOF_SSD) {
+        MasterMetricManager::instance().dec_allocated_nof_size(segment_name_,
+                                                               cur_size_);
+    }
 };
 
 std::unique_ptr<AllocatedBuffer> OffsetBufferAllocator::allocate(size_t size) {
@@ -244,7 +269,13 @@ std::unique_ptr<AllocatedBuffer> OffsetBufferAllocator::allocate(size_t size) {
     }
 
     cur_size_.fetch_add(size);
-    MasterMetricManager::instance().inc_allocated_mem_size(segment_name_, size);
+    if (replica_type_ == ReplicaType::MEMORY) {
+        MasterMetricManager::instance().inc_allocated_mem_size(segment_name_,
+                                                               size);
+    } else if (replica_type_ == ReplicaType::NOF_SSD) {
+        MasterMetricManager::instance().inc_allocated_nof_size(segment_name_,
+                                                               size);
+    }
     return allocated_buffer;
 }
 
@@ -255,8 +286,13 @@ void OffsetBufferAllocator::deallocate(AllocatedBuffer* handle) {
         size_t freed_size = handle->size();
         handle->offset_handle_.reset();
         cur_size_.fetch_sub(freed_size);
-        MasterMetricManager::instance().dec_allocated_mem_size(segment_name_,
-                                                               freed_size);
+        if (replica_type_ == ReplicaType::MEMORY) {
+            MasterMetricManager::instance().dec_allocated_mem_size(
+                segment_name_, freed_size);
+        } else if (replica_type_ == ReplicaType::NOF_SSD) {
+            MasterMetricManager::instance().dec_allocated_nof_size(
+                segment_name_, freed_size);
+        }
         VLOG(1) << "deallocation_succeeded address=" << handle->data()
                 << " size=" << freed_size << " segment=" << segment_name_;
     } catch (const std::exception& e) {
diff --git a/mooncake-store/src/client_buffer.cpp b/mooncake-store/src/client_buffer.cpp
index c1f62816..e846d4c0 100644
--- a/mooncake-store/src/client_buffer.cpp
+++ b/mooncake-store/src/client_buffer.cpp
@@ -13,9 +13,10 @@
 namespace mooncake {
 
 std::shared_ptr<ClientBufferAllocator> ClientBufferAllocator::create(
-    size_t size, const std::string& protocol, bool use_hugepage) {
+    size_t size, const std::string& protocol, bool use_hugepage,
+    bool use_spdk_dma) {
     return std::shared_ptr<ClientBufferAllocator>(
-        new ClientBufferAllocator(size, protocol, use_hugepage));
+        new ClientBufferAllocator(size, protocol, use_hugepage, use_spdk_dma));
 }
 
 std::shared_ptr<ClientBufferAllocator> ClientBufferAllocator::create(
@@ -26,8 +27,12 @@ std::shared_ptr<ClientBufferAllocator> ClientBufferAllocator::create(
 
 ClientBufferAllocator::ClientBufferAllocator(size_t size,
                                              const std::string& protocol,
-                                             bool use_hugepage)
-    : protocol(protocol), buffer_size_(size), use_hugepage_(use_hugepage) {
+                                             bool use_hugepage,
+                                             bool use_spdk_dma)
+    : buffer_size_(size),
+      use_hugepage_(use_hugepage),
+      protocol(protocol),
+      use_spdk_dma_(use_spdk_dma) {
     if (size == 0) {
         buffer_ = nullptr;
         allocator_ = nullptr;
@@ -38,7 +43,8 @@ ClientBufferAllocator::ClientBufferAllocator(size_t size,
     if (use_hugepage_) {
         buffer_ = allocate_buffer_mmap_memory(size, alignment);
     } else {
-        buffer_ = allocate_buffer_allocator_memory(size, protocol, alignment);
+        buffer_ = allocate_buffer_allocator_memory(size, protocol, alignment,
+                                                   use_spdk_dma_);
     }
     if (!buffer_) {
         throw std::bad_alloc();
@@ -50,7 +56,7 @@ ClientBufferAllocator::ClientBufferAllocator(size_t size,
 
 ClientBufferAllocator::ClientBufferAllocator(void* addr, size_t size,
                                              const std::string& protocol)
-    : protocol(protocol), buffer_size_(size) {
+    : buffer_size_(size), protocol(protocol) {
     buffer_ = addr;
     is_external_memory_ = true;
     allocator_ = mooncake::offset_allocator::OffsetAllocator::create(
@@ -132,6 +138,8 @@ uint64_t calculate_total_size(const Replica::Descriptor& replica) {
         total_length = disk_descriptor.object_size;
     } else if (replica.is_local_disk_replica()) {
         total_length = replica.get_local_disk_descriptor().object_size;
+    } else if (replica.is_nof_replica()) {
+        total_length = replica.get_nof_descriptor().buffer_descriptor.size_;
     } else {
         total_length = replica.get_memory_descriptor().buffer_descriptor.size_;
     }
@@ -153,6 +161,10 @@ int allocateSlices(std::vector<Slice>& slices,
     } else if (replica.is_local_disk_replica()) {
         slices.emplace_back(
             Slice{buffer_ptr, replica.get_local_disk_descriptor().object_size});
+    } else if (replica.is_nof_replica()) {
+        auto& handle = replica.get_nof_descriptor().buffer_descriptor;
+        void* chunk_ptr = buffer_ptr;
+        slices.emplace_back(Slice{chunk_ptr, handle.size_});
     } else {
         // For memory-based replica, split into slices based on buffer
         // descriptors
diff --git a/mooncake-store/src/client_service.cpp b/mooncake-store/src/client_service.cpp
index 16c7dfd0..fe41c86f 100644
--- a/mooncake-store/src/client_service.cpp
+++ b/mooncake-store/src/client_service.cpp
@@ -2,16 +2,26 @@
 
 #include <glog/logging.h>
 
+#include "allocator.h"
+#include "segment.h"
+
 #include <csignal>
 #include <algorithm>
 #include <cassert>
+#include <cerrno>
 #include <chrono>
 #include <cstdint>
 #include <cstring>
 #include <cstdlib>
 #include <iomanip>
+#include <limits>
+#ifdef USE_NOF
+#include <numa.h>
+#endif
 #include <optional>
 #include <ranges>
+#include <span>
+#include <sched.h>
 #include <thread>
 #include <set>
 #include <ylt/struct_json/json_reader.h>
@@ -35,6 +45,214 @@ using gpu_staging::CopyDeviceToHost;
 using gpu_staging::IsDevicePointer;
 using gpu_staging::SetDevice;
 
+namespace {
+
+#ifdef USE_NOF
+std::optional<int> GetConfiguredNumaSocketId() {
+    const char* raw_value = std::getenv("MC_STORE_NUMA_SOCKET_ID");
+    if (!raw_value || raw_value[0] == '\0') {
+        return std::nullopt;
+    }
+
+    char* end_ptr = nullptr;
+    errno = 0;
+    long parsed = std::strtol(raw_value, &end_ptr, 10);
+    if (errno != 0 || end_ptr == raw_value ||
+        (end_ptr != nullptr && *end_ptr != '\0') || parsed < 0 ||
+        parsed > std::numeric_limits<int>::max()) {
+        LOG(WARNING) << "Invalid MC_STORE_NUMA_SOCKET_ID=" << raw_value
+                     << ", falling back to auto-detect";
+        return std::nullopt;
+    }
+
+    return static_cast<int>(parsed);
+}
+
+int GetCurrentNumaSocketId() {
+    if (numa_available() < 0) {
+        return 0;
+    }
+    int cpu = sched_getcpu();
+    if (cpu < 0) {
+        return 0;
+    }
+    int node = numa_node_of_cpu(cpu);
+    return node < 0 ? 0 : node;
+}
+#endif
+
+struct ContiguousSliceRange {
+    void* ptr = nullptr;
+    size_t size = 0;
+};
+
+std::optional<ContiguousSliceRange> GetContiguousSliceRange(
+    std::span<const Slice> slices) {
+    if (slices.empty()) {
+        return std::nullopt;
+    }
+
+    uintptr_t expected_ptr = 0;
+    uintptr_t start_ptr = 0;
+    size_t total_size = 0;
+    for (size_t i = 0; i < slices.size(); ++i) {
+        const auto& slice = slices[i];
+        if (slice.ptr == nullptr) {
+            return std::nullopt;
+        }
+
+        const auto current_ptr = reinterpret_cast<uintptr_t>(slice.ptr);
+        if (i == 0) {
+            start_ptr = current_ptr;
+            expected_ptr = current_ptr;
+        } else if (current_ptr != expected_ptr) {
+            return std::nullopt;
+        }
+
+        if (slice.size > std::numeric_limits<size_t>::max() - total_size) {
+            return std::nullopt;
+        }
+        if (slice.size > std::numeric_limits<uintptr_t>::max() - current_ptr) {
+            return std::nullopt;
+        }
+
+        total_size += slice.size;
+        expected_ptr = current_ptr + slice.size;
+    }
+
+    if (total_size == 0) {
+        return std::nullopt;
+    }
+
+    return ContiguousSliceRange{.ptr = reinterpret_cast<void*>(start_ptr),
+                                .size = total_size};
+}
+
+struct ReplicaTransferSummary {
+    size_t allocated_memory_replicas = 0;
+    size_t allocated_nof_replicas = 0;
+    size_t successful_memory_transfers = 0;
+    size_t successful_nof_transfers = 0;
+    size_t failed_memory_transfers = 0;
+    size_t failed_nof_transfers = 0;
+    ErrorCode first_error = ErrorCode::OK;
+
+    void RecordAllocatedReplica(const Replica::Descriptor& replica) {
+        if (replica.is_memory_replica()) {
+            ++allocated_memory_replicas;
+        } else if (replica.is_nof_replica()) {
+            ++allocated_nof_replicas;
+        }
+    }
+
+    void RecordSuccess(ReplicaType replica_type) {
+        if (replica_type == ReplicaType::MEMORY) {
+            ++successful_memory_transfers;
+        } else if (replica_type == ReplicaType::NOF_SSD) {
+            ++successful_nof_transfers;
+        }
+    }
+
+    void RecordFailure(ReplicaType replica_type, ErrorCode error) {
+        if (replica_type == ReplicaType::MEMORY) {
+            ++failed_memory_transfers;
+        } else if (replica_type == ReplicaType::NOF_SSD) {
+            ++failed_nof_transfers;
+        }
+        if (first_error == ErrorCode::OK) {
+            first_error = error;
+        }
+    }
+};
+
+bool HasExpectedReplicaAllocation(const ReplicateConfig& config,
+                                  const ReplicaTransferSummary& summary) {
+    if (config.nof_replica_num == 0) {
+        return summary.allocated_memory_replicas > 0;
+    }
+    if (DetermineReplicaWriteMode(config) ==
+        ReplicaWriteMode::FLEXIBLE_DUAL_REPLICA) {
+        return summary.allocated_memory_replicas +
+                   summary.allocated_nof_replicas >
+               0;
+    }
+    return summary.allocated_memory_replicas == config.replica_num &&
+           summary.allocated_nof_replicas == config.nof_replica_num;
+}
+
+// success describes whether the overall put should succeed. Reliable modes
+// require all allocated replicas to complete. Flexible dual-replica mode only
+// requires one replica type to succeed, so success may be true while
+// revoke_type is also set for the failed side.
+struct FinalizeDecision {
+    std::optional<ReplicaType> end_type;
+    std::optional<ReplicaType> revoke_type;
+    bool success = false;
+    ErrorCode error = ErrorCode::OK;
+};
+
+FinalizeDecision DetermineFinalizeDecision(
+    const ReplicateConfig& config, const ReplicaTransferSummary& summary) {
+    const auto write_mode = DetermineReplicaWriteMode(config);
+    const bool allocation_satisfied =
+        HasExpectedReplicaAllocation(config, summary);
+
+    if (write_mode != ReplicaWriteMode::FLEXIBLE_DUAL_REPLICA) {
+        const bool all_transfers_succeeded =
+            summary.successful_memory_transfers ==
+                summary.allocated_memory_replicas &&
+            summary.successful_nof_transfers ==
+                summary.allocated_nof_replicas &&
+            summary.failed_memory_transfers == 0 &&
+            summary.failed_nof_transfers == 0;
+        if (allocation_satisfied && all_transfers_succeeded) {
+            return {.end_type = ReplicaType::ALL,
+                    .revoke_type = std::nullopt,
+                    .success = true,
+                    .error = ErrorCode::OK};
+        }
+        return {.end_type = std::nullopt,
+                .revoke_type = ReplicaType::ALL,
+                .success = false,
+                .error = allocation_satisfied
+                             ? (summary.first_error == ErrorCode::OK
+                                    ? ErrorCode::TRANSFER_FAIL
+                                    : summary.first_error)
+                             : ErrorCode::NO_AVAILABLE_HANDLE};
+    }
+
+    const bool memory_succeeded = summary.successful_memory_transfers > 0;
+    const bool nof_succeeded = summary.successful_nof_transfers > 0;
+
+    if (memory_succeeded && nof_succeeded) {
+        return {.end_type = ReplicaType::ALL,
+                .revoke_type = std::nullopt,
+                .success = true,
+                .error = ErrorCode::OK};
+    }
+    if (memory_succeeded) {
+        return {.end_type = ReplicaType::MEMORY,
+                .revoke_type = ReplicaType::NOF_SSD,
+                .success = true,
+                .error = ErrorCode::OK};
+    }
+    if (nof_succeeded) {
+        return {.end_type = ReplicaType::NOF_SSD,
+                .revoke_type = ReplicaType::MEMORY,
+                .success = true,
+                .error = ErrorCode::OK};
+    }
+
+    return {.end_type = std::nullopt,
+            .revoke_type = ReplicaType::ALL,
+            .success = false,
+            .error = summary.first_error == ErrorCode::OK
+                         ? ErrorCode::NO_AVAILABLE_HANDLE
+                         : summary.first_error};
+}
+
+}  // namespace
+
 [[nodiscard]] size_t CalculateSliceSize(const std::vector<Slice>& slices) {
     size_t slice_size = 0;
     for (const auto& slice : slices) {
@@ -54,11 +272,13 @@ using gpu_staging::SetDevice;
 Client::Client(const std::string& local_hostname,
                const std::string& metadata_connstring,
                const std::string& protocol,
-               const std::map<std::string, std::string>& labels)
+               const std::map<std::string, std::string>& labels,
+               const std::string& tenant_id)
     : client_id_(generate_uuid()),
       metrics_(ClientMetric::Create(merge_labels(labels))),
       master_client_(client_id_,
-                     metrics_ ? &metrics_->master_client_metric : nullptr),
+                     metrics_ ? &metrics_->master_client_metric : nullptr,
+                     tenant_id),
       local_hostname_(local_hostname),
       metadata_connstring_(metadata_connstring),
       protocol_(protocol),
@@ -98,55 +318,98 @@ Client::~Client() {
         leader_monitor_thread_.join();
     }
 
-    // Make a copy of mounted_segments_ to avoid modifying while iterating
-    std::vector<Segment> segments_to_unmount;
+    {
+        std::lock_guard<std::mutex> lock(graceful_unmount_timer_mutex_);
+        graceful_unmount_timer_stopping_ = true;
+    }
+    graceful_unmount_timer_cv_.notify_all();
+
+    // Stop queued timer/task callbacks before tearing down segment state.
+    task_running_ = false;
+    task_thread_pool_.stop();
+
+    // Make copies to avoid modifying while iterating
+    std::vector<Segment> mounted_segments_copy;
+    std::vector<Segment> gracefully_unmounting_segments_copy;
+    std::vector<std::pair<UUID, std::function<void(const UUID&)>>>
+        graceful_cleanup_callbacks;
     {
         std::lock_guard<std::mutex> lock(mounted_segments_mutex_);
-        segments_to_unmount.reserve(mounted_segments_.size());
+        mounted_segments_copy.reserve(mounted_segments_.size());
         for (auto& entry : mounted_segments_) {
-            segments_to_unmount.emplace_back(entry.second);
+            mounted_segments_copy.emplace_back(entry.second);
         }
+        for (auto& entry : gracefully_unmounting_segments_) {
+            gracefully_unmounting_segments_copy.emplace_back(entry.second);
+        }
+        graceful_cleanup_callbacks.reserve(
+            graceful_unmount_cleanup_callbacks_.size());
+        for (auto& entry : graceful_unmount_cleanup_callbacks_) {
+            graceful_cleanup_callbacks.emplace_back(entry.first,
+                                                    std::move(entry.second));
+        }
+        graceful_unmount_cleanup_callbacks_.clear();
     }
 
-    for (auto& segment : segments_to_unmount) {
+    // Unmount mounted segments: notify master + local cleanup
+    for (auto& segment : mounted_segments_copy) {
         auto result =
             UnmountSegment(reinterpret_cast<void*>(segment.base), segment.size);
         if (!result) {
-            LOG(ERROR) << "Failed to unmount segment: "
+            LOG(ERROR) << "Failed to unmount segment in destructor: "
                        << toString(result.error());
         }
     }
 
+    // Unregister gracefully unmounting segments: master already has timer
+    for (auto& segment : gracefully_unmounting_segments_copy) {
+        int rc = transfer_engine_->unregisterLocalMemory(
+            reinterpret_cast<void*>(segment.base));
+        if (rc != 0 && rc != ERR_ADDRESS_NOT_REGISTERED) {
+            LOG(ERROR) << "Failed to unregister transfer buffer in destructor: "
+                       << rc;
+        }
+    }
+
     // Clear any remaining segments
     {
         std::lock_guard<std::mutex> lock(mounted_segments_mutex_);
         mounted_segments_.clear();
+        gracefully_unmounting_segments_.clear();
     }
 
-    // Stop hot cache handler and hot cache
-    hot_cache_handler_.reset();
-    hot_cache_.reset();
+    for (auto& entry : graceful_cleanup_callbacks) {
+        if (entry.second) {
+            entry.second(entry.first);
+        }
+    }
 
-    // Stop task thread pool after task polling has stopped.
-    task_running_ = false;
-    task_thread_pool_.stop();
+    // Stop hot cache handler before unregistering/freeing its backing memory.
+    hot_cache_handler_.reset();
+    UnregisterLocalHotCacheMemory();
+    hot_cache_.reset();
 }
 
 static std::optional<bool> get_auto_discover() {
     const char* ev_ad = std::getenv("MC_MS_AUTO_DISC");
     if (ev_ad) {
-        int iv = std::stoi(ev_ad);
-        if (iv == 1) {
-            LOG(INFO) << "auto discovery set by env MC_MS_AUTO_DISC";
-            return true;
-        } else if (iv == 0) {
-            LOG(INFO) << "auto discovery not set by env MC_MS_AUTO_DISC";
-            return false;
-        } else {
-            LOG(WARNING)
-                << "invalid MC_MS_AUTO_DISC value: " << ev_ad
-                << ", should be 0 or 1, using default: auto discovery not set";
+        try {
+            int iv = std::stoi(ev_ad);
+            if (iv == 1) {
+                LOG(INFO) << "auto discovery set by env MC_MS_AUTO_DISC";
+                return true;
+            } else if (iv == 0) {
+                LOG(INFO) << "auto discovery not set by env MC_MS_AUTO_DISC";
+                return false;
+            }
+        } catch (const std::exception&) {
+            // A non-numeric or out-of-range value makes std::stoi throw; fall
+            // through to the warning below and use the default instead of
+            // letting the exception abort client initialization.
         }
+        LOG(WARNING)
+            << "invalid MC_MS_AUTO_DISC value: " << ev_ad
+            << ", should be 0 or 1, using default: auto discovery not set";
     }
     return std::nullopt;
 }
@@ -202,6 +465,11 @@ tl::expected<std::optional<ha::HABackendSpec>, ErrorCode> ParseHABackendSpec(
         return tl::make_unexpected(ErrorCode::INVALID_PARAMS);
     }
 
+    auto availability = ha::ValidateHABackendAvailability(backend_type.value());
+    if (availability != ErrorCode::OK) {
+        return tl::make_unexpected(availability);
+    }
+
     return std::optional<ha::HABackendSpec>{ha::HABackendSpec{
         .type = backend_type.value(),
         .connstring = master_server_entry.substr(delimiter_pos + 3),
@@ -314,7 +582,14 @@ ErrorCode Client::SwitchLeader(const ha::MasterView& target_view) {
 }
 
 void Client::LeaderMonitorThreadMain() {
-    constexpr auto kViewChangeTimeout = std::chrono::milliseconds(1000);
+    // WaitForViewChange now blocks on a backend watch and returns as soon as
+    // leadership actually changes, so this timeout no longer drives polling
+    // cadence -- it is only a periodic re-sync safety net that re-reads the
+    // current view in case a watch event was missed (e.g. a connection blip).
+    // Keeping it long (30s instead of 1s) collapses steady-state backend load
+    // from ~1 read/sec/client to ~1 read/30s/client without affecting failover
+    // responsiveness, which is driven by the watch.
+    constexpr auto kViewChangeTimeout = std::chrono::milliseconds(30000);
     constexpr auto kErrorRetryInterval = std::chrono::milliseconds(1000);
 
     while (leader_monitor_running_.load()) {
@@ -365,10 +640,23 @@ ErrorCode Client::InitTransferEngine(
     const std::string& local_hostname, const std::string& metadata_connstring,
     const std::string& protocol,
     const std::optional<std::string>& device_names) {
+    // TEs created through the Store entry (Client::Create ->
+    // InitTransferEngine) are tagged so ascend_direct can resolve a per-role
+    // link config (e.g. Store=RoCE/D2H, P2P=HCCS/D2D). The flag only needs to
+    // be live while the ascend transport is installed below; reset it on every
+    // exit path. Assumes TE inits are serialized within the process.
+    struct StoreTeInitGuard {
+        ~StoreTeInitGuard() { globalConfig().ascend_store_te_init = false; }
+    } store_te_init_guard;
+    if (protocol == "ascend") {
+        globalConfig().ascend_store_te_init = true;
+    }
+
     // Check if using TENT mode - TENT handles transport configuration
-    // internally
-    bool use_tent = (std::getenv("MC_USE_TENT") != nullptr) ||
-                    (std::getenv("MC_USE_TEV1") != nullptr);
+    // internally. Use the engine's own check rather than the raw env var:
+    // if Mooncake was built without USE_TENT, the env var has no effect on
+    // the TransferEngine, so the store must still install transports.
+    bool use_tent = transfer_engine_->isUsingTent();
 
     bool auto_discover = false;
     if (!use_tent) {
@@ -524,6 +812,30 @@ ErrorCode Client::InitTransferEngine(
                 LOG(ERROR) << "Failed to install CXL transport";
                 return ErrorCode::INTERNAL_ERROR;
             }
+        } else if (protocol == "cxi") {
+            try {
+                transport = transfer_engine_->installTransport("cxi", nullptr);
+            } catch (std::exception& e) {
+                LOG(ERROR) << "cxi_transport_install_failed error_message=\""
+                           << e.what() << "\"";
+            }
+        } else if (protocol == "ub") {
+            auto deviceName = device_names.value_or("bonding_dev_0");
+            LOG(ERROR) << "ub protocol entable device names is " << deviceName;
+            auto devices = splitString(deviceName, ',', true);
+            auto topology = transfer_engine_->getLocalTopology();
+            if (topology) {
+                topology->discover(devices);
+                LOG(INFO) << "Topology discovery complete with specified "
+                             "devices. Found "
+                          << topology->getHcaList().size() << " HCAs";
+            }
+            transport = transfer_engine_->installTransport("ub", nullptr);
+            if (!transport) {
+                LOG(ERROR) << "Failed to install ub transport with specified "
+                              "devices";
+                return ErrorCode::INTERNAL_ERROR;
+            }
         } else {
             LOG(ERROR) << "unsupported_protocol protocol=" << protocol;
             return ErrorCode::INVALID_PARAMS;
@@ -537,9 +849,17 @@ void Client::InitTransferSubmitter() {
     // Initialize TransferSubmitter after transfer engine is ready
     // Keep using logical local_hostname for name-based behaviors; endpoint is
     // used separately where needed.
+#ifdef USE_NOF
+    int numa_socket_id =
+        GetConfiguredNumaSocketId().value_or(GetCurrentNumaSocketId());
     transfer_submitter_ = std::make_unique<TransferSubmitter>(
-        *transfer_engine_, storage_backend_,
+        *transfer_engine_, storage_backend_, local_hostname_,
+        metrics_ ? &metrics_->transfer_metric : nullptr, numa_socket_id);
+#else
+    transfer_submitter_ = std::make_unique<TransferSubmitter>(
+        *transfer_engine_, storage_backend_, local_hostname_,
         metrics_ ? &metrics_->transfer_metric : nullptr);
+#endif
 }
 
 std::optional<std::shared_ptr<Client>> Client::Create(
@@ -547,9 +867,9 @@ std::optional<std::shared_ptr<Client>> Client::Create(
     const std::string& protocol, const std::optional<std::string>& device_names,
     const std::string& master_server_entry,
     const std::shared_ptr<TransferEngine>& transfer_engine,
-    std::map<std::string, std::string> labels) {
-    auto client = std::shared_ptr<Client>(
-        new Client(local_hostname, metadata_connstring, protocol, labels));
+    std::map<std::string, std::string> labels, const std::string& tenant_id) {
+    auto client = std::shared_ptr<Client>(new Client(
+        local_hostname, metadata_connstring, protocol, labels, tenant_id));
 
     ErrorCode err = client->ConnectToMaster(master_server_entry);
     if (err != ErrorCode::OK) {
@@ -727,9 +1047,14 @@ tl::expected<QueryResult, ErrorCode> Client::Query(
 
 std::vector<tl::expected<QueryResult, ErrorCode>> Client::BatchQuery(
     const std::vector<std::string>& object_keys) {
+    return BatchQuery(object_keys, master_client_.tenant_id());
+}
+
+std::vector<tl::expected<QueryResult, ErrorCode>> Client::BatchQuery(
+    const std::vector<std::string>& object_keys, const std::string& tenant_id) {
     std::chrono::steady_clock::time_point start_time =
         std::chrono::steady_clock::now();
-    auto response = master_client_.BatchGetReplicaList(object_keys);
+    auto response = master_client_.BatchGetReplicaList(object_keys, tenant_id);
 
     // Check if we got the expected number of responses
     if (response.size() != object_keys.size()) {
@@ -1067,8 +1392,21 @@ std::vector<tl::expected<void, ErrorCode>> Client::BatchGet(
         }
 
         // Submit transfer operation asynchronously
-        auto future = transfer_submitter_->submit(replica, slices_it->second,
-                                                  TransferRequest::READ);
+        std::optional<TransferFuture> future;
+        if (replica.is_nof_replica()) {
+            auto contiguous_range = GetContiguousSliceRange(slices_it->second);
+            if (!contiguous_range.has_value()) {
+                LOG(ERROR) << "NoF transfer requires contiguous slices";
+                results[i] = tl::unexpected(ErrorCode::INVALID_PARAMS);
+                continue;
+            }
+            future = transfer_submitter_->submit(
+                replica, slices_it->second, TransferRequest::READ,
+                contiguous_range->ptr, contiguous_range->size);
+        } else {
+            future = transfer_submitter_->submit(replica, slices_it->second,
+                                                 TransferRequest::READ);
+        }
         if (!future) {
             // Release cache block if submit failed
             if (hot_cache_ && cache_used) {
@@ -1184,6 +1522,10 @@ tl::expected<void, ErrorCode> Client::Put(const ObjectKey& key,
         client_cfg.preferred_segment = local_hostname_;
     }
 
+    if (hot_cache_) {
+        hot_cache_->RemoveHotKey(key);
+    }
+
     // Start put operation
     auto start_result = master_client_.PutStart(key, slice_lengths, client_cfg);
     if (!start_result) {
@@ -1202,6 +1544,11 @@ tl::expected<void, ErrorCode> Client::Put(const ObjectKey& key,
         return tl::unexpected(err);
     }
 
+    ReplicaTransferSummary transfer_summary;
+    for (const auto& replica : start_result.value()) {
+        transfer_summary.RecordAllocatedReplica(replica);
+    }
+
     // Record Put transfer latency (all replicas)
     auto t0_put = std::chrono::steady_clock::now();
 
@@ -1221,19 +1568,17 @@ tl::expected<void, ErrorCode> Client::Put(const ObjectKey& key,
     }
 
     for (const auto& replica : start_result.value()) {
-        if (replica.is_memory_replica()) {
+        if (replica.is_memory_replica() || replica.is_nof_replica()) {
             // Transfer data using allocated handles from all replicas
+            const auto replica_type = replica.is_memory_replica()
+                                          ? ReplicaType::MEMORY
+                                          : ReplicaType::NOF_SSD;
             ErrorCode transfer_err = TransferWrite(replica, slices);
             if (transfer_err != ErrorCode::OK) {
-                // Revoke put operation
-                auto revoke_result =
-                    master_client_.PutRevoke(key, ReplicaType::MEMORY);
-                if (!revoke_result) {
-                    LOG(ERROR) << "Failed to revoke put operation";
-                    return tl::unexpected(revoke_result.error());
-                }
-                return tl::unexpected(transfer_err);
+                transfer_summary.RecordFailure(replica_type, transfer_err);
+                continue;
             }
+            transfer_summary.RecordSuccess(replica_type);
         }
     }
 
@@ -1244,12 +1589,30 @@ tl::expected<void, ErrorCode> Client::Put(const ObjectKey& key,
         metrics_->transfer_metric.put_latency_us.observe(us_put);
     }
 
-    // End put operation
-    auto end_result = master_client_.PutEnd(key, ReplicaType::MEMORY);
-    if (!end_result) {
-        ErrorCode err = end_result.error();
-        LOG(ERROR) << "Failed to end put operation: " << err;
-        return tl::unexpected(err);
+    const auto finalize_decision =
+        DetermineFinalizeDecision(config, transfer_summary);
+
+    if (finalize_decision.end_type.has_value()) {
+        auto end_result =
+            master_client_.PutEnd(key, *finalize_decision.end_type);
+        if (!end_result) {
+            ErrorCode err = end_result.error();
+            LOG(ERROR) << "Failed to end put operation: " << err;
+            return tl::unexpected(err);
+        }
+    }
+
+    if (finalize_decision.revoke_type.has_value()) {
+        auto revoke_result =
+            master_client_.PutRevoke(key, *finalize_decision.revoke_type);
+        if (!revoke_result) {
+            LOG(ERROR) << "Failed to revoke put operation";
+            return tl::unexpected(revoke_result.error());
+        }
+    }
+
+    if (!finalize_decision.success) {
+        return tl::unexpected(finalize_decision.error);
     }
 
     return {};
@@ -1269,6 +1632,10 @@ tl::expected<void, ErrorCode> Client::Upsert(const ObjectKey& key,
         client_cfg.preferred_segment = local_hostname_;
     }
 
+    if (hot_cache_) {
+        hot_cache_->RemoveHotKey(key);
+    }
+
     // Start upsert operation
     auto start_result =
         master_client_.UpsertStart(key, slice_lengths, client_cfg);
@@ -1331,6 +1698,14 @@ tl::expected<void, ErrorCode> Client::Upsert(const ObjectKey& key,
         return tl::unexpected(err);
     }
 
+    // Success-side invalidation: a concurrent read between the pre-upsert
+    // RemoveHotKey() and UpsertEnd could have read the old value and submitted
+    // an async hot-cache fill with a still-valid token. Invalidate again now
+    // that the new value is committed so that stale fill cannot publish.
+    if (hot_cache_) {
+        hot_cache_->RemoveHotKey(key);
+    }
+
     return {};
 }
 
@@ -1376,23 +1751,34 @@ enum class PutOperationState {
 
 class PutOperation {
    public:
+    struct PendingTransferRecord {
+        ReplicaType replica_type;
+        TransferFuture future;
+
+        PendingTransferRecord(ReplicaType type,
+                              TransferFuture&& transfer_future)
+            : replica_type(type), future(std::move(transfer_future)) {}
+    };
+
     PutOperation(std::string_view k, const std::vector<Slice>& s)
         : key(k), slices(s) {
-        value_length = CalculateSliceSize(slices);
         // Initialize with a pending error state to ensure result is always set
         result = tl::unexpected(ErrorCode::INTERNAL_ERROR);
     }
 
     std::string key;
     std::vector<Slice> slices;
-    size_t value_length;
     std::vector<std::vector<Slice>> batched_slices;
 
     // Enhanced state tracking
     PutOperationState state = PutOperationState::PENDING;
     tl::expected<void, ErrorCode> result;
     std::vector<Replica::Descriptor> replicas;
-    std::vector<TransferFuture> pending_transfers;
+    std::vector<PendingTransferRecord> pending_transfers;
+
+    size_t requested_memory_replicas = 0;
+    size_t requested_nof_replicas = 0;
+    ReplicaTransferSummary transfer_summary;
 
     // Error context for debugging
     std::optional<std::string> failure_context;
@@ -1404,6 +1790,15 @@ class PutOperation {
         failure_context.reset();
     }
 
+    void SetTerminalError(ErrorCode error, PutOperationState terminal_state,
+                          const std::string& context = "") {
+        state = terminal_state;
+        result = tl::unexpected(error);
+        if (!context.empty()) {
+            failure_context = context;
+        }
+    }
+
     void SetError(ErrorCode error, const std::string& context = "") {
         result = tl::unexpected(error);
         if (!context.empty()) {
@@ -1421,6 +1816,36 @@ class PutOperation {
         }
     }
 
+    void AppendFailureContext(const std::string& context) {
+        if (context.empty()) {
+            return;
+        }
+        if (!failure_context.has_value()) {
+            failure_context = context;
+            return;
+        }
+        failure_context = *failure_context + "; " + context;
+    }
+
+    void InitializeRequestedReplicas(const ReplicateConfig& config) {
+        requested_memory_replicas = config.replica_num;
+        requested_nof_replicas = config.nof_replica_num;
+    }
+
+    ReplicateConfig ToReplicateConfig() const {
+        ReplicateConfig config;
+        config.replica_num = requested_memory_replicas;
+        config.nof_replica_num = requested_nof_replicas;
+        return config;
+    }
+
+    void RecordAllocatedReplicas() {
+        transfer_summary = ReplicaTransferSummary{};
+        for (const auto& replica : replicas) {
+            transfer_summary.RecordAllocatedReplica(replica);
+        }
+    }
+
     bool IsResolved() const { return state != PutOperationState::PENDING; }
 
     bool IsSuccessful() const {
@@ -1447,6 +1872,13 @@ void Client::StartBatchPut(std::vector<PutOperation>& ops,
     keys.reserve(ops.size());
     slice_lengths.reserve(ops.size());
 
+    if (hot_cache_) {
+        std::vector<std::string> hot_keys;
+        hot_keys.reserve(ops.size());
+        for (const auto& op : ops) hot_keys.emplace_back(op.key);
+        hot_cache_->RemoveHotKeys(hot_keys);
+    }
+
     for (const auto& op : ops) {
         keys.emplace_back(op.key);
 
@@ -1474,11 +1906,22 @@ void Client::StartBatchPut(std::vector<PutOperation>& ops,
 
     // Process individual responses with robust error handling
     for (size_t i = 0; i < ops.size(); ++i) {
+        ops[i].InitializeRequestedReplicas(config);
         if (!start_responses[i]) {
-            ops[i].SetError(start_responses[i].error(),
-                            "Master failed to start put operation");
+            ops[i].SetTerminalError(start_responses[i].error(),
+                                    PutOperationState::MASTER_FAILED,
+                                    "Master failed to start put operation");
         } else {
             ops[i].replicas = start_responses[i].value();
+            ops[i].RecordAllocatedReplicas();
+            if (!HasExpectedReplicaAllocation(config,
+                                              ops[i].transfer_summary)) {
+                ops[i].SetTerminalError(ErrorCode::NO_AVAILABLE_HANDLE,
+                                        PutOperationState::MASTER_FAILED,
+                                        "Allocated replicas do not satisfy "
+                                        "requested replica policy");
+                continue;
+            }
             // Operation continues to next stage - result remains INTERNAL_ERROR
             // until fully successful
             VLOG(1) << "Successfully started put for key " << ops[i].key
@@ -1495,6 +1938,13 @@ void Client::StartBatchUpsert(std::vector<PutOperation>& ops,
     keys.reserve(ops.size());
     slice_lengths.reserve(ops.size());
 
+    if (hot_cache_) {
+        std::vector<std::string> hot_keys;
+        hot_keys.reserve(ops.size());
+        for (const auto& op : ops) hot_keys.emplace_back(op.key);
+        hot_cache_->RemoveHotKeys(hot_keys);
+    }
+
     for (const auto& op : ops) {
         keys.emplace_back(op.key);
 
@@ -1537,8 +1987,9 @@ void Client::SubmitTransfers(std::vector<PutOperation>& ops) {
     if (!transfer_submitter_) {
         LOG(ERROR) << "TransferSubmitter not initialized";
         for (auto& op : ops) {
-            op.SetError(ErrorCode::INVALID_PARAMS,
-                        "TransferSubmitter not initialized");
+            op.SetTerminalError(ErrorCode::INVALID_PARAMS,
+                                PutOperationState::TRANSFER_FAILED,
+                                "TransferSubmitter not initialized");
         }
         return;
     }
@@ -1551,14 +2002,12 @@ void Client::SubmitTransfers(std::vector<PutOperation>& ops) {
 
         // Skip operations that don't have replicas (failed in StartBatchPut)
         if (op.replicas.empty()) {
-            op.SetError(ErrorCode::INTERNAL_ERROR,
-                        "No replicas available for transfer");
+            op.SetTerminalError(ErrorCode::INTERNAL_ERROR,
+                                PutOperationState::MASTER_FAILED,
+                                "No replicas available for transfer");
             continue;
         }
 
-        bool all_transfers_submitted = true;
-        std::string failure_context;
-
         // We must deal with disk replica first, then the disk putrevoke/putend
         // can be called surely
         if (storage_backend_) {
@@ -1576,31 +2025,48 @@ void Client::SubmitTransfers(std::vector<PutOperation>& ops) {
         for (size_t replica_idx = 0; replica_idx < op.replicas.size();
              ++replica_idx) {
             const auto& replica = op.replicas[replica_idx];
-            if (replica.is_memory_replica()) {
-                auto submit_result = transfer_submitter_->submit(
-                    replica, op.slices, TransferRequest::WRITE);
+            if (replica.is_memory_replica() || replica.is_nof_replica()) {
+                const auto replica_type = replica.is_memory_replica()
+                                              ? ReplicaType::MEMORY
+                                              : ReplicaType::NOF_SSD;
+                std::optional<TransferFuture> submit_result;
+                if (replica.is_nof_replica()) {
+                    auto contiguous_range = GetContiguousSliceRange(op.slices);
+                    if (!contiguous_range.has_value()) {
+                        std::string failure_context =
+                            "NoF transfer requires contiguous slices for "
+                            "replica " +
+                            std::to_string(replica_idx);
+                        op.transfer_summary.RecordFailure(
+                            replica_type, ErrorCode::INVALID_PARAMS);
+                        op.AppendFailureContext(failure_context);
+                        continue;
+                    }
+                    submit_result = transfer_submitter_->submit(
+                        replica, op.slices, TransferRequest::WRITE,
+                        contiguous_range->ptr, contiguous_range->size);
+                } else {
+                    submit_result = transfer_submitter_->submit(
+                        replica, op.slices, TransferRequest::WRITE);
+                }
 
                 if (!submit_result) {
-                    failure_context = "Failed to submit transfer for replica " +
-                                      std::to_string(replica_idx);
-                    all_transfers_submitted = false;
-                    break;
+                    std::string failure_context =
+                        "Failed to submit transfer for replica " +
+                        std::to_string(replica_idx);
+                    op.transfer_summary.RecordFailure(replica_type,
+                                                      ErrorCode::TRANSFER_FAIL);
+                    op.AppendFailureContext(failure_context);
+                    continue;
                 }
 
                 op.pending_transfers.emplace_back(
-                    std::move(submit_result.value()));
+                    replica_type, std::move(submit_result.value()));
             }
         }
 
-        if (!all_transfers_submitted) {
-            LOG(ERROR) << "Transfer submission failed for key " << op.key
-                       << ": " << failure_context;
-            op.SetError(ErrorCode::TRANSFER_FAIL, failure_context);
-            op.pending_transfers.clear();
-        } else {
-            VLOG(1) << "Successfully submitted " << op.pending_transfers.size()
-                    << " transfers for key " << op.key;
-        }
+        VLOG(1) << "Submitted " << op.pending_transfers.size()
+                << " transfers for key " << op.key;
     }
 }
 
@@ -1611,154 +2077,223 @@ void Client::WaitForTransfers(std::vector<PutOperation>& ops) {
             continue;
         }
 
-        // Skip operations with no pending transfers (failed in SubmitTransfers)
-        if (op.pending_transfers.empty()) {
-            op.SetError(ErrorCode::INTERNAL_ERROR,
-                        "No pending transfers to wait for");
-            continue;
-        }
-
-        bool all_transfers_succeeded = true;
-        ErrorCode first_error = ErrorCode::OK;
-        size_t failed_transfer_idx = 0;
-
         for (size_t i = 0; i < op.pending_transfers.size(); ++i) {
-            ErrorCode transfer_result = op.pending_transfers[i].get();
+            auto& pending_transfer = op.pending_transfers[i];
+            ErrorCode transfer_result = pending_transfer.future.get();
             if (transfer_result != ErrorCode::OK) {
-                if (all_transfers_succeeded) {
-                    // Record the first error for reporting
-                    first_error = transfer_result;
-                    failed_transfer_idx = i;
-                    all_transfers_succeeded = false;
-                }
-                // Continue waiting for all transfers to avoid resource leaks
+                op.transfer_summary.RecordFailure(pending_transfer.replica_type,
+                                                  transfer_result);
+                std::string error_context =
+                    "Transfer " + std::to_string(i) + " failed";
+                op.AppendFailureContext(error_context);
+            } else {
+                op.transfer_summary.RecordSuccess(
+                    pending_transfer.replica_type);
             }
         }
 
-        if (all_transfers_succeeded) {
-            VLOG(1) << "All transfers completed successfully for key "
-                    << op.key;
-            // Transfer phase successful - continue to finalization
-            // Note: Don't mark as SUCCESS yet, need to complete finalization
-        } else {
-            std::string error_context =
-                "Transfer " + std::to_string(failed_transfer_idx) + " failed";
-            LOG(ERROR) << "Transfer failed for key " << op.key << ": "
-                       << toString(first_error) << " (" << error_context << ")";
-            op.SetError(first_error, error_context);
-        }
+        VLOG(1) << "Transfers finished for key " << op.key << ", success(mem="
+                << op.transfer_summary.successful_memory_transfers
+                << ", nof=" << op.transfer_summary.successful_nof_transfers
+                << "), fail(mem=" << op.transfer_summary.failed_memory_transfers
+                << ", nof=" << op.transfer_summary.failed_nof_transfers << ")";
     }
 }
 
 void Client::FinalizeBatchPut(std::vector<PutOperation>& ops) {
-    // For each operation,
-    // If transfers completed successfully, we need to call BatchPutEnd
-    // If the operation failed but has allocated replicas, we need to call
-    // BatchPutRevoke
+    struct BatchFinalizeGroup {
+        std::vector<std::string> keys;
+        std::vector<size_t> indices;
+    };
 
-    std::vector<std::string> successful_keys;
-    std::vector<size_t> successful_indices;
-    std::vector<std::string> failed_keys;
-    std::vector<size_t> failed_indices;
+    BatchFinalizeGroup end_all_group;
+    BatchFinalizeGroup end_memory_group;
+    BatchFinalizeGroup end_nof_group;
+    BatchFinalizeGroup revoke_all_group;
+    BatchFinalizeGroup revoke_memory_group;
+    BatchFinalizeGroup revoke_nof_group;
 
-    // Reserve space to avoid reallocations
-    successful_keys.reserve(ops.size());
-    successful_indices.reserve(ops.size());
-    failed_keys.reserve(ops.size());
-    failed_indices.reserve(ops.size());
+    std::vector<size_t> pending_finalize_actions(ops.size(), 0);
+    std::vector<bool> should_succeed(ops.size(), false);
+    std::vector<ErrorCode> terminal_errors(ops.size(), ErrorCode::OK);
+    std::vector<std::optional<ErrorCode>> finalize_rpc_errors(ops.size());
+
+    auto add_group_entry = [](BatchFinalizeGroup& group, const std::string& key,
+                              size_t index) {
+        group.keys.emplace_back(key);
+        group.indices.emplace_back(index);
+    };
+
+    auto add_finalize_action =
+        [&](const std::optional<ReplicaType>& replica_type, bool is_end,
+            const std::string& key, size_t index) {
+            if (!replica_type.has_value()) {
+                return;
+            }
+            switch (*replica_type) {
+                case ReplicaType::ALL:
+                    add_group_entry(is_end ? end_all_group : revoke_all_group,
+                                    key, index);
+                    ++pending_finalize_actions[index];
+                    break;
+                case ReplicaType::MEMORY:
+                    add_group_entry(
+                        is_end ? end_memory_group : revoke_memory_group, key,
+                        index);
+                    ++pending_finalize_actions[index];
+                    break;
+                case ReplicaType::NOF_SSD:
+                    add_group_entry(is_end ? end_nof_group : revoke_nof_group,
+                                    key, index);
+                    ++pending_finalize_actions[index];
+                    break;
+                default:
+                    LOG(ERROR) << "Unexpected replica type in batch finalize: "
+                               << *replica_type;
+                    finalize_rpc_errors[index] = ErrorCode::INVALID_PARAMS;
+                    break;
+            }
+        };
+
+    auto complete_finalize_action = [&](size_t index) {
+        if (pending_finalize_actions[index] > 0) {
+            --pending_finalize_actions[index];
+        }
+    };
 
     for (size_t i = 0; i < ops.size(); ++i) {
         auto& op = ops[i];
-
-        // Check if operation completed transfers successfully and needs
-        // finalization
-        if (!op.IsResolved() && !op.replicas.empty() &&
-            !op.pending_transfers.empty()) {
-            // Transfers completed, needs BatchPutEnd
-            successful_keys.emplace_back(op.key);
-            successful_indices.emplace_back(i);
-        } else if (op.state != PutOperationState::PENDING &&
-                   !op.replicas.empty()) {
-            // Operation failed but has allocated replicas, needs BatchPutRevoke
-            failed_keys.emplace_back(op.key);
-            failed_indices.emplace_back(i);
+        if (op.IsResolved()) {
+            if (!op.IsSuccessful() && !op.replicas.empty()) {
+                terminal_errors[i] = op.result.has_value()
+                                         ? ErrorCode::INTERNAL_ERROR
+                                         : op.result.error();
+                add_finalize_action(ReplicaType::ALL, false, op.key, i);
+            }
+            continue;
         }
-        // Operations without replicas (early failures) don't need finalization
+        if (op.replicas.empty()) {
+            op.SetTerminalError(ErrorCode::INTERNAL_ERROR,
+                                PutOperationState::MASTER_FAILED,
+                                "Operation has no replicas to finalize");
+            continue;
+        }
+
+        const auto finalize_decision = DetermineFinalizeDecision(
+            op.ToReplicateConfig(), op.transfer_summary);
+        should_succeed[i] = finalize_decision.success;
+        terminal_errors[i] = finalize_decision.error;
+        add_finalize_action(finalize_decision.end_type, true, op.key, i);
+        add_finalize_action(finalize_decision.revoke_type, false, op.key, i);
     }
 
-    // Process successful operations
-    if (!successful_keys.empty()) {
-        auto end_responses = master_client_.BatchPutEnd(successful_keys);
-        if (end_responses.size() != successful_keys.size()) {
-            LOG(ERROR) << "BatchPutEnd response size mismatch: expected "
-                       << successful_keys.size() << ", got "
-                       << end_responses.size();
-            for (size_t idx : successful_indices) {
-                ops[idx].SetError(ErrorCode::RPC_FAIL,
-                                  "BatchPutEnd response size mismatch");
-            }
-        } else {
-            // Process individual responses
-            for (size_t i = 0; i < end_responses.size(); ++i) {
-                const size_t op_idx = successful_indices[i];
-                if (!end_responses[i]) {
-                    LOG(ERROR) << "Failed to finalize put for key "
-                               << successful_keys[i] << ": "
-                               << toString(end_responses[i].error());
-                    ops[op_idx].SetError(end_responses[i].error(),
-                                         "BatchPutEnd failed");
-                } else {
-                    // Operation fully successful
-                    ops[op_idx].SetSuccess();
-                    VLOG(1) << "Successfully completed put for key "
-                            << successful_keys[i];
-                }
-            }
+    auto process_end_group = [&](BatchFinalizeGroup& group,
+                                 ReplicaType replica_type) {
+        if (group.keys.empty()) {
+            return;
         }
-    }
+        auto responses = master_client_.BatchPutEnd(group.keys, replica_type);
+        if (responses.size() != group.keys.size()) {
+            for (size_t idx : group.indices) {
+                finalize_rpc_errors[idx] = ErrorCode::RPC_FAIL;
+                complete_finalize_action(idx);
+            }
+            return;
+        }
+        for (size_t i = 0; i < responses.size(); ++i) {
+            const size_t op_idx = group.indices[i];
+            if (!responses[i]) {
+                finalize_rpc_errors[op_idx] = responses[i].error();
+                LOG(ERROR) << "Failed to BatchPutEnd key " << group.keys[i]
+                           << ": " << toString(responses[i].error());
+                complete_finalize_action(op_idx);
+                continue;
+            }
+            complete_finalize_action(op_idx);
+        }
+    };
 
-    // Process failed operations that need cleanup
-    if (!failed_keys.empty()) {
-        auto revoke_responses = master_client_.BatchPutRevoke(failed_keys);
-        if (revoke_responses.size() != failed_keys.size()) {
-            LOG(ERROR) << "BatchPutRevoke response size mismatch: expected "
-                       << failed_keys.size() << ", got "
-                       << revoke_responses.size();
-            // Mark all failed operations with revoke RPC failure
-            for (size_t idx : failed_indices) {
-                ops[idx].SetError(ErrorCode::RPC_FAIL,
-                                  "BatchPutRevoke response size mismatch");
-            }
-        } else {
-            // Process individual revoke responses
-            for (size_t i = 0; i < revoke_responses.size(); ++i) {
-                const size_t op_idx = failed_indices[i];
-                if (!revoke_responses[i]) {
-                    LOG(ERROR)
-                        << "Failed to revoke put for key " << failed_keys[i]
-                        << ": " << toString(revoke_responses[i].error());
-                    // Preserve original error but note revoke failure in
-                    // context
-                    std::string original_context =
-                        ops[op_idx].failure_context.value_or("unknown error");
-                    ops[op_idx].failure_context =
-                        original_context + "; revoke also failed";
-                } else {
-                    LOG(INFO) << "Successfully revoked failed put for key "
-                              << failed_keys[i];
-                }
-            }
+    auto process_revoke_group = [&](BatchFinalizeGroup& group,
+                                    ReplicaType replica_type) {
+        if (group.keys.empty()) {
+            return;
         }
-    }
+        auto responses =
+            master_client_.BatchPutRevoke(group.keys, replica_type);
+        if (responses.size() != group.keys.size()) {
+            for (size_t idx : group.indices) {
+                finalize_rpc_errors[idx] = ErrorCode::RPC_FAIL;
+                complete_finalize_action(idx);
+            }
+            return;
+        }
+        for (size_t i = 0; i < responses.size(); ++i) {
+            const size_t op_idx = group.indices[i];
+            if (!responses[i]) {
+                finalize_rpc_errors[op_idx] = responses[i].error();
+                LOG(ERROR) << "Failed to BatchPutRevoke key " << group.keys[i]
+                           << ": " << toString(responses[i].error());
+                complete_finalize_action(op_idx);
+                continue;
+            }
+            complete_finalize_action(op_idx);
+        }
+    };
 
-    // Ensure all operations have definitive results
-    for (auto& op : ops) {
-        if (!op.IsResolved()) {
-            op.SetError(ErrorCode::INTERNAL_ERROR,
-                        "Operation not resolved after finalization");
-            LOG(ERROR) << "Operation for key " << op.key
-                       << " was not properly resolved";
+    process_end_group(end_all_group, ReplicaType::ALL);
+    process_end_group(end_memory_group, ReplicaType::MEMORY);
+    process_end_group(end_nof_group, ReplicaType::NOF_SSD);
+    process_revoke_group(revoke_all_group, ReplicaType::ALL);
+    process_revoke_group(revoke_memory_group, ReplicaType::MEMORY);
+    process_revoke_group(revoke_nof_group, ReplicaType::NOF_SSD);
+
+    auto append_finalize_error_context = [&](PutOperation& op, size_t index) {
+        if (finalize_rpc_errors[index].has_value()) {
+            op.AppendFailureContext("Batch finalization RPC failed: " +
+                                    toString(*finalize_rpc_errors[index]));
         }
+        if (pending_finalize_actions[index] != 0) {
+            op.AppendFailureContext(
+                "Operation has unfinished finalize actions");
+        }
+    };
+
+    for (size_t i = 0; i < ops.size(); ++i) {
+        auto& op = ops[i];
+        if (op.IsResolved()) {
+            if (!op.IsSuccessful()) {
+                append_finalize_error_context(op, i);
+            }
+            continue;
+        }
+        if (finalize_rpc_errors[i].has_value() ||
+            pending_finalize_actions[i] != 0) {
+            if (!should_succeed[i] && terminal_errors[i] != ErrorCode::OK) {
+                append_finalize_error_context(op, i);
+                op.SetTerminalError(
+                    terminal_errors[i], PutOperationState::TRANSFER_FAILED,
+                    op.failure_context.value_or(
+                        "Replica transfer failed before finalize"));
+            } else if (finalize_rpc_errors[i].has_value()) {
+                op.SetTerminalError(*finalize_rpc_errors[i],
+                                    PutOperationState::FINALIZE_FAILED,
+                                    "Batch finalization RPC failed");
+            } else {
+                op.SetTerminalError(
+                    ErrorCode::INTERNAL_ERROR,
+                    PutOperationState::FINALIZE_FAILED,
+                    "Operation has unfinished finalize actions");
+            }
+            continue;
+        }
+        if (should_succeed[i]) {
+            op.SetSuccess();
+            continue;
+        }
+        op.SetTerminalError(terminal_errors[i],
+                            PutOperationState::TRANSFER_FAILED,
+                            op.failure_context.value_or(
+                                "Replica transfer failed before finalize"));
     }
 }
 
@@ -1788,7 +2323,9 @@ void Client::FinalizeBatchUpsert(std::vector<PutOperation>& ops) {
     }
 
     // Process successful operations
+    std::vector<std::string> finalized_keys;
     if (!successful_keys.empty()) {
+        finalized_keys.reserve(successful_keys.size());
         auto end_responses = master_client_.BatchUpsertEnd(successful_keys);
         if (end_responses.size() != successful_keys.size()) {
             LOG(ERROR) << "BatchUpsertEnd response size mismatch: expected "
@@ -1809,6 +2346,7 @@ void Client::FinalizeBatchUpsert(std::vector<PutOperation>& ops) {
                                          "BatchUpsertEnd failed");
                 } else {
                     ops[op_idx].SetSuccess();
+                    finalized_keys.emplace_back(successful_keys[i]);
                     VLOG(1) << "Successfully completed upsert for key "
                             << successful_keys[i];
                 }
@@ -1816,6 +2354,15 @@ void Client::FinalizeBatchUpsert(std::vector<PutOperation>& ops) {
         }
     }
 
+    // Success-side invalidation for finalized upserts only: a concurrent read
+    // between StartBatchUpsert()'s pre-invalidation and BatchUpsertEnd could
+    // have read the old value and submitted an async hot-cache fill with a
+    // still-valid token. Invalidate again now that the new values are
+    // committed so those stale fills cannot publish.
+    if (hot_cache_ && !finalized_keys.empty()) {
+        hot_cache_->RemoveHotKeys(finalized_keys);
+    }
+
     // Process failed operations that need cleanup
     if (!failed_keys.empty()) {
         auto revoke_responses = master_client_.BatchUpsertRevoke(failed_keys);
@@ -1939,6 +2486,7 @@ std::vector<tl::expected<void, ErrorCode>> Client::BatchPutWhenPreferSameNode(
         merged_ops.emplace_back(op.key, op.slices);
         auto& merged_op = merged_ops.back();
         merged_op.replicas = op.replicas;
+        merged_op.transfer_summary.allocated_memory_replicas = 1;
         auto submit_result = transfer_submitter_->submit_batch(
             op.replicas, op.batched_slices, TransferRequest::WRITE);
         if (!submit_result) {
@@ -1946,12 +2494,14 @@ std::vector<tl::expected<void, ErrorCode>> Client::BatchPutWhenPreferSameNode(
             all_transfers_submitted = false;
         } else {
             merged_op.pending_transfers.emplace_back(
-                std::move(submit_result.value()));
+                ReplicaType::MEMORY, std::move(submit_result.value()));
         }
         if (!all_transfers_submitted) {
             LOG(ERROR) << "Transfer submission failed for key " << op.key
                        << ": " << failure_context;
-            merged_op.SetError(ErrorCode::TRANSFER_FAIL, failure_context);
+            merged_op.transfer_summary.RecordFailure(ReplicaType::MEMORY,
+                                                     ErrorCode::TRANSFER_FAIL);
+            merged_op.failure_context = failure_context;
             merged_op.pending_transfers.clear();
         } else {
             VLOG(1) << "Successfully submitted "
@@ -1964,7 +2514,8 @@ std::vector<tl::expected<void, ErrorCode>> Client::BatchPutWhenPreferSameNode(
         auto& memory_descriptor = op.replicas[0].get_memory_descriptor();
         auto& buffer_descriptor = memory_descriptor.buffer_descriptor;
         auto seg = buffer_descriptor.transport_endpoint_;
-        seg_to_ops.at(seg).state = op.state;
+        seg_to_ops.at(seg).transfer_summary = op.transfer_summary;
+        seg_to_ops.at(seg).failure_context = op.failure_context;
     }
     for (auto& op : ops) {
         if (op.IsResolved()) {
@@ -1973,10 +2524,16 @@ std::vector<tl::expected<void, ErrorCode>> Client::BatchPutWhenPreferSameNode(
         auto& memory_descriptor = op.replicas[0].get_memory_descriptor();
         auto& buffer_descriptor = memory_descriptor.buffer_descriptor;
         auto seg = buffer_descriptor.transport_endpoint_;
-        op.state = seg_to_ops.at(seg).state;
-        auto state = std::make_shared<EmptyOperationState>();
-        auto future = TransferFuture(state);
-        op.pending_transfers.emplace_back(std::move(future));
+        op.transfer_summary.successful_memory_transfers =
+            seg_to_ops.at(seg).transfer_summary.successful_memory_transfers > 0
+                ? 1
+                : 0;
+        op.transfer_summary.failed_memory_transfers =
+            seg_to_ops.at(seg).transfer_summary.failed_memory_transfers > 0 ? 1
+                                                                            : 0;
+        op.transfer_summary.first_error =
+            seg_to_ops.at(seg).transfer_summary.first_error;
+        op.failure_context = seg_to_ops.at(seg).failure_context;
     }
     auto us = std::chrono::duration_cast<std::chrono::microseconds>(
                   std::chrono::steady_clock::now() - t0)
@@ -1998,6 +2555,12 @@ std::vector<tl::expected<void, ErrorCode>> Client::BatchPut(
     }
     std::vector<PutOperation> ops = CreatePutOperations(keys, batched_slices);
     if (client_cfg.prefer_alloc_in_same_node) {
+        if (client_cfg.nof_replica_num > 0) {
+            LOG(ERROR) << "prefer_alloc_in_same_node is not supported with "
+                          "NoF replicas";
+            return std::vector<tl::expected<void, ErrorCode>>(
+                keys.size(), tl::unexpected(ErrorCode::INVALID_PARAMS));
+        }
         if (client_cfg.replica_num != 1) {
             LOG(ERROR) << "prefer_alloc_in_same_node is not supported with "
                           "replica_num != 1";
@@ -2024,6 +2587,10 @@ std::vector<tl::expected<void, ErrorCode>> Client::BatchPut(
 }
 
 tl::expected<void, ErrorCode> Client::Remove(const ObjectKey& key, bool force) {
+    if (hot_cache_) {
+        hot_cache_->BumpKeyGeneration(key);
+    }
+
     auto result = master_client_.Remove(key, force);
     // if (storage_backend_) {
     //     storage_backend_->RemoveFile(key);
@@ -2031,11 +2598,20 @@ tl::expected<void, ErrorCode> Client::Remove(const ObjectKey& key, bool force) {
     if (!result) {
         return tl::unexpected(result.error());
     }
+
+    if (hot_cache_) {
+        hot_cache_->RemoveHotKey(key);
+    }
+
     return {};
 }
 
 tl::expected<long, ErrorCode> Client::RemoveByRegex(const ObjectKey& str,
                                                     bool force) {
+    if (hot_cache_) {
+        hot_cache_->BumpCacheEpoch();
+    }
+
     auto result = master_client_.RemoveByRegex(str, force);
     // if (storage_backend_) {
     //     storage_backend_->RemoveByRegex(str);
@@ -2043,19 +2619,50 @@ tl::expected<long, ErrorCode> Client::RemoveByRegex(const ObjectKey& str,
     if (!result) {
         return tl::unexpected(result.error());
     }
+    if (result.value() > 0 && hot_cache_) {
+        hot_cache_->BumpCacheEpoch();
+        hot_cache_->RemoveHotKeysByRegex(str);
+    }
     return result.value();
 }
 
 tl::expected<long, ErrorCode> Client::RemoveAll(bool force) {
-    // if (storage_backend_) {
-    //     storage_backend_->RemoveAll();
-    // }
-    return master_client_.RemoveAll(force);
+    if (hot_cache_) {
+        hot_cache_->BumpCacheEpoch();
+    }
+
+    auto result = master_client_.RemoveAll(force);
+    if (result && storage_backend_) {
+        storage_backend_->RemoveAll();
+    }
+    if (result && result.value() > 0 && hot_cache_) {
+        hot_cache_->RemoveAllHotKeys();
+    }
+    return result;
 }
 
 std::vector<tl::expected<void, ErrorCode>> Client::BatchRemove(
     const std::vector<ObjectKey>& keys, bool force) {
-    return master_client_.BatchRemove(keys, force);
+    if (hot_cache_) {
+        hot_cache_->BumpKeyGenerations(keys);
+    }
+
+    auto results = master_client_.BatchRemove(keys, force);
+
+    if (hot_cache_) {
+        std::vector<std::string> removed_keys;
+        removed_keys.reserve(std::min(keys.size(), results.size()));
+        for (size_t i = 0; i < keys.size(); ++i) {
+            if (i < results.size() && results[i].has_value()) {
+                removed_keys.emplace_back(keys[i]);
+            }
+        }
+        if (!removed_keys.empty()) {
+            hot_cache_->RemoveHotKeys(removed_keys);
+        }
+    }
+
+    return results;
 }
 
 tl::expected<void, ErrorCode> Client::EvictDiskReplica(
@@ -2063,11 +2670,23 @@ tl::expected<void, ErrorCode> Client::EvictDiskReplica(
     return master_client_.EvictDiskReplica(key, replica_type);
 }
 
+tl::expected<void, ErrorCode> Client::EvictDiskReplica(
+    const std::string& key, const std::string& tenant_id,
+    ReplicaType replica_type) {
+    return master_client_.EvictDiskReplica(key, tenant_id, replica_type);
+}
+
 std::vector<tl::expected<void, ErrorCode>> Client::BatchEvictDiskReplica(
     const std::vector<std::string>& keys, ReplicaType replica_type) {
     return master_client_.BatchEvictDiskReplica(keys, replica_type);
 }
 
+std::vector<tl::expected<void, ErrorCode>> Client::BatchEvictDiskReplica(
+    const std::vector<std::string>& keys, const std::string& tenant_id,
+    ReplicaType replica_type) {
+    return master_client_.BatchEvictDiskReplica(keys, tenant_id, replica_type);
+}
+
 std::vector<int> Client::GetNicNumaNodes() const {
     std::set<int> nodes;
     if (!transfer_engine_) return {};
@@ -2082,6 +2701,64 @@ std::vector<int> Client::GetNicNumaNodes() const {
 }
 
 tl::expected<void, ErrorCode> Client::MountSegment(
+    const void* buffer, size_t size, const std::string& protocol,
+    const std::string& location) {
+    auto result = MountSegmentAndGetId(buffer, size, protocol, location);
+    if (!result) {
+        return tl::unexpected(result.error());
+    }
+    return {};
+}
+
+tl::expected<void, ErrorCode> Client::UnmountSegmentImpl(
+    std::unordered_map<UUID, Segment, boost::hash<UUID>>::iterator it) {
+    auto unmount_result = master_client_.UnmountSegment(it->second.id);
+    if (!unmount_result) {
+        ErrorCode err = unmount_result.error();
+        LOG(ERROR) << "Failed to unmount segment from master: "
+                   << toString(err);
+        return tl::unexpected(err);
+    }
+
+    int rc = transfer_engine_->unregisterLocalMemory(
+        reinterpret_cast<void*>(it->second.base));
+    if (rc != 0) {
+        LOG(ERROR) << "Failed to unregister transfer buffer with transfer "
+                      "engine ret is "
+                   << rc;
+        if (rc != ERR_ADDRESS_NOT_REGISTERED) {
+            return tl::unexpected(ErrorCode::INTERNAL_ERROR);
+        }
+        // Otherwise, the segment is already unregistered from transfer
+        // engine, we can continue
+    }
+
+    mounted_segments_.erase(it);
+    return {};
+}
+
+tl::expected<void, ErrorCode> Client::UnmountSegment(const void* buffer,
+                                                     size_t size) {
+    std::lock_guard<std::mutex> lock(mounted_segments_mutex_);
+    auto segment = mounted_segments_.end();
+
+    for (auto it = mounted_segments_.begin(); it != mounted_segments_.end();
+         ++it) {
+        if (it->second.base == reinterpret_cast<uintptr_t>(buffer) &&
+            it->second.size == size) {
+            segment = it;
+            break;
+        }
+    }
+    if (segment == mounted_segments_.end()) {
+        LOG(ERROR) << "segment_not_found base=" << buffer << " size=" << size;
+        return tl::unexpected(ErrorCode::INVALID_PARAMS);
+    }
+
+    return UnmountSegmentImpl(segment);
+}
+
+tl::expected<UUID, ErrorCode> Client::MountSegmentAndGetId(
     const void* buffer, size_t size, const std::string& protocol,
     const std::string& location) {
     auto check_result = CheckRegisterMemoryParams(buffer, size);
@@ -2089,6 +2766,7 @@ tl::expected<void, ErrorCode> Client::MountSegment(
         return tl::unexpected(check_result.error());
     }
 
+    UUID segment_id;
     {
         std::lock_guard<std::mutex> lock(mounted_segments_mutex_);
 
@@ -2115,17 +2793,12 @@ tl::expected<void, ErrorCode> Client::MountSegment(
             return tl::unexpected(ErrorCode::INVALID_PARAMS);
         }
 
-        // Build segment with logical name; attach TE endpoint for transport
         Segment segment;
         segment.id = generate_uuid();
         segment.name = local_hostname_;
         segment.base = reinterpret_cast<uintptr_t>(buffer);
         segment.size = size;
         segment.protocol = protocol;
-        // For P2P handshake mode, publish the actual transport endpoint that
-        // was negotiated by the transfer engine. Otherwise, keep the logical
-        // hostname so metadata backends (HTTP/etcd/redis) can resolve the
-        // segment by name.
         if (metadata_connstring_ == P2PHANDSHAKE) {
             segment.te_endpoint = transfer_engine_->getLocalIpAndPort();
         } else {
@@ -2140,54 +2813,133 @@ tl::expected<void, ErrorCode> Client::MountSegment(
             return tl::unexpected(err);
         }
 
-        mounted_segments_[segment.id] = segment;
+        segment_id = segment.id;
+        mounted_segments_[segment_id] = segment;
     }
 
     EnsureStorageControlPlaneStarted();
-    return {};
+    return segment_id;
 }
 
-tl::expected<void, ErrorCode> Client::UnmountSegment(const void* buffer,
-                                                     size_t size) {
+tl::expected<void, ErrorCode> Client::UnmountSegmentById(
+    const UUID& segment_id, uint64_t grace_period_ms,
+    std::function<void(const UUID&)> cleanup_callback) {
     std::lock_guard<std::mutex> lock(mounted_segments_mutex_);
-    auto segment = mounted_segments_.end();
-
-    for (auto it = mounted_segments_.begin(); it != mounted_segments_.end();
-         ++it) {
-        if (it->second.base == reinterpret_cast<uintptr_t>(buffer) &&
-            it->second.size == size) {
-            segment = it;
-            break;
-        }
-    }
+    auto segment = mounted_segments_.find(segment_id);
     if (segment == mounted_segments_.end()) {
-        LOG(ERROR) << "segment_not_found base=" << buffer << " size=" << size;
+        LOG(ERROR) << "segment_not_found id=" << UuidToString(segment_id);
         return tl::unexpected(ErrorCode::INVALID_PARAMS);
     }
 
-    auto unmount_result = master_client_.UnmountSegment(segment->second.id);
-    if (!unmount_result) {
-        ErrorCode err = unmount_result.error();
-        LOG(ERROR) << "Failed to unmount segment from master: "
+    if (grace_period_ms == 0) {
+        return UnmountSegmentImpl(segment);
+    }
+
+    auto result =
+        master_client_.GracefulUnmountSegment(segment_id, grace_period_ms);
+    if (!result) {
+        ErrorCode err = result.error();
+        LOG(ERROR) << "Failed to graceful unmount segment from master: "
                    << toString(err);
         return tl::unexpected(err);
     }
 
-    int rc = transfer_engine_->unregisterLocalMemory(
-        reinterpret_cast<void*>(segment->second.base));
-    if (rc != 0) {
-        LOG(ERROR) << "Failed to unregister transfer buffer with transfer "
-                      "engine ret is "
-                   << rc;
-        if (rc != ERR_ADDRESS_NOT_REGISTERED) {
-            return tl::unexpected(ErrorCode::INTERNAL_ERROR);
+    gracefully_unmounting_segments_.emplace(segment->first, segment->second);
+    if (cleanup_callback) {
+        graceful_unmount_cleanup_callbacks_[segment->first] =
+            std::move(cleanup_callback);
+    }
+    mounted_segments_.erase(segment);
+    StartGracefulUnmountTimer(segment_id, grace_period_ms);
+    return {};
+}
+
+bool Client::WaitForGracefulUnmountDelay(std::chrono::milliseconds delay) {
+    std::unique_lock<std::mutex> lock(graceful_unmount_timer_mutex_);
+    return graceful_unmount_timer_cv_.wait_for(
+        lock, delay, [this]() { return graceful_unmount_timer_stopping_; });
+}
+
+void Client::StartGracefulUnmountTimer(const UUID& segment_id,
+                                       uint64_t grace_period_ms) {
+    auto delay =
+        std::chrono::milliseconds(grace_period_ms) + std::chrono::seconds(10);
+    task_thread_pool_.enqueue([this, segment_id, delay]() {
+        if (this->WaitForGracefulUnmountDelay(delay)) {
+            return;
+        }
+        this->OnGracefulUnmountTimer(segment_id, /*retry_left=*/3);
+    });
+}
+
+void Client::OnGracefulUnmountTimer(const UUID& segment_id, int retry_left) {
+    // Query master to confirm the segment has been removed
+    {
+        std::lock_guard<std::mutex> lock(mounted_segments_mutex_);
+        auto it = gracefully_unmounting_segments_.find(segment_id);
+        if (it == gracefully_unmounting_segments_.end()) {
+            // Already cleaned up (e.g. by destructor)
+            return;
         }
-        // Otherwise, the segment is already unregistered from transfer
-        // engine, we can continue
     }
 
-    mounted_segments_.erase(segment);
-    return {};
+    auto status = master_client_.QuerySegmentStatusById(segment_id);
+    bool removed = false;
+    if (!status) {
+        if (status.error() == ErrorCode::SEGMENT_NOT_FOUND) {
+            removed = true;
+        } else {
+            LOG(WARNING) << "Failed to query graceful unmount segment status: "
+                         << toString(status.error());
+        }
+    } else if (status.value() == SegmentStatus::UNDEFINED) {
+        removed = true;
+    }
+
+    if (removed) {
+        std::function<void(const UUID&)> cleanup_callback;
+        {
+            std::lock_guard<std::mutex> lock(mounted_segments_mutex_);
+            auto it = gracefully_unmounting_segments_.find(segment_id);
+            if (it != gracefully_unmounting_segments_.end()) {
+                int rc = transfer_engine_->unregisterLocalMemory(
+                    reinterpret_cast<void*>(it->second.base));
+                if (rc != 0 && rc != ERR_ADDRESS_NOT_REGISTERED) {
+                    LOG(ERROR)
+                        << "Failed to unregister TE MR for graceful unmount: "
+                        << rc;
+                }
+                gracefully_unmounting_segments_.erase(it);
+            }
+            auto callback_it =
+                graceful_unmount_cleanup_callbacks_.find(segment_id);
+            if (callback_it != graceful_unmount_cleanup_callbacks_.end()) {
+                cleanup_callback = std::move(callback_it->second);
+                graceful_unmount_cleanup_callbacks_.erase(callback_it);
+            }
+        }
+        if (cleanup_callback) {
+            cleanup_callback(segment_id);
+        }
+        return;
+    }
+
+    if (retry_left > 0) {
+        try {
+            task_thread_pool_.enqueue([this, segment_id, retry_left]() {
+                if (this->WaitForGracefulUnmountDelay(
+                        std::chrono::seconds(10))) {
+                    return;
+                }
+                this->OnGracefulUnmountTimer(segment_id, retry_left - 1);
+            });
+        } catch (const std::runtime_error& e) {
+            VLOG(1) << "Skip graceful unmount retry enqueue: " << e.what();
+        }
+    } else {
+        LOG(WARNING) << "Graceful unmount cleanup timeout for segment "
+                     << UuidToString(segment_id);
+    }
 }
 
 tl::expected<void, ErrorCode> Client::RegisterLocalMemory(
@@ -2258,8 +3010,7 @@ tl::expected<void, ErrorCode> Client::MountLocalDiskSegment(
 }
 
 tl::expected<void, ErrorCode> Client::OffloadObjectHeartbeat(
-    bool enable_offloading,
-    std::unordered_map<std::string, int64_t>& offloading_objects) {
+    bool enable_offloading, std::vector<OffloadTaskItem>& offloading_objects) {
     auto response =
         master_client_.OffloadObjectHeartbeat(client_id_, enable_offloading);
     if (!response) {
@@ -2271,11 +3022,23 @@ tl::expected<void, ErrorCode> Client::OffloadObjectHeartbeat(
     return {};
 }
 
+tl::expected<void, ErrorCode> Client::ReportSsdCapacity(
+    int64_t ssd_total_capacity_bytes) {
+    auto response =
+        master_client_.ReportSsdCapacity(client_id_, ssd_total_capacity_bytes);
+    if (!response) {
+        LOG(ERROR) << "ReportSsdCapacity failed, error code is "
+                   << response.error();
+        return tl::make_unexpected(response.error());
+    }
+    return {};
+}
+
 tl::expected<void, ErrorCode> Client::BatchGetOffloadObject(
     const std::string& transfer_engine_addr,
     const std::vector<std::string>& keys,
     const std::vector<uintptr_t>& pointers,
-    const std::unordered_map<std::string, Slice>& batch_slices) {
+    const std::unordered_map<std::string, std::vector<Slice>>& batch_slices) {
     auto future = transfer_submitter_->submit_batch_get_offload_object(
         transfer_engine_addr, keys, pointers, batch_slices);
     if (!future) {
@@ -2299,17 +3062,86 @@ tl::expected<void, ErrorCode> Client::NotifyOffloadSuccess(
     return response;
 }
 
+tl::expected<void, ErrorCode> Client::NotifyOffloadSuccess(
+    const std::vector<OffloadTaskItem>& tasks,
+    const std::vector<StorageObjectMetadata>& metadatas) {
+    return master_client_.NotifyOffloadSuccess(client_id_, tasks, metadatas);
+}
+
+tl::expected<void, ErrorCode> Client::PromotionObjectHeartbeat(
+    std::vector<PromotionTaskItem>& promotion_objects) {
+    auto response = master_client_.PromotionObjectHeartbeat(client_id_);
+    if (!response) {
+        return tl::make_unexpected(response.error());
+    }
+    promotion_objects = std::move(response.value());
+    return {};
+}
+
+tl::expected<PromotionAllocStartResponse, ErrorCode>
+Client::PromotionAllocStart(
+    const std::string& key, uint64_t size,
+    const std::vector<std::string>& preferred_segments) {
+    return master_client_.PromotionAllocStart(client_id_, key, size,
+                                              preferred_segments);
+}
+
+tl::expected<PromotionAllocStartResponse, ErrorCode>
+Client::PromotionAllocStart(
+    const std::string& key, const std::string& tenant_id, uint64_t size,
+    const std::vector<std::string>& preferred_segments) {
+    return master_client_.PromotionAllocStart(client_id_, key, tenant_id, size,
+                                              preferred_segments);
+}
+
+tl::expected<void, ErrorCode> Client::NotifyPromotionSuccess(
+    const std::string& key) {
+    return master_client_.NotifyPromotionSuccess(client_id_, key);
+}
+
+tl::expected<void, ErrorCode> Client::NotifyPromotionSuccess(
+    const std::string& key, const std::string& tenant_id) {
+    return master_client_.NotifyPromotionSuccess(client_id_, key, tenant_id);
+}
+
+tl::expected<void, ErrorCode> Client::NotifyPromotionFailure(
+    const std::string& key) {
+    return master_client_.NotifyPromotionFailure(client_id_, key);
+}
+
+tl::expected<void, ErrorCode> Client::NotifyPromotionFailure(
+    const std::string& key, const std::string& tenant_id) {
+    return master_client_.NotifyPromotionFailure(client_id_, key, tenant_id);
+}
+
+ErrorCode Client::PromotionWrite(const Replica::Descriptor& memory_descriptor,
+                                 std::vector<Slice>& slices) {
+    return TransferWrite(memory_descriptor, slices);
+}
+
 tl::expected<UUID, ErrorCode> Client::CreateCopyTask(
     const std::string& key, const std::vector<std::string>& targets) {
     return master_client_.CreateCopyTask(key, targets);
 }
 
+tl::expected<UUID, ErrorCode> Client::CreateCopyTask(
+    const std::string& key, const std::string& tenant_id,
+    const std::vector<std::string>& targets) {
+    return master_client_.CreateCopyTask(key, tenant_id, targets);
+}
+
 tl::expected<UUID, ErrorCode> Client::CreateMoveTask(
     const std::string& key, const std::string& source,
     const std::string& target) {
     return master_client_.CreateMoveTask(key, source, target);
 }
 
+tl::expected<UUID, ErrorCode> Client::CreateMoveTask(
+    const std::string& key, const std::string& tenant_id,
+    const std::string& source, const std::string& target) {
+    return master_client_.CreateMoveTask(key, tenant_id, source, target);
+}
+
 tl::expected<void, ErrorCode> Client::ExecuteReplicaTransfer(
     const std::string& key, const std::string& action_name,
     std::function<tl::expected<void, ErrorCode>()> end_fn,
@@ -2370,11 +3202,18 @@ tl::expected<void, ErrorCode> Client::ExecuteReplicaTransfer(
 tl::expected<void, ErrorCode> Client::Copy(
     const std::string& key, const std::string& source,
     const std::vector<std::string>& targets) {
+    return Copy(key, master_client_.tenant_id(), source, targets);
+}
+
+tl::expected<void, ErrorCode> Client::Copy(
+    const std::string& key, const std::string& tenant_id,
+    const std::string& source, const std::vector<std::string>& targets) {
     LOG(INFO) << "action=replica_copy_start" << ", key=" << key
               << ", targets_count=" << targets.size();
 
     // Call CopyStart first - it validates existence and allocates replicas
-    auto start_result = master_client_.CopyStart(key, source, targets);
+    auto start_result =
+        master_client_.CopyStart(key, tenant_id, source, targets);
     if (!start_result.has_value()) {
         ErrorCode error = start_result.error();
         LOG(ERROR) << "action=replica_copy_failed" << ", key=" << key
@@ -2388,7 +3227,7 @@ tl::expected<void, ErrorCode> Client::Copy(
         LOG(INFO) << "action=replica_copy_skipped" << ", key=" << key
                   << ", info=target_replicas_already_exist";
         // Target replicas already exist, consider it success
-        auto copy_end_result = master_client_.CopyEnd(key);
+        auto copy_end_result = master_client_.CopyEnd(key, tenant_id);
         if (!copy_end_result.has_value()) {
             ErrorCode error = copy_end_result.error();
             LOG(ERROR) << "action=replica_copy_failed" << ", key=" << key
@@ -2399,9 +3238,9 @@ tl::expected<void, ErrorCode> Client::Copy(
     }
 
     auto result = ExecuteReplicaTransfer(
-        key, "copy", [&]() { return master_client_.CopyEnd(key); },
-        [&]() { return master_client_.CopyRevoke(key); }, response.source,
-        response.targets);
+        key, "copy", [&]() { return master_client_.CopyEnd(key, tenant_id); },
+        [&]() { return master_client_.CopyRevoke(key, tenant_id); },
+        response.source, response.targets);
 
     if (result.has_value()) {
         LOG(INFO) << "action=replica_copy_success" << ", key=" << key
@@ -2414,12 +3253,20 @@ tl::expected<void, ErrorCode> Client::Copy(
 tl::expected<void, ErrorCode> Client::Move(const std::string& key,
                                            const std::string& source,
                                            const std::string& target) {
+    return Move(key, master_client_.tenant_id(), source, target);
+}
+
+tl::expected<void, ErrorCode> Client::Move(const std::string& key,
+                                           const std::string& tenant_id,
+                                           const std::string& source,
+                                           const std::string& target) {
     LOG(INFO) << "action=replica_move_start" << ", key=" << key
               << ", source_segment=" << source << ", target_segment=" << target;
 
     // Call MoveStart first - it validates existence and allocates replica if
     // needed
-    auto move_start_result = master_client_.MoveStart(key, source, target);
+    auto move_start_result =
+        master_client_.MoveStart(key, tenant_id, source, target);
     if (!move_start_result.has_value()) {
         ErrorCode error = move_start_result.error();
         LOG(ERROR) << "action=replica_move_failed" << ", key=" << key
@@ -2433,7 +3280,7 @@ tl::expected<void, ErrorCode> Client::Move(const std::string& key,
         LOG(INFO) << "action=replica_move_skipped" << ", key=" << key
                   << ", info=target_replica_already_exists";
         // Target already exists, consider it success
-        auto move_end_result = master_client_.MoveEnd(key);
+        auto move_end_result = master_client_.MoveEnd(key, tenant_id);
         if (!move_end_result.has_value()) {
             ErrorCode error = move_end_result.error();
             LOG(ERROR) << "action=replica_move_failed" << ", key=" << key
@@ -2446,9 +3293,9 @@ tl::expected<void, ErrorCode> Client::Move(const std::string& key,
     std::vector<Replica::Descriptor> targets = {response.target.value()};
 
     auto result = ExecuteReplicaTransfer(
-        key, "move", [&]() { return master_client_.MoveEnd(key); },
-        [&]() { return master_client_.MoveRevoke(key); }, response.source,
-        targets);
+        key, "move", [&]() { return master_client_.MoveEnd(key, tenant_id); },
+        [&]() { return master_client_.MoveRevoke(key, tenant_id); },
+        response.source, targets);
 
     if (result.has_value()) {
         LOG(INFO) << "action=replica_move_success" << ", key=" << key
@@ -2518,7 +3365,12 @@ void Client::PutToLocalFile(const std::string& key,
                            << ", triggering PutRevoke for disk replica";
                 pinned_buffer_pool_->Release(buf);
                 // Must revoke to avoid phantom replica in master
-                master_client_.PutRevoke(key, ReplicaType::DISK);
+                auto revoke_result =
+                    master_client_.PutRevoke(key, ReplicaType::DISK);
+                if (!revoke_result) {
+                    LOG(ERROR)
+                        << "Failed to revoke put operation for key: " << key;
+                }
                 return;
             }
             value.append(buf.data, slice.size);
@@ -2576,8 +3428,20 @@ ErrorCode Client::TransferData(const Replica::Descriptor& replica_descriptor,
         return ErrorCode::INVALID_PARAMS;
     }
 
-    auto future =
-        transfer_submitter_->submit(replica_descriptor, slices, op_code);
+    std::optional<TransferFuture> future;
+    if (replica_descriptor.is_nof_replica()) {
+        auto contiguous_range = GetContiguousSliceRange(slices);
+        if (!contiguous_range.has_value()) {
+            LOG(ERROR) << "NoF transfer requires contiguous slices";
+            return ErrorCode::INVALID_PARAMS;
+        }
+        future = transfer_submitter_->submit(replica_descriptor, slices,
+                                             op_code, contiguous_range->ptr,
+                                             contiguous_range->size);
+    } else {
+        future =
+            transfer_submitter_->submit(replica_descriptor, slices, op_code);
+    }
     if (!future) {
         LOG(ERROR) << "Failed to submit transfer operation";
         return ErrorCode::TRANSFER_FAIL;
@@ -2619,9 +3483,15 @@ ErrorCode Client::TransferRead(const Replica::Descriptor& replica_descriptor,
     if (replica_descriptor.is_memory_replica()) {
         auto& mem_desc = replica_descriptor.get_memory_descriptor();
         total_size = mem_desc.buffer_descriptor.size_;
-    } else {
+    } else if (replica_descriptor.is_nof_replica()) {
+        auto& nof_desc = replica_descriptor.get_nof_descriptor();
+        total_size = nof_desc.buffer_descriptor.size_;
+    } else if (replica_descriptor.is_disk_replica()) {
         auto& disk_desc = replica_descriptor.get_disk_descriptor();
         total_size = disk_desc.object_size;
+    } else if (replica_descriptor.is_local_disk_replica()) {
+        auto& disk_desc = replica_descriptor.get_local_disk_descriptor();
+        total_size = disk_desc.object_size;
     }
 
     size_t slices_size = CalculateSliceSize(slices);
@@ -2698,8 +3568,8 @@ void Client::ExecuteTask(const ClientTask& client_task) {
             case TaskType::REPLICA_COPY: {
                 ReplicaCopyPayload payload;
                 struct_json::from_json(payload, assignment.payload);
-                auto copy_result =
-                    Copy(payload.key, payload.source, payload.targets);
+                auto copy_result = Copy(payload.key, payload.tenant_id,
+                                        payload.source, payload.targets);
                 if (copy_result.has_value()) {
                     result = ErrorCode::OK;
                 } else {
@@ -2710,8 +3580,8 @@ void Client::ExecuteTask(const ClientTask& client_task) {
             case TaskType::REPLICA_MOVE: {
                 ReplicaMovePayload payload;
                 struct_json::from_json(payload, assignment.payload);
-                auto move_result =
-                    Move(payload.key, payload.source, payload.target);
+                auto move_result = Move(payload.key, payload.tenant_id,
+                                        payload.source, payload.target);
                 if (move_result.has_value()) {
                     result = ErrorCode::OK;
                 } else {
@@ -2822,6 +3692,43 @@ void Client::StorageHeartbeatThreadMain() {
             ErrorCode err = remount_result.error();
             LOG(ERROR) << "Failed to remount segments: " << err;
         }
+        // Re-publish Transfer Engine segment descriptors to the HTTP
+        // metadata server.  When Master (which hosts the HTTP metadata
+        // server in the same process) is killed and restarted, all
+        // in-memory KV entries are lost.  ReMountSegment above only
+        // restores Master-side allocation state; it does NOT write back
+        // the transport-level segment descriptors.  Without this, remote
+        // peers get HTTP 404 when querying our segment descriptor and
+        // data transfers fail.
+        auto metadata = transfer_engine_->getMetadata();
+        if (metadata) {
+            int rc = metadata->updateLocalSegmentDesc();
+            if (rc != 0) {
+                LOG(ERROR) << "Failed to re-publish segment descriptor "
+                           << "to metadata server, rc=" << rc
+                           << ", will retry in next heartbeat cycle";
+                segment_desc_publish_pending_.store(true);
+            } else {
+                segment_desc_publish_pending_.store(false);
+            }
+            // Also re-publish RPC meta entry (mooncake/rpc_meta/<hostname>).
+            // Remote peers need this to locate our RDMA RPC port for
+            // handshake.  Like segment descriptors, this entry is lost
+            // when the HTTP metadata server is cleared on Master restart.
+            rc = metadata->rePublishRpcMetaEntry(local_hostname_);
+            if (rc != 0) {
+                LOG(ERROR) << "Failed to re-publish RPC meta entry "
+                           << "to metadata server, rc=" << rc
+                           << ", will retry in next heartbeat cycle";
+                rpc_meta_publish_pending_.store(true);
+            } else {
+                rpc_meta_publish_pending_.store(false);
+            }
+        }
+        // Note: LOCAL_DISK segment remount is NOT done here.
+        // It is handled by FileStorage::Heartbeat() when it detects
+        // SEGMENT_NOT_FOUND, which also triggers ScanMeta to
+        // re-register offloaded object metadata.
     };
     // Use another thread to remount segments to avoid blocking the ping
     // thread
@@ -2847,6 +3754,41 @@ void Client::StorageHeartbeatThreadMain() {
                 // Ensure at most one remount segment thread is running
                 remount_segment_future =
                     std::async(std::launch::async, remount_segment);
+            } else if (segment_desc_publish_pending_.load() &&
+                       !remount_segment_future.valid()) {
+                // Previous remount succeeded but updateLocalSegmentDesc()
+                // failed (e.g. transient HTTP error).  Retry it directly
+                // without re-running ReMountSegment.
+                auto metadata = transfer_engine_->getMetadata();
+                if (metadata) {
+                    int rc = metadata->updateLocalSegmentDesc();
+                    if (rc != 0) {
+                        LOG(ERROR)
+                            << "Retry: failed to re-publish segment "
+                            << "descriptor to metadata server, rc=" << rc;
+                    } else {
+                        LOG(INFO) << "Retry: successfully re-published "
+                                  << "segment descriptor to metadata server";
+                        segment_desc_publish_pending_.store(false);
+                    }
+                }
+            } else if (rpc_meta_publish_pending_.load() &&
+                       !remount_segment_future.valid()) {
+                // Previous remount succeeded but rePublishRpcMetaEntry()
+                // failed.  Retry it directly.
+                auto metadata = transfer_engine_->getMetadata();
+                if (metadata) {
+                    int rc = metadata->rePublishRpcMetaEntry(local_hostname_);
+                    if (rc != 0) {
+                        LOG(ERROR)
+                            << "Retry: failed to re-publish RPC "
+                            << "meta entry to metadata server, rc=" << rc;
+                    } else {
+                        LOG(INFO) << "Retry: successfully re-published "
+                                  << "RPC meta entry to metadata server";
+                        rpc_meta_publish_pending_.store(false);
+                    }
+                }
             }
 
             std::this_thread::sleep_for(
@@ -2951,6 +3893,7 @@ tl::expected<Replica::Descriptor, ErrorCode> Client::GetPreferredReplica(
         }
     }
 
+    // Prefer local MEMORY replicas first
     for (const auto& rep : replica_list) {
         if (rep.is_memory_replica()) {
             const auto& mem_desc = rep.get_memory_descriptor();
@@ -2962,6 +3905,18 @@ tl::expected<Replica::Descriptor, ErrorCode> Client::GetPreferredReplica(
         }
     }
 
+    // Then prefer local NOF_SSD replicas
+    for (const auto& rep : replica_list) {
+        if (rep.is_nof_replica()) {
+            const auto& nof_desc = rep.get_nof_descriptor();
+            const std::string& endpoint =
+                nof_desc.buffer_descriptor.transport_endpoint_;
+            if (local_endpoints.count(endpoint)) {
+                return rep;
+            }
+        }
+    }
+
     return replica_list[0];
 }
 
@@ -3020,6 +3975,11 @@ size_t Client::GetLocalHotBlockSizeFromEnv(size_t default_value) {
 }
 
 ErrorCode Client::InitLocalHotCache() {
+    hot_cache_handler_.reset();
+    UnregisterLocalHotCacheMemory();
+    hot_cache_.reset();
+    admission_sketch_.reset();
+
     // Defaults: hot cache is disabled unless MC_STORE_LOCAL_HOT_CACHE_SIZE is
     // set to a positive value; when enabled, default block size is 16MB and
     // thread_num is 2.
@@ -3030,9 +3990,6 @@ ErrorCode Client::InitLocalHotCache() {
     size_t total_cache = GetLocalHotCacheSizeFromEnv();
     if (total_cache == 0) {
         // Environment variable not set or invalid, disable cache
-        hot_cache_.reset();
-        hot_cache_handler_.reset();
-        admission_sketch_.reset();
         return ErrorCode::OK;
     }
 
@@ -3060,10 +4017,29 @@ ErrorCode Client::InitLocalHotCache() {
             admission_sketch_.reset();
             return ErrorCode::INVALID_PARAMS;
         }
+
+        int rc = transfer_engine_->registerLocalMemory(
+            hot_cache_->GetBaseAddress(), hot_cache_->GetTotalSize(),
+            kWildcardLocation, true, true);
+        if (rc != 0) {
+            LOG(ERROR)
+                << "Failed to register local hot cache memory with transfer "
+                   "engine, base="
+                << hot_cache_->GetBaseAddress()
+                << ", size=" << hot_cache_->GetTotalSize() << ", ret=" << rc;
+            hot_cache_.reset();
+            hot_cache_handler_.reset();
+            admission_sketch_.reset();
+            hot_cache_memory_registered_ = false;
+            return ErrorCode::INVALID_PARAMS;
+        }
+        hot_cache_memory_registered_ = true;
+
         LOG(INFO) << "Local hot cache enabled with cache size=" << total_cache
                   << ", block size=" << block_size
                   << ", block amount=" << hot_cache_->GetCacheSize()
-                  << ", shm=" << (use_shm ? "on" : "off");
+                  << ", shm=" << (use_shm ? "on" : "off")
+                  << ", transfer engine registered=on";
         // Create async handler with 2 worker threads
         hot_cache_handler_ =
             std::make_unique<LocalHotCacheHandler>(hot_cache_, thread_num);
@@ -3092,6 +4068,27 @@ ErrorCode Client::InitLocalHotCache() {
     return ErrorCode::OK;
 }
 
+void Client::UnregisterLocalHotCacheMemory() {
+    if (!(hot_cache_ && hot_cache_memory_registered_)) {
+        return;
+    }
+    if (!transfer_engine_) {
+        hot_cache_memory_registered_ = false;
+        return;
+    }
+
+    int rc = transfer_engine_->unregisterLocalMemory(
+        hot_cache_->GetBaseAddress(), true);
+    if (rc != 0 && rc != ERR_ADDRESS_NOT_REGISTERED) {
+        LOG(ERROR)
+            << "Failed to unregister local hot cache memory from transfer "
+               "engine, base="
+            << hot_cache_->GetBaseAddress()
+            << ", size=" << hot_cache_->GetTotalSize() << ", ret=" << rc;
+    }
+    hot_cache_memory_registered_ = false;
+}
+
 void Client::ProcessSlicesAsync(const std::string& key,
                                 const std::vector<Slice>& slices,
                                 const Replica::Descriptor& replica) {
diff --git a/mooncake-store/src/dummy_client.cpp b/mooncake-store/src/dummy_client.cpp
index 8c5365de..c1a84a5c 100644
--- a/mooncake-store/src/dummy_client.cpp
+++ b/mooncake-store/src/dummy_client.cpp
@@ -117,17 +117,30 @@ std::vector<std::vector<uint64_t>> void_ptr_rows_to_u64_nested(
 
 namespace mooncake {
 
+template <auto ServiceMethod>
+constexpr bool can_invoke_when_disconnected() {
+    using Method = std::remove_reference_t<decltype(ServiceMethod)>;
+    return std::is_same_v<
+               Method, std::remove_reference_t<decltype(&RealClient::ping)>> ||
+           std::is_same_v<Method,
+                          std::remove_reference_t<
+                              decltype(&RealClient::service_ready_internal)>> ||
+           std::is_same_v<Method,
+                          std::remove_reference_t<
+                              decltype(&RealClient::is_shm_mapped_internal)>> ||
+           std::is_same_v<Method,
+                          std::remove_reference_t<
+                              decltype(&RealClient::ascend_shm_internal)>> ||
+           std::is_same_v<Method,
+                          std::remove_reference_t<
+                              decltype(&RealClient::ascend_ipc_shm_internal)>>;
+}
+
 template <auto ServiceMethod, typename ReturnType, typename... Args>
 tl::expected<ReturnType, ErrorCode> DummyClient::invoke_rpc(Args&&... args) {
     auto pool = client_accessor_.GetClientPool();
 
-    if constexpr (!std::is_same_v<
-                      std::remove_reference_t<decltype(ServiceMethod)>,
-                      std::remove_reference_t<decltype(&RealClient::ping)>> &&
-                  !std::is_same_v<
-                      std::remove_reference_t<decltype(ServiceMethod)>,
-                      std::remove_reference_t<
-                          decltype(&RealClient::service_ready_internal)>>) {
+    if constexpr (!can_invoke_when_disconnected<ServiceMethod>()) {
         if (!connected_.load()) {
             LOG(ERROR) << "Dummy Client not connected";
             return tl::make_unexpected(ErrorCode::RPC_FAIL);
@@ -266,6 +279,20 @@ ErrorCode DummyClient::connect(const std::string& server_address) {
 int DummyClient::register_ascend_shm(const ShmHelper::ShmSegment* shm,
                                      bool is_local) {
 #ifdef USE_ASCEND_DIRECT
+    const auto dummy_base_addr = reinterpret_cast<uint64_t>(shm->base_addr);
+    auto mapped_result = invoke_rpc<&RealClient::is_shm_mapped_internal, bool>(
+        dummy_base_addr, client_id_);
+    if (!mapped_result.has_value()) {
+        LOG(WARNING) << "Failed to query real-side shared memory mapping, addr="
+                     << shm->base_addr;
+        return -1;
+    }
+    if (mapped_result.value()) {
+        LOG(INFO) << "Real-side shared memory mapping already exists, addr="
+                  << shm->base_addr << ", size=" << shm->size;
+        return 0;
+    }
+
     // Detect memory type: device memory uses IPC sharing
     aclrtPtrAttributes attributes;
     auto ret = aclrtPointerGetAttributes(shm->base_addr, &attributes);
@@ -288,8 +315,8 @@ int DummyClient::register_ascend_shm(const ShmHelper::ShmSegment* shm,
 
         std::string ipc_key_bytes(ipc_key, kIPCKeyLen);
         auto map_ret = invoke_rpc<&RealClient::ascend_ipc_shm_internal, void>(
-            reinterpret_cast<uint64_t>(shm->base_addr), shm->size, is_local,
-            ipc_key_bytes, device_id_, client_id_);
+            dummy_base_addr, shm->size, is_local, ipc_key_bytes, device_id_,
+            client_id_);
         if (!map_ret.has_value()) {
             LOG(ERROR) << "Failed to map IPC buffer on real side";
             return -1;
@@ -299,8 +326,15 @@ int DummyClient::register_ascend_shm(const ShmHelper::ShmSegment* shm,
         return 0;
     }
     if (!globalConfig().ascend_use_fabric_mem) {
-        LOG(ERROR) << "Host mem is only supported in fabric mem mode.";
-        return -1;
+        // Host: memfd + mmap (ShmHelper); register with Real like non-agent GPU
+        // path.
+        if (shm->fd < 0) {
+            LOG(ERROR)
+                << "Host POSIX shared memory requires memfd-backed allocation "
+                   "(use ShmHelper::allocate / alloc_from_mem_pool)";
+            return -1;
+        }
+        return register_shm_via_ipc(shm, is_local);
     }
 
     // Fabric host mem shared by vmm
@@ -324,8 +358,8 @@ int DummyClient::register_ascend_shm(const ShmHelper::ShmSegment* shm,
     std::string handle_bytes(reinterpret_cast<char*>(&export_handle),
                              sizeof(export_handle));
     auto map_ret = invoke_rpc<&RealClient::ascend_shm_internal, void>(
-        reinterpret_cast<uint64_t>(shm->base_addr), shm->size, is_local,
-        handle_bytes, device_id_, client_id_);
+        dummy_base_addr, shm->size, is_local, handle_bytes, device_id_,
+        client_id_);
     if (!map_ret.has_value()) {
         LOG(ERROR) << "Failed to map VMM buffer on real side";
         return -1;
@@ -385,6 +419,8 @@ int DummyClient::register_shm_via_ipc(const ShmHelper::ShmSegment* shm,
     req.client_id_second = client_id_.second;
     req.dummy_base_addr = reinterpret_cast<uintptr_t>(shm->base_addr);
     req.shm_size = shm->size;
+    req.device_id = globalConfig().ascend_agent_mode ? device_id_
+                                                     : kInvalidPhysicalDeviceId;
     req.is_local_buffer = is_local;
 
     if (ipc_send_fd(sock_fd, shm->fd, &req, sizeof(req)) < 0) {
@@ -482,16 +518,16 @@ int DummyClient::setup_dummy(size_t mem_pool_size, size_t local_buffer_size,
         }
         local_buffer_shm->registered = true;
         local_buffer_shm->is_local = true;
+
+        // Best-effort: request hot cache shm from real client
+        if (request_hot_cache_fd() != 0) {
+            LOG(INFO)
+                << "Hot cache shm not available (real client may not have it)";
+        }
     }
 
     ping_running_ = true;
     ping_thread_ = std::thread([this]() mutable { this->ping_thread_main(); });
-
-    // Best-effort: request hot cache shm from real client
-    if (request_hot_cache_fd() != 0) {
-        LOG(INFO)
-            << "Hot cache shm not available (real client may not have it)";
-    }
     return 0;
 }
 
@@ -517,6 +553,12 @@ int DummyClient::tearDownAll() {
     }
     connected_.store(false);
     last_ping_healthy_.store(false);
+#if defined(USE_ASCEND_DIRECT)
+    {
+        std::lock_guard<std::mutex> lock(registered_device_buffers_mutex_);
+        registered_device_buffers_.clear();
+    }
+#endif
     return 0;
 }
 
@@ -533,22 +575,103 @@ int64_t DummyClient::unregister_shm() {
         invoke_rpc<&RealClient::unmap_shm_internal, void>(client_id_));
 }
 
+#if defined(USE_ASCEND_DIRECT)
+int DummyClient::register_device_buffer_for_reconnect(void* buffer,
+                                                      size_t size) {
+    const auto buffer_addr = reinterpret_cast<uint64_t>(buffer);
+    {
+        std::lock_guard<std::mutex> lock(registered_device_buffers_mutex_);
+        auto it = registered_device_buffers_.find(buffer_addr);
+        if (it != registered_device_buffers_.end() && it->second != size) {
+            LOG(ERROR) << "Device buffer size mismatch for tracked buffer, "
+                       << "buffer=" << buffer << ", size=" << size
+                       << ", tracked_size=" << it->second;
+            return -1;
+        }
+    }
+
+    ShmHelper::ShmSegment shm{};
+    shm.base_addr = buffer;
+    shm.size = size;
+    if (register_ascend_shm(&shm, false) != 0) {
+        LOG(ERROR) << "Failed to register device buffer, buffer=" << buffer
+                   << ", size=" << size;
+        return -1;
+    }
+
+    {
+        std::lock_guard<std::mutex> lock(registered_device_buffers_mutex_);
+        registered_device_buffers_[buffer_addr] = size;
+    }
+    return 0;
+}
+
+int DummyClient::unregister_device_buffer_for_reconnect(void* buffer) {
+    const auto buffer_addr = reinterpret_cast<uint64_t>(buffer);
+    {
+        std::lock_guard<std::mutex> lock(registered_device_buffers_mutex_);
+        if (registered_device_buffers_.find(buffer_addr) ==
+            registered_device_buffers_.end()) {
+            LOG(ERROR) << "Device buffer is not registered with RealClient";
+            return -1;
+        }
+    }
+
+    auto ret = invoke_rpc<&RealClient::unregister_shm_buffer_internal, void>(
+        buffer_addr, client_id_);
+    if (ret.has_value()) {
+        std::lock_guard<std::mutex> lock(registered_device_buffers_mutex_);
+        registered_device_buffers_.erase(buffer_addr);
+    }
+    return to_py_ret(ret);
+}
+
+std::vector<ShmHelper::ShmSegment> DummyClient::get_registered_device_buffers()
+    const {
+    std::vector<ShmHelper::ShmSegment> device_buffers;
+    std::lock_guard<std::mutex> lock(registered_device_buffers_mutex_);
+    device_buffers.reserve(registered_device_buffers_.size());
+    for (const auto& [buffer_addr, size] : registered_device_buffers_) {
+        ShmHelper::ShmSegment shm{};
+        shm.base_addr = reinterpret_cast<void*>(buffer_addr);
+        shm.size = size;
+        device_buffers.push_back(std::move(shm));
+    }
+    return device_buffers;
+}
+#endif
+
 // Dummy only register buffer within the shared memory region
 int DummyClient::register_buffer(void* buffer, size_t size) {
     if (buffer == nullptr) {
         LOG(ERROR) << "Invalid buffer pointer";
         return -1;
     }
+#if defined(USE_ASCEND_DIRECT)
     if (globalConfig().ascend_agent_mode) {
-        auto shm = std::make_shared<ShmHelper::ShmSegment>();
-        shm->base_addr = buffer;
-        shm->size = size;
-        if (register_ascend_shm(shm.get(), false) != 0) {
-            LOG(ERROR) << "Failed to implicitly register new ascend shm.";
+        aclrtPtrAttributes attributes{};
+        auto acl_ret = aclrtPointerGetAttributes(buffer, &attributes);
+        if (acl_ret != ACL_ERROR_NONE) {
+            LOG(ERROR) << "Failed to get pointer attributes, ret=" << acl_ret;
             return -1;
         }
-        return 0;
+        if (attributes.location.type == ACL_MEM_LOCATION_TYPE_DEVICE) {
+            return register_device_buffer_for_reconnect(buffer, size);
+        }
+        if (globalConfig().ascend_use_fabric_mem) {
+            auto shm = std::make_shared<ShmHelper::ShmSegment>();
+            shm->base_addr = buffer;
+            shm->size = size;
+            if (register_ascend_shm(shm.get(), false) != 0) {
+                LOG(ERROR) << "Failed to register buffer, buffer=" << buffer
+                           << ", size=" << size;
+                return -1;
+            }
+            return 0;
+        }
+        // non-Fabric Host: same rules as shm
     }
+#endif
     // Find which shm this buffer belongs to
     auto shm = shm_helper_->get_shm(buffer);
     if (!shm) {
@@ -587,6 +710,17 @@ int DummyClient::unregister_buffer(void* buffer) {
         return -1;
     }
 
+#if defined(USE_ASCEND_DIRECT)
+    if (globalConfig().ascend_agent_mode) {
+        aclrtPtrAttributes attributes{};
+        auto acl_ret = aclrtPointerGetAttributes(buffer, &attributes);
+        if (acl_ret == ACL_ERROR_NONE &&
+            attributes.location.type == ACL_MEM_LOCATION_TYPE_DEVICE) {
+            return unregister_device_buffer_for_reconnect(buffer);
+        }
+    }
+#endif
+
     auto shm = shm_helper_->get_shm(buffer);
     if (!shm) {
         LOG(ERROR) << "Buffer is not in any registered shared memory";
@@ -906,15 +1040,18 @@ std::vector<std::vector<std::vector<int64_t>>> DummyClient::get_into_ranges(
     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) {
+    const std::vector<std::vector<std::vector<size_t>>>& all_sizes,
+    const QueryResultCache* query_result_cache) {
     std::vector<uint64_t> dummy_buffers = void_ptrs_to_u64(buffers);
+    auto cached_query_results =
+        build_cached_query_results_from_query_result_cache(query_result_cache);
     const auto start_time = std::chrono::steady_clock::now();
     auto internal_results =
         invoke_rpc<&RealClient::get_into_ranges_shm_helper,
                    std::vector<std::vector<
                        std::vector<tl::expected<int64_t, ErrorCode>>>>>(
             dummy_buffers, all_keys, all_dst_offsets, all_src_offsets,
-            all_sizes, device_id_, client_id_);
+            all_sizes, cached_query_results, device_id_, client_id_);
 
     if (!internal_results) {
         LOG(ERROR) << "get_into_ranges RPC failed";
@@ -931,6 +1068,39 @@ std::vector<std::vector<std::vector<int64_t>>> DummyClient::get_into_ranges(
     return results;
 }
 
+std::vector<tl::expected<QueryResult, ErrorCode>> DummyClient::batch_query(
+    const std::vector<std::string>& keys) {
+    auto cached_results =
+        invoke_rpc<&RealClient::batch_get_query_results,
+                   std::vector<CachedQueryResultResponse>>(keys);
+    if (!cached_results) {
+        return std::vector<tl::expected<QueryResult, ErrorCode>>(
+            keys.size(), tl::unexpected(cached_results.error()));
+    }
+    if (cached_results->size() != keys.size()) {
+        LOG(ERROR) << "BatchQuery response size mismatch: expected "
+                   << keys.size() << ", got " << cached_results->size();
+        return std::vector<tl::expected<QueryResult, ErrorCode>>(
+            keys.size(), tl::unexpected(ErrorCode::RPC_FAIL));
+    }
+
+    std::vector<tl::expected<QueryResult, ErrorCode>> results;
+    results.reserve(keys.size());
+    const auto now = std::chrono::steady_clock::now();
+    for (const auto& cached_result : *cached_results) {
+        if (!cached_result.success) {
+            results.emplace_back(tl::unexpected(cached_result.error));
+            continue;
+        }
+        results.emplace_back(QueryResult(
+            std::vector<Replica::Descriptor>(
+                cached_result.value.replicas.begin(),
+                cached_result.value.replicas.end()),
+            now + std::chrono::milliseconds(cached_result.value.lease_ttl_ms)));
+    }
+    return results;
+}
+
 std::string DummyClient::get_hostname() const {
     // Dummy client does not have a hostname
     return "";
@@ -1141,24 +1311,40 @@ void DummyClient::ping_thread_main() {
                         if (globalConfig().ascend_agent_mode) {
                             if (register_ascend_shm(shm_ptr.get(),
                                                     shm_ptr->is_local) != 0) {
-                                LOG(WARNING) << "Failed to re-register VMM "
-                                                "during reconnection";
-                                all_registered = false;
-                                break;
-                            }
-                        } else {
-                            if (register_shm_via_ipc(shm_ptr.get(),
-                                                     shm_ptr->is_local) != 0) {
                                 LOG(WARNING)
-                                    << "Failed to re-register shared memory "
-                                       "during reconnection";
+                                    << "Failed to re-register Ascend shared "
+                                       "memory during reconnection";
                                 all_registered = false;
                                 break;
                             }
+                        } else if (register_shm_via_ipc(
+                                       shm_ptr.get(), shm_ptr->is_local) != 0) {
+                            LOG(WARNING)
+                                << "Failed to re-register shared memory "
+                                   "during reconnection";
+                            all_registered = false;
+                            break;
                         }
                     }
                 }
 
+#if defined(USE_ASCEND_DIRECT)
+                if (all_registered && globalConfig().ascend_agent_mode) {
+                    auto device_buffers = get_registered_device_buffers();
+                    for (const auto& device_buffer : device_buffers) {
+                        if (register_ascend_shm(&device_buffer, false) != 0) {
+                            LOG(WARNING)
+                                << "Failed to re-register device buffer "
+                                   "during reconnection, buffer="
+                                << device_buffer.base_addr
+                                << ", size=" << device_buffer.size;
+                            all_registered = false;
+                            break;
+                        }
+                    }
+                }
+#endif
+
                 if (all_registered) {
                     LOG(INFO)
                         << "Re-registered all shared memorys successfully";
diff --git a/mooncake-store/src/engram/engram_store.cpp b/mooncake-store/src/engram/engram_store.cpp
new file mode 100644
index 00000000..7c6def08
--- /dev/null
+++ b/mooncake-store/src/engram/engram_store.cpp
@@ -0,0 +1,334 @@
+#include "engram/engram_store.h"
+
+#include <cstring>
+#include <limits>
+#include <sstream>
+#include <stdexcept>
+#include <vector>
+
+#include "pyclient.h"
+
+namespace mooncake {
+namespace engram {
+
+EngramStore::EngramStore(int layer_id, const EngramStoreConfig& config,
+                         std::shared_ptr<PyClient> store)
+    : store_(std::move(store)),
+      table_vocab_sizes_(config.table_vocab_sizes),
+      embedding_dim_(config.embedding_dim) {
+    if (table_vocab_sizes_.empty()) {
+        throw std::invalid_argument(
+            "EngramStoreConfig.table_vocab_sizes must not be empty");
+    }
+    if (embedding_dim_ <= 0) {
+        throw std::invalid_argument(
+            "EngramStoreConfig.embedding_dim must be positive");
+    }
+    for (int64_t vocab_size : table_vocab_sizes_) {
+        if (vocab_size <= 0) {
+            throw std::invalid_argument(
+                "EngramStoreConfig.table_vocab_sizes must contain only "
+                "positive values");
+        }
+    }
+
+    embed_keys_.reserve(table_vocab_sizes_.size());
+    for (size_t h = 0; h < table_vocab_sizes_.size(); ++h) {
+        std::ostringstream oss;
+        oss << "engram:l" << layer_id << ":h" << h;
+        embed_keys_.push_back(oss.str());
+    }
+}
+
+int EngramStore::lookup_rows_flat(const int64_t* row_ids, int B, int L,
+                                  void* output_buffer,
+                                  size_t output_size) const {
+    if (store_ == nullptr || row_ids == nullptr || output_buffer == nullptr ||
+        B <= 0 || L <= 0) {
+        return -1;
+    }
+
+    const int num_heads = static_cast<int>(table_vocab_sizes_.size());
+    const size_t row_bytes =
+        static_cast<size_t>(embedding_dim_) * sizeof(float);
+    const size_t max_size = std::numeric_limits<size_t>::max();
+    if (static_cast<size_t>(B) > max_size / static_cast<size_t>(L)) {
+        return -1;
+    }
+    const size_t token_count = static_cast<size_t>(B) * static_cast<size_t>(L);
+    if (token_count > max_size / static_cast<size_t>(num_heads) ||
+        token_count * static_cast<size_t>(num_heads) > max_size / row_bytes) {
+        return -1;
+    }
+    const size_t expected_size =
+        token_count * static_cast<size_t>(num_heads) * row_bytes;
+    if (output_size < expected_size) {
+        return -1;
+    }
+
+    auto fail_lookup = [&]() {
+        std::memset(output_buffer, 0, expected_size);
+        return -1;
+    };
+
+    std::vector<void*> buffers{output_buffer};
+    std::vector<std::vector<std::string>> all_keys(1);
+    std::vector<std::vector<std::vector<size_t>>> all_dst_offsets(1);
+    std::vector<std::vector<std::vector<size_t>>> all_src_offsets(1);
+    std::vector<std::vector<std::vector<size_t>>> all_sizes(1);
+
+    all_keys[0].reserve(static_cast<size_t>(num_heads));
+    all_dst_offsets[0].reserve(static_cast<size_t>(num_heads));
+    all_src_offsets[0].reserve(static_cast<size_t>(num_heads));
+    all_sizes[0].reserve(static_cast<size_t>(num_heads));
+
+    for (int h = 0; h < num_heads; ++h) {
+        all_keys[0].push_back(embed_keys_[h]);
+        all_dst_offsets[0].emplace_back();
+        all_src_offsets[0].emplace_back();
+        all_sizes[0].emplace_back();
+        all_dst_offsets[0].back().reserve(static_cast<size_t>(B) * L);
+        all_src_offsets[0].back().reserve(static_cast<size_t>(B) * L);
+        all_sizes[0].back().reserve(static_cast<size_t>(B) * L);
+    }
+
+    for (int b = 0; b < B; ++b) {
+        for (int l = 0; l < L; ++l) {
+            const size_t token_index = static_cast<size_t>(b) * L + l;
+            const size_t row_offset =
+                token_index * static_cast<size_t>(num_heads);
+            for (int h = 0; h < num_heads; ++h) {
+                const int64_t idx =
+                    row_ids[row_offset + static_cast<size_t>(h)];
+                if (idx < 0 || idx >= table_vocab_sizes_[h]) {
+                    return fail_lookup();
+                }
+                all_dst_offsets[0][h].push_back(
+                    (row_offset + static_cast<size_t>(h)) * row_bytes);
+                all_src_offsets[0][h].push_back(static_cast<size_t>(idx) *
+                                                row_bytes);
+                all_sizes[0][h].push_back(row_bytes);
+            }
+        }
+    }
+
+    const int register_ret =
+        store_->register_buffer(output_buffer, expected_size);
+    if (register_ret != 0) {
+        return fail_lookup();
+    }
+
+    mooncake::PyClient::QueryResultCache query_result_cache;
+    auto query_results = store_->batch_query(embed_keys_);
+    if (query_results.size() != embed_keys_.size()) {
+        store_->unregister_buffer(output_buffer);
+        return fail_lookup();
+    }
+    query_result_cache.reserve(embed_keys_.size());
+    for (size_t i = 0; i < embed_keys_.size(); ++i) {
+        query_result_cache.emplace(embed_keys_[i], query_results[i]);
+    }
+
+    auto results = store_->get_into_ranges(buffers, all_keys, all_dst_offsets,
+                                           all_src_offsets, all_sizes,
+                                           &query_result_cache);
+    const int unregister_ret = store_->unregister_buffer(output_buffer);
+    if (unregister_ret != 0) {
+        return fail_lookup();
+    }
+    if (results.size() != 1 ||
+        results[0].size() != static_cast<size_t>(num_heads)) {
+        return fail_lookup();
+    }
+    for (int h = 0; h < num_heads; ++h) {
+        if (results[0][h].size() != all_sizes[0][h].size()) {
+            return fail_lookup();
+        }
+        for (int64_t bytes_read : results[0][h]) {
+            if (bytes_read != static_cast<int64_t>(row_bytes)) {
+                return fail_lookup();
+            }
+        }
+    }
+
+    return 0;
+}
+
+int EngramStore::lookup_rows_contiguous(const int64_t* row_ids, int B, int L,
+                                        void* output_buffer,
+                                        size_t output_size) const {
+    return lookup_rows_flat(row_ids, B, L, output_buffer, output_size);
+}
+
+int EngramStore::lookup_rows(
+    const std::vector<std::vector<std::vector<int64_t>>>& row_ids,
+    void* output_buffer, size_t output_size) const {
+    if (store_ == nullptr || output_buffer == nullptr || row_ids.empty() ||
+        row_ids[0].empty()) {
+        return -1;
+    }
+
+    const int B = static_cast<int>(row_ids.size());
+    const int L = static_cast<int>(row_ids[0].size());
+    const int num_heads = static_cast<int>(table_vocab_sizes_.size());
+    std::vector<int64_t> flat_row_ids;
+    flat_row_ids.reserve(static_cast<size_t>(B) * L * num_heads);
+    for (int b = 0; b < B; ++b) {
+        if (static_cast<int>(row_ids[b].size()) != L) {
+            return -1;
+        }
+        for (int l = 0; l < L; ++l) {
+            if (static_cast<int>(row_ids[b][l].size()) != num_heads) {
+                return -1;
+            }
+            flat_row_ids.insert(flat_row_ids.end(), row_ids[b][l].begin(),
+                                row_ids[b][l].end());
+        }
+    }
+
+    return lookup_rows_flat(flat_row_ids.data(), B, L, output_buffer,
+                            output_size);
+}
+
+std::vector<int64_t> EngramStore::get_table_vocab_sizes() const {
+    return table_vocab_sizes_;
+}
+
+std::vector<std::string> EngramStore::get_store_keys() const {
+    return embed_keys_;
+}
+
+int EngramStore::get_num_heads() const {
+    return static_cast<int>(table_vocab_sizes_.size());
+}
+
+int EngramStore::get_embedding_dim() const { return embedding_dim_; }
+
+int EngramStore::remove_from_store(bool force) {
+    if (store_ == nullptr) {
+        return static_cast<int>(ErrorCode::INVALID_PARAMS);
+    }
+
+    constexpr int kObjectNotFound =
+        static_cast<int>(ErrorCode::OBJECT_NOT_FOUND);
+    int removed = 0;
+    int first_error = 0;
+
+    for (const auto& key : embed_keys_) {
+        int rc = store_->remove(key, force);
+        if (rc == 0) {
+            ++removed;
+            continue;
+        }
+        if (rc == kObjectNotFound) {
+            continue;
+        }
+        if (first_error == 0) {
+            first_error = rc;
+        }
+    }
+
+    return first_error != 0 ? first_error : removed;
+}
+
+int EngramStore::populate(const std::vector<void*>& embedding_buffers,
+                          const std::vector<size_t>& buffer_sizes) {
+    if (store_ == nullptr) {
+        return -1;
+    }
+    if (embedding_buffers.size() != embed_keys_.size() ||
+        buffer_sizes.size() != embed_keys_.size()) {
+        return -1;
+    }
+
+    for (size_t i = 0; i < buffer_sizes.size(); ++i) {
+        const size_t expected = static_cast<size_t>(table_vocab_sizes_[i]) *
+                                embedding_dim_ * sizeof(float);
+        if (embedding_buffers[i] == nullptr || buffer_sizes[i] != expected) {
+            return -1;
+        }
+    }
+
+    std::vector<int> exists_results = store_->batchIsExist(embed_keys_);
+    if (exists_results.size() != embed_keys_.size()) {
+        LOG(ERROR) << "Failed to preflight EngramStore populate key existence";
+        return -1;
+    }
+    for (size_t i = 0; i < exists_results.size(); ++i) {
+        const int exists = exists_results[i];
+        if (exists < 0) {
+            LOG(ERROR) << "Failed to query EngramStore key '" << embed_keys_[i]
+                       << "' before populate, rc=" << exists;
+            return -1;
+        }
+        if (exists != 0) {
+            LOG(ERROR)
+                << "EngramStore populate requires empty destination key '"
+                << embed_keys_[i] << "'. Remove the existing layer first.";
+            return -1;
+        }
+    }
+
+    auto cleanup_registered_buffers = [&](size_t count) {
+        bool cleanup_failed = false;
+        for (size_t i = 0; i < count; ++i) {
+            int rc = store_->unregister_buffer(embedding_buffers[i]);
+            if (rc != 0) {
+                cleanup_failed = true;
+                LOG(ERROR) << "Failed to unregister embedding buffer at index "
+                           << i << ", rc=" << rc;
+            }
+        }
+        return cleanup_failed;
+    };
+
+    for (size_t i = 0; i < embedding_buffers.size(); ++i) {
+        int ret =
+            store_->register_buffer(embedding_buffers[i], buffer_sizes[i]);
+        if (ret != 0) {
+            if (cleanup_registered_buffers(i)) {
+                LOG(ERROR) << "Failed to clean up registered embedding buffers "
+                              "after register_buffer error";
+            }
+            return -1;
+        }
+    }
+
+    std::vector<int> put_results =
+        store_->batch_put_from(embed_keys_, embedding_buffers, buffer_sizes);
+    const bool put_succeeded =
+        put_results.size() == embed_keys_.size() &&
+        std::all_of(put_results.begin(), put_results.end(),
+                    [](int result) { return result == 0; });
+
+    const bool unregister_failed =
+        cleanup_registered_buffers(embedding_buffers.size());
+
+    if (!put_succeeded || unregister_failed) {
+        const bool put_results_complete =
+            put_results.size() == embed_keys_.size();
+        for (size_t i = 0; i < embed_keys_.size(); ++i) {
+            if (put_results_complete && put_results[i] != 0) {
+                continue;
+            }
+            int rc = store_->remove(embed_keys_[i], true);
+            if (rc != 0 &&
+                rc != static_cast<int>(ErrorCode::OBJECT_NOT_FOUND)) {
+                LOG(ERROR)
+                    << "Failed to roll back partially populated EngramStore "
+                    << "key '" << embed_keys_[i] << "', rc=" << rc;
+            }
+        }
+        if (unregister_failed) {
+            LOG(ERROR)
+                << "Rolling back EngramStore populate because buffer cleanup "
+                   "failed after publish";
+        }
+        return -1;
+    }
+
+    return 0;
+}
+
+}  // namespace engram
+}  // namespace mooncake
diff --git a/mooncake-store/src/etcd_helper.cpp b/mooncake-store/src/etcd_helper.cpp
index 63a565a0..ba096c72 100644
--- a/mooncake-store/src/etcd_helper.cpp
+++ b/mooncake-store/src/etcd_helper.cpp
@@ -42,6 +42,26 @@ ErrorCode EtcdHelper::ConnectToEtcdStoreClient(
     }
 }
 
+ErrorCode EtcdHelper::ResetEtcdStoreClient(const std::string& etcd_endpoints) {
+    std::lock_guard<std::mutex> lock(etcd_mutex_);
+
+    char* err_msg = nullptr;
+    int ret = EtcdStoreResetClientWrapper(
+        const_cast<char*>(etcd_endpoints.c_str()), &err_msg);
+    if (ret != 0) {
+        LOG(ERROR) << "Failed to reset etcd store client: "
+                   << (err_msg == nullptr ? "" : err_msg);
+        if (err_msg != nullptr) {
+            free(err_msg);
+        }
+        return ErrorCode::ETCD_OPERATION_ERROR;
+    }
+
+    connected_endpoints_ = etcd_endpoints;
+    etcd_connected_ = true;
+    return ErrorCode::OK;
+}
+
 ErrorCode EtcdHelper::Get(const char* key, const size_t key_size,
                           std::string& value, EtcdRevisionId& revision_id) {
     char* err_msg = nullptr;
@@ -402,6 +422,11 @@ ErrorCode EtcdHelper::ConnectToEtcdStoreClient(
     return ErrorCode::ETCD_OPERATION_ERROR;
 }
 
+ErrorCode EtcdHelper::ResetEtcdStoreClient(const std::string& etcd_endpoints) {
+    (void)etcd_endpoints;
+    return ErrorCode::UNAVAILABLE_IN_CURRENT_MODE;
+}
+
 ErrorCode EtcdHelper::Get(const char* key, const size_t key_size,
                           std::string& value, EtcdRevisionId& revision_id) {
     (void)key;
diff --git a/mooncake-store/src/file_storage.cpp b/mooncake-store/src/file_storage.cpp
index 5b8cf452..4830c13e 100644
--- a/mooncake-store/src/file_storage.cpp
+++ b/mooncake-store/src/file_storage.cpp
@@ -18,6 +18,26 @@ using gpu_staging::CopyDeviceToHost;
 using gpu_staging::IsDevicePointer;
 using gpu_staging::SetDevice;
 
+namespace {
+
+std::vector<OffloadTaskItem> BuildOffloadTasksFromStorageKeys(
+    const std::vector<std::string>& storage_keys,
+    const std::vector<StorageObjectMetadata>& metadatas) {
+    std::vector<OffloadTaskItem> tasks;
+    tasks.reserve(storage_keys.size());
+    for (size_t i = 0; i < storage_keys.size(); ++i) {
+        auto [tenant_id, key] = ParseTenantScopedStorageKey(storage_keys[i]);
+        const int64_t size =
+            i < metadatas.size() ? metadatas[i].data_size : int64_t{0};
+        tasks.push_back(OffloadTaskItem{.tenant_id = std::move(tenant_id),
+                                        .key = std::move(key),
+                                        .size = size});
+    }
+    return tasks;
+}
+
+}  // namespace
+
 FileStorageConfig FileStorageConfig::FromEnvironment() {
     FileStorageConfig config;
 
@@ -32,6 +52,8 @@ FileStorageConfig FileStorageConfig::FromEnvironment() {
     } else if (storage_backend_descriptor ==
                "offset_allocator_storage_backend") {
         config.storage_backend_type = StorageBackendType::kOffsetAllocator;
+    } else if (storage_backend_descriptor == "distributed_storage_backend") {
+        config.storage_backend_type = StorageBackendType::kDistributed;
     } else {
         LOG(ERROR) << "Unknown storage backend.";
     }
@@ -58,7 +80,7 @@ FileStorageConfig FileStorageConfig::FromEnvironment() {
                            config.heartbeat_interval_seconds);
     config.client_buffer_gc_interval_seconds =
         GetEnvOr<uint32_t>("MOONCAKE_OFFLOAD_CLIENT_BUFFER_GC_INTERVAL_SECONDS",
-                           config.heartbeat_interval_seconds);
+                           config.client_buffer_gc_interval_seconds);
 
     config.client_buffer_gc_ttl_ms =
         GetEnvOr<uint64_t>("MOONCAKE_OFFLOAD_CLIENT_BUFFER_GC_TTL_MS",
@@ -161,8 +183,8 @@ FileStorage::FileStorage(const FileStorageConfig& config,
       ssd_metric_(ssd_metric),
       local_rpc_addr_(local_rpc_addr),
       pinned_buffer_pool_(std::make_unique<PinnedBufferPool>()),
-      client_buffer_allocator_(
-          AlignedClientBufferAllocator::create(config.local_buffer_size, "")) {
+      client_buffer_allocator_(AlignedClientBufferAllocator::create(
+          config.local_buffer_size, client ? client->GetProtocol() : "")) {
     if (!config.Validate()) {
         throw std::invalid_argument("Invalid FileStorage configuration");
     }
@@ -223,6 +245,13 @@ tl::expected<void, ErrorCode> FileStorage::Init() {
         return init_storage_backend_result;
     }
     auto enable_offloading_result = IsEnableOffloading();
+    if (enable_offloading_result.has_value()) {
+        LOG(INFO) << "IsEnableOffloading result: "
+                  << (enable_offloading_result.value() ? "true" : "false");
+    } else {
+        LOG(INFO) << "IsEnableOffloading result: error: "
+                  << enable_offloading_result.error();
+    }
     if (!enable_offloading_result) {
         LOG(ERROR) << "Failed to get enable persist result, error : "
                    << enable_offloading_result.error();
@@ -239,6 +268,17 @@ tl::expected<void, ErrorCode> FileStorage::Init() {
             return mount_file_storage_result;
         }
     }
+    // Report configured SSD capacity to Master so it can populate
+    // file_total_capacity_ (the denominator in "SSD Storage: X / Y").
+    // Called once at init; old Masters that lack this RPC will log an error
+    // but FileStorage continues normally.
+    if (config_.total_size_limit > 0) {
+        auto cap_result = client_->ReportSsdCapacity(config_.total_size_limit);
+        if (!cap_result) {
+            LOG(WARNING) << "ReportSsdCapacity failed (old Master?): "
+                         << cap_result.error();
+        }
+    }
 
     auto scan_meta_result = storage_backend_->ScanMeta(
         [this](const std::vector<std::string>& keys,
@@ -246,8 +286,9 @@ tl::expected<void, ErrorCode> FileStorage::Init() {
             for (auto& metadata : metadatas) {
                 metadata.transport_endpoint = local_rpc_addr_;
             }
+            auto tasks = BuildOffloadTasksFromStorageKeys(keys, metadatas);
             auto add_object_result =
-                client_->NotifyOffloadSuccess(keys, metadatas);
+                client_->NotifyOffloadSuccess(tasks, metadatas);
             if (!add_object_result) {
                 LOG(ERROR) << "Failed to add object to master: "
                            << add_object_result.error();
@@ -321,12 +362,26 @@ tl::expected<FileStorage::BatchGetResult, ErrorCode> FileStorage::BatchGet(
 }
 
 tl::expected<void, ErrorCode> FileStorage::OffloadObjects(
-    const std::unordered_map<std::string, int64_t>& offloading_objects) {
+    const std::vector<OffloadTaskItem>& offloading_objects) {
+    if (offloading_objects.empty()) {
+        return {};
+    }
+    std::unordered_map<std::string, int64_t> storage_object_sizes;
+    std::unordered_map<std::string, OffloadTaskItem> task_by_storage_key;
+    storage_object_sizes.reserve(offloading_objects.size());
+    task_by_storage_key.reserve(offloading_objects.size());
+    for (const auto& task : offloading_objects) {
+        const auto storage_key =
+            MakeTenantScopedStorageKey(task.tenant_id, task.key);
+        storage_object_sizes.emplace(storage_key, task.size);
+        task_by_storage_key.emplace(storage_key, task);
+    }
+
     std::vector<std::vector<std::string>> buckets_keys;
     if (auto bucket_backend =
             std::dynamic_pointer_cast<BucketStorageBackend>(storage_backend_)) {
         auto allocate_res = bucket_backend->AllocateOffloadingBuckets(
-            offloading_objects, buckets_keys);
+            storage_object_sizes, buckets_keys);
         if (!allocate_res) {
             LOG(ERROR) << "AllocateOffloadingBuckets failed with error: "
                        << allocate_res.error();
@@ -334,21 +389,32 @@ tl::expected<void, ErrorCode> FileStorage::OffloadObjects(
         }
     } else {
         std::vector<std::string> keys;
-        keys.reserve(offloading_objects.size());
-        for (const auto& it : offloading_objects) {
+        keys.reserve(storage_object_sizes.size());
+        for (const auto& it : storage_object_sizes) {
             keys.emplace_back(it.first);
         }
         buckets_keys.emplace_back(std::move(keys));
     }
 
     auto complete_handler =
-        [this](const std::vector<std::string>& keys,
-               std::vector<StorageObjectMetadata>& metadatas) -> ErrorCode {
+        [this, &task_by_storage_key](
+            const std::vector<std::string>& keys,
+            std::vector<StorageObjectMetadata>& metadatas) -> ErrorCode {
         VLOG(1) << "Success to store objects, keys count: " << keys.size();
         for (auto& metadata : metadatas) {
             metadata.transport_endpoint = local_rpc_addr_;
         }
-        auto result = client_->NotifyOffloadSuccess(keys, metadatas);
+        std::vector<OffloadTaskItem> tasks;
+        tasks.reserve(keys.size());
+        for (const auto& key : keys) {
+            auto it = task_by_storage_key.find(key);
+            if (it == task_by_storage_key.end()) {
+                LOG(ERROR) << "Offload task not found for storage key";
+                return ErrorCode::INVALID_KEY;
+            }
+            tasks.push_back(it->second);
+        }
+        auto result = client_->NotifyOffloadSuccess(tasks, metadatas);
         if (!result) {
             LOG(ERROR) << "NotifyOffloadSuccess failed with error: "
                        << result.error();
@@ -359,24 +425,63 @@ tl::expected<void, ErrorCode> FileStorage::OffloadObjects(
 
     for (const auto& keys : buckets_keys) {
         std::unordered_map<std::string, std::vector<Slice>> batch_object;
-        auto query_result = BatchQuerySegmentSlices(keys, batch_object);
-        if (!query_result) {
-            LOG(ERROR) << "BatchQuerySlices failed with error: "
-                       << query_result.error();
+        std::unordered_map<std::string, std::vector<std::string>>
+            storage_keys_by_tenant;
+        for (const auto& storage_key : keys) {
+            const auto it = task_by_storage_key.find(storage_key);
+            if (it != task_by_storage_key.end()) {
+                storage_keys_by_tenant[it->second.tenant_id].push_back(
+                    storage_key);
+            }
+        }
+        for (const auto& [tenant_id, storage_keys] : storage_keys_by_tenant) {
+            std::vector<std::string> user_keys;
+            user_keys.reserve(storage_keys.size());
+            for (const auto& storage_key : storage_keys) {
+                user_keys.push_back(task_by_storage_key[storage_key].key);
+            }
+            std::unordered_map<std::string, std::vector<Slice>>
+                user_batch_object;
+            auto query_result = BatchQuerySegmentSlices(user_keys, tenant_id,
+                                                        user_batch_object);
+            if (!query_result) {
+                LOG(ERROR) << "BatchQuerySlices failed with error: "
+                           << query_result.error();
+                continue;
+            }
+            for (size_t i = 0; i < storage_keys.size(); ++i) {
+                auto it = user_batch_object.find(user_keys[i]);
+                if (it != user_batch_object.end()) {
+                    batch_object.emplace(storage_keys[i],
+                                         std::move(it->second));
+                }
+            }
+        }
+        if (batch_object.empty()) {
             continue;
         }
 
         auto eviction_handler = [this](const std::vector<std::string>&
                                            evicted_keys) {
             if (evicted_keys.empty()) return;
-            auto results = client_->BatchEvictDiskReplica(
-                evicted_keys, ReplicaType::LOCAL_DISK);
-            for (size_t i = 0; i < results.size(); ++i) {
-                if (!results[i]) {
-                    LOG(WARNING)
-                        << "Failed to notify master about evicted local disk "
-                           "key: "
-                        << evicted_keys[i] << ", error: " << results[i].error();
+            std::unordered_map<std::string, std::vector<std::string>>
+                keys_by_tenant;
+            for (const auto& storage_key : evicted_keys) {
+                auto [tenant_id, key] =
+                    ParseTenantScopedStorageKey(storage_key);
+                keys_by_tenant[tenant_id].push_back(key);
+            }
+            for (const auto& [tenant_id, keys] : keys_by_tenant) {
+                auto results = client_->BatchEvictDiskReplica(
+                    keys, tenant_id, ReplicaType::LOCAL_DISK);
+                for (size_t i = 0; i < results.size(); ++i) {
+                    if (!results[i]) {
+                        LOG(WARNING)
+                            << "Failed to notify master about evicted local "
+                               "disk key: "
+                            << keys[i] << ", tenant_id=" << tenant_id
+                            << ", error: " << results[i].error();
+                    }
                 }
             }
         };
@@ -479,7 +584,28 @@ tl::expected<void, ErrorCode> FileStorage::Heartbeat() {
         LOG(ERROR) << "client is nullptr";
         return tl::make_unexpected(ErrorCode::INVALID_PARAMS);
     }
-    std::unordered_map<std::string, int64_t>
+
+    // Join previous rescan if completed
+    if (rescan_future_.valid() && rescan_future_.wait_for(std::chrono::seconds(
+                                      0)) == std::future_status::ready) {
+        rescan_future_ = std::future<void>();
+    }
+
+    // Retry metadata resync if previous attempt failed.
+    if (metadata_resync_pending_.load() && !rescan_future_.valid()) {
+        LOG(INFO) << "Retrying background metadata rescan";
+        rescan_future_ = std::async(std::launch::async, [this]() {
+            auto result = ReRegisterOffloadedObjects();
+            if (!result) {
+                LOG(ERROR) << "Background metadata rescan retry "
+                           << "failed: " << result.error();
+            } else {
+                metadata_resync_pending_.store(false);
+            }
+        });
+    }
+
+    std::vector<OffloadTaskItem>
         offloading_objects;  // Objects selected for offloading
 
     // === STEP 1: Send heartbeat and get offloading decisions ===
@@ -488,12 +614,58 @@ tl::expected<void, ErrorCode> FileStorage::Heartbeat() {
         auto heartbeat_result = client_->OffloadObjectHeartbeat(
             enable_offloading_, offloading_objects);
         if (!heartbeat_result) {
-            LOG(ERROR) << "Failed to send heartbeat with error: "
-                       << heartbeat_result.error();
-            return heartbeat_result;
+            ErrorCode err = heartbeat_result.error();
+            if (err == ErrorCode::SEGMENT_NOT_FOUND) {
+                // Master lost our LOCAL_DISK segment (likely restarted).
+                // Re-register the segment, retry the heartbeat, and
+                // trigger async ScanMeta to re-register object metadata.
+                LOG(WARNING) << "OffloadObjectHeartbeat returned "
+                             << "SEGMENT_NOT_FOUND, attempting to "
+                             << "re-register local disk segment and "
+                             << "re-register object metadata";
+                auto remount_result =
+                    client_->MountLocalDiskSegment(enable_offloading_);
+                if (remount_result) {
+                    heartbeat_result = client_->OffloadObjectHeartbeat(
+                        enable_offloading_, offloading_objects);
+                    if (!heartbeat_result) {
+                        LOG(ERROR) << "Heartbeat failed after re-registration: "
+                                   << heartbeat_result.error();
+                        return heartbeat_result;
+                    }
+                    // Master lost all object metadata on restart.
+                    // Trigger async ScanMeta to re-register them,
+                    // same as what Init() does on startup.
+                    if (!rescan_future_.valid()) {
+                        LOG(INFO) << "Triggering background metadata rescan "
+                                  << "after LOCAL_DISK segment re-registration";
+                        metadata_resync_pending_.store(true);
+                        rescan_future_ =
+                            std::async(std::launch::async, [this]() {
+                                auto result = ReRegisterOffloadedObjects();
+                                if (!result) {
+                                    LOG(ERROR) << "Background metadata rescan "
+                                               << "failed: " << result.error();
+                                } else {
+                                    metadata_resync_pending_.store(false);
+                                }
+                            });
+                    }
+                } else {
+                    LOG(ERROR) << "Failed to re-register local disk segment: "
+                               << remount_result.error();
+                    return tl::make_unexpected(remount_result.error());
+                }
+            } else {
+                LOG(ERROR) << "Failed to send heartbeat with error: " << err;
+                return heartbeat_result;
+            }
         }
     }
 
+    if (offloading_objects.empty()) {
+        return {};
+    }
     // === STEP 2: Persist offloaded objects (trigger actual data migration) ===
     auto offload_result = OffloadObjects(offloading_objects);
     if (!offload_result) {
@@ -502,11 +674,165 @@ tl::expected<void, ErrorCode> FileStorage::Heartbeat() {
         return offload_result;
     }
 
+    // Drive any pending L2->L1 promotion work for this client. Failures
+    // inside ProcessPromotionTasks are logged per-key and do not propagate;
+    // promotion is best-effort and must never break offload.
+    (void)ProcessPromotionTasks();
+
     // TODO(eviction): Implement an LRU eviction mechanism to manage local
     // storage capacity.
     return {};
 }
 
+tl::expected<void, ErrorCode> FileStorage::ProcessPromotionTasks() {
+    if (client_ == nullptr) {
+        return tl::make_unexpected(ErrorCode::INVALID_PARAMS);
+    }
+
+    std::vector<PromotionTaskItem> promotion_objects;
+    auto heartbeat_result =
+        client_->PromotionObjectHeartbeat(promotion_objects);
+    if (!heartbeat_result) {
+        // SEGMENT_NOT_FOUND happens between MountLocalDiskSegment and the
+        // first heartbeat tick if the master forgets us (e.g. across a master
+        // restart): benign no-op until next ReMount.
+        if (heartbeat_result.error() == ErrorCode::SEGMENT_NOT_FOUND) {
+            return {};
+        }
+        LOG(WARNING) << "PromotionObjectHeartbeat failed: "
+                     << heartbeat_result.error();
+        return tl::make_unexpected(heartbeat_result.error());
+    }
+    if (promotion_objects.empty()) {
+        return {};
+    }
+
+    VLOG(1) << "ProcessPromotionTasks pulled " << promotion_objects.size()
+            << " promotion candidate(s) from master";
+
+    // No segment preference from the client: let master pick from any
+    // DRAM segment.
+    const std::vector<std::string> preferred_segments;
+
+    // The master caps per-heartbeat work via PromotionObjectHeartbeat,
+    // returning at most one task per call so the heartbeat thread stays
+    // within the client-liveness window even for large objects. Leftover
+    // work stays queued in the master's promotion_objects map and is
+    // returned on subsequent heartbeats; we process whatever we received
+    // here without a second client-side cap.
+    for (const auto& task : promotion_objects) {
+        const auto& key = task.key;
+        const auto& tenant_id = task.tenant_id;
+        const int64_t size = task.size;
+        const auto storage_key = MakeTenantScopedStorageKey(tenant_id, key);
+        if (size <= 0) {
+            LOG(WARNING) << "Skipping promotion for key=" << key
+                         << " with non-positive size=" << size;
+            continue;
+        }
+
+        auto alloc_result = client_->PromotionAllocStart(
+            key, tenant_id, static_cast<uint64_t>(size), preferred_segments);
+        if (!alloc_result) {
+            // AllocStart failed (typically NO_AVAILABLE_HANDLE under
+            // DRAM pressure). No staged buffer to release, but the
+            // task entry already claimed a promotion_in_flight_ slot
+            // at admission. Notify the master to release it
+            // immediately; otherwise the slot stays pinned for the
+            // reaper TTL (~10 min default), turning transient DRAM
+            // pressure into a sustained outage of promotion_queue_limit_.
+            // Notify is idempotent and handles alloc_id == 0 correctly.
+            VLOG(1) << "PromotionAllocStart failed for key=" << key
+                    << ", error=" << alloc_result.error()
+                    << " (likely no free DRAM); releasing master slot";
+            auto release = client_->NotifyPromotionFailure(key, tenant_id);
+            if (!release) {
+                VLOG(1) << "Promotion: NotifyPromotionFailure failed for key="
+                        << key << ", error=" << release.error()
+                        << "; master reaper will reclaim on TTL expiry";
+            }
+            continue;
+        }
+
+        // Every failure path past this point has a master-side staged
+        // PROCESSING MEMORY buffer and an incremented in-flight slot.
+        // Eagerly notify the master on failure so the buffer is
+        // reclaimed and the slot is freed; otherwise transient SSD
+        // throttling or RDMA flakes saturate promotion_queue_limit_
+        // for the full reaper TTL. NotifyPromotionFailure is
+        // idempotent and best-effort — the reaper is the long-stop.
+        auto release_master_state = [this, &key, &tenant_id]() {
+            auto release = client_->NotifyPromotionFailure(key, tenant_id);
+            if (!release) {
+                VLOG(1) << "Promotion: NotifyPromotionFailure failed for key="
+                        << key << ", error=" << release.error()
+                        << "; master reaper will reclaim on TTL expiry";
+            }
+        };
+
+        // (a) Allocate an O_DIRECT-aligned staging buffer and read the bytes
+        // from the local SSD backend into it. AllocateBatch returns a
+        // shared_ptr<AllocatedBatch> whose BufferHandles RAII-release the
+        // staging space when the local goes out of scope.
+        std::vector<std::string> single_key{storage_key};
+        std::vector<int64_t> single_size{size};
+        auto allocate_res = AllocateBatch(single_key, single_size);
+        if (!allocate_res) {
+            LOG(WARNING) << "Promotion: AllocateBatch failed for key=" << key
+                         << ", error=" << allocate_res.error();
+            release_master_state();
+            continue;
+        }
+        auto staging = allocate_res.value();
+        auto load_res = BatchLoad(staging->slices);
+        if (!load_res) {
+            LOG(WARNING) << "Promotion: BatchLoad failed for key=" << key
+                         << ", error=" << load_res.error();
+            release_master_state();
+            continue;
+        }
+
+        // (b) TE-write from the staging slice into the freshly-allocated
+        // MEMORY replica. Slice ptr may have been bumped by O_DIRECT offset
+        // correction in BatchLoad, so re-read it from the slice map.
+        auto slice_it = staging->slices.find(storage_key);
+        if (slice_it == staging->slices.end()) {
+            LOG(WARNING) << "Promotion: staging slice missing for key=" << key;
+            release_master_state();
+            continue;
+        }
+        std::vector<Slice> tx_slices{slice_it->second};
+        ErrorCode write_err = client_->PromotionWrite(
+            alloc_result.value().memory_descriptor, tx_slices);
+        if (write_err != ErrorCode::OK) {
+            LOG(WARNING) << "Promotion: TransferWrite failed for key=" << key
+                         << ", error=" << write_err;
+            release_master_state();
+            continue;
+        }
+
+        // (c) Commit. Master flips the PROCESSING replica to COMPLETE and it
+        // becomes visible to readers.
+        auto notify_res = client_->NotifyPromotionSuccess(key, tenant_id);
+        if (!notify_res) {
+            // The write landed but the commit failed. We can't retry the
+            // commit (the success path is one-shot via alloc_id), and we
+            // don't know whether the failure was transient or structural.
+            // Release the master-side state so the slot is reusable; the
+            // bytes we wrote become stranded under a soon-to-be-erased
+            // PROCESSING replica, which is harmless.
+            LOG(WARNING) << "Promotion: NotifyPromotionSuccess failed for key="
+                         << key << ", error=" << notify_res.error();
+            release_master_state();
+            continue;
+        }
+
+        VLOG(1) << "Promotion completed for key=" << key << ", size=" << size;
+    }
+
+    return {};
+}
+
 tl::expected<void, ErrorCode> FileStorage::BatchLoad(
     std::unordered_map<std::string, Slice>& batch_object) {
     auto start_time = std::chrono::steady_clock::now();
@@ -537,9 +863,9 @@ tl::expected<void, ErrorCode> FileStorage::BatchLoad(
 }
 
 tl::expected<void, ErrorCode> FileStorage::BatchQuerySegmentSlices(
-    const std::vector<std::string>& keys,
+    const std::vector<std::string>& keys, const std::string& tenant_id,
     std::unordered_map<std::string, std::vector<Slice>>& batched_slices) {
-    auto batched_query_results = client_->BatchQuery(keys);
+    auto batched_query_results = client_->BatchQuery(keys, tenant_id);
     if (batched_query_results.empty())
         return tl::make_unexpected(ErrorCode::INVALID_REPLICA);
     for (size_t i = 0; i < batched_query_results.size(); ++i) {
@@ -573,7 +899,7 @@ tl::expected<void, ErrorCode> FileStorage::BatchQuerySegmentSlices(
 tl::expected<void, ErrorCode> FileStorage::RegisterLocalMemory() {
     auto error_code = client_->RegisterLocalMemory(
         client_buffer_allocator_->getBase(), config_.local_buffer_size,
-        kWildcardLocation, false, false);
+        kWildcardLocation, false, true);
     if (!error_code) {
         LOG(ERROR) << "Failed to register local memory: " << error_code.error();
         return error_code;
@@ -584,6 +910,11 @@ tl::expected<void, ErrorCode> FileStorage::RegisterLocalMemory() {
 tl::expected<std::shared_ptr<FileStorage::AllocatedBatch>, ErrorCode>
 FileStorage::AllocateBatch(const std::vector<std::string>& keys,
                            const std::vector<int64_t>& sizes) {
+    if (keys.size() != sizes.size()) {
+        LOG(ERROR) << "Mismatched keys and sizes count: keys=" << keys.size()
+                   << ", sizes=" << sizes.size();
+        return tl::make_unexpected(ErrorCode::INVALID_PARAMS);
+    }
     auto result = std::make_shared<AllocatedBatch>();
     result->batch_id = next_batch_id_.fetch_add(1, std::memory_order_relaxed);
     std::chrono::steady_clock::time_point now =
@@ -595,7 +926,11 @@ FileStorage::AllocateBatch(const std::vector<std::string>& keys,
     u_int64_t total_size = 0;
     bool gc_triggered = false;
     for (size_t i = 0; i < keys.size(); ++i) {
-        assert(sizes[i] <= kMaxSliceSize);
+        if (sizes[i] < 0 || sizes[i] > config_.local_buffer_size) {
+            LOG(ERROR) << "Invalid size for key " << keys[i] << ": " << sizes[i]
+                       << " (limit: " << config_.local_buffer_size << ")";
+            return tl::make_unexpected(ErrorCode::INVALID_PARAMS);
+        }
 
         // Allocate oversized buffer for O_DIRECT alignment:
         //   +4096 for aligning the ptr to 4096 boundary
@@ -688,4 +1023,59 @@ bool FileStorage::ReleaseBuffer(uint64_t batch_id) {
     return false;
 }
 
-}  // namespace mooncake
\ No newline at end of file
+tl::expected<void, ErrorCode> FileStorage::ReRegisterOffloadedObjects() {
+    LOG(INFO) << "ReRegisterOffloadedObjects: starting ScanMeta to re-register "
+              << "offloaded objects with master";
+    int total_keys = 0;
+    int total_batches = 0;
+    int total_failures = 0;
+    // Reset the scan iterator so ScanMeta starts from the beginning.
+    // BucketStorageBackend uses cursor-based iteration (next_bucket_);
+    // after Init() completes the cursor is 0 and HasNext() returns false,
+    // which would make ScanMeta skip all buckets.
+    storage_backend_->ResetScanIterator();
+    LOG(INFO) << "ReRegisterOffloadedObjects: about to call "
+                 "storage_backend_->ScanMeta()";
+    auto scan_meta_result =
+        storage_backend_->ScanMeta(
+            [this, &total_keys, &total_batches, &total_failures](
+                const std::vector<std::string>& keys,
+                std::vector<StorageObjectMetadata>& metadatas) {
+                total_batches++;
+                total_keys += keys.size();
+                for (auto& metadata : metadatas) {
+                    metadata.transport_endpoint = local_rpc_addr_;
+                }
+                auto tasks = BuildOffloadTasksFromStorageKeys(keys, metadatas);
+                auto add_object_result =
+                    client_->NotifyOffloadSuccess(tasks, metadatas);
+                if (!add_object_result) {
+                    total_failures++;
+                    LOG(ERROR)
+                        << "ReRegisterOffloadedObjects: NotifyOffloadSuccess "
+                        << "failed for batch " << total_batches << " with "
+                        << keys.size()
+                        << " keys, error: " << add_object_result.error();
+                    return add_object_result.error();
+                }
+                LOG(INFO) << "ReRegisterOffloadedObjects: NotifyOffloadSuccess "
+                          << "succeeded for batch " << total_batches << " with "
+                          << keys.size() << " keys";
+                return ErrorCode::OK;
+            });
+
+    LOG(INFO) << "ReRegisterOffloadedObjects: ScanMeta returned. success="
+              << scan_meta_result.has_value();
+    if (!scan_meta_result) {
+        LOG(ERROR) << "ReRegisterOffloadedObjects: ScanMeta failed: "
+                   << scan_meta_result.error();
+        return scan_meta_result;
+    }
+    LOG(INFO) << "ReRegisterOffloadedObjects: completed. "
+              << "total_keys=" << total_keys
+              << " total_batches=" << total_batches
+              << " total_failures=" << total_failures;
+    return {};
+}
+
+}  // namespace mooncake
diff --git a/mooncake-store/src/ha/leadership/backends/etcd/etcd_leader_coordinator.cpp b/mooncake-store/src/ha/leadership/backends/etcd/etcd_leader_coordinator.cpp
index 58c1e238..95b09a90 100644
--- a/mooncake-store/src/ha/leadership/backends/etcd/etcd_leader_coordinator.cpp
+++ b/mooncake-store/src/ha/leadership/backends/etcd/etcd_leader_coordinator.cpp
@@ -3,14 +3,18 @@
 #include <algorithm>
 #include <atomic>
 #include <chrono>
+#include <condition_variable>
 #include <cstdlib>
 #include <cstring>
+#include <mutex>
 #include <optional>
 #include <thread>
 
 #include <glog/logging.h>
 #include <ylt/util/tl/expected.hpp>
 
+#include "types.h"
+
 namespace mooncake {
 namespace ha {
 namespace backends {
@@ -19,7 +23,93 @@ namespace etcd {
 namespace {
 
 constexpr int kKeepAliveReadyTimeoutMs = 1000;
-constexpr auto kViewChangePollInterval = std::chrono::milliseconds(200);
+// Fallback poll interval used only when a watch cannot be armed (e.g. the watch
+// RPC fails, or the remaining timeout is too small to be worth a watch). The
+// steady-state path blocks on an etcd watch and does not poll.
+constexpr auto kViewChangeFallbackPollInterval = std::chrono::milliseconds(200);
+// Upper bound for waiting on the watch goroutine to fully stop before the
+// watch context is destroyed. Mirrors the value used by the oplog notifier.
+constexpr int kWatchStopTimeoutMs = 5000;
+
+// Shared state between WaitForViewChange and the etcd watch callback. The
+// callback only flips `changed` and wakes the waiter; the waiter decides what
+// the change means by re-reading the view.
+struct ViewChangeWatchState {
+    std::mutex mutex;
+    std::condition_variable cv;
+    bool changed = false;
+};
+
+// C-style trampoline invoked by the etcd watch goroutine for every event on the
+// watched prefix (PUT / DELETE / WATCH_BROKEN). Any event is treated as "the
+// view may have changed"; the waiter re-reads to find out what actually
+// happened. The watch context outlives every callback because the waiter calls
+// CancelWatchWithPrefix + WaitWatchWithPrefixStopped before it is destroyed.
+void ViewChangeWatchCallback(void* context, const char* /*key*/,
+                             size_t /*key_size*/, const char* /*value*/,
+                             size_t /*value_size*/, int /*event_type*/,
+                             int64_t /*mod_revision*/) {
+    auto* state = static_cast<ViewChangeWatchState*>(context);
+    if (state == nullptr) {
+        return;
+    }
+    {
+        std::lock_guard<std::mutex> lock(state->mutex);
+        state->changed = true;
+    }
+    state->cv.notify_all();
+}
+
+// RAII guard that tears the prefix watch down (and waits for the goroutine to
+// exit, so no further callbacks can run) before the heap-allocated
+// ViewChangeWatchState it owns is released.
+//
+// The watch state is HEAP-allocated (not stack-scoped) on purpose. The Go watch
+// goroutine holds a raw pointer to it and may still invoke the callback until
+// it has fully stopped. CancelWatchWithPrefix + WaitWatchWithPrefixStopped are
+// used to drain it, but WaitWatchWithPrefixStopped can time out. If it does, a
+// callback may still fire later, so freeing the state would be a
+// use-after-free. Mirroring the oplog notifier's policy, the guard therefore
+// only deletes the state when the stop wait SUCCEEDS; on timeout/failure it
+// intentionally LEAKS the state (and logs a warning) to avoid UAF.
+class PrefixWatchGuard {
+   public:
+    PrefixWatchGuard(std::string prefix, ViewChangeWatchState* state)
+        : prefix_(std::move(prefix)), state_(state) {}
+    ~PrefixWatchGuard() {
+        if (!active_) {
+            // Watch was never armed: no goroutine ever received the pointer, so
+            // it is safe to free the state unconditionally.
+            delete state_;
+            return;
+        }
+        EtcdHelper::CancelWatchWithPrefix(prefix_.c_str(), prefix_.size());
+        ErrorCode wait_err = EtcdHelper::WaitWatchWithPrefixStopped(
+            prefix_.c_str(), prefix_.size(), kWatchStopTimeoutMs);
+        if (wait_err == ErrorCode::OK) {
+            // Goroutine has fully exited; no callback can be in flight.
+            delete state_;
+        } else {
+            // The watch goroutine did not stop in time. It may still invoke the
+            // callback with `state_`, so we intentionally leak the state to
+            // avoid a use-after-free.
+            LOG(WARNING) << "Prefix watch goroutine did not stop in time for "
+                            "prefix "
+                         << prefix_
+                         << "; leaking ViewChangeWatchState to avoid UAF";
+        }
+    }
+
+    PrefixWatchGuard(const PrefixWatchGuard&) = delete;
+    PrefixWatchGuard& operator=(const PrefixWatchGuard&) = delete;
+
+    void Arm() { active_ = true; }
+
+   private:
+    std::string prefix_;
+    ViewChangeWatchState* state_ = nullptr;
+    bool active_ = false;
+};
 
 class EtcdLeadershipMonitorHandle final : public LeadershipMonitorHandle {
    public:
@@ -59,6 +149,16 @@ ErrorCode EtcdLeaderCoordinator::Connect() {
     return err;
 }
 
+ErrorCode EtcdLeaderCoordinator::ResetConnection() {
+    auto err = EtcdHelper::ResetEtcdStoreClient(spec_.connstring);
+    if (err != ErrorCode::OK) {
+        connected_ = false;
+        return err;
+    }
+    connected_ = true;
+    return ErrorCode::OK;
+}
+
 tl::expected<std::optional<MasterView>, ErrorCode>
 EtcdLeaderCoordinator::ReadCurrentView() {
     auto err = EnsureConnected();
@@ -92,6 +192,7 @@ EtcdLeaderCoordinator::TryAcquireLeadership(const std::string& leader_address) {
     EtcdLeaseId lease_id = 0;
     err = EtcdHelper::GrantLease(DEFAULT_MASTER_VIEW_LEASE_TTL_SEC, lease_id);
     if (err != ErrorCode::OK) {
+        (void)ResetConnection();
         return tl::make_unexpected(err);
     }
 
@@ -115,9 +216,13 @@ EtcdLeaderCoordinator::TryAcquireLeadership(const std::string& leader_address) {
         };
     }
     if (err != ErrorCode::OK) {
+        (void)ResetConnection();  // Reset first to get off the stale client
+
         auto revoke_err = EtcdHelper::RevokeLease(lease_id);
         if (revoke_err != ErrorCode::OK) {
-            return tl::make_unexpected(revoke_err);
+            LOG(WARNING)
+                << "Failed to revoke lease after CreateWithLease failure: "
+                << static_cast<int>(revoke_err);
         }
         return tl::make_unexpected(err);
     }
@@ -192,6 +297,9 @@ tl::expected<bool, ErrorCode> EtcdLeaderCoordinator::RenewLeadership(
             started_keepalive = true;
             keepalive_thread_ = std::thread([this, lease = lease_id.value()]() {
                 auto rc = EtcdHelper::KeepAlive(lease);
+                if (rc == ErrorCode::ETCD_OPERATION_ERROR) {
+                    (void)EtcdHelper::ResetEtcdStoreClient(spec_.connstring);
+                }
                 std::shared_ptr<std::atomic<bool>> monitor_armed;
                 LeadershipLostCallback on_leadership_lost;
                 LeadershipLossReason loss_reason =
@@ -245,21 +353,17 @@ EtcdLeaderCoordinator::WaitForViewChange(
         return tl::make_unexpected(err);
     }
 
+    // LIMITATION: prefix watches are keyed globally by prefix in the etcd Go
+    // wrapper (storePrefixWatchCtx is a map<prefix, ...>). Because every waiter
+    // here watches the same `master_view_key_` prefix, only a SINGLE waiter per
+    // process per prefix is supported: a second concurrent WaitForViewChange on
+    // the same prefix would collide on that registration. The intended
+    // deployment model is one waiter per process/prefix, which holds for the
+    // store client. Revisit (e.g. per-waiter watch IDs) if multiple
+    // same-process clients ever need to watch the same master_view
+    // concurrently.
     const auto deadline = std::chrono::steady_clock::now() + timeout;
     while (true) {
-        auto current_view = ReadCurrentView();
-        if (!current_view) {
-            return tl::make_unexpected(current_view.error());
-        }
-
-        if (!IsSameViewVersion(current_view.value(), known_version)) {
-            return ViewChangeResult{
-                .changed = true,
-                .timed_out = false,
-                .current_view = current_view.value(),
-            };
-        }
-
         if (timeout <= std::chrono::milliseconds::zero() ||
             std::chrono::steady_clock::now() >= deadline) {
             return ViewChangeResult{
@@ -272,8 +376,69 @@ EtcdLeaderCoordinator::WaitForViewChange(
         const auto remaining =
             std::chrono::duration_cast<std::chrono::milliseconds>(
                 deadline - std::chrono::steady_clock::now());
-        std::this_thread::sleep_for(
-            std::min(kViewChangePollInterval, remaining));
+
+        // Arm the watch BEFORE reading the view. The watch is established at
+        // some etcd revision R_watch; the subsequent read observes a revision
+        // R_read >= R_watch. Any change to master_view at revision C is then
+        // caught by exactly one of the two: C < R_read (the read sees it) or
+        // C >= R_watch (the watch delivers it), and the two ranges overlap, so
+        // a change happening between read and watch cannot be missed. The watch
+        // starts from the current revision (start_revision = 0), which also
+        // avoids depending on a possibly-compacted historical revision.
+        //
+        // The watch state is heap-allocated and owned by `guard`. The guard
+        // cancels the watch and waits for the goroutine to exit before
+        // releasing the state in its destructor; if the goroutine fails to stop
+        // in time it leaks the state instead of freeing it, so an in-flight
+        // callback can never reference freed memory (see PrefixWatchGuard).
+        auto* state = new ViewChangeWatchState();
+        PrefixWatchGuard guard(master_view_key_, state);
+
+        bool watching = false;
+        if (remaining > kViewChangeFallbackPollInterval) {
+            // Defensively clear any lingering watch on this key, then arm a new
+            // one. Both calls are no-ops when nothing is registered.
+            EtcdHelper::CancelWatchWithPrefix(master_view_key_.c_str(),
+                                              master_view_key_.size());
+            EtcdHelper::WaitWatchWithPrefixStopped(master_view_key_.c_str(),
+                                                   master_view_key_.size(),
+                                                   kWatchStopTimeoutMs);
+            auto watch_err = EtcdHelper::WatchWithPrefixFromRevision(
+                master_view_key_.c_str(), master_view_key_.size(),
+                /*start_revision=*/0, state, &ViewChangeWatchCallback);
+            if (watch_err == ErrorCode::OK) {
+                watching = true;
+                guard.Arm();
+            }
+        }
+
+        auto current_view = ReadCurrentView();
+        if (!current_view) {
+            return tl::make_unexpected(current_view.error());
+        }
+        if (!IsSameViewVersion(current_view.value(), known_version)) {
+            return ViewChangeResult{
+                .changed = true,
+                .timed_out = false,
+                .current_view = current_view.value(),
+            };
+        }
+
+        if (!watching) {
+            // Could not arm a watch (RPC failed, or too little time left).
+            // Fall back to a short poll so a change is still picked up.
+            std::this_thread::sleep_for(
+                std::min(kViewChangeFallbackPollInterval, remaining));
+            continue;
+        }
+
+        // Block until the watch reports an event or the caller's deadline
+        // elapses. Either way we loop and re-read: an event tells us the view
+        // changed (re-read returns it), a timeout falls through to the deadline
+        // check above and returns timed_out.
+        std::unique_lock<std::mutex> lock(state->mutex);
+        state->cv.wait_for(lock, remaining,
+                           [state]() { return state->changed; });
     }
 }
 
@@ -369,7 +534,7 @@ ClusterNamespace EtcdLeaderCoordinator::ResolveClusterNamespace(
     if (env_cluster_id != nullptr && std::strlen(env_cluster_id) > 0) {
         resolved_namespace = env_cluster_id;
     } else {
-        resolved_namespace = "mooncake";
+        resolved_namespace = DEFAULT_CLUSTER_ID;
     }
     return resolved_namespace;
 }
diff --git a/mooncake-store/src/ha/leadership/backends/k8s/k8s_leader_coordinator.cpp b/mooncake-store/src/ha/leadership/backends/k8s/k8s_leader_coordinator.cpp
new file mode 100644
index 00000000..ea24b1a4
--- /dev/null
+++ b/mooncake-store/src/ha/leadership/backends/k8s/k8s_leader_coordinator.cpp
@@ -0,0 +1,456 @@
+#include "ha/leadership/backends/k8s/k8s_leader_coordinator.h"
+
+#include <algorithm>
+#include <atomic>
+#include <chrono>
+#include <cstdint>
+#include <optional>
+#include <string>
+#include <thread>
+
+#include <glog/logging.h>
+#include <ylt/util/tl/expected.hpp>
+
+namespace mooncake {
+namespace ha {
+namespace backends {
+namespace k8s {
+
+namespace {
+
+constexpr int kDefaultLeaseDurationSec = 5;
+constexpr int kDefaultRenewDeadlineSec = 3;
+constexpr int kDefaultRetryPeriodSec = 1;
+constexpr auto kViewChangePollInterval = std::chrono::milliseconds(200);
+
+std::atomic<uint64_t> g_session_sequence{0};
+
+OwnerToken MakeOwnerToken(const std::string& namespace_name,
+                          const std::string& lease_name,
+                          const std::string& leader_address,
+                          int64_t lease_transitions) {
+    const auto sequence =
+        g_session_sequence.fetch_add(1, std::memory_order_relaxed) + 1;
+    return namespace_name + "/" + lease_name + "/" +
+           std::to_string(lease_transitions) + "/" + std::to_string(sequence) +
+           "/" + leader_address;
+}
+
+class K8sLeadershipMonitorHandle final : public LeadershipMonitorHandle {
+   public:
+    explicit K8sLeadershipMonitorHandle(
+        std::shared_ptr<std::atomic<bool>> armed)
+        : armed_(std::move(armed)) {}
+
+    void Stop() override {
+        if (armed_ != nullptr) {
+            armed_->store(false);
+        }
+    }
+
+   private:
+    std::shared_ptr<std::atomic<bool>> armed_;
+};
+
+}  // namespace
+
+K8sLeaderCoordinator::K8sLeaderCoordinator(const HABackendSpec& spec)
+    : spec_(spec) {}
+
+K8sLeaderCoordinator::~K8sLeaderCoordinator() { ShutdownElection(); }
+
+ErrorCode K8sLeaderCoordinator::Connect() {
+    if (connected_) {
+        return ErrorCode::OK;
+    }
+
+    auto parsed = ParseConnstring(spec_.connstring);
+    if (!parsed) {
+        return parsed.error();
+    }
+    namespace_ = std::move(parsed->first);
+    lease_name_ = std::move(parsed->second);
+
+    auto err = K8sLeaseHelper::Init();
+    if (err == ErrorCode::OK) {
+        connected_ = true;
+    }
+    return err;
+}
+
+tl::expected<std::optional<MasterView>, ErrorCode>
+K8sLeaderCoordinator::ReadCurrentView() {
+    auto err = EnsureConnected();
+    if (err != ErrorCode::OK) {
+        return tl::make_unexpected(err);
+    }
+
+    std::string holder;
+    int64_t transitions = 0;
+    err =
+        K8sLeaseHelper::GetHolder(namespace_, lease_name_, holder, transitions);
+    if (err == ErrorCode::K8S_LEASE_NOT_FOUND) {
+        return std::optional<MasterView>{std::nullopt};
+    }
+    if (err != ErrorCode::OK) {
+        return tl::make_unexpected(err);
+    }
+
+    if (holder.empty()) {
+        return std::optional<MasterView>{std::nullopt};
+    }
+
+    return std::optional<MasterView>{
+        MasterView{.leader_address = std::move(holder),
+                   .view_version = static_cast<ViewVersionId>(transitions)}};
+}
+
+tl::expected<AcquireLeadershipResult, ErrorCode>
+K8sLeaderCoordinator::TryAcquireLeadership(const std::string& leader_address) {
+    auto err = EnsureConnected();
+    if (err != ErrorCode::OK) {
+        return tl::make_unexpected(err);
+    }
+
+    // Start election goroutine
+    err = K8sLeaseHelper::RunElection(
+        namespace_, lease_name_, leader_address, kDefaultLeaseDurationSec,
+        kDefaultRenewDeadlineSec, kDefaultRetryPeriodSec);
+    if (err != ErrorCode::OK) {
+        return tl::make_unexpected(err);
+    }
+
+    constexpr int kElectionTimeoutSec = 2 * kDefaultLeaseDurationSec;
+    int64_t transitions = 0;
+    err = K8sLeaseHelper::WaitElected(namespace_, lease_name_,
+                                      kElectionTimeoutSec, transitions);
+    if (err != ErrorCode::OK) {
+        // Election failed — we are not the leader
+        K8sLeaseHelper::CancelElection(namespace_, lease_name_);
+
+        auto observed_view = ReadCurrentView();
+        if (!observed_view) {
+            return tl::make_unexpected(observed_view.error());
+        }
+        return AcquireLeadershipResult{
+            .status = AcquireLeadershipStatus::CONTENDED,
+            .session = std::nullopt,
+            .observed_view = observed_view.value(),
+        };
+    }
+
+    // We are the leader
+    OwnerToken token =
+        MakeOwnerToken(namespace_, lease_name_, leader_address, transitions);
+    LeadershipSession session{
+        .view =
+            MasterView{.leader_address = leader_address,
+                       .view_version = static_cast<ViewVersionId>(transitions)},
+        .owner_token = token,
+        .lease_ttl = std::chrono::seconds(kDefaultLeaseDurationSec),
+    };
+
+    std::thread thread_to_join;
+    {
+        std::lock_guard<std::mutex> lock(election_mutex_);
+        if (election_monitor_thread_.joinable() && election_monitor_stopped_) {
+            thread_to_join = std::move(election_monitor_thread_);
+        }
+        election_identity_ = leader_address;
+        election_owner_token_ = session.owner_token;
+        election_active_ = true;
+        election_monitor_stopped_ = false;
+        election_shutdown_requested_ = false;
+        ClearLeadershipMonitorStateLocked();
+    }
+
+    if (thread_to_join.joinable()) {
+        thread_to_join.join();
+    }
+
+    return AcquireLeadershipResult{
+        .status = AcquireLeadershipStatus::ACQUIRED,
+        .session = std::move(session),
+        .observed_view = std::nullopt,
+    };
+}
+
+tl::expected<bool, ErrorCode> K8sLeaderCoordinator::RenewLeadership(
+    const LeadershipSession& session) {
+    auto err = EnsureConnected();
+    if (err != ErrorCode::OK) {
+        return tl::make_unexpected(err);
+    }
+    if (session.owner_token.empty()) {
+        return tl::make_unexpected(ErrorCode::INVALID_PARAMS);
+    }
+
+    std::thread thread_to_join;
+    {
+        std::lock_guard<std::mutex> lock(election_mutex_);
+
+        if (election_shutdown_requested_) {
+            return false;
+        }
+        if (!election_active_) {
+            return false;
+        }
+        if (election_owner_token_ != session.owner_token) {
+            return false;
+        }
+
+        // client-go handles renewal internally. If election is still active,
+        // leadership is being renewed. Start the monitor thread if not running.
+        if (election_monitor_thread_.joinable()) {
+            if (!election_monitor_stopped_) {
+                return true;
+            }
+            thread_to_join = std::move(election_monitor_thread_);
+            election_monitor_stopped_ = false;
+        }
+
+        if (!election_monitor_thread_.joinable()) {
+            election_monitor_thread_ =
+                std::thread([this, token = session.owner_token]() {
+                    auto rc = K8sLeaseHelper::WaitLost(namespace_, lease_name_);
+
+                    std::shared_ptr<std::atomic<bool>> monitor_armed;
+                    LeadershipLostCallback on_leadership_lost;
+                    LeadershipLossReason loss_reason =
+                        (rc == ErrorCode::OK)
+                            ? LeadershipLossReason::kLostLeadership
+                            : LeadershipLossReason::kRenewError;
+                    {
+                        std::lock_guard<std::mutex> lock(election_mutex_);
+                        election_monitor_stopped_ = true;
+                        if (election_owner_token_ == token) {
+                            election_active_ = false;
+                        }
+                        if (!election_shutdown_requested_ &&
+                            leadership_monitor_owner_token_ == token) {
+                            monitor_armed = leadership_monitor_armed_;
+                            on_leadership_lost =
+                                std::move(leadership_monitor_callback_);
+                            leadership_monitor_armed_.reset();
+                            leadership_monitor_owner_token_.clear();
+                        }
+                    }
+
+                    if (monitor_armed != nullptr &&
+                        monitor_armed->exchange(false) &&
+                        on_leadership_lost != nullptr) {
+                        on_leadership_lost(loss_reason);
+                    }
+                });
+        }
+    }
+
+    if (thread_to_join.joinable()) {
+        thread_to_join.join();
+    }
+
+    return true;
+}
+
+tl::expected<ViewChangeResult, ErrorCode>
+K8sLeaderCoordinator::WaitForViewChange(
+    std::optional<ViewVersionId> known_version,
+    std::chrono::milliseconds timeout) {
+    auto err = EnsureConnected();
+    if (err != ErrorCode::OK) {
+        return tl::make_unexpected(err);
+    }
+
+    const auto deadline = std::chrono::steady_clock::now() + timeout;
+    while (true) {
+        auto current_view = ReadCurrentView();
+        if (!current_view) {
+            return tl::make_unexpected(current_view.error());
+        }
+
+        if (!IsSameViewVersion(current_view.value(), known_version)) {
+            return ViewChangeResult{
+                .changed = true,
+                .timed_out = false,
+                .current_view = current_view.value(),
+            };
+        }
+
+        if (timeout <= std::chrono::milliseconds::zero() ||
+            std::chrono::steady_clock::now() >= deadline) {
+            return ViewChangeResult{
+                .changed = false,
+                .timed_out = true,
+                .current_view = std::nullopt,
+            };
+        }
+
+        const auto remaining =
+            std::chrono::duration_cast<std::chrono::milliseconds>(
+                deadline - std::chrono::steady_clock::now());
+        std::this_thread::sleep_for(
+            std::min(kViewChangePollInterval, remaining));
+    }
+}
+
+tl::expected<std::unique_ptr<LeadershipMonitorHandle>, ErrorCode>
+K8sLeaderCoordinator::StartLeadershipMonitor(
+    const LeadershipSession& session,
+    LeadershipLostCallback on_leadership_lost) {
+    auto err = EnsureConnected();
+    if (err != ErrorCode::OK) {
+        return tl::make_unexpected(err);
+    }
+    if (!on_leadership_lost) {
+        return tl::make_unexpected(ErrorCode::INVALID_PARAMS);
+    }
+
+    std::lock_guard<std::mutex> lock(election_mutex_);
+    if (election_shutdown_requested_) {
+        return tl::make_unexpected(ErrorCode::UNAVAILABLE_IN_CURRENT_STATUS);
+    }
+    if (session.owner_token.empty()) {
+        return tl::make_unexpected(ErrorCode::INVALID_PARAMS);
+    }
+    if (!election_active_) {
+        return tl::make_unexpected(ErrorCode::UNAVAILABLE_IN_CURRENT_STATUS);
+    }
+    if (election_owner_token_ != session.owner_token) {
+        return tl::make_unexpected(ErrorCode::UNAVAILABLE_IN_CURRENT_STATUS);
+    }
+    if (leadership_monitor_armed_ != nullptr &&
+        leadership_monitor_armed_->load()) {
+        return tl::make_unexpected(ErrorCode::INVALID_PARAMS);
+    }
+
+    leadership_monitor_owner_token_ = session.owner_token;
+    leadership_monitor_callback_ = std::move(on_leadership_lost);
+    leadership_monitor_armed_ = std::make_shared<std::atomic<bool>>(true);
+    return std::unique_ptr<LeadershipMonitorHandle>(
+        std::make_unique<K8sLeadershipMonitorHandle>(
+            leadership_monitor_armed_));
+}
+
+ErrorCode K8sLeaderCoordinator::ReleaseLeadership(
+    const LeadershipSession& session) {
+    if (session.owner_token.empty()) {
+        return ErrorCode::INVALID_PARAMS;
+    }
+
+    std::thread thread_to_join;
+    {
+        std::lock_guard<std::mutex> lock(election_mutex_);
+        if (!election_owner_token_.empty() &&
+            election_owner_token_ != session.owner_token) {
+            return ErrorCode::INVALID_PARAMS;
+        }
+        election_shutdown_requested_ = true;
+        election_active_ = false;
+        election_monitor_stopped_ = true;
+        election_owner_token_.clear();
+        ClearLeadershipMonitorStateLocked();
+        if (election_monitor_thread_.joinable()) {
+            thread_to_join = std::move(election_monitor_thread_);
+        }
+    }
+
+    // Cancel the election goroutine (triggers WaitLost to return)
+    auto err = K8sLeaseHelper::CancelElection(namespace_, lease_name_);
+
+    if (thread_to_join.joinable()) {
+        thread_to_join.join();
+    }
+
+    return err;
+}
+
+ErrorCode K8sLeaderCoordinator::EnsureConnected() {
+    if (connected_) {
+        return ErrorCode::OK;
+    }
+    return Connect();
+}
+
+ErrorCode K8sLeaderCoordinator::ShutdownElection() {
+    std::thread thread_to_join;
+    bool should_cancel = false;
+    {
+        std::lock_guard<std::mutex> lock(election_mutex_);
+        election_shutdown_requested_ = true;
+        should_cancel = election_active_;
+        election_active_ = false;
+        election_monitor_stopped_ = true;
+        election_owner_token_.clear();
+        ClearLeadershipMonitorStateLocked();
+        if (election_monitor_thread_.joinable()) {
+            thread_to_join = std::move(election_monitor_thread_);
+        }
+    }
+
+    if (should_cancel) {
+        K8sLeaseHelper::CancelElection(namespace_, lease_name_);
+    }
+
+    if (thread_to_join.joinable()) {
+        thread_to_join.join();
+    }
+
+    return ErrorCode::OK;
+}
+
+void K8sLeaderCoordinator::ClearLeadershipMonitorStateLocked() {
+    if (leadership_monitor_armed_ != nullptr) {
+        leadership_monitor_armed_->store(false);
+        leadership_monitor_armed_.reset();
+    }
+    leadership_monitor_callback_ = nullptr;
+    leadership_monitor_owner_token_.clear();
+}
+
+bool K8sLeaderCoordinator::IsSameViewVersion(
+    const std::optional<MasterView>& current_view,
+    std::optional<ViewVersionId> known_version) const {
+    if (!current_view.has_value() && !known_version.has_value()) {
+        return true;
+    }
+    if (!current_view.has_value() || !known_version.has_value()) {
+        return false;
+    }
+    return current_view->view_version == known_version.value();
+}
+
+tl::expected<std::pair<std::string, std::string>, ErrorCode>
+K8sLeaderCoordinator::ParseConnstring(const std::string& connstring) {
+    // Format: "namespace/lease-name" or just "lease-name" (default namespace).
+    // Exactly zero or one '/' is allowed.
+    if (connstring.empty()) {
+        LOG(ERROR) << "K8s HA connstring is empty";
+        return tl::make_unexpected(ErrorCode::INVALID_PARAMS);
+    }
+
+    auto first_slash = connstring.find('/');
+    if (first_slash == std::string::npos) {
+        // No slash — use default namespace
+        return std::pair<std::string, std::string>{"default", connstring};
+    }
+
+    // Reject multiple slashes (e.g. "a/b/c")
+    if (connstring.find('/', first_slash + 1) != std::string::npos) {
+        LOG(ERROR) << "K8s HA connstring contains multiple '/': " << connstring;
+        return tl::make_unexpected(ErrorCode::INVALID_PARAMS);
+    }
+
+    auto ns = connstring.substr(0, first_slash);
+    auto name = connstring.substr(first_slash + 1);
+    if (ns.empty() || name.empty()) {
+        LOG(ERROR) << "K8s HA connstring has empty namespace or lease name: "
+                   << connstring;
+        return tl::make_unexpected(ErrorCode::INVALID_PARAMS);
+    }
+    return std::pair<std::string, std::string>{std::move(ns), std::move(name)};
+}
+
+}  // namespace k8s
+}  // namespace backends
+}  // namespace ha
+}  // namespace mooncake
diff --git a/mooncake-store/src/ha/leadership/backends/redis/redis_leader_coordinator.cpp b/mooncake-store/src/ha/leadership/backends/redis/redis_leader_coordinator.cpp
index 57a5ce3a..e8539c01 100644
--- a/mooncake-store/src/ha/leadership/backends/redis/redis_leader_coordinator.cpp
+++ b/mooncake-store/src/ha/leadership/backends/redis/redis_leader_coordinator.cpp
@@ -11,6 +11,8 @@
 #include <thread>
 
 #include <glog/logging.h>
+
+#include "types.h"
 #ifdef STORE_USE_REDIS
 #include <hiredis/hiredis.h>
 #endif
@@ -665,7 +667,7 @@ ClusterNamespace RedisLeaderCoordinator::ResolveClusterNamespace(
     if (env_cluster_id != nullptr && std::strlen(env_cluster_id) > 0) {
         resolved_namespace = env_cluster_id;
     } else {
-        resolved_namespace = "mooncake";
+        resolved_namespace = DEFAULT_CLUSTER_ID;
     }
     return resolved_namespace;
 }
diff --git a/mooncake-store/src/ha/leadership/leader_coordinator_factory.cpp b/mooncake-store/src/ha/leadership/leader_coordinator_factory.cpp
index e6bc49af..c8458691 100644
--- a/mooncake-store/src/ha/leadership/leader_coordinator_factory.cpp
+++ b/mooncake-store/src/ha/leadership/leader_coordinator_factory.cpp
@@ -2,6 +2,9 @@
 
 #include "ha/leadership/backends/etcd/etcd_leader_coordinator.h"
 #include "ha/leadership/backends/redis/redis_leader_coordinator.h"
+#ifdef STORE_USE_K8S_LEASE
+#include "ha/leadership/backends/k8s/k8s_leader_coordinator.h"
+#endif
 
 namespace mooncake {
 namespace ha {
@@ -29,8 +32,19 @@ CreateLeaderCoordinator(const HABackendSpec& spec) {
             }
             return std::unique_ptr<LeaderCoordinator>(std::move(coordinator));
         }
-        case HABackendType::K8S:
+        case HABackendType::K8S: {
+#ifdef STORE_USE_K8S_LEASE
+            auto coordinator =
+                std::make_unique<backends::k8s::K8sLeaderCoordinator>(spec);
+            auto err = coordinator->Connect();
+            if (err != ErrorCode::OK) {
+                return tl::make_unexpected(err);
+            }
+            return std::unique_ptr<LeaderCoordinator>(std::move(coordinator));
+#else
             return tl::make_unexpected(ErrorCode::UNAVAILABLE_IN_CURRENT_MODE);
+#endif
+        }
     }
 
     return tl::make_unexpected(ErrorCode::INVALID_PARAMS);
diff --git a/mooncake-store/src/ha/leadership/master_service_supervisor.cpp b/mooncake-store/src/ha/leadership/master_service_supervisor.cpp
index d38a6c31..4e7093b9 100644
--- a/mooncake-store/src/ha/leadership/master_service_supervisor.cpp
+++ b/mooncake-store/src/ha/leadership/master_service_supervisor.cpp
@@ -5,6 +5,7 @@
 #include <csignal>
 #include <cstdlib>
 #include <memory>
+#include <optional>
 #include <string_view>
 #include <thread>
 
@@ -12,7 +13,10 @@
 #include <ylt/coro_rpc/coro_rpc_server.hpp>
 
 #include "ha/leadership/leader_coordinator_factory.h"
+#include "ha/leadership/leader_label_reconciler.h"
 #include "ha/standby_controller.h"
+#include "k8s_lease_helper.h"
+#include "master_admin_service.h"
 #include "rpc_service.h"
 
 namespace mooncake {
@@ -23,13 +27,33 @@ namespace {
 constexpr auto kAcquireRetryInterval = std::chrono::seconds(1);
 constexpr auto kRenewCheckInterval = std::chrono::seconds(1);
 constexpr auto kSupervisorRetryInterval = std::chrono::seconds(1);
+constexpr auto kLabelReconcileRetryInterval = std::chrono::seconds(1);
+constexpr char kLeaderLabelKey[] = "mooncake.io/store-role";
+constexpr char kLeaderLabelValue[] = "leader";
+
+bool HasPodIdentity(const MasterServiceSupervisorConfig& config) {
+    return !config.pod_name.empty() && !config.pod_namespace.empty() &&
+           config.ha_backend_type == "k8s";
+}
+
+LeaderLabelReconciler MakeLeaderLabelReconciler(
+    const MasterServiceSupervisorConfig& config) {
+    return LeaderLabelReconciler(
+        HasPodIdentity(config),
+        [ns = config.pod_namespace, pod = config.pod_name](bool desired) {
+            return desired ? K8sLeaseHelper::SetPodLabel(
+                                 ns, pod, kLeaderLabelKey, kLeaderLabelValue)
+                           : K8sLeaseHelper::ClearPodLabel(ns, pod,
+                                                           kLeaderLabelKey);
+        },
+        kLabelReconcileRetryInterval);
+}
 
 std::string ResolveHABackendConnstring(
     const MasterServiceSupervisorConfig& config) {
-    if (!config.ha_backend_connstring.empty()) {
-        return config.ha_backend_connstring;
-    }
-    return config.etcd_endpoints;
+    return ResolveConfiguredHABackendConnstring(config.ha_backend_type,
+                                                config.ha_backend_connstring,
+                                                config.etcd_endpoints);
 }
 
 tl::expected<HABackendSpec, ErrorCode> BuildHABackendSpec(
@@ -39,6 +63,11 @@ tl::expected<HABackendSpec, ErrorCode> BuildHABackendSpec(
         return tl::make_unexpected(ErrorCode::INVALID_PARAMS);
     }
 
+    auto availability = ValidateHABackendAvailability(backend_type.value());
+    if (availability != ErrorCode::OK) {
+        return tl::make_unexpected(availability);
+    }
+
     auto connstring = ResolveHABackendConnstring(config);
     if (connstring.empty()) {
         return tl::make_unexpected(ErrorCode::INVALID_PARAMS);
@@ -126,17 +155,20 @@ void SetRuntimeState(MasterAdminServer& admin_server,
               << ", role=" << MasterRuntimeRoleToString(state);
 }
 
-void ActivateServingState(
-    MasterAdminServer& admin_server,
-    const std::shared_ptr<WrappedMasterService>& service) {
+void ActivateServingState(MasterAdminServer& admin_server,
+                          const std::shared_ptr<WrappedMasterService>& service,
+                          LeaderLabelReconciler& label_reconciler) {
     admin_server.SetServiceDelegate(service);
     admin_server.SetServiceAvailable(true);
     SetRuntimeState(admin_server, MasterRuntimeState::kServing);
+    label_reconciler.SetLeader(true);
 }
 
-void DeactivateServingState(MasterAdminServer& admin_server) {
+void DeactivateServingState(MasterAdminServer& admin_server,
+                            LeaderLabelReconciler& label_reconciler) {
     admin_server.SetServiceAvailable(false);
     admin_server.SetServiceDelegate(nullptr);
+    label_reconciler.SetLeader(false);
 }
 
 void StopLeadershipMonitor(std::unique_ptr<LeadershipMonitorHandle>& monitor) {
@@ -186,6 +218,8 @@ void EnterStandbyMode(MasterAdminServer& admin_server,
 int RunSupervisorLoop(const HABackendSpec& spec,
                       const MasterServiceSupervisorConfig& config,
                       MasterAdminServer& admin_server) {
+    auto label_reconciler = MakeLeaderLabelReconciler(config);
+    label_reconciler.SetLeader(false);
     SetRuntimeState(admin_server, MasterRuntimeState::kStarting);
     auto standby_controller = CreateStandbyController(spec, config);
     std::atomic<bool> accept_standby_runtime_updates{false};
@@ -345,15 +379,18 @@ int RunSupervisorLoop(const HABackendSpec& spec,
             server.init_ibv();
         }
 
+        // The serving primary handles heartbeats/unmounts, so forward the
+        // metadata cleanup config here like the non-HA path does.
         auto wrapped_master_service = std::make_shared<WrappedMasterService>(
             mooncake::WrappedMasterServiceConfig(
-                config, leadership_session->view.view_version));
+                config, leadership_session->view.view_version),
+            config.http_metadata_server, config.http_metadata_remote_url);
         mooncake::RegisterRpcService(server, *wrapped_master_service);
 
         auto serve_preflight =
             leader_coordinator.RenewLeadership(*leadership_session);
         if (!serve_preflight) {
-            DeactivateServingState(admin_server);
+            DeactivateServingState(admin_server, label_reconciler);
             EnterStandbyMode(admin_server, *standby_controller,
                              accept_standby_runtime_updates,
                              leadership_session->view);
@@ -366,7 +403,7 @@ int RunSupervisorLoop(const HABackendSpec& spec,
             continue;
         }
         if (!serve_preflight.value()) {
-            DeactivateServingState(admin_server);
+            DeactivateServingState(admin_server, label_reconciler);
             EnterStandbyMode(admin_server, *standby_controller,
                              accept_standby_runtime_updates, std::nullopt);
             LogLeadershipReleaseWarning(
@@ -380,16 +417,18 @@ int RunSupervisorLoop(const HABackendSpec& spec,
         std::atomic<bool> serve_shutdown_requested{false};
         auto leadership_monitor = leader_coordinator.StartLeadershipMonitor(
             *leadership_session,
-            [&server, &admin_server, &serve_shutdown_requested](auto reason) {
+            [&server, &admin_server, &serve_shutdown_requested,
+             &label_reconciler](auto reason) {
                 serve_shutdown_requested.store(true, std::memory_order_release);
                 admin_server.SetServiceAvailable(false);
+                label_reconciler.SetLeader(false);
                 SetRuntimeState(admin_server, MasterRuntimeState::kStandby);
                 LOG(INFO) << "Trying to stop server, reason="
                           << LeadershipLossReasonToString(reason);
                 server.stop();
             });
         if (!leadership_monitor) {
-            DeactivateServingState(admin_server);
+            DeactivateServingState(admin_server, label_reconciler);
             EnterStandbyMode(admin_server, *standby_controller,
                              accept_standby_runtime_updates,
                              leadership_session->view);
@@ -408,7 +447,7 @@ int RunSupervisorLoop(const HABackendSpec& spec,
             LOG(ERROR) << "Failed to start master service: "
                        << ec.result().value();
             StopLeadershipMonitor(leadership_monitor_handle);
-            DeactivateServingState(admin_server);
+            DeactivateServingState(admin_server, label_reconciler);
             EnterStandbyMode(admin_server, *standby_controller,
                              accept_standby_runtime_updates,
                              leadership_session->view);
@@ -421,14 +460,15 @@ int RunSupervisorLoop(const HABackendSpec& spec,
         }
 
         if (!serve_shutdown_requested.load(std::memory_order_acquire)) {
-            ActivateServingState(admin_server, wrapped_master_service);
+            ActivateServingState(admin_server, wrapped_master_service,
+                                 label_reconciler);
         }
 
         auto server_err = std::move(ec).get();
         LOG(ERROR) << "Master service stopped: " << server_err;
 
         StopLeadershipMonitor(leadership_monitor_handle);
-        DeactivateServingState(admin_server);
+        DeactivateServingState(admin_server, label_reconciler);
         auto err = leader_coordinator.ReleaseLeadership(*leadership_session);
         LOG(INFO) << "Release leadership: " << toString(err);
         auto current_view = leader_coordinator.ReadCurrentView();
diff --git a/mooncake-store/src/ha/oplog/etcd_oplog_store.cpp b/mooncake-store/src/ha/oplog/etcd_oplog_store.cpp
index 34b71f79..8a7cc1e1 100644
--- a/mooncake-store/src/ha/oplog/etcd_oplog_store.cpp
+++ b/mooncake-store/src/ha/oplog/etcd_oplog_store.cpp
@@ -38,8 +38,10 @@ ErrorCode EtcdOpLogStore::Init() {
     // This avoids "key not found" errors when querying the latest sequence ID.
     // Important: Only initialize if the key doesn't exist to avoid overwriting
     // existing data.
-    // Skip for read-only instances to avoid unnecessary etcd writes.
-    if (enable_batch_write_ && !cluster_id_.empty()) {
+    // Note: We allow any instance (reader or writer) to ensure the key exists,
+    // because snapshot resolution depends on it regardless of the HA role.
+    // The Create call uses CAS semantics so concurrent initializers are safe.
+    if (!cluster_id_.empty()) {
         std::string latest_key = BuildLatestKey();
         std::string existing_value;
         EtcdRevisionId revision_id;
diff --git a/mooncake-store/src/ha/snapshot/catalog/backends/embedded/embedded_snapshot_catalog_store.cpp b/mooncake-store/src/ha/snapshot/catalog/backends/embedded/embedded_snapshot_catalog_store.cpp
index 85cfb02a..a3f59e7d 100644
--- a/mooncake-store/src/ha/snapshot/catalog/backends/embedded/embedded_snapshot_catalog_store.cpp
+++ b/mooncake-store/src/ha/snapshot/catalog/backends/embedded/embedded_snapshot_catalog_store.cpp
@@ -19,10 +19,12 @@ ErrorCode ValidateObjectStore(SnapshotObjectStore* object_store) {
 }
 
 tl::expected<SnapshotDescriptor, ErrorCode> LoadSnapshotDescriptor(
-    SnapshotObjectStore* object_store, const SnapshotId& snapshot_id) {
+    SnapshotObjectStore* object_store, const std::string& snapshot_root,
+    const SnapshotId& snapshot_id) {
     std::string descriptor_payload;
     auto get_result = object_store->DownloadString(
-        snapshot_catalog_store_detail::BuildDescriptorKey(snapshot_id),
+        snapshot_catalog_store_detail::BuildDescriptorKey(snapshot_root,
+                                                          snapshot_id),
         descriptor_payload);
     if (!get_result) {
         return tl::make_unexpected(ErrorCode::PERSISTENT_FAIL);
@@ -30,7 +32,7 @@ tl::expected<SnapshotDescriptor, ErrorCode> LoadSnapshotDescriptor(
 
     auto descriptor =
         snapshot_catalog_store_detail::DeserializeSnapshotDescriptor(
-            snapshot_id, descriptor_payload);
+            snapshot_root, snapshot_id, descriptor_payload);
     if (!descriptor) {
         return tl::make_unexpected(descriptor.error());
     }
@@ -40,8 +42,10 @@ tl::expected<SnapshotDescriptor, ErrorCode> LoadSnapshotDescriptor(
 }  // namespace
 
 EmbeddedSnapshotCatalogStore::EmbeddedSnapshotCatalogStore(
-    SnapshotObjectStore* object_store)
-    : object_store_(object_store) {}
+    SnapshotObjectStore* object_store, const std::string& cluster_id)
+    : object_store_(object_store),
+      snapshot_root_(
+          snapshot_catalog_store_detail::BuildSnapshotRoot(cluster_id)) {}
 
 ErrorCode EmbeddedSnapshotCatalogStore::Publish(
     const SnapshotDescriptor& snapshot) {
@@ -55,14 +59,16 @@ ErrorCode EmbeddedSnapshotCatalogStore::Publish(
     }
 
     auto descriptor_result = object_store_->UploadString(
-        snapshot_catalog_store_detail::BuildDescriptorKey(snapshot.snapshot_id),
+        snapshot_catalog_store_detail::BuildDescriptorKey(snapshot_root_,
+                                                          snapshot.snapshot_id),
         snapshot_catalog_store_detail::SerializeSnapshotDescriptor(snapshot));
     if (!descriptor_result) {
         return ErrorCode::PERSISTENT_FAIL;
     }
 
     auto publish_result = object_store_->UploadString(
-        snapshot_catalog_store_detail::BuildLatestKey(), snapshot.snapshot_id);
+        snapshot_catalog_store_detail::BuildLatestKey(snapshot_root_),
+        snapshot.snapshot_id);
     if (!publish_result) {
         return ErrorCode::PERSISTENT_FAIL;
     }
@@ -79,7 +85,8 @@ EmbeddedSnapshotCatalogStore::GetLatest() {
 
     std::string latest_snapshot_id;
     auto get_result = object_store_->DownloadString(
-        snapshot_catalog_store_detail::BuildLatestKey(), latest_snapshot_id);
+        snapshot_catalog_store_detail::BuildLatestKey(snapshot_root_),
+        latest_snapshot_id);
     if (!get_result) {
         if (object_store_->IsNotFoundError(get_result.error())) {
             return std::optional<SnapshotDescriptor>();
@@ -98,7 +105,8 @@ EmbeddedSnapshotCatalogStore::GetLatest() {
 
     std::string descriptor_payload;
     auto descriptor_result = object_store_->DownloadString(
-        snapshot_catalog_store_detail::BuildDescriptorKey(latest_snapshot_id),
+        snapshot_catalog_store_detail::BuildDescriptorKey(snapshot_root_,
+                                                          latest_snapshot_id),
         descriptor_payload);
     if (!descriptor_result) {
         return tl::make_unexpected(ErrorCode::PERSISTENT_FAIL);
@@ -106,7 +114,7 @@ EmbeddedSnapshotCatalogStore::GetLatest() {
 
     auto descriptor =
         snapshot_catalog_store_detail::DeserializeSnapshotDescriptor(
-            latest_snapshot_id, descriptor_payload);
+            snapshot_root_, latest_snapshot_id, descriptor_payload);
     if (!descriptor) {
         return tl::make_unexpected(descriptor.error());
     }
@@ -121,22 +129,20 @@ EmbeddedSnapshotCatalogStore::List(size_t limit) {
     }
 
     std::vector<std::string> object_keys;
-    auto list_result = object_store_->ListObjectsWithPrefix(
-        std::string(snapshot_catalog_store_detail::kSnapshotRoot), object_keys);
+    auto list_result =
+        object_store_->ListObjectsWithPrefix(snapshot_root_, object_keys);
     if (!list_result) {
         return tl::make_unexpected(ErrorCode::PERSISTENT_FAIL);
     }
 
     std::set<SnapshotId, std::greater<>> snapshot_ids;
     for (const auto& object_key : object_keys) {
-        if (object_key.size() <=
-            snapshot_catalog_store_detail::kSnapshotRoot.size()) {
+        if (object_key.size() <= snapshot_root_.size()) {
             continue;
         }
 
         std::string_view suffix(object_key);
-        suffix.remove_prefix(
-            snapshot_catalog_store_detail::kSnapshotRoot.size());
+        suffix.remove_prefix(snapshot_root_.size());
         const size_t slash_pos = suffix.find('/');
         if (slash_pos == std::string_view::npos) {
             continue;
@@ -163,7 +169,8 @@ EmbeddedSnapshotCatalogStore::List(size_t limit) {
         if (limit != 0 && snapshots.size() >= limit) {
             break;
         }
-        auto descriptor = LoadSnapshotDescriptor(object_store_, snapshot_id);
+        auto descriptor =
+            LoadSnapshotDescriptor(object_store_, snapshot_root_, snapshot_id);
         if (!descriptor) {
             LOG(WARNING) << "Skipping unreadable embedded snapshot descriptor, "
                          << "snapshot_id=" << snapshot_id
@@ -209,7 +216,8 @@ ErrorCode EmbeddedSnapshotCatalogStore::Delete(const SnapshotId& snapshot_id) {
     }
 
     auto delete_result = object_store_->DeleteObjectsWithPrefix(
-        snapshot_catalog_store_detail::BuildSnapshotPrefix(snapshot_id));
+        snapshot_catalog_store_detail::BuildSnapshotPrefix(snapshot_root_,
+                                                           snapshot_id));
     if (!delete_result) {
         return ErrorCode::PERSISTENT_FAIL;
     }
@@ -220,7 +228,7 @@ ErrorCode EmbeddedSnapshotCatalogStore::Delete(const SnapshotId& snapshot_id) {
 
     if (next_latest.has_value()) {
         auto publish_result = object_store_->UploadString(
-            snapshot_catalog_store_detail::BuildLatestKey(),
+            snapshot_catalog_store_detail::BuildLatestKey(snapshot_root_),
             next_latest->snapshot_id);
         if (!publish_result) {
             return ErrorCode::PERSISTENT_FAIL;
@@ -229,7 +237,7 @@ ErrorCode EmbeddedSnapshotCatalogStore::Delete(const SnapshotId& snapshot_id) {
     }
 
     auto clear_result = object_store_->DeleteObjectsWithPrefix(
-        snapshot_catalog_store_detail::BuildLatestKey());
+        snapshot_catalog_store_detail::BuildLatestKey(snapshot_root_));
     if (!clear_result) {
         return ErrorCode::PERSISTENT_FAIL;
     }
diff --git a/mooncake-store/src/ha/snapshot/catalog/backends/redis/redis_snapshot_catalog_store.cpp b/mooncake-store/src/ha/snapshot/catalog/backends/redis/redis_snapshot_catalog_store.cpp
index 221cb57e..29e7371d 100644
--- a/mooncake-store/src/ha/snapshot/catalog/backends/redis/redis_snapshot_catalog_store.cpp
+++ b/mooncake-store/src/ha/snapshot/catalog/backends/redis/redis_snapshot_catalog_store.cpp
@@ -1,11 +1,14 @@
 #include "ha/snapshot/catalog/backends/redis/redis_snapshot_catalog_store.h"
 
+#include <cstdlib>
 #include <exception>
 #include <memory>
 #include <optional>
 #include <string_view>
 
 #include <glog/logging.h>
+
+#include "types.h"
 #ifdef STORE_USE_REDIS
 #include <hiredis/hiredis.h>
 #endif
@@ -24,6 +27,8 @@ using common::redis::IsStringReply;
 using common::redis::RedisReplyPtr;
 using common::redis::SanitizeHashTagComponent;
 
+#ifdef STORE_USE_REDIS
+
 tl::expected<long long, ErrorCode> ParseSnapshotScore(
     std::string_view snapshot_id) {
     if (!snapshot_catalog_store_detail::IsValidSnapshotId(snapshot_id)) {
@@ -46,14 +51,16 @@ tl::expected<long long, ErrorCode> ParseSnapshotScore(
 }
 
 tl::expected<SnapshotDescriptor, ErrorCode> LoadSnapshotDescriptor(
-    SnapshotObjectStore* object_store, const SnapshotId& snapshot_id) {
+    SnapshotObjectStore* object_store, const std::string& snapshot_root,
+    const SnapshotId& snapshot_id) {
     if (object_store == nullptr) {
         return tl::make_unexpected(ErrorCode::INVALID_PARAMS);
     }
 
     std::string descriptor_payload;
     auto get_result = object_store->DownloadString(
-        snapshot_catalog_store_detail::BuildDescriptorKey(snapshot_id),
+        snapshot_catalog_store_detail::BuildDescriptorKey(snapshot_root,
+                                                          snapshot_id),
         descriptor_payload);
     if (!get_result) {
         return tl::make_unexpected(ErrorCode::PERSISTENT_FAIL);
@@ -61,15 +68,13 @@ tl::expected<SnapshotDescriptor, ErrorCode> LoadSnapshotDescriptor(
 
     auto descriptor =
         snapshot_catalog_store_detail::DeserializeSnapshotDescriptor(
-            snapshot_id, descriptor_payload);
+            snapshot_root, snapshot_id, descriptor_payload);
     if (!descriptor) {
         return tl::make_unexpected(descriptor.error());
     }
     return descriptor.value();
 }
 
-#ifdef STORE_USE_REDIS
-
 constexpr char kPublishSnapshotScript[] = R"LUA(
 redis.call('ZADD', KEYS[1], ARGV[1], ARGV[2])
 redis.call('SET', KEYS[2], ARGV[2])
@@ -99,6 +104,8 @@ RedisSnapshotCatalogStore::RedisSnapshotCatalogStore(
     : object_store_(object_store),
       connstring_(std::move(connstring)),
       cluster_namespace_(ResolveClusterNamespace(cluster_namespace)),
+      snapshot_root_(
+          snapshot_catalog_store_detail::BuildSnapshotRoot(cluster_namespace_)),
       latest_key_(BuildLatestKey(cluster_namespace_)),
       index_key_(BuildIndexKey(cluster_namespace_)) {}
 
@@ -152,7 +159,8 @@ ErrorCode RedisSnapshotCatalogStore::Publish(
     }
 
     auto descriptor_result = object_store_->UploadString(
-        snapshot_catalog_store_detail::BuildDescriptorKey(snapshot.snapshot_id),
+        snapshot_catalog_store_detail::BuildDescriptorKey(snapshot_root_,
+                                                          snapshot.snapshot_id),
         snapshot_catalog_store_detail::SerializeSnapshotDescriptor(snapshot));
     if (!descriptor_result) {
         return ErrorCode::PERSISTENT_FAIL;
@@ -215,7 +223,8 @@ RedisSnapshotCatalogStore::GetLatest() {
 
     std::string descriptor_payload;
     auto descriptor_result = object_store_->DownloadString(
-        snapshot_catalog_store_detail::BuildDescriptorKey(latest_snapshot_id),
+        snapshot_catalog_store_detail::BuildDescriptorKey(snapshot_root_,
+                                                          latest_snapshot_id),
         descriptor_payload);
     if (!descriptor_result) {
         return tl::make_unexpected(ErrorCode::PERSISTENT_FAIL);
@@ -223,7 +232,7 @@ RedisSnapshotCatalogStore::GetLatest() {
 
     auto descriptor =
         snapshot_catalog_store_detail::DeserializeSnapshotDescriptor(
-            latest_snapshot_id, descriptor_payload);
+            snapshot_root_, latest_snapshot_id, descriptor_payload);
     if (!descriptor) {
         return tl::make_unexpected(descriptor.error());
     }
@@ -269,7 +278,8 @@ RedisSnapshotCatalogStore::List(size_t limit) {
             continue;
         }
 
-        auto descriptor = LoadSnapshotDescriptor(object_store_, snapshot_id);
+        auto descriptor =
+            LoadSnapshotDescriptor(object_store_, snapshot_root_, snapshot_id);
         if (!descriptor) {
             LOG(WARNING) << "Skipping unreadable Redis snapshot descriptor, "
                          << "snapshot_id=" << snapshot_id
@@ -302,7 +312,8 @@ ErrorCode RedisSnapshotCatalogStore::Delete(const SnapshotId& snapshot_id) {
     }
 
     auto delete_result = object_store_->DeleteObjectsWithPrefix(
-        snapshot_catalog_store_detail::BuildSnapshotPrefix(snapshot_id));
+        snapshot_catalog_store_detail::BuildSnapshotPrefix(snapshot_root_,
+                                                           snapshot_id));
     if (!delete_result) {
         LOG(ERROR) << "Failed to delete snapshot objects after Redis catalog "
                       "update, snapshot_id="
@@ -315,10 +326,15 @@ ErrorCode RedisSnapshotCatalogStore::Delete(const SnapshotId& snapshot_id) {
 
 ClusterNamespace RedisSnapshotCatalogStore::ResolveClusterNamespace(
     const ClusterNamespace& cluster_namespace) {
-    if (cluster_namespace.empty()) {
-        return "mooncake";
+    if (!cluster_namespace.empty()) {
+        return cluster_namespace;
     }
-    return cluster_namespace;
+
+    const char* env_cluster_id = std::getenv("MC_STORE_CLUSTER_ID");
+    if (env_cluster_id != nullptr && *env_cluster_id != '\0') {
+        return env_cluster_id;
+    }
+    return DEFAULT_CLUSTER_ID;
 }
 
 std::string RedisSnapshotCatalogStore::BuildLatestKey(
diff --git a/mooncake-store/src/ha/snapshot/catalog_backed_snapshot_provider.cpp b/mooncake-store/src/ha/snapshot/catalog_backed_snapshot_provider.cpp
index 21fffc2a..dc150e8e 100644
--- a/mooncake-store/src/ha/snapshot/catalog_backed_snapshot_provider.cpp
+++ b/mooncake-store/src/ha/snapshot/catalog_backed_snapshot_provider.cpp
@@ -131,14 +131,39 @@ DeserializeStandbyObjectMetadata(
         const auto soft_pin_timestamp_ms = array[index++].as<uint64_t>();
         const auto replica_count = array[index++].as<uint32_t>();
 
-        if (object.via.array.size != 7 + replica_count &&
-            object.via.array.size != 8 + replica_count) {
+        // Optional fields are decoded by type for backward/forward
+        // compatibility with MasterService::MetadataSerializer, which appends
+        // them over time:
+        //   data_type (positive int) appears before the replicas;
+        //   hard_pinned (bool) and group_id (str) trail them.
+        //   v1: 7 + replica_count, no optional fields
+        //   v2: 8 + replica_count, either data_type or hard_pinned
+        //   v3: 9 + replica_count, data_type + hard_pinned or
+        //       hard_pinned + group_id
+        //   v4: 10 + replica_count, data_type + hard_pinned + group_id
+        // 64-bit arithmetic keeps an attacker-controlled near-UINT32_MAX
+        // replica_count from wrapping the bounds and slipping an out-of-bounds
+        // index through.
+        constexpr uint64_t kBaseFieldCount = 7;
+        constexpr uint64_t kMaxOptionalFieldCount = 3;
+        const uint64_t total_elements = object.via.array.size;
+        const uint64_t min_elements = kBaseFieldCount + replica_count;
+        if (total_elements < min_elements ||
+            total_elements > min_elements + kMaxOptionalFieldCount) {
             LOG(ERROR) << "Snapshot metadata entry replica count mismatch, "
                        << "replicas=" << replica_count
-                       << ", total_fields=" << object.via.array.size;
+                       << ", total_fields=" << total_elements;
             return tl::make_unexpected(ErrorCode::DESERIALIZE_FAIL);
         }
 
+        // Skip the optional data_type; the standby restore path does not use
+        // it. A leading positive integer is data_type, whereas a replica is
+        // serialized as an array.
+        if (index < total_elements &&
+            array[index].type == msgpack::type::POSITIVE_INTEGER) {
+            ++index;  // data_type
+        }
+
         const auto lease_timeout = std::chrono::system_clock::time_point(
             std::chrono::milliseconds(lease_timestamp_ms));
         std::optional<std::chrono::system_clock::time_point> soft_pin_timeout;
@@ -156,6 +181,14 @@ DeserializeStandbyObjectMetadata(
         std::vector<Replica::Descriptor> replicas;
         replicas.reserve(replica_count);
         for (uint32_t i = 0; i < replica_count; ++i) {
+            // Defensive bound: a corrupt entry whose first post-count field
+            // looks like a data_type could otherwise read past the array.
+            if (index >= total_elements) {
+                LOG(ERROR) << "Snapshot metadata entry truncated, "
+                           << "replicas=" << replica_count
+                           << ", total_fields=" << total_elements;
+                return tl::make_unexpected(ErrorCode::DESERIALIZE_FAIL);
+            }
             auto replica_result =
                 Serializer<Replica>::deserialize(array[index++], segment_view);
             if (!replica_result) {
@@ -302,7 +335,7 @@ class CatalogBackedSnapshotProvider final : public SnapshotProvider {
             case SnapshotCatalogBackendKind::kEmbedded:
                 catalog_store_ = std::make_unique<
                     ha::backends::embedded::EmbeddedSnapshotCatalogStore>(
-                    object_store_.get());
+                    object_store_.get(), cluster_id_);
                 break;
             case SnapshotCatalogBackendKind::kRedis: {
 #ifndef STORE_USE_REDIS
@@ -352,7 +385,7 @@ class CatalogBackedSnapshotProvider final : public SnapshotProvider {
         if (object_prefix.empty()) {
             object_prefix =
                 ha::snapshot_catalog_store_detail::BuildSnapshotPrefix(
-                    descriptor.snapshot_id);
+                    catalog_store_->GetSnapshotRoot(), descriptor.snapshot_id);
         }
 
         std::string manifest_path = descriptor.manifest_key;
diff --git a/mooncake-store/src/hot_standby_service.cpp b/mooncake-store/src/hot_standby_service.cpp
index 35214dae..b733075b 100644
--- a/mooncake-store/src/hot_standby_service.cpp
+++ b/mooncake-store/src/hot_standby_service.cpp
@@ -16,10 +16,6 @@ namespace mooncake {
 
 HotStandbyService::HotStandbyService(const HotStandbyConfig& config)
     : config_(config) {
-    // Explicitly initialize HA metric manager to ensure thread safety
-    // during metric registration.
-    HAMetricManager::Init();
-
     metadata_store_ = std::make_unique<StandbyMetadataStore>();
     // OpLogApplier will be re-created in Start() with the resolved cluster_id
     // to enable etcd-based operations (e.g. requesting missing OpLog entries).
diff --git a/mooncake-store/src/http_metadata_server.cpp b/mooncake-store/src/http_metadata_server.cpp
index 350ecea6..ec52f041 100644
--- a/mooncake-store/src/http_metadata_server.cpp
+++ b/mooncake-store/src/http_metadata_server.cpp
@@ -57,14 +57,22 @@ void HttpMetadataServer::init_server() {
             std::string body(req.get_body());
             {
                 std::lock_guard<std::mutex> lock(store_mutex_);
-                if (key.find("rpc_meta") != std::string::npos &&
-                    store_.find(std::string(key)) != store_.end()) {
-                    resp.set_status_and_content(
-                        status_type::bad_request,
-                        "Duplicate rpc_meta key not allowed");
-                    return;
+                std::string key_str(key);
+                if (key_str.find("rpc_meta") != std::string::npos) {
+                    auto it = store_.find(key_str);
+                    if (it != store_.end()) {
+                        if (it->second == body) {
+                            resp.set_status_and_content(status_type::ok,
+                                                        "metadata unchanged");
+                            return;
+                        }
+                        resp.set_status_and_content(
+                            status_type::bad_request,
+                            "Duplicate rpc_meta key not allowed");
+                        return;
+                    }
                 }
-                store_[std::string(key)] = body;
+                store_[std::move(key_str)] = body;
             }
 
             resp.set_status_and_content(status_type::ok, "metadata updated");
@@ -127,4 +135,25 @@ KVPoll HttpMetadataServer::poll() const {
     return KVPoll::Success;
 }
 
+bool HttpMetadataServer::removeKey(const std::string& key) {
+    std::lock_guard<std::mutex> lock(store_mutex_);
+    if (store_.erase(key) > 0) {
+        LOG(INFO) << "HttpMetadataServer: removed key=" << key;
+        return true;
+    }
+    return false;
+}
+
+size_t HttpMetadataServer::removeKeys(const std::vector<std::string>& keys) {
+    std::lock_guard<std::mutex> lock(store_mutex_);
+    size_t removed = 0;
+    for (const auto& key : keys) {
+        if (store_.erase(key) > 0) {
+            LOG(INFO) << "HttpMetadataServer: removed key=" << key;
+            ++removed;
+        }
+    }
+    return removed;
+}
+
 }  // namespace mooncake
diff --git a/mooncake-store/src/k8s_lease_helper.cpp b/mooncake-store/src/k8s_lease_helper.cpp
new file mode 100644
index 00000000..301fc31d
--- /dev/null
+++ b/mooncake-store/src/k8s_lease_helper.cpp
@@ -0,0 +1,247 @@
+#include "k8s_lease_helper.h"
+
+#ifdef STORE_USE_K8S_LEASE
+#include "libk8s_lease_wrapper.h"
+#endif
+
+#include <glog/logging.h>
+
+namespace mooncake {
+
+std::mutex K8sLeaseHelper::init_mutex_;
+bool K8sLeaseHelper::initialized_ = false;
+
+#ifdef STORE_USE_K8S_LEASE
+
+ErrorCode K8sLeaseHelper::Init() {
+    std::lock_guard<std::mutex> lock(init_mutex_);
+    if (initialized_) {
+        return ErrorCode::OK;
+    }
+    char* err_msg = nullptr;
+    int ret = K8sLeaseInit(&err_msg);
+    if (ret != 0) {
+        LOG(ERROR) << "Failed to initialize K8s client: " << err_msg;
+        free(err_msg);
+        return ErrorCode::K8S_LEASE_OPERATION_ERROR;
+    }
+    initialized_ = true;
+    return ErrorCode::OK;
+}
+
+ErrorCode K8sLeaseHelper::RunElection(const std::string& ns,
+                                      const std::string& lease,
+                                      const std::string& identity,
+                                      int lease_dur, int renew_deadline,
+                                      int retry_period) {
+    char* err_msg = nullptr;
+    int ret = K8sLeaseRunElection(
+        const_cast<char*>(ns.c_str()), const_cast<char*>(lease.c_str()),
+        const_cast<char*>(identity.c_str()), lease_dur, renew_deadline,
+        retry_period, &err_msg);
+    if (ret != 0) {
+        LOG(ERROR) << "RunElection failed: " << err_msg;
+        free(err_msg);
+        return ErrorCode::K8S_LEASE_OPERATION_ERROR;
+    }
+    return ErrorCode::OK;
+}
+
+ErrorCode K8sLeaseHelper::WaitElected(const std::string& ns,
+                                      const std::string& lease, int timeout_sec,
+                                      int64_t& lease_transitions) {
+    char* err_msg = nullptr;
+    long long transitions = 0;
+    int ret = K8sLeaseWaitElected(const_cast<char*>(ns.c_str()),
+                                  const_cast<char*>(lease.c_str()), timeout_sec,
+                                  &transitions, &err_msg);
+    if (ret != 0) {
+        LOG(ERROR) << "WaitElected failed: " << err_msg;
+        free(err_msg);
+        return ErrorCode::K8S_LEASE_OPERATION_ERROR;
+    }
+    lease_transitions = static_cast<int64_t>(transitions);
+    return ErrorCode::OK;
+}
+
+ErrorCode K8sLeaseHelper::WaitLost(const std::string& ns,
+                                   const std::string& lease) {
+    char* err_msg = nullptr;
+    int ret = K8sLeaseWaitLost(const_cast<char*>(ns.c_str()),
+                               const_cast<char*>(lease.c_str()), &err_msg);
+    if (ret != 0) {
+        LOG(ERROR) << "WaitLost failed: " << err_msg;
+        free(err_msg);
+        return ErrorCode::K8S_LEASE_OPERATION_ERROR;
+    }
+    return ErrorCode::OK;
+}
+
+ErrorCode K8sLeaseHelper::CancelElection(const std::string& ns,
+                                         const std::string& lease) {
+    char* err_msg = nullptr;
+    int ret =
+        K8sLeaseCancelElection(const_cast<char*>(ns.c_str()),
+                               const_cast<char*>(lease.c_str()), &err_msg);
+    if (ret != 0) {
+        LOG(ERROR) << "CancelElection failed: " << err_msg;
+        free(err_msg);
+        return ErrorCode::K8S_LEASE_OPERATION_ERROR;
+    }
+    return ErrorCode::OK;
+}
+
+ErrorCode K8sLeaseHelper::GetHolder(const std::string& ns,
+                                    const std::string& lease,
+                                    std::string& holder,
+                                    int64_t& lease_transitions) {
+    char* err_msg = nullptr;
+    char* holder_ptr = nullptr;
+    long long transitions = 0;
+    int ret = K8sLeaseGetHolder(const_cast<char*>(ns.c_str()),
+                                const_cast<char*>(lease.c_str()), &holder_ptr,
+                                &transitions, &err_msg);
+    if (ret == 1) {
+        holder.clear();
+        lease_transitions = 0;
+        return ErrorCode::K8S_LEASE_NOT_FOUND;
+    }
+    if (ret != 0) {
+        LOG(ERROR) << "GetHolder failed: " << err_msg;
+        free(err_msg);
+        return ErrorCode::K8S_LEASE_OPERATION_ERROR;
+    }
+    if (holder_ptr != nullptr) {
+        holder = std::string(holder_ptr);
+        free(holder_ptr);
+    } else {
+        holder.clear();
+    }
+    lease_transitions = static_cast<int64_t>(transitions);
+    return ErrorCode::OK;
+}
+
+ErrorCode K8sLeaseHelper::WatchHolder(
+    const std::string& ns, const std::string& lease, void* callback_context,
+    void (*callback_func)(void*, const char*, size_t, int64_t)) {
+    char* err_msg = nullptr;
+    int ret = K8sLeaseWatchHolder(const_cast<char*>(ns.c_str()),
+                                  const_cast<char*>(lease.c_str()),
+                                  callback_context, callback_func, &err_msg);
+    if (ret != 0) {
+        LOG(ERROR) << "WatchHolder failed: " << err_msg;
+        free(err_msg);
+        return ErrorCode::K8S_LEASE_OPERATION_ERROR;
+    }
+    return ErrorCode::OK;
+}
+
+ErrorCode K8sLeaseHelper::CancelWatch(const std::string& ns,
+                                      const std::string& lease) {
+    char* err_msg = nullptr;
+    int ret = K8sLeaseCancelWatch(const_cast<char*>(ns.c_str()),
+                                  const_cast<char*>(lease.c_str()), &err_msg);
+    if (ret != 0) {
+        LOG(ERROR) << "CancelWatch failed: " << err_msg;
+        free(err_msg);
+        return ErrorCode::K8S_LEASE_OPERATION_ERROR;
+    }
+    return ErrorCode::OK;
+}
+
+ErrorCode K8sLeaseHelper::SetPodLabel(const std::string& ns,
+                                      const std::string& pod,
+                                      const std::string& key,
+                                      const std::string& value) {
+    char* err_msg = nullptr;
+    int ret = K8sPatchPodLabel(const_cast<char*>(ns.c_str()),
+                               const_cast<char*>(pod.c_str()),
+                               const_cast<char*>(key.c_str()),
+                               const_cast<char*>(value.c_str()), &err_msg);
+    if (ret != 0) {
+        LOG(ERROR) << "SetPodLabel failed: " << err_msg;
+        free(err_msg);
+        return ErrorCode::K8S_LEASE_OPERATION_ERROR;
+    }
+    return ErrorCode::OK;
+}
+
+ErrorCode K8sLeaseHelper::ClearPodLabel(const std::string& ns,
+                                        const std::string& pod,
+                                        const std::string& key) {
+    char* err_msg = nullptr;
+    int ret = K8sRemovePodLabel(const_cast<char*>(ns.c_str()),
+                                const_cast<char*>(pod.c_str()),
+                                const_cast<char*>(key.c_str()), &err_msg);
+    if (ret != 0) {
+        LOG(ERROR) << "ClearPodLabel failed: " << err_msg;
+        free(err_msg);
+        return ErrorCode::K8S_LEASE_OPERATION_ERROR;
+    }
+    return ErrorCode::OK;
+}
+
+#else  // !STORE_USE_K8S_LEASE
+
+ErrorCode K8sLeaseHelper::Init() {
+    LOG(FATAL) << "K8s Lease is not enabled in compilation";
+    return ErrorCode::K8S_LEASE_OPERATION_ERROR;
+}
+
+ErrorCode K8sLeaseHelper::RunElection(const std::string&, const std::string&,
+                                      const std::string&, int, int, int) {
+    LOG(FATAL) << "K8s Lease is not enabled in compilation";
+    return ErrorCode::K8S_LEASE_OPERATION_ERROR;
+}
+
+ErrorCode K8sLeaseHelper::WaitElected(const std::string&, const std::string&,
+                                      int, int64_t&) {
+    LOG(FATAL) << "K8s Lease is not enabled in compilation";
+    return ErrorCode::K8S_LEASE_OPERATION_ERROR;
+}
+
+ErrorCode K8sLeaseHelper::WaitLost(const std::string&, const std::string&) {
+    LOG(FATAL) << "K8s Lease is not enabled in compilation";
+    return ErrorCode::K8S_LEASE_OPERATION_ERROR;
+}
+
+ErrorCode K8sLeaseHelper::CancelElection(const std::string&,
+                                         const std::string&) {
+    LOG(FATAL) << "K8s Lease is not enabled in compilation";
+    return ErrorCode::K8S_LEASE_OPERATION_ERROR;
+}
+
+ErrorCode K8sLeaseHelper::GetHolder(const std::string&, const std::string&,
+                                    std::string&, int64_t&) {
+    LOG(FATAL) << "K8s Lease is not enabled in compilation";
+    return ErrorCode::K8S_LEASE_OPERATION_ERROR;
+}
+
+ErrorCode K8sLeaseHelper::WatchHolder(const std::string&, const std::string&,
+                                      void*,
+                                      void (*)(void*, const char*, size_t,
+                                               int64_t)) {
+    LOG(FATAL) << "K8s Lease is not enabled in compilation";
+    return ErrorCode::K8S_LEASE_OPERATION_ERROR;
+}
+
+ErrorCode K8sLeaseHelper::CancelWatch(const std::string&, const std::string&) {
+    LOG(FATAL) << "K8s Lease is not enabled in compilation";
+    return ErrorCode::K8S_LEASE_OPERATION_ERROR;
+}
+
+ErrorCode K8sLeaseHelper::SetPodLabel(const std::string&, const std::string&,
+                                      const std::string&, const std::string&) {
+    LOG(ERROR) << "K8s Lease is not enabled in compilation";
+    return ErrorCode::K8S_LEASE_OPERATION_ERROR;
+}
+
+ErrorCode K8sLeaseHelper::ClearPodLabel(const std::string&, const std::string&,
+                                        const std::string&) {
+    LOG(ERROR) << "K8s Lease is not enabled in compilation";
+    return ErrorCode::K8S_LEASE_OPERATION_ERROR;
+}
+
+#endif  // STORE_USE_K8S_LEASE
+
+}  // namespace mooncake
diff --git a/mooncake-store/src/local_hot_cache.cpp b/mooncake-store/src/local_hot_cache.cpp
index b4ad027f..c254862c 100644
--- a/mooncake-store/src/local_hot_cache.cpp
+++ b/mooncake-store/src/local_hot_cache.cpp
@@ -4,6 +4,7 @@
 #include <cstring>
 #include <cstdlib>
 #include <cstdint>
+#include <regex>
 #include <shared_mutex>
 #include <glog/logging.h>
 
@@ -75,7 +76,23 @@ LocalHotCache::~LocalHotCache() {
 
 bool LocalHotCache::PutHotKey(HotMemBlock* block) {
     std::unique_lock<std::shared_mutex> lk(lru_mutex_);
+    return putHotKeyLocked(block);
+}
 
+bool LocalHotCache::PutHotKey(HotMemBlock* block,
+                              const HotCachePutToken& token) {
+    std::unique_lock<std::shared_mutex> lk(lru_mutex_);
+    // Validate the token and publish under the same lock so a concurrent
+    // Remove/Bump cannot slip in between the check and the publish.
+    if (block && !block->key_.empty() &&
+        !isPutTokenValidLocked(block->key_, token)) {
+        // Stale async fill: drop the data and return the block to the pool.
+        block->key_.clear();
+    }
+    return putHotKeyLocked(block);
+}
+
+bool LocalHotCache::putHotKeyLocked(HotMemBlock* block) {
     // Drain deferred LRU touches
     drainDeferredTouches();
 
@@ -92,7 +109,8 @@ bool LocalHotCache::PutHotKey(HotMemBlock* block) {
 
     // Race condition check: did someone else insert this key while we were
     // copying
-    if (key_to_lru_it_.find(key) != key_to_lru_it_.end()) {
+    if (key_to_lru_it_.find(key) != key_to_lru_it_.end() ||
+        hasActiveBlockForKeyLocked(key)) {
         // Lost race -> Return to lru tail as free block
         block->key_.clear();
         block->ref_count = 0;
@@ -134,17 +152,211 @@ HotMemBlock* LocalHotCache::GetHotKey(const std::string& key) {
 }
 
 void LocalHotCache::ReleaseHotKey(const std::string& key) {
-    std::shared_lock<std::shared_mutex> lk(lru_mutex_);
+    std::unique_lock<std::shared_mutex> lk(lru_mutex_);
+    auto it = key_to_lru_it_.find(key);
+    if (it != key_to_lru_it_.end()) {
+        HotMemBlock* block = *(it->second);
+        if (block && block->ref_count > 0) {
+            block->ref_count--;
+            return;
+        }
+    }
+
+    // The entry may have been removed while a reader still holds the block.
+    // Keep the key on active removed blocks so release can find and retire it.
+    for (auto& owned_block : blocks_) {
+        HotMemBlock* block = owned_block.get();
+        if (block && block->key_ == key && block->ref_count > 0) {
+            block->ref_count--;
+            if (block->ref_count == 0 &&
+                key_to_lru_it_.find(key) == key_to_lru_it_.end()) {
+                block->key_.clear();
+                block->accessed.store(false, std::memory_order_relaxed);
+            }
+            return;
+        }
+    }
+}
+
+bool LocalHotCache::hasActiveBlockForKeyLocked(const std::string& key) const {
+    for (const auto& owned_block : blocks_) {
+        const HotMemBlock* block = owned_block.get();
+        if (block && block->key_ == key && block->ref_count > 0) {
+            return true;
+        }
+    }
+    return false;
+}
+
+bool LocalHotCache::removeHotKeyLocked(const std::string& key) {
     auto it = key_to_lru_it_.find(key);
     if (it == key_to_lru_it_.end()) {
+        return false;
+    }
+
+    HotMemBlock* block = *(it->second);
+    lru_queue_.erase(it->second);
+    key_to_lru_it_.erase(it);
+
+    if (block->ref_count == 0) {
+        block->key_.clear();
+        block->accessed.store(false, std::memory_order_relaxed);
+    }
+    lru_queue_.push_back(block);
+
+    return true;
+}
+
+bool LocalHotCache::RemoveHotKey(const std::string& key) {
+    std::unique_lock<std::shared_mutex> lk(lru_mutex_);
+
+    key_generation_[key]++;
+    drainDeferredTouches();
+
+    const bool removed = removeHotKeyLocked(key);
+    if (removed) {
+        VLOG(2) << "Removed hot key: " << key << " from hot cache";
+    }
+    return removed;
+}
+
+size_t LocalHotCache::RemoveHotKeys(const std::vector<std::string>& keys) {
+    if (keys.empty()) {
+        return 0;
+    }
+
+    std::unique_lock<std::shared_mutex> lk(lru_mutex_);
+    drainDeferredTouches();
+
+    size_t removed = 0;
+    for (const auto& key : keys) {
+        key_generation_[key]++;
+        if (removeHotKeyLocked(key)) {
+            ++removed;
+        }
+    }
+    return removed;
+}
+
+size_t LocalHotCache::RemoveHotKeysByRegex(const std::string& regex_pattern) {
+    std::regex pattern;
+    try {
+        pattern = std::regex(regex_pattern, std::regex::ECMAScript);
+    } catch (const std::regex_error& e) {
+        LOG(ERROR) << "RemoveHotKeysByRegex: invalid pattern: " << regex_pattern
+                   << ", error: " << e.what();
+        return 0;
+    }
+
+    std::vector<std::string> matching_keys;
+    {
+        std::shared_lock<std::shared_mutex> lk(lru_mutex_);
+        matching_keys.reserve(key_to_lru_it_.size());
+        for (const auto& [key, _] : key_to_lru_it_) {
+            if (std::regex_search(key, pattern)) {
+                matching_keys.emplace_back(key);
+            }
+        }
+    }
+
+    if (matching_keys.empty()) {
+        return 0;
+    }
+
+    std::unique_lock<std::shared_mutex> lk(lru_mutex_);
+    drainDeferredTouches();
+
+    size_t removed = 0;
+    for (const auto& key : matching_keys) {
+        key_generation_[key]++;
+        if (removeHotKeyLocked(key)) {
+            ++removed;
+        }
+    }
+    return removed;
+}
+
+size_t LocalHotCache::RemoveAllHotKeys() {
+    std::unique_lock<std::shared_mutex> lk(lru_mutex_);
+    cache_epoch_.fetch_add(1, std::memory_order_relaxed);
+
+    std::vector<std::string> keys;
+    keys.reserve(key_to_lru_it_.size());
+    for (const auto& [key, _] : key_to_lru_it_) {
+        keys.emplace_back(key);
+    }
+
+    size_t removed = 0;
+    for (const auto& key : keys) {
+        if (removeHotKeyLocked(key)) {
+            ++removed;
+        }
+    }
+    key_generation_.clear();
+    return removed;
+}
+
+void LocalHotCache::BumpKeyGeneration(const std::string& key) {
+    std::unique_lock<std::shared_mutex> lk(lru_mutex_);
+    key_generation_[key]++;
+}
+
+void LocalHotCache::BumpKeyGenerations(const std::vector<std::string>& keys) {
+    if (keys.empty()) {
         return;
     }
-    HotMemBlock* block = *(it->second);
-    if (block) {
-        block->ref_count--;
+
+    std::unique_lock<std::shared_mutex> lk(lru_mutex_);
+    for (const auto& key : keys) {
+        key_generation_[key]++;
     }
 }
 
+void LocalHotCache::BumpCacheEpoch() {
+    cache_epoch_.fetch_add(1, std::memory_order_relaxed);
+}
+
+void LocalHotCache::Clear() {
+    std::unique_lock<std::shared_mutex> lk(lru_mutex_);
+    cache_epoch_.fetch_add(1, std::memory_order_relaxed);
+
+    std::vector<std::string> keys;
+    keys.reserve(key_to_lru_it_.size());
+    for (const auto& [key, _] : key_to_lru_it_) {
+        keys.emplace_back(key);
+    }
+
+    for (const auto& key : keys) {
+        removeHotKeyLocked(key);
+    }
+    key_generation_.clear();
+}
+
+HotCachePutToken LocalHotCache::AcquirePutToken(const std::string& key) {
+    std::shared_lock<std::shared_mutex> lk(lru_mutex_);
+    HotCachePutToken token;
+    token.cache_epoch = cache_epoch_.load(std::memory_order_relaxed);
+    auto it = key_generation_.find(key);
+    token.key_generation = (it != key_generation_.end()) ? it->second : 0;
+    return token;
+}
+
+bool LocalHotCache::isPutTokenValidLocked(const std::string& key,
+                                          const HotCachePutToken& token) const {
+    if (token.cache_epoch != cache_epoch_.load(std::memory_order_relaxed)) {
+        return false;
+    }
+    auto it = key_generation_.find(key);
+    const uint64_t current_gen = (it != key_generation_.end()) ? it->second : 0;
+    return token.key_generation == current_gen;
+}
+
+bool LocalHotCache::IsPutTokenValid(const std::string& key,
+                                    const HotCachePutToken& token) const {
+    std::shared_lock<std::shared_mutex> lk(lru_mutex_);
+    return isPutTokenValidLocked(key, token);
+}
+
 bool LocalHotCache::TouchHotKey(const std::string& key) {
     std::shared_lock<std::shared_mutex> lk(lru_mutex_);
     auto it = key_to_lru_it_.find(key);
@@ -303,6 +515,7 @@ bool LocalHotCacheHandler::SubmitPutTask(const std::string& key,
     }
 
     // Try to get a free block (may evict from LRU tail)
+    HotCachePutToken token = hot_cache_->AcquirePutToken(key);
     HotMemBlock* block = hot_cache_->GetFreeBlock();
     if (!block) {
         LOG(ERROR) << "Hot cache is fully in-use, fail to get a free block: "
@@ -310,8 +523,14 @@ bool LocalHotCacheHandler::SubmitPutTask(const std::string& key,
         return false;
     }
 
-    // Check size compatibility
-    if (slice.size > block->size) {
+    // Check size compatibility against the block's physical capacity.
+    // block->size carries the *logical* length of the last object stored in
+    // the block (set below and relied upon as the object length on the read
+    // path), and GetFreeBlock() does not reset it on reuse. Comparing against
+    // block->size would therefore reject any object larger than the previous
+    // one a recycled block happened to hold, permanently shrinking the block's
+    // usable capacity. The true capacity is the fixed block size.
+    if (slice.size > hot_cache_->GetBlockSize()) {
         // Slice too big for block, return block to pool
         block->key_.clear();
         hot_cache_->PutHotKey(block);
@@ -324,7 +543,7 @@ bool LocalHotCacheHandler::SubmitPutTask(const std::string& key,
     block->size = slice.size;
     block->key_ = key;  // Set key for insertion
 
-    HotCachePutTask task(key, slice, block, hot_cache_);
+    HotCachePutTask task(key, slice, block, hot_cache_, token);
 
     {
         std::lock_guard<std::mutex> lock(queue_mutex_);
@@ -367,11 +586,12 @@ void LocalHotCacheHandler::workerThread() {
         // Execute the task if we have one
         if (task.hot_cache && task.block) {
             try {
-                // Insert the pre-filled block into LRU
-                if (task.hot_cache->PutHotKey(task.block)) {
+                // Validate token and publish atomically: a Remove/Bump that
+                // races after the check can no longer resurrect a stale fill.
+                if (task.hot_cache->PutHotKey(task.block, task.token)) {
                     VLOG(2) << "Put task completed: " << task.key;
                 } else {
-                    VLOG(2) << "Put task skipped: " << task.key;
+                    VLOG(2) << "Put task skipped or cancelled: " << task.key;
                 }
             } catch (const std::exception& e) {
                 LOG(ERROR) << "Exception during async hot cache put for key "
diff --git a/mooncake-store/src/master.cpp b/mooncake-store/src/master.cpp
index d8704903..e6970ba3 100644
--- a/mooncake-store/src/master.cpp
+++ b/mooncake-store/src/master.cpp
@@ -1,22 +1,33 @@
 #include <gflags/gflags.h>
 #include <glog/logging.h>
 
+#include <atomic>  // For std::atomic
 #include <chrono>  // For std::chrono
+#include <cmath>
+#include <cstdint>
 #include <csignal>
+#include <cstdlib>  // For std::getenv
+#include <fstream>  // For std::ifstream
+#include <limits>
 #include <memory>  // For std::unique_ptr
+#include <string>
 #include <thread>  // For std::thread
+#include <json/json.h>
 #include <ylt/coro_rpc/coro_rpc_server.hpp>
 #include <ylt/easylog/record.hpp>
 
 #include "default_config.h"
 #include "duration_utils.h"
 #include "ha/leadership/master_service_supervisor.h"
+
 #include "http_metadata_server.h"
+#include "master_admin_service.h"
 #include "rpc_service.h"
 #include "types.h"
 #include "utils.h"
 
 #include "master_config.h"
+#include "version.h"
 
 using namespace coro_rpc;
 using namespace async_simple;
@@ -56,6 +67,57 @@ uint64_t ParseDurationFlagOrDie(const char* flag_name,
     return parsed_value;
 }
 
+bool TryGetDoubleConfig(const mooncake::DefaultConfig& default_config,
+                        const std::string& key, double* value) {
+    const double missing = std::numeric_limits<double>::lowest();
+    default_config.GetDouble(key, value, missing);
+    return *value != missing;
+}
+
+// Derive the metadata server address for cleanup when it is deployed
+// separately. Priority: MOONCAKE_TE_META_DATA_SERVER, then the
+// "metadata_server" field of MOONCAKE_CONFIG_PATH. Returns "" if none found;
+// the caller validates the scheme (only http(s) is supported).
+std::string ResolveMetadataServerForCleanup() {
+    if (const char* env = std::getenv("MOONCAKE_TE_META_DATA_SERVER")) {
+        std::string value(env);
+        // P2PHANDSHAKE has no central metadata server, nothing to clean up.
+        if (!value.empty() && value != "P2PHANDSHAKE") {
+            return value;
+        }
+    }
+
+    if (const char* cfg = std::getenv("MOONCAKE_CONFIG_PATH")) {
+        if (cfg[0] != '\0') {
+            try {
+                std::ifstream fin(cfg);
+                if (fin) {
+                    Json::CharReaderBuilder builder;
+                    Json::Value root;
+                    std::string errs;
+                    if (Json::parseFromStream(builder, fin, &root, &errs) &&
+                        root.isMember("metadata_server") &&
+                        root["metadata_server"].isString()) {
+                        return root["metadata_server"].asString();
+                    }
+                    if (!errs.empty()) {
+                        LOG(WARNING) << "Failed to parse MOONCAKE_CONFIG_PATH ("
+                                     << cfg << "): " << errs;
+                    }
+                } else {
+                    LOG(WARNING) << "Cannot open MOONCAKE_CONFIG_PATH (" << cfg
+                                 << "): file does not exist or is not readable";
+                }
+            } catch (const std::exception& e) {
+                LOG(WARNING) << "Error reading MOONCAKE_CONFIG_PATH (" << cfg
+                             << "): " << e.what();
+            }
+        }
+    }
+
+    return {};
+}
+
 }  // namespace
 
 DEFINE_string(config_path, "", "master service config file path");
@@ -78,10 +140,15 @@ DEFINE_bool(allow_evict_soft_pinned_objects,
 DEFINE_validator(default_kv_lease_ttl, ValidateDurationFlag);
 DEFINE_validator(default_kv_soft_pin_ttl, ValidateDurationFlag);
 DEFINE_double(eviction_ratio, mooncake::DEFAULT_EVICTION_RATIO,
-              "Ratio of objects to evict when storage space is full");
+              "Ratio of objects to evict when Memory space is full");
 DEFINE_double(eviction_high_watermark_ratio,
               mooncake::DEFAULT_EVICTION_HIGH_WATERMARK_RATIO,
-              "Ratio of high watermark trigger eviction");
+              "Ratio of high watermark trigger eviction in Memory");
+DEFINE_double(nof_eviction_ratio, mooncake::DEFAULT_NOF_EVICTION_RATIO,
+              "Ratio of objects to evict when NoF SSD space is full");
+DEFINE_double(nof_eviction_high_watermark_ratio,
+              mooncake::DEFAULT_NOF_EVICTION_HIGH_WATERMARK_RATIO,
+              "Ratio of high watermark trigger eviction in NoF SSD");
 // RPC server configuration parameters (new, preferred)
 // TODO: deprecate port and max_threads in the future
 DEFINE_int32(rpc_thread_num, 0,
@@ -101,29 +168,97 @@ DEFINE_bool(rpc_enable_tcp_no_delay, true,
             "Enable TCP_NODELAY for RPC connections");
 DEFINE_validator(eviction_ratio, [](const char* flagname, double value) {
     if (value < 0.0 || value > 1.0) {
-        LOG(FATAL) << "Eviction ratio must be between 0.0 and 1.0";
+        LOG(FATAL) << "Mem eviction ratio must be between 0.0 and 1.0";
+        return false;
+    }
+    return true;
+});
+DEFINE_validator(nof_eviction_ratio, [](const char* flagname, double value) {
+    if (value < 0.0 || value > 1.0) {
+        LOG(FATAL) << "NoF eviction ratio must be between 0.0 and 1.0";
         return false;
     }
     return true;
 });
 DEFINE_bool(enable_ha, false,
-            "Enable high availability, which depends on etcd");
+            "Enable high availability using the configured HA backend");
 DEFINE_bool(enable_offload, false, "Enable offload availability");
 DEFINE_bool(offload_on_evict, false,
             "Defer LOCAL_DISK offload to eviction time instead of PutEnd");
 DEFINE_bool(offload_force_evict, false,
             "Force-evict objects exceeding offload cap without disk offload");
+DEFINE_uint64(offloading_queue_limit, 50000,
+              "Maximum number of objects allowed in the offloading queue per "
+              "local disk segment. Increase to allow more objects to be "
+              "offloaded to SSD before force-eviction kicks in");
+DEFINE_validator(offloading_queue_limit, [](const char* flagname,
+                                            uint64_t value) {
+    // Zero would cause PushOffloadingQueue to always return
+    // KEYS_ULTRA_LIMIT, disabling offload entirely. The upper
+    // bound (1e8) keeps `offloading_queue_limit_ *
+    // offload_cap_ratio_` well within signed long range to
+    // avoid overflow when computing offload_cap in
+    // BatchEvict / EvictTenantMemoryForQuota.
+    if (value == 0) {
+        LOG(FATAL) << "offloading_queue_limit must be greater than 0";
+        return false;
+    }
+    if (value > 100'000'000ULL) {
+        LOG(FATAL) << "offloading_queue_limit must be <= "
+                      "100000000 to avoid overflow";
+        return false;
+    }
+    return true;
+});
+DEFINE_double(offload_cap_ratio, 0.5,
+              "Per-cycle offload cap as a fraction of offloading_queue_limit. "
+              "Controls how many objects can be queued for offload in a single "
+              "eviction cycle before falling back to force-evict");
+DEFINE_validator(offload_cap_ratio, [](const char* flagname, double value) {
+    if (value < 0.0 || value > 1.0) {
+        LOG(FATAL) << "offload_cap_ratio must be between 0.0 and 1.0";
+        return false;
+    }
+    return true;
+});
+DEFINE_bool(promotion_on_hit, false,
+            "Promote LOCAL_DISK-only keys to MEMORY on read access (mirror of "
+            "offload_on_evict)");
+DEFINE_uint32(promotion_admission_threshold, 2,
+              "Min CountMinSketch count for a key before promotion fires "
+              "(set 1 to disable second-touch gating)");
+DEFINE_uint32(promotion_queue_limit, 50000,
+              "Max in-flight promotion tasks across all shards");
+DEFINE_uint32(promotion_max_per_heartbeat, 1,
+              "Max promotion tasks returned to a single client per "
+              "PromotionObjectHeartbeat call. Each task is a synchronous "
+              "SSD-read + RDMA-write on the client; serializing them avoids "
+              "blocking past the client-liveness window. Default 1 is "
+              "conservative.");
 DEFINE_string(ha_backend_type, "etcd",
               "HA backend type, e.g. etcd | redis | k8s");
 DEFINE_string(ha_backend_connstring, "",
-              "HA backend connection string. If unset, fallback to "
-              "etcd_endpoints for backward compatibility");
+              "HA backend connection string. If unset, only backend_type=etcd "
+              "falls back to etcd_endpoints for backward compatibility");
 DEFINE_string(
     etcd_endpoints, "",
     "Endpoints of ETCD server, separated by semicolon, required in HA mode");
-DEFINE_int64(client_ttl, mooncake::DEFAULT_CLIENT_LIVE_TTL_SEC,
-             "How long a client is considered alive after the last ping, only "
-             "used in HA mode");
+DEFINE_int64(
+    client_ttl, mooncake::DEFAULT_CLIENT_LIVE_TTL_SEC,
+    "Seconds a client stays considered alive after the last heartbeat. "
+    "If this TTL elapses without a refresh, the master treats the "
+    "client as disconnected and may unmount its segments");
+DEFINE_int64(nof_heartbeat_interval_sec,
+             mooncake::DEFAULT_NOF_HEARTBEAT_INTERVAL_SEC,
+             "How often master probes each mounted NoF segment");
+DEFINE_uint32(nof_heartbeat_probe_timeout_ms,
+              mooncake::DEFAULT_NOF_HEARTBEAT_PROBE_TIMEOUT_MS,
+              "Timeout in milliseconds for a single NoF heartbeat probe");
+DEFINE_uint32(
+    nof_heartbeat_failures_threshold,
+    mooncake::DEFAULT_NOF_HEARTBEAT_FAILURES_THRESHOLD,
+    "Consecutive NoF heartbeat failures required before unmounting a NoF "
+    "segment");
 
 DEFINE_string(root_fs_dir, mooncake::DEFAULT_ROOT_FS_DIR,
               "Root directory for storage backend, used in HA mode");
@@ -138,13 +273,26 @@ DEFINE_string(memory_allocator, "offset",
               "Memory allocator for global segments, cachelib | offset");
 DEFINE_string(
     allocation_strategy, "random",
-    "Allocation strategy for segments, random | free_ratio_first | cxl");
+    "Allocation strategy for segments, random | free_ratio_first | cxl | "
+    "ssd_free_ratio_first");
 DEFINE_bool(enable_http_metadata_server, false,
             "Enable HTTP metadata server instead of etcd");
 DEFINE_int32(http_metadata_server_port, 8080,
              "Port for HTTP metadata server to listen on");
 DEFINE_string(http_metadata_server_host, "0.0.0.0",
               "Host for HTTP metadata server to bind to");
+DEFINE_bool(
+    enable_metadata_cleanup_on_timeout, false,
+    "Enable cleanup of HTTP metadata (mooncake/ram/*, mooncake/rpc_meta/*) "
+    "when client heartbeat times out. Works in two modes: (1) co-located "
+    "(enable_http_metadata_server=true) via in-process removal, or "
+    "(2) separately-deployed metadata server via async HTTP DELETE.");
+
+DEFINE_string(pod_name, "",
+              "Pod name for K8s label-based routing (default: $POD_NAME)");
+DEFINE_string(pod_namespace, "",
+              "Pod namespace for K8s label-based routing "
+              "(default: $POD_NAMESPACE)");
 
 DEFINE_uint64(put_start_discard_timeout_sec,
               mooncake::DEFAULT_PUT_START_DISCARD_TIMEOUT,
@@ -158,6 +306,12 @@ DEFINE_bool(enable_disk_eviction, true,
 DEFINE_uint64(
     quota_bytes, 0,
     "Quota for storage backend in bytes (0 = use default 90% of capacity)");
+DEFINE_bool(enable_multi_tenants, false,
+            "Enable strict multi-tenant namespace and quota admission");
+DEFINE_string(tenant_quota_connector_type, "file",
+              "Tenant quota policy connector type");
+DEFINE_string(tenant_quota_connector_uri, "",
+              "Tenant quota policy connector URI");
 
 // Snapshot related configuration flags (migrated from global_flags)
 DEFINE_string(snapshot_backup_dir, "",
@@ -214,10 +368,9 @@ namespace {
 
 std::string ResolveHABackendConnstring(
     const mooncake::MasterConfig& master_config) {
-    if (!master_config.ha_backend_connstring.empty()) {
-        return master_config.ha_backend_connstring;
-    }
-    return master_config.etcd_endpoints;
+    return mooncake::ResolveConfiguredHABackendConnstring(
+        master_config.ha_backend_type, master_config.ha_backend_connstring,
+        master_config.etcd_endpoints);
 }
 
 void ResolveRpcAddressFromInterfaceOrDie(
@@ -284,14 +437,86 @@ void InitMasterConf(const mooncake::DefaultConfig& default_config,
     default_config.GetBool("allow_evict_soft_pinned_objects",
                            &master_config.allow_evict_soft_pinned_objects,
                            FLAGS_allow_evict_soft_pinned_objects);
+    for (uint16_t type_id = 0; type_id <= UINT8_MAX; ++type_id) {
+        const auto data_type = static_cast<mooncake::ObjectDataType>(
+            static_cast<uint8_t>(type_id));
+        const std::string type_key = std::to_string(type_id);
+
+        bool has_score_policy = false;
+        mooncake::ObjectTypeEvictionScorePolicy score_policy;
+        double reuse_scale = 0.0;
+        if (TryGetDoubleConfig(default_config,
+                               "object_type_eviction_score_policies." +
+                                   type_key + ".reuse_scale",
+                               &reuse_scale)) {
+            score_policy.reuse_scale = reuse_scale;
+            has_score_policy = true;
+        }
+
+        double soft_pin_weight = 0.0;
+        if (TryGetDoubleConfig(default_config,
+                               "object_type_eviction_score_policies." +
+                                   type_key + ".soft_pin_weight",
+                               &soft_pin_weight)) {
+            score_policy.soft_pin_weight = soft_pin_weight;
+            has_score_policy = true;
+        }
+
+        const int64_t missing_eviction_grace =
+            std::numeric_limits<int64_t>::min();
+        int64_t eviction_grace = missing_eviction_grace;
+        default_config.GetInt64("object_type_eviction_score_policies." +
+                                    type_key + ".eviction_grace",
+                                &eviction_grace, missing_eviction_grace);
+        if (eviction_grace != missing_eviction_grace) {
+            score_policy.eviction_grace = eviction_grace;
+            has_score_policy = true;
+        }
+
+        if (has_score_policy) {
+            mooncake::ValidateObjectTypeEvictionScorePolicy(score_policy);
+            master_config.object_type_eviction_score_policies[data_type] =
+                score_policy;
+        }
+
+        double budget_ratio = 0.0;
+        if (TryGetDoubleConfig(
+                default_config,
+                "object_type_eviction_policies." + type_key + ".budget_ratio",
+                &budget_ratio)) {
+            if (!std::isfinite(budget_ratio) || budget_ratio < 0.0 ||
+                budget_ratio > 1.0) {
+                LOG(FATAL) << "object_type_eviction_policies." << type_key
+                           << ".budget_ratio must be finite and between 0.0 "
+                              "and 1.0";
+            }
+            master_config.object_type_eviction_policies[data_type]
+                .budget_ratio = budget_ratio;
+        }
+    }
     default_config.GetDouble("eviction_ratio", &master_config.eviction_ratio,
                              FLAGS_eviction_ratio);
     default_config.GetDouble("eviction_high_watermark_ratio",
                              &master_config.eviction_high_watermark_ratio,
                              FLAGS_eviction_high_watermark_ratio);
+    default_config.GetDouble("nof_eviction_ratio",
+                             &master_config.nof_eviction_ratio,
+                             FLAGS_nof_eviction_ratio);
+    default_config.GetDouble("nof_eviction_high_watermark_ratio",
+                             &master_config.nof_eviction_high_watermark_ratio,
+                             FLAGS_nof_eviction_high_watermark_ratio);
     default_config.GetInt64("client_live_ttl_sec",
                             &master_config.client_live_ttl_sec,
                             FLAGS_client_ttl);
+    default_config.GetInt64("nof_heartbeat_interval_sec",
+                            &master_config.nof_heartbeat_interval_sec,
+                            FLAGS_nof_heartbeat_interval_sec);
+    default_config.GetUInt32("nof_heartbeat_probe_timeout_ms",
+                             &master_config.nof_heartbeat_probe_timeout_ms,
+                             FLAGS_nof_heartbeat_probe_timeout_ms);
+    default_config.GetUInt32("nof_heartbeat_failures_threshold",
+                             &master_config.nof_heartbeat_failures_threshold,
+                             FLAGS_nof_heartbeat_failures_threshold);
 
     default_config.GetBool("enable_ha", &master_config.enable_ha,
                            FLAGS_enable_ha);
@@ -302,6 +527,28 @@ void InitMasterConf(const mooncake::DefaultConfig& default_config,
     default_config.GetBool("offload_force_evict",
                            &master_config.offload_force_evict,
                            FLAGS_offload_force_evict);
+    {
+        uint64_t tmp_offloading_queue_limit = FLAGS_offloading_queue_limit;
+        default_config.GetUInt64("offloading_queue_limit",
+                                 &tmp_offloading_queue_limit,
+                                 FLAGS_offloading_queue_limit);
+        master_config.offloading_queue_limit =
+            static_cast<size_t>(tmp_offloading_queue_limit);
+    }
+    default_config.GetDouble("offload_cap_ratio",
+                             &master_config.offload_cap_ratio,
+                             FLAGS_offload_cap_ratio);
+    default_config.GetBool("promotion_on_hit", &master_config.promotion_on_hit,
+                           FLAGS_promotion_on_hit);
+    default_config.GetUInt32("promotion_admission_threshold",
+                             &master_config.promotion_admission_threshold,
+                             FLAGS_promotion_admission_threshold);
+    default_config.GetUInt32("promotion_queue_limit",
+                             &master_config.promotion_queue_limit,
+                             FLAGS_promotion_queue_limit);
+    default_config.GetUInt32("promotion_max_per_heartbeat",
+                             &master_config.promotion_max_per_heartbeat,
+                             FLAGS_promotion_max_per_heartbeat);
     default_config.GetString("ha_backend_type", &master_config.ha_backend_type,
                              FLAGS_ha_backend_type);
     default_config.GetString("ha_backend_connstring",
@@ -331,6 +578,13 @@ void InitMasterConf(const mooncake::DefaultConfig& default_config,
     default_config.GetString("http_metadata_server_host",
                              &master_config.http_metadata_server_host,
                              FLAGS_http_metadata_server_host);
+    default_config.GetString("pod_name", &master_config.pod_name,
+                             FLAGS_pod_name);
+    default_config.GetString("pod_namespace", &master_config.pod_namespace,
+                             FLAGS_pod_namespace);
+    default_config.GetBool("enable_metadata_cleanup_on_timeout",
+                           &master_config.enable_metadata_cleanup_on_timeout,
+                           FLAGS_enable_metadata_cleanup_on_timeout);
     default_config.GetUInt64("put_start_discard_timeout_sec",
                              &master_config.put_start_discard_timeout_sec,
                              FLAGS_put_start_discard_timeout_sec);
@@ -342,6 +596,15 @@ void InitMasterConf(const mooncake::DefaultConfig& default_config,
                            FLAGS_enable_disk_eviction);
     default_config.GetUInt64("quota_bytes", &master_config.quota_bytes,
                              FLAGS_quota_bytes);
+    default_config.GetBool("enable_multi_tenants",
+                           &master_config.enable_multi_tenants,
+                           FLAGS_enable_multi_tenants);
+    default_config.GetString("tenant_quota_connector_type",
+                             &master_config.tenant_quota_connector_type,
+                             FLAGS_tenant_quota_connector_type);
+    default_config.GetString("tenant_quota_connector_uri",
+                             &master_config.tenant_quota_connector_uri,
+                             FLAGS_tenant_quota_connector_uri);
 
     default_config.GetString("snapshot_backup_dir",
                              &master_config.snapshot_backup_dir,
@@ -533,6 +796,18 @@ void LoadConfigFromCmdline(mooncake::MasterConfig& master_config,
         master_config.eviction_high_watermark_ratio =
             FLAGS_eviction_high_watermark_ratio;
     }
+    if ((google::GetCommandLineFlagInfo("nof_eviction_ratio", &info) &&
+         !info.is_default) ||
+        !conf_set) {
+        master_config.nof_eviction_ratio = FLAGS_nof_eviction_ratio;
+    }
+    if ((google::GetCommandLineFlagInfo("nof_eviction_high_watermark_ratio",
+                                        &info) &&
+         !info.is_default) ||
+        !conf_set) {
+        master_config.nof_eviction_high_watermark_ratio =
+            FLAGS_nof_eviction_high_watermark_ratio;
+    }
     if ((google::GetCommandLineFlagInfo("enable_ha", &info) &&
          !info.is_default) ||
         !conf_set) {
@@ -543,6 +818,63 @@ void LoadConfigFromCmdline(mooncake::MasterConfig& master_config,
         !conf_set) {
         master_config.enable_offload = FLAGS_enable_offload;
     }
+    if ((google::GetCommandLineFlagInfo("offload_on_evict", &info) &&
+         !info.is_default) ||
+        !conf_set) {
+        master_config.offload_on_evict = FLAGS_offload_on_evict;
+    }
+    if ((google::GetCommandLineFlagInfo("offload_force_evict", &info) &&
+         !info.is_default) ||
+        !conf_set) {
+        master_config.offload_force_evict = FLAGS_offload_force_evict;
+    }
+    if ((google::GetCommandLineFlagInfo("offloading_queue_limit", &info) &&
+         !info.is_default) ||
+        !conf_set) {
+        master_config.offloading_queue_limit =
+            static_cast<size_t>(FLAGS_offloading_queue_limit);
+    }
+    if ((google::GetCommandLineFlagInfo("offload_cap_ratio", &info) &&
+         !info.is_default) ||
+        !conf_set) {
+        master_config.offload_cap_ratio = FLAGS_offload_cap_ratio;
+    }
+    if ((google::GetCommandLineFlagInfo("promotion_on_hit", &info) &&
+         !info.is_default) ||
+        !conf_set) {
+        master_config.promotion_on_hit = FLAGS_promotion_on_hit;
+    }
+    if ((google::GetCommandLineFlagInfo("promotion_admission_threshold",
+                                        &info) &&
+         !info.is_default) ||
+        !conf_set) {
+        master_config.promotion_admission_threshold =
+            FLAGS_promotion_admission_threshold;
+    }
+    if ((google::GetCommandLineFlagInfo("promotion_queue_limit", &info) &&
+         !info.is_default) ||
+        !conf_set) {
+        master_config.promotion_queue_limit = FLAGS_promotion_queue_limit;
+    }
+    if ((google::GetCommandLineFlagInfo("promotion_max_per_heartbeat", &info) &&
+         !info.is_default) ||
+        !conf_set) {
+        master_config.promotion_max_per_heartbeat =
+            FLAGS_promotion_max_per_heartbeat;
+    }
+    // Clamp promotion_admission_threshold into the sketch counter's
+    // representable range. The CountMinSketch uses 8-bit saturating
+    // counters (max 255) so any threshold beyond that would silently
+    // make the gate unreachable; clamping at parse time fails loudly and
+    // gives the gate a stable contract to compare uint8_t against.
+    if (master_config.promotion_admission_threshold > 255) {
+        LOG(WARNING) << "promotion_admission_threshold="
+                     << master_config.promotion_admission_threshold
+                     << " exceeds the CountMinSketch counter max (255). "
+                     << "Clamping to 255. Lower the configured value to "
+                     << "silence this warning.";
+        master_config.promotion_admission_threshold = 255;
+    }
     if ((google::GetCommandLineFlagInfo("ha_backend_type", &info) &&
          !info.is_default) ||
         !conf_set) {
@@ -563,6 +895,26 @@ void LoadConfigFromCmdline(mooncake::MasterConfig& master_config,
         !conf_set) {
         master_config.client_live_ttl_sec = FLAGS_client_ttl;
     }
+    if ((google::GetCommandLineFlagInfo("nof_heartbeat_interval_sec", &info) &&
+         !info.is_default) ||
+        !conf_set) {
+        master_config.nof_heartbeat_interval_sec =
+            FLAGS_nof_heartbeat_interval_sec;
+    }
+    if ((google::GetCommandLineFlagInfo("nof_heartbeat_probe_timeout_ms",
+                                        &info) &&
+         !info.is_default) ||
+        !conf_set) {
+        master_config.nof_heartbeat_probe_timeout_ms =
+            FLAGS_nof_heartbeat_probe_timeout_ms;
+    }
+    if ((google::GetCommandLineFlagInfo("nof_heartbeat_failures_threshold",
+                                        &info) &&
+         !info.is_default) ||
+        !conf_set) {
+        master_config.nof_heartbeat_failures_threshold =
+            FLAGS_nof_heartbeat_failures_threshold;
+    }
     if ((google::GetCommandLineFlagInfo("cluster_id", &info) &&
          !info.is_default) ||
         !conf_set) {
@@ -606,6 +958,23 @@ void LoadConfigFromCmdline(mooncake::MasterConfig& master_config,
         master_config.http_metadata_server_host =
             FLAGS_http_metadata_server_host;
     }
+    if ((google::GetCommandLineFlagInfo("pod_name", &info) &&
+         !info.is_default) ||
+        !conf_set) {
+        master_config.pod_name = FLAGS_pod_name;
+    }
+    if ((google::GetCommandLineFlagInfo("pod_namespace", &info) &&
+         !info.is_default) ||
+        !conf_set) {
+        master_config.pod_namespace = FLAGS_pod_namespace;
+    }
+    if ((google::GetCommandLineFlagInfo("enable_metadata_cleanup_on_timeout",
+                                        &info) &&
+         !info.is_default) ||
+        !conf_set) {
+        master_config.enable_metadata_cleanup_on_timeout =
+            FLAGS_enable_metadata_cleanup_on_timeout;
+    }
     if ((google::GetCommandLineFlagInfo("put_start_discard_timeout_sec",
                                         &info) &&
          !info.is_default) ||
@@ -630,6 +999,23 @@ void LoadConfigFromCmdline(mooncake::MasterConfig& master_config,
         !conf_set) {
         master_config.quota_bytes = FLAGS_quota_bytes;
     }
+    if ((google::GetCommandLineFlagInfo("enable_multi_tenants", &info) &&
+         !info.is_default) ||
+        !conf_set) {
+        master_config.enable_multi_tenants = FLAGS_enable_multi_tenants;
+    }
+    if ((google::GetCommandLineFlagInfo("tenant_quota_connector_type", &info) &&
+         !info.is_default) ||
+        !conf_set) {
+        master_config.tenant_quota_connector_type =
+            FLAGS_tenant_quota_connector_type;
+    }
+    if ((google::GetCommandLineFlagInfo("tenant_quota_connector_uri", &info) &&
+         !info.is_default) ||
+        !conf_set) {
+        master_config.tenant_quota_connector_uri =
+            FLAGS_tenant_quota_connector_uri;
+    }
     if ((google::GetCommandLineFlagInfo("max_total_finished_tasks", &info) &&
          !info.is_default) ||
         !conf_set) {
@@ -805,10 +1191,21 @@ std::unique_ptr<mooncake::HttpMetadataServer> StartHttpMetadataServer(
 int main(int argc, char* argv[]) {
     mooncake::init_ylt_log_level();
     // Initialize gflags
+    gflags::SetVersionString(mooncake::MOONCAKE_DISPLAY_VERSION);
     gflags::ParseCommandLineFlags(&argc, &argv, true);
 
     if (!FLAGS_log_dir.empty()) {
         google::InitGoogleLogging(argv[0]);
+        // Merge all master logs into a single journal file in --log_dir,
+        // reusing glog: every record is already written to its own severity
+        // file and all lower ones, so the INFO sink is a complete journal.
+        // Disable the higher-severity files so everything lands in one file.
+        const std::string log_base = FLAGS_log_dir + "/mooncake_master.";
+        google::SetLogDestination(google::GLOG_INFO, log_base.c_str());
+        google::SetLogDestination(google::GLOG_WARNING, "");
+        google::SetLogDestination(google::GLOG_ERROR, "");
+        google::SetLogDestination(google::GLOG_FATAL, "");
+        google::SetLogSymlink(google::GLOG_INFO, "mooncake_master");
     }
 
     // Initialize the master configuration
@@ -828,11 +1225,24 @@ int main(int argc, char* argv[]) {
     LoadConfigFromCmdline(master_config, !conf_path.empty());
     ResolveRpcAddressFromInterfaceOrDie(master_config);
 
+    // Fall back to environment variables for pod identity (K8s Downward API)
+    if (master_config.pod_name.empty()) {
+        const char* env = std::getenv("POD_NAME");
+        if (env) master_config.pod_name = env;
+    }
+    if (master_config.pod_namespace.empty()) {
+        const char* env = std::getenv("POD_NAMESPACE");
+        if (env) master_config.pod_namespace = env;
+    }
+
     const std::string ha_backend_connstring =
         ResolveHABackendConnstring(master_config);
     if (master_config.enable_ha && ha_backend_connstring.empty()) {
         LOG(FATAL) << "HA backend connection string must be set when "
-                   << "enable_ha is true";
+                   << "enable_ha is true, backend_type="
+                   << master_config.ha_backend_type
+                   << ". Only backend_type=etcd may fall back to "
+                   << "etcd_endpoints";
         return 1;
     }
     if (!master_config.enable_ha && (!ha_backend_connstring.empty() ||
@@ -862,6 +1272,41 @@ int main(int argc, char* argv[]) {
     if (value && std::string_view(value) == "rdma") {
         protocol = "rdma";
     }
+
+    // enable_metadata_cleanup_on_timeout requires a reachable HTTP metadata
+    // server. Two topologies are supported:
+    //   1) Co-located: enable_http_metadata_server=true -> the master cleans
+    //      up via the in-process server (no network overhead).
+    //   2) Separately deployed: the master derives the metadata server address
+    //      from the cluster's existing configuration
+    //      (MOONCAKE_TE_META_DATA_SERVER, or MOONCAKE_CONFIG_PATH json's
+    //      "metadata_server") and cleans up via HTTP DELETE. Only http(s)
+    //      endpoints are supported for now (etcd/redis left for future work).
+    // If neither is available, cleanup is disabled with a warning so the main
+    // process is never affected.
+    std::string http_metadata_remote_url;
+    if (master_config.enable_metadata_cleanup_on_timeout &&
+        !master_config.enable_http_metadata_server) {
+        std::string derived = ResolveMetadataServerForCleanup();
+        if (derived.rfind("http://", 0) == 0 ||
+            derived.rfind("https://", 0) == 0) {
+            http_metadata_remote_url = std::move(derived);
+            LOG(INFO) << "enable_metadata_cleanup_on_timeout: HTTP metadata "
+                         "server is deployed separately; cleanup will target "
+                      << http_metadata_remote_url;
+        } else {
+            LOG(WARNING)
+                << "enable_metadata_cleanup_on_timeout is set to true but "
+                   "enable_http_metadata_server is false and no HTTP metadata "
+                   "server address could be derived from the cluster config "
+                   "(set "
+                   "MOONCAKE_TE_META_DATA_SERVER=http://host:port/metadata "
+                   "or MOONCAKE_CONFIG_PATH). Disabling metadata cleanup on "
+                   "timeout.";
+            master_config.enable_metadata_cleanup_on_timeout = false;
+        }
+    }
+
     LOG(INFO)
         << "Master service started on port " << master_config.rpc_port
         << ", max_threads=" << master_config.rpc_thread_num
@@ -878,6 +1323,8 @@ int main(int argc, char* argv[]) {
         << ", enable_offload=" << master_config.enable_offload
         << ", offload_on_evict=" << master_config.offload_on_evict
         << ", offload_force_evict=" << master_config.offload_force_evict
+        << ", offloading_queue_limit=" << master_config.offloading_queue_limit
+        << ", offload_cap_ratio=" << master_config.offload_cap_ratio
         << ", ha_backend_type=" << master_config.ha_backend_type
         << ", ha_backend_connstring=" << ha_backend_connstring
         << ", etcd_endpoints=" << master_config.etcd_endpoints
@@ -901,6 +1348,8 @@ int main(int argc, char* argv[]) {
         << master_config.http_metadata_server_port
         << ", http_metadata_server_host="
         << master_config.http_metadata_server_host
+        << ", enable_metadata_cleanup_on_timeout="
+        << master_config.enable_metadata_cleanup_on_timeout
         << ", put_start_discard_timeout_sec="
         << master_config.put_start_discard_timeout_sec
         << ", put_start_release_timeout_sec="
@@ -946,9 +1395,20 @@ int main(int argc, char* argv[]) {
         std::this_thread::sleep_for(std::chrono::seconds(1));
     }
 
+    // Metadata cleanup on client timeout (used by both the HA and non-HA
+    // paths): prefer the co-located in-process server, else the separate URL.
+    mooncake::HttpMetadataServer* metadata_server_ptr = nullptr;
+    if (master_config.enable_metadata_cleanup_on_timeout &&
+        master_config.enable_http_metadata_server) {
+        metadata_server_ptr = http_metadata_server.get();
+    }
+
     if (master_config.enable_ha) {
-        mooncake::ha::MasterServiceSupervisor supervisor(
-            mooncake::MasterServiceSupervisorConfig{master_config});
+        mooncake::MasterServiceSupervisorConfig supervisor_config{
+            master_config};
+        supervisor_config.http_metadata_server = metadata_server_ptr;
+        supervisor_config.http_metadata_remote_url = http_metadata_remote_url;
+        mooncake::ha::MasterServiceSupervisor supervisor(supervisor_config);
         return supervisor.Start();
     } else {
         // version is not used in non-HA mode, just pass a dummy value
@@ -964,7 +1424,8 @@ int main(int argc, char* argv[]) {
         }
         auto wrapped_master_service =
             std::make_shared<mooncake::WrappedMasterService>(
-                mooncake::WrappedMasterServiceConfig(master_config, version));
+                mooncake::WrappedMasterServiceConfig(master_config, version),
+                metadata_server_ptr, http_metadata_remote_url);
         mooncake::MasterAdminServer admin_server(
             static_cast<uint16_t>(master_config.metrics_port),
             master_config.enable_metric_reporting);
@@ -978,6 +1439,32 @@ int main(int argc, char* argv[]) {
         admin_server.SetServiceAvailable(true);
 
         mooncake::RegisterRpcService(server, *wrapped_master_service);
-        return server.start();
+
+        static std::atomic<bool> shutdown_requested{false};
+        auto signal_handler = [](int /* signum */) {
+            shutdown_requested.store(true);
+        };
+        std::signal(SIGINT, signal_handler);
+        std::signal(SIGTERM, signal_handler);
+
+        int server_result = 0;
+        std::atomic<bool> server_done{false};
+        std::thread server_thread([&server, &server_result, &server_done]() {
+            server_result = server.start();
+            server_done.store(true);
+        });
+
+        while (!shutdown_requested.load() && !server_done.load()) {
+            std::this_thread::sleep_for(std::chrono::milliseconds(100));
+        }
+
+        server.stop();
+        server_thread.join();
+
+        if (shutdown_requested.load()) {
+            LOG(INFO) << "Shutdown signal received, exiting gracefully";
+            return 0;
+        }
+        return server_result;
     }
 }
diff --git a/mooncake-store/src/master_admin_service.cpp b/mooncake-store/src/master_admin_service.cpp
new file mode 100644
index 00000000..0898dadb
--- /dev/null
+++ b/mooncake-store/src/master_admin_service.cpp
@@ -0,0 +1,1193 @@
+#include "master_admin_service.h"
+
+#include <algorithm>
+#include <chrono>
+#include <cstdint>
+#include <exception>
+#include <map>
+#include <sstream>
+#include <string_view>
+#include <utility>
+#include <vector>
+
+#include <glog/logging.h>
+#include <ylt/reflection/user_reflect_macro.hpp>
+#include <ylt/struct_json/json_reader.h>
+#include <ylt/struct_json/json_writer.h>
+#include <ylt/util/tl/expected.hpp>
+
+#include "allocator.h"
+#include "ha_metric_manager.h"
+#include "master_metric_manager.h"
+#include "rpc_service.h"
+#include "types.h"
+
+namespace mooncake {
+
+const uint64_t kMetricReportIntervalSeconds = 10;
+
+namespace {
+
+std::string AppendMetricSections(std::string primary, std::string secondary) {
+    if (!primary.empty() && primary.back() != '\n') {
+        primary.push_back('\n');
+    }
+    primary += secondary;
+    return primary;
+}
+
+struct HttpErrorResponse {
+    bool success{false};
+    int32_t error_code{0};
+    std::string error_message;
+};
+YLT_REFL(HttpErrorResponse, success, error_code, error_message);
+
+struct HttpSimpleErrorResponse {
+    bool success{false};
+    std::string error;
+};
+YLT_REFL(HttpSimpleErrorResponse, success, error);
+
+template <typename T>
+void WriteJsonResponse(coro_http::coro_http_response& resp,
+                       coro_http::status_type status, const T& payload) {
+    std::string json;
+    struct_json::to_json(payload, json);
+    resp.add_header("Content-Type", "application/json; charset=utf-8");
+    resp.set_status_and_content(status, std::move(json));
+}
+
+template <typename T>
+std::string EnumToString(const T& value) {
+    std::ostringstream oss;
+    oss << value;
+    return oss.str();
+}
+
+coro_http::status_type ErrorCodeToHttpStatus(ErrorCode error) {
+    switch (error) {
+        case ErrorCode::INVALID_PARAMS:
+            return coro_http::status_type::bad_request;
+        case ErrorCode::JOB_NOT_FOUND:
+        case ErrorCode::SEGMENT_NOT_FOUND:
+        case ErrorCode::OBJECT_NOT_FOUND:
+        case ErrorCode::TENANT_NOT_REGISTERED:
+            return coro_http::status_type::not_found;
+        case ErrorCode::UNAVAILABLE_IN_CURRENT_MODE:
+        case ErrorCode::UNAVAILABLE_IN_CURRENT_STATUS:
+        case ErrorCode::TENANT_NOT_EMPTY:
+            return coro_http::status_type::conflict;
+        default:
+            return coro_http::status_type::internal_server_error;
+    }
+}
+
+void WriteErrorResponse(coro_http::coro_http_response& resp,
+                        coro_http::status_type status, ErrorCode error,
+                        std::string message = {}) {
+    HttpErrorResponse payload;
+    payload.error_code = toInt(error);
+    payload.error_message =
+        message.empty() ? toString(error) : std::move(message);
+    WriteJsonResponse(resp, status, payload);
+}
+
+void WriteSimpleErrorResponse(coro_http::coro_http_response& resp,
+                              coro_http::status_type status,
+                              std::string error) {
+    HttpSimpleErrorResponse payload;
+    payload.error = std::move(error);
+    WriteJsonResponse(resp, status, payload);
+}
+
+void SetServiceUnavailable(coro_http::coro_http_response& resp,
+                           std::string message) {
+    WriteErrorResponse(resp, coro_http::status_type::service_unavailable,
+                       ErrorCode::UNAVAILABLE_IN_CURRENT_MODE,
+                       std::move(message));
+}
+
+tl::expected<UUID, ErrorCode> ParseJobId(std::string_view job_id_view) {
+    if (job_id_view.empty()) {
+        return tl::make_unexpected(ErrorCode::INVALID_PARAMS);
+    }
+
+    UUID job_id;
+    if (!StringToUuid(std::string(job_id_view), job_id)) {
+        return tl::make_unexpected(ErrorCode::INVALID_PARAMS);
+    }
+
+    return job_id;
+}
+
+std::string EscapeJson(std::string_view input) {
+    std::string escaped;
+    escaped.reserve(input.size());
+    for (char ch : input) {
+        switch (ch) {
+            case '\\':
+                escaped += "\\\\";
+                break;
+            case '"':
+                escaped += "\\\"";
+                break;
+            case '\n':
+                escaped += "\\n";
+                break;
+            case '\r':
+                escaped += "\\r";
+                break;
+            case '\t':
+                escaped += "\\t";
+                break;
+            default:
+                escaped.push_back(ch);
+                break;
+        }
+    }
+    return escaped;
+}
+
+std::string EscapePrometheusLabel(std::string_view input) {
+    std::string escaped;
+    escaped.reserve(input.size());
+    for (char ch : input) {
+        switch (ch) {
+            case '\\':
+                escaped += "\\\\";
+                break;
+            case '"':
+                escaped += "\\\"";
+                break;
+            case '\n':
+                escaped += "\\n";
+                break;
+            default:
+                escaped.push_back(ch);
+                break;
+        }
+    }
+    return escaped;
+}
+
+struct HttpTenantQuotaSnapshot {
+    std::string tenant_id;
+    uint64_t requested_quota_bytes{0};
+    uint64_t effective_quota_bytes{0};
+    uint64_t used_bytes{0};
+    uint64_t reserved_bytes{0};
+    uint64_t committed_count{0};
+    uint64_t metadata_object_count{0};
+    bool over_quota{false};
+    bool has_explicit_policy{false};
+};
+YLT_REFL(HttpTenantQuotaSnapshot, tenant_id, requested_quota_bytes,
+         effective_quota_bytes, used_bytes, reserved_bytes, committed_count,
+         metadata_object_count, over_quota, has_explicit_policy);
+
+HttpTenantQuotaSnapshot ToHttpTenantQuotaSnapshot(
+    const TenantQuotaSnapshot& snapshot) {
+    return HttpTenantQuotaSnapshot{
+        .tenant_id = snapshot.tenant_id,
+        .requested_quota_bytes = snapshot.requested_quota_bytes,
+        .effective_quota_bytes = snapshot.effective_quota_bytes,
+        .used_bytes = snapshot.used_bytes,
+        .reserved_bytes = snapshot.reserved_bytes,
+        .committed_count = snapshot.committed_count,
+        .metadata_object_count = snapshot.metadata_object_count,
+        .over_quota = snapshot.over_quota,
+        .has_explicit_policy = snapshot.has_explicit_policy,
+    };
+}
+
+struct HttpTenantQuotaListResponse {
+    bool success{true};
+    std::vector<HttpTenantQuotaSnapshot> data;
+};
+YLT_REFL(HttpTenantQuotaListResponse, success, data);
+
+struct HttpTenantQuotaResponse {
+    bool success{true};
+    HttpTenantQuotaSnapshot data;
+};
+YLT_REFL(HttpTenantQuotaResponse, success, data);
+
+struct HttpTenantQuotaDeleteResponse {
+    bool success{true};
+    std::optional<HttpTenantQuotaSnapshot> data;
+};
+YLT_REFL(HttpTenantQuotaDeleteResponse, success, data);
+
+struct HttpTenantQuotaPolicyRequest {
+    uint64_t requested_quota_bytes{0};
+};
+YLT_REFL(HttpTenantQuotaPolicyRequest, requested_quota_bytes);
+
+tl::expected<std::string, ErrorCode> ParseAdminTenantId(
+    coro_http::coro_http_request& req) {
+    auto tenant_id_view = req.get_decode_query_value("tenant_id");
+    if (tenant_id_view.empty()) {
+        return tl::make_unexpected(ErrorCode::INVALID_PARAMS);
+    }
+    std::string tenant_id = NormalizeTenantId(std::string(tenant_id_view));
+    if (!IsValidTenantId(tenant_id)) {
+        return tl::make_unexpected(ErrorCode::INVALID_PARAMS);
+    }
+    return tenant_id;
+}
+
+tl::expected<HttpTenantQuotaPolicyRequest, std::string> ParseQuotaPolicyBody(
+    coro_http::coro_http_request& req) {
+    HttpTenantQuotaPolicyRequest request;
+    try {
+        struct_json::from_json(request, req.get_body());
+    } catch (const std::exception& e) {
+        return tl::make_unexpected(std::string("Invalid JSON body: ") +
+                                   e.what());
+    }
+    return request;
+}
+
+}  // namespace
+
+MasterAdminServer::MasterAdminServer(uint16_t http_port,
+                                     bool enable_metric_reporting)
+    : http_port_(http_port),
+      enable_metric_reporting_(enable_metric_reporting),
+      http_server_(4, http_port) {}
+
+MasterAdminServer::~MasterAdminServer() { Stop(); }
+
+bool MasterAdminServer::Start() {
+    if (started_.load()) {
+        return true;
+    }
+
+    RegisterHandler();
+
+    auto ec = http_server_.async_start();
+    if (ec.hasResult()) {
+        LOG(ERROR) << "Failed to start master admin server on port "
+                   << http_port_;
+        return false;
+    }
+
+    started_.store(true);
+    if (enable_metric_reporting_) {
+        metric_report_running_.store(true);
+        metric_report_thread_ = std::thread([this]() {
+            while (metric_report_running_.load()) {
+                const auto snapshot = SnapshotState();
+                std::ostringstream log_stream;
+                log_stream << "Master Admin Metrics: role="
+                           << ha::MasterRuntimeRoleToString(snapshot.state)
+                           << ", state="
+                           << ha::MasterRuntimeStateToString(snapshot.state)
+                           << ", service_ready="
+                           << (snapshot.service_available ? "true" : "false")
+                           << ", master={"
+                           << MasterMetricManager::instance()
+                                  .get_summary_string_and_update_snapshot()
+                           << "}"
+                           << ", ha={"
+                           << HAMetricManager::instance().get_summary_string()
+                           << "}";
+                if (snapshot.leader_view.has_value()) {
+                    log_stream
+                        << ", leader=" << snapshot.leader_view->leader_address
+                        << ", view_version="
+                        << snapshot.leader_view->view_version;
+                }
+                LOG(INFO) << log_stream.str();
+                if (metric_report_stop_sem_.try_acquire_for(
+                        std::chrono::seconds(kMetricReportIntervalSeconds))) {
+                    break;
+                }
+            }
+        });
+    }
+
+    LOG(INFO) << "Master admin server started on port " << http_server_.port();
+    return true;
+}
+
+void MasterAdminServer::Stop() {
+    metric_report_running_.store(false, std::memory_order_relaxed);
+    if (started_.exchange(false)) {
+        http_server_.stop();
+    }
+    if (metric_report_thread_.joinable()) {
+        metric_report_stop_sem_.release();
+        metric_report_thread_.join();
+    }
+}
+
+void MasterAdminServer::SetRuntimeState(ha::MasterRuntimeState state) {
+    std::lock_guard<std::mutex> lock(state_mutex_);
+    state_ = state;
+}
+
+void MasterAdminServer::SetObservedLeader(
+    const std::optional<ha::MasterView>& leader_view) {
+    std::lock_guard<std::mutex> lock(state_mutex_);
+    leader_view_ = leader_view;
+}
+
+void MasterAdminServer::SetServiceDelegate(
+    std::shared_ptr<WrappedMasterService> service) {
+    std::lock_guard<std::mutex> lock(state_mutex_);
+    service_ = std::move(service);
+    if (!service_) {
+        service_available_ = false;
+    }
+}
+
+void MasterAdminServer::SetServiceAvailable(bool available) {
+    std::lock_guard<std::mutex> lock(state_mutex_);
+    service_available_ = available && service_ != nullptr;
+}
+
+MasterAdminServer::RuntimeSnapshot MasterAdminServer::SnapshotState() const {
+    std::lock_guard<std::mutex> lock(state_mutex_);
+    return RuntimeSnapshot{
+        .state = state_,
+        .leader_view = leader_view_,
+        .service = service_,
+        .service_available = service_available_,
+    };
+}
+
+std::string MasterAdminServer::BuildMetricsText() const {
+    std::string metrics = AppendMetricSections(
+        MasterMetricManager::instance().serialize_metrics(),
+        HAMetricManager::instance().serialize_metrics());
+    auto tenant_metrics = BuildTenantQuotaMetricsText();
+    if (tenant_metrics.empty()) {
+        return metrics;
+    }
+    return AppendMetricSections(std::move(metrics), std::move(tenant_metrics));
+}
+
+std::string MasterAdminServer::BuildTenantQuotaMetricsText() const {
+    auto service = GetActiveService();
+    if (!service) {
+        return "";
+    }
+    auto snapshots_result = service->ListTenantQuotaSnapshots();
+    auto capacity_result = service->GetTenantQuotaAllocatableCapacityBytes();
+    if (!snapshots_result || !capacity_result) {
+        return "";
+    }
+
+    const auto& snapshots = snapshots_result.value();
+    uint64_t requested_sum = 0;
+    uint64_t effective_sum = 0;
+    std::ostringstream tenant_metrics;
+    tenant_metrics
+        << "# HELP mooncake_tenant_quota_requested_bytes Requested tenant "
+           "quota policy in bytes\n"
+        << "# TYPE mooncake_tenant_quota_requested_bytes gauge\n"
+        << "# HELP mooncake_tenant_quota_effective_bytes Effective tenant "
+           "quota in bytes\n"
+        << "# TYPE mooncake_tenant_quota_effective_bytes gauge\n"
+        << "# HELP mooncake_tenant_quota_used_bytes Tenant committed quota "
+           "usage in bytes\n"
+        << "# TYPE mooncake_tenant_quota_used_bytes gauge\n"
+        << "# HELP mooncake_tenant_quota_reserved_bytes Tenant reserved quota "
+           "usage in bytes\n"
+        << "# TYPE mooncake_tenant_quota_reserved_bytes gauge\n"
+        << "# HELP mooncake_tenant_quota_committed_count Tenant committed "
+           "object count\n"
+        << "# TYPE mooncake_tenant_quota_committed_count gauge\n"
+        << "# HELP mooncake_tenant_quota_metadata_object_count Tenant "
+           "metadata object count\n"
+        << "# TYPE mooncake_tenant_quota_metadata_object_count gauge\n"
+        << "# HELP mooncake_tenant_quota_over_quota Tenant over-quota flag\n"
+        << "# TYPE mooncake_tenant_quota_over_quota gauge\n"
+        << "# HELP mooncake_tenant_quota_explicit_policy Tenant explicit "
+           "policy flag\n"
+        << "# TYPE mooncake_tenant_quota_explicit_policy gauge\n";
+    for (const auto& snapshot : snapshots) {
+        requested_sum += snapshot.requested_quota_bytes;
+        effective_sum += snapshot.effective_quota_bytes;
+        const auto tenant = EscapePrometheusLabel(snapshot.tenant_id);
+        tenant_metrics << "mooncake_tenant_quota_requested_bytes{tenant_id=\""
+                       << tenant << "\"} " << snapshot.requested_quota_bytes
+                       << "\n";
+        tenant_metrics << "mooncake_tenant_quota_effective_bytes{tenant_id=\""
+                       << tenant << "\"} " << snapshot.effective_quota_bytes
+                       << "\n";
+        tenant_metrics << "mooncake_tenant_quota_used_bytes{tenant_id=\""
+                       << tenant << "\"} " << snapshot.used_bytes << "\n";
+        tenant_metrics << "mooncake_tenant_quota_reserved_bytes{tenant_id=\""
+                       << tenant << "\"} " << snapshot.reserved_bytes << "\n";
+        tenant_metrics << "mooncake_tenant_quota_committed_count{tenant_id=\""
+                       << tenant << "\"} " << snapshot.committed_count << "\n";
+        tenant_metrics
+            << "mooncake_tenant_quota_metadata_object_count{tenant_id=\""
+            << tenant << "\"} " << snapshot.metadata_object_count << "\n";
+        tenant_metrics << "mooncake_tenant_quota_over_quota{tenant_id=\""
+                       << tenant << "\"} " << (snapshot.over_quota ? 1 : 0)
+                       << "\n";
+        tenant_metrics << "mooncake_tenant_quota_explicit_policy{tenant_id=\""
+                       << tenant << "\"} "
+                       << (snapshot.has_explicit_policy ? 1 : 0) << "\n";
+    }
+    tenant_metrics
+        << "# HELP mooncake_tenant_quota_allocatable_capacity_bytes Global "
+           "tenant quota allocatable capacity in bytes\n"
+        << "# TYPE mooncake_tenant_quota_allocatable_capacity_bytes gauge\n"
+        << "mooncake_tenant_quota_allocatable_capacity_bytes "
+        << capacity_result.value() << "\n"
+        << "# HELP mooncake_tenant_quota_requested_bytes_sum Global requested "
+           "tenant quota sum in bytes\n"
+        << "# TYPE mooncake_tenant_quota_requested_bytes_sum gauge\n"
+        << "mooncake_tenant_quota_requested_bytes_sum " << requested_sum << "\n"
+        << "# HELP mooncake_tenant_quota_effective_bytes_sum Global effective "
+           "tenant quota sum in bytes\n"
+        << "# TYPE mooncake_tenant_quota_effective_bytes_sum gauge\n"
+        << "mooncake_tenant_quota_effective_bytes_sum " << effective_sum
+        << "\n";
+
+    return tenant_metrics.str();
+}
+
+std::string MasterAdminServer::BuildMetricsSummaryText() const {
+    const auto snapshot = SnapshotState();
+    std::ostringstream oss;
+    oss << "role=" << ha::MasterRuntimeRoleToString(snapshot.state)
+        << ", state=" << ha::MasterRuntimeStateToString(snapshot.state)
+        << ", service_ready=" << (snapshot.service_available ? "true" : "false")
+        << ", master={" << MasterMetricManager::instance().get_summary_string()
+        << "}, ha={" << HAMetricManager::instance().get_summary_string() << "}";
+    if (snapshot.leader_view.has_value()) {
+        oss << ", leader=" << snapshot.leader_view->leader_address
+            << ", view_version=" << snapshot.leader_view->view_version;
+    }
+    return oss.str();
+}
+
+struct HttpHealthResponse {
+    std::string status;
+    std::string role;
+    std::string ha_state;
+    bool service_ready{false};
+    std::optional<std::string> leader_address;
+    std::optional<uint64_t> view_version;
+};
+YLT_REFL(HttpHealthResponse, status, role, ha_state, service_ready,
+         leader_address, view_version);
+
+void MasterAdminServer::HandleHealth(coro_http::coro_http_request&,
+                                     coro_http::coro_http_response& resp) {
+    const auto snapshot = SnapshotState();
+    HttpHealthResponse payload;
+    payload.status = "ok";
+    payload.role = ha::MasterRuntimeRoleToString(snapshot.state);
+    payload.ha_state = ha::MasterRuntimeStateToString(snapshot.state);
+    payload.service_ready = snapshot.service_available;
+    if (snapshot.leader_view.has_value()) {
+        payload.leader_address = snapshot.leader_view->leader_address;
+        payload.view_version = snapshot.leader_view->view_version;
+    }
+    WriteJsonResponse(resp, coro_http::status_type::ok, payload);
+}
+
+struct HttpLeaderResponse {
+    bool present{false};
+    std::optional<std::string> leader_address;
+    std::optional<uint64_t> view_version;
+};
+YLT_REFL(HttpLeaderResponse, present, leader_address, view_version);
+
+void MasterAdminServer::HandleLeader(coro_http::coro_http_request&,
+                                     coro_http::coro_http_response& resp) {
+    const auto snapshot = SnapshotState();
+    HttpLeaderResponse payload;
+    payload.present = snapshot.leader_view.has_value();
+    if (snapshot.leader_view.has_value()) {
+        payload.leader_address = snapshot.leader_view->leader_address;
+        payload.view_version = snapshot.leader_view->view_version;
+    }
+    WriteJsonResponse(resp, coro_http::status_type::ok, payload);
+}
+
+std::shared_ptr<WrappedMasterService> MasterAdminServer::GetActiveService()
+    const {
+    const auto snapshot = SnapshotState();
+    if (!snapshot.service_available) {
+        return nullptr;
+    }
+    return snapshot.service;
+}
+
+template <typename Handler>
+void MasterAdminServer::WithActiveService(coro_http::coro_http_response& resp,
+                                          Handler&& handler) const {
+    auto service = GetActiveService();
+    if (!service) {
+        SetServiceUnavailable(resp, "service plane is not active");
+        return;
+    }
+    std::forward<Handler>(handler)(service);
+}
+
+void MasterAdminServer::HandleMetrics(coro_http::coro_http_request&,
+                                      coro_http::coro_http_response& resp) {
+    resp.add_header("Content-Type", "text/plain; version=0.0.4");
+    resp.set_status_and_content(coro_http::status_type::ok, BuildMetricsText());
+}
+
+void MasterAdminServer::HandleMetricsSummary(
+    coro_http::coro_http_request&, coro_http::coro_http_response& resp) {
+    resp.add_header("Content-Type", "text/plain; version=0.0.4");
+    resp.set_status_and_content(coro_http::status_type::ok,
+                                BuildMetricsSummaryText());
+}
+
+void MasterAdminServer::HandleRole(coro_http::coro_http_request&,
+                                   coro_http::coro_http_response& resp) {
+    const auto snapshot = SnapshotState();
+    resp.add_header("Content-Type", "text/plain; charset=utf-8");
+    resp.set_status_and_content(coro_http::status_type::ok,
+                                ha::MasterRuntimeRoleToString(snapshot.state));
+}
+
+void MasterAdminServer::HandleHaStatus(coro_http::coro_http_request&,
+                                       coro_http::coro_http_response& resp) {
+    const auto snapshot = SnapshotState();
+    resp.add_header("Content-Type", "text/plain; charset=utf-8");
+    resp.set_status_and_content(coro_http::status_type::ok,
+                                ha::MasterRuntimeStateToString(snapshot.state));
+}
+
+void MasterAdminServer::HandleQueryKey(coro_http::coro_http_request& req,
+                                       coro_http::coro_http_response& resp) {
+    WithActiveService(resp, [&](auto service) {
+        auto key = req.get_query_value("key");
+        auto get_result = service->GetReplicaList(std::string(key), "default");
+        resp.add_header("Content-Type", "application/json; charset=utf-8");
+        if (get_result) {
+            std::string ss = "{\"success\":true,\"data\":[";
+            const auto& replicas = get_result.value().replicas;
+            bool first = true;
+            for (const auto& replica : replicas) {
+                if (!replica.is_memory_replica()) {
+                    continue;
+                }
+                std::string tmp;
+                struct_json::to_json(
+                    replica.get_memory_descriptor().buffer_descriptor, tmp);
+                if (!first) {
+                    ss += ",";
+                }
+                ss += tmp;
+                first = false;
+            }
+            ss += "]}";
+            resp.set_status_and_content(coro_http::status_type::ok,
+                                        std::move(ss));
+            return;
+        }
+
+        resp.set_status_and_content(
+            coro_http::status_type::not_found,
+            std::string("{\"success\":false,\"error_code\":") +
+                std::to_string(toInt(get_result.error())) +
+                ",\"error_message\":\"" +
+                EscapeJson(toString(get_result.error())) + "\"}");
+    });
+}
+
+void MasterAdminServer::HandleGetAllKeys(coro_http::coro_http_request&,
+                                         coro_http::coro_http_response& resp) {
+    WithActiveService(resp, [&](auto service) {
+        resp.add_header("Content-Type", "text/plain; version=0.0.4");
+        auto result = service->GetAllKeysForAdmin();
+        if (!result) {
+            resp.set_status_and_content(
+                coro_http::status_type::internal_server_error,
+                "Failed to get all keys");
+            return;
+        }
+
+        std::string body;
+        for (const auto& key : result.value()) {
+            body += key;
+            body += "\n";
+        }
+        resp.set_status_and_content(coro_http::status_type::ok,
+                                    std::move(body));
+    });
+}
+
+void MasterAdminServer::HandleGetAllSegments(
+    coro_http::coro_http_request&, coro_http::coro_http_response& resp) {
+    WithActiveService(resp, [&](auto service) {
+        resp.add_header("Content-Type", "text/plain; version=0.0.4");
+        auto result = service->GetAllSegmentsForAdmin();
+        if (!result) {
+            resp.set_status_and_content(
+                coro_http::status_type::internal_server_error,
+                "Failed to get all segments");
+            return;
+        }
+
+        std::string body;
+        for (const auto& segment_name : result.value()) {
+            body += segment_name;
+            body += "\n";
+        }
+        resp.set_status_and_content(coro_http::status_type::ok,
+                                    std::move(body));
+    });
+}
+
+struct HttpSegmentDetailItem {
+    std::string segment_name;
+    std::string segment_id;
+    std::string client_id;
+    std::string base_address;
+    uint64_t size_bytes{0};
+    std::string size_human;
+    std::string te_endpoint;
+    std::string protocol;
+    std::string status;
+    uint64_t allocator_used_bytes{0};
+    uint64_t allocator_capacity_bytes{0};
+    double allocator_usage_percent{0.0};
+};
+YLT_REFL(HttpSegmentDetailItem, segment_name, segment_id, client_id,
+         base_address, size_bytes, size_human, te_endpoint, protocol, status,
+         allocator_used_bytes, allocator_capacity_bytes,
+         allocator_usage_percent);
+
+struct HttpSegmentsDetailResponse {
+    uint64_t total_segments{0};
+    std::vector<HttpSegmentDetailItem> segments;
+};
+YLT_REFL(HttpSegmentsDetailResponse, total_segments, segments);
+
+void MasterAdminServer::HandleGetSegmentsDetail(
+    coro_http::coro_http_request&, coro_http::coro_http_response& resp) {
+    WithActiveService(resp, [&](auto service) {
+        auto result = service->GetSegmentsDetailForAdmin();
+        if (!result) {
+            WriteErrorResponse(resp,
+                               coro_http::status_type::internal_server_error,
+                               result.error(), "Failed to get segments detail");
+            return;
+        }
+
+        HttpSegmentsDetailResponse payload;
+        payload.total_segments = result.value().size();
+        for (const auto& info : result.value()) {
+            HttpSegmentDetailItem item;
+            item.segment_name = info.segment_name;
+            item.segment_id = UuidToString(info.segment_id);
+            item.client_id = UuidToString(info.client_id);
+
+            std::ostringstream addr_oss;
+            addr_oss << "0x" << std::hex << info.base_address;
+            item.base_address = addr_oss.str();
+
+            item.size_bytes = info.size_bytes;
+            std::ostringstream size_oss;
+            size_oss << (info.size_bytes / 1024.0 / 1024.0 / 1024.0) << " GiB";
+            item.size_human = size_oss.str();
+
+            item.te_endpoint = info.te_endpoint;
+            item.protocol = info.protocol;
+            item.status = EnumToString(info.status);
+            item.allocator_used_bytes = info.allocator_used_bytes;
+            item.allocator_capacity_bytes = info.allocator_capacity_bytes;
+            item.allocator_usage_percent =
+                info.allocator_capacity_bytes > 0
+                    ? (static_cast<double>(info.allocator_used_bytes) /
+                       info.allocator_capacity_bytes * 100.0)
+                    : 0.0;
+            payload.segments.push_back(std::move(item));
+        }
+        WriteJsonResponse(resp, coro_http::status_type::ok, payload);
+    });
+}
+
+void MasterAdminServer::HandleQuerySegment(
+    coro_http::coro_http_request& req, coro_http::coro_http_response& resp) {
+    WithActiveService(resp, [&](auto service) {
+        auto segment = req.get_query_value("segment");
+        resp.add_header("Content-Type", "text/plain; version=0.0.4");
+        auto result = service->QuerySegmentForAdmin(std::string(segment));
+        if (!result) {
+            resp.set_status_and_content(
+                coro_http::status_type::internal_server_error,
+                "Failed to query segment");
+            return;
+        }
+
+        const auto [used, capacity] = result.value();
+        std::string body(segment);
+        body += "\nUsed(bytes): ";
+        body += std::to_string(used);
+        body += "\nCapacity(bytes) : ";
+        body += std::to_string(capacity);
+        body += "\n";
+        resp.set_status_and_content(coro_http::status_type::ok,
+                                    std::move(body));
+    });
+}
+
+struct HttpCreateDrainJobResponse {
+    bool success{false};
+    std::string job_id;
+    std::string status;
+    int32_t error_code{0};
+    std::string error_message;
+};
+YLT_REFL(HttpCreateDrainJobResponse, success, job_id, status, error_code,
+         error_message);
+
+void MasterAdminServer::HandleCreateDrainJob(
+    coro_http::coro_http_request& req, coro_http::coro_http_response& resp) {
+    CreateDrainJobRequest request;
+    try {
+        struct_json::from_json(request, req.get_body());
+    } catch (const std::exception& e) {
+        WriteErrorResponse(resp, coro_http::status_type::bad_request,
+                           ErrorCode::INVALID_PARAMS,
+                           std::string("Invalid JSON body: ") + e.what());
+        return;
+    }
+
+    WithActiveService(resp, [&](auto service) {
+        auto result = service->CreateDrainJob(request);
+        if (!result.has_value()) {
+            WriteErrorResponse(resp, ErrorCodeToHttpStatus(result.error()),
+                               result.error());
+            return;
+        }
+
+        HttpCreateDrainJobResponse payload;
+        payload.success = true;
+        payload.job_id = UuidToString(result.value());
+        payload.status = "CREATED";
+        WriteJsonResponse(resp, coro_http::status_type::ok, payload);
+    });
+}
+
+struct HttpQueryDrainJobResponse {
+    bool success{false};
+    std::string job_id;
+    int32_t type{0};
+    std::string type_name;
+    int32_t status{0};
+    std::string status_name;
+    int64_t created_at_ms_epoch{0};
+    int64_t last_updated_at_ms_epoch{0};
+    std::vector<std::string> segments;
+    uint64_t succeeded_units{0};
+    uint64_t failed_units{0};
+    uint64_t blocked_units{0};
+    uint64_t active_units{0};
+    uint64_t migrated_bytes{0};
+    std::string message;
+    int32_t error_code{0};
+    std::string error_message;
+};
+YLT_REFL(HttpQueryDrainJobResponse, success, job_id, type, type_name, status,
+         status_name, created_at_ms_epoch, last_updated_at_ms_epoch, segments,
+         succeeded_units, failed_units, blocked_units, active_units,
+         migrated_bytes, message, error_code, error_message);
+
+namespace {
+
+HttpQueryDrainJobResponse ToHttpQueryDrainJobResponse(
+    const QueryJobResponse& job) {
+    HttpQueryDrainJobResponse payload;
+    payload.success = true;
+    payload.job_id = UuidToString(job.id);
+    payload.type = static_cast<int32_t>(job.type);
+    payload.type_name = EnumToString(job.type);
+    payload.status = static_cast<int32_t>(job.status);
+    payload.status_name = EnumToString(job.status);
+    payload.created_at_ms_epoch = job.created_at_ms_epoch;
+    payload.last_updated_at_ms_epoch = job.last_updated_at_ms_epoch;
+    payload.segments = job.segments;
+    payload.succeeded_units = job.succeeded_units;
+    payload.failed_units = job.failed_units;
+    payload.blocked_units = job.blocked_units;
+    payload.active_units = job.active_units;
+    payload.migrated_bytes = job.migrated_bytes;
+    payload.message = job.message;
+    return payload;
+}
+
+}  // namespace
+
+void MasterAdminServer::HandleQueryDrainJob(
+    coro_http::coro_http_request& req, coro_http::coro_http_response& resp) {
+    auto job_id_result = ParseJobId(req.get_decode_query_value("job_id"));
+    if (!job_id_result.has_value()) {
+        WriteErrorResponse(resp, coro_http::status_type::bad_request,
+                           job_id_result.error(), "Missing or invalid job_id");
+        return;
+    }
+
+    WithActiveService(resp, [&](auto service) {
+        auto result = service->QueryDrainJob(job_id_result.value());
+        if (!result.has_value()) {
+            WriteErrorResponse(resp, ErrorCodeToHttpStatus(result.error()),
+                               result.error());
+            return;
+        }
+
+        WriteJsonResponse(resp, coro_http::status_type::ok,
+                          ToHttpQueryDrainJobResponse(result.value()));
+    });
+}
+
+struct HttpCancelDrainJobResponse {
+    bool success{false};
+    std::string job_id;
+    std::string status;
+    int32_t error_code{0};
+    std::string error_message;
+};
+YLT_REFL(HttpCancelDrainJobResponse, success, job_id, status, error_code,
+         error_message);
+
+void MasterAdminServer::HandleCancelDrainJob(
+    coro_http::coro_http_request& req, coro_http::coro_http_response& resp) {
+    auto job_id_result = ParseJobId(req.get_decode_query_value("job_id"));
+    if (!job_id_result.has_value()) {
+        WriteErrorResponse(resp, coro_http::status_type::bad_request,
+                           job_id_result.error(), "Missing or invalid job_id");
+        return;
+    }
+
+    WithActiveService(resp, [&](auto service) {
+        auto result = service->CancelDrainJob(job_id_result.value());
+        if (!result.has_value()) {
+            WriteErrorResponse(resp, ErrorCodeToHttpStatus(result.error()),
+                               result.error());
+            return;
+        }
+
+        HttpCancelDrainJobResponse payload;
+        payload.success = true;
+        payload.job_id = UuidToString(job_id_result.value());
+        payload.status = "CANCELED";
+        WriteJsonResponse(resp, coro_http::status_type::ok, payload);
+    });
+}
+
+struct HttpSegmentStatusResponse {
+    bool success{false};
+    std::string segment;
+    int32_t status{0};
+    std::string status_name;
+    int32_t error_code{0};
+    std::string error_message;
+};
+YLT_REFL(HttpSegmentStatusResponse, success, segment, status, status_name,
+         error_code, error_message);
+
+void MasterAdminServer::HandleSegmentStatus(
+    coro_http::coro_http_request& req, coro_http::coro_http_response& resp) {
+    auto segment_name = req.get_decode_query_value("segment");
+    if (segment_name.empty()) {
+        WriteErrorResponse(resp, coro_http::status_type::bad_request,
+                           ErrorCode::INVALID_PARAMS,
+                           "Missing segment query parameter");
+        return;
+    }
+
+    WithActiveService(resp, [&](auto service) {
+        auto result = service->QuerySegmentStatus(std::string(segment_name));
+        if (!result.has_value()) {
+            WriteErrorResponse(resp, ErrorCodeToHttpStatus(result.error()),
+                               result.error());
+            return;
+        }
+
+        HttpSegmentStatusResponse payload;
+        payload.success = true;
+        payload.segment = std::string(segment_name);
+        payload.status = static_cast<int32_t>(result.value());
+        payload.status_name = EnumToString(result.value());
+        WriteJsonResponse(resp, coro_http::status_type::ok, payload);
+    });
+}
+
+struct HttpBatchQueryKeyResult {
+    bool ok{false};
+    std::optional<std::string> error;
+    std::optional<std::vector<AllocatedBuffer::Descriptor>> values;
+};
+YLT_REFL(HttpBatchQueryKeyResult, ok, error, values);
+
+struct HttpBatchQueryKeysResponse {
+    bool success{false};
+    std::map<std::string, HttpBatchQueryKeyResult> data;
+};
+YLT_REFL(HttpBatchQueryKeysResponse, success, data);
+
+void MasterAdminServer::HandleBatchQueryKeys(
+    coro_http::coro_http_request& req, coro_http::coro_http_response& resp) {
+    auto service = GetActiveService();
+    if (!service) {
+        WriteSimpleErrorResponse(resp,
+                                 coro_http::status_type::service_unavailable,
+                                 "service plane is not active");
+        return;
+    }
+
+    auto keys_str = req.get_query_value("keys");
+    std::vector<std::string> keys;
+    if (!keys_str.empty()) {
+        std::string_view sv(keys_str);
+        size_t pos = 0;
+        while ((pos = sv.find(',')) != std::string_view::npos) {
+            keys.emplace_back(sv.substr(0, pos));
+            sv.remove_prefix(pos + 1);
+        }
+        keys.emplace_back(sv);
+    }
+
+    if (keys.empty()) {
+        WriteSimpleErrorResponse(resp, coro_http::status_type::bad_request,
+                                 "No keys provided. Use ?keys=key1,key2,...");
+        return;
+    }
+
+    auto results = service->BatchGetReplicaList(keys, "default");
+    const size_t n = std::min(keys.size(), results.size());
+    HttpBatchQueryKeysResponse payload;
+    payload.success = true;
+
+    for (size_t i = 0; i < n; ++i) {
+        const auto& result = results[i];
+        HttpBatchQueryKeyResult item;
+        if (!result.has_value()) {
+            item.error = toString(result.error());
+            payload.data.emplace(keys[i], std::move(item));
+            continue;
+        }
+
+        item.ok = true;
+        item.values = std::vector<AllocatedBuffer::Descriptor>{};
+        for (const auto& replica : result.value().replicas) {
+            if (!replica.is_memory_replica()) {
+                continue;
+            }
+            item.values->emplace_back(
+                replica.get_memory_descriptor().buffer_descriptor);
+        }
+        payload.data.emplace(keys[i], std::move(item));
+    }
+
+    if (results.size() != keys.size()) {
+        LOG(WARNING) << "BatchGetReplicaList size mismatch: keys="
+                     << keys.size() << " results=" << results.size();
+    }
+    WriteJsonResponse(resp, coro_http::status_type::ok, payload);
+}
+
+void MasterAdminServer::HandleGetTenantQuotas(
+    coro_http::coro_http_request& req, coro_http::coro_http_response& resp) {
+    auto tenant_id_view = req.get_decode_query_value("tenant_id");
+    WithActiveService(resp, [&](auto service) {
+        if (tenant_id_view.empty()) {
+            auto result = service->ListTenantQuotaSnapshots();
+            if (!result.has_value()) {
+                WriteErrorResponse(resp, ErrorCodeToHttpStatus(result.error()),
+                                   result.error());
+                return;
+            }
+            HttpTenantQuotaListResponse payload;
+            payload.data.reserve(result->size());
+            for (const auto& snapshot : result.value()) {
+                payload.data.push_back(ToHttpTenantQuotaSnapshot(snapshot));
+            }
+            WriteJsonResponse(resp, coro_http::status_type::ok, payload);
+            return;
+        }
+
+        auto tenant_id_result = ParseAdminTenantId(req);
+        if (!tenant_id_result.has_value()) {
+            WriteErrorResponse(resp, coro_http::status_type::bad_request,
+                               ErrorCode::INVALID_PARAMS, "Invalid tenant_id");
+            return;
+        }
+        auto result = service->GetTenantQuotaSnapshot(tenant_id_result.value());
+        if (!result.has_value()) {
+            WriteErrorResponse(resp, ErrorCodeToHttpStatus(result.error()),
+                               result.error());
+            return;
+        }
+        WriteJsonResponse(
+            resp, coro_http::status_type::ok,
+            HttpTenantQuotaResponse{
+                .data = ToHttpTenantQuotaSnapshot(result.value())});
+    });
+}
+
+void MasterAdminServer::HandleUpsertTenantQuota(
+    coro_http::coro_http_request& req, coro_http::coro_http_response& resp) {
+    auto tenant_id_result = ParseAdminTenantId(req);
+    if (!tenant_id_result.has_value()) {
+        WriteErrorResponse(resp, coro_http::status_type::bad_request,
+                           tenant_id_result.error(),
+                           "Missing or invalid tenant_id");
+        return;
+    }
+    auto body_result = ParseQuotaPolicyBody(req);
+    if (!body_result.has_value()) {
+        WriteErrorResponse(resp, coro_http::status_type::bad_request,
+                           ErrorCode::INVALID_PARAMS, body_result.error());
+        return;
+    }
+    if (body_result->requested_quota_bytes == 0) {
+        WriteErrorResponse(resp, coro_http::status_type::bad_request,
+                           ErrorCode::INVALID_PARAMS,
+                           "Tenant quota must be positive");
+        return;
+    }
+
+    WithActiveService(resp, [&](auto service) {
+        auto result = service->UpsertTenantQuotaPolicy(
+            tenant_id_result.value(), body_result->requested_quota_bytes);
+        if (!result.has_value()) {
+            WriteErrorResponse(resp, ErrorCodeToHttpStatus(result.error()),
+                               result.error());
+            return;
+        }
+        WriteJsonResponse(
+            resp, coro_http::status_type::ok,
+            HttpTenantQuotaResponse{
+                .data = ToHttpTenantQuotaSnapshot(result.value())});
+    });
+}
+
+void MasterAdminServer::HandleDeleteTenantQuota(
+    coro_http::coro_http_request& req, coro_http::coro_http_response& resp) {
+    auto tenant_id_result = ParseAdminTenantId(req);
+    if (!tenant_id_result.has_value()) {
+        WriteErrorResponse(resp, coro_http::status_type::bad_request,
+                           tenant_id_result.error(),
+                           "Missing or invalid tenant_id");
+        return;
+    }
+
+    WithActiveService(resp, [&](auto service) {
+        auto result =
+            service->DeleteTenantQuotaPolicy(tenant_id_result.value());
+        if (!result.has_value()) {
+            WriteErrorResponse(resp, ErrorCodeToHttpStatus(result.error()),
+                               result.error());
+            return;
+        }
+        HttpTenantQuotaDeleteResponse payload;
+        if (result.value().has_value()) {
+            payload.data = ToHttpTenantQuotaSnapshot(result.value().value());
+        }
+        WriteJsonResponse(resp, coro_http::status_type::ok, payload);
+    });
+}
+
+void MasterAdminServer::RegisterHandler() {
+    using namespace coro_http;
+
+    http_server_.set_http_handler<GET>(
+        "/metrics", [this](coro_http_request& req, coro_http_response& resp) {
+            HandleMetrics(req, resp);
+        });
+    http_server_.set_http_handler<GET>(
+        "/metrics/summary",
+        [this](coro_http_request& req, coro_http_response& resp) {
+            HandleMetricsSummary(req, resp);
+        });
+    http_server_.set_http_handler<GET>(
+        "/health", [this](coro_http_request& req, coro_http_response& resp) {
+            HandleHealth(req, resp);
+        });
+    http_server_.set_http_handler<GET>(
+        "/role", [this](coro_http_request& req, coro_http_response& resp) {
+            HandleRole(req, resp);
+        });
+    http_server_.set_http_handler<GET>(
+        "/ha_status", [this](coro_http_request& req, coro_http_response& resp) {
+            HandleHaStatus(req, resp);
+        });
+    http_server_.set_http_handler<GET>(
+        "/leader", [this](coro_http_request& req, coro_http_response& resp) {
+            HandleLeader(req, resp);
+        });
+    http_server_.set_http_handler<GET>(
+        "/query_key", [this](coro_http_request& req, coro_http_response& resp) {
+            HandleQueryKey(req, resp);
+        });
+    http_server_.set_http_handler<GET>(
+        "/get_all_keys",
+        [this](coro_http_request& req, coro_http_response& resp) {
+            HandleGetAllKeys(req, resp);
+        });
+    http_server_.set_http_handler<GET>(
+        "/get_all_segments",
+        [this](coro_http_request& req, coro_http_response& resp) {
+            HandleGetAllSegments(req, resp);
+        });
+    http_server_.set_http_handler<GET>(
+        "/get_segments_detail",
+        [this](coro_http_request& req, coro_http_response& resp) {
+            HandleGetSegmentsDetail(req, resp);
+        });
+
+    http_server_.set_http_handler<GET>(
+        "/query_segment",
+        [this](coro_http_request& req, coro_http_response& resp) {
+            HandleQuerySegment(req, resp);
+        });
+    http_server_.set_http_handler<POST>(
+        "/api/v1/drain_jobs",
+        [this](coro_http_request& req, coro_http_response& resp) {
+            HandleCreateDrainJob(req, resp);
+        });
+    http_server_.set_http_handler<GET>(
+        "/api/v1/drain_jobs/query",
+        [this](coro_http_request& req, coro_http_response& resp) {
+            HandleQueryDrainJob(req, resp);
+        });
+    http_server_.set_http_handler<POST>(
+        "/api/v1/drain_jobs/cancel",
+        [this](coro_http_request& req, coro_http_response& resp) {
+            HandleCancelDrainJob(req, resp);
+        });
+    http_server_.set_http_handler<GET>(
+        "/api/v1/segments/status",
+        [this](coro_http_request& req, coro_http_response& resp) {
+            HandleSegmentStatus(req, resp);
+        });
+    http_server_.set_http_handler<GET>(
+        "/api/v1/tenant_quotas",
+        [this](coro_http_request& req, coro_http_response& resp) {
+            HandleGetTenantQuotas(req, resp);
+        });
+    http_server_.set_http_handler<PUT>(
+        "/api/v1/tenant_quotas",
+        [this](coro_http_request& req, coro_http_response& resp) {
+            HandleUpsertTenantQuota(req, resp);
+        });
+    http_server_.set_http_handler<DEL>(
+        "/api/v1/tenant_quotas",
+        [this](coro_http_request& req, coro_http_response& resp) {
+            HandleDeleteTenantQuota(req, resp);
+        });
+    http_server_.set_http_handler<GET>(
+        "/batch_query_keys",
+        [this](coro_http_request& req, coro_http_response& resp) {
+            HandleBatchQueryKeys(req, resp);
+        });
+}
+}  // namespace mooncake
diff --git a/mooncake-store/src/master_client.cpp b/mooncake-store/src/master_client.cpp
index 44f562cc..e2d9db34 100644
--- a/mooncake-store/src/master_client.cpp
+++ b/mooncake-store/src/master_client.cpp
@@ -147,16 +147,46 @@ struct RpcNameTraits<&WrappedMasterService::MountSegment> {
     static constexpr const char* value = "MountSegment";
 };
 
+template <>
+struct RpcNameTraits<&WrappedMasterService::MountNoFSegment> {
+    static constexpr const char* value = "MountNoFSegment";
+};
+
 template <>
 struct RpcNameTraits<&WrappedMasterService::ReMountSegment> {
     static constexpr const char* value = "ReMountSegment";
 };
 
+template <>
+struct RpcNameTraits<&WrappedMasterService::ReMountNoFSegment> {
+    static constexpr const char* value = "ReMountNoFSegment";
+};
+
 template <>
 struct RpcNameTraits<&WrappedMasterService::UnmountSegment> {
     static constexpr const char* value = "UnmountSegment";
 };
 
+template <>
+struct RpcNameTraits<&WrappedMasterService::GracefulUnmountSegment> {
+    static constexpr const char* value = "GracefulUnmountSegment";
+};
+
+template <>
+struct RpcNameTraits<&WrappedMasterService::UnmountNoFSegment> {
+    static constexpr const char* value = "UnmountNoFSegment";
+};
+
+template <>
+struct RpcNameTraits<&WrappedMasterService::GetAllNoFSegments> {
+    static constexpr const char* value = "GetAllNoFSegments";
+};
+
+template <>
+struct RpcNameTraits<&WrappedMasterService::GetNoFSegmentsByName> {
+    static constexpr const char* value = "GetNoFSegmentsByName";
+};
+
 template <>
 struct RpcNameTraits<&WrappedMasterService::Ping> {
     static constexpr const char* value = "Ping";
@@ -167,6 +197,11 @@ struct RpcNameTraits<&WrappedMasterService::GetFsdir> {
     static constexpr const char* value = "GetFsdir";
 };
 
+template <>
+struct RpcNameTraits<&WrappedMasterService::QuerySegmentStatusById> {
+    static constexpr const char* value = "QuerySegmentStatusById";
+};
+
 template <>
 struct RpcNameTraits<&WrappedMasterService::GetStorageConfig> {
     static constexpr const char* value = "GetStorageConfig";
@@ -187,11 +222,36 @@ struct RpcNameTraits<&WrappedMasterService::OffloadObjectHeartbeat> {
     static constexpr const char* value = "OffloadObjectHeartbeat";
 };
 
+template <>
+struct RpcNameTraits<&WrappedMasterService::ReportSsdCapacity> {
+    static constexpr const char* value = "ReportSsdCapacity";
+};
+
 template <>
 struct RpcNameTraits<&WrappedMasterService::NotifyOffloadSuccess> {
     static constexpr const char* value = "NotifyOffloadSuccess";
 };
 
+template <>
+struct RpcNameTraits<&WrappedMasterService::PromotionObjectHeartbeat> {
+    static constexpr const char* value = "PromotionObjectHeartbeat";
+};
+
+template <>
+struct RpcNameTraits<&WrappedMasterService::PromotionAllocStart> {
+    static constexpr const char* value = "PromotionAllocStart";
+};
+
+template <>
+struct RpcNameTraits<&WrappedMasterService::NotifyPromotionSuccess> {
+    static constexpr const char* value = "NotifyPromotionSuccess";
+};
+
+template <>
+struct RpcNameTraits<&WrappedMasterService::NotifyPromotionFailure> {
+    static constexpr const char* value = "NotifyPromotionFailure";
+};
+
 template <>
 struct RpcNameTraits<&WrappedMasterService::CopyStart> {
     static constexpr const char* value = "CopyStart";
@@ -281,6 +341,10 @@ tl::expected<ReturnType, ErrorCode> MasterClient::invoke_rpc(Args&&... args) {
             }
             auto result = co_await std::move(ret.value());
             if (!result) {
+                if (result.error().code == coro_rpc::errc::timed_out) {
+                    LOG(ERROR) << "RPC call timed out: " << result.error().msg;
+                    co_return tl::make_unexpected(ErrorCode::RPC_TIMEOUT);
+                }
                 LOG(ERROR) << "RPC call failed: " << result.error().msg;
                 co_return tl::make_unexpected(ErrorCode::RPC_FAIL);
             }
@@ -323,12 +387,15 @@ std::vector<tl::expected<ResultType, ErrorCode>> MasterClient::invoke_batch_rpc(
             }
             auto result = co_await std::move(ret.value());
             if (!result) {
+                const ErrorCode err_code =
+                    (result.error().code == coro_rpc::errc::timed_out)
+                        ? ErrorCode::RPC_TIMEOUT
+                        : ErrorCode::RPC_FAIL;
                 LOG(ERROR) << "Batch RPC call failed: " << result.error().msg;
                 std::vector<tl::expected<ResultType, ErrorCode>> error_results;
                 error_results.reserve(input_size);
                 for (size_t i = 0; i < input_size; ++i) {
-                    error_results.emplace_back(
-                        tl::make_unexpected(ErrorCode::RPC_FAIL));
+                    error_results.emplace_back(tl::make_unexpected(err_code));
                 }
                 co_return error_results;
             }
@@ -385,7 +452,8 @@ tl::expected<bool, ErrorCode> MasterClient::ExistKey(
     ScopedVLogTimer timer(1, "MasterClient::ExistKey");
     timer.LogRequest("object_key=", object_key);
 
-    auto result = invoke_rpc<&WrappedMasterService::ExistKey, bool>(object_key);
+    auto result = invoke_rpc<&WrappedMasterService::ExistKey, bool>(object_key,
+                                                                    tenant_id_);
     timer.LogResponseExpected(result);
     return result;
 }
@@ -396,7 +464,7 @@ std::vector<tl::expected<bool, ErrorCode>> MasterClient::BatchExistKey(
     timer.LogRequest("keys_count=", object_keys.size());
 
     auto result = invoke_batch_rpc<&WrappedMasterService::BatchExistKey, bool>(
-        object_keys.size(), object_keys);
+        object_keys.size(), object_keys, tenant_id_);
     timer.LogResponse("result=", result.size(), " keys");
     return result;
 }
@@ -446,7 +514,8 @@ MasterClient::GetReplicaListByRegex(const std::string& str) {
 
     auto result = invoke_rpc<
         &WrappedMasterService::GetReplicaListByRegex,
-        std::unordered_map<std::string, std::vector<Replica::Descriptor>>>(str);
+        std::unordered_map<std::string, std::vector<Replica::Descriptor>>>(
+        str, tenant_id_);
 
     timer.LogResponseExpected(result);
     return result;
@@ -454,23 +523,35 @@ MasterClient::GetReplicaListByRegex(const std::string& str) {
 
 tl::expected<GetReplicaListResponse, ErrorCode> MasterClient::GetReplicaList(
     const std::string& object_key) {
+    return GetReplicaList(object_key, tenant_id_);
+}
+
+tl::expected<GetReplicaListResponse, ErrorCode> MasterClient::GetReplicaList(
+    const std::string& object_key, const std::string& tenant_id) {
     ScopedVLogTimer timer(1, "MasterClient::GetReplicaList");
-    timer.LogRequest("object_key=", object_key);
+    timer.LogRequest("object_key=", object_key, ", tenant_id=", tenant_id);
 
     auto result = invoke_rpc<&WrappedMasterService::GetReplicaList,
-                             GetReplicaListResponse>(object_key);
+                             GetReplicaListResponse>(object_key, tenant_id);
     timer.LogResponseExpected(result);
     return result;
 }
 
 std::vector<tl::expected<GetReplicaListResponse, ErrorCode>>
 MasterClient::BatchGetReplicaList(const std::vector<std::string>& object_keys) {
+    return BatchGetReplicaList(object_keys, tenant_id_);
+}
+
+std::vector<tl::expected<GetReplicaListResponse, ErrorCode>>
+MasterClient::BatchGetReplicaList(const std::vector<std::string>& object_keys,
+                                  const std::string& tenant_id) {
     ScopedVLogTimer timer(1, "MasterClient::BatchGetReplicaList");
-    timer.LogRequest("keys_count=", object_keys.size());
+    timer.LogRequest("keys_count=", object_keys.size(),
+                     ", tenant_id=", tenant_id);
 
     auto result = invoke_batch_rpc<&WrappedMasterService::BatchGetReplicaList,
-                                   GetReplicaListResponse>(object_keys.size(),
-                                                           object_keys);
+                                   GetReplicaListResponse>(
+        object_keys.size(), object_keys, tenant_id);
     timer.LogResponse("result=", result.size(), " operations");
     return result;
 }
@@ -489,7 +570,7 @@ MasterClient::PutStart(const std::string& key,
 
     auto result = invoke_rpc<&WrappedMasterService::PutStart,
                              std::vector<Replica::Descriptor>>(
-        client_id_, key, total_slice_length, config);
+        client_id_, key, total_slice_length, config, tenant_id_);
     timer.LogResponseExpected(result);
     return result;
 }
@@ -514,7 +595,7 @@ MasterClient::BatchPutStart(
 
     auto result = invoke_batch_rpc<&WrappedMasterService::BatchPutStart,
                                    std::vector<Replica::Descriptor>>(
-        keys.size(), client_id_, keys, total_slice_lengths, config);
+        keys.size(), client_id_, keys, total_slice_lengths, config, tenant_id_);
     timer.LogResponse("result=", result.size(), " operations");
     return result;
 }
@@ -525,18 +606,18 @@ tl::expected<void, ErrorCode> MasterClient::PutEnd(const std::string& key,
     timer.LogRequest("key=", key);
 
     auto result = invoke_rpc<&WrappedMasterService::PutEnd, void>(
-        client_id_, key, replica_type);
+        client_id_, key, replica_type, tenant_id_);
     timer.LogResponseExpected(result);
     return result;
 }
 
 std::vector<tl::expected<void, ErrorCode>> MasterClient::BatchPutEnd(
-    const std::vector<std::string>& keys) {
+    const std::vector<std::string>& keys, ReplicaType replica_type) {
     ScopedVLogTimer timer(1, "MasterClient::BatchPutEnd");
     timer.LogRequest("keys_count=", keys.size());
 
     auto result = invoke_batch_rpc<&WrappedMasterService::BatchPutEnd, void>(
-        keys.size(), client_id_, keys);
+        keys.size(), client_id_, keys, replica_type, tenant_id_);
     timer.LogResponse("result=", result.size(), " operations");
     return result;
 }
@@ -547,18 +628,18 @@ tl::expected<void, ErrorCode> MasterClient::PutRevoke(
     timer.LogRequest("key=", key);
 
     auto result = invoke_rpc<&WrappedMasterService::PutRevoke, void>(
-        client_id_, key, replica_type);
+        client_id_, key, replica_type, tenant_id_);
     timer.LogResponseExpected(result);
     return result;
 }
 
 std::vector<tl::expected<void, ErrorCode>> MasterClient::BatchPutRevoke(
-    const std::vector<std::string>& keys) {
+    const std::vector<std::string>& keys, ReplicaType replica_type) {
     ScopedVLogTimer timer(1, "MasterClient::BatchPutRevoke");
     timer.LogRequest("keys_count=", keys.size());
 
     auto result = invoke_batch_rpc<&WrappedMasterService::BatchPutRevoke, void>(
-        keys.size(), client_id_, keys);
+        keys.size(), client_id_, keys, replica_type, tenant_id_);
     timer.LogResponse("result=", result.size(), " operations");
     return result;
 }
@@ -577,7 +658,7 @@ MasterClient::UpsertStart(const std::string& key,
 
     auto result = invoke_rpc<&WrappedMasterService::UpsertStart,
                              std::vector<Replica::Descriptor>>(
-        client_id_, key, total_slice_length, config);
+        client_id_, key, total_slice_length, config, tenant_id_);
     timer.LogResponseExpected(result);
     return result;
 }
@@ -602,7 +683,7 @@ MasterClient::BatchUpsertStart(
 
     auto result = invoke_batch_rpc<&WrappedMasterService::BatchUpsertStart,
                                    std::vector<Replica::Descriptor>>(
-        keys.size(), client_id_, keys, total_slice_lengths, config);
+        keys.size(), client_id_, keys, total_slice_lengths, config, tenant_id_);
     timer.LogResponse("result=", result.size(), " operations");
     return result;
 }
@@ -613,7 +694,7 @@ tl::expected<void, ErrorCode> MasterClient::UpsertEnd(
     timer.LogRequest("key=", key);
 
     auto result = invoke_rpc<&WrappedMasterService::UpsertEnd, void>(
-        client_id_, key, replica_type);
+        client_id_, key, replica_type, tenant_id_);
     timer.LogResponseExpected(result);
     return result;
 }
@@ -624,7 +705,7 @@ std::vector<tl::expected<void, ErrorCode>> MasterClient::BatchUpsertEnd(
     timer.LogRequest("keys_count=", keys.size());
 
     auto result = invoke_batch_rpc<&WrappedMasterService::BatchUpsertEnd, void>(
-        keys.size(), client_id_, keys);
+        keys.size(), client_id_, keys, tenant_id_);
     timer.LogResponse("result=", result.size(), " operations");
     return result;
 }
@@ -635,7 +716,7 @@ tl::expected<void, ErrorCode> MasterClient::UpsertRevoke(
     timer.LogRequest("key=", key);
 
     auto result = invoke_rpc<&WrappedMasterService::UpsertRevoke, void>(
-        client_id_, key, replica_type);
+        client_id_, key, replica_type, tenant_id_);
     timer.LogResponseExpected(result);
     return result;
 }
@@ -647,7 +728,7 @@ std::vector<tl::expected<void, ErrorCode>> MasterClient::BatchUpsertRevoke(
 
     auto result =
         invoke_batch_rpc<&WrappedMasterService::BatchUpsertRevoke, void>(
-            keys.size(), client_id_, keys);
+            keys.size(), client_id_, keys, tenant_id_);
     timer.LogResponse("result=", result.size(), " operations");
     return result;
 }
@@ -657,7 +738,8 @@ tl::expected<void, ErrorCode> MasterClient::Remove(const std::string& key,
     ScopedVLogTimer timer(1, "MasterClient::Remove");
     timer.LogRequest("key=", key, ", force=", force);
 
-    auto result = invoke_rpc<&WrappedMasterService::Remove, void>(key, force);
+    auto result =
+        invoke_rpc<&WrappedMasterService::Remove, void>(key, force, tenant_id_);
     timer.LogResponseExpected(result);
     return result;
 }
@@ -667,8 +749,8 @@ tl::expected<long, ErrorCode> MasterClient::RemoveByRegex(
     ScopedVLogTimer timer(1, "MasterClient::RemoveByRegex");
     timer.LogRequest("key=", str, ", force=", force);
 
-    auto result =
-        invoke_rpc<&WrappedMasterService::RemoveByRegex, long>(str, force);
+    auto result = invoke_rpc<&WrappedMasterService::RemoveByRegex, long>(
+        str, force, tenant_id_);
     timer.LogResponseExpected(result);
     return result;
 }
@@ -677,7 +759,8 @@ tl::expected<long, ErrorCode> MasterClient::RemoveAll(bool force) {
     ScopedVLogTimer timer(1, "MasterClient::RemoveAll");
     timer.LogRequest("action=remove_all_objects, force=", force);
 
-    auto result = invoke_rpc<&WrappedMasterService::RemoveAll, long>(force);
+    auto result =
+        invoke_rpc<&WrappedMasterService::RemoveAll, long>(force, tenant_id_);
     timer.LogResponseExpected(result);
     return result;
 }
@@ -688,7 +771,7 @@ std::vector<tl::expected<void, ErrorCode>> MasterClient::BatchRemove(
     timer.LogRequest("keys_count=", keys.size(), ", force=", force);
 
     auto result = invoke_batch_rpc<&WrappedMasterService::BatchRemove, void>(
-        keys.size(), keys, force);
+        keys.size(), keys, force, tenant_id_);
     timer.LogResponse("result=", result.size(), " operations");
     return result;
 }
@@ -706,6 +789,19 @@ tl::expected<void, ErrorCode> MasterClient::MountSegment(
     return result;
 }
 
+tl::expected<void, ErrorCode> MasterClient::MountNoFSegment(
+    const NoFSegment& segment) {
+    ScopedVLogTimer timer(1, "MasterClient::MountNofSegment");
+    timer.LogRequest("NoF segment mount: ", "base=", segment.base,
+                     ", size=", segment.size, ", name=", segment.name,
+                     ", id=", segment.id, ", client_id=", client_id_);
+
+    auto result = invoke_rpc<&WrappedMasterService::MountNoFSegment, void>(
+        segment, client_id_);
+    timer.LogResponseExpected(result);
+    return result;
+}
+
 tl::expected<void, ErrorCode> MasterClient::ReMountSegment(
     const std::vector<Segment>& segments) {
     ScopedVLogTimer timer(1, "MasterClient::ReMountSegment");
@@ -718,6 +814,18 @@ tl::expected<void, ErrorCode> MasterClient::ReMountSegment(
     return result;
 }
 
+tl::expected<void, ErrorCode> MasterClient::ReMountNoFSegment(
+    const std::vector<NoFSegment>& segments) {
+    ScopedVLogTimer timer(1, "MasterClient::ReMountNofSegment");
+    timer.LogRequest("NoF segment remount: ", "segments_num=", segments.size(),
+                     ", client_id=", client_id_);
+
+    auto result = invoke_rpc<&WrappedMasterService::ReMountNoFSegment, void>(
+        segments, client_id_);
+    timer.LogResponseExpected(result);
+    return result;
+}
+
 tl::expected<void, ErrorCode> MasterClient::UnmountSegment(
     const UUID& segment_id) {
     ScopedVLogTimer timer(1, "MasterClient::UnmountSegment");
@@ -729,6 +837,53 @@ tl::expected<void, ErrorCode> MasterClient::UnmountSegment(
     return result;
 }
 
+tl::expected<void, ErrorCode> MasterClient::GracefulUnmountSegment(
+    const UUID& segment_id, uint64_t grace_period_ms) {
+    ScopedVLogTimer timer(1, "MasterClient::GracefulUnmountSegment");
+    timer.LogRequest("segment_id=", segment_id, ", client_id=", client_id_,
+                     ", grace_period_ms=", grace_period_ms);
+
+    auto result =
+        invoke_rpc<&WrappedMasterService::GracefulUnmountSegment, void>(
+            segment_id, client_id_, grace_period_ms);
+    timer.LogResponseExpected(result);
+    return result;
+}
+
+tl::expected<void, ErrorCode> MasterClient::UnmountNoFSegment(
+    const UUID& segment_id) {
+    ScopedVLogTimer timer(1, "MasterClient::UnmountNoFSegment");
+    timer.LogRequest("NoF segment unmount: ", "segment_id=", segment_id,
+                     ", client_id=", client_id_);
+
+    auto result = invoke_rpc<&WrappedMasterService::UnmountNoFSegment, void>(
+        segment_id, client_id_);
+    timer.LogResponseExpected(result);
+    return result;
+}
+
+tl::expected<std::vector<NoFSegment>, ErrorCode>
+MasterClient::GetAllNoFSegments() {
+    ScopedVLogTimer timer(1, "MasterClient::GetAllNoFSegments");
+    timer.LogRequest("Get all NoF segments, client_id=", client_id_);
+
+    auto result = invoke_rpc<&WrappedMasterService::GetAllNoFSegments,
+                             std::vector<NoFSegment>>();
+    timer.LogResponseExpected(result);
+    return result;
+}
+
+tl::expected<std::vector<NoFSegmentOwnerInfo>, ErrorCode>
+MasterClient::GetNoFSegmentsByName(const std::string& segment_name) {
+    ScopedVLogTimer timer(1, "MasterClient::GetNoFSegmentsByName");
+    timer.LogRequest("segment_name=", segment_name, ", client_id=", client_id_);
+
+    auto result = invoke_rpc<&WrappedMasterService::GetNoFSegmentsByName,
+                             std::vector<NoFSegmentOwnerInfo>>(segment_name);
+    timer.LogResponseExpected(result);
+    return result;
+}
+
 tl::expected<PingResponse, ErrorCode> MasterClient::Ping() {
     ScopedVLogTimer timer(1, "MasterClient::Ping");
     timer.LogRequest("client_id=", client_id_);
@@ -748,6 +903,17 @@ tl::expected<std::string, ErrorCode> MasterClient::GetFsdir() {
     return result;
 }
 
+tl::expected<SegmentStatus, ErrorCode> MasterClient::QuerySegmentStatusById(
+    const UUID& segment_id) {
+    ScopedVLogTimer timer(1, "MasterClient::QuerySegmentStatusById");
+    timer.LogRequest("segment_id=", segment_id);
+
+    auto result = invoke_rpc<&WrappedMasterService::QuerySegmentStatusById,
+                             SegmentStatus>(segment_id);
+    timer.LogResponseExpected(result);
+    return result;
+}
+
 tl::expected<GetStorageConfigResponse, ErrorCode>
 MasterClient::GetStorageConfig() {
     ScopedVLogTimer timer(1, "MasterClient::GetStorageConfig");
@@ -774,11 +940,18 @@ tl::expected<void, ErrorCode> MasterClient::MountLocalDiskSegment(
 
 tl::expected<UUID, ErrorCode> MasterClient::CreateCopyTask(
     const std::string& key, const std::vector<std::string>& targets) {
-    ScopedVLogTimer timer(1, "MasterClient::CreateCopyTask");
-    timer.LogRequest("key=", key, ", targets_size=", targets.size());
+    return CreateCopyTask(key, tenant_id_, targets);
+}
 
-    auto result =
-        invoke_rpc<&WrappedMasterService::CreateCopyTask, UUID>(key, targets);
+tl::expected<UUID, ErrorCode> MasterClient::CreateCopyTask(
+    const std::string& key, const std::string& tenant_id,
+    const std::vector<std::string>& targets) {
+    ScopedVLogTimer timer(1, "MasterClient::CreateCopyTask");
+    timer.LogRequest("key=", key, ", tenant_id=", tenant_id,
+                     ", targets_size=", targets.size());
+
+    auto result = invoke_rpc<&WrappedMasterService::CreateCopyTask, UUID>(
+        key, tenant_id, targets);
     timer.LogResponseExpected(result);
     return result;
 }
@@ -786,37 +959,132 @@ tl::expected<UUID, ErrorCode> MasterClient::CreateCopyTask(
 tl::expected<UUID, ErrorCode> MasterClient::CreateMoveTask(
     const std::string& key, const std::string& source,
     const std::string& target) {
+    return CreateMoveTask(key, tenant_id_, source, target);
+}
+
+tl::expected<UUID, ErrorCode> MasterClient::CreateMoveTask(
+    const std::string& key, const std::string& tenant_id,
+    const std::string& source, const std::string& target) {
     ScopedVLogTimer timer(1, "MasterClient::CreateMoveTask");
-    timer.LogRequest("key=", key, ", source=", source, ", target=", target);
+    timer.LogRequest("key=", key, ", tenant_id=", tenant_id,
+                     ", source=", source, ", target=", target);
 
     auto result = invoke_rpc<&WrappedMasterService::CreateMoveTask, UUID>(
-        key, source, target);
+        key, tenant_id, source, target);
     timer.LogResponseExpected(result);
     return result;
 }
 
-tl::expected<std::unordered_map<std::string, int64_t>, ErrorCode>
+tl::expected<std::vector<OffloadTaskItem>, ErrorCode>
 MasterClient::OffloadObjectHeartbeat(const UUID& client_id,
                                      bool enable_offloading) {
     ScopedVLogTimer timer(1, "MasterClient::OffloadObjectHeartbeat");
     timer.LogRequest("client_id=", client_id,
                      ", enable_offloading=", enable_offloading);
 
-    auto result = invoke_rpc<&WrappedMasterService::OffloadObjectHeartbeat,
-                             std::unordered_map<std::string, int64_t>>(
-        client_id, enable_offloading);
+    auto result =
+        invoke_rpc<&WrappedMasterService::OffloadObjectHeartbeat,
+                   std::vector<OffloadTaskItem>>(client_id, enable_offloading);
     return result;
 }
 
+tl::expected<void, ErrorCode> MasterClient::ReportSsdCapacity(
+    const UUID& client_id, int64_t ssd_total_capacity_bytes) {
+    ScopedVLogTimer timer(1, "MasterClient::ReportSsdCapacity");
+    timer.LogRequest("client_id=", client_id,
+                     ", ssd_total_capacity_bytes=", ssd_total_capacity_bytes);
+    return invoke_rpc<&WrappedMasterService::ReportSsdCapacity, void>(
+        client_id, ssd_total_capacity_bytes);
+}
+
 tl::expected<void, ErrorCode> MasterClient::NotifyOffloadSuccess(
     const UUID& client_id, const std::vector<std::string>& keys,
     const std::vector<StorageObjectMetadata>& metadatas) {
+    std::vector<OffloadTaskItem> tasks;
+    tasks.reserve(keys.size());
+    for (const auto& key : keys) {
+        tasks.push_back(
+            OffloadTaskItem{.tenant_id = tenant_id_, .key = key, .size = 0});
+    }
+    return NotifyOffloadSuccess(client_id, tasks, metadatas);
+}
+
+tl::expected<void, ErrorCode> MasterClient::NotifyOffloadSuccess(
+    const UUID& client_id, const std::vector<OffloadTaskItem>& tasks,
+    const std::vector<StorageObjectMetadata>& metadatas) {
     ScopedVLogTimer timer(1, "MasterClient::NotifyOffloadSuccess");
-    timer.LogRequest("client_id=", client_id, ", keys_count=", keys.size(),
+    timer.LogRequest("client_id=", client_id, ", tasks_count=", tasks.size(),
                      ", metadatas_count=", metadatas.size());
 
     auto result = invoke_rpc<&WrappedMasterService::NotifyOffloadSuccess, void>(
-        client_id, keys, metadatas);
+        client_id, tasks, metadatas);
+    timer.LogResponseExpected(result);
+    return result;
+}
+
+tl::expected<std::vector<PromotionTaskItem>, ErrorCode>
+MasterClient::PromotionObjectHeartbeat(const UUID& client_id) {
+    ScopedVLogTimer timer(1, "MasterClient::PromotionObjectHeartbeat");
+    timer.LogRequest("client_id=", client_id);
+    return invoke_rpc<&WrappedMasterService::PromotionObjectHeartbeat,
+                      std::vector<PromotionTaskItem>>(client_id);
+}
+
+tl::expected<PromotionAllocStartResponse, ErrorCode>
+MasterClient::PromotionAllocStart(
+    const UUID& client_id, const std::string& key, uint64_t size,
+    const std::vector<std::string>& preferred_segments) {
+    return PromotionAllocStart(client_id, key, tenant_id_, size,
+                               preferred_segments);
+}
+
+tl::expected<PromotionAllocStartResponse, ErrorCode>
+MasterClient::PromotionAllocStart(
+    const UUID& client_id, const std::string& key, const std::string& tenant_id,
+    uint64_t size, const std::vector<std::string>& preferred_segments) {
+    ScopedVLogTimer timer(1, "MasterClient::PromotionAllocStart");
+    timer.LogRequest("client_id=", client_id, ", key=", key,
+                     ", tenant_id=", tenant_id, ", size=", size,
+                     ", preferred_count=", preferred_segments.size());
+    auto result = invoke_rpc<&WrappedMasterService::PromotionAllocStart,
+                             PromotionAllocStartResponse>(
+        client_id, key, tenant_id, size, preferred_segments);
+    timer.LogResponseExpected(result);
+    return result;
+}
+
+tl::expected<void, ErrorCode> MasterClient::NotifyPromotionSuccess(
+    const UUID& client_id, const std::string& key) {
+    return NotifyPromotionSuccess(client_id, key, tenant_id_);
+}
+
+tl::expected<void, ErrorCode> MasterClient::NotifyPromotionSuccess(
+    const UUID& client_id, const std::string& key,
+    const std::string& tenant_id) {
+    ScopedVLogTimer timer(1, "MasterClient::NotifyPromotionSuccess");
+    timer.LogRequest("client_id=", client_id, ", key=", key,
+                     ", tenant_id=", tenant_id);
+    auto result =
+        invoke_rpc<&WrappedMasterService::NotifyPromotionSuccess, void>(
+            client_id, key, tenant_id);
+    timer.LogResponseExpected(result);
+    return result;
+}
+
+tl::expected<void, ErrorCode> MasterClient::NotifyPromotionFailure(
+    const UUID& client_id, const std::string& key) {
+    return NotifyPromotionFailure(client_id, key, tenant_id_);
+}
+
+tl::expected<void, ErrorCode> MasterClient::NotifyPromotionFailure(
+    const UUID& client_id, const std::string& key,
+    const std::string& tenant_id) {
+    ScopedVLogTimer timer(1, "MasterClient::NotifyPromotionFailure");
+    timer.LogRequest("client_id=", client_id, ", key=", key,
+                     ", tenant_id=", tenant_id);
+    auto result =
+        invoke_rpc<&WrappedMasterService::NotifyPromotionFailure, void>(
+            client_id, key, tenant_id);
     timer.LogResponseExpected(result);
     return result;
 }
@@ -824,13 +1092,21 @@ tl::expected<void, ErrorCode> MasterClient::NotifyOffloadSuccess(
 tl::expected<CopyStartResponse, ErrorCode> MasterClient::CopyStart(
     const std::string& key, const std::string& src_segment,
     const std::vector<std::string>& tgt_segments) {
+    return CopyStart(key, tenant_id_, src_segment, tgt_segments);
+}
+
+tl::expected<CopyStartResponse, ErrorCode> MasterClient::CopyStart(
+    const std::string& key, const std::string& tenant_id,
+    const std::string& src_segment,
+    const std::vector<std::string>& tgt_segments) {
     ScopedVLogTimer timer(1, "MasterClient::CopyStart");
-    timer.LogRequest("key=", key, ", src_segment=", src_segment,
+    timer.LogRequest("key=", key, ", tenant_id=", tenant_id,
+                     ", src_segment=", src_segment,
                      ", tgt_segments_count=", tgt_segments.size());
 
     auto result =
         invoke_rpc<&WrappedMasterService::CopyStart, CopyStartResponse>(
-            client_id_, key, src_segment, tgt_segments);
+            client_id_, key, tenant_id, src_segment, tgt_segments);
     timer.LogResponseExpected(result);
     return result;
 }
@@ -848,11 +1124,16 @@ tl::expected<QueryTaskResponse, ErrorCode> MasterClient::QueryTask(
 }
 
 tl::expected<void, ErrorCode> MasterClient::CopyEnd(const std::string& key) {
-    ScopedVLogTimer timer(1, "MasterClient::CopyEnd");
-    timer.LogRequest("key=", key);
+    return CopyEnd(key, tenant_id_);
+}
 
-    auto result =
-        invoke_rpc<&WrappedMasterService::CopyEnd, void>(client_id_, key);
+tl::expected<void, ErrorCode> MasterClient::CopyEnd(
+    const std::string& key, const std::string& tenant_id) {
+    ScopedVLogTimer timer(1, "MasterClient::CopyEnd");
+    timer.LogRequest("key=", key, ", tenant_id=", tenant_id);
+
+    auto result = invoke_rpc<&WrappedMasterService::CopyEnd, void>(
+        client_id_, key, tenant_id);
     timer.LogResponseExpected(result);
     return result;
 }
@@ -869,11 +1150,16 @@ tl::expected<std::vector<TaskAssignment>, ErrorCode> MasterClient::FetchTasks(
 }
 
 tl::expected<void, ErrorCode> MasterClient::CopyRevoke(const std::string& key) {
-    ScopedVLogTimer timer(1, "MasterClient::CopyRevoke");
-    timer.LogRequest("key=", key);
+    return CopyRevoke(key, tenant_id_);
+}
 
-    auto result =
-        invoke_rpc<&WrappedMasterService::CopyRevoke, void>(client_id_, key);
+tl::expected<void, ErrorCode> MasterClient::CopyRevoke(
+    const std::string& key, const std::string& tenant_id) {
+    ScopedVLogTimer timer(1, "MasterClient::CopyRevoke");
+    timer.LogRequest("key=", key, ", tenant_id=", tenant_id);
+
+    auto result = invoke_rpc<&WrappedMasterService::CopyRevoke, void>(
+        client_id_, key, tenant_id);
     timer.LogResponseExpected(result);
     return result;
 }
@@ -881,33 +1167,50 @@ tl::expected<void, ErrorCode> MasterClient::CopyRevoke(const std::string& key) {
 tl::expected<MoveStartResponse, ErrorCode> MasterClient::MoveStart(
     const std::string& key, const std::string& src_segment,
     const std::string& tgt_segment) {
+    return MoveStart(key, tenant_id_, src_segment, tgt_segment);
+}
+
+tl::expected<MoveStartResponse, ErrorCode> MasterClient::MoveStart(
+    const std::string& key, const std::string& tenant_id,
+    const std::string& src_segment, const std::string& tgt_segment) {
     ScopedVLogTimer timer(1, "MasterClient::MoveStart");
-    timer.LogRequest("key=", key, ", src_segment=", src_segment,
+    timer.LogRequest("key=", key, ", tenant_id=", tenant_id,
+                     ", src_segment=", src_segment,
                      ", tgt_segment=", tgt_segment);
 
     auto result =
         invoke_rpc<&WrappedMasterService::MoveStart, MoveStartResponse>(
-            client_id_, key, src_segment, tgt_segment);
+            client_id_, key, tenant_id, src_segment, tgt_segment);
     timer.LogResponseExpected(result);
     return result;
 }
 
 tl::expected<void, ErrorCode> MasterClient::MoveEnd(const std::string& key) {
-    ScopedVLogTimer timer(1, "MasterClient::MoveEnd");
-    timer.LogRequest("key=", key);
+    return MoveEnd(key, tenant_id_);
+}
 
-    auto result =
-        invoke_rpc<&WrappedMasterService::MoveEnd, void>(client_id_, key);
+tl::expected<void, ErrorCode> MasterClient::MoveEnd(
+    const std::string& key, const std::string& tenant_id) {
+    ScopedVLogTimer timer(1, "MasterClient::MoveEnd");
+    timer.LogRequest("key=", key, ", tenant_id=", tenant_id);
+
+    auto result = invoke_rpc<&WrappedMasterService::MoveEnd, void>(
+        client_id_, key, tenant_id);
     timer.LogResponseExpected(result);
     return result;
 }
 
 tl::expected<void, ErrorCode> MasterClient::MoveRevoke(const std::string& key) {
-    ScopedVLogTimer timer(1, "MasterClient::MoveRevoke");
-    timer.LogRequest("key=", key);
+    return MoveRevoke(key, tenant_id_);
+}
 
-    auto result =
-        invoke_rpc<&WrappedMasterService::MoveRevoke, void>(client_id_, key);
+tl::expected<void, ErrorCode> MasterClient::MoveRevoke(
+    const std::string& key, const std::string& tenant_id) {
+    ScopedVLogTimer timer(1, "MasterClient::MoveRevoke");
+    timer.LogRequest("key=", key, ", tenant_id=", tenant_id);
+
+    auto result = invoke_rpc<&WrappedMasterService::MoveRevoke, void>(
+        client_id_, key, tenant_id);
     timer.LogResponseExpected(result);
     return result;
 }
@@ -924,24 +1227,37 @@ tl::expected<void, ErrorCode> MasterClient::MarkTaskToComplete(
 
 tl::expected<void, ErrorCode> MasterClient::EvictDiskReplica(
     const std::string& key, ReplicaType replica_type) {
+    return EvictDiskReplica(key, tenant_id_, replica_type);
+}
+
+tl::expected<void, ErrorCode> MasterClient::EvictDiskReplica(
+    const std::string& key, const std::string& tenant_id,
+    ReplicaType replica_type) {
     ScopedVLogTimer timer(1, "MasterClient::EvictDiskReplica");
-    timer.LogRequest("key=", key, ", replica_type=", replica_type);
+    timer.LogRequest("key=", key, ", tenant_id=", tenant_id,
+                     ", replica_type=", replica_type);
 
     auto result = invoke_rpc<&WrappedMasterService::EvictDiskReplica, void>(
-        client_id_, key, replica_type);
+        client_id_, key, tenant_id, replica_type);
     timer.LogResponseExpected(result);
     return result;
 }
 
 std::vector<tl::expected<void, ErrorCode>> MasterClient::BatchEvictDiskReplica(
     const std::vector<std::string>& keys, ReplicaType replica_type) {
+    return BatchEvictDiskReplica(keys, tenant_id_, replica_type);
+}
+
+std::vector<tl::expected<void, ErrorCode>> MasterClient::BatchEvictDiskReplica(
+    const std::vector<std::string>& keys, const std::string& tenant_id,
+    ReplicaType replica_type) {
     ScopedVLogTimer timer(1, "MasterClient::BatchEvictDiskReplica");
-    timer.LogRequest("keys_count=", keys.size(),
+    timer.LogRequest("keys_count=", keys.size(), ", tenant_id=", tenant_id,
                      ", replica_type=", replica_type);
 
     auto result =
         invoke_batch_rpc<&WrappedMasterService::BatchEvictDiskReplica, void>(
-            keys.size(), client_id_, keys, replica_type);
+            keys.size(), client_id_, keys, tenant_id, replica_type);
     timer.LogResponse("result=", result.size(), " operations");
     return result;
 }
diff --git a/mooncake-store/src/master_metric_manager.cpp b/mooncake-store/src/master_metric_manager.cpp
index 1f01da47..c5675cf1 100644
--- a/mooncake-store/src/master_metric_manager.cpp
+++ b/mooncake-store/src/master_metric_manager.cpp
@@ -2,6 +2,7 @@
 
 #include <glog/logging.h>
 #include <iomanip>  // For std::fixed, std::setprecision
+#include <limits>   // For std::numeric_limits
 #include <sstream>  // For string building during serialization
 #include <vector>   // Required by histogram serialization
 #include <cmath>
@@ -31,6 +32,18 @@ MasterMetricManager::MasterMetricManager()
       mem_total_capacity_per_segment_(
           "segment_total_capacity_bytes",
           "Total memory capacity of the mounted segment", {"segment"}),
+      nof_allocated_size_(
+          "master_nof_allocated_bytes",
+          "Total nof ssd bytes currently allocated across all segments"),
+      nof_total_capacity_("master_total_nof_capacity_bytes",
+                          "Total nof ssd capacity across all mounted segments"),
+      nof_allocated_size_per_segment_(
+          "nof_segment_allocated_bytes",
+          "Total nof ssd bytes currently allocated of the segment",
+          {"segment"}),
+      nof_total_capacity_per_segment_(
+          "nof_segment_total_capacity_bytes",
+          "Total nof ssd capacity of the mounted segment", {"segment"}),
       file_allocated_size_(
           "master_allocated_file_size_bytes",
           "Total bytes currently allocated for file storage in 3fs/nfs"),
@@ -54,6 +67,10 @@ MasterMetricManager::MasterMetricManager()
                           "Total number of PutStart requests received"),
       put_start_failures_("master_put_start_failures_total",
                           "Total number of failed PutStart requests"),
+      put_start_alloc_failures_(
+          "master_put_start_alloc_failures_total",
+          "Total number of PutStart failures caused by replica allocation "
+          "failure"),
       put_end_requests_("master_put_end_requests_total",
                         "Total number of PutEnd requests received"),
       put_end_failures_("master_put_end_failures_total",
@@ -109,10 +126,44 @@ MasterMetricManager::MasterMetricManager()
       remount_segment_failures_(
           "master_remount_segment_failures_total",
           "Total number of failed RemountSegment requests"),
+      mount_nof_segment_requests_(
+          "master_mount_nof_segment_requests_total",
+          "Total number of MountNoFSegment requests received"),
+      mount_nof_segment_failures_(
+          "master_mount_nof_segment_failures_total",
+          "Total number of failed MountNoFSegment requests"),
+      unmount_nof_segment_requests_(
+          "master_unmount_nof_segment_requests_total",
+          "Total number of UnmountNoFSegment requests received"),
+      unmount_nof_segment_failures_(
+          "master_unmount_nof_segment_failures_total",
+          "Total number of failed UnmountNoFSegment requests"),
+      remount_nof_segment_requests_(
+          "master_remount_nof_segment_requests_total",
+          "Total number of RemountNoFSegment requests received"),
+      remount_nof_segment_failures_(
+          "master_remount_nof_segment_failures_total",
+          "Total number of failed RemountNoFSegment requests"),
       ping_requests_("master_ping_requests_total",
                      "Total number of ping requests received"),
       ping_failures_("master_ping_failures_total",
                      "Total number of failed ping requests"),
+      nof_heartbeat_success_total_(
+          "master_nof_heartbeat_success_total",
+          "Total number of successful NoF heartbeat probes"),
+      nof_heartbeat_failure_total_(
+          "master_nof_heartbeat_failure_total",
+          "Total number of failed NoF heartbeat probes"),
+      nof_heartbeat_timeout_total_(
+          "master_nof_heartbeat_timeout_total",
+          "Total number of timed out NoF heartbeat probes"),
+      nof_segments_unmounted_by_heartbeat_total_(
+          "master_nof_segments_unmounted_by_heartbeat_total",
+          "Total number of NoF segments unmounted due to heartbeat failures"),
+      nof_heartbeat_probe_latency_ms_(
+          "master_nof_heartbeat_probe_latency_ms",
+          "Latency distribution of NoF heartbeat probes in milliseconds",
+          {1, 2, 5, 10, 20, 50, 100, 200, 500, 1000, 2000}),
 
       // Initialize Batch Request Counters
       batch_exist_key_requests_(
@@ -218,20 +269,26 @@ MasterMetricManager::MasterMetricManager()
           "master_batch_put_revoke_failed_items_total",
           "Total number of failed items in BatchPutRevoke requests"),
 
-      // Initialize cache hit rate metrics
+      // Initialize Store-observed cache reuse metrics. These are not
+      // end-to-end request/token-level cache hit ratio metrics.
       mem_cache_hit_nums_("mem_cache_hit_nums_",
-                          "Total number of cache hits in the memory pool"),
+                          "Total number of GetReplicaList results served from "
+                          "the memory pool"),
       file_cache_hit_nums_("file_cache_hit_nums_",
-                           "Total number of cache hits in the ssd"),
+                           "Total number of GetReplicaList results served from "
+                           "the SSD cache"),
       mem_cache_nums_("mem_cache_nums_",
-                      "Total number of cached values in the memory pool"),
+                      "Current number of cached values in the memory pool"),
       file_cache_nums_("file_cache_nums_",
-                       "Total number of cached values in the ssd"),
+                       "Current number of cached values in the SSD cache"),
       valid_get_nums_("valid_get_nums_",
-                      "Total number of valid get operations"),
-      total_get_nums_("total_get_nums_", "Total number of get operations"),
+                      "Total number of GetReplicaList operations that returned "
+                      "at least one completed replica"),
+      total_get_nums_("total_get_nums_",
+                      "Total number of GetReplicaList operations"),
 
       // Initialize Eviction Counters
+      // total eviction
       eviction_success_("master_successful_evictions_total",
                         "Total number of successful eviction operations"),
       eviction_attempts_("master_attempted_evictions_total",
@@ -240,6 +297,28 @@ MasterMetricManager::MasterMetricManager()
                          "Total number of keys evicted"),
       evicted_size_("master_evicted_size_bytes",
                     "Total bytes of evicted objects"),
+      // mem eviction
+      mem_eviction_success_(
+          "master_successful_evictions_mem",
+          "Total number of successful eviction operations in mem"),
+      mem_eviction_attempts_(
+          "master_attempted_evictions_mem",
+          "Total number of attempted eviction operations in mem"),
+      mem_evicted_key_count_("master_evicted_key_count_mem",
+                             "Total number of keys evicted in mem"),
+      mem_evicted_size_("master_evicted_size_bytes_mem",
+                        "Total bytes of evicted objects in mem"),
+      // nof eviction
+      nof_eviction_success_(
+          "master_successful_evictions_nof",
+          "Total number of successful eviction operations in nof"),
+      nof_eviction_attempts_(
+          "master_attempted_evictions_nof",
+          "Total number of attempted eviction operations in nof"),
+      nof_evicted_key_count_("master_evicted_key_count_nof",
+                             "Total number of keys evicted in nof"),
+      nof_evicted_size_("master_evicted_size_bytes_nof",
+                        "Total bytes of evicted objects in nof"),
 
       // Initialize Discarded Replicas Counters
       put_start_discard_cnt_("master_put_start_discard_cnt",
@@ -250,6 +329,53 @@ MasterMetricManager::MasterMetricManager()
           "master_put_start_discarded_staging_size",
           "Total size of memory replicas in discarded but not yet released "
           "PutStart operations"),
+
+      // Promotion-on-hit Metrics
+      promotion_in_flight_metric_(
+          "master_promotion_in_flight",
+          "Current number of in-flight L2->L1 promotion tasks"),
+      promotion_admitted_(
+          "master_promotion_admitted_total",
+          "Total promotion tasks admitted past all gates and enqueued"),
+      promotion_completed_(
+          "master_promotion_completed_total",
+          "Total promotion tasks committed via NotifyPromotionSuccess"),
+      promotion_completed_bytes_(
+          "master_promotion_completed_bytes_total",
+          "Total bytes promoted from LOCAL_DISK to MEMORY"),
+      promotion_expired_("master_promotion_expired_total",
+                         "Total promotion tasks expired via the reaper "
+                         "(put_start_release_timeout_sec)"),
+      promotion_failed_(
+          "master_promotion_failed_total",
+          "Total promotion tasks aborted by holder via "
+          "NotifyPromotionFailure (holder reported a downstream failure)"),
+      promotion_cancelled_(
+          "master_promotion_cancelled_total",
+          "Total promotion tasks removed because the prerequisite went "
+          "away: object removal mid-flight (Remove / UpsertStart / etc.), "
+          "holder-client expiry (ClearInvalidHandles), or staged replica "
+          "lost (NotifyPromotionSuccess committed=false)"),
+      promotion_rejected_frequency_(
+          "master_promotion_rejected_frequency_total",
+          "Promotion attempts rejected because CountMinSketch frequency "
+          "was below promotion_admission_threshold"),
+      promotion_rejected_watermark_(
+          "master_promotion_rejected_watermark_total",
+          "Promotion attempts rejected because DRAM was at or above the "
+          "eviction high watermark"),
+      promotion_rejected_cap_(
+          "master_promotion_rejected_cap_total",
+          "Promotion attempts rejected because promotion_in_flight was at "
+          "promotion_queue_limit"),
+      tenant_quota_reject_total_(
+          "mooncake_tenant_quota_reject_total",
+          "Total number of tenant quota admission rejects",
+          {"tenant_id", "reason"}),
+      tenant_evict_bytes_total_(
+          "mooncake_tenant_evict_bytes_total",
+          "Total bytes evicted by tenant-scoped quota eviction", {"tenant_id"}),
+
       // Snapshot Metrics
       snapshot_duration_ms_(
           "master_snapshot_duration_ms",
@@ -338,10 +464,21 @@ void MasterMetricManager::update_metrics_for_zero_output() {
     mem_cache_nums_.update(0);
     file_cache_nums_.update(0);
     put_start_discarded_staging_size_.update(0);
+    promotion_in_flight_metric_.update(0);
 
     // Update Counters (use inc(0) to mark as changed)
+    promotion_admitted_.inc(0);
+    promotion_completed_.inc(0);
+    promotion_completed_bytes_.inc(0);
+    promotion_expired_.inc(0);
+    promotion_failed_.inc(0);
+    promotion_cancelled_.inc(0);
+    promotion_rejected_frequency_.inc(0);
+    promotion_rejected_watermark_.inc(0);
+    promotion_rejected_cap_.inc(0);
     put_start_requests_.inc(0);
     put_start_failures_.inc(0);
+    put_start_alloc_failures_.inc(0);
     put_end_requests_.inc(0);
     put_end_failures_.inc(0);
     put_revoke_requests_.inc(0);
@@ -431,7 +568,7 @@ void MasterMetricManager::update_metrics_for_zero_output() {
     batch_put_revoke_items_.inc(0);
     batch_put_revoke_failed_items_.inc(0);
 
-    // Update cache hit rate metrics
+    // Update Store-observed cache reuse metrics
     mem_cache_hit_nums_.inc(0);
     file_cache_hit_nums_.inc(0);
     valid_get_nums_.inc(0);
@@ -449,6 +586,7 @@ void MasterMetricManager::update_metrics_for_zero_output() {
 
     // Update Histogram (use observe(0) to mark as changed)
     value_size_distribution_.observe(0);
+    nof_heartbeat_probe_latency_ms_.observe(0);
 
     // Note: dynamic_gauge_1t (mem_allocated_size_per_segment_ and
     // mem_total_capacity_per_segment_) are not initialized here because they
@@ -536,6 +674,76 @@ double MasterMetricManager::get_segment_mem_used_ratio(
     return allocated / capacity;
 }
 
+// NoF segment Metrics
+void MasterMetricManager::inc_allocated_nof_size(const std::string& segment,
+                                                 int64_t val) {
+    nof_allocated_size_.inc(val);
+    if (!segment.empty()) nof_allocated_size_per_segment_.inc({segment}, val);
+}
+
+void MasterMetricManager::dec_allocated_nof_size(const std::string& segment,
+                                                 int64_t val) {
+    nof_allocated_size_.dec(val);
+    if (!segment.empty()) nof_allocated_size_per_segment_.dec({segment}, val);
+}
+
+void MasterMetricManager::reset_allocated_nof_size() {
+    nof_allocated_size_.reset();
+}
+
+void MasterMetricManager::inc_total_nof_capacity(const std::string& segment,
+                                                 int64_t val) {
+    nof_total_capacity_.inc(val);
+    if (!segment.empty()) nof_total_capacity_per_segment_.inc({segment}, val);
+}
+
+void MasterMetricManager::dec_total_nof_capacity(const std::string& segment,
+                                                 int64_t val) {
+    nof_total_capacity_.dec(val);
+    if (!segment.empty()) nof_total_capacity_per_segment_.dec({segment}, val);
+}
+
+void MasterMetricManager::reset_total_nof_capacity() {
+    nof_total_capacity_.reset();
+}
+
+int64_t MasterMetricManager::get_allocated_nof_size() {
+    return nof_allocated_size_.value();
+}
+
+int64_t MasterMetricManager::get_total_nof_capacity() {
+    return nof_total_capacity_.value();
+}
+
+double MasterMetricManager::get_global_nof_used_ratio(void) {
+    double allocated = nof_allocated_size_.value();
+    double capacity = nof_total_capacity_.value();
+    if (capacity == 0) {
+        return 0.0;
+    }
+    return allocated / capacity;
+}
+
+int64_t MasterMetricManager::get_segment_allocated_nof_size(
+    const std::string& segment) {
+    return nof_allocated_size_per_segment_.value({segment});
+}
+
+int64_t MasterMetricManager::get_segment_total_nof_capacity(
+    const std::string& segment) {
+    return nof_total_capacity_per_segment_.value({segment});
+}
+
+double MasterMetricManager::get_segment_nof_used_ratio(
+    const std::string& segment) {
+    double allocated = get_segment_allocated_nof_size(segment);
+    double capacity = get_segment_total_nof_capacity(segment);
+    if (capacity == 0) {
+        return 0.0;
+    }
+    return allocated / capacity;
+}
+
 // File Storage Metrics
 void MasterMetricManager::inc_allocated_file_size(int64_t val) {
     file_allocated_size_.inc(val);
@@ -551,15 +759,29 @@ void MasterMetricManager::dec_total_file_capacity(int64_t val) {
     file_total_capacity_.dec(val);
 }
 
+void MasterMetricManager::set_dfs_capacity_unlimited(bool unlimited) {
+    dfs_capacity_unlimited_ = unlimited;
+}
+
+bool MasterMetricManager::is_dfs_capacity_unlimited() const {
+    return dfs_capacity_unlimited_;
+}
+
 int64_t MasterMetricManager::get_allocated_file_size() {
     return file_allocated_size_.value();
 }
 
 int64_t MasterMetricManager::get_total_file_capacity() {
+    if (dfs_capacity_unlimited_ && file_total_capacity_.value() == 0) {
+        return std::numeric_limits<int64_t>::max();
+    }
     return file_total_capacity_.value();
 }
 
 double MasterMetricManager::get_global_file_used_ratio(void) {
+    if (dfs_capacity_unlimited_ && file_total_capacity_.value() == 0) {
+        return 0.0;
+    }
     double allocated = file_allocated_size_.value();
     double capacity = file_total_capacity_.value();
     if (capacity == 0) {
@@ -602,7 +824,7 @@ int64_t MasterMetricManager::get_active_clients() {
     return active_clients_.value();
 }
 
-// cache hit rate metrics
+// Store-observed cache reuse metrics
 void MasterMetricManager::inc_mem_cache_hit_nums(int64_t val) {
     mem_cache_hit_nums_.inc(val);
 }
@@ -621,6 +843,10 @@ void MasterMetricManager::dec_mem_cache_nums(int64_t val) {
 void MasterMetricManager::dec_file_cache_nums(int64_t val) {
     file_cache_nums_.dec(val);
 }
+void MasterMetricManager::reset_cache_total_nums() {
+    mem_cache_nums_.reset();
+    file_cache_nums_.reset();
+}
 void MasterMetricManager::inc_valid_get_nums(int64_t val) {
     valid_get_nums_.inc(val);
 }
@@ -641,6 +867,9 @@ void MasterMetricManager::inc_put_start_requests(int64_t val) {
 void MasterMetricManager::inc_put_start_failures(int64_t val) {
     put_start_failures_.inc(val);
 }
+void MasterMetricManager::inc_put_start_alloc_failures(int64_t val) {
+    put_start_alloc_failures_.inc(val);
+}
 void MasterMetricManager::inc_put_end_requests(int64_t val) {
     put_end_requests_.inc(val);
 }
@@ -689,18 +918,36 @@ void MasterMetricManager::inc_mount_segment_requests(int64_t val) {
 void MasterMetricManager::inc_mount_segment_failures(int64_t val) {
     mount_segment_failures_.inc(val);
 }
+void MasterMetricManager::inc_mount_nof_segment_requests(int64_t val) {
+    mount_nof_segment_requests_.inc(val);
+}
+void MasterMetricManager::inc_mount_nof_segment_failures(int64_t val) {
+    mount_nof_segment_failures_.inc(val);
+}
 void MasterMetricManager::inc_unmount_segment_requests(int64_t val) {
     unmount_segment_requests_.inc(val);
 }
 void MasterMetricManager::inc_unmount_segment_failures(int64_t val) {
     unmount_segment_failures_.inc(val);
 }
+void MasterMetricManager::inc_unmount_nof_segment_requests(int64_t val) {
+    unmount_nof_segment_requests_.inc(val);
+}
+void MasterMetricManager::inc_unmount_nof_segment_failures(int64_t val) {
+    unmount_nof_segment_failures_.inc(val);
+}
 void MasterMetricManager::inc_remount_segment_requests(int64_t val) {
     remount_segment_requests_.inc(val);
 }
 void MasterMetricManager::inc_remount_segment_failures(int64_t val) {
     remount_segment_failures_.inc(val);
 }
+void MasterMetricManager::inc_remount_nof_segment_requests(int64_t val) {
+    remount_nof_segment_requests_.inc(val);
+}
+void MasterMetricManager::inc_remount_nof_segment_failures(int64_t val) {
+    remount_nof_segment_failures_.inc(val);
+}
 void MasterMetricManager::inc_ping_requests(int64_t val) {
     ping_requests_.inc(val);
 }
@@ -708,6 +955,28 @@ void MasterMetricManager::inc_ping_failures(int64_t val) {
     ping_failures_.inc(val);
 }
 
+void MasterMetricManager::inc_nof_heartbeat_success_total(int64_t val) {
+    nof_heartbeat_success_total_.inc(val);
+}
+
+void MasterMetricManager::inc_nof_heartbeat_failure_total(int64_t val) {
+    nof_heartbeat_failure_total_.inc(val);
+}
+
+void MasterMetricManager::inc_nof_heartbeat_timeout_total(int64_t val) {
+    nof_heartbeat_timeout_total_.inc(val);
+}
+
+void MasterMetricManager::inc_nof_segments_unmounted_by_heartbeat_total(
+    int64_t val) {
+    nof_segments_unmounted_by_heartbeat_total_.inc(val);
+}
+
+void MasterMetricManager::observe_nof_heartbeat_probe_latency_ms(
+    int64_t latency_ms) {
+    nof_heartbeat_probe_latency_ms_.observe(latency_ms);
+}
+
 // Batch Operation Statistics (Counters)
 void MasterMetricManager::inc_batch_exist_key_requests(int64_t items) {
     batch_exist_key_requests_.inc(1);
@@ -816,6 +1085,52 @@ void MasterMetricManager::inc_put_start_release_cnt(int64_t count,
     put_start_discarded_staging_size_.dec(size);
 }
 
+// --- Promotion-on-hit Metrics ---
+void MasterMetricManager::inc_promotion_in_flight(int64_t val) {
+    promotion_in_flight_metric_.inc(val);
+}
+void MasterMetricManager::dec_promotion_in_flight(int64_t val) {
+    promotion_in_flight_metric_.dec(val);
+}
+void MasterMetricManager::inc_promotion_admitted(int64_t val) {
+    promotion_admitted_.inc(val);
+}
+void MasterMetricManager::inc_promotion_completed(int64_t val) {
+    promotion_completed_.inc(val);
+}
+void MasterMetricManager::inc_promotion_completed_bytes(int64_t bytes) {
+    promotion_completed_bytes_.inc(bytes);
+}
+void MasterMetricManager::inc_promotion_expired(int64_t val) {
+    promotion_expired_.inc(val);
+}
+void MasterMetricManager::inc_promotion_failed(int64_t val) {
+    promotion_failed_.inc(val);
+}
+void MasterMetricManager::inc_promotion_cancelled(int64_t val) {
+    promotion_cancelled_.inc(val);
+}
+void MasterMetricManager::inc_promotion_rejected_frequency(int64_t val) {
+    promotion_rejected_frequency_.inc(val);
+}
+void MasterMetricManager::inc_promotion_rejected_watermark(int64_t val) {
+    promotion_rejected_watermark_.inc(val);
+}
+void MasterMetricManager::inc_promotion_rejected_cap(int64_t val) {
+    promotion_rejected_cap_.inc(val);
+}
+
+void MasterMetricManager::inc_tenant_quota_reject(const std::string& tenant_id,
+                                                  const std::string& reason,
+                                                  int64_t val) {
+    tenant_quota_reject_total_.inc({tenant_id, reason}, val);
+}
+
+void MasterMetricManager::inc_tenant_evict_bytes(const std::string& tenant_id,
+                                                 int64_t bytes) {
+    tenant_evict_bytes_total_.inc({tenant_id}, bytes);
+}
+
 void MasterMetricManager::set_snapshot_duration_ms(int64_t size) {
     snapshot_duration_ms_.observe(size);
 }
@@ -832,6 +1147,10 @@ int64_t MasterMetricManager::get_put_start_failures() {
     return put_start_failures_.value();
 }
 
+int64_t MasterMetricManager::get_put_start_alloc_failures() {
+    return put_start_alloc_failures_.value();
+}
+
 int64_t MasterMetricManager::get_put_end_requests() {
     return put_end_requests_.value();
 }
@@ -1079,6 +1398,30 @@ void MasterMetricManager::inc_eviction_success(int64_t key_count,
 
 void MasterMetricManager::inc_eviction_fail() { eviction_attempts_.inc(); }
 
+void MasterMetricManager::inc_mem_eviction_success(int64_t key_count,
+                                                   int64_t size) {
+    mem_evicted_key_count_.inc(key_count);
+    mem_evicted_size_.inc(size);
+    mem_eviction_success_.inc();
+    mem_eviction_attempts_.inc();
+}
+
+void MasterMetricManager::inc_mem_eviction_fail() {
+    mem_eviction_attempts_.inc();
+}
+
+void MasterMetricManager::inc_nof_eviction_success(int64_t key_count,
+                                                   int64_t size) {
+    nof_evicted_key_count_.inc(key_count);
+    nof_evicted_size_.inc(size);
+    nof_eviction_success_.inc();
+    nof_eviction_attempts_.inc();
+}
+
+void MasterMetricManager::inc_nof_eviction_fail() {
+    nof_eviction_attempts_.inc();
+}
+
 int64_t MasterMetricManager::get_eviction_success() {
     return eviction_success_.value();
 }
@@ -1095,6 +1438,38 @@ int64_t MasterMetricManager::get_evicted_size() {
     return evicted_size_.value();
 }
 
+int64_t MasterMetricManager::get_mem_eviction_success() {
+    return mem_eviction_success_.value();
+}
+
+int64_t MasterMetricManager::get_mem_eviction_attempts() {
+    return mem_eviction_attempts_.value();
+}
+
+int64_t MasterMetricManager::get_mem_evicted_key_count() {
+    return mem_evicted_key_count_.value();
+}
+
+int64_t MasterMetricManager::get_mem_evicted_size() {
+    return mem_evicted_size_.value();
+}
+
+int64_t MasterMetricManager::get_nof_eviction_success() {
+    return nof_eviction_success_.value();
+}
+
+int64_t MasterMetricManager::get_nof_eviction_attempts() {
+    return nof_eviction_attempts_.value();
+}
+
+int64_t MasterMetricManager::get_nof_evicted_key_count() {
+    return nof_evicted_key_count_.value();
+}
+
+int64_t MasterMetricManager::get_nof_evicted_size() {
+    return nof_evicted_size_.value();
+}
+
 // PutStart Discard Metrics Getters
 int64_t MasterMetricManager::get_put_start_discard_cnt() {
     return put_start_discard_cnt_.value();
@@ -1108,6 +1483,38 @@ int64_t MasterMetricManager::get_put_start_discarded_staging_size() {
     return put_start_discarded_staging_size_.value();
 }
 
+// --- Promotion-on-hit Metrics Getters ---
+int64_t MasterMetricManager::get_promotion_in_flight() {
+    return promotion_in_flight_metric_.value();
+}
+int64_t MasterMetricManager::get_promotion_admitted() {
+    return promotion_admitted_.value();
+}
+int64_t MasterMetricManager::get_promotion_completed() {
+    return promotion_completed_.value();
+}
+int64_t MasterMetricManager::get_promotion_completed_bytes() {
+    return promotion_completed_bytes_.value();
+}
+int64_t MasterMetricManager::get_promotion_expired() {
+    return promotion_expired_.value();
+}
+int64_t MasterMetricManager::get_promotion_failed() {
+    return promotion_failed_.value();
+}
+int64_t MasterMetricManager::get_promotion_cancelled() {
+    return promotion_cancelled_.value();
+}
+int64_t MasterMetricManager::get_promotion_rejected_frequency() {
+    return promotion_rejected_frequency_.value();
+}
+int64_t MasterMetricManager::get_promotion_rejected_watermark() {
+    return promotion_rejected_watermark_.value();
+}
+int64_t MasterMetricManager::get_promotion_rejected_cap() {
+    return promotion_rejected_cap_.value();
+}
+
 // CopyStart, CopyEnd, CopyRevoke, MoveStart, MoveEnd, MoveRevoke Metrics
 void MasterMetricManager::inc_copy_start_requests(int64_t val) {
     copy_start_requests_.inc(val);
@@ -1294,6 +1701,7 @@ std::string MasterMetricManager::serialize_metrics() {
     serialize_metric(exist_key_failures_);
     serialize_metric(put_start_requests_);
     serialize_metric(put_start_failures_);
+    serialize_metric(put_start_alloc_failures_);
     serialize_metric(put_end_requests_);
     serialize_metric(put_end_failures_);
     serialize_metric(put_revoke_requests_);
@@ -1316,6 +1724,11 @@ std::string MasterMetricManager::serialize_metrics() {
     serialize_metric(remount_segment_failures_);
     serialize_metric(ping_requests_);
     serialize_metric(ping_failures_);
+    serialize_metric(nof_heartbeat_success_total_);
+    serialize_metric(nof_heartbeat_failure_total_);
+    serialize_metric(nof_heartbeat_timeout_total_);
+    serialize_metric(nof_segments_unmounted_by_heartbeat_total_);
+    serialize_metric(nof_heartbeat_probe_latency_ms_);
 
     // Serialize CopyStart, CopyEnd, CopyRevoke, MoveStart, MoveEnd, MoveRevoke
     // Counters
@@ -1374,6 +1787,20 @@ std::string MasterMetricManager::serialize_metrics() {
     serialize_metric(put_start_release_cnt_);
     serialize_metric(put_start_discarded_staging_size_);
 
+    // Serialize Promotion-on-hit Metrics
+    serialize_metric(promotion_in_flight_metric_);
+    serialize_metric(promotion_admitted_);
+    serialize_metric(promotion_completed_);
+    serialize_metric(promotion_completed_bytes_);
+    serialize_metric(promotion_expired_);
+    serialize_metric(promotion_failed_);
+    serialize_metric(promotion_cancelled_);
+    serialize_metric(promotion_rejected_frequency_);
+    serialize_metric(promotion_rejected_watermark_);
+    serialize_metric(promotion_rejected_cap_);
+    serialize_metric(tenant_quota_reject_total_);
+    serialize_metric(tenant_evict_bytes_total_);
+
     // Serialize Snapshot Metrics
     serialize_metric(snapshot_duration_ms_);
     serialize_metric(snapshot_success_);
@@ -1396,25 +1823,33 @@ MasterMetricManager::calculate_cache_stats() {
     int64_t valid_get_nums = valid_get_nums_.value();
     int64_t total_get_nums = total_get_nums_.value();
 
-    double mem_hit_rate = 0.0;
+    // These values divide cumulative Store-observed hits by the current cached
+    // object count. They are not bounded request/token-level hit ratios; that
+    // end-to-end metric belongs to Conductor or the inference engine.
+    double mem_hits_per_current_cached_object = 0.0;
     if (mem_total_cache > 0) {
-        mem_hit_rate = static_cast<double>(mem_cache_hits) /
-                       static_cast<double>(mem_total_cache);
-        mem_hit_rate = std::round(mem_hit_rate * 100.0) / 100.0;
+        mem_hits_per_current_cached_object =
+            static_cast<double>(mem_cache_hits) /
+            static_cast<double>(mem_total_cache);
+        mem_hits_per_current_cached_object =
+            std::round(mem_hits_per_current_cached_object * 100.0) / 100.0;
     }
 
-    double ssd_hit_rate = 0.0;
+    double ssd_hits_per_current_cached_object = 0.0;
     if (ssd_total_cache > 0) {
-        ssd_hit_rate = static_cast<double>(ssd_cache_hits) /
-                       static_cast<double>(ssd_total_cache);
-        ssd_hit_rate = std::round(ssd_hit_rate * 100.0) / 100.0;
+        ssd_hits_per_current_cached_object =
+            static_cast<double>(ssd_cache_hits) /
+            static_cast<double>(ssd_total_cache);
+        ssd_hits_per_current_cached_object =
+            std::round(ssd_hits_per_current_cached_object * 100.0) / 100.0;
     }
 
-    double total_hit_rate = 0.0;
+    double overall_hits_per_current_cached_object = 0.0;
     if (total_cache > 0) {
-        total_hit_rate =
+        overall_hits_per_current_cached_object =
             static_cast<double>(total_hits) / static_cast<double>(total_cache);
-        total_hit_rate = std::round(total_hit_rate * 100.0) / 100.0;
+        overall_hits_per_current_cached_object =
+            std::round(overall_hits_per_current_cached_object * 100.0) / 100.0;
     }
 
     double valid_get_rate = 0.0;
@@ -1428,10 +1863,12 @@ MasterMetricManager::calculate_cache_stats() {
     add_stat_to_dict(stats_dict, CacheHitStat::SSD_HITS, ssd_cache_hits);
     add_stat_to_dict(stats_dict, CacheHitStat::MEMORY_TOTAL, mem_total_cache);
     add_stat_to_dict(stats_dict, CacheHitStat::SSD_TOTAL, ssd_total_cache);
-    add_stat_to_dict(stats_dict, CacheHitStat::MEMORY_HIT_RATE, mem_hit_rate);
-    add_stat_to_dict(stats_dict, CacheHitStat::SSD_HIT_RATE, ssd_hit_rate);
+    add_stat_to_dict(stats_dict, CacheHitStat::MEMORY_HIT_RATE,
+                     mem_hits_per_current_cached_object);
+    add_stat_to_dict(stats_dict, CacheHitStat::SSD_HIT_RATE,
+                     ssd_hits_per_current_cached_object);
     add_stat_to_dict(stats_dict, CacheHitStat::OVERALL_HIT_RATE,
-                     total_hit_rate);
+                     overall_hits_per_current_cached_object);
     add_stat_to_dict(stats_dict, CacheHitStat::VALID_GET_RATE, valid_get_rate);
     return stats_dict;
 }
@@ -1447,11 +1884,22 @@ void MasterMetricManager::add_stat_to_dict(
 
 // --- Human-Readable Summary ---
 std::string MasterMetricManager::get_summary_string() {
+    return get_summary_string(false);
+}
+
+std::string MasterMetricManager::get_summary_string_and_update_snapshot() {
+    return get_summary_string(true);
+}
+
+std::string MasterMetricManager::get_summary_string(
+    bool update_summary_snapshot) {
     std::stringstream ss;
 
     // --- Get current values ---
     int64_t mem_allocated = mem_allocated_size_.value();
     int64_t mem_capacity = mem_total_capacity_.value();
+    int64_t nof_allocated = nof_allocated_size_.value();
+    int64_t nof_capacity = nof_total_capacity_.value();
     int64_t file_allocated = file_allocated_size_.value();
     int64_t file_capacity = file_total_capacity_.value();
     int64_t keys = key_count_.value();
@@ -1463,6 +1911,7 @@ std::string MasterMetricManager::get_summary_string() {
     int64_t exist_key_fails = exist_key_failures_.value();
     int64_t put_starts = put_start_requests_.value();
     int64_t put_start_fails = put_start_failures_.value();
+    int64_t put_start_alloc_fails = put_start_alloc_failures_.value();
     int64_t put_ends = put_end_requests_.value();
     int64_t put_end_fails = put_end_failures_.value();
     int64_t put_revoke_requests = put_revoke_requests_.value();
@@ -1552,10 +2001,21 @@ std::string MasterMetricManager::get_summary_string() {
         batch_replica_clear_failed_items_.value();
 
     // Eviction counters
+    // Total counters
     int64_t eviction_success = eviction_success_.value();
     int64_t eviction_attempts = eviction_attempts_.value();
     int64_t evicted_key_count = evicted_key_count_.value();
     int64_t evicted_size = evicted_size_.value();
+    // Mem eviction counters
+    int64_t mem_eviction_success = mem_eviction_success_.value();
+    int64_t mem_eviction_attempts = mem_eviction_attempts_.value();
+    int64_t mem_evicted_key_count = mem_evicted_key_count_.value();
+    int64_t mem_evicted_size = mem_evicted_size_.value();
+    // NoF eviction counters
+    int64_t nof_eviction_success = nof_eviction_success_.value();
+    int64_t nof_eviction_attempts = nof_eviction_attempts_.value();
+    int64_t nof_evicted_key_count = nof_evicted_key_count_.value();
+    int64_t nof_evicted_size = nof_evicted_size_.value();
 
     // Ping counters
     int64_t ping = ping_requests_.value();
@@ -1567,6 +2027,162 @@ std::string MasterMetricManager::get_summary_string() {
     int64_t put_start_discarded_staging_size =
         put_start_discarded_staging_size_.value();
 
+    SummaryCounters current_counters;
+    current_counters.exist_keys = exist_keys;
+    current_counters.exist_key_fails = exist_key_fails;
+    current_counters.put_starts = put_starts;
+    current_counters.put_start_fails = put_start_fails;
+    current_counters.put_start_alloc_fails = put_start_alloc_fails;
+    current_counters.put_ends = put_ends;
+    current_counters.put_end_fails = put_end_fails;
+    current_counters.put_revoke_requests = put_revoke_requests;
+    current_counters.put_revoke_fails = put_revoke_fails;
+    current_counters.get_replicas = get_replicas;
+    current_counters.get_replica_fails = get_replica_fails;
+    current_counters.removes = removes;
+    current_counters.remove_fails = remove_fails;
+    current_counters.remove_all = remove_all;
+    current_counters.remove_all_fails = remove_all_fails;
+    current_counters.create_move_tasks = create_move_tasks;
+    current_counters.create_move_task_fails = create_move_task_fails;
+    current_counters.create_copy_tasks = create_copy_tasks;
+    current_counters.create_copy_task_fails = create_copy_task_fails;
+    current_counters.query_tasks = query_tasks;
+    current_counters.query_task_fails = query_task_fails;
+    current_counters.fetch_tasks = fetch_tasks;
+    current_counters.fetch_task_fails = fetch_task_fails;
+    current_counters.copy_starts = copy_starts;
+    current_counters.copy_start_fails = copy_start_fails;
+    current_counters.copy_ends = copy_ends;
+    current_counters.copy_end_fails = copy_end_fails;
+    current_counters.copy_revokes = copy_revokes;
+    current_counters.copy_revoke_fails = copy_revoke_fails;
+    current_counters.move_starts = move_starts;
+    current_counters.move_start_fails = move_start_fails;
+    current_counters.move_ends = move_ends;
+    current_counters.move_end_fails = move_end_fails;
+    current_counters.move_revokes = move_revokes;
+    current_counters.move_revoke_fails = move_revoke_fails;
+    current_counters.evict_disk_replicas = evict_disk_replicas;
+    current_counters.evict_disk_replica_fails = evict_disk_replica_fails;
+    current_counters.batch_put_start_requests = batch_put_start_requests;
+    current_counters.batch_put_start_fails = batch_put_start_fails;
+    current_counters.batch_put_start_partial_successes =
+        batch_put_start_partial_successes;
+    current_counters.batch_put_start_items = batch_put_start_items;
+    current_counters.batch_put_start_failed_items =
+        batch_put_start_failed_items;
+    current_counters.batch_put_end_requests = batch_put_end_requests;
+    current_counters.batch_put_end_fails = batch_put_end_fails;
+    current_counters.batch_put_end_partial_successes =
+        batch_put_end_partial_successes;
+    current_counters.batch_put_end_items = batch_put_end_items;
+    current_counters.batch_put_end_failed_items = batch_put_end_failed_items;
+    current_counters.batch_put_revoke_requests = batch_put_revoke_requests;
+    current_counters.batch_put_revoke_fails = batch_put_revoke_fails;
+    current_counters.batch_put_revoke_partial_successes =
+        batch_put_revoke_partial_successes;
+    current_counters.batch_put_revoke_items = batch_put_revoke_items;
+    current_counters.batch_put_revoke_failed_items =
+        batch_put_revoke_failed_items;
+    current_counters.batch_get_replica_list_requests =
+        batch_get_replica_list_requests;
+    current_counters.batch_get_replica_list_fails =
+        batch_get_replica_list_fails;
+    current_counters.batch_get_replica_list_partial_successes =
+        batch_get_replica_list_partial_successes;
+    current_counters.batch_get_replica_list_items =
+        batch_get_replica_list_items;
+    current_counters.batch_get_replica_list_failed_items =
+        batch_get_replica_list_failed_items;
+    current_counters.batch_exist_key_requests = batch_exist_key_requests;
+    current_counters.batch_exist_key_fails = batch_exist_key_fails;
+    current_counters.batch_exist_key_partial_successes =
+        batch_exist_key_partial_successes;
+    current_counters.batch_exist_key_items = batch_exist_key_items;
+    current_counters.batch_exist_key_failed_items =
+        batch_exist_key_failed_items;
+    current_counters.batch_query_ip_requests = batch_query_ip_requests;
+    current_counters.batch_query_ip_fails = batch_query_ip_fails;
+    current_counters.batch_query_ip_partial_successes =
+        batch_query_ip_partial_successes;
+    current_counters.batch_query_ip_items = batch_query_ip_items;
+    current_counters.batch_query_ip_failed_items = batch_query_ip_failed_items;
+    current_counters.batch_replica_clear_requests =
+        batch_replica_clear_requests;
+    current_counters.batch_replica_clear_fails = batch_replica_clear_fails;
+    current_counters.batch_replica_clear_partial_successes =
+        batch_replica_clear_partial_successes;
+    current_counters.batch_replica_clear_items = batch_replica_clear_items;
+    current_counters.batch_replica_clear_failed_items =
+        batch_replica_clear_failed_items;
+    current_counters.eviction_success = eviction_success;
+    current_counters.eviction_attempts = eviction_attempts;
+    current_counters.evicted_key_count = evicted_key_count;
+    current_counters.evicted_size = evicted_size;
+    current_counters.mem_eviction_success = mem_eviction_success;
+    current_counters.mem_eviction_attempts = mem_eviction_attempts;
+    current_counters.mem_evicted_key_count = mem_evicted_key_count;
+    current_counters.mem_evicted_size = mem_evicted_size;
+    current_counters.nof_eviction_success = nof_eviction_success;
+    current_counters.nof_eviction_attempts = nof_eviction_attempts;
+    current_counters.nof_evicted_key_count = nof_evicted_key_count;
+    current_counters.nof_evicted_size = nof_evicted_size;
+    current_counters.ping = ping;
+    current_counters.ping_fails = ping_fails;
+    current_counters.mark_task_to_complete_requests =
+        mark_task_to_complete_requests_.value();
+    current_counters.mark_task_to_complete_fails =
+        mark_task_to_complete_failures_.value();
+
+    SummaryCounters previous_counters = current_counters;
+    bool has_previous_summary = false;
+    double elapsed_seconds = 0.0;
+    const auto now = std::chrono::steady_clock::now();
+    {
+        std::lock_guard<std::mutex> lock(summary_snapshot_mutex_);
+        has_previous_summary = summary_snapshot_.initialized;
+        if (has_previous_summary) {
+            previous_counters = summary_snapshot_.counters;
+            elapsed_seconds =
+                std::chrono::duration<double>(now - summary_snapshot_.timestamp)
+                    .count();
+        }
+        if (update_summary_snapshot) {
+            summary_snapshot_.initialized = true;
+            summary_snapshot_.timestamp = now;
+            summary_snapshot_.counters = current_counters;
+        }
+    }
+
+    auto delta = [&](int64_t SummaryCounters::* field) {
+        if (!has_previous_summary) {
+            return int64_t{0};
+        }
+        int64_t value = current_counters.*field - previous_counters.*field;
+        return value > 0 ? value : int64_t{0};
+    };
+    auto rate = [&](int64_t value) {
+        if (elapsed_seconds <= 0.0) {
+            return 0.0;
+        }
+        return static_cast<double>(value) / elapsed_seconds;
+    };
+    auto format_rate_value = [&](int64_t value) {
+        std::ostringstream rate_stream;
+        rate_stream << std::fixed << std::setprecision(2) << rate(value);
+        return rate_stream.str();
+    };
+    auto format_rate_pair = [&](int64_t success, int64_t total) {
+        return format_rate_value(success) + "/" + format_rate_value(total);
+    };
+    auto format_rate_triple = [&](int64_t success, int64_t partial_success,
+                                  int64_t total) {
+        return format_rate_value(success) + "/" +
+               format_rate_value(partial_success) + "/" +
+               format_rate_value(total);
+    };
+
     // --- Format the summary string ---
     ss << "Mem Storage: " << byte_size_to_string(mem_allocated) << " / "
        << byte_size_to_string(mem_capacity);
@@ -1574,109 +2190,241 @@ std::string MasterMetricManager::get_summary_string() {
         ss << " (" << std::fixed << std::setprecision(1)
            << ((double)mem_allocated / (double)mem_capacity * 100.0) << "%)";
     }
+    int64_t file_display_capacity = get_total_file_capacity();
+    ss << " | NVMe-oF SSD: " << byte_size_to_string(nof_allocated) << " / "
+       << byte_size_to_string(nof_capacity);
+    if (nof_capacity > 0) {
+        ss << " (" << std::fixed << std::setprecision(1)
+           << ((double)nof_allocated / (double)nof_capacity * 100.0) << "%)";
+    }
     ss << " | SSD Storage: " << byte_size_to_string(file_allocated) << " / "
-       << byte_size_to_string(file_capacity);
+       << byte_size_to_string(file_display_capacity);
     ss << " | Keys: " << keys << " (soft-pinned: " << soft_pin_keys << ")";
     ss << " | Clients: " << active_clients;
 
-    // Request summary - focus on the most important metrics
-    ss << " | Requests (Success/Total): ";
-    ss << "PutStart=" << put_starts - put_start_fails << "/" << put_starts
+    // Request summary - rate per second of the last window
+    ss << " | Requests (Success/Total per sec): ";
+    ss << "PutStart="
+       << format_rate_pair(delta(&SummaryCounters::put_starts) -
+                               delta(&SummaryCounters::put_start_fails),
+                           delta(&SummaryCounters::put_starts))
        << ", ";
-    ss << "PutEnd=" << put_ends - put_end_fails << "/" << put_ends << ", ";
-    ss << "PutRevoke=" << put_revoke_requests - put_revoke_fails << "/"
-       << put_revoke_requests << ", ";
-    ss << "Get=" << get_replicas - get_replica_fails << "/" << get_replicas
+    ss << "PutEnd="
+       << format_rate_pair(delta(&SummaryCounters::put_ends) -
+                               delta(&SummaryCounters::put_end_fails),
+                           delta(&SummaryCounters::put_ends))
        << ", ";
-    ss << "Exist=" << exist_keys - exist_key_fails << "/" << exist_keys << ", ";
-    ss << "Del=" << removes - remove_fails << "/" << removes << ", ";
-    ss << "DelAll=" << remove_all - remove_all_fails << "/" << remove_all
+    ss << "PutRevoke="
+       << format_rate_pair(delta(&SummaryCounters::put_revoke_requests) -
+                               delta(&SummaryCounters::put_revoke_fails),
+                           delta(&SummaryCounters::put_revoke_requests))
        << ", ";
-    ss << "Ping=" << ping - ping_fails << "/" << ping << ", ";
-    ss << "CopyStart=" << copy_starts - copy_start_fails << "/" << copy_starts
+    ss << "Get="
+       << format_rate_pair(delta(&SummaryCounters::get_replicas) -
+                               delta(&SummaryCounters::get_replica_fails),
+                           delta(&SummaryCounters::get_replicas))
        << ", ";
-    ss << "CopyEnd=" << copy_ends - copy_end_fails << "/" << copy_ends << ", ";
-    ss << "CopyRevoke=" << copy_revokes - copy_revoke_fails << "/"
-       << copy_revokes << ", ";
-    ss << "MoveStart=" << move_starts - move_start_fails << "/" << move_starts
+    ss << "Exist="
+       << format_rate_pair(delta(&SummaryCounters::exist_keys) -
+                               delta(&SummaryCounters::exist_key_fails),
+                           delta(&SummaryCounters::exist_keys))
        << ", ";
-    ss << "MoveEnd=" << move_ends - move_end_fails << "/" << move_ends << ", ";
-    ss << "MoveRevoke=" << move_revokes - move_revoke_fails << "/"
-       << move_revokes << ", ";
-    ss << "EvictDiskReplica=" << evict_disk_replicas - evict_disk_replica_fails
-       << "/" << evict_disk_replicas;
+    ss << "Del="
+       << format_rate_pair(delta(&SummaryCounters::removes) -
+                               delta(&SummaryCounters::remove_fails),
+                           delta(&SummaryCounters::removes))
+       << ", ";
+    ss << "DelAll="
+       << format_rate_pair(delta(&SummaryCounters::remove_all) -
+                               delta(&SummaryCounters::remove_all_fails),
+                           delta(&SummaryCounters::remove_all))
+       << ", ";
+    ss << "Ping="
+       << format_rate_pair(delta(&SummaryCounters::ping) -
+                               delta(&SummaryCounters::ping_fails),
+                           delta(&SummaryCounters::ping))
+       << ", ";
+    ss << "CopyStart="
+       << format_rate_pair(delta(&SummaryCounters::copy_starts) -
+                               delta(&SummaryCounters::copy_start_fails),
+                           delta(&SummaryCounters::copy_starts))
+       << ", ";
+    ss << "CopyEnd="
+       << format_rate_pair(delta(&SummaryCounters::copy_ends) -
+                               delta(&SummaryCounters::copy_end_fails),
+                           delta(&SummaryCounters::copy_ends))
+       << ", ";
+    ss << "CopyRevoke="
+       << format_rate_pair(delta(&SummaryCounters::copy_revokes) -
+                               delta(&SummaryCounters::copy_revoke_fails),
+                           delta(&SummaryCounters::copy_revokes))
+       << ", ";
+    ss << "MoveStart="
+       << format_rate_pair(delta(&SummaryCounters::move_starts) -
+                               delta(&SummaryCounters::move_start_fails),
+                           delta(&SummaryCounters::move_starts))
+       << ", ";
+    ss << "MoveEnd="
+       << format_rate_pair(delta(&SummaryCounters::move_ends) -
+                               delta(&SummaryCounters::move_end_fails),
+                           delta(&SummaryCounters::move_ends))
+       << ", ";
+    ss << "MoveRevoke="
+       << format_rate_pair(delta(&SummaryCounters::move_revokes) -
+                               delta(&SummaryCounters::move_revoke_fails),
+                           delta(&SummaryCounters::move_revokes))
+       << ", ";
+    ss << "EvictDiskReplica="
+       << format_rate_pair(
+              delta(&SummaryCounters::evict_disk_replicas) -
+                  delta(&SummaryCounters::evict_disk_replica_fails),
+              delta(&SummaryCounters::evict_disk_replicas));
 
     // Batch request summary
     ss << " | Batch Requests "
-          "(Req=Success/PartialSuccess/Total, Item=Success/Total): ";
+          "(per sec, Req=Success/PartialSuccess/Total, "
+          "Item=Success/Total): ";
     ss << "PutStart:(Req="
-       << batch_put_start_requests - batch_put_start_fails -
-              batch_put_start_partial_successes
-       << "/" << batch_put_start_partial_successes << "/"
-       << batch_put_start_requests
-       << ", Item=" << batch_put_start_items - batch_put_start_failed_items
-       << "/" << batch_put_start_items << "), ";
+       << format_rate_triple(
+              delta(&SummaryCounters::batch_put_start_requests) -
+                  delta(&SummaryCounters::batch_put_start_fails) -
+                  delta(&SummaryCounters::batch_put_start_partial_successes),
+              delta(&SummaryCounters::batch_put_start_partial_successes),
+              delta(&SummaryCounters::batch_put_start_requests))
+       << ", Item="
+       << format_rate_pair(
+              delta(&SummaryCounters::batch_put_start_items) -
+                  delta(&SummaryCounters::batch_put_start_failed_items),
+              delta(&SummaryCounters::batch_put_start_items))
+       << "), ";
     ss << "PutEnd:(Req="
-       << batch_put_end_requests - batch_put_end_fails -
-              batch_put_end_partial_successes
-       << "/" << batch_put_end_partial_successes << "/"
-       << batch_put_end_requests
-       << ", Item=" << batch_put_end_items - batch_put_end_failed_items << "/"
-       << batch_put_end_items << "), ";
+       << format_rate_triple(
+              delta(&SummaryCounters::batch_put_end_requests) -
+                  delta(&SummaryCounters::batch_put_end_fails) -
+                  delta(&SummaryCounters::batch_put_end_partial_successes),
+              delta(&SummaryCounters::batch_put_end_partial_successes),
+              delta(&SummaryCounters::batch_put_end_requests))
+       << ", Item="
+       << format_rate_pair(
+              delta(&SummaryCounters::batch_put_end_items) -
+                  delta(&SummaryCounters::batch_put_end_failed_items),
+              delta(&SummaryCounters::batch_put_end_items))
+       << "), ";
     ss << "PutRevoke:(Req="
-       << batch_put_revoke_requests - batch_put_revoke_fails -
-              batch_put_revoke_partial_successes
-       << "/" << batch_put_revoke_partial_successes << "/"
-       << batch_put_revoke_requests
-       << ", Item=" << batch_put_revoke_items - batch_put_revoke_failed_items
-       << "/" << batch_put_revoke_items << "), ";
+       << format_rate_triple(
+              delta(&SummaryCounters::batch_put_revoke_requests) -
+                  delta(&SummaryCounters::batch_put_revoke_fails) -
+                  delta(&SummaryCounters::batch_put_revoke_partial_successes),
+              delta(&SummaryCounters::batch_put_revoke_partial_successes),
+              delta(&SummaryCounters::batch_put_revoke_requests))
+       << ", Item="
+       << format_rate_pair(
+              delta(&SummaryCounters::batch_put_revoke_items) -
+                  delta(&SummaryCounters::batch_put_revoke_failed_items),
+              delta(&SummaryCounters::batch_put_revoke_items))
+       << "), ";
     ss << "Get:(Req="
-       << batch_get_replica_list_requests - batch_get_replica_list_fails -
-              batch_get_replica_list_partial_successes
-       << "/" << batch_get_replica_list_partial_successes << "/"
-       << batch_get_replica_list_requests << ", Item="
-       << batch_get_replica_list_items - batch_get_replica_list_failed_items
-       << "/" << batch_get_replica_list_items << "), ";
+       << format_rate_triple(
+              delta(&SummaryCounters::batch_get_replica_list_requests) -
+                  delta(&SummaryCounters::batch_get_replica_list_fails) -
+                  delta(&SummaryCounters::
+                            batch_get_replica_list_partial_successes),
+              delta(&SummaryCounters::batch_get_replica_list_partial_successes),
+              delta(&SummaryCounters::batch_get_replica_list_requests))
+       << ", Item="
+       << format_rate_pair(
+              delta(&SummaryCounters::batch_get_replica_list_items) -
+                  delta(&SummaryCounters::batch_get_replica_list_failed_items),
+              delta(&SummaryCounters::batch_get_replica_list_items))
+       << "), ";
     ss << "ExistKey:(Req="
-       << batch_exist_key_requests - batch_exist_key_fails -
-              batch_exist_key_partial_successes
-       << "/" << batch_exist_key_partial_successes << "/"
-       << batch_exist_key_requests
-       << ", Item=" << batch_exist_key_items - batch_exist_key_failed_items
-       << "/" << batch_exist_key_items << "), ";
+       << format_rate_triple(
+              delta(&SummaryCounters::batch_exist_key_requests) -
+                  delta(&SummaryCounters::batch_exist_key_fails) -
+                  delta(&SummaryCounters::batch_exist_key_partial_successes),
+              delta(&SummaryCounters::batch_exist_key_partial_successes),
+              delta(&SummaryCounters::batch_exist_key_requests))
+       << ", Item="
+       << format_rate_pair(
+              delta(&SummaryCounters::batch_exist_key_items) -
+                  delta(&SummaryCounters::batch_exist_key_failed_items),
+              delta(&SummaryCounters::batch_exist_key_items))
+       << "), ";
     ss << "QueryIp:(Req="
-       << batch_query_ip_requests - batch_query_ip_fails -
-              batch_query_ip_partial_successes
-       << "/" << batch_query_ip_partial_successes << "/"
-       << batch_query_ip_requests
-       << ", Item=" << batch_query_ip_items - batch_query_ip_failed_items << "/"
-       << batch_query_ip_items << "), ";
+       << format_rate_triple(
+              delta(&SummaryCounters::batch_query_ip_requests) -
+                  delta(&SummaryCounters::batch_query_ip_fails) -
+                  delta(&SummaryCounters::batch_query_ip_partial_successes),
+              delta(&SummaryCounters::batch_query_ip_partial_successes),
+              delta(&SummaryCounters::batch_query_ip_requests))
+       << ", Item="
+       << format_rate_pair(
+              delta(&SummaryCounters::batch_query_ip_items) -
+                  delta(&SummaryCounters::batch_query_ip_failed_items),
+              delta(&SummaryCounters::batch_query_ip_items))
+       << "), ";
     ss << "Clear:(Req="
-       << batch_replica_clear_requests - batch_replica_clear_fails -
-              batch_replica_clear_partial_successes
-       << "/" << batch_replica_clear_partial_successes << "/"
-       << batch_replica_clear_requests << ", Item="
-       << batch_replica_clear_items - batch_replica_clear_failed_items << "/"
-       << batch_replica_clear_items << "), ";
+       << format_rate_triple(
+              delta(&SummaryCounters::batch_replica_clear_requests) -
+                  delta(&SummaryCounters::batch_replica_clear_fails) -
+                  delta(
+                      &SummaryCounters::batch_replica_clear_partial_successes),
+              delta(&SummaryCounters::batch_replica_clear_partial_successes),
+              delta(&SummaryCounters::batch_replica_clear_requests))
+       << ", Item="
+       << format_rate_pair(
+              delta(&SummaryCounters::batch_replica_clear_items) -
+                  delta(&SummaryCounters::batch_replica_clear_failed_items),
+              delta(&SummaryCounters::batch_replica_clear_items))
+       << "), ";
 
-    ss << "CreateMoveTask:(Req=" << create_move_tasks - create_move_task_fails
-       << "/" << create_move_tasks << "), ";
-    ss << "CreateCopyTask:(Req=" << create_copy_tasks - create_copy_task_fails
-       << "/" << create_copy_tasks << "), ";
-    ss << "QueryTask=(Req=" << query_tasks - query_task_fails << "/"
-       << query_tasks << "), ";
-    ss << "FetchTasks=(Req=" << fetch_tasks - fetch_task_fails << "/"
-       << fetch_tasks << "), ";
-    ss << "MarkTaskToComplete= (Req="
-       << mark_task_to_complete_requests_.value() -
-              mark_task_to_complete_failures_.value()
-       << "/" << mark_task_to_complete_requests_.value() << "), ";
-    // Eviction summary
+    ss << "CreateMoveTask:(Req="
+       << format_rate_pair(delta(&SummaryCounters::create_move_tasks) -
+                               delta(&SummaryCounters::create_move_task_fails),
+                           delta(&SummaryCounters::create_move_tasks))
+       << "), ";
+    ss << "CreateCopyTask:(Req="
+       << format_rate_pair(delta(&SummaryCounters::create_copy_tasks) -
+                               delta(&SummaryCounters::create_copy_task_fails),
+                           delta(&SummaryCounters::create_copy_tasks))
+       << "), ";
+    ss << "QueryTask:(Req="
+       << format_rate_pair(delta(&SummaryCounters::query_tasks) -
+                               delta(&SummaryCounters::query_task_fails),
+                           delta(&SummaryCounters::query_tasks))
+       << "), ";
+    ss << "FetchTasks:(Req="
+       << format_rate_pair(delta(&SummaryCounters::fetch_tasks) -
+                               delta(&SummaryCounters::fetch_task_fails),
+                           delta(&SummaryCounters::fetch_tasks))
+       << "), ";
+    ss << "MarkTaskToComplete:(Req="
+       << format_rate_pair(
+              delta(&SummaryCounters::mark_task_to_complete_requests) -
+                  delta(&SummaryCounters::mark_task_to_complete_fails),
+              delta(&SummaryCounters::mark_task_to_complete_requests))
+       << ")";
+    // Eviction counters are cumulative. Request counters above are reported as
+    // window rates, but eviction totals are used to track long-running pressure
+    // and should not reset after each admin metrics snapshot.
     ss << " | Eviction: "
        << "Success/Attempts=" << eviction_success << "/" << eviction_attempts
        << ", "
+       << "AllocFail=" << delta(&SummaryCounters::put_start_alloc_fails) << ", "
        << "keys=" << evicted_key_count << ", "
        << "size=" << byte_size_to_string(evicted_size);
+    // mem eviction
+    ss << " | Mem Eviction: "
+       << "Success/Attempts=" << mem_eviction_success << "/"
+       << mem_eviction_attempts << ", "
+       << "keys=" << mem_evicted_key_count << ", "
+       << "size=" << byte_size_to_string(mem_evicted_size);
+    // nof eviction
+    ss << " | NoF Eviction: "
+       << "Success/Attempts=" << nof_eviction_success << "/"
+       << nof_eviction_attempts << ", "
+       << "keys=" << nof_evicted_key_count << ", "
+       << "size=" << byte_size_to_string(nof_evicted_size);
 
     // Discard summary
     ss << " | Discard: "
@@ -1684,6 +2432,21 @@ std::string MasterMetricManager::get_summary_string() {
        << put_start_discard_cnt << ", StagingSize="
        << byte_size_to_string(put_start_discarded_staging_size);
 
+    // Promotion-on-hit summary (counters are cumulative, not deltas; the
+    // gauge is current-state).
+    ss << " | Promotion: "
+       << "in_flight=" << promotion_in_flight_metric_.value() << ", "
+       << "admitted=" << promotion_admitted_.value() << ", "
+       << "completed=" << promotion_completed_.value() << ", "
+       << "failed=" << promotion_failed_.value() << ", "
+       << "cancelled=" << promotion_cancelled_.value() << ", "
+       << "expired=" << promotion_expired_.value() << ", "
+       << "bytes=" << byte_size_to_string(promotion_completed_bytes_.value())
+       << ", "
+       << "rejected(freq/wm/cap)=" << promotion_rejected_frequency_.value()
+       << "/" << promotion_rejected_watermark_.value() << "/"
+       << promotion_rejected_cap_.value();
+
     // Snapshot summary
     ss << " | Snapshots: "
        << "Success=" << snapshot_success_.value() << ", "
diff --git a/mooncake-store/src/master_service.cpp b/mooncake-store/src/master_service.cpp
index 799a3cdc..034edbb6 100644
--- a/mooncake-store/src/master_service.cpp
+++ b/mooncake-store/src/master_service.cpp
@@ -1,9 +1,15 @@
 #include "master_service.h"
 
+#include <array>
 #include <cassert>
+#include <cmath>
 #include <cstdint>
+#include <cstring>
 #include <limits>
+#include <random>
 #include <shared_mutex>
+#include <sstream>
+#include <thread>
 #include <regex>
 #include <unordered_set>
 #include <unistd.h>
@@ -12,8 +18,16 @@
 #include <ylt/util/tl/expected.hpp>
 #include <boost/algorithm/string.hpp>
 
+#include "http_metadata_server.h"
 #include "master_metric_manager.h"
+#include "common.h"
 #include "segment.h"
+#ifdef USE_HTTP
+#include "transfer_metadata_plugin.h"
+#endif
+#ifdef USE_NOF
+#include "spdk/spdk_wrapper.h"
+#endif
 #ifdef STORE_USE_ETCD
 #include "etcd_helper.h"
 #include "ha/oplog/etcd_oplog_store.h"
@@ -36,7 +50,6 @@ static const std::string SNAPSHOT_SEGMENTS_FILE = "segments";
 static const std::string SNAPSHOT_TASK_MANAGER_FILE = "task_manager";
 static const std::string SNAPSHOT_MANIFEST_FILE = "manifest.txt";
 static const std::string SNAPSHOT_LATEST_FILE = "latest.txt";
-static const std::string SNAPSHOT_ROOT = "mooncake_master_snapshot";
 static const std::string SNAPSHOT_BACKUP_SAVE_DIR =
     "mooncake_snapshot_save_backup";
 static const std::string SNAPSHOT_BACKUP_RESTORE_DIR =
@@ -47,13 +60,14 @@ static const std::string SNAPSHOT_SERIALIZER_TYPE = "messagepack";
 namespace {
 
 constexpr size_t kUnlimitedSnapshotList = 0;
+constexpr int kMaxTenantQuotaEvictionRetries = 2;
 
 // Per-cycle offload cap as a fraction of `offloading_queue_limit_`. Used only
 // when offload-on-evict mode is active. Defers memory eviction for at most
 // this fraction of the queue limit per BatchEvict cycle; beyond that, eviction
-// falls back according to `offload_force_evict_`. A future change may expose
-// this as a configurable parameter if workloads demand tuning.
-constexpr double kOffloadCapRatio = 0.5;
+// falls back according to `offload_force_evict_`.
+// NOTE: Both offloading_queue_limit_ and offload_cap_ratio_ are now
+// configurable via --offloading_queue_limit and --offload_cap_ratio flags.
 
 enum class SnapshotCatalogBackendKind {
     kEmbedded,
@@ -79,18 +93,114 @@ int64_t CurrentTimeMs() {
         .count();
 }
 
+size_t RandomIndex(size_t upper_bound) {
+    static thread_local std::mt19937 generator(std::random_device{}());
+    std::uniform_int_distribution<size_t> dist(0, upper_bound - 1);
+    return dist(generator);
+}
+
+uint64_t SaturatingAdd(uint64_t lhs, uint64_t rhs) {
+    if (lhs > std::numeric_limits<uint64_t>::max() - rhs) {
+        return std::numeric_limits<uint64_t>::max();
+    }
+    return lhs + rhs;
+}
+
+uint64_t SaturatingMultiply(uint64_t lhs, uint64_t rhs) {
+    if (lhs != 0 && rhs > std::numeric_limits<uint64_t>::max() / lhs) {
+        return std::numeric_limits<uint64_t>::max();
+    }
+    return lhs * rhs;
+}
+
+bool HasExpectedReplicaAllocation(const ReplicateConfig& config,
+                                  size_t allocated_memory_replicas,
+                                  size_t allocated_nof_replicas) {
+    if (config.nof_replica_num == 0) {
+        return allocated_memory_replicas > 0;
+    }
+    if (DetermineReplicaWriteMode(config) ==
+        ReplicaWriteMode::FLEXIBLE_DUAL_REPLICA) {
+        return allocated_memory_replicas + allocated_nof_replicas > 0;
+    }
+    return allocated_memory_replicas == config.replica_num &&
+           allocated_nof_replicas == config.nof_replica_num;
+}
+
+bool IsLazyEmptyTenantQuotaState(const TenantQuotaState& state) {
+    return !state.has_explicit_policy && state.used_bytes == 0 &&
+           state.reserved_bytes == 0 && state.committed_count == 0 &&
+           state.metadata_object_count == 0;
+}
+
+void RefreshTenantQuotaOverQuota(TenantQuotaState& state) {
+    state.over_quota =
+        (!state.has_explicit_policy && state.metadata_object_count > 0) ||
+        static_cast<unsigned __int128>(state.used_bytes) +
+                state.reserved_bytes >
+            state.effective_quota_bytes;
+}
+
+TenantQuotaSnapshot MakeTenantQuotaSnapshot(const std::string& tenant_id,
+                                            const TenantQuotaState& state) {
+    return TenantQuotaSnapshot{
+        .tenant_id = tenant_id,
+        .requested_quota_bytes = state.requested_quota_bytes,
+        .effective_quota_bytes = state.effective_quota_bytes,
+        .used_bytes = state.used_bytes,
+        .reserved_bytes = state.reserved_bytes,
+        .committed_count = state.committed_count,
+        .metadata_object_count = state.metadata_object_count,
+        .has_explicit_policy = state.has_explicit_policy,
+        .over_quota = state.over_quota};
+}
+
+tl::expected<std::string, ErrorCode> GetGroupIdForKey(
+    const ReplicateConfig& config, size_t key_count, size_t key_index) {
+    if (!config.group_ids.has_value()) {
+        return "";
+    }
+    if (config.group_ids->size() != key_count || key_index >= key_count) {
+        LOG(ERROR) << "group_ids.size()=" << config.group_ids->size()
+                   << ", key_count=" << key_count
+                   << ", error=invalid_group_ids";
+        return tl::make_unexpected(ErrorCode::INVALID_PARAMS);
+    }
+    return config.group_ids->at(key_index);
+}
+
 }  // namespace
 
 MasterService::MasterService() : MasterService(MasterServiceConfig()) {}
 
 MasterService::MasterService(const MasterServiceConfig& config)
-    : default_kv_lease_ttl_(config.default_kv_lease_ttl),
+    : graceful_unmount_scheduler_(
+          [this](const GracefulUnmountDeadlineRecord& record) {
+              auto result =
+                  this->UnmountSegment(record.segment_id, record.client_id);
+              if (!result.has_value()) {
+                  LOG(WARNING)
+                      << "Failed to complete graceful unmount, segment_id="
+                      << record.segment_id << ", client_id=" << record.client_id
+                      << ", error=" << toString(result.error());
+              }
+          }),
+      default_kv_lease_ttl_(config.default_kv_lease_ttl),
       default_kv_soft_pin_ttl_(config.default_kv_soft_pin_ttl),
       allow_evict_soft_pinned_objects_(config.allow_evict_soft_pinned_objects),
       eviction_ratio_(config.eviction_ratio),
       eviction_high_watermark_ratio_(config.eviction_high_watermark_ratio),
+      nof_eviction_ratio_(config.nof_eviction_ratio),
+      nof_eviction_high_watermark_ratio_(
+          config.nof_eviction_high_watermark_ratio),
       view_version_(config.view_version),
       client_live_ttl_sec_(config.client_live_ttl_sec),
+      nof_heartbeat_interval_sec_(
+          std::chrono::seconds(config.nof_heartbeat_interval_sec)),
+      nof_heartbeat_probe_timeout_ms_(
+          std::chrono::milliseconds(config.nof_heartbeat_probe_timeout_ms)),
+      nof_heartbeat_failures_threshold_(
+          config.nof_heartbeat_failures_threshold),
       enable_ha_(config.enable_ha),
       enable_offload_(config.enable_offload),
       ha_backend_type_(config.ha_backend_type),
@@ -100,10 +210,16 @@ MasterService::MasterService(const MasterServiceConfig& config)
       global_file_segment_size_(config.global_file_segment_size),
       enable_disk_eviction_(config.enable_disk_eviction),
       quota_bytes_(config.quota_bytes),
+      enable_multi_tenants_(config.enable_multi_tenants),
+      tenant_quota_connector_type_(config.tenant_quota_connector_type),
+      tenant_quota_connector_uri_(config.tenant_quota_connector_uri),
       segment_manager_(config.memory_allocator, config.enable_cxl),
+      nof_segment_manager_(config.memory_allocator),
       memory_allocator_type_(config.memory_allocator),
-      allocation_strategy_(
-          CreateAllocationStrategy(config.allocation_strategy_type)),
+      allocation_strategy_type_(config.enable_cxl
+                                    ? AllocationStrategyType::CXL
+                                    : config.allocation_strategy_type),
+      allocation_strategy_(CreateAllocationStrategy(allocation_strategy_type_)),
       enable_snapshot_restore_(config.enable_snapshot_restore),
       enable_snapshot_(config.enable_snapshot),
       snapshot_backup_dir_(config.snapshot_backup_dir),
@@ -118,7 +234,39 @@ MasterService::MasterService(const MasterServiceConfig& config)
       task_manager_(config.task_manager_config),
       cxl_path_(config.cxl_path),
       cxl_size_(config.cxl_size),
-      enable_cxl_(config.enable_cxl) {
+      enable_cxl_(config.enable_cxl),
+      offloading_queue_limit_(config.offloading_queue_limit),
+      offload_cap_ratio_(config.offload_cap_ratio) {
+    object_type_reuse_scales_.fill(1.0);
+    object_type_soft_pin_weights_.fill(1.0);
+    object_type_eviction_graces_.fill(0);
+    object_type_budget_ratios_.fill(1.0);
+    for (const auto& [data_type, policy] :
+         config.object_type_eviction_score_policies) {
+        ValidateObjectTypeEvictionScorePolicy(policy);
+        const auto idx = static_cast<uint8_t>(data_type);
+        object_type_reuse_scales_[idx] = policy.reuse_scale;
+        object_type_soft_pin_weights_[idx] = policy.soft_pin_weight;
+        object_type_eviction_graces_[idx] = policy.eviction_grace;
+    }
+    for (const auto& [data_type, policy] :
+         config.object_type_eviction_policies) {
+        ValidateObjectTypeEvictionPolicy(policy);
+        const auto idx = static_cast<uint8_t>(data_type);
+        object_type_budget_ratios_[idx] = policy.budget_ratio;
+    }
+
+    // Initialize HTTP metadata key prefix (read env var once at startup)
+    const char* custom_prefix = std::getenv("MC_METADATA_CLUSTER_ID");
+    if (custom_prefix && std::strlen(custom_prefix) > 0) {
+        http_metadata_prefix_ = "mooncake/" + std::string(custom_prefix);
+        if (http_metadata_prefix_.back() != '/') {
+            http_metadata_prefix_ += '/';
+        }
+    } else {
+        http_metadata_prefix_ = "mooncake/";
+    }
+
     if (enable_snapshot_ || enable_snapshot_restore_) {
         try {
             auto object_store_type =
@@ -136,9 +284,22 @@ MasterService::MasterService(const MasterServiceConfig& config)
         }
     }
 
+    if (enable_multi_tenants_) {
+        auto store = CreateTenantQuotaPolicyStore(tenant_quota_connector_type_,
+                                                  tenant_quota_connector_uri_);
+        if (!store) {
+            throw std::invalid_argument(store.error());
+        }
+        tenant_quota_policy_store_ = std::move(store.value());
+    }
+
     if (enable_snapshot_restore_) {
         RestoreState();
     }
+    if (enable_multi_tenants_) {
+        LoadTenantQuotaPoliciesFromStoreOrThrow();
+        RebuildTenantQuotaUsageFromMetadata();
+    }
     if (enable_snapshot_ && snapshot_retention_count_ == 0) {
         LOG(ERROR) << "snapshot_retention_count must be greater than 0";
         throw std::invalid_argument("snapshot_retention_count must be > 0");
@@ -156,6 +317,25 @@ MasterService::MasterService(const MasterServiceConfig& config)
         throw std::invalid_argument("Invalid eviction high watermark ratio");
     }
 
+    // Validate offload tuning knobs here (not only via gflags validator),
+    // because values loaded from a configuration file bypass the gflags
+    // validator chain.
+    if (offload_cap_ratio_ < 0.0 || offload_cap_ratio_ > 1.0) {
+        LOG(ERROR) << "offload_cap_ratio must be between 0.0 and 1.0, "
+                   << "current value: " << offload_cap_ratio_;
+        throw std::invalid_argument("Invalid offload_cap_ratio");
+    }
+    if (offloading_queue_limit_ == 0) {
+        LOG(ERROR) << "offloading_queue_limit must be greater than 0";
+        throw std::invalid_argument("Invalid offloading_queue_limit");
+    }
+    if (offloading_queue_limit_ > 100'000'000ULL) {
+        LOG(ERROR) << "offloading_queue_limit must be <= 100000000 to avoid "
+                   << "overflow when computing offload_cap, current value: "
+                   << offloading_queue_limit_;
+        throw std::invalid_argument("Invalid offloading_queue_limit");
+    }
+
     if (put_start_release_timeout_sec_ <= put_start_discard_timeout_sec_) {
         LOG(ERROR) << "put_start_release_timeout="
                    << put_start_release_timeout_sec_.count()
@@ -166,6 +346,29 @@ MasterService::MasterService(const MasterServiceConfig& config)
             "put_start_discard_timeout_sec");
     }
 
+#ifdef USE_NOF
+    if (nof_heartbeat_interval_sec_.count() <= 0) {
+        LOG(ERROR) << "nof_heartbeat_interval_sec must be positive, current "
+                   << nof_heartbeat_interval_sec_.count();
+        throw std::invalid_argument("Invalid nof heartbeat interval");
+    }
+    if (nof_heartbeat_probe_timeout_ms_.count() <= 0) {
+        LOG(ERROR) << "nof_heartbeat_probe_timeout_ms must be positive, "
+                   << "current " << nof_heartbeat_probe_timeout_ms_.count();
+        throw std::invalid_argument("Invalid nof heartbeat probe timeout");
+    }
+    if (nof_heartbeat_failures_threshold_ == 0) {
+        LOG(ERROR) << "nof_heartbeat_failures_threshold must be positive";
+        throw std::invalid_argument("Invalid nof heartbeat failure threshold");
+    }
+
+    nof_probe_fn_ = [](const std::string& te_endpoint, uint32_t timeout_ms,
+                       std::string* error_reason) {
+        return SpdkWrapper::GetInstance().ProbeNofSegment(
+            te_endpoint, timeout_ms, error_reason);
+    };
+#endif
+
     // Offload-on-evict: defer LOCAL_DISK offload to eviction time
     offload_on_evict_ = enable_offload_ && config.offload_on_evict;
     if (offload_on_evict_) {
@@ -179,6 +382,45 @@ MasterService::MasterService(const MasterServiceConfig& config)
         }
     }
 
+    // Promotion-on-hit: when Get observes a LOCAL_DISK-only key, queue an
+    // async copy back to MEMORY. Only meaningful when offload is enabled
+    // (otherwise no LOCAL_DISK replicas exist in the first place).
+    promotion_on_hit_ = enable_offload_ && config.promotion_on_hit;
+    promotion_admission_threshold_ = config.promotion_admission_threshold;
+    promotion_queue_limit_ = config.promotion_queue_limit;
+    promotion_max_per_heartbeat_ = config.promotion_max_per_heartbeat;
+    // Clamp to >=1: 0 would make PromotionObjectHeartbeat return an empty
+    // batch every call, silently disabling promotion delivery.
+    if (promotion_max_per_heartbeat_ == 0) {
+        promotion_max_per_heartbeat_ = 1;
+    }
+    // Defense-in-depth clamp: master.cpp clamps threshold into [1, 255]
+    // at flag-parse time, but direct MasterServiceConfig construction
+    // (tests, embedded users) bypasses that. Without the clamp here,
+    // threshold=0 would silently bypass the frequency gate entirely
+    // (freq < 0 is never true for uint8_t).
+    if (promotion_admission_threshold_ == 0) {
+        promotion_admission_threshold_ = 1;
+    } else if (promotion_admission_threshold_ > 255) {
+        promotion_admission_threshold_ = 255;
+    }
+    if (config.promotion_on_hit && !enable_offload_) {
+        LOG(WARNING) << "promotion_on_hit=true was requested but "
+                     << "enable_offload=false; promotion is silently "
+                     << "disabled because it requires offload to produce "
+                     << "LOCAL_DISK replicas. Set enable_offload=true to "
+                     << "use this feature.";
+    }
+    if (promotion_on_hit_) {
+        promotion_sketch_ = std::make_unique<CountMinSketch>();
+        LOG(INFO) << "Promotion-on-hit mode enabled: LOCAL_DISK-only Gets "
+                     "will queue async promotion to MEMORY (threshold="
+                  << promotion_admission_threshold_
+                  << ", queue_limit=" << promotion_queue_limit_
+                  << ", max_per_heartbeat=" << promotion_max_per_heartbeat_
+                  << ")";
+    }
+
     eviction_running_ = true;
     eviction_thread_ = std::thread(&MasterService::EvictionThreadFunc, this);
     VLOG(1) << "action=start_eviction_thread";
@@ -189,12 +431,24 @@ MasterService::MasterService(const MasterServiceConfig& config)
         std::thread(&MasterService::ClientMonitorFunc, this);
     VLOG(1) << "action=start_client_monitor_thread";
 
+#ifdef USE_NOF
+    nof_heartbeat_running_ = true;
+    nof_heartbeat_thread_ =
+        std::thread(&MasterService::NofHeartbeatThreadFunc, this);
+    VLOG(1) << "action=start_nof_heartbeat_thread";
+#endif
+
     // Start task cleanup thread
     task_cleanup_running_ = true;
     task_cleanup_thread_ =
         std::thread(&MasterService::TaskCleanupThreadFunc, this);
     VLOG(1) << "action=start_task_cleanup_thread";
 
+    // NOTE: The async HTTP metadata cleanup worker is started lazily in
+    // setHttpMetadataRemoteUrl() once http_metadata_remote_ is initialized,
+    // since that happens after this constructor returns (in
+    // WrappedMasterService).
+
     job_dispatch_running_ = true;
     job_dispatch_thread_ =
         std::thread(&MasterService::JobDispatchThreadFunc, this);
@@ -202,8 +456,12 @@ MasterService::MasterService(const MasterServiceConfig& config)
 
     if (!root_fs_dir_.empty()) {
         use_disk_replica_ = true;
-        MasterMetricManager::instance().inc_total_file_capacity(
-            global_file_segment_size_);
+        if (global_file_segment_size_ == std::numeric_limits<int64_t>::max()) {
+            MasterMetricManager::instance().set_dfs_capacity_unlimited(true);
+        } else {
+            MasterMetricManager::instance().inc_total_file_capacity(
+                global_file_segment_size_);
+        }
     }
 
     if (enable_snapshot_) {
@@ -232,7 +490,7 @@ MasterService::CreateSnapshotCatalogStore() {
         case SnapshotCatalogBackendKind::kEmbedded:
             return std::make_unique<
                 ha::backends::embedded::EmbeddedSnapshotCatalogStore>(
-                snapshot_object_store_.get());
+                snapshot_object_store_.get(), cluster_id_);
         case SnapshotCatalogBackendKind::kRedis: {
 #ifndef STORE_USE_REDIS
             throw std::invalid_argument(
@@ -261,12 +519,22 @@ MasterService::~MasterService() {
     // Stop and join the threads
     eviction_running_ = false;
     client_monitor_running_ = false;
-    snapshot_running_ = false;
+    {
+        std::lock_guard<std::mutex> lk(snapshot_thread_mutex_);
+        snapshot_running_ = false;
+    }
     task_cleanup_running_ = false;
     job_dispatch_running_ = false;
+    http_metadata_cleanup_running_ = false;
+    graceful_unmount_scheduler_.Stop();
+#ifdef USE_NOF
+    nof_heartbeat_running_ = false;
+#endif
 
     // Wake sleepers so join() doesn't block for long sleep intervals.
+    snapshot_thread_cv_.notify_all();
     task_cleanup_cv_.notify_all();
+    http_metadata_cleanup_cv_.notify_all();
 
     if (eviction_thread_.joinable()) {
         eviction_thread_.join();
@@ -274,48 +542,285 @@ MasterService::~MasterService() {
     if (client_monitor_thread_.joinable()) {
         client_monitor_thread_.join();
     }
+#ifdef USE_NOF
+    if (nof_heartbeat_thread_.joinable()) {
+        nof_heartbeat_thread_.join();
+    }
+#endif
     if (snapshot_thread_.joinable()) {
         snapshot_thread_.join();
     }
     if (task_cleanup_thread_.joinable()) {
         task_cleanup_thread_.join();
     }
+    if (http_metadata_cleanup_thread_.joinable()) {
+        http_metadata_cleanup_thread_.join();
+    }
     if (job_dispatch_thread_.joinable()) {
         job_dispatch_thread_.join();
     }
 }
 
+void MasterService::SetNoFProbeFnForTesting(NoFProbeFn fn) {
+#ifdef USE_NOF
+    std::lock_guard<std::mutex> lock(nof_probe_fn_mutex_);
+    if (fn) {
+        nof_probe_fn_ = std::move(fn);
+        return;
+    }
+    nof_probe_fn_ = [](const std::string& te_endpoint, uint32_t timeout_ms,
+                       std::string* error_reason) {
+        return SpdkWrapper::GetInstance().ProbeNofSegment(
+            te_endpoint, timeout_ms, error_reason);
+    };
+#else
+    (void)fn;
+#endif
+}
+
+size_t MasterService::GetMountedNoFSegmentCountForTesting() {
+    std::vector<MountedNoFSegmentSnapshot> mounted_segments;
+    nof_segment_manager_.GetMountedSegmentsSnapshot(mounted_segments);
+    return mounted_segments.size();
+}
+
+bool MasterService::IsNoFSegmentMountedForTesting(const UUID& segment_id) {
+    std::vector<MountedNoFSegmentSnapshot> mounted_segments;
+    nof_segment_manager_.GetMountedSegmentsSnapshot(mounted_segments);
+    return std::any_of(
+        mounted_segments.begin(), mounted_segments.end(),
+        [&segment_id](const MountedNoFSegmentSnapshot& snapshot) {
+            return snapshot.segment_id == segment_id &&
+                   snapshot.status == SegmentStatus::OK;
+        });
+}
+
+std::optional<uint32_t> MasterService::GetNoFHeartbeatFailureCountForTesting(
+    const UUID& segment_id) {
+    std::lock_guard<std::mutex> lock(nof_heartbeat_mutex_);
+    auto it = nof_heartbeat_states_.find(segment_id);
+    if (it == nof_heartbeat_states_.end()) {
+        return std::nullopt;
+    }
+    return it->second.consecutive_failures;
+}
+
+std::optional<TenantQuotaSnapshot>
+MasterService::GetTenantQuotaSnapshotForTesting(
+    const std::string& tenant_id) const {
+    const auto normalized_tenant = NormalizeTenantId(tenant_id);
+    const auto shard_idx = getTenantQuotaShardIndex(normalized_tenant);
+    const auto& shard = tenant_quota_shards_[shard_idx];
+    std::lock_guard<std::mutex> lock(shard.mutex);
+    auto it = shard.tenants.find(normalized_tenant);
+    if (it == shard.tenants.end()) {
+        return std::nullopt;
+    }
+    return MakeTenantQuotaSnapshot(normalized_tenant, it->second);
+}
+
+bool MasterService::IsTenantQuotaEnabled() const {
+    return enable_multi_tenants_;
+}
+
+std::vector<TenantQuotaSnapshot> MasterService::ListTenantQuotaSnapshots()
+    const {
+    std::vector<TenantQuotaSnapshot> snapshots;
+    for (size_t i = 0; i < kNumTenantQuotaShards; ++i) {
+        const auto& shard = tenant_quota_shards_[i];
+        std::lock_guard<std::mutex> lock(shard.mutex);
+        for (const auto& [tenant_id, state] : shard.tenants) {
+            if (IsLazyEmptyTenantQuotaState(state)) {
+                continue;
+            }
+            snapshots.push_back(MakeTenantQuotaSnapshot(tenant_id, state));
+        }
+    }
+    std::sort(snapshots.begin(), snapshots.end(),
+              [](const auto& lhs, const auto& rhs) {
+                  return lhs.tenant_id < rhs.tenant_id;
+              });
+    return snapshots;
+}
+
+std::optional<TenantQuotaSnapshot> MasterService::GetTenantQuotaSnapshot(
+    const std::string& tenant_id) const {
+    return GetTenantQuotaSnapshotForTesting(tenant_id);
+}
+
+tl::expected<TenantQuotaSnapshot, ErrorCode>
+MasterService::UpsertTenantQuotaPolicy(const std::string& tenant_id,
+                                       uint64_t requested_quota_bytes) {
+    if (!enable_multi_tenants_) {
+        return tl::make_unexpected(ErrorCode::UNAVAILABLE_IN_CURRENT_MODE);
+    }
+    if (requested_quota_bytes == 0) {
+        return tl::make_unexpected(ErrorCode::INVALID_PARAMS);
+    }
+
+    if (tenant_id.empty()) {
+        return tl::make_unexpected(ErrorCode::INVALID_PARAMS);
+    }
+    const auto normalized_tenant = NormalizeTenantId(tenant_id);
+    if (!IsValidTenantId(normalized_tenant)) {
+        return tl::make_unexpected(ErrorCode::INVALID_PARAMS);
+    }
+
+    std::lock_guard<std::mutex> policy_lock(tenant_quota_policy_mutex_);
+    auto policy = BuildTenantQuotaPolicySnapshot();
+    policy.tenant_quotas[normalized_tenant] = requested_quota_bytes;
+    auto save_result = tenant_quota_policy_store_->Save(policy);
+    if (!save_result) {
+        LOG(ERROR) << "failed to save tenant quota policy: "
+                   << save_result.error();
+        return tl::make_unexpected(ErrorCode::PERSISTENT_FAIL);
+    }
+    ApplyTenantQuotaPolicies(policy);
+    auto result_snapshot = GetTenantQuotaSnapshot(normalized_tenant);
+    if (!result_snapshot.has_value()) {
+        return tl::make_unexpected(ErrorCode::INTERNAL_ERROR);
+    }
+    return result_snapshot.value();
+}
+
+tl::expected<std::optional<TenantQuotaSnapshot>, ErrorCode>
+MasterService::DeleteTenantQuotaPolicy(const std::string& tenant_id) {
+    if (!enable_multi_tenants_) {
+        return tl::make_unexpected(ErrorCode::UNAVAILABLE_IN_CURRENT_MODE);
+    }
+
+    if (tenant_id.empty()) {
+        return tl::make_unexpected(ErrorCode::INVALID_PARAMS);
+    }
+    const auto normalized_tenant = NormalizeTenantId(tenant_id);
+    if (!IsValidTenantId(normalized_tenant)) {
+        return tl::make_unexpected(ErrorCode::INVALID_PARAMS);
+    }
+
+    std::lock_guard<std::mutex> policy_lock(tenant_quota_policy_mutex_);
+    auto policy = BuildTenantQuotaPolicySnapshot();
+    auto policy_it = policy.tenant_quotas.find(normalized_tenant);
+    if (policy_it == policy.tenant_quotas.end()) {
+        return tl::make_unexpected(ErrorCode::OBJECT_NOT_FOUND);
+    }
+    const uint64_t requested_quota_bytes = policy_it->second;
+
+    auto restore_policy = [&] {
+        auto& shard =
+            tenant_quota_shards_[getTenantQuotaShardIndex(normalized_tenant)];
+        {
+            std::lock_guard<std::mutex> lock(shard.mutex);
+            auto& state = shard.tenants[normalized_tenant];
+            state.requested_quota_bytes = requested_quota_bytes;
+            state.has_explicit_policy = true;
+            RefreshTenantQuotaOverQuota(state);
+        }
+        RecomputeTenantEffectiveQuotas();
+    };
+
+    {
+        auto& shard =
+            tenant_quota_shards_[getTenantQuotaShardIndex(normalized_tenant)];
+        std::lock_guard<std::mutex> lock(shard.mutex);
+        auto quota_it = shard.tenants.find(normalized_tenant);
+        if (quota_it == shard.tenants.end() ||
+            !quota_it->second.has_explicit_policy) {
+            return tl::make_unexpected(ErrorCode::OBJECT_NOT_FOUND);
+        }
+        auto& state = quota_it->second;
+        if (state.used_bytes != 0 || state.reserved_bytes != 0 ||
+            state.committed_count != 0 || state.metadata_object_count != 0) {
+            return tl::make_unexpected(ErrorCode::TENANT_NOT_EMPTY);
+        }
+        state.has_explicit_policy = false;
+        state.requested_quota_bytes = 0;
+        state.effective_quota_bytes = 0;
+        RefreshTenantQuotaOverQuota(state);
+    }
+
+    auto post_mark_snapshot = GetTenantQuotaSnapshot(normalized_tenant);
+    if (TenantHasObjects(normalized_tenant) ||
+        (post_mark_snapshot.has_value() &&
+         (post_mark_snapshot->used_bytes != 0 ||
+          post_mark_snapshot->reserved_bytes != 0 ||
+          post_mark_snapshot->committed_count != 0 ||
+          post_mark_snapshot->metadata_object_count != 0))) {
+        restore_policy();
+        return tl::make_unexpected(ErrorCode::TENANT_NOT_EMPTY);
+    }
+
+    policy.tenant_quotas.erase(policy_it);
+    auto save_result = tenant_quota_policy_store_->Save(policy);
+    if (!save_result) {
+        restore_policy();
+        LOG(ERROR) << "failed to save tenant quota policy: "
+                   << save_result.error();
+        return tl::make_unexpected(ErrorCode::PERSISTENT_FAIL);
+    }
+    ApplyTenantQuotaPolicies(policy);
+    return GetTenantQuotaSnapshot(normalized_tenant);
+}
+
 auto MasterService::MountSegment(const Segment& segment, const UUID& client_id)
     -> tl::expected<void, ErrorCode> {
     std::shared_lock<std::shared_mutex> shared_lock(snapshot_mutex_);
-    ScopedSegmentAccess segment_access = segment_manager_.getSegmentAccess();
-
-    // Tell the client monitor thread to start timing for this client. To
-    // avoid the following undesired situations, this message must be sent
-    // after locking the segment mutex and before the mounting operation
-    // completes:
-    // 1. Sending the message before the lock: the client expires and
-    // unmouting invokes before this mounting are completed, which prevents
-    // this segment being able to be unmounted forever;
-    // 2. Sending the message after mounting the segment: After mounting
-    // this segment, when trying to push id to the queue, the queue is
-    // already full. However, at this point, the message must be sent,
-    // otherwise this client cannot be monitored and expired.
     {
-        PodUUID pod_client_id;
-        pod_client_id.first = client_id.first;
-        pod_client_id.second = client_id.second;
-        if (!client_ping_queue_.push(pod_client_id)) {
-            LOG(ERROR) << "segment_name=" << segment.name
-                       << ", error=client_ping_queue_full";
-            return tl::make_unexpected(ErrorCode::INTERNAL_ERROR);
+        ScopedSegmentAccess segment_access =
+            segment_manager_.getSegmentAccess();
+
+        // Tell the client monitor thread to start timing for this client. To
+        // avoid the following undesired situations, this message must be sent
+        // after locking the segment mutex and before the mounting operation
+        // completes:
+        // 1. Sending the message before the lock: the client expires and
+        // unmouting invokes before this mounting are completed, which prevents
+        // this segment being able to be unmounted forever;
+        // 2. Sending the message after mounting the segment: After mounting
+        // this segment, when trying to push id to the queue, the queue is
+        // already full. However, at this point, the message must be sent,
+        // otherwise this client cannot be monitored and expired.
+        {
+            PodUUID pod_client_id;
+            pod_client_id.first = client_id.first;
+            pod_client_id.second = client_id.second;
+            if (!client_ping_queue_.push(pod_client_id)) {
+                LOG(ERROR) << "segment_name=" << segment.name
+                           << ", error=client_ping_queue_full";
+                return tl::make_unexpected(ErrorCode::INTERNAL_ERROR);
+            }
+        }
+
+        LOG(INFO) << "client_id=" << client_id
+                  << ", action=mount_segment, segment_name=" << segment.name;
+
+        auto err = segment_access.MountSegment(segment, client_id);
+        if (err == ErrorCode::SEGMENT_ALREADY_EXISTS) {
+            // Return OK because this is an idempotent operation
+            return {};
+        } else if (err != ErrorCode::OK) {
+            return tl::make_unexpected(err);
         }
     }
+    RecomputeTenantEffectiveQuotas();
+    return {};
+}
 
-    LOG(INFO) << "client_id=" << client_id
+auto MasterService::MountNoFSegment(const NoFSegment& segment,
+                                    const UUID& client_id)
+    -> tl::expected<void, ErrorCode> {
+#ifndef USE_NOF
+    LOG(ERROR) << "client_id=" << client_id << ", segment_name=" << segment.name
+               << ", error=nof_pool_disabled";
+    return tl::make_unexpected(ErrorCode::UNAVAILABLE_IN_CURRENT_MODE);
+#else
+    ScopedNoFSegmentAccess nof_segment_access =
+        nof_segment_manager_.getNoFSegmentAccess();
+
+    LOG(INFO) << "NoF segment mount: "
+              << "client_id=" << client_id
               << ", action=mount_segment, segment_name=" << segment.name;
 
-    auto err = segment_access.MountSegment(segment, client_id);
+    auto err = nof_segment_access.MountSegment(segment, client_id);
     if (err == ErrorCode::SEGMENT_ALREADY_EXISTS) {
         // Return OK because this is an idempotent operation
         return {};
@@ -323,52 +828,79 @@ auto MasterService::MountSegment(const Segment& segment, const UUID& client_id)
         return tl::make_unexpected(err);
     }
     return {};
+#endif
 }
 
 auto MasterService::ReMountSegment(const std::vector<Segment>& segments,
                                    const UUID& client_id)
     -> tl::expected<void, ErrorCode> {
     std::shared_lock<std::shared_mutex> shared_lock(snapshot_mutex_);
-    std::unique_lock<std::shared_mutex> lock(client_mutex_);
-    if (ok_client_.contains(client_id)) {
-        LOG(WARNING) << "client_id=" << client_id
-                     << ", warn=client_already_remounted";
-        // Return OK because this is an idempotent operation
-        return {};
+    {
+        std::unique_lock<std::shared_mutex> lock(client_mutex_);
+        if (ok_client_.contains(client_id)) {
+            LOG(WARNING) << "client_id=" << client_id
+                         << ", warn=client_already_remounted";
+            // Return OK because this is an idempotent operation
+            return {};
+        }
+
+        {
+            ScopedSegmentAccess segment_access =
+                segment_manager_.getSegmentAccess();
+
+            // Tell the client monitor thread to start timing for this client.
+            // To avoid the following undesired situations, this message must be
+            // sent after locking the segment mutex or client mutex and before
+            // the remounting operation completes:
+            // 1. Sending the message before the lock: the client expires and
+            // unmouting invokes before this remounting are completed, which
+            // prevents this segment being able to be unmounted forever;
+            // 2. Sending the message after remounting the segments: After
+            // remounting these segments, when trying to push id to the queue,
+            // the queue is already full. However, at this point, the message
+            // must be sent, otherwise this client cannot be monitored and
+            // expired.
+            PodUUID pod_client_id;
+            pod_client_id.first = client_id.first;
+            pod_client_id.second = client_id.second;
+            if (!client_ping_queue_.push(pod_client_id)) {
+                LOG(ERROR) << "client_id=" << client_id
+                           << ", error=client_ping_queue_full";
+                return tl::make_unexpected(ErrorCode::INTERNAL_ERROR);
+            }
+
+            ErrorCode err = segment_access.ReMountSegment(segments, client_id);
+            if (err != ErrorCode::OK) {
+                return tl::make_unexpected(err);
+            }
+        }
+
+        // Change the client status to OK
+        ok_client_.insert(client_id);
+        MasterMetricManager::instance().inc_active_clients();
     }
+    RecomputeTenantEffectiveQuotas();
 
-    ScopedSegmentAccess segment_access = segment_manager_.getSegmentAccess();
+    return {};
+}
 
-    // Tell the client monitor thread to start timing for this client. To
-    // avoid the following undesired situations, this message must be sent
-    // after locking the segment mutex or client mutex and before the remounting
-    // operation completes:
-    // 1. Sending the message before the lock: the client expires and
-    // unmouting invokes before this remounting are completed, which prevents
-    // this segment being able to be unmounted forever;
-    // 2. Sending the message after remounting the segments: After remounting
-    // these segments, when trying to push id to the queue, the queue is
-    // already full. However, at this point, the message must be sent,
-    // otherwise this client cannot be monitored and expired.
-    PodUUID pod_client_id;
-    pod_client_id.first = client_id.first;
-    pod_client_id.second = client_id.second;
-    if (!client_ping_queue_.push(pod_client_id)) {
-        LOG(ERROR) << "client_id=" << client_id
-                   << ", error=client_ping_queue_full";
-        return tl::make_unexpected(ErrorCode::INTERNAL_ERROR);
-    }
-
-    ErrorCode err = segment_access.ReMountSegment(segments, client_id);
+auto MasterService::ReMountNoFSegment(const std::vector<NoFSegment>& segments,
+                                      const UUID& client_id)
+    -> tl::expected<void, ErrorCode> {
+#ifndef USE_NOF
+    LOG(ERROR) << "client_id=" << client_id
+               << ", segments_count=" << segments.size()
+               << ", error=nof_pool_disabled";
+    return tl::make_unexpected(ErrorCode::UNAVAILABLE_IN_CURRENT_MODE);
+#else
+    ScopedNoFSegmentAccess nof_segment_access =
+        nof_segment_manager_.getNoFSegmentAccess();
+    ErrorCode err = nof_segment_access.ReMountSegment(segments, client_id);
     if (err != ErrorCode::OK) {
         return tl::make_unexpected(err);
     }
-
-    // Change the client status to OK
-    ok_client_.insert(client_id);
-    MasterMetricManager::instance().inc_active_clients();
-
     return {};
+#endif
 }
 
 std::unordered_set<UUID, boost::hash<UUID>>
@@ -377,6 +909,922 @@ MasterService::getAliveClientsSnapshot() const {
     return ok_client_;
 }
 
+size_t MasterService::getMetadataShardIndex(const std::string& tenant_id,
+                                            const std::string& key) const {
+    const auto normalized_tenant = NormalizeTenantId(tenant_id);
+    std::shared_lock<std::shared_mutex> lock(group_routing_mutex_);
+    auto it =
+        object_group_ids_.find(MakeTenantScopedKey(normalized_tenant, key));
+    if (it == object_group_ids_.end()) {
+        return getShardIndex(normalized_tenant, key);
+    }
+    return getShardIndex(it->second);
+}
+
+size_t MasterService::getTenantQuotaShardIndex(
+    const std::string& tenant_id) const {
+    return std::hash<std::string>{}(NormalizeTenantId(tenant_id)) %
+           kNumTenantQuotaShards;
+}
+
+std::string MasterService::NormalizeRequestTenantId(
+    const std::string& tenant_id) const {
+    if (!enable_multi_tenants_) {
+        return "default";
+    }
+    return NormalizeTenantId(tenant_id);
+}
+
+MasterService::ObjectIdentity MasterService::MakeObjectIdentityForRequest(
+    const std::string& user_key, const std::string& tenant_id) const {
+    return {NormalizeRequestTenantId(tenant_id), user_key};
+}
+
+bool MasterService::IsTenantRegistered(const std::string& tenant_id) const {
+    if (!enable_multi_tenants_) {
+        return true;
+    }
+    const auto normalized_tenant = NormalizeTenantId(tenant_id);
+    const auto& shard =
+        tenant_quota_shards_[getTenantQuotaShardIndex(normalized_tenant)];
+    std::lock_guard<std::mutex> lock(shard.mutex);
+    auto it = shard.tenants.find(normalized_tenant);
+    return it != shard.tenants.end() && it->second.has_explicit_policy;
+}
+
+tl::expected<std::string, ErrorCode> MasterService::NormalizeTenantIdForWrite(
+    const std::string& tenant_id) const {
+    if (!enable_multi_tenants_) {
+        return std::string("default");
+    }
+    std::lock_guard<std::mutex> policy_lock(tenant_quota_policy_mutex_);
+    return NormalizeTenantIdForWriteLocked(tenant_id);
+}
+
+tl::expected<std::string, ErrorCode>
+MasterService::NormalizeTenantIdForWriteLocked(
+    const std::string& tenant_id) const {
+    if (!enable_multi_tenants_) {
+        return std::string("default");
+    }
+    if (tenant_id.empty()) {
+        return tl::make_unexpected(ErrorCode::TENANT_NOT_REGISTERED);
+    }
+    const auto normalized_tenant = NormalizeTenantId(tenant_id);
+    if (!IsValidTenantId(normalized_tenant)) {
+        return tl::make_unexpected(ErrorCode::TENANT_NOT_REGISTERED);
+    }
+    if (!IsTenantRegistered(normalized_tenant)) {
+        return tl::make_unexpected(ErrorCode::TENANT_NOT_REGISTERED);
+    }
+    return normalized_tenant;
+}
+
+bool MasterService::TenantHasObjects(const std::string& tenant_id) const {
+    const auto normalized_tenant = NormalizeTenantId(tenant_id);
+    for (size_t i = 0; i < kNumShards; ++i) {
+        MetadataShardAccessorRO shard(this, i);
+        auto tenant_it = shard->tenants.find(normalized_tenant);
+        if (tenant_it != shard->tenants.end() &&
+            !tenant_it->second.metadata.empty()) {
+            return true;
+        }
+    }
+    return false;
+}
+
+TenantQuotaPolicySnapshot MasterService::BuildTenantQuotaPolicySnapshot()
+    const {
+    TenantQuotaPolicySnapshot snapshot;
+    for (size_t i = 0; i < kNumTenantQuotaShards; ++i) {
+        const auto& shard = tenant_quota_shards_[i];
+        std::lock_guard<std::mutex> lock(shard.mutex);
+        for (const auto& [tenant_id, state] : shard.tenants) {
+            if (state.has_explicit_policy) {
+                snapshot.tenant_quotas[tenant_id] = state.requested_quota_bytes;
+            }
+        }
+    }
+    return snapshot;
+}
+
+void MasterService::ApplyTenantQuotaPolicies(
+    const TenantQuotaPolicySnapshot& snapshot) {
+    std::array<std::vector<std::pair<std::string, uint64_t>>,
+               kNumTenantQuotaShards>
+        grouped_quotas;
+    for (const auto& [tenant_id, requested_quota_bytes] :
+         snapshot.tenant_quotas) {
+        grouped_quotas[getTenantQuotaShardIndex(tenant_id)].emplace_back(
+            tenant_id, requested_quota_bytes);
+    }
+
+    for (size_t i = 0; i < kNumTenantQuotaShards; ++i) {
+        auto& shard = tenant_quota_shards_[i];
+        std::lock_guard<std::mutex> lock(shard.mutex);
+        for (auto it = shard.tenants.begin(); it != shard.tenants.end();) {
+            const auto policy_it = snapshot.tenant_quotas.find(it->first);
+            auto& state = it->second;
+            if (policy_it != snapshot.tenant_quotas.end()) {
+                state.requested_quota_bytes = policy_it->second;
+                state.has_explicit_policy = true;
+                RefreshTenantQuotaOverQuota(state);
+                ++it;
+            } else {
+                state.has_explicit_policy = false;
+                state.requested_quota_bytes = 0;
+                state.effective_quota_bytes = 0;
+                if (IsLazyEmptyTenantQuotaState(state)) {
+                    it = shard.tenants.erase(it);
+                } else {
+                    RefreshTenantQuotaOverQuota(state);
+                    ++it;
+                }
+            }
+        }
+
+        for (const auto& [tenant_id, requested_quota_bytes] :
+             grouped_quotas[i]) {
+            auto [tenant_it, inserted] = shard.tenants.try_emplace(tenant_id);
+            if (!inserted) {
+                continue;
+            }
+            auto& state = tenant_it->second;
+            state.requested_quota_bytes = requested_quota_bytes;
+            state.has_explicit_policy = true;
+            RefreshTenantQuotaOverQuota(state);
+        }
+    }
+    RecomputeTenantEffectiveQuotas();
+}
+
+void MasterService::LoadTenantQuotaPoliciesFromStoreOrThrow() {
+    if (!enable_multi_tenants_) {
+        return;
+    }
+    if (!tenant_quota_policy_store_) {
+        throw std::runtime_error(
+            "tenant quota policy store is not initialized");
+    }
+    std::lock_guard<std::mutex> policy_lock(tenant_quota_policy_mutex_);
+    auto snapshot = tenant_quota_policy_store_->Load();
+    if (!snapshot) {
+        throw std::runtime_error("failed to load tenant quota policy: " +
+                                 snapshot.error());
+    }
+    ApplyTenantQuotaPolicies(snapshot.value());
+}
+
+uint64_t MasterService::CompletedMemoryQuotaCharge(
+    const ObjectMetadata& metadata) const {
+    return static_cast<uint64_t>(metadata.size) *
+           metadata.CountReplicas([](const Replica& replica) {
+               return replica.is_memory_replica() && replica.is_completed();
+           });
+}
+
+uint64_t MasterService::RequestedMemoryQuotaCharge(
+    uint64_t value_length, const ReplicateConfig& config) const {
+    const unsigned __int128 charge =
+        static_cast<unsigned __int128>(value_length) * config.replica_num;
+    if (charge > std::numeric_limits<uint64_t>::max()) {
+        return std::numeric_limits<uint64_t>::max();
+    }
+    return static_cast<uint64_t>(charge);
+}
+
+bool MasterService::ShouldProtectZeroChargeMetadataCreate(
+    uint64_t requested_quota_charge) const {
+    return enable_multi_tenants_ && requested_quota_charge == 0;
+}
+
+uint64_t MasterService::ComputeTenantQuotaDeficit(
+    const std::string& tenant_id, uint64_t incoming_quota_charge) {
+    uint64_t deficit = incoming_quota_charge;
+    const auto normalized_tenant = NormalizeTenantId(tenant_id);
+    auto& quota_shard =
+        tenant_quota_shards_[getTenantQuotaShardIndex(normalized_tenant)];
+    std::lock_guard<std::mutex> lock(quota_shard.mutex);
+    auto quota_it = quota_shard.tenants.find(normalized_tenant);
+    if (quota_it == quota_shard.tenants.end()) {
+        return deficit;
+    }
+
+    const auto& state = quota_it->second;
+    const unsigned __int128 demand =
+        static_cast<unsigned __int128>(state.used_bytes) +
+        state.reserved_bytes + incoming_quota_charge;
+    if (demand <= state.effective_quota_bytes) {
+        return 0;
+    }
+
+    const unsigned __int128 raw_deficit = demand - state.effective_quota_bytes;
+    return raw_deficit > std::numeric_limits<uint64_t>::max()
+               ? std::numeric_limits<uint64_t>::max()
+               : static_cast<uint64_t>(raw_deficit);
+}
+
+uint64_t MasterService::GetTenantQuotaCapacityBytes() {
+    uint64_t capacity = 0;
+    ScopedSegmentAccess segment_access = segment_manager_.getSegmentAccess();
+    std::vector<std::pair<Segment, UUID>> segments;
+    if (segment_access.GetAllSegments(segments) != ErrorCode::OK) {
+        return 0;
+    }
+    for (const auto& [segment, _] : segments) {
+        if (capacity > std::numeric_limits<uint64_t>::max() - segment.size) {
+            return std::numeric_limits<uint64_t>::max();
+        }
+        capacity += segment.size;
+    }
+    return capacity;
+}
+
+uint64_t MasterService::GetTenantQuotaAllocatableCapacityBytes() {
+    return GetTenantQuotaCapacityBytes();
+}
+
+void MasterService::RecomputeTenantEffectiveQuotas() {
+    if (!enable_multi_tenants_) {
+        return;
+    }
+    std::lock_guard<std::mutex> recompute_lock(tenant_quota_recompute_mutex_);
+    const uint64_t capacity = GetTenantQuotaAllocatableCapacityBytes();
+
+    std::map<std::string, TenantQuotaState> active_tenants;
+    for (size_t i = 0; i < kNumTenantQuotaShards; ++i) {
+        auto& shard = tenant_quota_shards_[i];
+        std::lock_guard<std::mutex> lock(shard.mutex);
+        for (auto it = shard.tenants.begin(); it != shard.tenants.end();) {
+            const auto& tenant_id = it->first;
+            auto& state = it->second;
+            if (IsLazyEmptyTenantQuotaState(state)) {
+                it = shard.tenants.erase(it);
+                continue;
+            }
+            if (!state.has_explicit_policy) {
+                state.requested_quota_bytes = 0;
+                state.effective_quota_bytes = 0;
+                RefreshTenantQuotaOverQuota(state);
+            }
+            active_tenants.emplace(tenant_id, state);
+            ++it;
+        }
+    }
+
+    for (const auto& assignment :
+         BuildEffectiveQuotaAssignments(active_tenants, capacity)) {
+        auto& shard = tenant_quota_shards_[getTenantQuotaShardIndex(
+            assignment.tenant_id)];
+        std::lock_guard<std::mutex> lock(shard.mutex);
+        auto it = shard.tenants.find(assignment.tenant_id);
+        if (it == shard.tenants.end()) {
+            continue;
+        }
+        auto& state = it->second;
+        state.effective_quota_bytes = assignment.effective_quota_bytes;
+        RefreshTenantQuotaOverQuota(state);
+    }
+}
+
+tl::expected<void, ErrorCode> MasterService::ReserveTenantQuota(
+    const std::string& tenant_id, uint64_t bytes) {
+    if (!enable_multi_tenants_) {
+        return {};
+    }
+    const auto normalized_tenant = NormalizeTenantId(tenant_id);
+    auto exceeds_effective_quota = [](const TenantQuotaState& state,
+                                      uint64_t additional_bytes) {
+        return static_cast<unsigned __int128>(state.used_bytes) +
+                   state.reserved_bytes + additional_bytes >
+               state.effective_quota_bytes;
+    };
+    auto& shard =
+        tenant_quota_shards_[getTenantQuotaShardIndex(normalized_tenant)];
+    {
+        std::lock_guard<std::mutex> lock(shard.mutex);
+        auto it = shard.tenants.find(normalized_tenant);
+        if (it == shard.tenants.end() || !it->second.has_explicit_policy) {
+            return tl::make_unexpected(ErrorCode::TENANT_NOT_REGISTERED);
+        }
+
+        auto& state = it->second;
+        if (bytes == 0) {
+            return {};
+        }
+        if (exceeds_effective_quota(state, bytes)) {
+            return tl::make_unexpected(ErrorCode::TENANT_QUOTA_EXCEEDED);
+        }
+
+        state.reserved_bytes += bytes;
+        RefreshTenantQuotaOverQuota(state);
+        return {};
+    }
+}
+
+void MasterService::CommitTenantQuota(const std::string& tenant_id,
+                                      uint64_t bytes) {
+    if (!enable_multi_tenants_ || bytes == 0) {
+        return;
+    }
+    const auto normalized_tenant = NormalizeTenantId(tenant_id);
+    auto& shard =
+        tenant_quota_shards_[getTenantQuotaShardIndex(normalized_tenant)];
+    std::lock_guard<std::mutex> lock(shard.mutex);
+    auto it = shard.tenants.find(normalized_tenant);
+    if (it == shard.tenants.end()) {
+        LOG(ERROR) << "tenant quota commit mismatch tenant="
+                   << normalized_tenant << ", bytes=" << bytes
+                   << ", reserved=0";
+        return;
+    }
+    auto& state = it->second;
+    if (state.reserved_bytes < bytes) {
+        LOG(ERROR) << "tenant quota commit mismatch tenant="
+                   << normalized_tenant << ", bytes=" << bytes
+                   << ", reserved=" << state.reserved_bytes;
+        return;
+    }
+    state.reserved_bytes -= bytes;
+    if (state.used_bytes > std::numeric_limits<uint64_t>::max() - bytes) {
+        state.used_bytes = std::numeric_limits<uint64_t>::max();
+    } else {
+        state.used_bytes += bytes;
+    }
+    ++state.committed_count;
+    RefreshTenantQuotaOverQuota(state);
+}
+
+void MasterService::AbortTenantQuota(const std::string& tenant_id,
+                                     uint64_t bytes) {
+    if (!enable_multi_tenants_ || bytes == 0) {
+        return;
+    }
+    const auto normalized_tenant = NormalizeTenantId(tenant_id);
+    auto& shard =
+        tenant_quota_shards_[getTenantQuotaShardIndex(normalized_tenant)];
+    bool recompute_needed = false;
+    {
+        std::lock_guard<std::mutex> lock(shard.mutex);
+        auto it = shard.tenants.find(normalized_tenant);
+        if (it == shard.tenants.end()) {
+            LOG(ERROR) << "tenant quota abort mismatch tenant="
+                       << normalized_tenant << ", bytes=" << bytes
+                       << ", reserved=0";
+            return;
+        }
+        auto& state = it->second;
+        if (state.reserved_bytes < bytes) {
+            LOG(ERROR) << "tenant quota abort mismatch tenant="
+                       << normalized_tenant << ", bytes=" << bytes
+                       << ", reserved=" << state.reserved_bytes;
+            return;
+        }
+        state.reserved_bytes -= bytes;
+        if (IsLazyEmptyTenantQuotaState(state)) {
+            shard.tenants.erase(it);
+            recompute_needed = true;
+        } else {
+            RefreshTenantQuotaOverQuota(state);
+        }
+    }
+    if (recompute_needed) {
+        RecomputeTenantEffectiveQuotas();
+    }
+}
+
+void MasterService::ReleaseTenantQuota(const std::string& tenant_id,
+                                       uint64_t bytes) {
+    if (!enable_multi_tenants_ || bytes == 0) {
+        return;
+    }
+    const auto normalized_tenant = NormalizeTenantId(tenant_id);
+    auto& shard =
+        tenant_quota_shards_[getTenantQuotaShardIndex(normalized_tenant)];
+    bool recompute_needed = false;
+    {
+        std::lock_guard<std::mutex> lock(shard.mutex);
+        auto it = shard.tenants.find(normalized_tenant);
+        if (it == shard.tenants.end()) {
+            LOG(ERROR) << "tenant quota release mismatch tenant="
+                       << normalized_tenant << ", bytes=" << bytes
+                       << ", used=0";
+            return;
+        }
+        auto& state = it->second;
+        if (state.used_bytes < bytes) {
+            LOG(ERROR) << "tenant quota release mismatch tenant="
+                       << normalized_tenant << ", bytes=" << bytes
+                       << ", used=" << state.used_bytes;
+            return;
+        }
+        state.used_bytes -= bytes;
+        if (state.committed_count > 0) {
+            --state.committed_count;
+        }
+        if (IsLazyEmptyTenantQuotaState(state)) {
+            shard.tenants.erase(it);
+            recompute_needed = true;
+        } else {
+            RefreshTenantQuotaOverQuota(state);
+        }
+    }
+    if (recompute_needed) {
+        RecomputeTenantEffectiveQuotas();
+    }
+}
+
+void MasterService::ReleaseTenantQuotaPartial(const std::string& tenant_id,
+                                              uint64_t bytes) {
+    if (!enable_multi_tenants_ || bytes == 0) {
+        return;
+    }
+    const auto normalized_tenant = NormalizeTenantId(tenant_id);
+    auto& shard =
+        tenant_quota_shards_[getTenantQuotaShardIndex(normalized_tenant)];
+    std::lock_guard<std::mutex> lock(shard.mutex);
+    auto it = shard.tenants.find(normalized_tenant);
+    if (it == shard.tenants.end()) {
+        LOG(ERROR) << "tenant quota partial release mismatch tenant="
+                   << normalized_tenant << ", bytes=" << bytes << ", used=0";
+        return;
+    }
+    auto& state = it->second;
+    if (state.used_bytes < bytes) {
+        LOG(ERROR) << "tenant quota partial release mismatch tenant="
+                   << normalized_tenant << ", bytes=" << bytes
+                   << ", used=" << state.used_bytes;
+        return;
+    }
+    state.used_bytes -= bytes;
+    RefreshTenantQuotaOverQuota(state);
+}
+
+void MasterService::CommitAdditionalTenantQuota(const std::string& tenant_id,
+                                                uint64_t bytes) {
+    if (!enable_multi_tenants_ || bytes == 0) {
+        return;
+    }
+    const auto normalized_tenant = NormalizeTenantId(tenant_id);
+    auto& shard =
+        tenant_quota_shards_[getTenantQuotaShardIndex(normalized_tenant)];
+    std::lock_guard<std::mutex> lock(shard.mutex);
+    auto it = shard.tenants.find(normalized_tenant);
+    if (it == shard.tenants.end()) {
+        LOG(ERROR) << "tenant quota additional commit mismatch tenant="
+                   << normalized_tenant << ", bytes=" << bytes
+                   << ", reserved=0";
+        return;
+    }
+    auto& state = it->second;
+    if (state.reserved_bytes < bytes) {
+        LOG(ERROR) << "tenant quota additional commit mismatch tenant="
+                   << normalized_tenant << ", bytes=" << bytes
+                   << ", reserved=" << state.reserved_bytes;
+        return;
+    }
+    state.reserved_bytes -= bytes;
+    if (state.used_bytes > std::numeric_limits<uint64_t>::max() - bytes) {
+        state.used_bytes = std::numeric_limits<uint64_t>::max();
+    } else {
+        state.used_bytes += bytes;
+    }
+    RefreshTenantQuotaOverQuota(state);
+}
+
+void MasterService::AbortReplicationTaskQuota(const std::string& tenant_id,
+                                              const ReplicationTask& task) {
+    AbortTenantQuota(tenant_id, task.reserved_quota_charge_bytes);
+}
+
+void MasterService::IncrementTenantMetadataObjectCount(
+    const std::string& tenant_id) {
+    if (!enable_multi_tenants_) {
+        return;
+    }
+    const auto normalized_tenant = NormalizeTenantId(tenant_id);
+    auto& shard =
+        tenant_quota_shards_[getTenantQuotaShardIndex(normalized_tenant)];
+    std::lock_guard<std::mutex> lock(shard.mutex);
+    auto& state = shard.tenants[normalized_tenant];
+    if (state.metadata_object_count < std::numeric_limits<uint64_t>::max()) {
+        ++state.metadata_object_count;
+    }
+    RefreshTenantQuotaOverQuota(state);
+}
+
+void MasterService::DecrementTenantMetadataObjectCount(
+    const std::string& tenant_id) {
+    if (!enable_multi_tenants_) {
+        return;
+    }
+    const auto normalized_tenant = NormalizeTenantId(tenant_id);
+    auto& shard =
+        tenant_quota_shards_[getTenantQuotaShardIndex(normalized_tenant)];
+    bool recompute_needed = false;
+    {
+        std::lock_guard<std::mutex> lock(shard.mutex);
+        auto it = shard.tenants.find(normalized_tenant);
+        if (it == shard.tenants.end()) {
+            LOG(WARNING) << "tenant metadata object count decrement mismatch "
+                         << "tenant=" << normalized_tenant;
+            return;
+        }
+        auto& state = it->second;
+        if (state.metadata_object_count > 0) {
+            --state.metadata_object_count;
+        } else {
+            LOG(WARNING) << "tenant metadata object count underflow tenant="
+                         << normalized_tenant;
+        }
+        if (IsLazyEmptyTenantQuotaState(state)) {
+            shard.tenants.erase(it);
+            recompute_needed = true;
+        } else {
+            RefreshTenantQuotaOverQuota(state);
+        }
+    }
+    if (recompute_needed) {
+        RecomputeTenantEffectiveQuotas();
+    }
+}
+
+void MasterService::ReleaseCommittedQuotaCharge(ObjectMetadata& metadata,
+                                                uint64_t bytes) {
+    if (!enable_multi_tenants_ || bytes == 0) {
+        return;
+    }
+    const uint64_t release_bytes =
+        std::min(bytes, metadata.committed_quota_charge_bytes);
+    if (release_bytes == metadata.committed_quota_charge_bytes) {
+        ReleaseTenantQuota(metadata.tenant_id, release_bytes);
+    } else {
+        ReleaseTenantQuotaPartial(metadata.tenant_id, release_bytes);
+    }
+    metadata.committed_quota_charge_bytes -= release_bytes;
+}
+
+void MasterService::RebuildTenantQuotaUsageFromMetadata() {
+    if (!enable_multi_tenants_) {
+        return;
+    }
+
+    std::unordered_set<std::string> metadata_tenants;
+    std::unordered_map<std::string, uint64_t> used_by_tenant;
+    std::unordered_map<std::string, uint64_t> committed_count_by_tenant;
+    std::unordered_map<std::string, uint64_t> metadata_count_by_tenant;
+    for (size_t i = 0; i < kNumShards; ++i) {
+        MetadataShardAccessorRW shard(this, i);
+        for (auto& [tenant_id, tenant_state] : shard->tenants) {
+            metadata_tenants.insert(tenant_id);
+            for (auto& [_, metadata] : tenant_state.metadata) {
+                metadata_count_by_tenant[tenant_id]++;
+                const uint64_t charge = CompletedMemoryQuotaCharge(metadata);
+                metadata.reserved_quota_charge_bytes = 0;
+                metadata.committed_quota_charge_bytes = charge;
+                metadata.pending_replaced_quota_charge_bytes = 0;
+                if (charge == 0) {
+                    continue;
+                }
+                used_by_tenant[tenant_id] += charge;
+                committed_count_by_tenant[tenant_id]++;
+            }
+        }
+    }
+
+    for (size_t i = 0; i < kNumTenantQuotaShards; ++i) {
+        auto& shard = tenant_quota_shards_[i];
+        std::lock_guard<std::mutex> lock(shard.mutex);
+        for (auto& [_, state] : shard.tenants) {
+            state.used_bytes = 0;
+            state.reserved_bytes = 0;
+            state.committed_count = 0;
+            state.metadata_object_count = 0;
+        }
+    }
+    for (const auto& tenant_id : metadata_tenants) {
+        auto& shard = tenant_quota_shards_[getTenantQuotaShardIndex(tenant_id)];
+        std::lock_guard<std::mutex> lock(shard.mutex);
+        auto [it, inserted] = shard.tenants.try_emplace(tenant_id);
+        auto& state = it->second;
+        if (inserted || !state.has_explicit_policy) {
+            state.requested_quota_bytes = 0;
+            state.effective_quota_bytes = 0;
+            state.has_explicit_policy = false;
+            LOG(WARNING)
+                << "tenant " << tenant_id
+                << " exists in metadata but has no connector quota policy; "
+                   "creating orphan quota state";
+        }
+        state.used_bytes = used_by_tenant[tenant_id];
+        state.committed_count = committed_count_by_tenant[tenant_id];
+        state.metadata_object_count = metadata_count_by_tenant[tenant_id];
+        RefreshTenantQuotaOverQuota(state);
+    }
+    RecomputeTenantEffectiveQuotas();
+}
+
+std::optional<std::string> MasterService::GetGroupRoute(
+    const std::string& tenant_id, const std::string& key) const {
+    const auto normalized_tenant = NormalizeTenantId(tenant_id);
+    std::shared_lock<std::shared_mutex> lock(group_routing_mutex_);
+    auto it =
+        object_group_ids_.find(MakeTenantScopedKey(normalized_tenant, key));
+    if (it == object_group_ids_.end()) {
+        return std::nullopt;
+    }
+    return it->second;
+}
+
+MasterService::ObjectOperationLock MasterService::AcquireObjectOperationLock(
+    const std::string& tenant_id, const std::string& key) {
+    const auto scoped_key = MakeTenantScopedKey(tenant_id, key);
+    const auto stripe_idx =
+        std::hash<std::string>{}(scoped_key) % kObjectOperationLockStripes;
+    return {std::unique_lock<std::mutex>(object_operation_locks_[stripe_idx])};
+}
+
+void MasterService::RegisterGroupMember(TenantState& tenant_state,
+                                        const std::string& tenant_id,
+                                        const std::string& key,
+                                        const std::string& group_id) {
+    if (group_id.empty()) {
+        return;
+    }
+    const auto normalized_tenant = NormalizeTenantId(tenant_id);
+    std::unique_lock<std::shared_mutex> lock(group_routing_mutex_);
+    object_group_ids_[MakeTenantScopedKey(normalized_tenant, key)] = group_id;
+    groups_needing_lease_refresh_.insert(
+        MakeTenantScopedKey(normalized_tenant, group_id));
+    tenant_state.group_members[group_id].insert(key);
+}
+
+void MasterService::UnregisterGroupMember(TenantState& tenant_state,
+                                          const std::string& tenant_id,
+                                          const std::string& key,
+                                          const std::string& group_id) {
+    if (group_id.empty()) {
+        return;
+    }
+    const auto normalized_tenant = NormalizeTenantId(tenant_id);
+    bool group_empty = false;
+    auto group_it = tenant_state.group_members.find(group_id);
+    if (group_it != tenant_state.group_members.end()) {
+        group_it->second.erase(key);
+        if (group_it->second.empty()) {
+            tenant_state.group_members.erase(group_it);
+            group_empty = true;
+        }
+    }
+    std::unique_lock<std::shared_mutex> lock(group_routing_mutex_);
+    auto route_it =
+        object_group_ids_.find(MakeTenantScopedKey(normalized_tenant, key));
+    if (route_it != object_group_ids_.end() && route_it->second == group_id) {
+        object_group_ids_.erase(route_it);
+    }
+    if (group_empty) {
+        groups_needing_lease_refresh_.erase(
+            MakeTenantScopedKey(normalized_tenant, group_id));
+    }
+}
+
+bool MasterService::HasCompletedMemoryCacheReplica(
+    const ObjectMetadata& metadata) {
+    return metadata.HasReplica([](const Replica& replica) {
+        return replica.is_memory_replica() && replica.is_completed();
+    });
+}
+
+bool MasterService::HasCompletedDiskCacheReplica(
+    const ObjectMetadata& metadata) {
+    return metadata.HasReplica([](const Replica& replica) {
+        return replica.is_disk_replica() && replica.is_completed();
+    });
+}
+
+void MasterService::SyncCacheTotalAccounting(ObjectMetadata& metadata) {
+    const bool has_memory_cache_replica =
+        HasCompletedMemoryCacheReplica(metadata);
+    const bool has_disk_cache_replica = HasCompletedDiskCacheReplica(metadata);
+
+    if (!metadata.memory_cache_total_accounted && has_memory_cache_replica) {
+        MasterMetricManager::instance().inc_mem_cache_nums();
+        metadata.memory_cache_total_accounted = true;
+    } else if (metadata.memory_cache_total_accounted &&
+               !has_memory_cache_replica) {
+        MasterMetricManager::instance().dec_mem_cache_nums();
+        metadata.memory_cache_total_accounted = false;
+    }
+
+    if (!metadata.disk_cache_total_accounted && has_disk_cache_replica) {
+        MasterMetricManager::instance().inc_file_cache_nums();
+        metadata.disk_cache_total_accounted = true;
+    } else if (metadata.disk_cache_total_accounted && !has_disk_cache_replica) {
+        MasterMetricManager::instance().dec_file_cache_nums();
+        metadata.disk_cache_total_accounted = false;
+    }
+}
+
+void MasterService::AccountCacheTotalRemoval(ObjectMetadata& metadata) {
+    if (metadata.memory_cache_total_accounted) {
+        MasterMetricManager::instance().dec_mem_cache_nums();
+        metadata.memory_cache_total_accounted = false;
+    }
+    if (metadata.disk_cache_total_accounted) {
+        MasterMetricManager::instance().dec_file_cache_nums();
+        metadata.disk_cache_total_accounted = false;
+    }
+}
+
+void MasterService::RebuildCacheTotalAccounting() {
+    MasterMetricManager::instance().reset_cache_total_nums();
+    for (auto& shard : metadata_shards_) {
+        for (auto& tenant_entry : shard.tenants) {
+            for (auto& metadata_entry : tenant_entry.second.metadata) {
+                SyncCacheTotalAccounting(metadata_entry.second);
+            }
+        }
+    }
+}
+
+std::vector<Replica> MasterService::PopReplicasWithCacheTotalAccounting(
+    ObjectMetadata& metadata,
+    const std::function<bool(const Replica&)>& pred_fn) {
+    auto replicas = metadata.PopReplicas(pred_fn);
+    SyncCacheTotalAccounting(metadata);
+    return replicas;
+}
+
+std::vector<Replica> MasterService::PopReplicasWithCacheTotalAccounting(
+    ObjectMetadata& metadata) {
+    auto replicas = metadata.PopReplicas();
+    SyncCacheTotalAccounting(metadata);
+    return replicas;
+}
+
+size_t MasterService::EraseReplicasWithCacheTotalAccounting(
+    ObjectMetadata& metadata,
+    const std::function<bool(const Replica&)>& pred_fn) {
+    auto erased_replicas =
+        PopReplicasWithCacheTotalAccounting(metadata, pred_fn);
+    // Release SSD/local-disk usage for any local-disk replicas being removed.
+    // No-op for memory/noF replicas, so it is safe to call unconditionally.
+    ReleaseLocalDiskUsage(erased_replicas);
+    return erased_replicas.size();
+}
+
+std::unordered_map<std::string, MasterService::ObjectMetadata>::iterator
+MasterService::EraseMetadata(
+    TenantState& tenant_state,
+    std::unordered_map<std::string, ObjectMetadata>::iterator it,
+    const std::string& tenant_id) {
+    return EraseMetadata(tenant_state, it, tenant_id, QuotaEraseMode::kFull);
+}
+
+std::unordered_map<std::string, MasterService::ObjectMetadata>::iterator
+MasterService::EraseMetadata(
+    TenantState& tenant_state,
+    std::unordered_map<std::string, ObjectMetadata>::iterator it,
+    const std::string& tenant_id, QuotaEraseMode quota_mode) {
+    return EraseMetadata(tenant_state, it, tenant_id, quota_mode, nullptr);
+}
+
+std::unordered_map<std::string, MasterService::ObjectMetadata>::iterator
+MasterService::EraseMetadata(
+    TenantState& tenant_state,
+    std::unordered_map<std::string, ObjectMetadata>::iterator it,
+    const std::string& tenant_id, QuotaEraseMode quota_mode,
+    MetadataShardAccessorRW* shard) {
+    bool had_completed_disk = it->second.HasReplica([](const Replica& r) {
+        return r.is_local_disk_replica() && r.is_completed();
+    });
+    const std::string key = it->first;
+    const std::string group_id = it->second.group_id;
+    auto& metadata = it->second;
+    ReleaseLocalDiskUsage(metadata.GetAllReplicas());
+    AccountCacheTotalRemoval(metadata);
+    switch (quota_mode) {
+        case QuotaEraseMode::kFull:
+            AbortTenantQuota(tenant_id, metadata.reserved_quota_charge_bytes);
+            ReleaseTenantQuota(tenant_id,
+                               metadata.committed_quota_charge_bytes);
+            ReleaseTenantQuota(tenant_id,
+                               metadata.pending_replaced_quota_charge_bytes);
+            break;
+        case QuotaEraseMode::kPreserveOld:
+            AbortTenantQuota(tenant_id, metadata.reserved_quota_charge_bytes);
+            break;
+        case QuotaEraseMode::kAbortOnly:
+            AbortTenantQuota(tenant_id, metadata.reserved_quota_charge_bytes);
+            break;
+    }
+    auto next = tenant_state.metadata.erase(it);
+    DecrementTenantMetadataObjectCount(tenant_id);
+    if (had_completed_disk && shard) {
+        shard->OnDiskReplicaRemoved(had_completed_disk);
+    }
+    UnregisterGroupMember(tenant_state, tenant_id, key, group_id);
+    return next;
+}
+
+void MasterService::ReleaseLocalDiskUsage(
+    const std::vector<Replica>& replicas) {
+    std::unordered_map<UUID, int64_t, boost::hash<UUID>> bytes_by_client;
+    for (const auto& replica : replicas) {
+        if (!replica.is_local_disk_replica()) {
+            continue;
+        }
+        const auto descriptor =
+            replica.get_descriptor().get_local_disk_descriptor();
+        if (descriptor.object_size > 0) {
+            bytes_by_client[descriptor.client_id] += descriptor.object_size;
+        }
+    }
+    if (bytes_by_client.empty()) {
+        return;
+    }
+
+    ScopedLocalDiskSegmentAccess ssd_access =
+        segment_manager_.getLocalDiskSegmentAccess();
+    auto& client_segments = ssd_access.getClientLocalDiskSegment();
+    for (const auto& [client_id, bytes] : bytes_by_client) {
+        auto disk_it = client_segments.find(client_id);
+        if (disk_it != client_segments.end()) {
+            disk_it->second->ssd_used_bytes.fetch_sub(
+                bytes, std::memory_order_relaxed);
+        }
+    }
+}
+
+void MasterService::RebuildGroupRoutingIndex() {
+    std::unordered_map<std::string, std::string> rebuilt_group_ids;
+    std::unordered_set<std::string> groups_needing_refresh;
+    for (size_t shard_idx = 0; shard_idx < kNumShards; ++shard_idx) {
+        MetadataShardAccessorRW shard(this, shard_idx);
+        for (auto& [tenant_id, tenant_state] : shard->tenants) {
+            tenant_state.group_members.clear();
+            for (const auto& [key, metadata] : tenant_state.metadata) {
+                if (!metadata.IsGrouped()) {
+                    continue;
+                }
+                tenant_state.group_members[metadata.group_id].insert(key);
+                rebuilt_group_ids[MakeTenantScopedKey(tenant_id, key)] =
+                    metadata.group_id;
+                groups_needing_refresh.insert(
+                    MakeTenantScopedKey(tenant_id, metadata.group_id));
+            }
+        }
+    }
+    {
+        std::unique_lock<std::shared_mutex> lock(group_routing_mutex_);
+        object_group_ids_ = std::move(rebuilt_group_ids);
+        groups_needing_lease_refresh_ = std::move(groups_needing_refresh);
+    }
+}
+
+void MasterService::GrantLeaseForGroup(const TenantState& tenant_state,
+                                       const std::string& key,
+                                       const ObjectMetadata& metadata) const {
+    if (!metadata.IsGrouped()) {
+        metadata.GrantLease(default_kv_lease_ttl_, default_kv_soft_pin_ttl_);
+        return;
+    }
+
+    bool needs_refresh = metadata.NeedsLeaseRefresh(default_kv_lease_ttl_,
+                                                    default_kv_soft_pin_ttl_);
+    if (!needs_refresh) {
+        std::shared_lock<std::shared_mutex> lock(group_routing_mutex_);
+        needs_refresh = groups_needing_lease_refresh_.find(MakeTenantScopedKey(
+                            metadata.tenant_id, metadata.group_id)) !=
+                        groups_needing_lease_refresh_.end();
+    }
+    if (!needs_refresh) {
+        return;
+    }
+
+    auto group_it = tenant_state.group_members.find(metadata.group_id);
+    if (group_it == tenant_state.group_members.end()) {
+        metadata.GrantLease(default_kv_lease_ttl_, default_kv_soft_pin_ttl_);
+        return;
+    }
+
+    for (const auto& member_key : group_it->second) {
+        auto mit = tenant_state.metadata.find(member_key);
+        if (mit != tenant_state.metadata.end()) {
+            mit->second.GrantLease(default_kv_lease_ttl_,
+                                   default_kv_soft_pin_ttl_);
+        }
+    }
+    if (group_it->second.find(key) == group_it->second.end()) {
+        metadata.GrantLease(default_kv_lease_ttl_, default_kv_soft_pin_ttl_);
+    }
+    {
+        std::unique_lock<std::shared_mutex> lock(group_routing_mutex_);
+        groups_needing_lease_refresh_.erase(
+            MakeTenantScopedKey(metadata.tenant_id, metadata.group_id));
+    }
+}
+
 void MasterService::ClearInvalidHandles() {
     ClearInvalidHandles(getAliveClientsSnapshot());
 }
@@ -385,18 +1833,33 @@ void MasterService::ClearInvalidHandles(
     const std::unordered_set<UUID, boost::hash<UUID>>& alive_clients) {
     for (size_t i = 0; i < kNumShards; i++) {
         MetadataShardAccessorRW shard(this, i);
-        auto it = shard->metadata.begin();
-        while (it != shard->metadata.end()) {
-            if (CleanupStaleHandles(it->second, alive_clients)) {
-                // If the object is empty, we need to erase the iterator and
-                // also erase the key from processing_keys,
-                // replication_tasks, and offloading_tasks.
-                shard->processing_keys.erase(it->first);
-                shard->replication_tasks.erase(it->first);
-                shard->offloading_tasks.erase(it->first);
-                it = shard->metadata.erase(it);
+        for (auto tenant_it = shard->tenants.begin();
+             tenant_it != shard->tenants.end();) {
+            auto& tenant_state = tenant_it->second;
+            auto it = tenant_state.metadata.begin();
+            while (it != tenant_state.metadata.end()) {
+                if (CleanupStaleHandles(it->second, alive_clients, &shard)) {
+                    tenant_state.processing_keys.erase(it->first);
+                    auto task_it =
+                        tenant_state.replication_tasks.find(it->first);
+                    if (task_it != tenant_state.replication_tasks.end()) {
+                        AbortReplicationTaskQuota(tenant_it->first,
+                                                  task_it->second);
+                        tenant_state.replication_tasks.erase(task_it);
+                    }
+                    tenant_state.offloading_tasks.erase(it->first);
+                    ErasePromotionTaskIfPresent(tenant_state, it->first,
+                                                tenant_it->first);
+                    it = EraseMetadata(tenant_state, it, tenant_it->first,
+                                       QuotaEraseMode::kFull, &shard);
+                } else {
+                    ++it;
+                }
+            }
+            if (tenant_state.Empty()) {
+                tenant_it = shard->tenants.erase(tenant_it);
             } else {
-                ++it;
+                ++tenant_it;
             }
         }
     }
@@ -451,19 +1914,108 @@ auto MasterService::UnmountSegment(const UUID& segment_id,
     ClearInvalidHandles();
 
     // 3. Commit the unmount operation
+    {
+        ScopedSegmentAccess segment_access =
+            segment_manager_.getSegmentAccess();
+        auto err = segment_access.CommitUnmountSegment(segment_id, client_id,
+                                                       metrics_dec_capacity);
+        if (err != ErrorCode::OK) {
+            return tl::make_unexpected(err);
+        }
+    }
+    RecomputeTenantEffectiveQuotas();
+    return {};
+}
+
+auto MasterService::GracefulUnmountSegment(const UUID& segment_id,
+                                           const UUID& client_id,
+                                           uint64_t grace_period_ms)
+    -> tl::expected<void, ErrorCode> {
+    std::unique_lock<std::shared_mutex> lock(snapshot_mutex_);
     ScopedSegmentAccess segment_access = segment_manager_.getSegmentAccess();
+
+    // Verify ownership: the segment must belong to the calling client
+    std::vector<Segment> client_segments;
+    auto err = segment_access.GetClientSegments(client_id, client_segments);
+    if (err != ErrorCode::OK) {
+        return tl::make_unexpected(err);
+    }
+    bool owned = false;
+    for (auto& seg : client_segments) {
+        if (seg.id == segment_id) {
+            owned = true;
+            break;
+        }
+    }
+    if (!owned) {
+        return tl::make_unexpected(ErrorCode::INVALID_PARAMS);
+    }
+
+    err = segment_access.PrepareGracefulUnmountSegment(segment_id);
+    if (err == ErrorCode::SEGMENT_NOT_FOUND) {
+        return {};
+    }
+    if (err != ErrorCode::OK) {
+        return tl::make_unexpected(err);
+    }
+
+    auto expire_time = std::chrono::steady_clock::now() +
+                       std::chrono::milliseconds(grace_period_ms);
+    graceful_unmount_scheduler_.Schedule({segment_id, client_id}, expire_time);
+    return {};
+}
+
+auto MasterService::UnmountNoFSegment(const UUID& segment_id,
+                                      const UUID& client_id)
+    -> tl::expected<void, ErrorCode> {
+#ifndef USE_NOF
+    LOG(ERROR) << "client_id=" << client_id << ", segment_id=" << segment_id
+               << ", error=nof_pool_disabled";
+    return tl::make_unexpected(ErrorCode::UNAVAILABLE_IN_CURRENT_MODE);
+#else
+    size_t metrics_dec_capacity = 0;  // to update the metrics
+
+    // 1. Prepare to unmount the segment by deleting its allocator
+    {
+        ScopedNoFSegmentAccess segment_access =
+            nof_segment_manager_.getNoFSegmentAccess();
+        ErrorCode err = segment_access.PrepareUnmountSegment(
+            segment_id, metrics_dec_capacity);
+        if (err == ErrorCode::SEGMENT_NOT_FOUND) {
+            // Return OK because this is an idempotent operation
+            return {};
+        }
+        if (err != ErrorCode::OK) {
+            return tl::make_unexpected(err);
+        }
+    }  // Release the segment mutex before long-running step 2 and avoid
+       // deadlocks
+
+    // 2. Remove the metadata of the related objects
+    ClearInvalidHandles();
+
+    // 3. Commit the unmount operation
+    ScopedNoFSegmentAccess segment_access =
+        nof_segment_manager_.getNoFSegmentAccess();
     auto err = segment_access.CommitUnmountSegment(segment_id, client_id,
                                                    metrics_dec_capacity);
     if (err != ErrorCode::OK) {
         return tl::make_unexpected(err);
     }
+    {
+        std::lock_guard<std::mutex> lock(nof_heartbeat_mutex_);
+        nof_heartbeat_states_.erase(segment_id);
+    }
     return {};
+#endif
 }
 
-auto MasterService::ExistKey(const std::string& key)
+auto MasterService::ExistKey(const std::string& key,
+                             const std::string& tenant_id)
     -> tl::expected<bool, ErrorCode> {
     std::shared_lock<std::shared_mutex> shared_lock(snapshot_mutex_);
-    MetadataAccessorRO accessor(this, key);
+    MetadataAccessorRO accessor(this,
+                                MakeObjectIdentityForRequest(key, tenant_id));
     if (!accessor.Exists()) {
         VLOG(1) << "key=" << key << ", info=object_not_found";
         return false;
@@ -473,7 +2025,13 @@ auto MasterService::ExistKey(const std::string& key)
     if (metadata.HasReplica(&Replica::fn_is_completed)) {
         // Grant a lease to the object as it may be further used by the
         // client.
-        metadata.GrantLease(default_kv_lease_ttl_, default_kv_soft_pin_ttl_);
+        auto* ts = accessor.GetTenantState();
+        if (ts) {
+            GrantLeaseForGroup(*ts, key, metadata);
+        } else {
+            metadata.GrantLease(default_kv_lease_ttl_,
+                                default_kv_soft_pin_ttl_);
+        }
         return true;
     }
 
@@ -481,22 +2039,87 @@ auto MasterService::ExistKey(const std::string& key)
 }
 
 std::vector<tl::expected<bool, ErrorCode>> MasterService::BatchExistKey(
-    const std::vector<std::string>& keys) {
-    std::vector<tl::expected<bool, ErrorCode>> results;
-    results.reserve(keys.size());
-    for (const auto& key : keys) {
-        results.emplace_back(ExistKey(key));
+    const std::vector<std::string>& keys, const std::string& tenant_id) {
+    const std::string normalized_tenant = NormalizeRequestTenantId(tenant_id);
+    std::vector<tl::expected<bool, ErrorCode>> results(keys.size());
+    if (keys.empty()) {
+        return results;
+    }
+
+    std::vector<std::vector<size_t>> indices_by_shard(kNumShards);
+    {
+        std::shared_lock<std::shared_mutex> group_routing_lock(
+            group_routing_mutex_);
+        for (size_t i = 0; i < keys.size(); ++i) {
+            auto route_it = object_group_ids_.find(
+                MakeTenantScopedKey(normalized_tenant, keys[i]));
+            const size_t shard_idx =
+                route_it == object_group_ids_.end()
+                    ? getShardIndex(normalized_tenant, keys[i])
+                    : getShardIndex(route_it->second);
+            indices_by_shard[shard_idx].push_back(i);
+        }
+    }
+
+    const size_t start_shard = RandomIndex(kNumShards);
+    for (size_t scanned = 0; scanned < kNumShards; ++scanned) {
+        const size_t shard_idx =
+            (start_shard + kNumShards - scanned) % kNumShards;
+        const auto& key_indices = indices_by_shard[shard_idx];
+        if (key_indices.empty()) {
+            continue;
+        }
+
+        std::shared_lock<std::shared_mutex> shared_lock(snapshot_mutex_);
+        MetadataShardAccessorRO shard(this, shard_idx);
+        auto tenant_it = shard->tenants.find(normalized_tenant);
+        if (tenant_it == shard->tenants.end()) {
+            for (const size_t i : key_indices) {
+                VLOG(1) << "key=" << keys[i]
+                        << ", tenant_id=" << normalized_tenant
+                        << ", info=object_not_found";
+                results[i] = false;
+            }
+            continue;
+        }
+
+        const auto& tenant_state = tenant_it->second;
+        for (const size_t i : key_indices) {
+            const auto& key = keys[i];
+            auto it = tenant_state.metadata.find(key);
+            if (it == tenant_state.metadata.end() || !it->second.IsValid()) {
+                VLOG(1) << "key=" << key << ", tenant_id=" << normalized_tenant
+                        << ", info=object_not_found";
+                results[i] = false;
+                continue;
+            }
+
+            const auto& metadata = it->second;
+            if (metadata.HasReplica(&Replica::fn_is_completed)) {
+                GrantLeaseForGroup(tenant_state, key, metadata);
+                results[i] = true;
+            } else {
+                results[i] = false;
+            }
+        }
     }
     return results;
 }
 
-auto MasterService::GetAllKeys()
+auto MasterService::GetAllKeys(const std::string& tenant_id)
     -> tl::expected<std::vector<std::string>, ErrorCode> {
     std::vector<std::string> all_keys;
+    const auto normalized_tenant = NormalizeRequestTenantId(tenant_id);
     for (size_t i = 0; i < kNumShards; i++) {
         MetadataShardAccessorRO shard(this, i);
-        for (const auto& item : shard->metadata) {
-            all_keys.push_back(item.first);
+        auto tenant_it = shard->tenants.find(normalized_tenant);
+        if (tenant_it == shard->tenants.end()) {
+            continue;
+        }
+        for (const auto& item : tenant_it->second.metadata) {
+            all_keys.push_back(item.second.user_key.empty()
+                                   ? item.first
+                                   : item.second.user_key);
         }
     }
     return all_keys;
@@ -513,6 +2136,63 @@ auto MasterService::GetAllSegments()
     return all_segments;
 }
 
+auto MasterService::GetAllNoFSegments()
+    -> tl::expected<std::vector<NoFSegment>, ErrorCode> {
+    std::vector<MountedNoFSegmentSnapshot> mounted_segments;
+    nof_segment_manager_.GetMountedSegmentsSnapshot(mounted_segments);
+
+    std::vector<NoFSegment> result;
+    for (const auto& segment : mounted_segments) {
+        result.push_back(segment.segment);
+    }
+
+    return result;
+}
+
+auto MasterService::GetNoFSegmentsByName(const std::string& segment_name)
+    -> tl::expected<std::vector<NoFSegmentOwnerInfo>, ErrorCode> {
+    return nof_segment_manager_.GetSegmentsByName(segment_name);
+}
+
+auto MasterService::GetSegmentsDetail()
+    -> tl::expected<std::vector<SegmentDetailInfo>, ErrorCode> {
+    ScopedSegmentAccess segment_access = segment_manager_.getSegmentAccess();
+
+    // Get full info of all segments (including Segment and client_id)
+    std::vector<std::pair<Segment, UUID>> all_segments;
+    auto err = segment_access.GetAllSegments(all_segments);
+    if (err != ErrorCode::OK) {
+        return tl::make_unexpected(err);
+    }
+
+    std::vector<SegmentDetailInfo> result;
+    result.reserve(all_segments.size());
+
+    for (const auto& [segment, client_id] : all_segments) {
+        SegmentDetailInfo info;
+        info.segment_name = segment.name;
+        info.segment_id = segment.id;
+        info.client_id = client_id;
+        info.base_address = segment.base;
+        info.size_bytes = segment.size;
+        info.te_endpoint = segment.te_endpoint;
+        info.protocol = segment.protocol;
+
+        // Query segment status
+        segment_access.GetSegmentStatusByName(segment.name, info.status);
+
+        // Query allocator used/capacity
+        size_t used = 0, capacity = 0;
+        segment_access.QuerySegments(segment.name, used, capacity);
+        info.allocator_used_bytes = used;
+        info.allocator_capacity_bytes = capacity;
+
+        result.push_back(std::move(info));
+    }
+
+    return result;
+}
+
 auto MasterService::QuerySegments(const std::string& segment)
     -> tl::expected<std::pair<size_t, size_t>, ErrorCode> {
     ScopedSegmentAccess segment_access = segment_manager_.getSegmentAccess();
@@ -535,6 +2215,17 @@ auto MasterService::QuerySegmentStatus(const std::string& segment_name)
     return status;
 }
 
+auto MasterService::QuerySegmentStatusById(const UUID& segment_id)
+    -> tl::expected<SegmentStatus, ErrorCode> {
+    ScopedSegmentAccess segment_access = segment_manager_.getSegmentAccess();
+    SegmentStatus status = SegmentStatus::UNDEFINED;
+    auto err = segment_access.GetSegmentStatusById(segment_id, status);
+    if (err != ErrorCode::OK) {
+        return tl::make_unexpected(err);
+    }
+    return status;
+}
+
 auto MasterService::QueryIp(const UUID& client_id)
     -> tl::expected<std::vector<std::string>, ErrorCode> {
     ScopedSegmentAccess segment_access = segment_manager_.getSegmentAccess();
@@ -557,13 +2248,7 @@ auto MasterService::QueryIp(const UUID& client_id)
     unique_ips.reserve(segments.size());
     for (const auto& segment : segments) {
         if (!segment.te_endpoint.empty()) {
-            size_t colon_pos = segment.te_endpoint.find(':');
-            if (colon_pos != std::string::npos) {
-                std::string ip = segment.te_endpoint.substr(0, colon_pos);
-                unique_ips.emplace(ip);
-            } else {
-                unique_ips.emplace(segment.te_endpoint);
-            }
+            unique_ips.emplace(getHostNameWithoutPort(segment.te_endpoint));
         }
     }
 
@@ -606,7 +2291,8 @@ auto MasterService::BatchReplicaClear(
             LOG(WARNING) << "BatchReplicaClear: empty key, skipping";
             continue;
         }
-        MetadataAccessorRW accessor(this, key);
+        // BatchReplicaClear is a default-tenant compatibility/admin helper.
+        MetadataAccessorRW accessor(this, MakeObjectIdentity(key, "default"));
         if (!accessor.Exists()) {
             LOG(WARNING) << "BatchReplicaClear: key=" << key
                          << " not found, skipping";
@@ -641,16 +2327,9 @@ auto MasterService::BatchReplicaClear(
                 continue;
             }
 
-            metadata.VisitReplicas(
-                &Replica::fn_is_completed, [](Replica& replica) {
-                    if (replica.is_memory_replica()) {
-                        MasterMetricManager::instance().dec_mem_cache_nums();
-                    } else if (replica.is_disk_replica()) {
-                        MasterMetricManager::instance().dec_file_cache_nums();
-                    }
-                });
-
             // Erase the entire metadata (all replicas will be deallocated)
+            // accessor.Erase() internally calls EraseMetadata which already
+            // decrements disk_object_count via OnDiskReplicaRemoved.
             accessor.Erase();
             cleared_keys.emplace_back(key);
             VLOG(1) << "BatchReplicaClear: successfully cleared all replicas "
@@ -674,15 +2353,8 @@ auto MasterService::BatchReplicaClear(
                 return false;
             };
 
-            metadata.VisitReplicas(
-                match_replica_on_segment, [&](Replica& replica) {
-                    has_replica_on_segment = true;
-                    if (replica.is_memory_replica()) {
-                        MasterMetricManager::instance().dec_mem_cache_nums();
-                    } else if (replica.is_disk_replica()) {
-                        MasterMetricManager::instance().dec_file_cache_nums();
-                    }
-                });
+            has_replica_on_segment =
+                metadata.HasReplica(match_replica_on_segment);
 
             if (!has_replica_on_segment) {
                 LOG(WARNING)
@@ -692,9 +2364,32 @@ auto MasterService::BatchReplicaClear(
                 continue;
             }
 
-            metadata.EraseReplicas(match_replica_on_segment);
+            bool had_completed_disk_on_segment =
+                metadata.HasReplica([&segment_name](const Replica& r) {
+                    if (!r.is_local_disk_replica() || !r.is_completed())
+                        return false;
+                    for (const auto& name : r.get_segment_names()) {
+                        if (name.has_value() && name.value() == segment_name)
+                            return true;
+                    }
+                    return false;
+                });
+
+            EraseReplicasWithCacheTotalAccounting(metadata,
+                                                  match_replica_on_segment);
+
+            if (had_completed_disk_on_segment &&
+                !metadata.HasReplica([](const Replica& r) {
+                    return r.is_local_disk_replica() && r.is_completed();
+                })) {
+                auto& shard = accessor.GetShard();
+                shard.OnDiskReplicaRemoved(had_completed_disk_on_segment,
+                                           metadata);
+            }
 
             // If no valid replicas remain, erase the entire metadata
+            // accessor.Erase() internally calls EraseMetadata which already
+            // decrements disk_object_count via OnDiskReplicaRemoved.
             if (!metadata.IsValid()) {
                 accessor.Erase();
             }
@@ -710,7 +2405,8 @@ auto MasterService::BatchReplicaClear(
     return cleared_keys;
 }
 
-auto MasterService::GetReplicaListByRegex(const std::string& regex_pattern)
+auto MasterService::GetReplicaListByRegex(const std::string& regex_pattern,
+                                          const std::string& tenant_id)
     -> tl::expected<
         std::unordered_map<std::string, std::vector<Replica::Descriptor>>,
         ErrorCode> {
@@ -726,10 +2422,14 @@ auto MasterService::GetReplicaListByRegex(const std::string& regex_pattern)
     }
 
     std::shared_lock<std::shared_mutex> shared_lock(snapshot_mutex_);
+    const auto normalized_tenant = NormalizeRequestTenantId(tenant_id);
     for (size_t i = 0; i < kNumShards; ++i) {
         MetadataShardAccessorRO shard(this, i);
-
-        for (const auto& [key, metadata] : shard->metadata) {
+        auto tenant_it = shard->tenants.find(normalized_tenant);
+        if (tenant_it == shard->tenants.end()) {
+            continue;
+        }
+        for (const auto& [key, metadata] : tenant_it->second.metadata) {
             if (std::regex_search(key, pattern)) {
                 std::vector<Replica::Descriptor> replica_list;
                 metadata.VisitReplicas(
@@ -746,8 +2446,7 @@ auto MasterService::GetReplicaListByRegex(const std::string& regex_pattern)
                 }
 
                 results.emplace(key, std::move(replica_list));
-                metadata.GrantLease(default_kv_lease_ttl_,
-                                    default_kv_soft_pin_ttl_);
+                GrantLeaseForGroup(tenant_it->second, key, metadata);
             }
         }
     }
@@ -755,52 +2454,223 @@ auto MasterService::GetReplicaListByRegex(const std::string& regex_pattern)
     return results;
 }
 
-auto MasterService::GetReplicaList(const std::string& key)
+auto MasterService::GetReplicaList(const std::string& key,
+                                   const std::string& tenant_id)
     -> tl::expected<GetReplicaListResponse, ErrorCode> {
     std::shared_lock<std::shared_mutex> shared_lock(snapshot_mutex_);
-    MetadataAccessorRO accessor(this, key);
+    const auto object_id = MakeObjectIdentityForRequest(key, tenant_id);
 
-    MasterMetricManager::instance().inc_total_get_nums();
+    GetReplicaListResponse resp({}, default_kv_lease_ttl_);
+    bool promotion_eligible = false;
+    {
+        MetadataAccessorRO accessor(this, object_id);
 
-    if (!accessor.Exists()) {
-        VLOG(1) << "key=" << key << ", info=object_not_found";
-        return tl::make_unexpected(ErrorCode::OBJECT_NOT_FOUND);
+        MasterMetricManager::instance().inc_total_get_nums();
+
+        if (!accessor.Exists()) {
+            VLOG(1) << "key=" << key << ", info=object_not_found";
+            return tl::make_unexpected(ErrorCode::OBJECT_NOT_FOUND);
+        }
+        const auto& metadata = accessor.Get();
+
+        std::vector<Replica::Descriptor> replica_list;
+        metadata.VisitReplicas(
+            &Replica::fn_is_completed, [&replica_list](const Replica& replica) {
+                replica_list.emplace_back(replica.get_descriptor());
+            });
+
+        if (replica_list.empty()) {
+            LOG(WARNING) << "key=" << key << ", error=replica_not_ready";
+            return tl::make_unexpected(ErrorCode::REPLICA_IS_NOT_READY);
+        }
+
+        // TODO: NoF SSD support (ranhaojia)
+        if (replica_list[0].is_memory_replica()) {
+            MasterMetricManager::instance().inc_mem_cache_hit_nums();
+        } else if (replica_list[0].is_disk_replica()) {
+            MasterMetricManager::instance().inc_file_cache_hit_nums();
+        }
+        MasterMetricManager::instance().inc_valid_get_nums();
+        // Grant a lease to the object so it will not be removed
+        // when the client is reading it.
+        auto* ts = accessor.GetTenantState();
+        if (ts) {
+            GrantLeaseForGroup(*ts, key, metadata);
+        } else {
+            metadata.GrantLease(default_kv_lease_ttl_,
+                                default_kv_soft_pin_ttl_);
+        }
+
+        // Promotion-on-hit eligibility: only when no MEMORY replica is
+        // present but at least one LOCAL_DISK replica is. Decided here while
+        // we hold the RO accessor; the actual enqueue happens after we
+        // release the accessor below to avoid lock-upgrade complexity.
+        if (promotion_on_hit_) {
+            const bool any_memory =
+                metadata.HasReplica(&Replica::fn_is_memory_replica);
+            const bool any_local_disk =
+                metadata.HasReplica(&Replica::fn_is_local_disk_replica);
+            promotion_eligible = !any_memory && any_local_disk;
+        }
+
+        resp = GetReplicaListResponse(std::move(replica_list),
+                                      default_kv_lease_ttl_);
     }
-    const auto& metadata = accessor.Get();
+    // RO accessor released. Safe to take a fresh RW accessor now.
+    if (promotion_eligible) {
+        TryPushPromotionQueue(object_id);
+    }
+    return resp;
+}
 
-    std::vector<Replica::Descriptor> replica_list;
-    metadata.VisitReplicas(
-        &Replica::fn_is_completed, [&replica_list](const Replica& replica) {
-            replica_list.emplace_back(replica.get_descriptor());
-        });
+std::vector<tl::expected<GetReplicaListResponse, ErrorCode>>
+MasterService::BatchGetReplicaList(const std::vector<std::string>& keys,
+                                   const std::string& tenant_id) {
+    using GetResult = tl::expected<GetReplicaListResponse, ErrorCode>;
 
-    if (replica_list.empty()) {
-        LOG(WARNING) << "key=" << key << ", error=replica_not_ready";
-        return tl::make_unexpected(ErrorCode::REPLICA_IS_NOT_READY);
+    std::vector<GetResult> results(
+        keys.size(), tl::make_unexpected(ErrorCode::OBJECT_NOT_FOUND));
+    if (keys.empty()) {
+        return results;
     }
 
-    if (replica_list[0].is_memory_replica()) {
-        MasterMetricManager::instance().inc_mem_cache_hit_nums();
-    } else if (replica_list[0].is_disk_replica()) {
-        MasterMetricManager::instance().inc_file_cache_hit_nums();
+    const auto normalized_tenant = NormalizeRequestTenantId(tenant_id);
+    constexpr size_t kInvalidKeyIndex = std::numeric_limits<size_t>::max();
+    std::array<size_t, kNumShards> key_list_heads;
+    key_list_heads.fill(kInvalidKeyIndex);
+    std::vector<size_t> next_key_indexes(keys.size(), kInvalidKeyIndex);
+    {
+        std::shared_lock<std::shared_mutex> lock(group_routing_mutex_);
+        for (size_t i = keys.size(); i > 0; --i) {
+            const size_t original_idx = i - 1;
+            const auto scoped_key =
+                MakeTenantScopedKey(normalized_tenant, keys[original_idx]);
+            const auto route_it = object_group_ids_.find(scoped_key);
+            const size_t shard_idx =
+                route_it == object_group_ids_.end()
+                    ? getShardIndex(normalized_tenant, keys[original_idx])
+                    : getShardIndex(route_it->second);
+            next_key_indexes[original_idx] = key_list_heads[shard_idx];
+            key_list_heads[shard_idx] = original_idx;
+        }
     }
-    MasterMetricManager::instance().inc_valid_get_nums();
-    // Grant a lease to the object so it will not be removed
-    // when the client is reading it.
-    metadata.GrantLease(default_kv_lease_ttl_, default_kv_soft_pin_ttl_);
 
-    return GetReplicaListResponse(std::move(replica_list),
-                                  default_kv_lease_ttl_);
+    const size_t start_shard = RandomIndex(kNumShards);
+    for (size_t scanned = 0; scanned < kNumShards; ++scanned) {
+        const size_t shard_idx =
+            (start_shard + kNumShards - scanned) % kNumShards;
+        if (key_list_heads[shard_idx] == kInvalidKeyIndex) {
+            continue;
+        }
+
+        std::vector<ObjectIdentity> promotion_candidates;
+        std::shared_lock<std::shared_mutex> shared_lock(snapshot_mutex_);
+        {
+            MetadataShardAccessorRO shard(this, shard_idx);
+            const auto tenant_it = shard->tenants.find(normalized_tenant);
+            for (size_t original_idx = key_list_heads[shard_idx];
+                 original_idx != kInvalidKeyIndex;
+                 original_idx = next_key_indexes[original_idx]) {
+                const std::string& key = keys[original_idx];
+                MasterMetricManager::instance().inc_total_get_nums();
+
+                if (tenant_it == shard->tenants.end()) {
+                    VLOG(1) << "key=" << key << ", info=object_not_found";
+                    results[original_idx] =
+                        tl::make_unexpected(ErrorCode::OBJECT_NOT_FOUND);
+                    continue;
+                }
+
+                const auto& tenant_state = tenant_it->second;
+                const auto metadata_it = tenant_state.metadata.find(key);
+                if (metadata_it == tenant_state.metadata.end() ||
+                    !metadata_it->second.IsValid()) {
+                    VLOG(1) << "key=" << key << ", info=object_not_found";
+                    results[original_idx] =
+                        tl::make_unexpected(ErrorCode::OBJECT_NOT_FOUND);
+                    continue;
+                }
+
+                const auto& metadata = metadata_it->second;
+                std::vector<Replica::Descriptor> replica_list;
+                metadata.VisitReplicas(
+                    &Replica::fn_is_completed,
+                    [&replica_list](const Replica& replica) {
+                        replica_list.emplace_back(replica.get_descriptor());
+                    });
+
+                if (replica_list.empty()) {
+                    LOG(WARNING)
+                        << "key=" << key << ", error=replica_not_ready";
+                    results[original_idx] =
+                        tl::make_unexpected(ErrorCode::REPLICA_IS_NOT_READY);
+                    continue;
+                }
+
+                if (replica_list[0].is_memory_replica()) {
+                    MasterMetricManager::instance().inc_mem_cache_hit_nums();
+                } else if (replica_list[0].is_disk_replica()) {
+                    MasterMetricManager::instance().inc_file_cache_hit_nums();
+                }
+                MasterMetricManager::instance().inc_valid_get_nums();
+                GrantLeaseForGroup(tenant_state, key, metadata);
+
+                if (promotion_on_hit_) {
+                    const bool any_memory =
+                        metadata.HasReplica(&Replica::fn_is_memory_replica);
+                    const bool any_local_disk =
+                        metadata.HasReplica(&Replica::fn_is_local_disk_replica);
+                    if (!any_memory && any_local_disk) {
+                        promotion_candidates.push_back(
+                            MakeObjectIdentity(key, normalized_tenant));
+                    }
+                }
+
+                results[original_idx] = GetReplicaListResponse(
+                    std::move(replica_list), default_kv_lease_ttl_);
+            }
+        }
+
+        for (const auto& object_id : promotion_candidates) {
+            TryPushPromotionQueue(object_id);
+        }
+    }
+
+    return results;
 }
 
 auto MasterService::AllocateAndInsertMetadata(
     MetadataShardAccessorRW& shard, const UUID& client_id,
     const std::string& key, uint64_t value_length,
-    const ReplicateConfig& config,
+    const ReplicateConfig& config, const std::string& group_id,
+    const std::string& tenant_id,
     const std::chrono::system_clock::time_point& now)
     -> tl::expected<std::vector<Replica::Descriptor>, ErrorCode> {
+    auto& tenant_state = shard->tenants[tenant_id];
+    if (tenant_state.metadata.contains(key)) {
+        LOG(INFO) << "key=" << key << ", info=object_already_exists";
+        return tl::make_unexpected(ErrorCode::OBJECT_ALREADY_EXISTS);
+    }
+    if (GetGroupRoute(tenant_id, key).has_value()) {
+        LOG(INFO) << "key=" << key << ", info=object_already_exists";
+        return tl::make_unexpected(ErrorCode::OBJECT_ALREADY_EXISTS);
+    }
+
+    const uint64_t reserved_quota_charge =
+        RequestedMemoryQuotaCharge(value_length, config);
+    auto quota_result = ReserveTenantQuota(tenant_id, reserved_quota_charge);
+    if (!quota_result) {
+        return tl::make_unexpected(quota_result.error());
+    }
+    auto abort_reserved_quota = [&] {
+        AbortTenantQuota(tenant_id, reserved_quota_charge);
+    };
+
     std::vector<Replica> replicas;
-    {
+    const auto write_mode = DetermineReplicaWriteMode(config);
+    size_t allocated_memory_replicas = 0;
+    size_t allocated_nof_replicas = 0;
+    if (config.replica_num > 0) {
         ScopedAllocatorAccess allocator_access =
             segment_manager_.getAllocatorAccess();
         const auto& allocator_manager = allocator_access.getAllocatorManager();
@@ -812,21 +2682,97 @@ auto MasterService::AllocateAndInsertMetadata(
             preferred_segments = config.preferred_segments;
         }
 
+        const SsdMetricsProvider* ssd_provider = nullptr;
+        std::optional<ScopedLocalDiskSegmentAccess> ssd_access;
+        if (allocation_strategy_type_ ==
+            AllocationStrategyType::SSD_FREE_RATIO_FIRST) {
+            ssd_access.emplace(segment_manager_.getLocalDiskSegmentAccess());
+            ssd_provider = &*ssd_access;
+        }
+
         auto allocation_result = allocation_strategy_->Allocate(
             allocator_manager, value_length, config.replica_num,
-            preferred_segments);
+            preferred_segments, std::set<std::string>(), ReplicaType::MEMORY,
+            ssd_provider);
 
         if (!allocation_result.has_value()) {
             VLOG(1) << "Failed to allocate replicas for key=" << key
                     << ", error: " << allocation_result.error();
             if (allocation_result.error() == ErrorCode::INVALID_PARAMS) {
+                abort_reserved_quota();
                 return tl::make_unexpected(ErrorCode::INVALID_PARAMS);
             }
-            need_eviction_ = true;
-            return tl::make_unexpected(ErrorCode::NO_AVAILABLE_HANDLE);
+            if (write_mode != ReplicaWriteMode::FLEXIBLE_DUAL_REPLICA) {
+                MasterMetricManager::instance().inc_put_start_alloc_failures();
+                need_mem_eviction_ = true;
+                abort_reserved_quota();
+                return tl::make_unexpected(ErrorCode::NO_AVAILABLE_HANDLE);
+            }
+        } else {
+            allocated_memory_replicas = allocation_result->size();
+            replicas = std::move(allocation_result.value());
         }
+    }
 
-        replicas = std::move(allocation_result.value());
+#ifdef USE_NOF
+    if (config.nof_replica_num > 0 &&
+        nof_segment_manager_.getMountedSegmentCount() > 0) {
+        ScopedAllocatorAccess allocator_access =
+            nof_segment_manager_.getAllocatorAccess();
+        const auto& allocator_manager = allocator_access.getAllocatorManager();
+
+        std::vector<std::string> preferred_segments =
+            config.preferred_nof_segments;
+
+        auto allocation_result = allocation_strategy_->Allocate(
+            allocator_manager, value_length, config.nof_replica_num,
+            preferred_segments, std::set<std::string>(), ReplicaType::NOF_SSD);
+
+        if (!allocation_result.has_value()) {
+            VLOG(1) << "Failed to allocate nof replicas for key=" << key
+                    << ", error: " << allocation_result.error();
+            if (allocation_result.error() == ErrorCode::INVALID_PARAMS) {
+                abort_reserved_quota();
+                return tl::make_unexpected(ErrorCode::INVALID_PARAMS);
+            }
+            if (write_mode != ReplicaWriteMode::FLEXIBLE_DUAL_REPLICA) {
+                MasterMetricManager::instance().inc_put_start_alloc_failures();
+                need_nof_eviction_ = true;
+                abort_reserved_quota();
+                return tl::make_unexpected(ErrorCode::NO_AVAILABLE_HANDLE);
+            }
+        } else {
+            allocated_nof_replicas = allocation_result->size();
+            for (auto& replica : allocation_result.value()) {
+                replicas.push_back(std::move(replica));
+            }
+        }
+    }
+#endif
+
+    if (!HasExpectedReplicaAllocation(config, allocated_memory_replicas,
+                                      allocated_nof_replicas)) {
+        if ((config.replica_num > 0 &&
+             allocated_memory_replicas != config.replica_num) ||
+            (config.nof_replica_num > 0 &&
+             allocated_nof_replicas != config.nof_replica_num)) {
+            MasterMetricManager::instance().inc_put_start_alloc_failures();
+            if (config.replica_num > 0 &&
+                allocated_memory_replicas != config.replica_num) {
+                need_mem_eviction_ = true;
+            }
+            if (config.nof_replica_num > 0 &&
+                allocated_nof_replicas != config.nof_replica_num) {
+                need_nof_eviction_ = true;
+            }
+        }
+        VLOG(1) << "Failed to satisfy replica allocation requirement for key="
+                << key << ", requested_memory_replicas=" << config.replica_num
+                << ", allocated_memory_replicas=" << allocated_memory_replicas
+                << ", requested_nof_replicas=" << config.nof_replica_num
+                << ", allocated_nof_replicas=" << allocated_nof_replicas;
+        abort_reserved_quota();
+        return tl::make_unexpected(ErrorCode::NO_AVAILABLE_HANDLE);
     }
 
     if (use_disk_replica_) {
@@ -838,29 +2784,80 @@ auto MasterService::AllocateAndInsertMetadata(
 
     std::vector<Replica::Descriptor> replica_list;
     replica_list.reserve(replicas.size());
+    int i = 0;
+    VLOG(1) << "PutStart, create replicas: client_id=" << client_id
+            << ", key=" << key << ", value_length=" << value_length;
     for (const auto& replica : replicas) {
-        replica_list.emplace_back(replica.get_descriptor());
+        const auto desc = replica.get_descriptor();
+        replica_list.emplace_back(desc);
+
+        if (replica.is_memory_replica()) {
+            const auto& mem_desc = desc.get_memory_descriptor();
+            VLOG(1) << "Replica #" << ++i << ": buffer_address="
+                    << mem_desc.buffer_descriptor.buffer_address_
+                    << ", transport_endpoint="
+                    << mem_desc.buffer_descriptor.transport_endpoint_;
+        } else if (replica.is_nof_replica()) {
+            const auto& nof_desc = desc.get_nof_descriptor();
+            VLOG(1) << "Replica #" << ++i << ": buffer_address="
+                    << nof_desc.buffer_descriptor.buffer_address_
+                    << ", transport_endpoint="
+                    << nof_desc.buffer_descriptor.transport_endpoint_;
+        }
     }
 
-    shard->metadata.emplace(
+    auto [it, inserted] = tenant_state.metadata.emplace(
         std::piecewise_construct, std::forward_as_tuple(key),
         std::forward_as_tuple(client_id, now, value_length, std::move(replicas),
-                              config.with_soft_pin, config.with_hard_pin));
-    shard->processing_keys.insert(key);
+                              config.with_soft_pin, config.with_hard_pin,
+                              config.data_type, group_id, tenant_id, key));
+    if (!inserted) {
+        LOG(INFO) << "key=" << key << ", info=object_already_exists";
+        abort_reserved_quota();
+        return tl::make_unexpected(ErrorCode::OBJECT_ALREADY_EXISTS);
+    }
+    IncrementTenantMetadataObjectCount(tenant_id);
+    it->second.reserved_quota_charge_bytes = reserved_quota_charge;
+    RegisterGroupMember(tenant_state, tenant_id, key, group_id);
+    tenant_state.processing_keys.insert(key);
 
     return replica_list;
 }
 
 auto MasterService::PutStart(const UUID& client_id, const std::string& key,
+                             const std::string& tenant_id,
                              const uint64_t slice_length,
                              const ReplicateConfig& config)
     -> tl::expected<std::vector<Replica::Descriptor>, ErrorCode> {
-    if (config.replica_num == 0 || key.empty() || slice_length == 0) {
+    auto normalized_tenant_result = NormalizeTenantIdForWrite(tenant_id);
+    if (!normalized_tenant_result) {
+        return tl::make_unexpected(normalized_tenant_result.error());
+    }
+    const ObjectIdentity object_id{normalized_tenant_result.value(), key};
+    if ((config.replica_num == 0 && config.nof_replica_num == 0) ||
+        key.empty() || slice_length == 0) {
         LOG(ERROR) << "key=" << key << ", replica_num=" << config.replica_num
+                   << ", nof_replica_num=" << config.nof_replica_num
                    << ", slice_length=" << slice_length
                    << ", key_size=" << key.size() << ", error=invalid_params";
         return tl::make_unexpected(ErrorCode::INVALID_PARAMS);
     }
+    if (config.prefer_alloc_in_same_node && config.nof_replica_num > 0) {
+        LOG(ERROR) << "key=" << key
+                   << ", nof_replica_num=" << config.nof_replica_num
+                   << ", prefer_alloc_in_same_node="
+                   << config.prefer_alloc_in_same_node
+                   << ", error=nof_not_supported_with_prefer_same_node";
+        return tl::make_unexpected(ErrorCode::INVALID_PARAMS);
+    }
+#ifndef USE_NOF
+    if (config.nof_replica_num > 0) {
+        LOG(ERROR) << "key=" << key
+                   << ", nof_replica_num=" << config.nof_replica_num
+                   << ", error=nof_pool_disabled";
+        return tl::make_unexpected(ErrorCode::INVALID_PARAMS);
+    }
+#endif
 
     if ((memory_allocator_type_ == BufferAllocatorType::CACHELIB) &&
         (slice_length > kMaxSliceSize)) {
@@ -873,45 +2870,142 @@ auto MasterService::PutStart(const UUID& client_id, const std::string& key,
     VLOG(1) << "key=" << key << ", value_length=" << slice_length
             << ", config=" << config << ", action=put_start_begin";
 
-    auto alive_clients = getAliveClientsSnapshot();
-    std::shared_lock<std::shared_mutex> shared_lock(snapshot_mutex_);
-    // Lock the shard and check if object already exists
-    MetadataShardAccessorRW shard(this, getShardIndex(key));
+    auto group_id_result = GetGroupIdForKey(config, 1, 0);
+    if (!group_id_result) {
+        return tl::make_unexpected(group_id_result.error());
+    }
+    const std::string group_id = group_id_result.value();
 
-    const auto now = std::chrono::system_clock::now();
-    auto it = shard->metadata.find(key);
-    if (it != shard->metadata.end() &&
-        !CleanupStaleHandles(it->second, alive_clients)) {
-        auto& metadata = it->second;
-        // If the object's PutStart expired and has not completed any
-        // replicas, we can discard it and allow the new PutStart to
-        // go.
-        if (!metadata.HasReplica(&Replica::fn_is_completed) &&
-            metadata.put_start_time + put_start_discard_timeout_sec_ < now) {
-            auto replicas = metadata.PopReplicas(&Replica::fn_is_processing);
-            if (!replicas.empty()) {
-                std::lock_guard lock(discarded_replicas_mutex_);
-                discarded_replicas_.emplace_back(
-                    std::move(replicas),
-                    metadata.put_start_time + put_start_release_timeout_sec_);
+    [[maybe_unused]] auto object_operation_lock =
+        AcquireObjectOperationLock(object_id.tenant_id, object_id.user_key);
+    const uint64_t requested_quota_charge =
+        RequestedMemoryQuotaCharge(slice_length, config);
+
+    auto attempt_once =
+        [&]() -> tl::expected<std::vector<Replica::Descriptor>, ErrorCode> {
+        std::unique_lock<std::mutex> zero_charge_policy_lock(
+            tenant_quota_policy_mutex_, std::defer_lock);
+        if (ShouldProtectZeroChargeMetadataCreate(requested_quota_charge)) {
+            zero_charge_policy_lock.lock();
+            auto latest_tenant_result =
+                NormalizeTenantIdForWriteLocked(tenant_id);
+            if (!latest_tenant_result) {
+                return tl::make_unexpected(latest_tenant_result.error());
             }
-            shard->processing_keys.erase(key);
-            shard->metadata.erase(it);
-        } else {
+        }
+
+        auto now = std::chrono::system_clock::now();
+        std::optional<size_t> retry_shard_idx;
+        {
+            auto alive_clients = getAliveClientsSnapshot();
+            std::shared_lock<std::shared_mutex> shared_lock(snapshot_mutex_);
+            const size_t lookup_shard_idx =
+                getMetadataShardIndex(object_id.tenant_id, object_id.user_key);
+            MetadataShardAccessorRW shard(this, lookup_shard_idx);
+            auto& tenant_state = shard->tenants[object_id.tenant_id];
+
+            auto it = tenant_state.metadata.find(key);
+            if (it != tenant_state.metadata.end()) {
+                if (CleanupStaleHandles(it->second, alive_clients, &shard)) {
+                    tenant_state.processing_keys.erase(key);
+                    auto task_it = tenant_state.replication_tasks.find(key);
+                    if (task_it != tenant_state.replication_tasks.end()) {
+                        AbortReplicationTaskQuota(object_id.tenant_id,
+                                                  task_it->second);
+                        tenant_state.replication_tasks.erase(task_it);
+                    }
+                    tenant_state.offloading_tasks.erase(key);
+                    ErasePromotionTaskIfPresent(tenant_state, key,
+                                                object_id.tenant_id);
+                    EraseMetadata(tenant_state, it, object_id.tenant_id,
+                                  QuotaEraseMode::kFull, &shard);
+                    it = tenant_state.metadata.end();
+                } else {
+                    auto& metadata = it->second;
+                    if (metadata.HasReplica(&Replica::fn_is_completed) ||
+                        metadata.put_start_time +
+                                put_start_discard_timeout_sec_ >=
+                            now) {
+                        LOG(INFO)
+                            << "key=" << key << ", info=object_already_exists";
+                        return tl::make_unexpected(
+                            ErrorCode::OBJECT_ALREADY_EXISTS);
+                    }
+                    auto replicas =
+                        metadata.PopReplicas(&Replica::fn_is_processing);
+                    if (!replicas.empty()) {
+                        std::lock_guard lock(discarded_replicas_mutex_);
+                        discarded_replicas_.emplace_back(
+                            std::move(replicas),
+                            metadata.put_start_time +
+                                put_start_release_timeout_sec_);
+                    }
+                    tenant_state.processing_keys.erase(key);
+                    EraseMetadata(tenant_state, it, object_id.tenant_id,
+                                  QuotaEraseMode::kFull, &shard);
+                    it = tenant_state.metadata.end();
+                }
+            }
+
+            if (it == tenant_state.metadata.end()) {
+                const size_t target_shard_idx =
+                    group_id.empty()
+                        ? getShardIndex(object_id.tenant_id, object_id.user_key)
+                        : getShardIndex(group_id);
+                if (target_shard_idx != lookup_shard_idx) {
+                    retry_shard_idx = target_shard_idx;
+                    if (tenant_state.Empty()) {
+                        shard->tenants.erase(object_id.tenant_id);
+                    }
+                } else {
+                    return AllocateAndInsertMetadata(
+                        shard, client_id, key, slice_length, config, group_id,
+                        object_id.tenant_id, now);
+                }
+            }
+        }
+
+        std::shared_lock<std::shared_mutex> shared_lock(snapshot_mutex_);
+        MetadataShardAccessorRW shard(this, retry_shard_idx.value());
+        auto& retry_tenant_state = shard->tenants[object_id.tenant_id];
+        if (GetGroupRoute(object_id.tenant_id, object_id.user_key)
+                .has_value() ||
+            retry_tenant_state.metadata.contains(key)) {
             LOG(INFO) << "key=" << key << ", info=object_already_exists";
             return tl::make_unexpected(ErrorCode::OBJECT_ALREADY_EXISTS);
         }
-    }
+        return AllocateAndInsertMetadata(shard, client_id, key, slice_length,
+                                         config, group_id, object_id.tenant_id,
+                                         now);
+    };
 
-    return AllocateAndInsertMetadata(shard, client_id, key, slice_length,
-                                     config, now);
+    for (int attempt = 0; attempt <= kMaxTenantQuotaEvictionRetries;
+         ++attempt) {
+        auto result = attempt_once();
+        if (result.has_value() ||
+            result.error() != ErrorCode::TENANT_QUOTA_EXCEEDED) {
+            return result;
+        }
+        if (attempt == kMaxTenantQuotaEvictionRetries) {
+            MasterMetricManager::instance().inc_tenant_quota_reject(
+                object_id.tenant_id, "quota_exceeded");
+            return result;
+        }
+        EvictTenantMemoryForQuota(
+            object_id.tenant_id,
+            ComputeTenantQuotaDeficit(object_id.tenant_id,
+                                      requested_quota_charge));
+    }
+    return tl::make_unexpected(ErrorCode::TENANT_QUOTA_EXCEEDED);
 }
 
 auto MasterService::PutEnd(const UUID& client_id, const std::string& key,
+                           const std::string& tenant_id,
                            ReplicaType replica_type)
     -> tl::expected<void, ErrorCode> {
     std::shared_lock<std::shared_mutex> shared_lock(snapshot_mutex_);
-    MetadataAccessorRW accessor(this, key);
+    const auto object_id = MakeObjectIdentityForRequest(key, tenant_id);
+    MetadataAccessorRW accessor(this, object_id);
     if (!accessor.Exists()) {
         LOG(ERROR) << "key=" << key << ", error=object_not_found";
         return tl::make_unexpected(ErrorCode::OBJECT_NOT_FOUND);
@@ -926,20 +3020,62 @@ auto MasterService::PutEnd(const UUID& client_id, const std::string& key,
 
     metadata.VisitReplicas(
         [replica_type](const Replica& replica) {
+            if (replica_type == ReplicaType::ALL) {
+                return (replica.is_memory_replica() &&
+                        !replica.has_invalid_mem_handle()) ||
+                       (replica.is_nof_replica() &&
+                        !replica.has_invalid_nof_handle());
+            }
+            if (replica_type == ReplicaType::MEMORY) {
+                return replica.is_memory_replica() &&
+                       !replica.has_invalid_mem_handle();
+            }
+            if (replica_type == ReplicaType::NOF_SSD) {
+                return replica.is_nof_replica() &&
+                       !replica.has_invalid_nof_handle();
+            }
             return replica.type() == replica_type;
         },
         [](Replica& replica) { replica.mark_complete(); });
 
+    const bool has_memory_replica = metadata.HasMemReplica();
+    const bool should_settle_quota =
+        replica_type == ReplicaType::MEMORY ||
+        (replica_type == ReplicaType::ALL && has_memory_replica) ||
+        !has_memory_replica;
+    if (metadata.reserved_quota_charge_bytes > 0 && should_settle_quota) {
+        const uint64_t actual_charge = CompletedMemoryQuotaCharge(metadata);
+        const uint64_t commit_charge =
+            actual_charge > metadata.committed_quota_charge_bytes
+                ? actual_charge - metadata.committed_quota_charge_bytes
+                : 0;
+        const uint64_t abort_charge =
+            metadata.reserved_quota_charge_bytes > commit_charge
+                ? metadata.reserved_quota_charge_bytes - commit_charge
+                : 0;
+        CommitTenantQuota(object_id.tenant_id, commit_charge);
+        AbortTenantQuota(object_id.tenant_id, abort_charge);
+        metadata.reserved_quota_charge_bytes = 0;
+        metadata.committed_quota_charge_bytes = actual_charge;
+        ReleaseTenantQuota(object_id.tenant_id,
+                           metadata.pending_replaced_quota_charge_bytes);
+        metadata.pending_replaced_quota_charge_bytes = 0;
+    }
+
     if (enable_offload_ && !offload_on_evict_) {
-        auto& shard = accessor.GetShard();
+        auto& tenant_state = accessor.GetTenantState();
         metadata.VisitReplicas(
-            &Replica::fn_is_completed, [this, &key, &shard](Replica& replica) {
-                auto result = PushOffloadingQueue(key, replica);
+            [](const Replica& replica) {
+                return replica.is_completed() && replica.is_memory_replica();
+            },
+            [this, &object_id, &tenant_state](Replica& replica) {
+                auto result = PushOffloadingQueue(object_id, replica);
                 if (result) {
                     replica.inc_refcnt();
-                    shard->offloading_tasks.emplace(
-                        key, OffloadingTask{replica.id(),
-                                            std::chrono::system_clock::now()});
+                    tenant_state.offloading_tasks.emplace(
+                        object_id.user_key,
+                        OffloadingTask{replica.id(),
+                                       std::chrono::system_clock::now()});
                 }
             });
     }
@@ -950,11 +3086,8 @@ auto MasterService::PutEnd(const UUID& client_id, const std::string& key,
         accessor.EraseFromProcessing();
     }
 
-    if (replica_type == ReplicaType::MEMORY) {
-        MasterMetricManager::instance().inc_mem_cache_nums();
-    } else if (replica_type == ReplicaType::DISK) {
-        MasterMetricManager::instance().inc_file_cache_nums();
-    }
+    SyncCacheTotalAccounting(metadata);
+    // TODO: add inc_nof_cache_nums() (ranhaojia)
     // 1. Set lease timeout to now, indicating that the object has no lease
     // at beginning. 2. If this object has soft pin enabled, set it to be soft
     // pinned.
@@ -963,10 +3096,23 @@ auto MasterService::PutEnd(const UUID& client_id, const std::string& key,
 }
 
 auto MasterService::AddReplica(const UUID& client_id, const std::string& key,
-                               Replica& replica)
-    -> tl::expected<void, ErrorCode> {
+                               const std::string& tenant_id, Replica& replica)
+    -> tl::expected<bool, ErrorCode> {
+    std::string normalized_tenant = "default";
+    std::unique_lock<std::mutex> policy_lock(tenant_quota_policy_mutex_,
+                                             std::defer_lock);
+    if (enable_multi_tenants_) {
+        policy_lock.lock();
+        auto normalized_tenant_result =
+            NormalizeTenantIdForWriteLocked(tenant_id);
+        if (!normalized_tenant_result) {
+            return tl::make_unexpected(normalized_tenant_result.error());
+        }
+        normalized_tenant = std::move(normalized_tenant_result.value());
+    }
     std::shared_lock<std::shared_mutex> shared_lock(snapshot_mutex_);
-    MetadataAccessorRW accessor(this, key);
+    const ObjectIdentity object_id{normalized_tenant, key};
+    MetadataAccessorRW accessor(this, object_id);
     if (!accessor.Exists()) {
         accessor.Create(
             client_id,
@@ -984,7 +3130,9 @@ auto MasterService::AddReplica(const UUID& client_id, const std::string& key,
         std::vector<Replica> replicas;
         replicas.emplace_back(std::move(replica));
         metadata.AddReplicas(std::move(replicas));
-        return {};
+        auto& shard = accessor.GetShard();
+        shard.OnDiskReplicaAdded(metadata);
+        return true;
     }
 
     metadata.VisitReplicas(
@@ -1004,14 +3152,16 @@ auto MasterService::AddReplica(const UUID& client_id, const std::string& key,
                     .get_local_disk_descriptor()
                     .object_size;
         });
-    return {};
+    return false;
 }
 
 auto MasterService::PutRevoke(const UUID& client_id, const std::string& key,
+                              const std::string& tenant_id,
                               ReplicaType replica_type)
     -> tl::expected<void, ErrorCode> {
     std::shared_lock<std::shared_mutex> shared_lock(snapshot_mutex_);
-    MetadataAccessorRW accessor(this, key);
+    const auto object_id = MakeObjectIdentityForRequest(key, tenant_id);
+    MetadataAccessorRW accessor(this, object_id);
     if (!accessor.Exists()) {
         LOG(INFO) << "key=" << key << ", info=object_not_found";
         return tl::make_unexpected(ErrorCode::OBJECT_NOT_FOUND);
@@ -1024,25 +3174,45 @@ auto MasterService::PutRevoke(const UUID& client_id, const std::string& key,
         return tl::make_unexpected(ErrorCode::ILLEGAL_CLIENT);
     }
 
-    auto processing_rep =
-        metadata.GetFirstReplica([replica_type](const Replica& replica) {
-            return replica.type() == replica_type && !replica.is_processing();
-        });
+    auto processing_rep = metadata.GetFirstReplica([replica_type](
+                                                       const Replica& replica) {
+        if (replica_type == ReplicaType::ALL) {
+            return (replica.is_memory_replica() || replica.is_nof_replica()) &&
+                   !replica.is_processing();
+        }
+        return replica.type() == replica_type && !replica.is_processing();
+    });
     if (processing_rep != nullptr) {
         LOG(ERROR) << "key=" << key << ", status=" << processing_rep->status()
                    << ", error=invalid_replica_status";
         return tl::make_unexpected(ErrorCode::INVALID_WRITE);
     }
 
-    if (replica_type == ReplicaType::MEMORY) {
-        MasterMetricManager::instance().dec_mem_cache_nums();
-    } else if (replica_type == ReplicaType::DISK) {
-        MasterMetricManager::instance().dec_file_cache_nums();
+    const uint64_t before_charge = CompletedMemoryQuotaCharge(metadata);
+    EraseReplicasWithCacheTotalAccounting(
+        metadata, [replica_type](const Replica& replica) {
+            if (replica_type == ReplicaType::ALL) {
+                return replica.is_memory_replica() || replica.is_nof_replica();
+            }
+            return replica.type() == replica_type;
+        });
+    const uint64_t after_charge = CompletedMemoryQuotaCharge(metadata);
+    if (before_charge > after_charge) {
+        ReleaseCommittedQuotaCharge(metadata, before_charge - after_charge);
+    }
+    if (!metadata.HasReplica(&Replica::fn_is_memory_replica)) {
+        AbortTenantQuota(object_id.tenant_id,
+                         metadata.reserved_quota_charge_bytes);
+        metadata.reserved_quota_charge_bytes = 0;
     }
 
-    metadata.EraseReplicas([replica_type](const Replica& replica) {
-        return replica.type() == replica_type;
-    });
+    EraseReplicasWithCacheTotalAccounting(
+        metadata, [replica_type](const Replica& replica) {
+            if (replica_type == ReplicaType::ALL) {
+                return replica.is_memory_replica() || replica.is_nof_replica();
+            }
+            return replica.type() == replica_type;
+        });
 
     // If the object is completed, remove it from the processing set.
     if (metadata.AllReplicas(&Replica::fn_is_completed) &&
@@ -1057,21 +3227,24 @@ auto MasterService::PutRevoke(const UUID& client_id, const std::string& key,
 }
 
 std::vector<tl::expected<void, ErrorCode>> MasterService::BatchPutEnd(
-    const UUID& client_id, const std::vector<std::string>& keys) {
+    const UUID& client_id, const std::vector<std::string>& keys,
+    const std::string& tenant_id, ReplicaType replica_type) {
     std::vector<tl::expected<void, ErrorCode>> results;
     results.reserve(keys.size());
     for (const auto& key : keys) {
-        results.emplace_back(PutEnd(client_id, key, ReplicaType::MEMORY));
+        results.emplace_back(PutEnd(client_id, key, tenant_id, replica_type));
     }
     return results;
 }
 
 std::vector<tl::expected<void, ErrorCode>> MasterService::BatchPutRevoke(
-    const UUID& client_id, const std::vector<std::string>& keys) {
+    const UUID& client_id, const std::vector<std::string>& keys,
+    const std::string& tenant_id, ReplicaType replica_type) {
     std::vector<tl::expected<void, ErrorCode>> results;
     results.reserve(keys.size());
     for (const auto& key : keys) {
-        results.emplace_back(PutRevoke(client_id, key, ReplicaType::MEMORY));
+        results.emplace_back(
+            PutRevoke(client_id, key, tenant_id, replica_type));
     }
     return results;
 }
@@ -1091,16 +3264,40 @@ std::vector<tl::expected<void, ErrorCode>> MasterService::BatchPutRevoke(
 // Note: during Case B the key is temporarily unreadable (all replicas are
 // PROCESSING).  Readers will get REPLICA_IS_NOT_READY until UpsertEnd.
 auto MasterService::UpsertStart(const UUID& client_id, const std::string& key,
+                                const std::string& tenant_id,
                                 const uint64_t slice_length,
                                 const ReplicateConfig& config)
     -> tl::expected<std::vector<Replica::Descriptor>, ErrorCode> {
+    auto normalized_tenant_result = NormalizeTenantIdForWrite(tenant_id);
+    if (!normalized_tenant_result) {
+        return tl::make_unexpected(normalized_tenant_result.error());
+    }
+    const ObjectIdentity object_id{normalized_tenant_result.value(), key};
     // --- Parameter validation (same as PutStart) ---
-    if (config.replica_num == 0 || key.empty() || slice_length == 0) {
+    if ((config.replica_num == 0 && config.nof_replica_num == 0) ||
+        key.empty() || slice_length == 0) {
         LOG(ERROR) << "key=" << key << ", replica_num=" << config.replica_num
+                   << ", nof_replica_num=" << config.nof_replica_num
                    << ", slice_length=" << slice_length
                    << ", key_size=" << key.size() << ", error=invalid_params";
         return tl::make_unexpected(ErrorCode::INVALID_PARAMS);
     }
+    if (config.prefer_alloc_in_same_node && config.nof_replica_num > 0) {
+        LOG(ERROR) << "key=" << key
+                   << ", nof_replica_num=" << config.nof_replica_num
+                   << ", prefer_alloc_in_same_node="
+                   << config.prefer_alloc_in_same_node
+                   << ", error=nof_not_supported_with_prefer_same_node";
+        return tl::make_unexpected(ErrorCode::INVALID_PARAMS);
+    }
+#ifndef USE_NOF
+    if (config.nof_replica_num > 0) {
+        LOG(ERROR) << "key=" << key
+                   << ", nof_replica_num=" << config.nof_replica_num
+                   << ", error=nof_pool_disabled";
+        return tl::make_unexpected(ErrorCode::INVALID_PARAMS);
+    }
+#endif
 
     if ((memory_allocator_type_ == BufferAllocatorType::CACHELIB) &&
         (slice_length > kMaxSliceSize)) {
@@ -1113,173 +3310,302 @@ auto MasterService::UpsertStart(const UUID& client_id, const std::string& key,
     VLOG(1) << "key=" << key << ", value_length=" << slice_length
             << ", config=" << config << ", action=upsert_start_begin";
 
-    // --- Lock acquisition ---
-    // snapshot_mutex_ (shared): allows concurrent reads/writes, blocks only
-    //   during full metadata snapshots.
-    // shard lock (exclusive via MetadataShardAccessorRW): serializes all
-    //   operations on keys that hash to the same shard.
-    auto alive_clients = getAliveClientsSnapshot();
-    std::shared_lock<std::shared_mutex> shared_lock(snapshot_mutex_);
-    MetadataShardAccessorRW shard(this, getShardIndex(key));
-
-    const auto now = std::chrono::system_clock::now();
-    auto it = shard->metadata.find(key);
-
-    // --- Step 0: stale handle cleanup ---
-    // If all memory replicas point to unmounted segments (node crashed and
-    // restarted), the metadata is useless — erase it and treat as new key.
-    // Also clean up local_disk replicas whose owner client has expired.
-    if (it != shard->metadata.end() &&
-        CleanupStaleHandles(it->second, alive_clients)) {
-        shard->processing_keys.erase(key);
-        shard->metadata.erase(it);
-        it = shard->metadata.end();
+    auto group_id_result = GetGroupIdForKey(config, 1, 0);
+    if (!group_id_result) {
+        return tl::make_unexpected(group_id_result.error());
     }
+    const std::string group_id = group_id_result.value();
 
-    // --- Step 1: safety checks and preemption (only if key exists) ---
-    if (it != shard->metadata.end()) {
-        auto& metadata = it->second;
+    [[maybe_unused]] auto object_operation_lock =
+        AcquireObjectOperationLock(object_id.tenant_id, object_id.user_key);
+    const uint64_t requested_quota_charge =
+        RequestedMemoryQuotaCharge(slice_length, config);
 
-        // Reject if a Copy/Move task is actively reading this key's replicas.
-        // Writing during replication would corrupt the copy.
-        if (shard->replication_tasks.count(key) > 0) {
-            LOG(INFO) << "key=" << key << ", error=object_has_replication_task";
-            return tl::make_unexpected(ErrorCode::OBJECT_HAS_REPLICATION_TASK);
-        }
-
-        // Reject if an offload-to-disk task is in progress (same reason).
-        if (shard->offloading_tasks.count(key) > 0) {
-            LOG(INFO) << "key=" << key << ", error=object_has_offloading_task";
-            return tl::make_unexpected(ErrorCode::OBJECT_HAS_REPLICATION_TASK);
-        }
-
-        // Preempt an in-progress Put/Upsert on the same key.  The previous
-        // writer's PROCESSING replicas are moved to discarded_replicas_ with a
-        // TTL so they are not freed while the old writer may still be doing
-        // RDMA writes.  Unlike PutStart (which only preempts after a timeout),
-        // UpsertStart preempts immediately.
-        if (shard->processing_keys.count(key) > 0) {
-            auto processing_replicas =
-                metadata.PopReplicas(&Replica::fn_is_processing);
-            if (!processing_replicas.empty()) {
-                std::lock_guard lock(discarded_replicas_mutex_);
-                discarded_replicas_.emplace_back(
-                    std::move(processing_replicas),
-                    now + put_start_release_timeout_sec_);
-            }
-            shard->processing_keys.erase(key);
-
-            // If no COMPLETE replicas survive the preemption, this key
-            // effectively does not exist — fall through to Case A.
-            if (!metadata.HasReplica(&Replica::fn_is_completed)) {
-                shard->metadata.erase(it);
-                it = shard->metadata.end();
+    auto attempt_once =
+        [&]() -> tl::expected<std::vector<Replica::Descriptor>, ErrorCode> {
+        std::unique_lock<std::mutex> zero_charge_policy_lock(
+            tenant_quota_policy_mutex_, std::defer_lock);
+        if (ShouldProtectZeroChargeMetadataCreate(requested_quota_charge)) {
+            zero_charge_policy_lock.lock();
+            auto latest_tenant_result =
+                NormalizeTenantIdForWriteLocked(tenant_id);
+            if (!latest_tenant_result) {
+                return tl::make_unexpected(latest_tenant_result.error());
             }
         }
-    }
 
-    // --- Case A: key does not exist (or was erased above) ---
-    // Allocate fresh buffers, identical to PutStart.
-    if (it == shard->metadata.end()) {
-        VLOG(1) << "key=" << key << ", action=upsert_start_case_a";
-        return AllocateAndInsertMetadata(shard, client_id, key, slice_length,
-                                         config, now);
-    }
-
-    // --- Step 2: key exists with COMPLETE replicas → Case B or C ---
-    auto& metadata = it->second;
-
-    // Reject if any reader holds a reference (refcnt > 0).  Overwriting a
-    // buffer that an RDMA read is streaming from would cause data corruption.
-    // The client should retry after readers finish.
-    if (metadata.HasReplica(&Replica::fn_is_busy)) {
-        LOG(INFO) << "key=" << key << ", error=object_replica_busy";
-        return tl::make_unexpected(ErrorCode::OBJECT_REPLICA_BUSY);
-    }
-
-    if (metadata.size == slice_length) {
-        // --- Case B: same size — in-place update ---
-        // Reuse existing buffer addresses.  No allocation or deallocation.
-        // The client will RDMA-write new data to the same addresses.
-        //
-        // hard_pinned is const and preserved automatically — upsert does not
-        // change the eviction protection level of an existing object.
-        metadata.client_id = client_id;
-        metadata.put_start_time = now;
-
-        // Reconcile soft_pin state with the incoming config.
+        auto now = std::chrono::system_clock::now();
+        std::optional<size_t> case_a_retry_shard_idx;
         {
-            SpinLocker locker(&metadata.lock);
-            if (config.with_soft_pin && !metadata.soft_pin_timeout) {
-                metadata.soft_pin_timeout.emplace();
-                MasterMetricManager::instance().inc_soft_pin_key_count(1);
-            } else if (!config.with_soft_pin && metadata.soft_pin_timeout) {
-                metadata.soft_pin_timeout.reset();
-                MasterMetricManager::instance().dec_soft_pin_key_count(1);
+            // --- Lock acquisition ---
+            auto alive_clients = getAliveClientsSnapshot();
+            std::shared_lock<std::shared_mutex> shared_lock(snapshot_mutex_);
+            // Use getMetadataShardIndex to find the object at its current shard
+            // (handles both grouped and ungrouped routing).
+            const size_t lookup_shard_idx =
+                getMetadataShardIndex(object_id.tenant_id, object_id.user_key);
+            MetadataShardAccessorRW shard(this, lookup_shard_idx);
+            auto& tenant_state = shard->tenants[object_id.tenant_id];
+
+            auto it = tenant_state.metadata.find(key);
+
+            // --- Step 0: stale handle cleanup ---
+            if (it != tenant_state.metadata.end() &&
+                CleanupStaleHandles(it->second, alive_clients, &shard)) {
+                tenant_state.processing_keys.erase(key);
+                ErasePromotionTaskIfPresent(tenant_state, key,
+                                            object_id.tenant_id);
+                EraseMetadata(tenant_state, it, object_id.tenant_id,
+                              QuotaEraseMode::kFull, &shard);
+                it = tenant_state.metadata.end();
+            }
+
+            // --- Step 1: safety checks and preemption (only if key exists) ---
+            if (it != tenant_state.metadata.end()) {
+                auto& metadata = it->second;
+
+                // Reject if the caller tries to change group membership.
+                // Group membership is immutable while an object exists.
+                if (config.group_ids.has_value() &&
+                    metadata.group_id != group_id) {
+                    LOG(ERROR) << "key=" << key
+                               << ", error=group_membership_is_immutable";
+                    return tl::make_unexpected(ErrorCode::INVALID_PARAMS);
+                }
+
+                // Reject if a Copy/Move task is actively reading this key's
+                // replicas.
+                if (tenant_state.replication_tasks.count(key) > 0) {
+                    LOG(INFO) << "key=" << key
+                              << ", error=object_has_replication_task";
+                    return tl::make_unexpected(
+                        ErrorCode::OBJECT_HAS_REPLICATION_TASK);
+                }
+
+                // Reject if an offload-to-disk task is in progress (same
+                // reason).
+                if (tenant_state.offloading_tasks.count(key) > 0) {
+                    LOG(INFO) << "key=" << key
+                              << ", error=object_has_offloading_task";
+                    return tl::make_unexpected(
+                        ErrorCode::OBJECT_HAS_REPLICATION_TASK);
+                }
+
+                // Preempt an in-progress Put/Upsert on the same key.  The
+                // previous writer's PROCESSING replicas are moved to
+                // discarded_replicas_ with a TTL so they are not freed while
+                // the old writer may still be doing RDMA writes.  Unlike
+                // PutStart (which only preempts after a timeout), UpsertStart
+                // preempts immediately.
+                if (tenant_state.processing_keys.count(key) > 0) {
+                    auto processing_replicas =
+                        metadata.PopReplicas(&Replica::fn_is_processing);
+                    if (!processing_replicas.empty()) {
+                        std::lock_guard lock(discarded_replicas_mutex_);
+                        discarded_replicas_.emplace_back(
+                            std::move(processing_replicas),
+                            now + put_start_release_timeout_sec_);
+                    }
+                    tenant_state.processing_keys.erase(key);
+
+                    // If no COMPLETE replicas survive the preemption, this key
+                    // effectively does not exist — fall through to Case A.
+                    if (!metadata.HasReplica(&Replica::fn_is_completed)) {
+                        ErasePromotionTaskIfPresent(tenant_state, key,
+                                                    object_id.tenant_id);
+                        EraseMetadata(tenant_state, it, object_id.tenant_id,
+                                      QuotaEraseMode::kFull, &shard);
+                        it = tenant_state.metadata.end();
+                    }
+                }
+            }
+
+            // --- Case A: key does not exist (or was erased above) ---
+            // Allocate fresh buffers, identical to PutStart.
+            if (it == tenant_state.metadata.end()) {
+                VLOG(1) << "key=" << key << ", action=upsert_start_case_a";
+                const size_t case_a_shard_idx =
+                    group_id.empty()
+                        ? getShardIndex(object_id.tenant_id, object_id.user_key)
+                        : getShardIndex(group_id);
+                if (case_a_shard_idx != lookup_shard_idx) {
+                    case_a_retry_shard_idx = case_a_shard_idx;
+                    if (tenant_state.Empty()) {
+                        shard->tenants.erase(object_id.tenant_id);
+                    }
+                } else {
+                    return AllocateAndInsertMetadata(
+                        shard, client_id, key, slice_length, config, group_id,
+                        object_id.tenant_id, now);
+                }
+            } else {
+                // --- Step 2: key exists with COMPLETE replicas → Case B or C
+                // ---
+                auto& metadata = it->second;
+
+                // Reject if any reader holds a reference (refcnt > 0).
+                // Overwriting a buffer that an RDMA read is streaming from
+                // would cause data corruption. The client should retry after
+                // readers finish.
+                if (metadata.HasReplica(&Replica::fn_is_busy)) {
+                    LOG(INFO) << "key=" << key << ", error=object_replica_busy";
+                    return tl::make_unexpected(ErrorCode::OBJECT_REPLICA_BUSY);
+                }
+
+                if (metadata.size == slice_length) {
+                    // --- Case B: same size — in-place update ---
+                    // Reuse existing buffer addresses.  No allocation or
+                    // deallocation. The client will RDMA-write new data to the
+                    // same addresses.
+                    //
+                    // hard_pinned is const and preserved automatically — upsert
+                    // does not change the eviction protection level of an
+                    // existing object.
+                    metadata.client_id = client_id;
+                    metadata.put_start_time = now;
+
+                    // Reconcile soft_pin state with the incoming config.
+                    {
+                        SpinLocker locker(&metadata.lock);
+                        if (config.with_soft_pin &&
+                            !metadata.soft_pin_timeout) {
+                            metadata.soft_pin_timeout.emplace();
+                            MasterMetricManager::instance()
+                                .inc_soft_pin_key_count(1);
+                        } else if (!config.with_soft_pin &&
+                                   metadata.soft_pin_timeout) {
+                            metadata.soft_pin_timeout.reset();
+                            MasterMetricManager::instance()
+                                .dec_soft_pin_key_count(1);
+                        }
+                    }
+
+                    // Mark COMPLETE → PROCESSING so readers won't see stale
+                    // data mid-transfer.  The key becomes unreadable until
+                    // UpsertEnd.
+                    metadata.VisitReplicas(
+                        &Replica::fn_is_completed,
+                        [](Replica& replica) { replica.mark_processing(); });
+                    SyncCacheTotalAccounting(metadata);
+
+                    tenant_state.processing_keys.insert(key);
+
+                    // Return the existing descriptors — same buffer addresses
+                    // as before.
+                    std::vector<Replica::Descriptor> replica_list;
+                    const auto& all_replicas = metadata.GetAllReplicas();
+                    replica_list.reserve(all_replicas.size());
+                    for (const auto& replica : all_replicas) {
+                        replica_list.emplace_back(replica.get_descriptor());
+                    }
+
+                    VLOG(1) << "key=" << key
+                            << ", action=upsert_start_case_b_inplace";
+                    return replica_list;
+                }
+
+                // --- Case C: different size — discard old replicas and
+                // reallocate
+                // --- Old buffers cannot be reused.  Move them to
+                // discarded_replicas_ for delayed release (readers may still
+                // hold descriptors without refcnt), then allocate fresh buffers
+                // at the new size.
+                //
+                // Preserve hard_pin and soft_pin from the old metadata so that
+                // eviction protection survives a size-changing upsert (RFC
+                // §2.2.2).
+                ReplicateConfig merged_config = config;
+                merged_config.with_hard_pin =
+                    merged_config.with_hard_pin || metadata.IsHardPinned();
+                merged_config.with_soft_pin =
+                    merged_config.with_soft_pin || metadata.IsSoftPinned();
+
+                const std::string existing_group_id = metadata.group_id;
+                const uint64_t old_quota_charge =
+                    metadata.committed_quota_charge_bytes != 0
+                        ? metadata.committed_quota_charge_bytes
+                        : CompletedMemoryQuotaCharge(metadata);
+                auto old_replicas =
+                    PopReplicasWithCacheTotalAccounting(metadata);
+                if (!old_replicas.empty()) {
+                    std::lock_guard lock(discarded_replicas_mutex_);
+                    discarded_replicas_.emplace_back(
+                        std::move(old_replicas),
+                        now + put_start_release_timeout_sec_);
+                }
+                EraseMetadata(tenant_state, it, object_id.tenant_id,
+                              QuotaEraseMode::kPreserveOld, &shard);
+
+                VLOG(1) << "key=" << key
+                        << ", action=upsert_start_case_c_reallocate";
+                auto allocate_result = AllocateAndInsertMetadata(
+                    shard, client_id, key, slice_length, merged_config,
+                    existing_group_id, object_id.tenant_id, now);
+                if (!allocate_result) {
+                    ReleaseTenantQuota(object_id.tenant_id, old_quota_charge);
+                    return allocate_result;
+                }
+                auto new_it = tenant_state.metadata.find(key);
+                if (new_it != tenant_state.metadata.end()) {
+                    new_it->second.pending_replaced_quota_charge_bytes =
+                        old_quota_charge;
+                }
+                return allocate_result;
             }
         }
-
-        // Mark COMPLETE → PROCESSING so readers won't see stale data
-        // mid-transfer.  The key becomes unreadable until UpsertEnd.
-        metadata.VisitReplicas(&Replica::fn_is_completed, [](Replica& replica) {
-            replica.mark_processing();
-        });
-
-        shard->processing_keys.insert(key);
-
-        // Return the existing descriptors — same buffer addresses as before.
-        std::vector<Replica::Descriptor> replica_list;
-        const auto& all_replicas = metadata.GetAllReplicas();
-        replica_list.reserve(all_replicas.size());
-        for (const auto& replica : all_replicas) {
-            replica_list.emplace_back(replica.get_descriptor());
+        std::shared_lock<std::shared_mutex> shared_lock(snapshot_mutex_);
+        MetadataShardAccessorRW shard(this, case_a_retry_shard_idx.value());
+        auto& retry_tenant_state = shard->tenants[object_id.tenant_id];
+        const auto current_route =
+            GetGroupRoute(object_id.tenant_id, object_id.user_key);
+        if (current_route.has_value() ||
+            retry_tenant_state.metadata.contains(key)) {
+            LOG(INFO) << "key=" << key << ", info=object_already_exists";
+            return tl::make_unexpected(ErrorCode::OBJECT_ALREADY_EXISTS);
         }
+        return AllocateAndInsertMetadata(shard, client_id, key, slice_length,
+                                         config, group_id, object_id.tenant_id,
+                                         now);
+    };
 
-        VLOG(1) << "key=" << key << ", action=upsert_start_case_b_inplace";
-        return replica_list;
+    for (int attempt = 0; attempt <= kMaxTenantQuotaEvictionRetries;
+         ++attempt) {
+        auto result = attempt_once();
+        if (result.has_value() ||
+            result.error() != ErrorCode::TENANT_QUOTA_EXCEEDED) {
+            return result;
+        }
+        if (attempt == kMaxTenantQuotaEvictionRetries) {
+            MasterMetricManager::instance().inc_tenant_quota_reject(
+                object_id.tenant_id, "quota_exceeded");
+            return result;
+        }
+        EvictTenantMemoryForQuota(
+            object_id.tenant_id,
+            ComputeTenantQuotaDeficit(object_id.tenant_id,
+                                      requested_quota_charge));
     }
-
-    // --- Case C: different size — discard old replicas and reallocate ---
-    // Old buffers cannot be reused.  Move them to discarded_replicas_ for
-    // delayed release (readers may still hold descriptors without refcnt),
-    // then allocate fresh buffers at the new size.
-    //
-    // Preserve hard_pin and soft_pin from the old metadata so that eviction
-    // protection survives a size-changing upsert (RFC §2.2.2).
-    ReplicateConfig merged_config = config;
-    merged_config.with_hard_pin =
-        merged_config.with_hard_pin || metadata.IsHardPinned();
-    merged_config.with_soft_pin =
-        merged_config.with_soft_pin || metadata.IsSoftPinned();
-
-    auto old_replicas = metadata.PopReplicas();
-    if (!old_replicas.empty()) {
-        std::lock_guard lock(discarded_replicas_mutex_);
-        discarded_replicas_.emplace_back(std::move(old_replicas),
-                                         now + put_start_release_timeout_sec_);
-    }
-    shard->metadata.erase(it);
-
-    VLOG(1) << "key=" << key << ", action=upsert_start_case_c_reallocate";
-    return AllocateAndInsertMetadata(shard, client_id, key, slice_length,
-                                     merged_config, now);
+    return tl::make_unexpected(ErrorCode::TENANT_QUOTA_EXCEEDED);
 }
 
 auto MasterService::UpsertEnd(const UUID& client_id, const std::string& key,
+                              const std::string& tenant_id,
                               ReplicaType replica_type)
     -> tl::expected<void, ErrorCode> {
-    return PutEnd(client_id, key, replica_type);
+    return PutEnd(client_id, key, tenant_id, replica_type);
 }
 
 auto MasterService::UpsertRevoke(const UUID& client_id, const std::string& key,
+                                 const std::string& tenant_id,
                                  ReplicaType replica_type)
     -> tl::expected<void, ErrorCode> {
-    return PutRevoke(client_id, key, replica_type);
+    return PutRevoke(client_id, key, tenant_id, replica_type);
 }
 
 std::vector<tl::expected<std::vector<Replica::Descriptor>, ErrorCode>>
 MasterService::BatchUpsertStart(const UUID& client_id,
                                 const std::vector<std::string>& keys,
+                                const std::string& tenant_id,
                                 const std::vector<uint64_t>& slice_lengths,
                                 const ReplicateConfig& config) {
     if (keys.size() != slice_lengths.size()) {
@@ -1289,49 +3615,75 @@ MasterService::BatchUpsertStart(const UUID& client_id,
             tl::expected<std::vector<Replica::Descriptor>, ErrorCode>>(
             keys.size(), tl::make_unexpected(ErrorCode::INVALID_PARAMS));
     }
+    if (config.group_ids.has_value() &&
+        config.group_ids->size() != keys.size()) {
+        LOG(ERROR) << "BatchUpsertStart: group_ids.size()="
+                   << config.group_ids->size()
+                   << " != keys.size()=" << keys.size();
+        return std::vector<
+            tl::expected<std::vector<Replica::Descriptor>, ErrorCode>>(
+            keys.size(), tl::make_unexpected(ErrorCode::INVALID_PARAMS));
+    }
     std::vector<tl::expected<std::vector<Replica::Descriptor>, ErrorCode>>
         results;
     results.reserve(keys.size());
     for (size_t i = 0; i < keys.size(); ++i) {
-        results.emplace_back(
-            UpsertStart(client_id, keys[i], slice_lengths[i], config));
+        auto key_config = config.ForSingleKey(i);
+        results.emplace_back(UpsertStart(client_id, keys[i], tenant_id,
+                                         slice_lengths[i], key_config));
     }
     return results;
 }
 
 std::vector<tl::expected<void, ErrorCode>> MasterService::BatchUpsertEnd(
-    const UUID& client_id, const std::vector<std::string>& keys) {
-    return BatchPutEnd(client_id, keys);
+    const UUID& client_id, const std::vector<std::string>& keys,
+    const std::string& tenant_id) {
+    return BatchPutEnd(client_id, keys, tenant_id);
 }
 
 std::vector<tl::expected<void, ErrorCode>> MasterService::BatchUpsertRevoke(
-    const UUID& client_id, const std::vector<std::string>& keys) {
-    return BatchPutRevoke(client_id, keys);
+    const UUID& client_id, const std::vector<std::string>& keys,
+    const std::string& tenant_id) {
+    return BatchPutRevoke(client_id, keys, tenant_id);
 }
 
 auto MasterService::EvictDiskReplica(const UUID& client_id,
                                      const std::string& key,
+                                     const std::string& tenant_id,
                                      ReplicaType replica_type)
     -> tl::expected<void, ErrorCode> {
-    MetadataAccessorRW accessor(this, key);
+    const auto object_id = MakeObjectIdentityForRequest(key, tenant_id);
+    MetadataAccessorRW accessor(this, object_id);
     if (!accessor.Exists()) {
-        LOG(INFO) << "key=" << key << ", info=object_not_found_for_eviction";
+        LOG(INFO) << "key=" << key << ", tenant_id=" << object_id.tenant_id
+                  << ", info=object_not_found_for_eviction";
         return tl::make_unexpected(ErrorCode::OBJECT_NOT_FOUND);
     }
 
     auto& metadata = accessor.Get();
 
     if (replica_type == ReplicaType::DISK) {
-        metadata.EraseReplicas(
+        EraseReplicasWithCacheTotalAccounting(
+            metadata,
             [](const Replica& replica) { return replica.is_disk_replica(); });
-        MasterMetricManager::instance().dec_file_cache_nums();
     } else if (replica_type == ReplicaType::LOCAL_DISK) {
-        metadata.EraseReplicas([&client_id](const Replica& replica) {
-            return replica.is_local_disk_replica() &&
-                   replica.get_descriptor()
-                           .get_local_disk_descriptor()
-                           .client_id == client_id;
+        bool had_completed_disk = metadata.HasReplica([](const Replica& r) {
+            return r.is_local_disk_replica() && r.is_completed();
         });
+        // PopReplicas (vs EraseReplicas) so we can release the SSD usage of
+        // the removed local-disk replicas; both remove the same matching set.
+        auto erased_replicas =
+            metadata.PopReplicas([&client_id](const Replica& replica) {
+                return replica.is_local_disk_replica() &&
+                       replica.get_descriptor()
+                               .get_local_disk_descriptor()
+                               .client_id == client_id;
+            });
+        ReleaseLocalDiskUsage(erased_replicas);
+        if (had_completed_disk) {
+            auto& shard = accessor.GetShard();
+            shard.OnDiskReplicaRemoved(had_completed_disk, metadata);
+        }
     } else {
         LOG(ERROR) << "key=" << key
                    << ", error=invalid_replica_type_for_eviction";
@@ -1346,19 +3698,25 @@ auto MasterService::EvictDiskReplica(const UUID& client_id,
 
 std::vector<tl::expected<void, ErrorCode>> MasterService::BatchEvictDiskReplica(
     const UUID& client_id, const std::vector<std::string>& keys,
-    ReplicaType replica_type) {
+    const std::string& tenant_id, ReplicaType replica_type) {
     std::vector<tl::expected<void, ErrorCode>> results;
     results.reserve(keys.size());
     for (const auto& key : keys) {
-        results.push_back(EvictDiskReplica(client_id, key, replica_type));
+        results.push_back(
+            EvictDiskReplica(client_id, key, tenant_id, replica_type));
     }
     return results;
 }
 
 tl::expected<CopyStartResponse, ErrorCode> MasterService::CopyStart(
-    const UUID& client_id, const std::string& key,
+    const UUID& client_id, const std::string& key, const std::string& tenant_id,
     const std::string& src_segment,
     const std::vector<std::string>& tgt_segments) {
+    auto normalized_tenant_result = NormalizeTenantIdForWrite(tenant_id);
+    if (!normalized_tenant_result) {
+        return tl::make_unexpected(normalized_tenant_result.error());
+    }
+    const ObjectIdentity object_id{normalized_tenant_result.value(), key};
     std::shared_lock<std::shared_mutex> shared_lock(snapshot_mutex_);
     {
         ScopedSegmentAccess segment_access =
@@ -1377,7 +3735,7 @@ tl::expected<CopyStartResponse, ErrorCode> MasterService::CopyStart(
             }
         }
     }
-    MetadataAccessorRW accessor(this, key);
+    MetadataAccessorRW accessor(this, object_id);
     if (!accessor.Exists()) {
         LOG(ERROR) << "key=" << key << ", object not found";
         return tl::make_unexpected(ErrorCode::OBJECT_NOT_FOUND);
@@ -1398,8 +3756,31 @@ tl::expected<CopyStartResponse, ErrorCode> MasterService::CopyStart(
         return tl::make_unexpected(ErrorCode::REPLICA_NOT_FOUND);
     }
 
+    size_t new_replica_count = 0;
+    for (const auto& tgt_segment : tgt_segments) {
+        if (metadata.GetReplicaBySegmentName(tgt_segment) == nullptr) {
+            ++new_replica_count;
+        }
+    }
+
+    const uint64_t reserved_quota_charge =
+        SaturatingMultiply(static_cast<uint64_t>(metadata.size),
+                           static_cast<uint64_t>(new_replica_count));
+    auto quota_result =
+        ReserveTenantQuota(object_id.tenant_id, reserved_quota_charge);
+    if (!quota_result) {
+        if (quota_result.error() == ErrorCode::TENANT_QUOTA_EXCEEDED) {
+            MasterMetricManager::instance().inc_tenant_quota_reject(
+                object_id.tenant_id, "quota_exceeded");
+        }
+        return tl::make_unexpected(quota_result.error());
+    }
+    auto abort_reserved_quota = [&] {
+        AbortTenantQuota(object_id.tenant_id, reserved_quota_charge);
+    };
+
     std::vector<Replica> replicas;
-    replicas.reserve(tgt_segments.size());
+    replicas.reserve(new_replica_count);
     {
         ScopedAllocatorAccess allocator_access =
             segment_manager_.getAllocatorAccess();
@@ -1416,6 +3797,7 @@ tl::expected<CopyStartResponse, ErrorCode> MasterService::CopyStart(
             if (!replica.has_value()) {
                 LOG(ERROR) << "key=" << key << ", tgt_segment=" << tgt_segment
                            << ", failed to allocate replica";
+                abort_reserved_quota();
                 return tl::make_unexpected(replica.error());
             }
             replicas.push_back(std::move(*replica));
@@ -1434,12 +3816,16 @@ tl::expected<CopyStartResponse, ErrorCode> MasterService::CopyStart(
     }
 
     // Create replication task for tracking.
-    auto& shard = accessor.GetShard();
-    shard->replication_tasks.emplace(
+    auto& tenant_state = accessor.GetTenantState();
+    auto task_insert = tenant_state.replication_tasks.emplace(
         std::piecewise_construct, std::forward_as_tuple(key),
         std::forward_as_tuple(client_id, std::chrono::system_clock::now(),
                               ReplicationTask::Type::COPY, source->id(),
-                              std::move(replica_ids)));
+                              std::move(replica_ids), reserved_quota_charge));
+    if (!task_insert.second) {
+        abort_reserved_quota();
+        return tl::make_unexpected(ErrorCode::OBJECT_HAS_REPLICATION_TASK);
+    }
 
     // Increase source refcnt to protect it from eviction.
     source->inc_refcnt();
@@ -1451,10 +3837,12 @@ tl::expected<CopyStartResponse, ErrorCode> MasterService::CopyStart(
     return response;
 }
 
-tl::expected<void, ErrorCode> MasterService::CopyEnd(const UUID& client_id,
-                                                     const std::string& key) {
+tl::expected<void, ErrorCode> MasterService::CopyEnd(
+    const UUID& client_id, const std::string& key,
+    const std::string& tenant_id) {
     std::shared_lock<std::shared_mutex> shared_lock(snapshot_mutex_);
-    MetadataAccessorRW accessor(this, key);
+    MetadataAccessorRW accessor(this,
+                                MakeObjectIdentityForRequest(key, tenant_id));
     if (!accessor.Exists()) {
         LOG(ERROR) << "key=" << key << ", error=object_not_found";
         return tl::make_unexpected(ErrorCode::OBJECT_NOT_FOUND);
@@ -1486,11 +3874,20 @@ tl::expected<void, ErrorCode> MasterService::CopyEnd(const UUID& client_id,
         LOG(ERROR) << "key=" << key << ", source_id=" << source_id
                    << ", status=" << (source == nullptr ? "nullptr" : "invalid")
                    << ", copy source becomes invalid during data transfer";
+        // Release the refcnt taken in CopyStart. The success path below does
+        // this once the copy completes; this error path must do it too, or the
+        // source replica stays pinned and can never be evicted.
+        if (source != nullptr) {
+            source->dec_refcnt();
+        }
         // Discard target replicas and clear the replication task.
-        metadata.EraseReplicas([&task](const Replica& replica) {
-            return std::find(task.replica_ids.begin(), task.replica_ids.end(),
-                             replica.id()) != task.replica_ids.end();
-        });
+        EraseReplicasWithCacheTotalAccounting(
+            metadata, [&task](const Replica& replica) {
+                return std::find(task.replica_ids.begin(),
+                                 task.replica_ids.end(),
+                                 replica.id()) != task.replica_ids.end();
+            });
+        AbortReplicationTaskQuota(metadata.tenant_id, task);
         accessor.EraseReplicationTask();
         if (!metadata.IsValid()) {
             // Remove the object if it does not have any replicas.
@@ -1504,6 +3901,7 @@ tl::expected<void, ErrorCode> MasterService::CopyEnd(const UUID& client_id,
 
     // Mark all replica_ids as complete
     bool all_complete = true;
+    uint64_t completed_quota_charge = 0;
     for (const auto& replica_id : task.replica_ids) {
         auto replica = metadata.GetReplicaByID(replica_id);
         if (replica == nullptr || replica->has_invalid_mem_handle()) {
@@ -1513,8 +3911,20 @@ tl::expected<void, ErrorCode> MasterService::CopyEnd(const UUID& client_id,
             all_complete = false;
         } else {
             replica->mark_complete();
+            completed_quota_charge = SaturatingAdd(
+                completed_quota_charge, static_cast<uint64_t>(metadata.size));
         }
     }
+    SyncCacheTotalAccounting(metadata);
+
+    const uint64_t commit_charge =
+        std::min(completed_quota_charge, task.reserved_quota_charge_bytes);
+    const uint64_t abort_charge =
+        task.reserved_quota_charge_bytes - commit_charge;
+    CommitAdditionalTenantQuota(metadata.tenant_id, commit_charge);
+    AbortTenantQuota(metadata.tenant_id, abort_charge);
+    metadata.committed_quota_charge_bytes =
+        SaturatingAdd(metadata.committed_quota_charge_bytes, commit_charge);
 
     accessor.EraseReplicationTask();
 
@@ -1523,9 +3933,11 @@ tl::expected<void, ErrorCode> MasterService::CopyEnd(const UUID& client_id,
 }
 
 tl::expected<void, ErrorCode> MasterService::CopyRevoke(
-    const UUID& client_id, const std::string& key) {
+    const UUID& client_id, const std::string& key,
+    const std::string& tenant_id) {
     std::shared_lock<std::shared_mutex> shared_lock(snapshot_mutex_);
-    MetadataAccessorRW accessor(this, key);
+    MetadataAccessorRW accessor(this,
+                                MakeObjectIdentityForRequest(key, tenant_id));
     if (!accessor.Exists()) {
         LOG(ERROR) << "key=" << key << ", error=object_not_found";
         return tl::make_unexpected(ErrorCode::OBJECT_NOT_FOUND);
@@ -1562,9 +3974,13 @@ tl::expected<void, ErrorCode> MasterService::CopyRevoke(
 
     // Erase all replica_ids
     for (const auto& replica_id : task.replica_ids) {
-        metadata.EraseReplicaByID(replica_id);
+        EraseReplicasWithCacheTotalAccounting(
+            metadata, [&replica_id](const Replica& replica) {
+                return replica.id() == replica_id;
+            });
     }
 
+    AbortReplicationTaskQuota(metadata.tenant_id, task);
     accessor.EraseReplicationTask();
 
     if (!metadata.IsValid()) {
@@ -1576,8 +3992,13 @@ tl::expected<void, ErrorCode> MasterService::CopyRevoke(
 }
 
 tl::expected<MoveStartResponse, ErrorCode> MasterService::MoveStart(
-    const UUID& client_id, const std::string& key,
+    const UUID& client_id, const std::string& key, const std::string& tenant_id,
     const std::string& src_segment, const std::string& tgt_segment) {
+    auto normalized_tenant_result = NormalizeTenantIdForWrite(tenant_id);
+    if (!normalized_tenant_result) {
+        return tl::make_unexpected(normalized_tenant_result.error());
+    }
+    const ObjectIdentity object_id{normalized_tenant_result.value(), key};
     std::shared_lock<std::shared_mutex> shared_lock(snapshot_mutex_);
     if (src_segment == tgt_segment) {
         LOG(ERROR) << "key=" << key << ", move_tgt=" << tgt_segment
@@ -1600,7 +4021,7 @@ tl::expected<MoveStartResponse, ErrorCode> MasterService::MoveStart(
         }
     }
 
-    MetadataAccessorRW accessor(this, key);
+    MetadataAccessorRW accessor(this, object_id);
     if (!accessor.Exists()) {
         LOG(ERROR) << "key=" << key << ", object not found";
         return tl::make_unexpected(ErrorCode::OBJECT_NOT_FOUND);
@@ -1623,6 +4044,21 @@ tl::expected<MoveStartResponse, ErrorCode> MasterService::MoveStart(
 
     std::vector<Replica> replicas;
     if (metadata.GetReplicaBySegmentName(tgt_segment) == nullptr) {
+        const uint64_t reserved_quota_charge =
+            SaturatingMultiply(static_cast<uint64_t>(metadata.size), 1);
+        auto quota_result =
+            ReserveTenantQuota(object_id.tenant_id, reserved_quota_charge);
+        if (!quota_result) {
+            if (quota_result.error() == ErrorCode::TENANT_QUOTA_EXCEEDED) {
+                MasterMetricManager::instance().inc_tenant_quota_reject(
+                    object_id.tenant_id, "quota_exceeded");
+            }
+            return tl::make_unexpected(quota_result.error());
+        }
+        auto abort_reserved_quota = [&] {
+            AbortTenantQuota(object_id.tenant_id, reserved_quota_charge);
+        };
+
         ScopedAllocatorAccess allocator_access =
             segment_manager_.getAllocatorAccess();
         const auto& allocator_manager = allocator_access.getAllocatorManager();
@@ -1632,11 +4068,22 @@ tl::expected<MoveStartResponse, ErrorCode> MasterService::MoveStart(
         if (!replica.has_value()) {
             LOG(ERROR) << "key=" << key << ", tgt_segment=" << tgt_segment
                        << ", failed to allocate replica";
+            abort_reserved_quota();
             return tl::make_unexpected(replica.error());
         }
         replicas.push_back(std::move(*replica));
+    } else {
+        auto quota_result = ReserveTenantQuota(object_id.tenant_id, 0);
+        if (!quota_result) {
+            return tl::make_unexpected(quota_result.error());
+        }
     }
 
+    const uint64_t reserved_quota_charge =
+        replicas.empty()
+            ? 0
+            : SaturatingMultiply(static_cast<uint64_t>(metadata.size), 1);
+
     MoveStartResponse response;
     std::vector<ReplicaID> replica_ids;
 
@@ -1649,12 +4096,16 @@ tl::expected<MoveStartResponse, ErrorCode> MasterService::MoveStart(
     }
 
     // Create replication task for tracking.
-    auto& shard = accessor.GetShard();
-    shard->replication_tasks.emplace(
+    auto& tenant_state = accessor.GetTenantState();
+    auto task_insert = tenant_state.replication_tasks.emplace(
         std::piecewise_construct, std::forward_as_tuple(key),
         std::forward_as_tuple(client_id, std::chrono::system_clock::now(),
                               ReplicationTask::Type::MOVE, source->id(),
-                              std::move(replica_ids)));
+                              std::move(replica_ids), reserved_quota_charge));
+    if (!task_insert.second) {
+        AbortTenantQuota(object_id.tenant_id, reserved_quota_charge);
+        return tl::make_unexpected(ErrorCode::OBJECT_HAS_REPLICATION_TASK);
+    }
 
     // Increase source refcnt to protect it from eviction.
     source->inc_refcnt();
@@ -1666,10 +4117,12 @@ tl::expected<MoveStartResponse, ErrorCode> MasterService::MoveStart(
     return response;
 }
 
-tl::expected<void, ErrorCode> MasterService::MoveEnd(const UUID& client_id,
-                                                     const std::string& key) {
+tl::expected<void, ErrorCode> MasterService::MoveEnd(
+    const UUID& client_id, const std::string& key,
+    const std::string& tenant_id) {
     std::shared_lock<std::shared_mutex> shared_lock(snapshot_mutex_);
-    MetadataAccessorRW accessor(this, key);
+    MetadataAccessorRW accessor(this,
+                                MakeObjectIdentityForRequest(key, tenant_id));
     if (!accessor.Exists()) {
         LOG(ERROR) << "key=" << key << ", error=object_not_found";
         return tl::make_unexpected(ErrorCode::OBJECT_NOT_FOUND);
@@ -1701,11 +4154,20 @@ tl::expected<void, ErrorCode> MasterService::MoveEnd(const UUID& client_id,
         LOG(ERROR) << "key=" << key << ", source_id=" << source_id
                    << ", status=" << (source == nullptr ? "nullptr" : "invalid")
                    << ", move source becomes invalid during data transfer";
+        // Release the refcnt taken in MoveStart. The success path below does
+        // this once the move completes; this error path must do it too, or the
+        // source replica stays pinned and can never be evicted.
+        if (source != nullptr) {
+            source->dec_refcnt();
+        }
         // Discard target replica and clear the replication task.
-        metadata.EraseReplicas([&task](const Replica& replica) {
-            return std::find(task.replica_ids.begin(), task.replica_ids.end(),
-                             replica.id()) != task.replica_ids.end();
-        });
+        EraseReplicasWithCacheTotalAccounting(
+            metadata, [&task](const Replica& replica) {
+                return std::find(task.replica_ids.begin(),
+                                 task.replica_ids.end(),
+                                 replica.id()) != task.replica_ids.end();
+            });
+        AbortReplicationTaskQuota(metadata.tenant_id, task);
         accessor.EraseReplicationTask();
         if (!metadata.IsValid()) {
             // Remove the object if it does not have any replicas.
@@ -1727,17 +4189,19 @@ tl::expected<void, ErrorCode> MasterService::MoveEnd(const UUID& client_id,
             LOG(WARNING)
                 << "key=" << key << ", replica_id=" << replica_id
                 << ", move target becomes invalid during data transfer";
+            AbortReplicationTaskQuota(metadata.tenant_id, task);
             accessor.EraseReplicationTask();
             return tl::make_unexpected(ErrorCode::REPLICA_IS_GONE);
         }
 
         // Mark replica as complete
         replica->mark_complete();
+        SyncCacheTotalAccounting(metadata);
     }
 
     // Remove the source replica and release its space later.
-    auto source_replica =
-        metadata.PopReplicas([&source_id](const Replica& replica) {
+    auto source_replica = PopReplicasWithCacheTotalAccounting(
+        metadata, [&source_id](const Replica& replica) {
             return replica.id() == source_id;
         });
     if (!source_replica.empty()) {
@@ -1747,15 +4211,18 @@ tl::expected<void, ErrorCode> MasterService::MoveEnd(const UUID& client_id,
             std::chrono::system_clock::now() + put_start_release_timeout_sec_);
     }
 
+    AbortReplicationTaskQuota(metadata.tenant_id, task);
     accessor.EraseReplicationTask();
 
     return {};
 }
 
 tl::expected<void, ErrorCode> MasterService::MoveRevoke(
-    const UUID& client_id, const std::string& key) {
+    const UUID& client_id, const std::string& key,
+    const std::string& tenant_id) {
     std::shared_lock<std::shared_mutex> shared_lock(snapshot_mutex_);
-    MetadataAccessorRW accessor(this, key);
+    MetadataAccessorRW accessor(this,
+                                MakeObjectIdentityForRequest(key, tenant_id));
     if (!accessor.Exists()) {
         LOG(ERROR) << "key=" << key << ", error=object_not_found";
         return tl::make_unexpected(ErrorCode::OBJECT_NOT_FOUND);
@@ -1792,9 +4259,13 @@ tl::expected<void, ErrorCode> MasterService::MoveRevoke(
 
     // Erase all replica_ids (in MOVE operation, there should be at most one)
     for (const auto& replica_id : task.replica_ids) {
-        metadata.EraseReplicaByID(replica_id);
+        EraseReplicasWithCacheTotalAccounting(
+            metadata, [&replica_id](const Replica& replica) {
+                return replica.id() == replica_id;
+            });
     }
 
+    AbortReplicationTaskQuota(metadata.tenant_id, task);
     accessor.EraseReplicationTask();
 
     if (!metadata.IsValid()) {
@@ -1805,10 +4276,11 @@ tl::expected<void, ErrorCode> MasterService::MoveRevoke(
     return {};
 }
 
-auto MasterService::Remove(const std::string& key, bool force)
-    -> tl::expected<void, ErrorCode> {
+auto MasterService::Remove(const std::string& key, const std::string& tenant_id,
+                           bool force) -> tl::expected<void, ErrorCode> {
     std::shared_lock<std::shared_mutex> shared_lock(snapshot_mutex_);
-    MetadataAccessorRW accessor(this, key);
+    const auto object_id = MakeObjectIdentityForRequest(key, tenant_id);
+    MetadataAccessorRW accessor(this, object_id);
     if (!accessor.Exists()) {
         VLOG(1) << "key=" << key << ", error=object_not_found";
         return tl::make_unexpected(ErrorCode::OBJECT_NOT_FOUND);
@@ -1837,12 +4309,14 @@ auto MasterService::Remove(const std::string& key, bool force)
         return tl::make_unexpected(ErrorCode::OBJECT_HAS_REPLICATION_TASK);
     }
 
-    // Remove object metadata
+    auto& tenant_state = accessor.GetTenantState();
+    ErasePromotionTaskIfPresent(tenant_state, key, object_id.tenant_id);
     accessor.Erase();
     return {};
 }
 
-auto MasterService::RemoveByRegex(const std::string& regex_pattern, bool force)
+auto MasterService::RemoveByRegex(const std::string& regex_pattern,
+                                  const std::string& tenant_id, bool force)
     -> tl::expected<long, ErrorCode> {
     long removed_count = 0;
     std::regex pattern;
@@ -1856,10 +4330,17 @@ auto MasterService::RemoveByRegex(const std::string& regex_pattern, bool force)
     }
 
     std::shared_lock<std::shared_mutex> shared_lock(snapshot_mutex_);
+    const auto normalized_tenant = NormalizeRequestTenantId(tenant_id);
     for (size_t i = 0; i < kNumShards; ++i) {
         MetadataShardAccessorRW shard(this, i);
+        auto tenant_it = shard->tenants.find(normalized_tenant);
+        if (tenant_it == shard->tenants.end()) {
+            continue;
+        }
+        auto& tenant_state = tenant_it->second;
 
-        for (auto it = shard->metadata.begin(); it != shard->metadata.end();) {
+        for (auto it = tenant_state.metadata.begin();
+             it != tenant_state.metadata.end();) {
             if (std::regex_search(it->first, pattern)) {
                 if (!force && !it->second.IsLeaseExpired()) {
                     VLOG(1) << "key=" << it->first
@@ -1882,7 +4363,7 @@ auto MasterService::RemoveByRegex(const std::string& regex_pattern, bool force)
                     ++it;
                     continue;
                 }
-                if (metadata_shards_[i].replication_tasks.contains(it->first)) {
+                if (tenant_state.replication_tasks.contains(it->first)) {
                     LOG(WARNING) << "key=" << it->first
                                  << ", matched by regex, but has replication "
                                     "task. Skipping removal.";
@@ -1892,12 +4373,18 @@ auto MasterService::RemoveByRegex(const std::string& regex_pattern, bool force)
 
                 VLOG(1) << "key=" << it->first
                         << " matched by regex. Removing.";
-                it = shard->metadata.erase(it);
+                ErasePromotionTaskIfPresent(tenant_state, it->first,
+                                            normalized_tenant);
+                it = EraseMetadata(tenant_state, it, normalized_tenant,
+                                   QuotaEraseMode::kFull, &shard);
                 removed_count++;
             } else {
                 ++it;
             }
         }
+        if (tenant_state.Empty()) {
+            shard->tenants.erase(tenant_it);
+        }
     }
 
     VLOG(1) << "action=remove_by_regex, pattern=" << regex_pattern
@@ -1908,36 +4395,35 @@ auto MasterService::RemoveByRegex(const std::string& regex_pattern, bool force)
 long MasterService::RemoveAll(bool force) {
     long removed_count = 0;
     uint64_t total_freed_size = 0;
-    // Store the current time to avoid repeatedly
-    // calling std::chrono::steady_clock::now()
     std::shared_lock<std::shared_mutex> shared_lock(snapshot_mutex_);
     auto now = std::chrono::system_clock::now();
 
     for (size_t i = 0; i < kNumShards; i++) {
         MetadataShardAccessorRW shard(this, i);
-        if (shard->metadata.empty()) {
-            continue;
-        }
-
-        // Only remove completed objects with expired leases (unless force=true)
-        auto it = shard->metadata.begin();
-        while (it != shard->metadata.end()) {
-            /**
-             * The reason the force operation here does not bypass the replica
-             * check is that put operations (which could also be copy or move)
-             * and remove operations might be happening concurrently, making it
-             * extremely dangerous to perform a direct removal at this point.
-             */
-            if ((force || it->second.IsLeaseExpired(now)) &&
-                it->second.AllReplicas(&Replica::fn_is_completed) &&
-                !shard->replication_tasks.contains(it->first)) {
-                auto mem_rep_count =
-                    it->second.CountReplicas(&Replica::fn_is_memory_replica);
-                total_freed_size += it->second.size * mem_rep_count;
-                it = shard->metadata.erase(it);
-                removed_count++;
+        for (auto tenant_it = shard->tenants.begin();
+             tenant_it != shard->tenants.end();) {
+            auto& tenant_state = tenant_it->second;
+            auto it = tenant_state.metadata.begin();
+            while (it != tenant_state.metadata.end()) {
+                if ((force || it->second.IsLeaseExpired(now)) &&
+                    it->second.AllReplicas(&Replica::fn_is_completed) &&
+                    !tenant_state.replication_tasks.contains(it->first)) {
+                    auto mem_rep_count = it->second.CountReplicas(
+                        &Replica::fn_is_memory_replica);
+                    total_freed_size += it->second.size * mem_rep_count;
+                    ErasePromotionTaskIfPresent(tenant_state, it->first,
+                                                tenant_it->first);
+                    it = EraseMetadata(tenant_state, it, tenant_it->first,
+                                       QuotaEraseMode::kFull, &shard);
+                    removed_count++;
+                } else {
+                    ++it;
+                }
+            }
+            if (tenant_state.Empty()) {
+                tenant_it = shard->tenants.erase(tenant_it);
             } else {
-                ++it;
+                ++tenant_it;
             }
         }
     }
@@ -1948,10 +4434,56 @@ long MasterService::RemoveAll(bool force) {
     return removed_count;
 }
 
+long MasterService::RemoveAll(const std::string& tenant_id, bool force) {
+    long removed_count = 0;
+    uint64_t total_freed_size = 0;
+    // Store the current time to avoid repeatedly
+    // calling std::chrono::steady_clock::now()
+    std::shared_lock<std::shared_mutex> shared_lock(snapshot_mutex_);
+    auto now = std::chrono::system_clock::now();
+    const auto normalized_tenant = NormalizeRequestTenantId(tenant_id);
+
+    for (size_t i = 0; i < kNumShards; i++) {
+        MetadataShardAccessorRW shard(this, i);
+        auto tenant_it = shard->tenants.find(normalized_tenant);
+        if (tenant_it == shard->tenants.end()) {
+            continue;
+        }
+        auto& tenant_state = tenant_it->second;
+        auto it = tenant_state.metadata.begin();
+        while (it != tenant_state.metadata.end()) {
+            if ((force || it->second.IsLeaseExpired(now)) &&
+                it->second.AllReplicas(&Replica::fn_is_completed) &&
+                !tenant_state.replication_tasks.contains(it->first)) {
+                auto mem_rep_count =
+                    it->second.CountReplicas(&Replica::fn_is_memory_replica);
+                total_freed_size += it->second.size * mem_rep_count;
+                ErasePromotionTaskIfPresent(tenant_state, it->first,
+                                            normalized_tenant);
+                it = EraseMetadata(tenant_state, it, normalized_tenant,
+                                   QuotaEraseMode::kFull, &shard);
+                removed_count++;
+            } else {
+                ++it;
+            }
+        }
+        if (tenant_state.Empty()) {
+            shard->tenants.erase(tenant_it);
+        }
+    }
+
+    VLOG(1) << "action=remove_all_objects"
+            << ", tenant_id=" << normalized_tenant
+            << ", removed_count=" << removed_count
+            << ", total_freed_size=" << total_freed_size;
+    return removed_count;
+}
+
 auto MasterService::BatchRemove(const std::vector<std::string>& keys,
-                                bool force)
+                                const std::string& tenant_id, bool force)
     -> std::vector<tl::expected<void, ErrorCode>> {
     std::vector<tl::expected<void, ErrorCode>> results(keys.size());
+    const auto normalized_tenant = NormalizeRequestTenantId(tenant_id);
 
     // Group keys by shard to reduce lock contention
     std::unordered_map<size_t,
@@ -1961,7 +4493,7 @@ auto MasterService::BatchRemove(const std::vector<std::string>& keys,
         std::min(keys.size(), static_cast<size_t>(kNumShards)));
 
     for (size_t i = 0; i < keys.size(); ++i) {
-        size_t shard_idx = getShardIndex(keys[i]);
+        size_t shard_idx = getMetadataShardIndex(normalized_tenant, keys[i]);
         keys_by_shard[shard_idx].emplace_back(i, &keys[i]);
     }
 
@@ -1976,9 +4508,17 @@ auto MasterService::BatchRemove(const std::vector<std::string>& keys,
 
         for (const auto& [original_idx, key_ptr] : key_group) {
             const std::string& key = *key_ptr;
-            auto it = shard->metadata.find(key);
+            auto tenant_it = shard->tenants.find(normalized_tenant);
+            if (tenant_it == shard->tenants.end()) {
+                VLOG(1) << "key=" << key << ", error=object_not_found";
+                results[original_idx] =
+                    tl::make_unexpected(ErrorCode::OBJECT_NOT_FOUND);
+                continue;
+            }
+            auto& tenant_state = tenant_it->second;
+            auto it = tenant_state.metadata.find(key);
 
-            if (it == shard->metadata.end()) {
+            if (it == tenant_state.metadata.end()) {
                 VLOG(1) << "key=" << key << ", error=object_not_found";
                 results[original_idx] =
                     tl::make_unexpected(ErrorCode::OBJECT_NOT_FOUND);
@@ -1986,11 +4526,22 @@ auto MasterService::BatchRemove(const std::vector<std::string>& keys,
             }
 
             // Clean up stale replica handles (consistent with single Remove)
-            if (CleanupStaleHandles(it->second, alive_clients)) {
-                shard->processing_keys.erase(key);
-                shard->replication_tasks.erase(key);
-                shard->offloading_tasks.erase(key);
-                shard->metadata.erase(it);
+            if (CleanupStaleHandles(it->second, alive_clients, &shard)) {
+                tenant_state.processing_keys.erase(key);
+                auto task_it = tenant_state.replication_tasks.find(key);
+                if (task_it != tenant_state.replication_tasks.end()) {
+                    AbortReplicationTaskQuota(normalized_tenant,
+                                              task_it->second);
+                    tenant_state.replication_tasks.erase(task_it);
+                }
+                tenant_state.offloading_tasks.erase(key);
+                ErasePromotionTaskIfPresent(tenant_state, key,
+                                            normalized_tenant);
+                EraseMetadata(tenant_state, it, normalized_tenant,
+                              QuotaEraseMode::kFull, &shard);
+                if (tenant_state.Empty()) {
+                    shard->tenants.erase(tenant_it);
+                }
                 results[original_idx] =
                     tl::make_unexpected(ErrorCode::OBJECT_NOT_FOUND);
                 continue;
@@ -2017,7 +4568,7 @@ auto MasterService::BatchRemove(const std::vector<std::string>& keys,
                 continue;
             }
 
-            if (shard->replication_tasks.contains(key)) {
+            if (tenant_state.replication_tasks.contains(key)) {
                 LOG(ERROR) << "key=" << key
                            << ", error=object_has_replication_task";
                 results[original_idx] =
@@ -2026,7 +4577,12 @@ auto MasterService::BatchRemove(const std::vector<std::string>& keys,
             }
 
             // Remove object metadata
-            shard->metadata.erase(it);
+            ErasePromotionTaskIfPresent(tenant_state, key, normalized_tenant);
+            EraseMetadata(tenant_state, it, normalized_tenant,
+                          QuotaEraseMode::kFull, &shard);
+            if (tenant_state.Empty()) {
+                shard->tenants.erase(tenant_it);
+            }
             results[original_idx] = {};  // Success
         }
     }
@@ -2036,13 +4592,31 @@ auto MasterService::BatchRemove(const std::vector<std::string>& keys,
 
 bool MasterService::CleanupStaleHandles(
     ObjectMetadata& metadata,
-    const std::unordered_set<UUID, boost::hash<UUID>>& alive_clients) {
+    const std::unordered_set<UUID, boost::hash<UUID>>& alive_clients,
+    MetadataShardAccessorRW* shard) {
+    bool had_completed_disk = metadata.HasReplica([](const Replica& r) {
+        return r.is_local_disk_replica() && r.is_completed();
+    });
     // Remove those with invalid allocators (memory replicas on unmounted
     // segments) and local_disk replicas whose owner client is no longer alive.
-    metadata.EraseReplicas([&alive_clients](const Replica& replica) {
-        return replica.has_invalid_mem_handle() ||
-               replica.has_stale_local_disk_client(alive_clients);
-    });
+    const uint64_t before_charge = CompletedMemoryQuotaCharge(metadata);
+    EraseReplicasWithCacheTotalAccounting(
+        metadata, [&alive_clients](const Replica& replica) {
+            return (replica.has_invalid_mem_handle() ||
+                    replica.has_invalid_nof_handle() ||
+                    replica.has_stale_local_disk_client(alive_clients)) &&
+                   replica.is_completed();
+        });
+    const uint64_t after_charge = CompletedMemoryQuotaCharge(metadata);
+    if (before_charge > after_charge) {
+        ReleaseCommittedQuotaCharge(metadata, before_charge - after_charge);
+    }
+    if (had_completed_disk && shard &&
+        !metadata.HasReplica([](const Replica& r) {
+            return r.is_local_disk_replica() && r.is_completed();
+        })) {
+        shard->OnDiskReplicaRemoved(had_completed_disk, metadata);
+    }
 
     // Return true if no valid replicas remain after cleanup
     return !metadata.IsValid();
@@ -2052,7 +4626,9 @@ size_t MasterService::GetKeyCount() const {
     size_t total = 0;
     for (size_t i = 0; i < kNumShards; i++) {
         MetadataShardAccessorRO shard(this, i);
-        total += shard->metadata.size();
+        for (const auto& [tenant_id, tenant_state] : shard->tenants) {
+            total += tenant_state.metadata.size();
+        }
     }
     return total;
 }
@@ -2118,12 +4694,91 @@ auto MasterService::MountLocalDiskSegment(const UUID& client_id,
     } else if (err != ErrorCode::OK) {
         return tl::make_unexpected(err);
     }
+
+    // Notify the client monitor thread to start tracking this client's TTL.
+    // Without this, a client that only mounts a LOCAL_DISK segment (and
+    // doesn't ping) would be considered expired by ClientMonitorFunc, which
+    // would then clear all its LOCAL_DISK replicas.
+    PodUUID pod_client_id;
+    pod_client_id.first = client_id.first;
+    pod_client_id.second = client_id.second;
+    if (!client_ping_queue_.push(pod_client_id)) {
+        LOG(ERROR) << "client_id=" << client_id
+                   << ", error=client_ping_queue_full";
+        return tl::make_unexpected(ErrorCode::INTERNAL_ERROR);
+    }
+
     return {};
 }
 
 auto MasterService::OffloadObjectHeartbeat(const UUID& client_id,
                                            bool enable_offloading)
-    -> tl::expected<std::unordered_map<std::string, int64_t>, ErrorCode> {
+    -> tl::expected<std::vector<OffloadTaskItem>, ErrorCode> {
+    std::shared_lock<std::shared_mutex> shared_lock(snapshot_mutex_);
+    ScopedLocalDiskSegmentAccess local_disk_segment_access =
+        segment_manager_.getLocalDiskSegmentAccess();
+    auto& client_local_disk_segment =
+        local_disk_segment_access.getClientLocalDiskSegment();
+    auto local_disk_segment_it = client_local_disk_segment.find(client_id);
+    if (local_disk_segment_it == client_local_disk_segment.end()) {
+        LOG(ERROR) << "Local disk segment not found with client id = "
+                   << client_id;
+        return tl::make_unexpected(ErrorCode::SEGMENT_NOT_FOUND);
+    }
+    std::unordered_map<std::string, OffloadTaskItem> offloading_objects_copy;
+    {
+        MutexLocker locker(&local_disk_segment_it->second->offloading_mutex_);
+        local_disk_segment_it->second->enable_offloading = enable_offloading;
+        if (enable_offloading) {
+            std::vector<OffloadTaskItem> result;
+            result.reserve(
+                local_disk_segment_it->second->offloading_objects.size());
+            for (const auto& [_, task] :
+                 local_disk_segment_it->second->offloading_objects) {
+                result.push_back(task);
+            }
+            local_disk_segment_it->second->offloading_objects.clear();
+            return result;
+        }
+        // Offloading is disabled: clear the pending queue to prevent
+        // unbounded growth that would trigger KEYS_ULTRA_LIMIT in
+        // PushOffloadingQueue. We must also clean up corresponding
+        // offloading_tasks and decrement source replica refcounts to avoid
+        // resource leaks and blocked writes (OBJECT_HAS_REPLICATION_TASK).
+        // Copy keys out before releasing the mutex to avoid lock order
+        // violation: the lock order is Shard Lock -> offloading_mutex_, so we
+        // must release offloading_mutex_ before taking shard locks via
+        // MetadataAccessorRW.
+        offloading_objects_copy =
+            std::move(local_disk_segment_it->second->offloading_objects);
+    }
+
+    for (auto& [_, task] : offloading_objects_copy) {
+        const auto object_id = MakeObjectIdentity(task.key, task.tenant_id);
+        MetadataAccessorRW accessor(this, object_id);
+        if (accessor.Exists()) {
+            auto& tenant_state = accessor.GetTenantState();
+            auto task_it =
+                tenant_state.offloading_tasks.find(object_id.user_key);
+            if (task_it != tenant_state.offloading_tasks.end()) {
+                auto source =
+                    accessor.Get().GetReplicaByID(task_it->second.source_id);
+                if (source) {
+                    source->dec_refcnt();
+                }
+                tenant_state.offloading_tasks.erase(task_it);
+            }
+        }
+    }
+    return {};
+}
+
+auto MasterService::ReportSsdCapacity(const UUID& client_id,
+                                      int64_t ssd_total_capacity_bytes)
+    -> tl::expected<void, ErrorCode> {
+    if (ssd_total_capacity_bytes < 0) {
+        return tl::make_unexpected(ErrorCode::INVALID_PARAMS);
+    }
     std::shared_lock<std::shared_mutex> shared_lock(snapshot_mutex_);
     ScopedLocalDiskSegmentAccess local_disk_segment_access =
         segment_manager_.getLocalDiskSegmentAccess();
@@ -2136,54 +4791,148 @@ auto MasterService::OffloadObjectHeartbeat(const UUID& client_id,
         return tl::make_unexpected(ErrorCode::SEGMENT_NOT_FOUND);
     }
     MutexLocker locker(&local_disk_segment_it->second->offloading_mutex_);
-    local_disk_segment_it->second->enable_offloading = enable_offloading;
-    if (enable_offloading) {
-        return std::move(local_disk_segment_it->second->offloading_objects);
+    int64_t old_capacity =
+        local_disk_segment_it->second->ssd_total_capacity_bytes;
+    if (ssd_total_capacity_bytes != old_capacity) {
+        local_disk_segment_it->second->ssd_total_capacity_bytes =
+            ssd_total_capacity_bytes;
+        if (old_capacity > 0) {
+            MasterMetricManager::instance().dec_total_file_capacity(
+                old_capacity);
+        }
+        if (ssd_total_capacity_bytes > 0) {
+            MasterMetricManager::instance().inc_total_file_capacity(
+                ssd_total_capacity_bytes);
+        }
     }
     return {};
 }
 
 auto MasterService::NotifyOffloadSuccess(
-    const UUID& client_id, const std::vector<std::string>& keys,
+    const UUID& client_id, const std::vector<OffloadTaskItem>& tasks,
     const std::vector<StorageObjectMetadata>& metadatas)
     -> tl::expected<void, ErrorCode> {
-    for (size_t i = 0; i < keys.size(); ++i) {
-        const auto& key = keys[i];
-        const auto& metadata = metadatas[i];
+    if (tasks.size() != metadatas.size()) {
+        return tl::make_unexpected(ErrorCode::INVALID_PARAMS);
+    }
+    std::shared_ptr<LocalDiskSegment> local_disk_segment;
+    {
+        ScopedLocalDiskSegmentAccess ssd_access =
+            segment_manager_.getLocalDiskSegmentAccess();
+        auto& client_segments = ssd_access.getClientLocalDiskSegment();
+        auto disk_it = client_segments.find(client_id);
+        if (disk_it != client_segments.end()) {
+            local_disk_segment = disk_it->second;
+        }
+    }
 
-        // Release refcnt and clear offloading task.
+    for (size_t i = 0; i < tasks.size(); ++i) {
+        const auto& task = tasks[i];
+        const auto& metadata = metadatas[i];
+        const auto request_object_id =
+            MakeObjectIdentityForRequest(task.key, task.tenant_id);
+
+        Replica replica(client_id, metadata.data_size,
+                        metadata.transport_endpoint, ReplicaStatus::COMPLETE);
+        bool handled_existing_object = false;
+        bool added_new_local_disk_replica = false;
         {
-            MetadataAccessorRW accessor(this, key);
+            std::shared_lock<std::shared_mutex> shared_lock(snapshot_mutex_);
+            MetadataAccessorRW accessor(this, request_object_id);
             if (accessor.Exists()) {
                 auto& obj_metadata = accessor.Get();
-                auto& shard = accessor.GetShard();
-                auto task_it = shard->offloading_tasks.find(key);
-                if (task_it != shard->offloading_tasks.end()) {
+                auto& tenant_state = accessor.GetTenantState();
+                auto task_it = tenant_state.offloading_tasks.find(
+                    request_object_id.user_key);
+                if (task_it != tenant_state.offloading_tasks.end() &&
+                    replica.type() != ReplicaType::LOCAL_DISK) {
+                    LOG(ERROR) << "Invalid replica type: " << replica.type()
+                               << ". Expected ReplicaType::LOCAL_DISK.";
+                    return tl::make_unexpected(ErrorCode::INVALID_PARAMS);
+                }
+
+                // Existing orphan objects can only bypass tenant registration
+                // for a master-admitted offload completion. Without this task
+                // marker, fall through to the regular registration check.
+                if (task_it != tenant_state.offloading_tasks.end()) {
                     auto source =
                         obj_metadata.GetReplicaByID(task_it->second.source_id);
                     if (source != nullptr) {
                         source->dec_refcnt();
                     }
-                    shard->offloading_tasks.erase(task_it);
+                    tenant_state.offloading_tasks.erase(task_it);
+
+                    if (!obj_metadata.HasReplica(
+                            &Replica::fn_is_local_disk_replica)) {
+                        std::vector<Replica> replicas;
+                        replicas.emplace_back(std::move(replica));
+                        obj_metadata.AddReplicas(std::move(replicas));
+                        auto& shard = accessor.GetShard();
+                        shard.OnDiskReplicaAdded(obj_metadata);
+                        added_new_local_disk_replica = true;
+                    } else {
+                        obj_metadata.VisitReplicas(
+                            [client_id](const Replica& rep) {
+                                return rep.type() == ReplicaType::LOCAL_DISK &&
+                                       rep.get_descriptor()
+                                               .get_local_disk_descriptor()
+                                               .client_id == client_id;
+                            },
+                            [&replica](Replica& rep) {
+                                rep.get_descriptor()
+                                    .get_local_disk_descriptor()
+                                    .transport_endpoint =
+                                    replica.get_descriptor()
+                                        .get_local_disk_descriptor()
+                                        .transport_endpoint;
+                                rep.get_descriptor()
+                                    .get_local_disk_descriptor()
+                                    .object_size =
+                                    replica.get_descriptor()
+                                        .get_local_disk_descriptor()
+                                        .object_size;
+                            });
+                    }
+                    handled_existing_object = true;
                 }
             }
         }
 
-        // Add LOCAL_DISK replica.
-        Replica replica(client_id, metadata.data_size,
-                        metadata.transport_endpoint, ReplicaStatus::COMPLETE);
-        auto res = AddReplica(client_id, key, replica);
-        if (!res && res.error() != ErrorCode::OBJECT_NOT_FOUND) {
-            LOG(ERROR) << "Failed to add replica: error=" << res.error()
-                       << ", client_id=" << client_id << ", key=" << key;
-            return tl::make_unexpected(res.error());
+        if (!handled_existing_object) {
+            auto normalized_tenant_result =
+                NormalizeTenantIdForWrite(task.tenant_id);
+            if (!normalized_tenant_result) {
+                return tl::make_unexpected(normalized_tenant_result.error());
+            }
+            const ObjectIdentity object_id{normalized_tenant_result.value(),
+                                           task.key};
+
+            auto res = AddReplica(client_id, object_id.user_key,
+                                  object_id.tenant_id, replica);
+            if (!res) {
+                if (res.error() == ErrorCode::OBJECT_NOT_FOUND) {
+                    continue;
+                }
+                LOG(ERROR) << "Failed to add replica: error=" << res.error()
+                           << ", client_id=" << client_id
+                           << ", tenant_id=" << object_id.tenant_id
+                           << ", key=" << object_id.user_key;
+                return tl::make_unexpected(res.error());
+            }
+            added_new_local_disk_replica = res.value();
+        }
+        if (local_disk_segment && metadata.data_size > 0 &&
+            added_new_local_disk_replica) {
+            local_disk_segment->ssd_used_bytes.fetch_add(
+                metadata.data_size, std::memory_order_relaxed);
         }
     }
+
     return {};
 }
 
 tl::expected<void, ErrorCode> MasterService::PushOffloadingQueue(
-    const std::string& key, Replica& replica) {
+    const ObjectIdentity& object_id, Replica& replica) {
     const auto& segment_names = replica.get_segment_names();
     if (segment_names.empty()) {
         return {};
@@ -2216,10 +4965,14 @@ tl::expected<void, ErrorCode> MasterService::PushOffloadingQueue(
             offloading_queue_limit_) {
             return tl::make_unexpected(ErrorCode::KEYS_ULTRA_LIMIT);
         }
+        const int64_t size = replica.get_descriptor()
+                                 .get_memory_descriptor()
+                                 .buffer_descriptor.size_;
         auto res = local_disk_segment_it->second->offloading_objects.emplace(
-            key, replica.get_descriptor()
-                     .get_memory_descriptor()
-                     .buffer_descriptor.size_);
+            MakeTenantScopedStorageKey(object_id.tenant_id, object_id.user_key),
+            OffloadTaskItem{.tenant_id = object_id.tenant_id,
+                            .key = object_id.user_key,
+                            .size = size});
         if (!res.second) {
             return tl::make_unexpected(ErrorCode::OBJECT_ALREADY_EXISTS);
         }
@@ -2227,6 +4980,456 @@ tl::expected<void, ErrorCode> MasterService::PushOffloadingQueue(
     return {};
 }
 
+// Promotion-on-hit
+
+// Push a key onto the holder client's promotion_objects map. Resolves the
+// holder via the LOCAL_DISK replica's embedded client_id rather than via
+// the segment-name reverse lookup.
+tl::expected<void, ErrorCode> MasterService::PushPromotionQueue(
+    const ObjectIdentity& object_id, Replica& source_replica) {
+    auto holder_id = source_replica.get_local_disk_client_id();
+    if (!holder_id.has_value()) {
+        return tl::make_unexpected(ErrorCode::INVALID_PARAMS);
+    }
+    ScopedLocalDiskSegmentAccess local_disk_segment_access =
+        segment_manager_.getLocalDiskSegmentAccess();
+    auto& client_local_disk_segment =
+        local_disk_segment_access.getClientLocalDiskSegment();
+    auto local_disk_segment_it =
+        client_local_disk_segment.find(holder_id.value());
+    if (local_disk_segment_it == client_local_disk_segment.end()) {
+        // Holder client expired or never had a LocalDiskSegment registered;
+        // the LOCAL_DISK replica will be cleaned up by ClientMonitorFunc on
+        // its own schedule.
+        return tl::make_unexpected(ErrorCode::SEGMENT_NOT_FOUND);
+    }
+    MutexLocker locker(&local_disk_segment_it->second->offloading_mutex_);
+    auto res = local_disk_segment_it->second->promotion_objects.emplace(
+        MakeTenantScopedStorageKey(object_id.tenant_id, object_id.user_key),
+        PromotionTaskItem{
+            .tenant_id = object_id.tenant_id,
+            .key = object_id.user_key,
+            .size = static_cast<int64_t>(source_replica.get_descriptor()
+                                             .get_local_disk_descriptor()
+                                             .object_size)});
+    if (!res.second) {
+        return tl::make_unexpected(ErrorCode::OBJECT_ALREADY_EXISTS);
+    }
+    return {};
+}
+
+void MasterService::TryPushPromotionQueue(const ObjectIdentity& object_id) {
+    if (!promotion_on_hit_ || !promotion_sketch_) {
+        return;
+    }
+    const auto& key = object_id.user_key;
+    const auto admission_key =
+        MakeTenantScopedStorageKey(object_id.tenant_id, key);
+
+    // Frequency gate: bump and compare against the threshold. The sketch
+    // returns uint8_t (saturating at 255); promotion_admission_threshold_
+    // is clamped into [1, 255] at config parse time (see master.cpp), so
+    // direct comparison is well-defined and threshold=0 (which would
+    // bypass the gate entirely since freq is uint8_t) cannot reach here.
+    const uint8_t freq = promotion_sketch_->increment(admission_key);
+    if (freq < promotion_admission_threshold_) {
+        MasterMetricManager::instance().inc_promotion_rejected_frequency();
+        return;
+    }
+
+    // Watermark gate: don't promote if DRAM is already under eviction
+    // pressure. The check is best-effort (state can change between this
+    // sample and the actual allocation in PromotionAllocStart).
+    const double used_ratio =
+        MasterMetricManager::instance().get_global_mem_used_ratio();
+    if (used_ratio >= eviction_high_watermark_ratio_) {
+        MasterMetricManager::instance().inc_promotion_rejected_watermark();
+        return;
+    }
+
+    // Acquire a fresh RW shard accessor for dedup, refcnt-pin, and task
+    // record. Safe to call here because GetReplicaList has already released
+    // its RO accessor.
+    MetadataAccessorRW accessor(this, object_id);
+    if (!accessor.Exists()) {
+        return;
+    }
+    auto& metadata = accessor.Get();
+    auto& tenant_state = accessor.GetTenantState();
+
+    // Dedup: don't queue twice if a promotion is already in flight or if a
+    // MEMORY replica has appeared since GetReplicaList observed only-disk.
+    if (tenant_state.promotion_tasks.count(key) > 0) {
+        return;
+    }
+    if (metadata.HasReplica(&Replica::fn_is_memory_replica)) {
+        return;
+    }
+
+    // Cap gate: read the cluster-wide in-flight count. Soft cap — a
+    // benign TOCTOU race between this load and the emplace below can let
+    // a few extra tasks slip in, but the per-shard mutex already
+    // serializes inserts within a shard and the dedup gate above prevents
+    // duplicate work, so the worst case is N concurrent inserters across
+    // distinct shards each admitting one extra task. Atomic load is
+    // relaxed because the value is purely advisory.
+    if (promotion_in_flight_.load(std::memory_order_relaxed) >=
+        promotion_queue_limit_) {
+        MasterMetricManager::instance().inc_promotion_rejected_cap();
+        return;
+    }
+
+    // Find the LOCAL_DISK source replica.
+    Replica* source = nullptr;
+    metadata.VisitReplicas(&Replica::fn_is_local_disk_replica,
+                           [&source](Replica& r) {
+                               if (source == nullptr) source = &r;
+                           });
+    if (source == nullptr) {
+        return;
+    }
+
+    // Pin the source replica.
+    source->inc_refcnt();
+    const uint64_t object_size =
+        source->get_descriptor().get_local_disk_descriptor().object_size;
+
+    // Try to enqueue on the holder client. On failure, drop the refcnt back.
+    auto push_result = PushPromotionQueue(object_id, *source);
+    if (!push_result) {
+        source->dec_refcnt();
+        VLOG(1) << "promotion_push_failed key=" << key
+                << " error=" << push_result.error();
+        return;
+    }
+
+    // Capture the holder client_id so NotifyPromotionSuccess can reject
+    // calls from other clients. PushPromotionQueue already validated
+    // get_local_disk_client_id() returns a value, so .value() is safe.
+    const UUID holder_id = source->get_local_disk_client_id().value();
+
+    // Record the in-flight task. alloc_id is filled in by
+    // PromotionAllocStart once the new MEMORY replica is staged.
+    tenant_state.promotion_tasks.emplace(
+        key, PromotionTask{.source_id = source->id(),
+                           .alloc_id = 0,
+                           .object_size = object_size,
+                           .start_time = std::chrono::system_clock::now(),
+                           .holder_id = holder_id});
+    promotion_in_flight_.fetch_add(1, std::memory_order_relaxed);
+    MasterMetricManager::instance().inc_promotion_in_flight();
+    MasterMetricManager::instance().inc_promotion_admitted();
+    VLOG(1) << "promotion_queued key=" << key << " size=" << object_size;
+}
+
+auto MasterService::PromotionObjectHeartbeat(const UUID& client_id)
+    -> tl::expected<std::vector<PromotionTaskItem>, ErrorCode> {
+    std::shared_lock<std::shared_mutex> shared_lock(snapshot_mutex_);
+    ScopedLocalDiskSegmentAccess local_disk_segment_access =
+        segment_manager_.getLocalDiskSegmentAccess();
+    auto& client_local_disk_segment =
+        local_disk_segment_access.getClientLocalDiskSegment();
+    auto local_disk_segment_it = client_local_disk_segment.find(client_id);
+    if (local_disk_segment_it == client_local_disk_segment.end()) {
+        return tl::make_unexpected(ErrorCode::SEGMENT_NOT_FOUND);
+    }
+    MutexLocker locker(&local_disk_segment_it->second->offloading_mutex_);
+    // Return at most promotion_max_per_heartbeat_ tasks. Each task does
+    // a synchronous SSD read + RDMA write on the client side; allowing
+    // more than one per heartbeat risks blocking past the client-
+    // liveness window and the master marking the client dead. The rest
+    // stay queued in promotion_objects for subsequent heartbeats. The
+    // cap must live here (server side) rather than on the client so
+    // leftover work isn't silently dropped.
+    auto& src = local_disk_segment_it->second->promotion_objects;
+    std::vector<PromotionTaskItem> result;
+    while (result.size() < promotion_max_per_heartbeat_ && !src.empty()) {
+        auto node = src.extract(src.begin());
+        result.push_back(std::move(node.mapped()));
+    }
+    return result;
+}
+
+auto MasterService::PromotionAllocStart(
+    const UUID& client_id, const std::string& key, const std::string& tenant_id,
+    uint64_t size, const std::vector<std::string>& preferred_segments)
+    -> tl::expected<PromotionAllocStartResponse, ErrorCode> {
+    auto normalized_tenant_result = NormalizeTenantIdForWrite(tenant_id);
+    if (!normalized_tenant_result) {
+        return tl::make_unexpected(normalized_tenant_result.error());
+    }
+    const ObjectIdentity object_id{normalized_tenant_result.value(), key};
+    std::shared_lock<std::shared_mutex> shared_lock(snapshot_mutex_);
+    MetadataAccessorRW accessor(this, object_id);
+    if (!accessor.Exists()) {
+        return tl::make_unexpected(ErrorCode::OBJECT_NOT_FOUND);
+    }
+    auto& metadata = accessor.Get();
+
+    // Verify the in-flight task still exists before allocating. The
+    // reaper can sweep it between the holder's heartbeat and this
+    // AllocStart call (a hung client, GC pause, or HA failover can
+    // stall AllocStart past put_start_release_timeout_sec_). If we
+    // allocated and AddReplicas'd anyway, the staged PROCESSING MEMORY
+    // replica would have no PromotionTask pointing at it: the generic
+    // PROCESSING reaper iterates tenant_state.processing_keys (never
+    // populated by promotion) and the promotion-task reaper would have
+    // nothing left to iterate, leaking the buffer until the object is
+    // removed or evicted. The shard mutex is held for the rest of this
+    // function, so the iterator stays valid across the allocation step.
+    auto& tenant_state = accessor.GetTenantState();
+    auto task_it = tenant_state.promotion_tasks.find(object_id.user_key);
+    if (task_it == tenant_state.promotion_tasks.end()) {
+        return tl::make_unexpected(ErrorCode::REPLICA_IS_NOT_READY);
+    }
+
+    // Holder-only gate (see PromotionTask::holder_id doc).
+    if (task_it->second.holder_id != client_id) {
+        return tl::make_unexpected(ErrorCode::INVALID_PARAMS);
+    }
+
+    // Defensive size check: must match the source LOCAL_DISK
+    // descriptor's object_size captured at admission. A mismatch would
+    // let a buggy caller request a wrong-sized allocation — smaller
+    // risks RDMA overflow, larger wastes DRAM pinned until reaper TTL.
+    if (task_it->second.object_size != size) {
+        return tl::make_unexpected(ErrorCode::INVALID_PARAMS);
+    }
+    if (metadata.HasReplica(&Replica::fn_is_memory_replica)) {
+        return tl::make_unexpected(ErrorCode::REPLICA_IS_NOT_READY);
+    }
+    if (task_it->second.alloc_id != 0 ||
+        task_it->second.reserved_quota_charge_bytes != 0) {
+        return tl::make_unexpected(ErrorCode::REPLICA_IS_NOT_READY);
+    }
+
+    const uint64_t reserved_quota_charge = size;
+    auto quota_result =
+        ReserveTenantQuota(object_id.tenant_id, reserved_quota_charge);
+    if (!quota_result) {
+        return tl::make_unexpected(quota_result.error());
+    }
+    auto abort_reserved_quota = [&] {
+        AbortTenantQuota(object_id.tenant_id, reserved_quota_charge);
+    };
+
+    // Allocate a single MEMORY replica via the existing strategy, biased to
+    // the holder's mem segment when possible.
+    ReplicateConfig config;
+    config.replica_num = 1;
+    if (!preferred_segments.empty()) {
+        config.preferred_segments = preferred_segments;
+    }
+
+    std::vector<Replica> staged_replicas;
+    {
+        ScopedAllocatorAccess allocator_access =
+            segment_manager_.getAllocatorAccess();
+        const auto& allocator_manager = allocator_access.getAllocatorManager();
+        auto allocation_result = allocation_strategy_->Allocate(
+            allocator_manager, size, config.replica_num, preferred_segments);
+        if (!allocation_result) {
+            abort_reserved_quota();
+            return tl::make_unexpected(allocation_result.error());
+        }
+        staged_replicas = std::move(allocation_result.value());
+    }
+    if (staged_replicas.empty()) {
+        abort_reserved_quota();
+        return tl::make_unexpected(ErrorCode::NO_AVAILABLE_HANDLE);
+    }
+
+    // Append the new PROCESSING MEMORY replica to the existing object's
+    // metadata. Visible only after NotifyPromotionSuccess flips it COMPLETE.
+    Replica::Descriptor desc = staged_replicas[0].get_descriptor();
+    const ReplicaID new_id = staged_replicas[0].id();
+    std::vector<Replica> to_add;
+    to_add.push_back(std::move(staged_replicas[0]));
+    metadata.AddReplicas(std::move(to_add));
+
+    // Record the new replica's ID on the in-flight PromotionTask so
+    // NotifyPromotionSuccess knows exactly which replica to commit. A
+    // concurrent Put on this key may stage other PROCESSING MEMORY
+    // replicas; using alloc_id avoids the "first PROCESSING memory"
+    // ambiguity.
+    //
+    // Also reset start_time so the reaper TTL covers the active-
+    // transfer phase (AllocStart -> SSD read -> RDMA write -> Notify)
+    // measured from when a master-allocated buffer becomes vulnerable,
+    // rather than being consumed by queue-waiting. Without the reset,
+    // a backlogged task could enter active transfer with little TTL
+    // remaining and the reaper could free the staged replica via
+    // EraseReplicaByID mid-RDMA-write. The queue-waiting phase
+    // (alloc_id == 0) is bounded by its own original start_time window
+    // during which the reaper's EraseReplicaByID branch is a no-op.
+    task_it->second.alloc_id = new_id;
+    task_it->second.reserved_quota_charge_bytes = reserved_quota_charge;
+    task_it->second.start_time = std::chrono::system_clock::now();
+    return PromotionAllocStartResponse{std::move(desc)};
+}
+
+auto MasterService::NotifyPromotionSuccess(const UUID& client_id,
+                                           const std::string& key,
+                                           const std::string& tenant_id)
+    -> tl::expected<void, ErrorCode> {
+    std::shared_lock<std::shared_mutex> shared_lock(snapshot_mutex_);
+    const auto object_id = MakeObjectIdentityForRequest(key, tenant_id);
+    MetadataAccessorRW accessor(this, object_id);
+    if (!accessor.Exists()) {
+        return tl::make_unexpected(ErrorCode::OBJECT_NOT_FOUND);
+    }
+    auto& metadata = accessor.Get();
+    auto& tenant_state = accessor.GetTenantState();
+
+    // Look up the in-flight task to find the exact replica we staged. A
+    // concurrent Put on this key may have created other PROCESSING MEMORY
+    // replicas, so we must not just "mark first PROCESSING memory
+    // complete" — that would risk committing someone else's half-written
+    // replica.
+    auto task_it = tenant_state.promotion_tasks.find(object_id.user_key);
+    if (task_it == tenant_state.promotion_tasks.end() ||
+        task_it->second.alloc_id == 0) {
+        return tl::make_unexpected(ErrorCode::REPLICA_IS_NOT_READY);
+    }
+
+    // Holder-only gate (see PromotionTask::holder_id doc).
+    if (task_it->second.holder_id != client_id) {
+        return tl::make_unexpected(ErrorCode::INVALID_PARAMS);
+    }
+
+    bool committed = false;
+    Replica* staged = metadata.GetReplicaByID(task_it->second.alloc_id);
+    if (staged != nullptr && staged->is_memory_replica() &&
+        staged->is_processing()) {
+        staged->mark_complete();
+        committed = true;
+    }
+
+    // Drop the source LOCAL_DISK replica's refcnt and erase the task.
+    auto* source = metadata.GetReplicaByID(task_it->second.source_id);
+    if (source != nullptr) {
+        source->dec_refcnt();
+    }
+    const uint64_t completed_bytes = task_it->second.object_size;
+    const uint64_t reserved_quota_charge =
+        task_it->second.reserved_quota_charge_bytes;
+    if (committed) {
+        const uint64_t actual_charge = CompletedMemoryQuotaCharge(metadata);
+        const uint64_t commit_charge =
+            actual_charge > metadata.committed_quota_charge_bytes
+                ? actual_charge - metadata.committed_quota_charge_bytes
+                : 0;
+        const uint64_t abort_charge =
+            reserved_quota_charge > commit_charge
+                ? reserved_quota_charge - commit_charge
+                : 0;
+        CommitTenantQuota(object_id.tenant_id, commit_charge);
+        AbortTenantQuota(object_id.tenant_id, abort_charge);
+        metadata.committed_quota_charge_bytes = actual_charge;
+    } else {
+        AbortTenantQuota(object_id.tenant_id, reserved_quota_charge);
+    }
+    task_it->second.reserved_quota_charge_bytes = 0;
+    tenant_state.promotion_tasks.erase(task_it);
+    promotion_in_flight_.fetch_sub(1, std::memory_order_relaxed);
+    MasterMetricManager::instance().dec_promotion_in_flight();
+    if (committed) {
+        SyncCacheTotalAccounting(metadata);
+        MasterMetricManager::instance().inc_promotion_completed();
+        MasterMetricManager::instance().inc_promotion_completed_bytes(
+            static_cast<int64_t>(completed_bytes));
+    } else {
+        MasterMetricManager::instance().inc_promotion_cancelled();
+    }
+
+    // Erase the per-client promotion_objects entry (best-effort; the
+    // heartbeat may have already drained it).
+    {
+        ScopedLocalDiskSegmentAccess local_disk_segment_access =
+            segment_manager_.getLocalDiskSegmentAccess();
+        auto& client_local_disk_segment =
+            local_disk_segment_access.getClientLocalDiskSegment();
+        auto it = client_local_disk_segment.find(client_id);
+        if (it != client_local_disk_segment.end()) {
+            MutexLocker locker(&it->second->offloading_mutex_);
+            it->second->promotion_objects.erase(MakeTenantScopedStorageKey(
+                object_id.tenant_id, object_id.user_key));
+        }
+    }
+
+    if (!committed) {
+        return tl::make_unexpected(ErrorCode::REPLICA_IS_NOT_READY);
+    }
+    return {};
+}
+
+auto MasterService::NotifyPromotionFailure(const UUID& client_id,
+                                           const std::string& key,
+                                           const std::string& tenant_id)
+    -> tl::expected<void, ErrorCode> {
+    std::shared_lock<std::shared_mutex> shared_lock(snapshot_mutex_);
+    const auto object_id = MakeObjectIdentityForRequest(key, tenant_id);
+    MetadataAccessorRW accessor(this, object_id);
+    if (!accessor.Exists()) {
+        return tl::make_unexpected(ErrorCode::OBJECT_NOT_FOUND);
+    }
+    auto& metadata = accessor.Get();
+    auto& tenant_state = accessor.GetTenantState();
+
+    auto task_it = tenant_state.promotion_tasks.find(object_id.user_key);
+    if (task_it == tenant_state.promotion_tasks.end()) {
+        // No task to release. Either the reaper already swept it, or the
+        // client never had a task here. Return OK to keep this RPC
+        // idempotent — repeated failure notifications on the same key
+        // should be safe.
+        return {};
+    }
+
+    // Holder-only gate (see PromotionTask::holder_id doc).
+    if (task_it->second.holder_id != client_id) {
+        return tl::make_unexpected(ErrorCode::INVALID_PARAMS);
+    }
+
+    // Mirror the reaper's expiry path; see DiscardExpiredProcessingReplicas
+    // Part 4 for the full rationale on each step.
+    auto* source = metadata.GetReplicaByID(task_it->second.source_id);
+    if (source != nullptr) {
+        source->dec_refcnt();
+    }
+    if (task_it->second.alloc_id != 0) {
+        const ReplicaID alloc_id = task_it->second.alloc_id;
+        EraseReplicasWithCacheTotalAccounting(
+            metadata, [alloc_id](const Replica& replica) {
+                return replica.id() == alloc_id;
+            });
+    }
+    AbortTenantQuota(object_id.tenant_id,
+                     task_it->second.reserved_quota_charge_bytes);
+    task_it->second.reserved_quota_charge_bytes = 0;
+    tenant_state.promotion_tasks.erase(task_it);
+    promotion_in_flight_.fetch_sub(1, std::memory_order_relaxed);
+    MasterMetricManager::instance().dec_promotion_in_flight();
+    MasterMetricManager::instance().inc_promotion_failed();
+
+    // Clear the holder's per-client promotion_objects entry. Same
+    // best-effort cleanup pattern as NotifyPromotionSuccess — the
+    // heartbeat may have already drained it.
+    {
+        ScopedLocalDiskSegmentAccess local_disk_segment_access =
+            segment_manager_.getLocalDiskSegmentAccess();
+        auto& client_local_disk_segment =
+            local_disk_segment_access.getClientLocalDiskSegment();
+        auto it = client_local_disk_segment.find(client_id);
+        if (it != client_local_disk_segment.end()) {
+            MutexLocker locker(&it->second->offloading_mutex_);
+            it->second->promotion_objects.erase(MakeTenantScopedStorageKey(
+                object_id.tenant_id, object_id.user_key));
+        }
+    }
+
+    return {};
+}
+
 void MasterService::EvictionThreadFunc() {
     VLOG(1) << "action=eviction_thread_started";
 
@@ -2236,10 +5439,10 @@ void MasterService::EvictionThreadFunc() {
         double used_ratio =
             MasterMetricManager::instance().get_global_mem_used_ratio();
         if (used_ratio > eviction_high_watermark_ratio_ ||
-            (need_eviction_ && eviction_ratio_ > 0.0)) {
+            (need_mem_eviction_ && eviction_ratio_ > 0.0)) {
             LOG(INFO) << "[EVICT-TRIGGER] memory_ratio=" << used_ratio
                       << " high_watermark=" << eviction_high_watermark_ratio_
-                      << " need_eviction=" << need_eviction_
+                      << " need_mem_eviction=" << need_mem_eviction_
                       << " eviction_ratio=" << eviction_ratio_;
             double evict_ratio_target = std::max(
                 eviction_ratio_,
@@ -2265,6 +5468,22 @@ void MasterService::EvictionThreadFunc() {
             last_discard_time = now;
         }
 
+#ifdef USE_NOF
+        double nof_used_ratio =
+            MasterMetricManager::instance().get_global_nof_used_ratio();
+        if (nof_used_ratio > nof_eviction_high_watermark_ratio_ ||
+            (need_nof_eviction_ && nof_eviction_ratio_ > 0.0)) {
+            double nof_evict_ratio_target =
+                std::max(nof_eviction_ratio_,
+                         nof_used_ratio - nof_eviction_high_watermark_ratio_ +
+                             nof_eviction_ratio_);
+            double nof_evict_ratio_lowerbound =
+                std::max(nof_evict_ratio_target * 0.5,
+                         nof_used_ratio - nof_eviction_high_watermark_ratio_);
+            NoFBatchEvict(nof_evict_ratio_target, nof_evict_ratio_lowerbound);
+        }
+#endif
+
         std::this_thread::sleep_for(
             std::chrono::milliseconds(kEvictionThreadSleepMs));
     }
@@ -2277,127 +5496,152 @@ void MasterService::DiscardExpiredProcessingReplicas(
     const std::chrono::system_clock::time_point& now) {
     std::list<DiscardedReplicas> discarded_replicas;
 
-    // Part 1: Discard expired PutStart operations.
-    for (auto key_it = shard->processing_keys.begin();
-         key_it != shard->processing_keys.end();) {
-        auto it = shard->metadata.find(*key_it);
-        if (it == shard->metadata.end()) {
-            // The key has been removed from metadata. This should be
-            // impossible.
-            LOG(ERROR) << "Key " << *key_it
-                       << " was removed while in processing";
-            key_it = shard->processing_keys.erase(key_it);
-            continue;
-        }
+    for (auto tenant_it = shard->tenants.begin();
+         tenant_it != shard->tenants.end();) {
+        auto& tenant_state = tenant_it->second;
 
-        auto& metadata = it->second;
-        // If the object is not valid or not in processing state, just
-        // remove it from the processing set.
-        if (!metadata.IsValid() ||
-            metadata.AllReplicas(&Replica::fn_is_completed)) {
-            if (!metadata.IsValid()) {
-                shard->metadata.erase(it);
-            }
-            key_it = shard->processing_keys.erase(key_it);
-            continue;
-        }
-
-        // If the object's PutStart timedout, discard and release it's
-        // space. Note that instead of releasing the space directly, we
-        // insert the replicas into the discarded list so that the
-        // discarding and releasing operations can be recorded in
-        // statistics.
-        const auto ttl =
-            metadata.put_start_time + put_start_release_timeout_sec_;
-        if (ttl < now) {
-            auto replicas = metadata.PopReplicas(&Replica::fn_is_processing);
-            if (!replicas.empty()) {
-                discarded_replicas.emplace_back(std::move(replicas), ttl);
+        for (auto key_it = tenant_state.processing_keys.begin();
+             key_it != tenant_state.processing_keys.end();) {
+            auto it = tenant_state.metadata.find(*key_it);
+            if (it == tenant_state.metadata.end()) {
+                LOG(ERROR) << "Key " << *key_it
+                           << " was removed while in processing";
+                key_it = tenant_state.processing_keys.erase(key_it);
+                continue;
             }
 
-            if (!metadata.IsValid()) {
-                // All replicas of this object are discarded, just
-                // remove the whole object.
-                shard->metadata.erase(it);
+            auto& metadata = it->second;
+            if (!metadata.IsValid() ||
+                metadata.AllReplicas(&Replica::fn_is_completed)) {
+                if (!metadata.IsValid()) {
+                    EraseMetadata(tenant_state, it, tenant_it->first,
+                                  QuotaEraseMode::kFull, &shard);
+                }
+                key_it = tenant_state.processing_keys.erase(key_it);
+                continue;
             }
 
-            key_it = shard->processing_keys.erase(key_it);
-            continue;
+            const auto ttl =
+                metadata.put_start_time + put_start_release_timeout_sec_;
+            if (ttl < now) {
+                auto replicas =
+                    metadata.PopReplicas(&Replica::fn_is_processing);
+                if (!replicas.empty()) {
+                    discarded_replicas.emplace_back(std::move(replicas), ttl);
+                }
+                if (!metadata.IsValid()) {
+                    EraseMetadata(tenant_state, it, tenant_it->first,
+                                  QuotaEraseMode::kFull, &shard);
+                }
+                key_it = tenant_state.processing_keys.erase(key_it);
+                continue;
+            }
+            key_it++;
         }
 
-        key_it++;
-    }
+        for (auto task_it = tenant_state.replication_tasks.begin();
+             task_it != tenant_state.replication_tasks.end();) {
+            auto metadata_it = tenant_state.metadata.find(task_it->first);
+            if (metadata_it == tenant_state.metadata.end()) {
+                LOG(ERROR) << "Key " << task_it->first
+                           << " was removed with ongoing replication task";
+                AbortReplicationTaskQuota(tenant_it->first, task_it->second);
+                task_it = tenant_state.replication_tasks.erase(task_it);
+                continue;
+            }
 
-    // Part 2: Discard expired CopyStart/MoveStart operations.
-    for (auto task_it = shard->replication_tasks.begin();
-         task_it != shard->replication_tasks.end();) {
-        auto metadata_it = shard->metadata.find(task_it->first);
-        if (metadata_it == shard->metadata.end()) {
-            // The key has been removed from metadata. This should be
-            // impossible.
-            LOG(ERROR) << "Key " << task_it->first
-                       << " was removed with ongoing replication task";
-            task_it = shard->replication_tasks.erase(task_it);
-            continue;
-        }
+            const auto ttl =
+                task_it->second.start_time + put_start_release_timeout_sec_;
+            if (ttl > now) {
+                task_it++;
+                continue;
+            }
 
-        const auto ttl =
-            task_it->second.start_time + put_start_release_timeout_sec_;
-        if (ttl > now) {
-            // The task is not expired, skip it.
-            task_it++;
-            continue;
-        }
-
-        auto& metadata = metadata_it->second;
-
-        // Release source refcnt.
-        auto source = metadata.GetReplicaByID(task_it->second.source_id);
-        if (source != nullptr) {
-            source->dec_refcnt();
-        }
-
-        // Discard allocated replicas.
-        auto& replica_ids = task_it->second.replica_ids;
-        auto replicas =
-            metadata.PopReplicas([&replica_ids](const Replica& replica) {
-                auto it = std::find(replica_ids.begin(), replica_ids.end(),
-                                    replica.id());
-                return it != replica_ids.end();
-            });
-        if (!replicas.empty()) {
-            discarded_replicas.emplace_back(std::move(replicas), ttl);
-        }
-
-        // Check whether the object is still valid.
-        if (!metadata.IsValid()) {
-            shard->metadata.erase(metadata_it);
-        }
-
-        task_it = shard->replication_tasks.erase(task_it);
-    }
-
-    // Part 3: Discard expired offloading operations.
-    for (auto task_it = shard->offloading_tasks.begin();
-         task_it != shard->offloading_tasks.end();) {
-        const auto ttl =
-            task_it->second.start_time + put_start_release_timeout_sec_;
-        if (ttl > now) {
-            task_it++;
-            continue;
-        }
-
-        auto metadata_it = shard->metadata.find(task_it->first);
-        if (metadata_it != shard->metadata.end()) {
-            auto source =
-                metadata_it->second.GetReplicaByID(task_it->second.source_id);
+            auto& metadata = metadata_it->second;
+            auto source = metadata.GetReplicaByID(task_it->second.source_id);
             if (source != nullptr) {
                 source->dec_refcnt();
             }
+
+            auto& replica_ids = task_it->second.replica_ids;
+            auto replicas = PopReplicasWithCacheTotalAccounting(
+                metadata, [&replica_ids](const Replica& replica) {
+                    auto it = std::find(replica_ids.begin(), replica_ids.end(),
+                                        replica.id());
+                    return it != replica_ids.end();
+                });
+            if (!replicas.empty()) {
+                discarded_replicas.emplace_back(std::move(replicas), ttl);
+            }
+            if (!metadata.IsValid()) {
+                EraseMetadata(tenant_state, metadata_it, tenant_it->first,
+                              QuotaEraseMode::kFull, &shard);
+            }
+            AbortReplicationTaskQuota(tenant_it->first, task_it->second);
+            task_it = tenant_state.replication_tasks.erase(task_it);
         }
 
-        LOG(WARNING) << "Offloading task expired for key: " << task_it->first;
-        task_it = shard->offloading_tasks.erase(task_it);
+        for (auto task_it = tenant_state.offloading_tasks.begin();
+             task_it != tenant_state.offloading_tasks.end();) {
+            const auto ttl =
+                task_it->second.start_time + put_start_release_timeout_sec_;
+            if (ttl > now) {
+                task_it++;
+                continue;
+            }
+            auto metadata_it = tenant_state.metadata.find(task_it->first);
+            if (metadata_it != tenant_state.metadata.end()) {
+                auto source = metadata_it->second.GetReplicaByID(
+                    task_it->second.source_id);
+                if (source != nullptr) {
+                    source->dec_refcnt();
+                }
+            }
+            LOG(WARNING) << "Offloading task expired for key: "
+                         << task_it->first;
+            task_it = tenant_state.offloading_tasks.erase(task_it);
+        }
+
+        for (auto task_it = tenant_state.promotion_tasks.begin();
+             task_it != tenant_state.promotion_tasks.end();) {
+            const auto ttl =
+                task_it->second.start_time + put_start_release_timeout_sec_;
+            if (ttl > now) {
+                task_it++;
+                continue;
+            }
+            auto metadata_it = tenant_state.metadata.find(task_it->first);
+            if (metadata_it != tenant_state.metadata.end()) {
+                auto source = metadata_it->second.GetReplicaByID(
+                    task_it->second.source_id);
+                if (source != nullptr) {
+                    source->dec_refcnt();
+                }
+                if (task_it->second.alloc_id != 0) {
+                    const ReplicaID alloc_id = task_it->second.alloc_id;
+                    EraseReplicasWithCacheTotalAccounting(
+                        metadata_it->second,
+                        [alloc_id](const Replica& replica) {
+                            return replica.id() == alloc_id;
+                        });
+                }
+            }
+            AbortTenantQuota(tenant_it->first,
+                             task_it->second.reserved_quota_charge_bytes);
+            task_it->second.reserved_quota_charge_bytes = 0;
+            LOG(WARNING) << "Promotion task expired for key: "
+                         << task_it->first;
+            task_it = tenant_state.promotion_tasks.erase(task_it);
+            promotion_in_flight_.fetch_sub(1, std::memory_order_relaxed);
+            MasterMetricManager::instance().dec_promotion_in_flight();
+            MasterMetricManager::instance().inc_promotion_expired();
+        }
+
+        if (tenant_state.Empty()) {
+            tenant_it = shard->tenants.erase(tenant_it);
+        } else {
+            ++tenant_it;
+        }
     }
 
     if (!discarded_replicas.empty()) {
@@ -2425,8 +5669,18 @@ uint64_t MasterService::ReleaseExpiredDiscardedReplicas(
 void MasterService::SnapshotThreadFunc() {
     LOG(INFO) << "[Snapshot] snapshot_thread started";
     while (snapshot_running_) {
-        std::this_thread::sleep_for(
-            std::chrono::seconds(snapshot_interval_seconds_));
+        // Wait for the next snapshot cycle, but allow fast shutdown.
+        {
+            std::unique_lock<std::mutex> lk(snapshot_thread_mutex_);
+            snapshot_thread_cv_.wait_for(
+                lk, std::chrono::seconds(snapshot_interval_seconds_),
+                [&] { return !snapshot_running_.load(); });
+        }
+
+        if (!snapshot_running_) {
+            break;
+        }
+
         if (!enable_snapshot_) {
             // Snapshot is disabled
             LOG(INFO)
@@ -2448,7 +5702,9 @@ void MasterService::SnapshotThreadFunc() {
             continue;
         }
 
-        const std::string path_prefix = SNAPSHOT_ROOT + "/" + snapshot_id + "/";
+        const std::string& snapshot_root =
+            snapshot_catalog_store_->GetSnapshotRoot();
+        const std::string path_prefix = snapshot_root + snapshot_id + "/";
         const std::string manifest_path = path_prefix + SNAPSHOT_MANIFEST_FILE;
         auto descriptor =
             BuildSnapshotDescriptor(snapshot_id, manifest_path, path_prefix);
@@ -2698,7 +5954,7 @@ MasterService::ResolveSnapshotSequenceId() const {
 
     uint64_t sequence_id = 0;
     auto err = oplog_store.value()->GetLatestSequenceId(sequence_id);
-    if (err == ErrorCode::ETCD_KEY_NOT_EXIST) {
+    if (err == ErrorCode::OPLOG_ENTRY_NOT_FOUND) {
         return ha::OpLogSequenceId{0};
     }
     if (err != ErrorCode::OK) {
@@ -2756,8 +6012,10 @@ MasterService::BuildSnapshotDescriptor(const std::string& snapshot_id,
         return tl::make_unexpected(sequence_id.error());
     }
 
-    auto descriptor =
-        ha::snapshot_catalog_store_detail::MakeSnapshotDescriptor(snapshot_id);
+    const std::string& snapshot_root =
+        snapshot_catalog_store_->GetSnapshotRoot();
+    auto descriptor = ha::snapshot_catalog_store_detail::MakeSnapshotDescriptor(
+        snapshot_root, snapshot_id);
     descriptor.last_included_seq = sequence_id.value();
     descriptor.producer_view_version = view_version_;
     descriptor.manifest_key = manifest_path;
@@ -2768,7 +6026,9 @@ MasterService::BuildSnapshotDescriptor(const std::string& snapshot_id,
 
 tl::expected<void, SerializationError> MasterService::PersistState(
     const std::string& snapshot_id) {
-    const std::string path_prefix = SNAPSHOT_ROOT + "/" + snapshot_id + "/";
+    const std::string& snapshot_root =
+        snapshot_catalog_store_->GetSnapshotRoot();
+    const std::string path_prefix = snapshot_root + snapshot_id + "/";
     const std::string manifest_path = path_prefix + SNAPSHOT_MANIFEST_FILE;
     auto descriptor =
         BuildSnapshotDescriptor(snapshot_id, manifest_path, path_prefix);
@@ -2936,7 +6196,8 @@ tl::expected<void, SerializationError> MasterService::PersistState(
         }
 
         // Publish snapshot catalog entry and advance the latest marker.
-        std::string latest_path = SNAPSHOT_ROOT + "/" + SNAPSHOT_LATEST_FILE;
+        std::string latest_path =
+            snapshot_catalog_store->GetSnapshotRoot() + SNAPSHOT_LATEST_FILE;
         std::string latest_content = snapshot_id;
 
         auto publish_result = snapshot_catalog_store->Publish(descriptor);
@@ -3166,7 +6427,8 @@ bool MasterService::TryRestoreStateFromSnapshot(
     const std::string& state_id = snapshot.snapshot_id;
     std::string path_prefix = snapshot.object_prefix;
     if (path_prefix.empty()) {
-        path_prefix = SNAPSHOT_ROOT + "/" + state_id + "/";
+        path_prefix =
+            snapshot_catalog_store_->GetSnapshotRoot() + state_id + "/";
     }
 
     std::string manifest_path = snapshot.manifest_key;
@@ -3336,69 +6598,66 @@ bool MasterService::TryRestoreStateFromSnapshot(
             if (!skip_cleanup) {
                 auto cleanup_now = now;
                 for (auto& shard : metadata_shards_) {
-                    for (auto it = shard.metadata.begin();
-                         it != shard.metadata.end();) {
-                        if (it->second.HasDiffRepStatus(
-                                ReplicaStatus::COMPLETE) ||
-                            (it->second.IsLeaseExpired(cleanup_now) &&
-                             !it->second.IsSoftPinned(cleanup_now))) {
-                            VLOG(1)
-                                << "clear metadata key=" << it->first
-                                << " ,lease_timeout="
-                                << std::chrono::duration_cast<
-                                       std::chrono::milliseconds>(
-                                       it->second.lease_timeout
-                                           .time_since_epoch())
-                                       .count()
-                                << " ,soft_pin_timeout="
-                                << (it->second.soft_pin_timeout.has_value()
-                                        ? std::to_string(
-                                              std::chrono::duration_cast<
-                                                  std::chrono::milliseconds>(
-                                                  it->second.soft_pin_timeout
-                                                      .value()
-                                                      .time_since_epoch())
-                                                  .count())
-                                        : "null");
-                            it = shard.metadata.erase(it);
+                    for (auto tenant_it = shard.tenants.begin();
+                         tenant_it != shard.tenants.end();) {
+                        auto& tenant_state = tenant_it->second;
+                        for (auto it = tenant_state.metadata.begin();
+                             it != tenant_state.metadata.end();) {
+                            if (it->second.HasDiffRepStatus(
+                                    ReplicaStatus::COMPLETE) ||
+                                (it->second.IsLeaseExpired(cleanup_now) &&
+                                 !it->second.IsSoftPinned(cleanup_now))) {
+                                VLOG(1) << "clear metadata key=" << it->first;
+                                it = EraseMetadata(tenant_state, it,
+                                                   tenant_it->first);
+                            } else {
+                                ++it;
+                            }
+                        }
+                        if (tenant_state.Empty()) {
+                            tenant_it = shard.tenants.erase(tenant_it);
                         } else {
-                            ++it;
+                            ++tenant_it;
                         }
                     }
                 }
             }
 
             MasterMetricManager::instance().reset_allocated_mem_size();
+            RebuildCacheTotalAccounting();
             for (auto& segment_name : segment_names) {
                 MasterMetricManager::instance()
                     .reset_segment_allocated_mem_size(segment_name);
             }
 
             for (auto& shard : metadata_shards_) {
-                for (auto it = shard.metadata.begin();
-                     it != shard.metadata.end();) {
-                    for (auto& replica : it->second.GetAllReplicas()) {
-                        if (!replica.get_descriptor().is_memory_replica()) {
-                            continue;
+                for (auto& [tenant_id, tenant_state] : shard.tenants) {
+                    for (auto it = tenant_state.metadata.begin();
+                         it != tenant_state.metadata.end();) {
+                        for (auto& replica : it->second.GetAllReplicas()) {
+                            if (!replica.get_descriptor().is_memory_replica()) {
+                                continue;
+                            }
+                            auto temp_segment_names =
+                                replica.get_segment_names();
+                            if (temp_segment_names.empty()) {
+                                continue;
+                            }
+                            if (!temp_segment_names[0].has_value()) {
+                                continue;
+                            }
+                            auto buffer_descriptor =
+                                replica.get_descriptor()
+                                    .get_memory_descriptor()
+                                    .buffer_descriptor;
+                            MasterMetricManager::instance()
+                                .inc_allocated_mem_size(
+                                    temp_segment_names[0].value(),
+                                    static_cast<int64_t>(
+                                        buffer_descriptor.size_));
                         }
-                        auto temp_segment_names = replica.get_segment_names();
-                        if (temp_segment_names.empty()) {
-                            continue;
-                        }
-
-                        std::string temp_segment_name;
-                        if (temp_segment_names[0].has_value()) {
-                            temp_segment_name = temp_segment_names[0].value();
-                        }
-
-                        auto buffer_descriptor = replica.get_descriptor()
-                                                     .get_memory_descriptor()
-                                                     .buffer_descriptor;
-                        MasterMetricManager::instance().inc_allocated_mem_size(
-                            temp_segment_name,
-                            static_cast<int64_t>(buffer_descriptor.size_));
+                        ++it;
                     }
-                    ++it;
                 }
             }
 
@@ -3473,12 +6732,245 @@ void MasterService::ResetStateAfterFailedRestoreAttempt() {
 
     MasterMetricManager::instance().reset_allocated_mem_size();
     MasterMetricManager::instance().reset_total_mem_capacity();
+    MasterMetricManager::instance().reset_cache_total_nums();
 }
 
 ha::SnapshotCatalogStore* MasterService::GetSnapshotCatalogStore() {
     return snapshot_catalog_store_.get();
 }
 
+MasterService::TenantQuotaEvictionResult
+MasterService::EvictTenantMemoryForQuota(const std::string& tenant_id,
+                                         uint64_t target_bytes) {
+    TenantQuotaEvictionResult total;
+    if (!enable_multi_tenants_ || target_bytes == 0) {
+        return total;
+    }
+
+    const auto normalized_tenant = NormalizeTenantId(tenant_id);
+    auto now = std::chrono::system_clock::now();
+    std::shared_lock<std::shared_mutex> shared_lock(snapshot_mutex_);
+
+    auto is_evictable_memory_replica = [](const Replica& replica) {
+        return replica.is_memory_replica() && replica.is_completed() &&
+               replica.get_refcnt() == 0;
+    };
+    auto can_evict_replicas = [&](const ObjectMetadata& metadata) {
+        return metadata.HasReplica(is_evictable_memory_replica);
+    };
+    auto has_local_disk_replica = [](const ObjectMetadata& metadata) {
+        return metadata.HasReplica(&Replica::fn_is_local_disk_replica);
+    };
+    auto evict_replicas =
+        [&, this](ObjectMetadata& metadata,
+                  std::vector<std::vector<Replica>>& deferred_replicas) {
+            const uint64_t before_charge = CompletedMemoryQuotaCharge(metadata);
+            auto replicas = PopReplicasWithCacheTotalAccounting(
+                metadata, is_evictable_memory_replica);
+            const uint64_t replica_count = replicas.size();
+            if (!replicas.empty()) {
+                deferred_replicas.emplace_back(std::move(replicas));
+            }
+            const uint64_t after_charge = CompletedMemoryQuotaCharge(metadata);
+            if (before_charge > after_charge) {
+                ReleaseCommittedQuotaCharge(metadata,
+                                            before_charge - after_charge);
+            }
+            return metadata.size * replica_count;
+        };
+    long offload_queued_this_call = 0;
+    long offload_deferred_count = 0;
+    long offload_cap_forced_count = 0;
+    long offload_push_failed_forced = 0;
+    const long offload_cap =
+        offload_on_evict_
+            ? static_cast<long>(offloading_queue_limit_ * offload_cap_ratio_)
+            : 0;
+
+    auto try_evict_or_offload =
+        [&, this](const std::string& key, ObjectMetadata& metadata,
+                  TenantState& tenant_state,
+                  std::vector<std::vector<Replica>>& deferred_replicas) {
+            if (!offload_on_evict_) {
+                return evict_replicas(metadata, deferred_replicas);
+            }
+
+            if (has_local_disk_replica(metadata)) {
+                return evict_replicas(metadata, deferred_replicas);
+            }
+
+            if (offload_force_evict_ &&
+                offload_queued_this_call >= offload_cap) {
+                ++offload_cap_forced_count;
+                return evict_replicas(metadata, deferred_replicas);
+            }
+
+            bool queued = false;
+            metadata.VisitReplicas(
+                is_evictable_memory_replica,
+                [this, &key, &normalized_tenant, &tenant_state, &queued,
+                 &now](Replica& replica) {
+                    if (queued) {
+                        return;
+                    }
+                    auto result = PushOffloadingQueue(
+                        MakeObjectIdentity(key, normalized_tenant), replica);
+                    if (result) {
+                        replica.inc_refcnt();
+                        tenant_state.offloading_tasks.emplace(
+                            key, OffloadingTask{replica.id(), now});
+                        queued = true;
+                    }
+                });
+
+            if (queued) {
+                ++offload_queued_this_call;
+                ++offload_deferred_count;
+                return evict_replicas(metadata, deferred_replicas);
+            }
+
+            if (offload_force_evict_) {
+                ++offload_push_failed_forced;
+                return evict_replicas(metadata, deferred_replicas);
+            }
+            return uint64_t{0};
+        };
+
+    auto try_evict_group_or_object =
+        [&, this](const std::string& key, ObjectMetadata& metadata,
+                  TenantState& tenant_state,
+                  std::vector<std::vector<Replica>>& deferred_replicas,
+                  bool allow_soft_pinned) -> TenantQuotaEvictionResult {
+        if (!metadata.IsGrouped()) {
+            uint64_t freed = try_evict_or_offload(key, metadata, tenant_state,
+                                                  deferred_replicas);
+            return {.freed_bytes = freed,
+                    .evicted_objects = freed > 0 ? 1U : 0U};
+        }
+
+        auto group_it = tenant_state.group_members.find(metadata.group_id);
+        if (group_it == tenant_state.group_members.end()) {
+            uint64_t freed = try_evict_or_offload(key, metadata, tenant_state,
+                                                  deferred_replicas);
+            return {.freed_bytes = freed,
+                    .evicted_objects = freed > 0 ? 1U : 0U};
+        }
+
+        for (const auto& member_key : group_it->second) {
+            auto member_it = tenant_state.metadata.find(member_key);
+            if (member_it != tenant_state.metadata.end() &&
+                !member_it->second.IsLeaseExpired(now)) {
+                return {};
+            }
+        }
+
+        TenantQuotaEvictionResult result;
+        std::vector<std::string> member_keys(group_it->second.begin(),
+                                             group_it->second.end());
+        for (const auto& member_key : member_keys) {
+            auto member_it = tenant_state.metadata.find(member_key);
+            if (member_it == tenant_state.metadata.end()) {
+                continue;
+            }
+            auto& member_metadata = member_it->second;
+            if (member_metadata.IsHardPinned() ||
+                !member_metadata.IsLeaseExpired(now) ||
+                (!allow_soft_pinned && member_metadata.IsSoftPinned(now)) ||
+                !can_evict_replicas(member_metadata)) {
+                continue;
+            }
+
+            const uint64_t freed = try_evict_or_offload(
+                member_key, member_metadata, tenant_state, deferred_replicas);
+            result.freed_bytes += freed;
+            if (freed > 0) {
+                ++result.evicted_objects;
+            }
+            if (member_key != key && !member_metadata.IsValid()) {
+                EraseMetadata(tenant_state, member_it, normalized_tenant);
+            }
+        }
+        return result;
+    };
+
+    auto pass = [&](bool allow_soft_pinned) {
+        const size_t start_shard = RandomIndex(kNumShards);
+        for (size_t scanned = 0;
+             scanned < kNumShards && total.freed_bytes < target_bytes;
+             ++scanned) {
+            const size_t shard_idx = (start_shard + scanned) % kNumShards;
+            std::vector<std::vector<Replica>> deferred_replicas;
+            {
+                MetadataShardAccessorRW shard(this, shard_idx);
+                auto tenant_it = shard->tenants.find(normalized_tenant);
+                if (tenant_it == shard->tenants.end()) {
+                    continue;
+                }
+                auto& tenant_state = tenant_it->second;
+                for (auto it = tenant_state.metadata.begin();
+                     it != tenant_state.metadata.end() &&
+                     total.freed_bytes < target_bytes;) {
+                    auto& metadata = it->second;
+                    if (metadata.IsHardPinned() ||
+                        !metadata.IsLeaseExpired(now) ||
+                        (!allow_soft_pinned && metadata.IsSoftPinned(now)) ||
+                        !can_evict_replicas(metadata)) {
+                        ++it;
+                        continue;
+                    }
+
+                    auto evict_result = try_evict_group_or_object(
+                        it->first, metadata, tenant_state, deferred_replicas,
+                        allow_soft_pinned);
+                    total.freed_bytes += evict_result.freed_bytes;
+                    total.evicted_objects += evict_result.evicted_objects;
+                    if (!metadata.IsValid()) {
+                        it = EraseMetadata(tenant_state, it, normalized_tenant);
+                    } else {
+                        ++it;
+                    }
+                }
+                if (tenant_state.Empty()) {
+                    shard->tenants.erase(tenant_it);
+                }
+            }
+        }
+    };
+
+    pass(/*allow_soft_pinned=*/false);
+    if (allow_evict_soft_pinned_objects_ && total.freed_bytes < target_bytes) {
+        pass(/*allow_soft_pinned=*/true);
+    }
+
+    if (total.freed_bytes > 0) {
+        MasterMetricManager::instance().inc_tenant_evict_bytes(
+            normalized_tenant,
+            static_cast<int64_t>(std::min<uint64_t>(
+                total.freed_bytes,
+                static_cast<uint64_t>(std::numeric_limits<int64_t>::max()))));
+    }
+    if (offload_on_evict_ && total.freed_bytes == 0 &&
+        offload_deferred_count > 0) {
+        LOG(WARNING) << "[TENANT-EVICT] No memory freed for tenant "
+                     << normalized_tenant << "; " << offload_deferred_count
+                     << " object(s) deferred for disk offload.";
+    }
+    if (offload_cap_forced_count > 0) {
+        LOG(WARNING) << "[TENANT-EVICT] Offload cap (" << offload_cap
+                     << ") reached for tenant " << normalized_tenant
+                     << "; force-evicted " << offload_cap_forced_count
+                     << " object(s) without disk offload.";
+    }
+    if (offload_push_failed_forced > 0) {
+        LOG(WARNING) << "[TENANT-EVICT] PushOffloadingQueue failed for tenant "
+                     << normalized_tenant << " on "
+                     << offload_push_failed_forced
+                     << " object(s); force-evicted without disk offload "
+                        "(offload_force_evict=true).";
+    }
+    return total;
+}
+
 void MasterService::BatchEvict(double evict_ratio_target,
                                double evict_ratio_lowerbound) {
     if (evict_ratio_target < evict_ratio_lowerbound) {
@@ -3489,28 +6981,34 @@ void MasterService::BatchEvict(double evict_ratio_target,
     }
 
     auto now = std::chrono::system_clock::now();
-    long evicted_count = 0;
-    long object_count = 0;
-    uint64_t total_freed_size = 0;
 
-    // Candidates for second pass eviction
-    std::vector<std::chrono::system_clock::time_point> no_pin_objects;
-    std::vector<std::chrono::system_clock::time_point> soft_pin_objects;
-
-    auto can_evict_replicas = [](const ObjectMetadata& metadata) {
-        return metadata.HasReplica([](const Replica& replica) {
-            return replica.is_memory_replica() && replica.is_completed() &&
-                   replica.get_refcnt() == 0;
-        });
+    auto is_evictable_memory_replica = [](const Replica& replica) {
+        return replica.is_memory_replica() && replica.is_completed() &&
+               replica.get_refcnt() == 0;
     };
 
-    auto evict_replicas = [](ObjectMetadata& metadata) {
-        return metadata.EraseReplicas([](const Replica& replica) {
-            return replica.is_memory_replica() && replica.is_completed() &&
-                   replica.get_refcnt() == 0;
-        });
+    auto can_evict_replicas = [&](const ObjectMetadata& metadata) {
+        return metadata.HasReplica(is_evictable_memory_replica);
     };
 
+    auto evict_replicas =
+        [&, this](ObjectMetadata& metadata,
+                  std::vector<std::vector<Replica>>& deferred_replicas) {
+            const uint64_t before_charge = CompletedMemoryQuotaCharge(metadata);
+            auto replicas = PopReplicasWithCacheTotalAccounting(
+                metadata, is_evictable_memory_replica);
+            const size_t replica_count = replicas.size();
+            if (!replicas.empty()) {
+                deferred_replicas.emplace_back(std::move(replicas));
+            }
+            const uint64_t after_charge = CompletedMemoryQuotaCharge(metadata);
+            if (before_charge > after_charge) {
+                ReleaseCommittedQuotaCharge(metadata,
+                                            before_charge - after_charge);
+            }
+            return metadata.size * replica_count;
+        };
+
     // --- Offload-on-evict support ---
     long offload_queued_this_cycle = 0;
     long offload_deferred_count = 0;
@@ -3518,7 +7016,7 @@ void MasterService::BatchEvict(double evict_ratio_target,
     long offload_push_failed_forced = 0;  // #keys force-evicted on push fail
     const long offload_cap =
         offload_on_evict_
-            ? static_cast<long>(offloading_queue_limit_ * kOffloadCapRatio)
+            ? static_cast<long>(offloading_queue_limit_ * offload_cap_ratio_)
             : 0;
 
     auto has_local_disk_replica = [](const ObjectMetadata& metadata) {
@@ -3528,16 +7026,18 @@ void MasterService::BatchEvict(double evict_ratio_target,
     // Returns freed bytes. Returns 0 if offload-queued and no additional
     // replicas were evicted (all MEMORY replicas of the key are now pinned).
     auto try_evict_or_offload =
-        [&, this](const std::string& key, ObjectMetadata& metadata,
-                  MetadataShardAccessorRW& shard) -> uint64_t {
+        [&, this](
+            const std::string& tenant_id, const std::string& key,
+            ObjectMetadata& metadata, TenantState& tenant_state,
+            std::vector<std::vector<Replica>>& deferred_replicas) -> uint64_t {
         if (!offload_on_evict_) {
             // Original behavior
-            return metadata.size * evict_replicas(metadata);
+            return evict_replicas(metadata, deferred_replicas);
         }
 
         // LOCAL_DISK replica already exists — safe to delete MEMORY immediately
         if (has_local_disk_replica(metadata)) {
-            return metadata.size * evict_replicas(metadata);
+            return evict_replicas(metadata, deferred_replicas);
         }
 
         // Force-evict cap: if force_evict enabled and cap reached, force
@@ -3545,7 +7045,7 @@ void MasterService::BatchEvict(double evict_ratio_target,
         // flooding.
         if (offload_force_evict_ && offload_queued_this_cycle >= offload_cap) {
             offload_cap_forced_count++;
-            return metadata.size * evict_replicas(metadata);
+            return evict_replicas(metadata, deferred_replicas);
         }
 
         // Queue one MEMORY replica for offload; others will be evicted below.
@@ -3555,12 +7055,14 @@ void MasterService::BatchEvict(double evict_ratio_target,
                 return r.is_memory_replica() && r.is_completed() &&
                        r.get_refcnt() == 0;
             },
-            [this, &key, &shard, &queued, &now](Replica& replica) {
+            [this, &tenant_id, &key, &tenant_state, &queued,
+             &now](Replica& replica) {
                 if (queued) return;  // only need to pin one replica for offload
-                auto result = PushOffloadingQueue(key, replica);
+                auto result = PushOffloadingQueue(
+                    MakeObjectIdentity(key, tenant_id), replica);
                 if (result) {
                     replica.inc_refcnt();
-                    shard->offloading_tasks.emplace(
+                    tenant_state.offloading_tasks.emplace(
                         key, OffloadingTask{replica.id(), now});
                     queued = true;
                 }
@@ -3572,7 +7074,7 @@ void MasterService::BatchEvict(double evict_ratio_target,
             // Any remaining MEMORY replicas with refcnt==0 are redundant copies
             // (data survives via the pinned replica → disk). Evict them now to
             // reclaim memory immediately rather than waiting another cycle.
-            return metadata.size * evict_replicas(metadata);
+            return evict_replicas(metadata, deferred_replicas);
         }
 
         // PushOffloadingQueue failed. Default (data-preserving) behavior is to
@@ -3581,102 +7083,382 @@ void MasterService::BatchEvict(double evict_ratio_target,
         // prevent silent data loss when the queue is unavailable.
         if (offload_force_evict_) {
             offload_push_failed_forced++;
-            return metadata.size * evict_replicas(metadata);
+            return evict_replicas(metadata, deferred_replicas);
         }
         return 0;
     };
 
-    // Randomly select a starting shard to avoid imbalance eviction between
-    // shards. No need to use expensive random_device here.
-    size_t start_idx = rand() % kNumShards;
-    std::shared_lock<std::shared_mutex> shared_lock(snapshot_mutex_);
+    struct EvictionResult {
+        uint64_t freed_bytes{0};
+        long evicted_objects{0};
+    };
 
-    // First pass: evict objects without soft pin and lease expired
-    for (size_t i = 0; i < kNumShards; i++) {
-        MetadataShardAccessorRW shard(this, (start_idx + i) % kNumShards);
+    auto try_evict_group_or_object =
+        [&, this](const std::string& tenant_id, const std::string& key,
+                  ObjectMetadata& metadata, MetadataShardAccessorRW& shard,
+                  TenantState& tenant_state,
+                  std::vector<std::vector<Replica>>& deferred_replicas,
+                  bool allow_soft_pinned) -> EvictionResult {
+        if (!metadata.IsGrouped()) {
+            uint64_t freed = try_evict_or_offload(
+                tenant_id, key, metadata, tenant_state, deferred_replicas);
+            return {.freed_bytes = freed, .evicted_objects = freed > 0 ? 1 : 0};
+        }
 
-        // Discard expired processing keys first so that they won't be counted
-        // in later evictions.
-        DiscardExpiredProcessingReplicas(shard, now);
+        auto group_it = tenant_state.group_members.find(metadata.group_id);
+        if (group_it == tenant_state.group_members.end()) {
+            uint64_t freed = try_evict_or_offload(
+                tenant_id, key, metadata, tenant_state, deferred_replicas);
+            return {.freed_bytes = freed, .evicted_objects = freed > 0 ? 1 : 0};
+        }
 
-        // object_count must be updated at beginning as it will be used later
-        // to compute ideal_evict_num
-        object_count += shard->metadata.size();
-
-        // To achieve evicted_count / object_count = evict_ratio_target,
-        // ideally how many object should be evicted in this shard
-        const long ideal_evict_num =
-            std::ceil(object_count * evict_ratio_target) - evicted_count;
-
-        std::vector<std::chrono::system_clock::time_point>
-            candidates;  // can be removed
-        for (auto it = shard->metadata.begin(); it != shard->metadata.end();
-             it++) {
-            // Hard-pinned objects are never evicted
-            if (it->second.IsHardPinned()) {
-                continue;
-            }
-            // Skip objects that are not expired or have incomplete replicas
-            if (!it->second.IsLeaseExpired(now) ||
-                !can_evict_replicas(it->second)) {
-                continue;
-            }
-            if (!it->second.IsSoftPinned(now)) {
-                if (ideal_evict_num > 0) {
-                    // first pass candidates
-                    candidates.push_back(it->second.lease_timeout);
-                } else {
-                    // No need to evict any object in this shard, put to
-                    // second pass candidates
-                    no_pin_objects.push_back(it->second.lease_timeout);
-                }
-            } else if (allow_evict_soft_pinned_objects_) {
-                // second pass candidates, only if
-                // allow_evict_soft_pinned_objects_ is true
-                soft_pin_objects.push_back(it->second.lease_timeout);
+        for (const auto& member_key : group_it->second) {
+            auto member_it = tenant_state.metadata.find(member_key);
+            if (member_it != tenant_state.metadata.end() &&
+                !member_it->second.IsLeaseExpired(now)) {
+                return {};
             }
         }
 
-        if (ideal_evict_num > 0 && !candidates.empty()) {
-            long evict_num = std::min(ideal_evict_num, (long)candidates.size());
-            long shard_evicted_count =
-                0;  // number of objects evicted from this shard
-            std::nth_element(candidates.begin(),
-                             candidates.begin() + (evict_num - 1),
-                             candidates.end());
-            auto target_timeout = candidates[evict_num - 1];
-            // Evict objects with lease timeout less than or equal to target.
-            auto it = shard->metadata.begin();
-            while (it != shard->metadata.end()) {
-                // Skip objects that are not allowed to be evicted in the first
-                // pass
-                if (it->second.IsHardPinned() ||
-                    !it->second.IsLeaseExpired(now) ||
-                    it->second.IsSoftPinned(now) ||
-                    !can_evict_replicas(it->second)) {
-                    ++it;
+        EvictionResult result;
+        std::vector<std::string> member_keys(group_it->second.begin(),
+                                             group_it->second.end());
+        for (const auto& member_key : member_keys) {
+            auto member_it = tenant_state.metadata.find(member_key);
+            if (member_it == tenant_state.metadata.end()) {
+                continue;
+            }
+            auto& member_metadata = member_it->second;
+            if (member_metadata.IsHardPinned() ||
+                !member_metadata.IsLeaseExpired(now) ||
+                (!allow_soft_pinned && member_metadata.IsSoftPinned(now)) ||
+                !can_evict_replicas(member_metadata)) {
+                continue;
+            }
+
+            uint64_t freed =
+                try_evict_or_offload(tenant_id, member_key, member_metadata,
+                                     tenant_state, deferred_replicas);
+            result.freed_bytes += freed;
+            if (freed > 0) {
+                result.evicted_objects++;
+            }
+            if (member_key != key && !member_metadata.IsValid()) {
+                EraseMetadata(tenant_state, member_it, tenant_id,
+                              QuotaEraseMode::kFull, &shard);
+            }
+        }
+        return result;
+    };
+
+    // Candidate carries key for safe lookup after releasing shard lock.
+    // Iterators would be invalid if the shard is modified between phases.
+    struct Candidate {
+        size_t shard_idx;
+        std::string tenant_id;
+        std::string key;
+        ObjectDataType data_type;
+        std::chrono::system_clock::time_point lease_timeout;
+        int64_t adjusted_age;
+    };
+
+    const auto rank_reference_time = now;
+    auto adjusted_age = [this, rank_reference_time](
+                            const ObjectMetadata& metadata,
+                            bool is_soft_pinned) {
+        const auto idx = static_cast<uint8_t>(metadata.data_type);
+        const double reuse_scale = object_type_reuse_scales_[idx];
+        const double weight =
+            is_soft_pinned ? object_type_soft_pin_weights_[idx] : 1.0;
+        const int64_t eviction_grace = object_type_eviction_graces_[idx];
+        const auto expired_age = std::max<int64_t>(
+            std::chrono::duration_cast<std::chrono::milliseconds>(
+                rank_reference_time - metadata.lease_timeout)
+                    .count() -
+                eviction_grace,
+            0);
+        const double score = expired_age * weight / reuse_scale;
+        if (!std::isfinite(score) ||
+            score >= static_cast<double>(std::numeric_limits<int64_t>::max())) {
+            return std::numeric_limits<int64_t>::max();
+        }
+        return static_cast<int64_t>(score);
+    };
+
+    // Randomly select a starting shard to avoid imbalance eviction between
+    // shards.
+    size_t start_idx = RandomIndex(kNumShards);
+    std::shared_lock<std::shared_mutex> shared_lock(snapshot_mutex_);
+
+    // ===== Phase 1: Parallel candidate collection =====
+    // N threads each scan a batch of shards, collecting Candidates with
+    // shard_idx + tenant_id + key for safe re-lookup in Phase 2.
+    int num_threads = std::min((int)kNumShards, 16);
+    size_t shards_per_thread = (kNumShards + num_threads - 1) / num_threads;
+
+    std::vector<std::vector<Candidate>> local_candidates(num_threads);
+    std::vector<long> local_eviction_base(num_threads, 0);
+    std::vector<long> local_object_count(num_threads, 0);
+    std::vector<std::array<long, UINT8_MAX + 1>> local_per_type_eviction_base(
+        num_threads);
+    std::vector<std::array<uint64_t, UINT8_MAX + 1>> local_per_type_used_bytes(
+        num_threads);
+    for (auto& counts : local_per_type_eviction_base) counts.fill(0);
+    for (auto& used_bytes : local_per_type_used_bytes) used_bytes.fill(0);
+    std::vector<std::vector<int64_t>> local_soft_pin(num_threads);
+
+    std::vector<std::thread> threads;
+    for (int t = 0; t < num_threads; t++) {
+        threads.emplace_back([&, t] {
+            size_t s_start = t * shards_per_thread;
+            size_t s_end = std::min(s_start + shards_per_thread, kNumShards);
+            for (size_t s = s_start; s < s_end; s++) {
+                MetadataShardAccessorRW shard(this, s);
+                DiscardExpiredProcessingReplicas(shard, now);
+
+                size_t shard_metadata_count = 0;
+                size_t shard_evictable_count = 0;
+                for (const auto& [tenant_id, tenant_state] : shard->tenants) {
+                    shard_metadata_count += tenant_state.metadata.size();
+                    for (auto it = tenant_state.metadata.begin();
+                         it != tenant_state.metadata.end(); ++it) {
+                        const ObjectDataType data_type = it->second.data_type;
+                        const auto data_type_idx =
+                            static_cast<size_t>(data_type);
+                        const uint64_t charge =
+                            CompletedMemoryQuotaCharge(it->second);
+                        if (charge > 0) {
+                            auto& used_bytes =
+                                local_per_type_used_bytes[t][data_type_idx];
+                            used_bytes = SaturatingAdd(used_bytes, charge);
+                        }
+                        if (it->second.IsHardPinned()) continue;
+                        bool has_evictable = can_evict_replicas(it->second);
+                        if (has_evictable) {
+                            shard_evictable_count++;
+                            local_per_type_eviction_base[t][data_type_idx]++;
+                        }
+                        if (!it->second.IsLeaseExpired(now) || !has_evictable)
+                            continue;
+                        if (!it->second.IsSoftPinned(now)) {
+                            local_candidates[t].push_back(
+                                {s, tenant_id, it->first, data_type,
+                                 it->second.lease_timeout,
+                                 adjusted_age(it->second,
+                                              /*is_soft_pinned=*/false)});
+                        } else if (allow_evict_soft_pinned_objects_) {
+                            local_soft_pin[t].push_back(
+                                adjusted_age(it->second,
+                                             /*is_soft_pinned=*/true));
+                        }
+                    }
+                }
+                local_object_count[t] += shard_metadata_count;
+                local_eviction_base[t] += shard_evictable_count;
+            }
+        });
+    }
+    for (auto& t : threads) t.join();
+
+    // Merge per-thread results
+    long total_eviction_base = 0;
+    for (auto v : local_eviction_base) total_eviction_base += v;
+
+    std::array<long, UINT8_MAX + 1> per_type_eviction_base{};
+    for (const auto& local_counts : local_per_type_eviction_base) {
+        for (size_t i = 0; i < per_type_eviction_base.size(); ++i) {
+            per_type_eviction_base[i] += local_counts[i];
+        }
+    }
+
+    long object_count = 0;
+    for (auto v : local_object_count) object_count += v;
+
+    std::array<uint64_t, UINT8_MAX + 1> per_type_used_bytes{};
+    for (const auto& local_used_bytes : local_per_type_used_bytes) {
+        for (size_t i = 0; i < per_type_used_bytes.size(); ++i) {
+            per_type_used_bytes[i] =
+                SaturatingAdd(per_type_used_bytes[i], local_used_bytes[i]);
+        }
+    }
+
+    std::vector<Candidate> candidates;
+    std::array<std::vector<size_t>, UINT8_MAX + 1> per_type_candidate_indices;
+    {
+        size_t total = 0;
+        for (auto& v : local_candidates) total += v.size();
+        candidates.reserve(total);
+    }
+    for (auto& v : local_candidates) {
+        for (auto& candidate : v) {
+            const size_t candidate_idx = candidates.size();
+            const auto data_type_idx = static_cast<size_t>(candidate.data_type);
+            per_type_candidate_indices[data_type_idx].push_back(candidate_idx);
+            candidates.push_back(std::move(candidate));
+        }
+    }
+
+    std::vector<int64_t> soft_pin_objects;
+    {
+        size_t total = 0;
+        for (auto& v : local_soft_pin) total += v.size();
+        soft_pin_objects.reserve(total);
+    }
+    for (auto& v : local_soft_pin) {
+        soft_pin_objects.insert(soft_pin_objects.end(),
+                                std::make_move_iterator(v.begin()),
+                                std::make_move_iterator(v.end()));
+    }
+
+    if (total_eviction_base == 0) {
+        need_mem_eviction_ = false;
+        VLOG(1) << "[EVICT-DIAG] object_count=" << object_count
+                << " eviction_base=0 (no evictable memory objects)";
+        return;
+    }
+
+    // ===== Phase 2: Serial eviction via key lookup =====
+    long evicted_count = 0;
+    uint64_t total_freed_size = 0;
+    std::vector<int64_t> no_pin_objects;
+    std::vector<std::vector<Replica>> deferred_replicas;
+
+    // First pass: evict candidates with no soft pin
+    if (!candidates.empty()) {
+        const int64_t total_mem_capacity =
+            MasterMetricManager::instance().get_total_mem_capacity();
+        if (total_mem_capacity > 0) {
+            for (size_t data_type_idx = 0;
+                 data_type_idx < per_type_used_bytes.size(); ++data_type_idx) {
+                const uint64_t type_used_bytes =
+                    per_type_used_bytes[data_type_idx];
+                const double budget_ratio =
+                    object_type_budget_ratios_[data_type_idx];
+                const double type_budget_bytes =
+                    budget_ratio * static_cast<double>(total_mem_capacity);
+                if (static_cast<double>(type_used_bytes) <= type_budget_bytes) {
                     continue;
                 }
-                if (it->second.lease_timeout <= target_timeout) {
-                    // Evict this object (or defer for offload)
-                    uint64_t freed =
-                        try_evict_or_offload(it->first, it->second, shard);
-                    total_freed_size += freed;
-                    if (it->second.IsValid() == false) {
-                        it = shard->metadata.erase(it);
-                    } else {
-                        ++it;
+
+                auto& type_candidate_indices =
+                    per_type_candidate_indices[data_type_idx];
+                const double over_budget_ratio =
+                    static_cast<double>(type_used_bytes) /
+                        static_cast<double>(total_mem_capacity) -
+                    budget_ratio;
+                const double type_evict_ratio =
+                    std::min(evict_ratio_target, over_budget_ratio);
+                long type_evict_num = std::ceil(
+                    per_type_eviction_base[data_type_idx] * type_evict_ratio);
+                type_evict_num = std::min(type_evict_num,
+                                          (long)type_candidate_indices.size());
+                if (type_evict_num <= 0) {
+                    continue;
+                }
+
+                std::nth_element(
+                    type_candidate_indices.begin(),
+                    type_candidate_indices.begin() + (type_evict_num - 1),
+                    type_candidate_indices.end(), [&](size_t a, size_t b) {
+                        return candidates[a].adjusted_age >
+                               candidates[b].adjusted_age;
+                    });
+                auto target_adjusted_age =
+                    candidates[type_candidate_indices[type_evict_num - 1]]
+                        .adjusted_age;
+
+                long evicted_this_type = 0;
+                for (const auto candidate_idx : type_candidate_indices) {
+                    const auto& c = candidates[candidate_idx];
+                    if (evicted_this_type >= type_evict_num &&
+                        c.adjusted_age < target_adjusted_age) {
+                        continue;
                     }
-                    if (freed > 0) {
-                        shard_evicted_count++;
+                    {
+                        MetadataShardAccessorRW shard(this, c.shard_idx);
+                        auto tenant_it = shard->tenants.find(c.tenant_id);
+                        if (tenant_it == shard->tenants.end()) continue;
+                        auto& tenant_state = tenant_it->second;
+                        auto it = tenant_state.metadata.find(c.key);
+                        if (it == tenant_state.metadata.end()) continue;
+                        if (!it->second.IsLeaseExpired(now) ||
+                            it->second.IsSoftPinned(now) ||
+                            !can_evict_replicas(it->second)) {
+                            continue;
+                        }
+                        auto evict_result = try_evict_group_or_object(
+                            c.tenant_id, c.key, it->second, shard, tenant_state,
+                            deferred_replicas,
+                            /*allow_soft_pinned=*/false);
+                        total_freed_size += evict_result.freed_bytes;
+                        if (!it->second.IsValid()) {
+                            EraseMetadata(tenant_state, it, c.tenant_id,
+                                          QuotaEraseMode::kFull, &shard);
+                        }
+                        if (tenant_state.Empty()) {
+                            shard->tenants.erase(tenant_it);
+                        }
+                        evicted_count += evict_result.evicted_objects;
+                        evicted_this_type += evict_result.evicted_objects;
                     }
-                } else {
-                    // second pass candidates
-                    no_pin_objects.push_back(it->second.lease_timeout);
-                    ++it;
+                    deferred_replicas.clear();
                 }
             }
-            evicted_count += shard_evicted_count;
+        }
+    }
+
+    const long remaining_evict_num =
+        std::ceil(total_eviction_base * evict_ratio_target) - evicted_count;
+    if (!candidates.empty() && remaining_evict_num > 0) {
+        long evict_num = std::min(remaining_evict_num, (long)candidates.size());
+
+        std::nth_element(candidates.begin(),
+                         candidates.begin() + (evict_num - 1), candidates.end(),
+                         [](const Candidate& a, const Candidate& b) {
+                             return a.adjusted_age > b.adjusted_age;
+                         });
+        auto target_adjusted_age = candidates[evict_num - 1].adjusted_age;
+
+        // Treat evict_num as a minimum: if re-validation skips a candidate,
+        // continue trying the next one so actual evicted count reaches
+        // evict_num. This matches the old per-shard over-eviction behavior.
+        long evicted_this_pass = 0;
+        for (auto& c : candidates) {
+            if (evicted_this_pass >= evict_num &&
+                c.adjusted_age < target_adjusted_age) {
+                no_pin_objects.push_back(c.adjusted_age);
+                continue;
+            }
+            {
+                MetadataShardAccessorRW shard(this, c.shard_idx);
+                auto tenant_it = shard->tenants.find(c.tenant_id);
+                if (tenant_it == shard->tenants.end()) continue;
+                auto& tenant_state = tenant_it->second;
+                auto it = tenant_state.metadata.find(c.key);
+                if (it == tenant_state.metadata.end()) continue;
+                // Re-validate: state may have changed since Phase 1
+                if (!it->second.IsLeaseExpired(now) ||
+                    it->second.IsSoftPinned(now) ||
+                    !can_evict_replicas(it->second)) {
+                    no_pin_objects.push_back(c.adjusted_age);
+                    continue;
+                }
+                auto evict_result = try_evict_group_or_object(
+                    c.tenant_id, c.key, it->second, shard, tenant_state,
+                    deferred_replicas,
+                    /*allow_soft_pinned=*/false);
+                total_freed_size += evict_result.freed_bytes;
+                if (!it->second.IsValid()) {
+                    EraseMetadata(tenant_state, it, c.tenant_id,
+                                  QuotaEraseMode::kFull, &shard);
+                }
+                if (tenant_state.Empty()) {
+                    shard->tenants.erase(tenant_it);
+                }
+                evicted_count += evict_result.evicted_objects;
+                evicted_this_pass += evict_result.evicted_objects;
+            }
+            deferred_replicas.clear();
         }
     }
 
@@ -3685,8 +7467,9 @@ void MasterService::BatchEvict(double evict_ratio_target,
     uint64_t released_discarded_cnt = ReleaseExpiredDiscardedReplicas(now);
 
     // The ideal number of objects to evict in the second pass
-    long target_evict_num = std::ceil(object_count * evict_ratio_lowerbound) -
-                            evicted_count - released_discarded_cnt;
+    long target_evict_num =
+        std::ceil(total_eviction_base * evict_ratio_lowerbound) -
+        evicted_count - released_discarded_cnt;
     // The actual number of objects we can evict in the second pass
     target_evict_num =
         std::min(target_evict_num,
@@ -3696,134 +7479,184 @@ void MasterService::BatchEvict(double evict_ratio_target,
     // evicted AND 2). The evicted number in the first pass is less than
     // evict_ratio_lowerbound.
     if (target_evict_num > 0) {
-        // If 1). there are enough candidates without soft pin OR 2). soft pin
-        // candidates are empty, then do second pass A. Otherwise, do second
-        // pass B. Note that the second condition is ensured implicitly by the
-        // calculation of target_evict_num.
         if (target_evict_num <= static_cast<long>(no_pin_objects.size())) {
-            // Second pass A: only evict objects without soft pin. The following
-            // code is error-prone if target_evict_num > no_pin_objects.size().
-
+            // Second pass A: only evict objects without soft pin.
             std::nth_element(no_pin_objects.begin(),
                              no_pin_objects.begin() + (target_evict_num - 1),
-                             no_pin_objects.end());
-            auto target_timeout = no_pin_objects[target_evict_num - 1];
+                             no_pin_objects.end(),
+                             [](int64_t a, int64_t b) { return a > b; });
+            auto target_adjusted_age = no_pin_objects[target_evict_num - 1];
 
-            // Evict objects with lease timeout less than or equal to target.
-            // Stop when the target is reached.
+            // Evict via key lookup — avoid full metadata traversal
             for (size_t i = 0; i < kNumShards && target_evict_num > 0; i++) {
-                MetadataShardAccessorRW shard(this,
-                                              (start_idx + i) % kNumShards);
-                auto it = shard->metadata.begin();
-                while (it != shard->metadata.end() && target_evict_num > 0) {
-                    if (!it->second.IsHardPinned() &&
-                        it->second.IsLeaseExpired(now) &&
-                        it->second.lease_timeout <= target_timeout &&
-                        !it->second.IsSoftPinned(now) &&
-                        can_evict_replicas(it->second)) {
-                        // Evict this object (or defer for offload)
-                        uint64_t freed =
-                            try_evict_or_offload(it->first, it->second, shard);
-                        total_freed_size += freed;
-                        if (it->second.IsValid() == false) {
-                            it = shard->metadata.erase(it);
+                {
+                    MetadataShardAccessorRW shard(this,
+                                                  (start_idx + i) % kNumShards);
+                    for (auto tenant_it = shard->tenants.begin();
+                         tenant_it != shard->tenants.end() &&
+                         target_evict_num > 0;) {
+                        auto& tenant_state = tenant_it->second;
+                        auto it = tenant_state.metadata.begin();
+                        while (it != tenant_state.metadata.end() &&
+                               target_evict_num > 0) {
+                            if (!it->second.IsHardPinned() &&
+                                it->second.IsLeaseExpired(now) &&
+                                adjusted_age(it->second,
+                                             /*is_soft_pinned=*/false) >=
+                                    target_adjusted_age &&
+                                !it->second.IsSoftPinned(now) &&
+                                can_evict_replicas(it->second)) {
+                                auto evict_result = try_evict_group_or_object(
+                                    tenant_it->first, it->first, it->second,
+                                    shard, tenant_state, deferred_replicas,
+                                    /*allow_soft_pinned=*/false);
+                                total_freed_size += evict_result.freed_bytes;
+                                if (!it->second.IsValid()) {
+                                    it = EraseMetadata(
+                                        tenant_state, it, tenant_it->first,
+                                        QuotaEraseMode::kFull, &shard);
+                                } else {
+                                    ++it;
+                                }
+                                evicted_count += evict_result.evicted_objects;
+                                target_evict_num -=
+                                    evict_result.evicted_objects;
+                            } else {
+                                ++it;
+                            }
+                        }
+                        if (tenant_state.Empty()) {
+                            tenant_it = shard->tenants.erase(tenant_it);
                         } else {
-                            ++it;
+                            ++tenant_it;
                         }
-                        if (freed > 0) {
-                            evicted_count++;
-                        }
-                        target_evict_num--;
-                    } else {
-                        ++it;
                     }
                 }
+                deferred_replicas.clear();
             }
         } else if (!soft_pin_objects.empty()) {
-            // Second pass B: Prioritize evicting objects without soft pin, but
-            // also allow to evict soft pinned objects. The following code is
-            // error-prone if the soft pin objects are empty.
-
+            // Second pass B: Prioritize evicting objects without soft pin,
+            // but also allow evicting soft pinned objects.
             const long soft_pin_evict_num =
                 target_evict_num - static_cast<long>(no_pin_objects.size());
-            // For soft pin objects, prioritize to evict the ones with smaller
-            // lease timeout.
             std::nth_element(
                 soft_pin_objects.begin(),
                 soft_pin_objects.begin() + (soft_pin_evict_num - 1),
-                soft_pin_objects.end());
-            auto soft_target_timeout = soft_pin_objects[soft_pin_evict_num - 1];
+                soft_pin_objects.end(),
+                [](int64_t a, int64_t b) { return a > b; });
+            auto soft_target_adjusted_age =
+                soft_pin_objects[soft_pin_evict_num - 1];
 
-            // Stop when the target is reached.
             for (size_t i = 0; i < kNumShards && target_evict_num > 0; i++) {
-                MetadataShardAccessorRW shard(this,
-                                              (start_idx + i) % kNumShards);
+                {
+                    MetadataShardAccessorRW shard(this,
+                                                  (start_idx + i) % kNumShards);
 
-                auto it = shard->metadata.begin();
-                while (it != shard->metadata.end() && target_evict_num > 0) {
-                    // Skip hard-pinned or not-yet-expired objects
-                    if (it->second.IsHardPinned() ||
-                        !it->second.IsLeaseExpired(now) ||
-                        !can_evict_replicas(it->second)) {
-                        ++it;
-                        continue;
-                    }
-                    // Evict objects with 1). no soft pin OR 2). with soft pin
-                    // and lease timeout less than or equal to target.
-                    if (!it->second.IsSoftPinned(now) ||
-                        it->second.lease_timeout <= soft_target_timeout) {
-                        // Evict this object (or defer for offload)
-                        uint64_t freed =
-                            try_evict_or_offload(it->first, it->second, shard);
-                        total_freed_size += freed;
-                        if (it->second.IsValid() == false) {
-                            it = shard->metadata.erase(it);
+                    for (auto tenant_it = shard->tenants.begin();
+                         tenant_it != shard->tenants.end() &&
+                         target_evict_num > 0;) {
+                        auto& tenant_state = tenant_it->second;
+                        auto it = tenant_state.metadata.begin();
+                        while (it != tenant_state.metadata.end() &&
+                               target_evict_num > 0) {
+                            if (it->second.IsHardPinned() ||
+                                !it->second.IsLeaseExpired(now) ||
+                                !can_evict_replicas(it->second)) {
+                                ++it;
+                                continue;
+                            }
+                            if (!it->second.IsSoftPinned(now) ||
+                                adjusted_age(it->second,
+                                             /*is_soft_pinned=*/true) >=
+                                    soft_target_adjusted_age) {
+                                auto evict_result = try_evict_group_or_object(
+                                    tenant_it->first, it->first, it->second,
+                                    shard, tenant_state, deferred_replicas,
+                                    /*allow_soft_pinned=*/true);
+                                total_freed_size += evict_result.freed_bytes;
+                                if (!it->second.IsValid()) {
+                                    it = EraseMetadata(
+                                        tenant_state, it, tenant_it->first,
+                                        QuotaEraseMode::kFull, &shard);
+                                } else {
+                                    ++it;
+                                }
+                                evicted_count += evict_result.evicted_objects;
+                                target_evict_num -=
+                                    evict_result.evicted_objects;
+                            } else {
+                                ++it;
+                            }
+                        }
+                        if (tenant_state.Empty()) {
+                            tenant_it = shard->tenants.erase(tenant_it);
                         } else {
-                            ++it;
+                            ++tenant_it;
                         }
-                        if (freed > 0) {
-                            evicted_count++;
-                        }
-                        target_evict_num--;
-                    } else {
-                        ++it;
                     }
                 }
+                deferred_replicas.clear();
             }
         } else {
-            // This should not happen.
             LOG(ERROR) << "Error in second pass eviction: target_evict_num="
                        << target_evict_num
                        << ", no_pin_objects.size()=" << no_pin_objects.size()
                        << ", soft_pin_objects.size()="
                        << soft_pin_objects.size()
                        << ", evicted_count=" << evicted_count
-                       << ", object_count=" << object_count
+                       << ", eviction_base=" << total_eviction_base
                        << ", evict_ratio_target=" << evict_ratio_target
                        << ", evict_ratio_lowerbound=" << evict_ratio_lowerbound;
         }
     }
 
-    if (evicted_count > 0 || released_discarded_cnt > 0 ||
-        offload_deferred_count > 0) {
-        // Offload-deferred counts as partial success: work was done (objects
-        // queued for disk offload), so suppress re-triggering until the next
-        // watermark breach or explicit need_eviction_ signal.
-        need_eviction_ = false;
+    if (evicted_count > 0 || released_discarded_cnt > 0) {
+        need_mem_eviction_ = false;
         MasterMetricManager::instance().inc_eviction_success(evicted_count,
                                                              total_freed_size);
+        MasterMetricManager::instance().inc_mem_eviction_success(
+            evicted_count, total_freed_size);
+    } else if (offload_deferred_count > 0) {
+        need_mem_eviction_ = false;
+        MasterMetricManager::instance().inc_eviction_success(0, 0);
+        MasterMetricManager::instance().inc_mem_eviction_success(0, 0);
     } else {
-        if (object_count == 0) {
-            // No objects to evict, no need to check again
-            need_eviction_ = false;
+        if (total_eviction_base == 0) {
+            need_mem_eviction_ = false;
         }
         MasterMetricManager::instance().inc_eviction_fail();
+        MasterMetricManager::instance().inc_mem_eviction_fail();
     }
-    VLOG(1) << "action=evict_objects" << ", evicted_count=" << evicted_count
+    VLOG(1) << "action=evict_objects"
+            << ", evicted_count=" << evicted_count
             << ", offload_deferred=" << offload_deferred_count
             << ", offload_cap_forced=" << offload_cap_forced_count
             << ", offload_push_failed_forced=" << offload_push_failed_forced
-            << ", total_freed_size=" << total_freed_size;
+            << ", total_freed_size=" << total_freed_size
+            << ", eviction_base=" << total_eviction_base
+            << ", actual_evict_ratio="
+            << (total_eviction_base > 0
+                    ? (double)evicted_count / total_eviction_base
+                    : 0.0)
+            << ", target_evict_ratio=" << evict_ratio_target;
+    VLOG(1) << "[EVICT-DIAG] object_count=" << object_count
+            << " disk_object_count=" << (object_count - total_eviction_base)
+            << " eviction_base=" << total_eviction_base << " disk_ratio="
+            << (object_count > 0
+                    ? (double)(object_count - total_eviction_base) /
+                          object_count
+                    : 0.0)
+            << " ideal_evict_num_inflated="
+            << (long)std::ceil(object_count * evict_ratio_target)
+            << " ideal_evict_num_correct="
+            << (long)std::ceil(total_eviction_base * evict_ratio_target);
+    LOG(INFO) << "[EVICT-RESULT] evicted_count=" << evicted_count
+              << ", eviction_base=" << total_eviction_base
+              << ", actual_evict_ratio="
+              << (total_eviction_base > 0
+                      ? (double)evicted_count / total_eviction_base
+                      : 0.0)
+              << ", target_evict_ratio=" << evict_ratio_target;
     if (offload_on_evict_ && evicted_count == 0 && offload_deferred_count > 0) {
         LOG(WARNING) << "[EVICT] No memory freed this cycle; "
                      << offload_deferred_count
@@ -3843,6 +7676,98 @@ void MasterService::BatchEvict(double evict_ratio_target,
     }
 }
 
+void MasterService::NoFBatchEvict(double evict_ratio_target,
+                                  double evict_ratio_lowerbound) {
+    if (evict_ratio_target < evict_ratio_lowerbound) {
+        LOG(ERROR) << "nof_evict_ratio_target=" << evict_ratio_target
+                   << ", nof_evict_ratio_lowerbound=" << evict_ratio_lowerbound
+                   << ", error=invalid_params";
+        evict_ratio_lowerbound = evict_ratio_target;
+    }
+
+    auto now = std::chrono::system_clock::now();
+    long evicted_count = 0;
+    long object_count = 0;
+    uint64_t total_freed_size = 0;
+
+    size_t start_idx = RandomIndex(metadata_shards_.size());
+    for (size_t i = 0; i < metadata_shards_.size(); i++) {
+        MetadataShardAccessorRW shard(
+            this, (start_idx + i) % metadata_shards_.size());
+        DiscardExpiredProcessingReplicas(shard, now);
+        for (const auto& [tenant_id, tenant_state] : shard->tenants) {
+            object_count += tenant_state.metadata.size();
+        }
+
+        const long ideal_evict_num =
+            std::ceil(object_count * evict_ratio_target) - evicted_count;
+        if (ideal_evict_num <= 0) {
+            continue;
+        }
+
+        long shard_evicted_count = 0;
+        for (auto tenant_it = shard->tenants.begin();
+             tenant_it != shard->tenants.end() &&
+             shard_evicted_count < ideal_evict_num;) {
+            auto& tenant_state = tenant_it->second;
+            for (auto it = tenant_state.metadata.begin();
+                 it != tenant_state.metadata.end() &&
+                 shard_evicted_count < ideal_evict_num;) {
+                auto& metadata = it->second;
+                if (metadata.IsHardPinned() || !metadata.IsLeaseExpired(now) ||
+                    metadata.IsSoftPinned(now)) {
+                    ++it;
+                    continue;
+                }
+
+                const size_t erased =
+                    metadata.EraseReplicas([](const Replica& replica) {
+                        return replica.is_nof_replica() &&
+                               replica.is_completed() &&
+                               replica.get_refcnt() == 0;
+                    });
+                if (erased == 0) {
+                    ++it;
+                    continue;
+                }
+
+                total_freed_size += metadata.size * erased;
+                shard_evicted_count++;
+                if (!metadata.IsValid()) {
+                    it = EraseMetadata(tenant_state, it, tenant_it->first,
+                                       QuotaEraseMode::kFull, &shard);
+                } else {
+                    ++it;
+                }
+            }
+            if (tenant_state.Empty()) {
+                tenant_it = shard->tenants.erase(tenant_it);
+            } else {
+                ++tenant_it;
+            }
+        }
+        evicted_count += shard_evicted_count;
+    }
+
+    if (evicted_count > 0) {
+        need_nof_eviction_ = false;
+        MasterMetricManager::instance().inc_eviction_success(evicted_count,
+                                                             total_freed_size);
+        MasterMetricManager::instance().inc_nof_eviction_success(
+            evicted_count, total_freed_size);
+    } else {
+        if (object_count == 0) {
+            need_nof_eviction_ = false;
+        }
+        MasterMetricManager::instance().inc_eviction_fail();
+        MasterMetricManager::instance().inc_nof_eviction_fail();
+    }
+
+    VLOG(1) << "action=evict_nof_replicas"
+            << ", evicted_count=" << evicted_count
+            << ", total_freed_size=" << total_freed_size;
+}
+
 void MasterService::ClientMonitorFunc() {
     std::unordered_map<UUID, std::chrono::steady_clock::time_point,
                        boost::hash<UUID>>
@@ -3873,6 +7798,15 @@ void MasterService::ClientMonitorFunc() {
 
         // Update the client status to NEED_REMOUNT
         if (!expired_clients.empty()) {
+            // Notify graceful unmount scheduler to drop pending records
+            // for expired clients. The actual unmount is handled below.
+            for (auto& cid : expired_clients) {
+                graceful_unmount_scheduler_.RemoveIf(
+                    [&cid](const GracefulUnmountDeadlineRecord& record) {
+                        return record.client_id == cid;
+                    });
+            }
+
             // Record which segments are unmounted, will be used in the commit
             // phase.
             std::vector<UUID> unmount_segments;
@@ -3894,6 +7828,7 @@ void MasterService::ClientMonitorFunc() {
                 ScopedSegmentAccess segment_access =
                     segment_manager_.getSegmentAccess();
                 for (auto& client_id : expired_clients) {
+                    // mounted mem segemtns of this expired client
                     std::vector<Segment> segments;
                     segment_access.GetClientSegments(client_id, segments);
                     for (auto& seg : segments) {
@@ -3910,7 +7845,7 @@ void MasterService::ClientMonitorFunc() {
                                        << ", segment_name=" << seg.name
                                        << ", "
                                           "error=prepare_unmount_expired_"
-                                          "segment_failed";
+                                          "mem_segment_failed";
                         }
                     }
                 }
@@ -3933,12 +7868,15 @@ void MasterService::ClientMonitorFunc() {
                         unmount_segments[i], client_ids[i], dec_capacities[i]);
                     LOG(INFO) << "client_id=" << client_ids[i]
                               << ", segment_name=" << segment_names[i]
-                              << ", action=unmount_expired_segment";
+                              << ", action=unmount_expired_mem_segment";
+                    // Clean up HTTP metadata if enabled
+                    cleanupHttpMetadata(segment_names[i]);
                 }
                 for (auto& client_id : expired_clients) {
                     segment_access.UnmountLocalDiskSegment(client_id);
                 }
             }
+            RecomputeTenantEffectiveQuotas();
         }
 
         std::this_thread::sleep_for(
@@ -3946,6 +7884,242 @@ void MasterService::ClientMonitorFunc() {
     }
 }
 
+bool MasterService::ProbeNoFSegment(const std::string& te_endpoint,
+                                    std::string* error_reason) {
+#ifndef USE_NOF
+    if (error_reason) {
+        *error_reason = "nof_pool_disabled";
+    }
+    return false;
+#else
+    NoFProbeFn probe_fn;
+    {
+        std::lock_guard<std::mutex> lock(nof_probe_fn_mutex_);
+        probe_fn = nof_probe_fn_;
+    }
+    if (!probe_fn) {
+        if (error_reason) {
+            *error_reason = "probe_not_configured";
+        }
+        return false;
+    }
+    return probe_fn(
+        te_endpoint,
+        static_cast<uint32_t>(nof_heartbeat_probe_timeout_ms_.count()),
+        error_reason);
+#endif
+}
+
+bool MasterService::TryUnmountNoFSegmentByHeartbeat(
+    const MountedNoFSegmentSnapshot& snapshot,
+    const std::string& error_reason) {
+    size_t metrics_dec_capacity = 0;
+    {
+        auto nof_segment_access = nof_segment_manager_.getNoFSegmentAccess();
+        ErrorCode err = nof_segment_access.PrepareUnmountSegment(
+            snapshot.segment_id, metrics_dec_capacity);
+        if (err == ErrorCode::SEGMENT_NOT_FOUND ||
+            err == ErrorCode::UNAVAILABLE_IN_CURRENT_STATUS) {
+            std::lock_guard<std::mutex> lock(nof_heartbeat_mutex_);
+            nof_heartbeat_states_.erase(snapshot.segment_id);
+            VLOG(1) << "segment_id=" << snapshot.segment_id
+                    << ", action=skip_nof_heartbeat_unmount"
+                    << ", reason=" << toString(err);
+            return false;
+        }
+        if (err != ErrorCode::OK) {
+            LOG(ERROR) << "segment_id=" << snapshot.segment_id
+                       << ", segment_name=" << snapshot.segment.name
+                       << ", error=prepare_unmount_nof_segment_by_"
+                          "heartbeat_failed"
+                       << ", reason=" << err;
+            return false;
+        }
+    }
+
+    ClearInvalidHandles();
+
+    {
+        auto nof_segment_access = nof_segment_manager_.getNoFSegmentAccess();
+        ErrorCode err = nof_segment_access.CommitUnmountSegment(
+            snapshot.segment_id, snapshot.client_id, metrics_dec_capacity);
+        if (err != ErrorCode::OK && err != ErrorCode::SEGMENT_NOT_FOUND) {
+            LOG(ERROR) << "segment_id=" << snapshot.segment_id
+                       << ", segment_name=" << snapshot.segment.name
+                       << ", error=commit_unmount_nof_segment_by_"
+                          "heartbeat_failed"
+                       << ", reason=" << err;
+            return false;
+        }
+    }
+
+    {
+        std::lock_guard<std::mutex> lock(nof_heartbeat_mutex_);
+        nof_heartbeat_states_.erase(snapshot.segment_id);
+    }
+    MasterMetricManager::instance()
+        .inc_nof_segments_unmounted_by_heartbeat_total();
+    LOG(INFO) << "segment_id=" << snapshot.segment_id
+              << ", client_id=" << snapshot.client_id
+              << ", segment_name=" << snapshot.segment.name
+              << ", endpoint=" << snapshot.segment.te_endpoint
+              << ", action=unmount_nof_segment_by_heartbeat"
+              << ", last_error_reason=" << error_reason;
+    return true;
+}
+
+void MasterService::NofHeartbeatThreadFunc() {
+    size_t next_probe_index = 0;
+    while (nof_heartbeat_running_) {
+        auto now = std::chrono::steady_clock::now();
+        std::vector<MountedNoFSegmentSnapshot> mounted_segments;
+        nof_segment_manager_.GetMountedSegmentsSnapshot(mounted_segments);
+
+        std::vector<MountedNoFSegmentSnapshot> ok_segments;
+        ok_segments.reserve(mounted_segments.size());
+        for (const auto& snapshot : mounted_segments) {
+            if (snapshot.status == SegmentStatus::OK) {
+                ok_segments.push_back(snapshot);
+            }
+        }
+
+        std::optional<MountedNoFSegmentSnapshot> probe_target;
+        {
+            std::lock_guard<std::mutex> lock(nof_heartbeat_mutex_);
+            std::unordered_set<UUID, boost::hash<UUID>> live_segment_ids;
+            live_segment_ids.reserve(ok_segments.size());
+
+            const auto interval_ms =
+                std::chrono::duration_cast<std::chrono::milliseconds>(
+                    nof_heartbeat_interval_sec_);
+            for (size_t i = 0; i < ok_segments.size(); ++i) {
+                const auto& snapshot = ok_segments[i];
+                live_segment_ids.insert(snapshot.segment_id);
+                auto [it, inserted] =
+                    nof_heartbeat_states_.try_emplace(snapshot.segment_id);
+                auto& state = it->second;
+                state.owner_client_id = snapshot.client_id;
+                state.segment_name = snapshot.segment.name;
+                state.te_endpoint = snapshot.segment.te_endpoint;
+                if (inserted) {
+                    int64_t spread_ms = 0;
+                    if (!ok_segments.empty()) {
+                        spread_ms = static_cast<int64_t>(
+                            (interval_ms.count() * i) / ok_segments.size());
+                    }
+                    state.last_success_at = now;
+                    state.next_probe_at = now + nof_heartbeat_interval_sec_ +
+                                          std::chrono::milliseconds(spread_ms);
+                }
+            }
+
+            for (auto it = nof_heartbeat_states_.begin();
+                 it != nof_heartbeat_states_.end();) {
+                if (!live_segment_ids.contains(it->first)) {
+                    it = nof_heartbeat_states_.erase(it);
+                } else {
+                    ++it;
+                }
+            }
+
+            if (!ok_segments.empty()) {
+                next_probe_index %= ok_segments.size();
+                for (size_t offset = 0; offset < ok_segments.size(); ++offset) {
+                    const auto& candidate =
+                        ok_segments[(next_probe_index + offset) %
+                                    ok_segments.size()];
+                    auto state_it =
+                        nof_heartbeat_states_.find(candidate.segment_id);
+                    if (state_it == nof_heartbeat_states_.end()) {
+                        continue;
+                    }
+                    if (state_it->second.next_probe_at <= now) {
+                        probe_target = candidate;
+                        next_probe_index = (next_probe_index + offset + 1) %
+                                           ok_segments.size();
+                        break;
+                    }
+                }
+            }
+        }
+
+        if (!probe_target.has_value()) {
+            std::this_thread::sleep_for(
+                std::chrono::milliseconds(kNoFHeartbeatThreadSleepMs));
+            continue;
+        }
+
+        auto probe_start = std::chrono::steady_clock::now();
+        std::string error_reason;
+        bool probe_success =
+            ProbeNoFSegment(probe_target->segment.te_endpoint, &error_reason);
+        auto latency_ms = std::chrono::duration_cast<std::chrono::milliseconds>(
+                              std::chrono::steady_clock::now() - probe_start)
+                              .count();
+        MasterMetricManager::instance().observe_nof_heartbeat_probe_latency_ms(
+            latency_ms);
+
+        if (probe_success) {
+            MasterMetricManager::instance().inc_nof_heartbeat_success_total();
+            auto success_time = std::chrono::steady_clock::now();
+            {
+                std::lock_guard<std::mutex> lock(nof_heartbeat_mutex_);
+                auto it = nof_heartbeat_states_.find(probe_target->segment_id);
+                if (it != nof_heartbeat_states_.end()) {
+                    it->second.consecutive_failures = 0;
+                    it->second.last_success_at = success_time;
+                    it->second.last_error_reason.clear();
+                    it->second.next_probe_at =
+                        success_time + nof_heartbeat_interval_sec_;
+                }
+            }
+            VLOG(1) << "segment_id=" << probe_target->segment_id
+                    << ", segment_name=" << probe_target->segment.name
+                    << ", endpoint=" << probe_target->segment.te_endpoint
+                    << ", action=nof_heartbeat_success"
+                    << ", latency_ms=" << latency_ms;
+            continue;
+        }
+
+        MasterMetricManager::instance().inc_nof_heartbeat_failure_total();
+        if (error_reason == "completion_timeout") {
+            MasterMetricManager::instance().inc_nof_heartbeat_timeout_total();
+        }
+
+        bool should_unmount = false;
+        uint32_t failure_count = 0;
+        auto failure_time = std::chrono::steady_clock::now();
+        auto alive_timeout =
+            nof_heartbeat_interval_sec_ *
+            static_cast<int64_t>(nof_heartbeat_failures_threshold_);
+        {
+            std::lock_guard<std::mutex> lock(nof_heartbeat_mutex_);
+            auto it = nof_heartbeat_states_.find(probe_target->segment_id);
+            if (it != nof_heartbeat_states_.end()) {
+                it->second.consecutive_failures++;
+                failure_count = it->second.consecutive_failures;
+                it->second.last_error_reason = error_reason;
+                it->second.next_probe_at =
+                    failure_time + nof_heartbeat_interval_sec_;
+                should_unmount =
+                    failure_time - it->second.last_success_at >= alive_timeout;
+            }
+        }
+
+        LOG(WARNING) << "segment_id=" << probe_target->segment_id
+                     << ", segment_name=" << probe_target->segment.name
+                     << ", endpoint=" << probe_target->segment.te_endpoint
+                     << ", action=nof_heartbeat_failure"
+                     << ", failure_count=" << failure_count
+                     << ", latency_ms=" << latency_ms
+                     << ", reason=" << error_reason;
+
+        if (should_unmount) {
+            TryUnmountNoFSegmentByHeartbeat(*probe_target, error_reason);
+        }
+    }
+}
+
 tl::expected<std::vector<uint8_t>, SerializationError>
 MasterService::MetadataSerializer::Serialize() {
     msgpack::sbuffer sbuf;
@@ -3958,10 +8132,17 @@ MasterService::MetadataSerializer::Serialize() {
     // 1. Serialize metadata shards
     packer.pack("shards");
 
-    // First count non-empty shards
+    // First count shards that have actual metadata entries.
+    // A shard may have empty tenants left after eviction erased all
+    // metadata but didn't clean up the tenant map; using metadata_count
+    // (not tenants.empty()) ensures the count matches the skip logic below.
     size_t valid_shards = 0;
     for (size_t i = 0; i < kNumShards; ++i) {
-        if (!service_->metadata_shards_[i].metadata.empty()) {
+        size_t metadata_count = 0;
+        for (const auto& [tid, ts] : service_->metadata_shards_[i].tenants) {
+            metadata_count += ts.metadata.size();
+        }
+        if (metadata_count > 0) {
             valid_shards++;
         }
     }
@@ -3973,8 +8154,16 @@ MasterService::MetadataSerializer::Serialize() {
     for (size_t shard_idx = 0; shard_idx < kNumShards; ++shard_idx) {
         const auto& shard = service_->metadata_shards_[shard_idx];
 
-        // Skip if shard is empty
-        if (shard.metadata.empty()) {
+        // Skip shards with no actual metadata entries.
+        // A shard may have empty tenants left after eviction erased all
+        // metadata but didn't clean up the tenant map; serializing those
+        // would produce an entry that deserialization never recreates,
+        // breaking the snapshot round-trip comparison.
+        size_t metadata_count = 0;
+        for (const auto& [tid, ts] : shard.tenants) {
+            metadata_count += ts.metadata.size();
+        }
+        if (metadata_count == 0) {
             continue;
         }
 
@@ -4145,13 +8334,19 @@ MasterService::MetadataSerializer::Deserialize(
     auto next_id = replica_next_id_obj->as<uint64_t>();
     Replica::next_id_.store(next_id);
     LOG(INFO) << "Restored Replica::next_id_ to " << next_id;
-
+    service_->RebuildGroupRoutingIndex();
     return {};
 }
 
 void MasterService::MetadataSerializer::Reset() {
     for (auto& shard : service_->metadata_shards_) {
-        shard.metadata.clear();
+        shard.tenants.clear();
+    }
+    {
+        std::unique_lock<std::shared_mutex> lock(
+            service_->group_routing_mutex_);
+        service_->object_group_ids_.clear();
+        service_->groups_needing_lease_refresh_.clear();
     }
     {
         std::lock_guard lock(service_->discarded_replicas_mutex_);
@@ -4168,29 +8363,46 @@ MasterService::MetadataSerializer::SerializeShard(const MetadataShard& shard,
 
     // Serialize metadata
     packer.pack("metadata");
-    packer.pack_array(shard.metadata.size());
-
-    // Sort keys to ensure consistent serialization order.
-    // NOTE: sort may be slow for large shards.
-    std::vector<std::string> sorted_keys;
-    sorted_keys.reserve(shard.metadata.size());
-    for (const auto& [key, metadata] : shard.metadata) {
-        sorted_keys.push_back(key);
+    size_t metadata_count = 0;
+    for (const auto& [tenant_id, tenant_state] : shard.tenants) {
+        metadata_count += tenant_state.metadata.size();
     }
-    std::sort(sorted_keys.begin(), sorted_keys.end());
+    packer.pack_array(metadata_count);
 
-    for (const auto& key : sorted_keys) {
-        const auto& metadata = shard.metadata.at(key);
-        // Each metadata item format: [key, metadata_object]
-        packer.pack_array(2);
-        packer.pack(key);
+    // Sort tenant/key pairs to ensure consistent serialization order.
+    // NOTE: sort may be slow for large shards.
+    struct SortedEntry {
+        std::string tenant_id;
+        std::string key;
+        const ObjectMetadata* metadata;
+    };
+    std::vector<SortedEntry> sorted_entries;
+    sorted_entries.reserve(metadata_count);
+    for (const auto& [tenant_id, tenant_state] : shard.tenants) {
+        for (const auto& [key, metadata] : tenant_state.metadata) {
+            sorted_entries.push_back({tenant_id, key, &metadata});
+        }
+    }
+    std::sort(sorted_entries.begin(), sorted_entries.end(),
+              [](const SortedEntry& lhs, const SortedEntry& rhs) {
+                  if (lhs.tenant_id != rhs.tenant_id) {
+                      return lhs.tenant_id < rhs.tenant_id;
+                  }
+                  return lhs.key < rhs.key;
+              });
 
-        auto result = SerializeMetadata(metadata, packer);
+    for (const auto& entry : sorted_entries) {
+        // Each metadata item format: [tenant_id, key, metadata_object].
+        packer.pack_array(3);
+        packer.pack(entry.tenant_id);
+        packer.pack(entry.key);
+
+        auto result = SerializeMetadata(*entry.metadata, packer);
         if (!result) {
             return tl::make_unexpected(SerializationError(
                 result.error().code,
                 fmt::format("Failed to serialize metadata for key '{}': {}",
-                            key, result.error().message)));
+                            entry.key, result.error().message)));
         }
     }
 
@@ -4219,7 +8431,7 @@ MasterService::MetadataSerializer::DeserializeShard(const msgpack::object& obj,
     }
 
     // Clear existing data
-    shard.metadata.clear();
+    shard.tenants.clear();
 
     // Deserialize metadata
     if (metadata_array == nullptr ||
@@ -4229,21 +8441,33 @@ MasterService::MetadataSerializer::DeserializeShard(const msgpack::object& obj,
                                "Missing or invalid 'metadata' field in shard"));
     }
 
-    shard.metadata.reserve(metadata_array->via.array.size);
+    shard.tenants.reserve(metadata_array->via.array.size);
 
     for (uint32_t j = 0; j < metadata_array->via.array.size; ++j) {
         const msgpack::object& item = metadata_array->via.array.ptr[j];
 
-        if (item.type != msgpack::type::ARRAY || item.via.array.size != 2) {
+        if (item.type != msgpack::type::ARRAY ||
+            (item.via.array.size != 2 && item.via.array.size != 3)) {
             return tl::make_unexpected(SerializationError(
                 ErrorCode::DESERIALIZE_FAIL,
-                "Invalid metadata item format: expected [key, metadata]"));
+                "Invalid metadata item format: expected [key, metadata] or "
+                "[tenant_id, key, metadata]"));
         }
 
-        std::string key = item.via.array.ptr[0].as<std::string>();
-        const msgpack::object& value_obj = item.via.array.ptr[1];
+        std::string tenant_id = "default";
+        std::string key;
+        const msgpack::object* value_obj = nullptr;
+        if (item.via.array.size == 2) {
+            key = item.via.array.ptr[0].as<std::string>();
+            value_obj = &item.via.array.ptr[1];
+        } else {
+            tenant_id =
+                NormalizeTenantId(item.via.array.ptr[0].as<std::string>());
+            key = item.via.array.ptr[1].as<std::string>();
+            value_obj = &item.via.array.ptr[2];
+        }
 
-        auto metadata_result = DeserializeMetadata(value_obj);
+        auto metadata_result = DeserializeMetadata(*value_obj);
         if (!metadata_result) {
             LOG(ERROR) << "Failed to deserialize metadata for key: " << key
                        << ": " << metadata_result.error().message;
@@ -4251,16 +8475,26 @@ MasterService::MetadataSerializer::DeserializeShard(const msgpack::object& obj,
         }
 
         auto metadata_ptr = std::move(metadata_result.value());
-        auto [it, inserted] = shard.metadata.emplace(
+        auto& tenant_state = shard.tenants[tenant_id];
+        const std::string user_key = key;
+        auto [it, inserted] = tenant_state.metadata.emplace(
             std::piecewise_construct, std::forward_as_tuple(std::move(key)),
             std::forward_as_tuple(
                 metadata_ptr->client_id, metadata_ptr->put_start_time,
                 metadata_ptr->size, metadata_ptr->PopReplicas(),
                 metadata_ptr->soft_pin_timeout.has_value(),
-                metadata_ptr->IsHardPinned()));
+                metadata_ptr->IsHardPinned(), metadata_ptr->data_type,
+                metadata_ptr->group_id, tenant_id, user_key));
 
         it->second.lease_timeout = metadata_ptr->lease_timeout;
         it->second.soft_pin_timeout = metadata_ptr->soft_pin_timeout;
+
+        // Recompute disk_object_count for restored metadata
+        if (it->second.HasReplica([](const Replica& r) {
+                return r.is_local_disk_replica() && r.is_completed();
+            })) {
+            shard.disk_object_count++;
+        }
     }
 
     return {};
@@ -4272,12 +8506,12 @@ MasterService::MetadataSerializer::SerializeMetadata(
     MsgpackPacker& packer) const {
     // Pack ObjectMetadata using array structure for efficiency
     // Format: [client_id, put_start_time, size, lease_timeout,
-    // has_soft_pin_timeout, soft_pin_timeout, replicas_count, replicas...,
-    // hard_pinned]
+    // has_soft_pin_timeout, soft_pin_timeout, replicas_count, data_type,
+    // replicas..., hard_pinned, group_id]
 
-    size_t array_size = 8;  // client_id, put_start_time, size, lease_timeout,
-                            // has_soft_pin_timeout, soft_pin_timeout,
-                            // replicas_count + hard_pinned
+    size_t array_size = 10;  // client_id, put_start_time, size, lease_timeout,
+                             // has_soft_pin_timeout, soft_pin_timeout,
+                             // replicas_count, data_type, hard_pinned, group_id
     array_size += metadata.CountReplicas();  // One element per replica
     packer.pack_array(array_size);
 
@@ -4317,6 +8551,9 @@ MasterService::MetadataSerializer::SerializeMetadata(
     // Serialize replicas count
     packer.pack(static_cast<uint32_t>(metadata.CountReplicas()));
 
+    // Serialize data_type
+    packer.pack(static_cast<uint8_t>(metadata.data_type));
+
     // Serialize replicas
     for (const auto& replica : metadata.GetAllReplicas()) {
         auto result = Serializer<Replica>::serialize(
@@ -4327,6 +8564,7 @@ MasterService::MetadataSerializer::SerializeMetadata(
     }
 
     packer.pack(metadata.IsHardPinned());
+    packer.pack(metadata.group_id);
 
     return {};
 }
@@ -4343,7 +8581,6 @@ MasterService::MetadataSerializer::DeserializeMetadata(
 
     // Need at least 7 elements: client_id, put_start_time, size, lease_timeout,
     // has_soft_pin_timeout, soft_pin_timeout, replicas_count
-    // (8th element = hard_pinned is optional for backward compat)
     if (obj.via.array.size < 7) {
         return tl::unexpected(SerializationError(
             ErrorCode::DESERIALIZE_FAIL,
@@ -4376,20 +8613,45 @@ MasterService::MetadataSerializer::DeserializeMetadata(
     // Deserialize replicas count
     uint32_t replicas_count = array[index++].as<uint32_t>();
 
-    // Array size: 7 + replicas_count (old format) or 8 + replicas_count (new
-    // format with hard_pinned)
-    if (obj.via.array.size != 7 + replicas_count &&
-        obj.via.array.size != 8 + replicas_count) {
+    // Format detection (decode optional fields by type for back-compat):
+    //   v1: 7 + replicas_count, no optional fields
+    //   v2: 8 + replicas_count, either data_type or hard_pinned
+    //   v3: 9 + replicas_count, data_type + hard_pinned or hard_pinned +
+    //   group_id v4: 10 + replicas_count, data_type + hard_pinned + group_id
+    // 64-bit arithmetic keeps an attacker-controlled near-UINT32_MAX
+    // replicas_count from wrapping the bounds and slipping an out-of-bounds
+    // index past the size check.
+    constexpr uint64_t kBaseFieldCount = 7;
+    constexpr uint64_t kMaxOptionalFieldCount = 3;
+    const uint64_t total_elements = obj.via.array.size;
+    const uint64_t min_elements = kBaseFieldCount + replicas_count;
+    if (total_elements < min_elements ||
+        total_elements > min_elements + kMaxOptionalFieldCount) {
         return tl::unexpected(SerializationError(
             ErrorCode::DESERIALIZE_FAIL,
             "deserialize ObjectMetadata array size mismatch"));
     }
 
+    ObjectDataType data_type = ObjectDataType::UNKNOWN;
+    if (index < total_elements &&
+        array[index].type == msgpack::type::POSITIVE_INTEGER) {
+        data_type = static_cast<ObjectDataType>(array[index++].as<uint8_t>());
+    }
+
     // Deserialize replicas
     std::vector<Replica> replicas;
     replicas.reserve(replicas_count);
 
     for (uint32_t i = 0; i < replicas_count; i++) {
+        // Defensive bound: the data_type skip above can consume a slot the
+        // size check counted on, so a crafted entry whose first post-count
+        // field looks like a data_type could otherwise read past the array.
+        // Mirrors the standby reader in catalog_backed_snapshot_provider.cpp.
+        if (index >= total_elements) {
+            return tl::unexpected(
+                SerializationError(ErrorCode::DESERIALIZE_FAIL,
+                                   "deserialize ObjectMetadata truncated"));
+        }
         auto result = Serializer<Replica>::deserialize(
             array[index++], service_->segment_manager_.getView());
         if (!result) {
@@ -4400,17 +8662,24 @@ MasterService::MetadataSerializer::DeserializeMetadata(
 
     // Deserialize hard_pinned (if present, otherwise default to false)
     bool is_hard_pinned = false;
-    if (index < obj.via.array.size) {
+    if (index < obj.via.array.size &&
+        array[index].type == msgpack::type::BOOLEAN) {
         is_hard_pinned = array[index++].as<bool>();
     }
 
+    std::string group_id;
+    if (index < obj.via.array.size && array[index].type == msgpack::type::STR) {
+        group_id = array[index++].as<std::string>();
+    }
+
     // Create ObjectMetadata instance
     bool enable_soft_pin = has_soft_pin_timeout;
     auto metadata = std::make_unique<ObjectMetadata>(
         client_id,
         std::chrono::system_clock::time_point(
             std::chrono::milliseconds(put_start_time_timestamp)),
-        size, std::move(replicas), enable_soft_pin, is_hard_pinned);
+        size, std::move(replicas), enable_soft_pin, is_hard_pinned, data_type,
+        group_id);
     metadata->lease_timeout = std::chrono::system_clock::time_point(
         std::chrono::milliseconds(lease_timestamp));
 
@@ -4442,13 +8711,19 @@ std::string MasterService::FormatTimestamp(
 }
 
 tl::expected<UUID, ErrorCode> MasterService::CreateCopyTask(
-    const std::string& key, const std::vector<std::string>& targets) {
+    const std::string& key, const std::string& tenant_id,
+    const std::vector<std::string>& targets) {
+    auto normalized_tenant_result = NormalizeTenantIdForWrite(tenant_id);
+    if (!normalized_tenant_result) {
+        return tl::make_unexpected(normalized_tenant_result.error());
+    }
+    const ObjectIdentity object_id{normalized_tenant_result.value(), key};
     std::shared_lock<std::shared_mutex> shared_lock(snapshot_mutex_);
     if (targets.empty()) {
         LOG(ERROR) << "key=" << key << ", error=empty_targets";
         return tl::make_unexpected(ErrorCode::INVALID_PARAMS);
     }
-    MetadataAccessorRO accessor(this, key);
+    MetadataAccessorRO accessor(this, object_id);
     if (!accessor.Exists()) {
         VLOG(1) << "key=" << key << ", info=object_not_found";
         return tl::make_unexpected(ErrorCode::OBJECT_NOT_FOUND);
@@ -4491,16 +8766,22 @@ tl::expected<UUID, ErrorCode> MasterService::CreateCopyTask(
     }
     return task_manager_.get_write_access()
         .submit_task_typed<TaskType::REPLICA_COPY>(
-            select_client, {.key = key,
+            select_client, {.tenant_id = object_id.tenant_id,
+                            .key = object_id.user_key,
                             .source = selected_source_segment,
                             .targets = targets});
 }
 
 tl::expected<UUID, ErrorCode> MasterService::CreateMoveTask(
-    const std::string& key, const std::string& source,
-    const std::string& target) {
+    const std::string& key, const std::string& tenant_id,
+    const std::string& source, const std::string& target) {
+    auto normalized_tenant_result = NormalizeTenantIdForWrite(tenant_id);
+    if (!normalized_tenant_result) {
+        return tl::make_unexpected(normalized_tenant_result.error());
+    }
+    const ObjectIdentity object_id{normalized_tenant_result.value(), key};
     std::shared_lock<std::shared_mutex> shared_lock(snapshot_mutex_);
-    MetadataAccessorRO accessor(this, key);
+    MetadataAccessorRO accessor(this, object_id);
     if (!accessor.Exists()) {
         VLOG(1) << "key=" << key << ", info=object_not_found";
         return tl::make_unexpected(ErrorCode::OBJECT_NOT_FOUND);
@@ -4546,7 +8827,10 @@ tl::expected<UUID, ErrorCode> MasterService::CreateMoveTask(
 
     return task_manager_.get_write_access()
         .submit_task_typed<TaskType::REPLICA_MOVE>(
-            select_client, {.key = key, .source = source, .target = target});
+            select_client, {.tenant_id = object_id.tenant_id,
+                            .key = object_id.user_key,
+                            .source = source,
+                            .target = target});
 }
 
 tl::expected<QueryTaskResponse, ErrorCode> MasterService::QueryTask(
@@ -4753,8 +9037,11 @@ tl::expected<void, ErrorCode> MasterService::CancelDrainJob(
 }
 
 std::string MasterService::MakeDrainUnitKey(
-    const std::string& key, const std::string& source_segment) const {
-    return std::to_string(key.size()) + ":" + key + ":" + source_segment;
+    const std::string& tenant_id, const std::string& key,
+    const std::string& source_segment) const {
+    const auto normalized_tenant = NormalizeTenantId(tenant_id);
+    return std::to_string(normalized_tenant.size()) + ":" + normalized_tenant +
+           ":" + std::to_string(key.size()) + ":" + key + ":" + source_segment;
 }
 
 std::optional<std::string> MasterService::SelectDrainTargetForKey(
@@ -4849,6 +9136,7 @@ void MasterService::ScheduleDrainJobTasks(DrainJob& job) {
     }
 
     struct DrainPlan {
+        std::string tenant_id;
         std::string key;
         std::string source_segment;
         std::string target_segment;
@@ -4869,40 +9157,45 @@ void MasterService::ScheduleDrainJobTasks(DrainJob& job) {
         std::shared_lock<std::shared_mutex> shared_lock(snapshot_mutex_);
         for (size_t i = 0; i < kNumShards; ++i) {
             MetadataShardAccessorRO shard(this, i);
-            for (const auto& [key, metadata] : shard->metadata) {
-                for (const auto& source_segment : job.request.segments) {
-                    const auto unit_key = MakeDrainUnitKey(key, source_segment);
-                    if (job.completed_unit_keys.contains(unit_key) ||
-                        active_unit_keys.contains(unit_key) ||
-                        job.terminal_failed_unit_keys.contains(unit_key)) {
-                        continue;
-                    }
+            for (const auto& [tenant_id, tenant_state] : shard->tenants) {
+                for (const auto& [key, metadata] : tenant_state.metadata) {
+                    for (const auto& source_segment : job.request.segments) {
+                        const auto unit_key =
+                            MakeDrainUnitKey(tenant_id, key, source_segment);
+                        if (job.completed_unit_keys.contains(unit_key) ||
+                            active_unit_keys.contains(unit_key) ||
+                            job.terminal_failed_unit_keys.contains(unit_key)) {
+                            continue;
+                        }
 
-                    const auto replica_segments =
-                        metadata.GetReplicaSegmentNames();
-                    if (std::find(replica_segments.begin(),
-                                  replica_segments.end(),
-                                  source_segment) == replica_segments.end()) {
-                        continue;
-                    }
+                        const auto replica_segments =
+                            metadata.GetReplicaSegmentNames();
+                        if (std::find(replica_segments.begin(),
+                                      replica_segments.end(), source_segment) ==
+                            replica_segments.end()) {
+                            continue;
+                        }
 
-                    if (metadata.IsHardPinned() || !metadata.IsLeaseExpired() ||
-                        !metadata.AllReplicas(&Replica::fn_is_completed) ||
-                        shard->replication_tasks.contains(key)) {
-                        blocked_unit_keys.insert(unit_key);
-                        continue;
-                    }
+                        if (metadata.IsHardPinned() ||
+                            !metadata.IsLeaseExpired() ||
+                            !metadata.AllReplicas(&Replica::fn_is_completed) ||
+                            tenant_state.replication_tasks.contains(key)) {
+                            blocked_unit_keys.insert(unit_key);
+                            continue;
+                        }
 
-                    auto target = SelectDrainTargetForKey(
-                        metadata, source_segment, job.request.target_segments);
-                    if (!target.has_value()) {
-                        blocked_unit_keys.insert(unit_key);
-                        continue;
-                    }
+                        auto target = SelectDrainTargetForKey(
+                            metadata, source_segment,
+                            job.request.target_segments);
+                        if (!target.has_value()) {
+                            blocked_unit_keys.insert(unit_key);
+                            continue;
+                        }
 
-                    if (plans.size() < slots) {
-                        plans.push_back({key, source_segment, *target,
-                                         metadata.size, unit_key});
+                        if (plans.size() < slots) {
+                            plans.push_back({tenant_id, key, source_segment,
+                                             *target, metadata.size, unit_key});
+                        }
                     }
                 }
             }
@@ -4912,11 +9205,12 @@ void MasterService::ScheduleDrainJobTasks(DrainJob& job) {
     job.blocked_units = blocked_unit_keys.size();
 
     for (const auto& plan : plans) {
-        auto task_id =
-            CreateMoveTask(plan.key, plan.source_segment, plan.target_segment);
+        auto task_id = CreateMoveTask(plan.key, plan.tenant_id,
+                                      plan.source_segment, plan.target_segment);
         if (task_id.has_value()) {
             ActiveDrainTask active_task;
             active_task.task_id = task_id.value();
+            active_task.tenant_id = plan.tenant_id;
             active_task.key = plan.key;
             active_task.source_segment = plan.source_segment;
             active_task.target_segment = plan.target_segment;
@@ -4954,15 +9248,18 @@ bool MasterService::MaybeCompleteDrainJob(DrainJob& job) {
         std::shared_lock<std::shared_mutex> shared_lock(snapshot_mutex_);
         for (size_t i = 0; i < kNumShards; ++i) {
             MetadataShardAccessorRO shard(this, i);
-            for (const auto& [key, metadata] : shard->metadata) {
-                const auto replica_segments = metadata.GetReplicaSegmentNames();
-                for (const auto& source_segment : job.request.segments) {
-                    if (std::find(replica_segments.begin(),
-                                  replica_segments.end(),
-                                  source_segment) != replica_segments.end()) {
-                        remaining_segments.insert(source_segment);
-                        remaining_unit_keys.insert(
-                            MakeDrainUnitKey(key, source_segment));
+            for (const auto& [tenant_id, tenant_state] : shard->tenants) {
+                for (const auto& [key, metadata] : tenant_state.metadata) {
+                    const auto replica_segments =
+                        metadata.GetReplicaSegmentNames();
+                    for (const auto& source_segment : job.request.segments) {
+                        if (std::find(replica_segments.begin(),
+                                      replica_segments.end(), source_segment) !=
+                            replica_segments.end()) {
+                            remaining_segments.insert(source_segment);
+                            remaining_unit_keys.insert(MakeDrainUnitKey(
+                                tenant_id, key, source_segment));
+                        }
                     }
                 }
             }
@@ -5164,4 +9461,132 @@ MasterService::MetadataSerializer::DeserializeDiscardedReplicas(
     return {};
 }
 
+void MasterService::setHttpMetadataServer(HttpMetadataServer* server) {
+    http_metadata_server_ = server;
+    if (server) {
+        LOG(INFO) << "HTTP metadata cleanup on client timeout: enabled "
+                     "(co-located metadata server)";
+    }
+}
+
+void MasterService::setHttpMetadataRemoteUrl(
+    const std::string& metadata_connstring) {
+#ifdef USE_HTTP
+    // Only http(s) is supported; guard the scheme to avoid
+    // MetadataStoragePlugin::Create()'s LOG(FATAL) on other backends.
+    if (metadata_connstring.rfind("http://", 0) == 0 ||
+        metadata_connstring.rfind("https://", 0) == 0) {
+        try {
+            http_metadata_remote_ =
+                MetadataStoragePlugin::Create(metadata_connstring);
+            LOG(INFO) << "HTTP metadata cleanup on client timeout: enabled "
+                         "(remote metadata server "
+                      << metadata_connstring << ")";
+            // Start async cleanup worker now that http_metadata_remote_ is
+            // ready
+            http_metadata_cleanup_running_ = true;
+            http_metadata_cleanup_thread_ = std::thread(
+                &MasterService::HttpMetadataCleanupThreadFunc, this);
+            LOG(INFO) << "HTTP metadata cleanup worker thread started";
+        } catch (const std::exception& e) {
+            LOG(WARNING) << "Failed to initialize remote HTTP metadata client "
+                            "for "
+                         << metadata_connstring << ": " << e.what()
+                         << ". Metadata cleanup on timeout disabled.";
+            http_metadata_remote_.reset();
+        }
+        return;
+    }
+    LOG(WARNING) << "enable_metadata_cleanup_on_timeout is set but the "
+                    "configured metadata server '"
+                 << metadata_connstring
+                 << "' is not an HTTP endpoint; remote cleanup currently "
+                    "supports only http(s). Metadata cleanup on timeout "
+                    "disabled.";
+#else
+    (void)metadata_connstring;
+    LOG(WARNING) << "enable_metadata_cleanup_on_timeout is set but this build "
+                    "has no HTTP metadata support (USE_HTTP=OFF); metadata "
+                    "cleanup on timeout disabled.";
+#endif
+}
+
+void MasterService::cleanupHttpMetadata(const std::string& segment_name) {
+    // Co-located: remove in-process, safe to run inline (no network I/O).
+    if (http_metadata_server_) {
+        const std::string ram_key =
+            http_metadata_prefix_ + "ram/" + segment_name;
+        const std::string rpc_key =
+            http_metadata_prefix_ + "rpc_meta/" + segment_name;
+        bool ram_removed = http_metadata_server_->removeKey(ram_key);
+        bool rpc_removed = http_metadata_server_->removeKey(rpc_key);
+        LOG(INFO) << "Cleaned up HTTP metadata for segment: " << segment_name
+                  << ", ram_key_removed=" << ram_removed
+                  << ", rpc_key_removed=" << rpc_removed;
+        return;
+    }
+
+    // Separately-deployed: enqueue for async cleanup so a slow/unreachable
+    // server never blocks the client monitor thread.
+    if (http_metadata_remote_) {
+        {
+            std::lock_guard<std::mutex> lk(http_metadata_cleanup_mutex_);
+            http_metadata_cleanup_queue_.push_back(segment_name);
+        }
+        http_metadata_cleanup_cv_.notify_one();
+        return;
+    }
+
+    // Neither configured: cleanup is disabled, nothing to do.
+}
+
+void MasterService::HttpMetadataCleanupThreadFunc() {
+    LOG(INFO) << "HTTP metadata cleanup worker started";
+    while (http_metadata_cleanup_running_) {
+        std::vector<std::string> batch;
+        {
+            std::unique_lock<std::mutex> lk(http_metadata_cleanup_mutex_);
+            http_metadata_cleanup_cv_.wait(lk, [&] {
+                return !http_metadata_cleanup_queue_.empty() ||
+                       !http_metadata_cleanup_running_.load();
+            });
+            if (!http_metadata_cleanup_running_ &&
+                http_metadata_cleanup_queue_.empty()) {
+                break;
+            }
+            batch.swap(http_metadata_cleanup_queue_);
+        }
+
+        for (const auto& segment_name : batch) {
+            const std::string ram_key =
+                http_metadata_prefix_ + "ram/" + segment_name;
+            const std::string rpc_key =
+                http_metadata_prefix_ + "rpc_meta/" + segment_name;
+
+            // Each key attempted independently so one failure does not
+            // prevent cleanup of the other.
+            bool ram_removed = false;
+            bool rpc_removed = false;
+            try {
+                ram_removed = http_metadata_remote_->remove(ram_key);
+            } catch (const std::exception& e) {
+                LOG(WARNING)
+                    << "Remote HTTP metadata cleanup failed for ram_key: "
+                    << ram_key << ": " << e.what();
+            }
+            try {
+                rpc_removed = http_metadata_remote_->remove(rpc_key);
+            } catch (const std::exception& e) {
+                LOG(WARNING)
+                    << "Remote HTTP metadata cleanup failed for rpc_key: "
+                    << rpc_key << ": " << e.what();
+            }
+            LOG(INFO) << "Cleaned up remote HTTP metadata for segment: "
+                      << segment_name << ", ram_key_removed=" << ram_removed
+                      << ", rpc_key_removed=" << rpc_removed;
+        }
+    }
+    LOG(INFO) << "HTTP metadata cleanup worker stopped";
+}
+
 }  // namespace mooncake
diff --git a/mooncake-store/src/memory_alloc.cpp b/mooncake-store/src/memory_alloc.cpp
new file mode 100644
index 00000000..74947a4d
--- /dev/null
+++ b/mooncake-store/src/memory_alloc.cpp
@@ -0,0 +1,19 @@
+#include "memory_alloc.h"
+
+#ifdef USE_NOF
+#include "spdk/spdk_wrapper.h"
+#endif
+
+void *hugepage_memory_alloc(size_t size) {
+#ifndef USE_NOF
+    return nullptr;
+#else
+    return mooncake::SpdkWrapper::GetInstance().Alloc(size, 0x1000, -1);
+#endif
+}
+
+void hugepage_memory_free(void *ptr) {
+#ifdef USE_NOF
+    mooncake::SpdkWrapper::GetInstance().Free(ptr);
+#endif
+}
diff --git a/mooncake-store/src/mmap_arena.cpp b/mooncake-store/src/mmap_arena.cpp
new file mode 100644
index 00000000..a0fc41e2
--- /dev/null
+++ b/mooncake-store/src/mmap_arena.cpp
@@ -0,0 +1,274 @@
+// Copyright 2026 KVCache.AI
+// Simple arena allocator implementation
+
+#include "mmap_arena.h"
+#include "utils.h"
+#include <sys/mman.h>
+#include <glog/logging.h>
+#include <cerrno>
+#include <cstring>
+#include <algorithm>
+
+namespace mooncake {
+
+// Safe alignment with overflow detection
+// Returns false if overflow would occur, true on success
+static inline bool safe_align_up(size_t size, size_t alignment,
+                                 size_t* result) {
+    if (size == 0) {
+        *result = 0;
+        return true;
+    }
+    if (alignment == 0 || (alignment & (alignment - 1)) != 0) {
+        // Must be power of 2
+        return false;
+    }
+    // Check if size + alignment - 1 would overflow
+    if (size > SIZE_MAX - alignment + 1) {
+        return false;  // Overflow would occur
+    }
+    *result = (size + alignment - 1) & ~(alignment - 1);
+    return true;
+}
+
+MmapArena::MmapArena()
+    : pool_base_(nullptr),
+      pool_size_(0),
+      alignment_(kMinAlignment),
+      alloc_cursor_(0),
+      peak_reserved_(0),
+      num_allocations_(0),
+      num_failed_allocs_(0) {}
+
+MmapArena::~MmapArena() {
+    void* pool_base = pool_base_.load(std::memory_order_acquire);
+    if (pool_base != nullptr) {
+        size_t pool_size = pool_size_.load(std::memory_order_acquire);
+        if (munmap(pool_base, pool_size) != 0) {
+            LOG(ERROR) << "Arena munmap failed: " << strerror(errno);
+        }
+        pool_base_.store(nullptr, std::memory_order_release);
+    }
+}
+
+bool MmapArena::initialize(size_t pool_size, size_t alignment,
+                           bool allow_regular_page_fallback) {
+    // Mutex serializes concurrent initialize() calls so that exactly one
+    // thread performs the mmap and publishes the pool.  This avoids the
+    // metadata-overwrite race that existed with the old CAS approach
+    // (losing threads could clobber alignment_/pool_size_ before CAS).
+    std::lock_guard<std::mutex> lock(init_mutex_);
+
+    if (pool_size == 0) {
+        LOG(ERROR) << "Arena pool size must be > 0";
+        return false;
+    }
+
+    if (alignment != 0 && (alignment & (alignment - 1)) != 0) {
+        LOG(ERROR) << "Arena alignment must be a power of 2, got " << alignment;
+        return false;
+    }
+
+    if (pool_base_.load(std::memory_order_acquire) != nullptr) {
+        LOG(WARNING) << "Arena already initialized";
+        return false;
+    }
+
+    const size_t actual_alignment = std::max(alignment, kMinAlignment);
+
+    // Align pool size to 2MB for huge pages with overflow protection
+    size_t aligned_pool_size;
+    if (!safe_align_up(pool_size, SZ_2MB, &aligned_pool_size)) {
+        LOG(ERROR) << "Arena pool size overflow: requested=" << pool_size;
+        return false;
+    }
+
+    // Allocate pool with mmap.  Use MAP_POPULATE so that all huge pages
+    // are faulted upfront.  Without it, lazy page-faults during GPU DMA
+    // (e.g. Mooncake transfer engine ↔ CUDA) cause cudaErrorIllegalAddress
+    // on some platforms (observed on H100 80GB).
+    int flags = MAP_PRIVATE | MAP_ANONYMOUS | MAP_POPULATE;
+
+// Try huge pages for better TLB performance
+#ifdef MAP_HUGETLB
+    flags |= MAP_HUGETLB;
+#endif
+
+    void* pool_base =
+        mmap(nullptr, aligned_pool_size, PROT_READ | PROT_WRITE, flags, -1, 0);
+
+    if (pool_base == MAP_FAILED) {
+        const int mmap_errno = errno;
+        if (!allow_regular_page_fallback) {
+            LOG(ERROR) << "Arena hugepage mmap failed for pool_size="
+                       << aligned_pool_size << " bytes"
+                       << ", errno=" << mmap_errno << " ("
+                       << strerror(mmap_errno)
+                       << "); regular-page fallback disabled because "
+                          "MC_STORE_USE_HUGEPAGE explicitly requested "
+                          "hugepages";
+            return false;
+        }
+
+        // Retry without huge pages
+        flags &= ~MAP_HUGETLB;
+        LOG(WARNING) << "Arena hugepage mmap failed for pool_size="
+                     << aligned_pool_size << " bytes"
+                     << ", errno=" << mmap_errno << " (" << strerror(mmap_errno)
+                     << "); retrying without huge pages";
+        pool_base = mmap(nullptr, aligned_pool_size, PROT_READ | PROT_WRITE,
+                         flags, -1, 0);
+
+        if (pool_base == MAP_FAILED) {
+            LOG(ERROR) << "Arena mmap failed: size=" << aligned_pool_size
+                       << ", errno=" << errno << " (" << strerror(errno) << ")";
+            return false;
+        }
+        LOG(WARNING) << "Arena initialized without huge pages for pool_size="
+                     << aligned_pool_size
+                     << " bytes; MAP_POPULATE will prefault regular pages";
+    } else {
+        LOG(INFO) << "Arena initialized with huge pages";
+    }
+
+    // Prevent child processes from inheriting this mapping on fork().
+    // With MAP_POPULATE, all pages are physically backed; fork() would
+    // create a massive copy-on-write mapping (up to 64GB) in the child,
+    // potentially doubling memory usage and triggering OOM killer.
+    // Child processes don't need the parent's DMA buffers — they exec()
+    // or initialize their own arena.
+    if (madvise(pool_base, aligned_pool_size, MADV_DONTFORK) != 0) {
+        LOG(WARNING) << "madvise(MADV_DONTFORK) failed: " << strerror(errno)
+                     << " for pool_size=" << aligned_pool_size
+                     << " bytes (fork safety not guaranteed)";
+    }
+
+    // Store metadata BEFORE publishing pool_base_.
+    // The release store on pool_base_ ensures these are visible to any
+    // thread that loads pool_base_ with acquire in allocate().
+    alignment_.store(actual_alignment, std::memory_order_relaxed);
+    pool_size_.store(aligned_pool_size, std::memory_order_relaxed);
+    pool_base_.store(pool_base, std::memory_order_release);
+
+    LOG(INFO) << "Arena initialized: " << (aligned_pool_size / BYTES_PER_GIB)
+              << " GiB, alignment=" << actual_alignment << " bytes";
+
+    return true;
+}
+
+void* MmapArena::allocate(size_t size, size_t alignment) {
+    void* pool_base = pool_base_.load(std::memory_order_acquire);
+    if (pool_base == nullptr) {
+        LOG(ERROR) << "Arena not initialized";
+        return nullptr;
+    }
+
+    if (size == 0) {
+        return nullptr;
+    }
+
+    // Effective alignment: max of arena default and caller's request.
+    // This honors the caller's alignment contract without weakening
+    // the arena's minimum guarantee.
+    // pool_base_ was acquired above; this relaxed load observes the
+    // alignment value published before that release-store in initialize().
+    size_t base_alignment = alignment_.load(std::memory_order_relaxed);
+    size_t effective_alignment = std::max(base_alignment, alignment);
+
+    // Align allocation size with overflow check
+    size_t aligned_size;
+    if (!safe_align_up(size, effective_alignment, &aligned_size)) {
+        num_failed_allocs_.fetch_add(1, std::memory_order_relaxed);
+        LOG(ERROR) << "Arena allocation size overflow: size=" << size
+                   << ", alignment=" << effective_alignment;
+        return nullptr;
+    }
+
+    size_t pool_size = pool_size_.load(std::memory_order_acquire);
+
+    // CAS loop: Reserve aligned space atomically with bounds check.
+    // We align the OFFSET (not just the size) so the returned pointer
+    // honours the caller's alignment contract even when the cursor sits
+    // at a non-aligned position from a previous smaller-alignment alloc.
+    size_t aligned_offset;
+    size_t next;
+    while (true) {
+        size_t raw = alloc_cursor_.load(std::memory_order_relaxed);
+
+        // Align the offset up to effective_alignment
+        if (!safe_align_up(raw, effective_alignment, &aligned_offset)) {
+            num_failed_allocs_.fetch_add(1, std::memory_order_relaxed);
+            LOG(ERROR) << "Arena offset alignment overflow: raw=" << raw
+                       << ", alignment=" << effective_alignment;
+            return nullptr;
+        }
+
+        next = aligned_offset + aligned_size;
+
+        // Check for overflow (next wrapped) and OOM BEFORE modifying cursor
+        if (next < aligned_offset || next > pool_size) {
+            num_failed_allocs_.fetch_add(1, std::memory_order_relaxed);
+            LOG(ERROR) << "Arena OOM: requested=" << size
+                       << ", aligned_size=" << aligned_size
+                       << ", aligned_offset=" << aligned_offset
+                       << ", pool_size=" << pool_size;
+            return nullptr;
+        }
+
+        // Try to reserve [aligned_offset, next) atomically.
+        // CAS from raw (not aligned_offset) — another thread may have
+        // bumped the cursor since we loaded it.
+        if (alloc_cursor_.compare_exchange_weak(raw, next,
+                                                std::memory_order_relaxed,
+                                                std::memory_order_relaxed)) {
+            break;  // Success - space reserved
+        }
+        // CAS failed, retry with new raw value
+    }
+
+    // Space successfully reserved at [aligned_offset, next)
+    num_allocations_.fetch_add(1, std::memory_order_relaxed);
+
+    // Update peak statistics using `next` (the actual end of reservation,
+    // including any alignment padding before aligned_offset)
+    size_t old_peak = peak_reserved_.load(std::memory_order_relaxed);
+    while (next > old_peak && !peak_reserved_.compare_exchange_weak(
+                                  old_peak, next, std::memory_order_relaxed,
+                                  std::memory_order_relaxed)) {
+        // CAS loop for peak tracking
+    }
+
+    void* ptr = static_cast<char*>(pool_base) + aligned_offset;
+
+    VLOG(2) << "[ARENA] Allocated: size=" << size
+            << ", aligned_size=" << aligned_size
+            << ", aligned_offset=" << aligned_offset << ", ptr=" << ptr
+            << ", utilization=" << (100.0 * next / pool_size) << "%";
+
+    return ptr;
+}
+
+MmapArena::Stats MmapArena::getStats() const {
+    Stats stats;
+    stats.pool_size = pool_size_.load(std::memory_order_relaxed);
+    stats.reserved_bytes = alloc_cursor_.load(std::memory_order_relaxed);
+    stats.peak_reserved_bytes = peak_reserved_.load(std::memory_order_relaxed);
+    stats.num_allocations = num_allocations_.load(std::memory_order_relaxed);
+    stats.num_failed_allocs =
+        num_failed_allocs_.load(std::memory_order_relaxed);
+    return stats;
+}
+
+bool MmapArena::owns(const void* ptr) const {
+    void* pool_base = pool_base_.load(std::memory_order_acquire);
+    if (!ptr || !pool_base) {
+        return false;
+    }
+
+    size_t pool_size = pool_size_.load(std::memory_order_acquire);
+    uintptr_t addr = reinterpret_cast<uintptr_t>(ptr);
+    uintptr_t base = reinterpret_cast<uintptr_t>(pool_base);
+    return addr >= base && addr < base + pool_size;
+}
+
+}  // namespace mooncake
diff --git a/mooncake-store/src/real_client.cpp b/mooncake-store/src/real_client.cpp
index f2f25b99..6439bf61 100644
--- a/mooncake-store/src/real_client.cpp
+++ b/mooncake-store/src/real_client.cpp
@@ -2,6 +2,8 @@
 #include <sys/socket.h>
 #include <sys/un.h>
 #include <unistd.h>
+#include <fcntl.h>
+#include <sys/mman.h>
 #include <numa.h>
 #include <numaif.h>
 #include <pthread.h>
@@ -13,21 +15,27 @@
 #include <cstdlib>  // for atexit
 #include <algorithm>
 #include <cctype>
+#include <functional>
 #include <limits>
 #include <optional>
 #include <vector>
 
 #include "real_client.h"
 #include "client_buffer.hpp"
+#include "common.h"
 #include "config.h"
 #include "mutex.h"
 #include "types.h"
 #include "utils.h"
 #include "rpc_types.h"
 #include "file_storage.h"
+#include "gpu_staging_utils.h"
 #include "default_config.h"
 #include "shm_helper.h"
 #include "memory_location.h"
+#ifdef USE_NOF
+#include "spdk/spdk_wrapper.h"
+#endif
 #ifdef USE_ASCEND_DIRECT
 #include "acl/acl_rt.h"
 #include "transport/ascend_transport/ascend_direct_transport/context_manager.h"
@@ -51,6 +59,25 @@ bool checkAcl(aclError result, const char *message) {
     }
     return true;
 }
+
+tl::expected<void, ErrorCode> set_context_if_needed(const std::string &protocol,
+                                                    int32_t physical_device_id,
+                                                    const char *action) {
+    if (protocol != "ascend" || !globalConfig().ascend_agent_mode) {
+        return {};
+    }
+    if (physical_device_id == kInvalidPhysicalDeviceId) {
+        LOG(ERROR) << "Missing physical device id for " << action;
+        return tl::make_unexpected(ErrorCode::INVALID_PARAMS);
+    }
+    if (!ContextManager::getInstance().setCurrentContextByPhysicalId(
+            physical_device_id)) {
+        LOG(ERROR) << "Failed to set current context for physical device "
+                   << physical_device_id << " during " << action;
+        return tl::make_unexpected(ErrorCode::INVALID_PARAMS);
+    }
+    return {};
+}
 #endif
 
 struct PreparedRangedReadRequest {
@@ -235,6 +262,76 @@ void fill_ranged_read_results_with_error(
         }
     }
 }
+
+// Scatter host (CPU) memory to a destination that may be GPU or host.
+// Returns tl::expected<void, ErrorCode> for use in functions returning
+// tl::expected<int64_t, ErrorCode>.
+inline tl::expected<void, ErrorCode> scatter_host_to_maybe_device(
+    void *dst, const void *src, size_t size, const std::string &context) {
+    int device_id = -1;
+    if (gpu_staging::IsDevicePointer(dst, &device_id)) {
+        gpu_staging::SetDevice(device_id);
+        if (!gpu_staging::CopyHostToDevice(dst, src, size)) {
+            LOG(ERROR) << "H2D copy failed: " << context;
+            return tl::unexpected(ErrorCode::TRANSFER_FAIL);
+        }
+    } else if (gpu_staging::IsHostPointer(dst)) {
+        memcpy(dst, src, size);
+    } else {
+        LOG(ERROR) << "Unknown memory type for dst buffer: " << context;
+        return tl::unexpected(ErrorCode::INVALID_PARAMS);
+    }
+    return {};
+}
+
+// Select the best replica from a list: prefer local MEMORY, then any
+// MEMORY, then LOCAL_DISK, then DISK.  Master may return replicas in any
+// order, so we always scan.
+inline const Replica::Descriptor *SelectBestReplica(
+    const std::vector<Replica::Descriptor> &replicas,
+    const std::unordered_set<std::string> &local_endpoints) {
+    const Replica::Descriptor *first_memory = nullptr;
+    const Replica::Descriptor *first_nof = nullptr;
+    for (const auto &r : replicas) {
+        if (r.status != ReplicaStatus::COMPLETE) continue;
+        if (r.is_memory_replica()) {
+            if (local_endpoints.count(
+                    r.get_memory_descriptor()
+                        .buffer_descriptor.transport_endpoint_)) {
+                return &r;  // local MEMORY — best case
+            }
+            if (!first_memory) first_memory = &r;
+        } else if (r.is_nof_replica()) {
+            if (local_endpoints.count(
+                    r.get_nof_descriptor()
+                        .buffer_descriptor.transport_endpoint_)) {
+                return &r;  // local NOF_SSD — also good
+            }
+            if (!first_nof) first_nof = &r;
+        }
+    }
+    if (first_memory) return first_memory;
+    if (first_nof) return first_nof;
+
+    const Replica::Descriptor *best = nullptr;
+    for (const auto &r : replicas) {
+        if (r.status != ReplicaStatus::COMPLETE) continue;
+        if (r.is_local_disk_replica()) {
+            best = &r;  // LOCAL_DISK always overrides DISK
+        } else if (r.is_disk_replica() && !best) {
+            best = &r;
+        }
+    }
+    return best;
+}
+
+// Build a QueryResult containing only the chosen replica so that
+// Client::Get / Client::BatchGet (which internally call
+// FindFirstCompleteReplica) cannot pick a different replica type.
+inline QueryResult FilterQueryResult(const QueryResult &qr,
+                                     const Replica::Descriptor &replica) {
+    return QueryResult({replica}, qr.lease_timeout);
+}
 }  // namespace
 
 PyClient::~PyClient() {}
@@ -536,12 +633,11 @@ tl::expected<void, ErrorCode> RealClient::setup_internal(
     const std::shared_ptr<TransferEngine> &transfer_engine,
     const std::string &ipc_socket_path, int local_rpc_port,
     bool enable_ssd_offload, bool start_offload_rpc_server,
-    const std::string &ssd_offload_path) {
+    const std::string &ssd_offload_path, const std::string &tenant_id) {
     this->protocol = protocol;
     this->ipc_socket_path_ = ipc_socket_path;
-    const bool should_use_hugepage = use_hugepage_ &&
-                                     this->protocol != "ascend" &&
-                                     this->protocol != "ubshmem";
+    const bool should_use_hugepage =
+        use_hugepage_ && this->protocol != "ubshmem";
 #ifdef USE_ASCEND_DIRECT
     if (protocol == "ascend" && globalConfig().ascend_agent_mode) {
         auto ascend_setup = setup_ascend_internal(local_buffer_size);
@@ -553,9 +649,17 @@ tl::expected<void, ErrorCode> RealClient::setup_internal(
     }
 #endif
 
+#ifdef USE_NOF
+    if (!SpdkWrapper::GetInstance().InitializeEnv()) {
+        LOG(ERROR) << "spdk env init fail";
+        return tl::unexpected(ErrorCode::INTERNAL_ERROR);
+    }
+#endif
+
     std::optional<std::string> device_name =
-        (rdma_devices.empty() ? std::nullopt
-                              : std::make_optional(rdma_devices));
+        ((rdma_devices.empty() || rdma_devices == "auto-discovery")
+             ? std::nullopt
+             : std::make_optional(rdma_devices));
 
     // Validate required parameters
     if (local_hostname.empty()) {
@@ -565,17 +669,17 @@ tl::expected<void, ErrorCode> RealClient::setup_internal(
 
     // Check if hostname already contains a port
     const std::string &hostname = local_hostname;
-    size_t colon_pos = hostname.find(':');
-    bool user_specified_port = (colon_pos != std::string::npos);
+    bool user_specified_port = hasExplicitPort(hostname);
 
     if (user_specified_port) {
         // User specified port, no retry needed
         this->local_hostname = local_hostname;
-        this->local_rpc_addr =
-            hostname.substr(0, colon_pos + 1) + std::to_string(local_rpc_port);
+        this->local_rpc_addr = buildHostNameWithPort(
+            getHostNameWithoutPort(hostname), local_rpc_port);
         auto client_opt = mooncake::Client::Create(
             this->local_hostname, metadata_server, protocol, device_name,
-            master_server_addr, transfer_engine, {{"client_mode", "real"}});
+            master_server_addr, transfer_engine, {{"client_mode", "real"}},
+            tenant_id);
         if (!client_opt) {
             LOG(ERROR) << "Failed to create client";
             return tl::unexpected(ErrorCode::INVALID_PARAMS);
@@ -585,11 +689,17 @@ tl::expected<void, ErrorCode> RealClient::setup_internal(
         // Auto port binding with retry on metadata registration failure
         const int kMaxRetries =
             GetEnvOr<int>("MC_STORE_CLIENT_SETUP_RETRIES", 20);
+        const int rawMinPort = GetEnvOr<int>("MC_STORE_CLIENT_MIN_PORT", 12300);
+        const int rawMaxPort = GetEnvOr<int>("MC_STORE_CLIENT_MAX_PORT", 14300);
+        constexpr int kDefaultMinPort = 12300;
+        constexpr int kDefaultMaxPort = 14300;
+        auto [minPort, maxPort] = ValidatePortRange(
+            rawMinPort, rawMaxPort, kDefaultMinPort, kDefaultMaxPort);
         bool success = false;
 
         for (int retry = 0; retry < kMaxRetries; ++retry) {
             // Create port binder to hold a port
-            port_binder_ = std::make_unique<AutoPortBinder>();
+            port_binder_ = std::make_unique<AutoPortBinder>(minPort, maxPort);
             int port = port_binder_->getPort();
             if (port < 0) {
                 LOG(WARNING) << "Failed to bind available port, retry "
@@ -599,12 +709,13 @@ tl::expected<void, ErrorCode> RealClient::setup_internal(
                 continue;
             }
 
-            this->local_hostname = hostname + ":" + std::to_string(port);
+            this->local_hostname = buildHostNameWithPort(hostname, port);
             this->local_rpc_addr =
-                hostname + ":" + std::to_string(local_rpc_port);
+                buildHostNameWithPort(hostname, local_rpc_port);
             auto client_opt = mooncake::Client::Create(
                 this->local_hostname, metadata_server, protocol, device_name,
-                master_server_addr, transfer_engine, {{"client_mode", "real"}});
+                master_server_addr, transfer_engine, {{"client_mode", "real"}},
+                tenant_id);
             if (client_opt) {
                 client_ = *client_opt;
                 success = true;
@@ -633,8 +744,13 @@ tl::expected<void, ErrorCode> RealClient::setup_internal(
     // fail in some rdma implementations.
     // Dummy Client can create shm and share it with Real Client, so Real Client
     // can create client buffer allocator on the shared memory later.
+    bool use_spdk_dma_for_client_buffer = false;
+#ifdef USE_NOF
+    use_spdk_dma_for_client_buffer = true;
+#endif
     client_buffer_allocator_ = ClientBufferAllocator::create(
-        local_buffer_size, this->protocol, should_use_hugepage);
+        local_buffer_size, this->protocol, should_use_hugepage,
+        use_spdk_dma_for_client_buffer);
     if (local_buffer_size > 0 && protocol != "cxl") {
         LOG(INFO) << "Registering local memory: " << local_buffer_size
                   << " bytes";
@@ -646,6 +762,14 @@ tl::expected<void, ErrorCode> RealClient::setup_internal(
                        << toString(result.error());
             return tl::unexpected(result.error());
         }
+        {
+            std::unique_lock<std::shared_mutex> lock(registered_buffer_mutex_);
+            local_buffer_region_ = WritableBufferRegion{
+                .base = client_buffer_allocator_->getBase(),
+                .size = local_buffer_size,
+                .offset = 0,
+            };
+        }
     } else {
         LOG(INFO) << "Local buffer size is 0, skip registering local memory";
     }
@@ -737,6 +861,9 @@ tl::expected<void, ErrorCode> RealClient::setup_internal(
             if (this->protocol == "ascend" || this->protocol == "ubshmem") {
                 ascend_segment_ptrs_.emplace_back(
                     ptr, AscendSegmentDeleter{this->protocol});
+            } else if (this->protocol == "ub") {
+                ub_segment_ptrs_.emplace_back(ptr,
+                                              UbSegmentDeleter{mapped_size});
             } else if (!seg_numa_nodes.empty() || should_use_hugepage) {
                 // NUMA-segmented or hugepage: track as mmap allocation for
                 // munmap cleanup
@@ -787,13 +914,8 @@ tl::expected<void, ErrorCode> RealClient::setup_internal(
         LOG(INFO) << "Offload RPC server started on port " << offload_rpc_port_;
 
         // Build local_rpc_addr from hostname + auto-allocated port
-        std::string rpc_host = this->local_hostname;
-        auto pos = rpc_host.find(':');
-        if (pos != std::string::npos) {
-            rpc_host = rpc_host.substr(0, pos);
-        }
-        this->local_rpc_addr =
-            rpc_host + ":" + std::to_string(offload_rpc_port_);
+        this->local_rpc_addr = buildHostNameWithPort(
+            getHostNameWithoutPort(this->local_hostname), offload_rpc_port_);
     }
     if (enable_ssd_offload) {
         auto file_storage_config = FileStorageConfig::FromEnvironment();
@@ -828,11 +950,12 @@ int RealClient::setup_real(
     const std::string &master_server_addr,
     const std::shared_ptr<TransferEngine> &transfer_engine,
     const std::string &ipc_socket_path, bool enable_ssd_offload,
-    const std::string &ssd_offload_path) {
+    const std::string &ssd_offload_path, const std::string &tenant_id) {
     return to_py_ret(setup_internal(
         local_hostname, metadata_server, global_segment_size, local_buffer_size,
         protocol, rdma_devices, master_server_addr, transfer_engine,
-        ipc_socket_path, 50052, enable_ssd_offload, true, ssd_offload_path));
+        ipc_socket_path, 50052, enable_ssd_offload, true, ssd_offload_path,
+        tenant_id));
 }
 
 namespace {
@@ -843,30 +966,21 @@ inline std::string get_config(const ConfigDict &config, const std::string &key,
     return (it != config.end()) ? it->second : default_value;
 }
 
-inline size_t get_config_size(const ConfigDict &config, const std::string &key,
-                              size_t default_value) {
+inline std::optional<size_t> get_config_size(const ConfigDict &config,
+                                             const std::string &key,
+                                             size_t default_value) {
     auto it = config.find(key);
     if (it == config.end()) {
         return default_value;
     }
-    const std::string &value = it->second;
-    // Check for negative numbers (stoull incorrectly parses "-1" as large val)
-    if (!value.empty() && value[0] == '-') {
-        LOG(WARNING) << "Invalid negative value for config key '" << key
-                     << "': " << value << ", using default: " << default_value;
-        return default_value;
-    }
-    try {
-        return std::stoull(value);
-    } catch (const std::invalid_argument &e) {
-        LOG(WARNING) << "Invalid non-numeric value for config key '" << key
-                     << "': " << value << ", using default: " << default_value;
-        return default_value;
-    } catch (const std::out_of_range &e) {
-        LOG(WARNING) << "Value out of range for config key '" << key
-                     << "': " << value << ", using default: " << default_value;
-        return default_value;
+
+    auto parsed_size_opt = try_string_to_byte_size(it->second);
+    if (!parsed_size_opt.has_value()) {
+        LOG(ERROR) << "Invalid size value for config key '" << key
+                   << "': " << it->second;
+        return std::nullopt;
     }
+    return static_cast<size_t>(parsed_size_opt.value());
 }
 }  // namespace
 
@@ -888,10 +1002,18 @@ tl::expected<void, ErrorCode> RealClient::setup_internal(
     }
 
     // Extract optional parameters with defaults
-    size_t global_segment_size = get_config_size(
+    auto global_segment_size_opt = get_config_size(
         config, CONFIG_KEY_GLOBAL_SEGMENT_SIZE, DEFAULT_GLOBAL_SEGMENT_SIZE);
-    size_t local_buffer_size = get_config_size(
+    if (!global_segment_size_opt.has_value()) {
+        return tl::unexpected(ErrorCode::INVALID_PARAMS);
+    }
+    auto local_buffer_size_opt = get_config_size(
         config, CONFIG_KEY_LOCAL_BUFFER_SIZE, DEFAULT_LOCAL_BUFFER_SIZE);
+    if (!local_buffer_size_opt.has_value()) {
+        return tl::unexpected(ErrorCode::INVALID_PARAMS);
+    }
+    size_t global_segment_size = global_segment_size_opt.value();
+    size_t local_buffer_size = local_buffer_size_opt.value();
     std::string protocol =
         get_config(config, CONFIG_KEY_PROTOCOL, DEFAULT_PROTOCOL);
     std::string rdma_devices = get_config(config, CONFIG_KEY_RDMA_DEVICES);
@@ -900,19 +1022,19 @@ tl::expected<void, ErrorCode> RealClient::setup_internal(
     std::string ipc_socket_path =
         get_config(config, CONFIG_KEY_IPC_SOCKET_PATH);
 
-    // Validate size parameters are within acceptable ranges
-    if (global_segment_size < MIN_SEGMENT_SIZE ||
-        global_segment_size > MAX_SEGMENT_SIZE) {
-        LOG(ERROR) << "Invalid " << CONFIG_KEY_GLOBAL_SEGMENT_SIZE << ": "
-                   << global_segment_size << ", must be between "
-                   << MIN_SEGMENT_SIZE << " and " << MAX_SEGMENT_SIZE;
-        return tl::unexpected(ErrorCode::INVALID_PARAMS);
-    }
-    if (local_buffer_size < MIN_SEGMENT_SIZE ||
-        local_buffer_size > MAX_SEGMENT_SIZE) {
-        LOG(ERROR) << "Invalid " << CONFIG_KEY_LOCAL_BUFFER_SIZE << ": "
-                   << local_buffer_size << ", must be between "
-                   << MIN_SEGMENT_SIZE << " and " << MAX_SEGMENT_SIZE;
+    // A size of 0 keeps the pure client/server setup semantics.
+    auto validate_size = [](const char *key, size_t value) {
+        if ((value != 0 && value < MIN_SEGMENT_SIZE) ||
+            value > MAX_SEGMENT_SIZE) {
+            LOG(ERROR) << "Invalid " << key << ": " << value
+                       << ", must be 0 or between " << MIN_SEGMENT_SIZE
+                       << " and " << MAX_SEGMENT_SIZE;
+            return false;
+        }
+        return true;
+    };
+    if (!validate_size(CONFIG_KEY_GLOBAL_SEGMENT_SIZE, global_segment_size) ||
+        !validate_size(CONFIG_KEY_LOCAL_BUFFER_SIZE, local_buffer_size)) {
         return tl::unexpected(ErrorCode::INVALID_PARAMS);
     }
 
@@ -924,6 +1046,7 @@ tl::expected<void, ErrorCode> RealClient::setup_internal(
     }
 
     std::string ssd_offload_path = get_config(config, "ssd_offload_path");
+    std::string tenant_id = get_config(config, CONFIG_KEY_TENANT_ID, "default");
 
     std::string enable_ssd_offload_str =
         get_config(config, "enable_ssd_offload", "false");
@@ -933,10 +1056,10 @@ tl::expected<void, ErrorCode> RealClient::setup_internal(
     bool enable_ssd_offload =
         (enable_ssd_offload_str == "true" || enable_ssd_offload_str == "1");
 
-    return setup_internal(local_hostname, metadata_server, global_segment_size,
-                          local_buffer_size, protocol, rdma_devices,
-                          master_server_addr, nullptr, ipc_socket_path, 50052,
-                          enable_ssd_offload, true, ssd_offload_path);
+    return setup_internal(
+        local_hostname, metadata_server, global_segment_size, local_buffer_size,
+        protocol, rdma_devices, master_server_addr, nullptr, ipc_socket_path,
+        50052, enable_ssd_offload, true, ssd_offload_path, tenant_id);
 }
 
 tl::expected<void, ErrorCode> RealClient::initAll_internal(
@@ -968,6 +1091,7 @@ tl::expected<void, ErrorCode> RealClient::tearDownAll_internal() {
     }
 
     stop_ipc_server();
+    stop_dummy_client_monitor();
     stop_http_server();
 
     if (!client_) {
@@ -983,12 +1107,18 @@ tl::expected<void, ErrorCode> RealClient::tearDownAll_internal() {
                 << "Failed to unregister client local buffer on tear down: "
                 << toString(unregister_result.error());
         }
+        std::unique_lock<std::shared_mutex> lock(registered_buffer_mutex_);
+        local_buffer_region_.reset();
     }
+
     // Reset all resources
     client_.reset();
+    ReleaseAllMountedSegmentRecords();
+    ReleaseAllAllocatedSegmentRecords();
     client_buffer_allocator_.reset();
     port_binder_.reset();
     hugepage_segment_ptrs_.clear();
+    ub_segment_ptrs_.clear();
     segment_ptrs_.clear();
     local_hostname = "";
     device_name = "";
@@ -1021,6 +1151,432 @@ tl::expected<void, ErrorCode> RealClient::tearDownAll_internal() {
 
 int RealClient::tearDownAll() { return to_py_ret(tearDownAll_internal()); }
 
+int RealClient::mountSegment(const std::string &path, size_t offset,
+                             size_t size, const std::string &protocol,
+                             const std::string &location,
+                             std::vector<std::string> &out_segment_ids) {
+    if (!client_) {
+        LOG(ERROR) << "Client not initialized";
+        return -1;
+    }
+
+    size_t max_mr_size = globalConfig().max_mr_size;
+    if (max_mr_size == 0) {
+        LOG(ERROR) << "Invalid max_mr_size: 0";
+        return -1;
+    }
+
+    size_t page_size = sysconf(_SC_PAGESIZE);
+    if (max_mr_size < page_size) {
+        LOG(ERROR) << "max_mr_size " << max_mr_size
+                   << " is smaller than page_size " << page_size;
+        return -1;
+    }
+
+    int fd = open(path.c_str(), O_RDWR);
+    if (fd < 0) {
+        LOG(ERROR) << "open failed for " << path << ": " << strerror(errno);
+        return -1;
+    }
+
+    struct stat st;
+    if (fstat(fd, &st) != 0) {
+        LOG(ERROR) << "fstat failed for " << path << ": " << strerror(errno);
+        close(fd);
+        return -1;
+    }
+
+    if (offset > (size_t)st.st_size || size > (size_t)st.st_size - offset) {
+        LOG(ERROR) << "requested range [offset=" << offset << ", size=" << size
+                   << "] exceeds file size " << st.st_size;
+        close(fd);
+        return -1;
+    }
+
+    if (offset % page_size != 0) {
+        LOG(ERROR) << "offset " << offset
+                   << " is not page-aligned (page_size=" << page_size << ")";
+        close(fd);
+        return -1;
+    }
+
+    size_t aligned_max_chunk = (max_mr_size / page_size) * page_size;
+    size_t remaining = size;
+    size_t current_offset = offset;
+    std::vector<std::string> mounted_ids;
+    std::vector<MountedSegmentRecord> mounted_records;
+
+    while (remaining > 0) {
+        size_t chunk_size = std::min(remaining, aligned_max_chunk);
+        if (chunk_size == 0) break;
+
+        void *ptr = mmap(nullptr, chunk_size, PROT_READ | PROT_WRITE,
+                         MAP_SHARED, fd, current_offset);
+        if (ptr == MAP_FAILED) {
+            LOG(ERROR) << "mmap failed for " << path << ": " << strerror(errno);
+            break;
+        }
+
+        auto result =
+            client_->MountSegmentAndGetId(ptr, chunk_size, protocol, location);
+        if (!result.has_value()) {
+            LOG(ERROR) << "MountSegmentAndGetId failed";
+            munmap(ptr, chunk_size);
+            break;
+        }
+
+        std::string segment_id = UuidToString(result.value());
+        mounted_ids.push_back(segment_id);
+        mounted_records.push_back({ptr, chunk_size, path});
+
+        remaining -= chunk_size;
+        current_offset += chunk_size;
+    }
+
+    close(fd);
+
+    // If not all chunks were mounted, roll back successfully mounted ones
+    if (remaining > 0) {
+        for (size_t i = 0; i < mounted_ids.size(); ++i) {
+            UUID id;
+            if (StringToUuid(mounted_ids[i], id)) {
+                client_->UnmountSegmentById(id);
+            }
+            if (mounted_records[i].mmap_base) {
+                munmap(mounted_records[i].mmap_base, mounted_records[i].size);
+            }
+        }
+        out_segment_ids.clear();
+        return -1;
+    }
+
+    {
+        std::lock_guard<std::mutex> lock(mounted_segment_records_mutex_);
+        for (size_t i = 0; i < mounted_ids.size(); ++i) {
+            mounted_segment_records_[mounted_ids[i]] = mounted_records[i];
+        }
+    }
+    out_segment_ids = std::move(mounted_ids);
+    return 0;
+}
+
+void RealClient::ReleaseMountedSegmentRecord(const std::string &segment_id) {
+    MountedSegmentRecord record;
+    bool found = false;
+    {
+        std::lock_guard<std::mutex> lock(mounted_segment_records_mutex_);
+        auto it = mounted_segment_records_.find(segment_id);
+        if (it != mounted_segment_records_.end()) {
+            record = it->second;
+            mounted_segment_records_.erase(it);
+            found = true;
+        }
+    }
+    if (found && record.mmap_base) {
+        munmap(record.mmap_base, record.size);
+    }
+}
+
+void RealClient::ReleaseAllMountedSegmentRecords() {
+    std::vector<MountedSegmentRecord> records;
+    {
+        std::lock_guard<std::mutex> lock(mounted_segment_records_mutex_);
+        records.reserve(mounted_segment_records_.size());
+        for (auto &entry : mounted_segment_records_) {
+            records.push_back(entry.second);
+        }
+        mounted_segment_records_.clear();
+    }
+    for (auto &record : records) {
+        if (record.mmap_base) {
+            munmap(record.mmap_base, record.size);
+        }
+    }
+}
+
+void RealClient::ReleaseAllocatedSegmentRecord(const std::string &segment_id) {
+    AllocatedSegmentRecord record;
+    bool found = false;
+    {
+        std::lock_guard<std::mutex> lock(allocated_segment_records_mutex_);
+        auto it = allocated_segment_records_.find(segment_id);
+        if (it != allocated_segment_records_.end()) {
+            record = it->second;
+            allocated_segment_records_.erase(it);
+            found = true;
+        }
+    }
+    if (found && record.base) {
+        free_memory(record.protocol, record.base);
+    }
+}
+
+void RealClient::ReleaseAllAllocatedSegmentRecords() {
+    std::unordered_map<std::string, AllocatedSegmentRecord> records;
+    {
+        std::lock_guard<std::mutex> lock(allocated_segment_records_mutex_);
+        records.swap(allocated_segment_records_);
+    }
+    for (auto &entry : records) {
+        if (entry.second.base) {
+            free_memory(entry.second.protocol, entry.second.base);
+        }
+    }
+}
+
+int RealClient::unmountSegment(const std::vector<std::string> &segment_ids,
+                               uint64_t grace_period_seconds) {
+    if (!client_) {
+        LOG(ERROR) << "Client not initialized";
+        return -1;
+    }
+
+    uint64_t grace_period_ms = grace_period_seconds * 1000;
+    int first_error = 0;
+    struct SegmentToUnmount {
+        std::string segment_id;
+        UUID id;
+    };
+    std::vector<SegmentToUnmount> to_unmount;
+    std::vector<std::pair<std::string, MountedSegmentRecord>> to_cleanup;
+    {
+        std::lock_guard<std::mutex> lock(mounted_segment_records_mutex_);
+        for (const auto &segment_id : segment_ids) {
+            UUID id;
+            if (!StringToUuid(segment_id, id)) {
+                LOG(ERROR) << "Invalid segment_id: " << segment_id;
+                if (first_error == 0) first_error = -1;
+                continue;
+            }
+
+            auto it = mounted_segment_records_.find(segment_id);
+            if (it == mounted_segment_records_.end()) {
+                LOG(ERROR) << "segment_id not found in mounted records: "
+                           << segment_id;
+                if (first_error == 0) first_error = -1;
+                continue;
+            }
+
+            to_unmount.push_back({segment_id, id});
+        }
+    }
+
+    for (auto &entry : to_unmount) {
+        std::function<void(const UUID &)> cleanup_callback;
+        if (grace_period_ms != 0) {
+            cleanup_callback = [this](const UUID &cleanup_id) {
+                ReleaseMountedSegmentRecord(UuidToString(cleanup_id));
+            };
+        }
+        auto result = client_->UnmountSegmentById(entry.id, grace_period_ms,
+                                                  std::move(cleanup_callback));
+        if (!result.has_value()) {
+            LOG(ERROR) << "UnmountSegmentById failed for " << entry.segment_id;
+            if (first_error == 0) {
+                first_error = static_cast<int>(result.error());
+            }
+            continue;  // Don't release local resources on failure
+        }
+
+        // For immediate unmount, clean up local mmap/fd right away.
+        // For graceful unmount, local mmap/fd is kept until the segment
+        // is actually removed or the client destructor runs.
+        if (grace_period_ms == 0) {
+            std::lock_guard<std::mutex> lock(mounted_segment_records_mutex_);
+            auto it = mounted_segment_records_.find(entry.segment_id);
+            if (it != mounted_segment_records_.end()) {
+                to_cleanup.emplace_back(entry.segment_id, it->second);
+                mounted_segment_records_.erase(it);
+            }
+        }
+    }
+
+    for (auto &p : to_cleanup) {
+        if (p.second.mmap_base) {
+            munmap(p.second.mmap_base, p.second.size);
+        }
+    }
+
+    return first_error;
+}
+
+int RealClient::allocateAndMountSegment(
+    size_t size, const std::string &protocol, const std::string &location,
+    std::vector<std::string> &out_segment_ids, size_t *out_allocated_size) {
+    if (!client_) {
+        LOG(ERROR) << "Client not initialized";
+        return -1;
+    }
+
+    size_t max_mr_size = globalConfig().max_mr_size;
+    if (max_mr_size == 0) {
+        LOG(ERROR) << "Invalid max_mr_size: 0";
+        return -1;
+    }
+
+    if (size == 0) {
+        LOG(ERROR) << "size is 0";
+        return -1;
+    }
+
+    const size_t slab_size = facebook::cachelib::Slab::kSize;
+    size_t page_size = sysconf(_SC_PAGESIZE);
+    if (max_mr_size < page_size) {
+        LOG(ERROR) << "max_mr_size " << max_mr_size
+                   << " is smaller than page_size " << page_size;
+        return -1;
+    }
+
+    size_t aligned_max_chunk = (max_mr_size / page_size) * page_size;
+    if (aligned_max_chunk < slab_size) {
+        LOG(ERROR) << "max_mr_size " << max_mr_size
+                   << " is smaller than slab_size " << slab_size;
+        return -1;
+    }
+    // Round down chunk size to slab_size multiple
+    aligned_max_chunk = (aligned_max_chunk / slab_size) * slab_size;
+
+    // Check overflow before aligning up to slab_size
+    if (size > std::numeric_limits<size_t>::max() - (slab_size - 1)) {
+        LOG(ERROR) << "size " << size
+                   << " overflows when aligning to slab_size";
+        return -1;
+    }
+    // Round up total size to slab_size multiple
+    size_t aligned_total_size =
+        ((size + slab_size - 1) / slab_size) * slab_size;
+    size_t remaining = aligned_total_size;
+    std::vector<std::string> mounted_ids;
+    std::vector<AllocatedSegmentRecord> allocated_records;
+
+    while (remaining > 0) {
+        size_t chunk_size = std::min(remaining, aligned_max_chunk);
+        if (chunk_size == 0) break;
+
+        void *ptr = allocate_buffer_allocator_memory(chunk_size, protocol);
+        if (!ptr) {
+            LOG(ERROR) << "allocate_buffer_allocator_memory failed for size "
+                       << chunk_size;
+            break;
+        }
+
+        auto result =
+            client_->MountSegmentAndGetId(ptr, chunk_size, protocol, location);
+        if (!result.has_value()) {
+            LOG(ERROR) << "MountSegmentAndGetId failed";
+            free_memory(protocol, ptr);
+            break;
+        }
+
+        std::string segment_id = UuidToString(result.value());
+        mounted_ids.push_back(segment_id);
+        allocated_records.push_back({ptr, chunk_size, protocol});
+
+        remaining -= chunk_size;
+    }
+
+    if (remaining > 0) {
+        for (size_t i = 0; i < mounted_ids.size(); ++i) {
+            UUID id;
+            if (StringToUuid(mounted_ids[i], id)) {
+                client_->UnmountSegmentById(id);
+            }
+            if (allocated_records[i].base) {
+                free_memory(allocated_records[i].protocol,
+                            allocated_records[i].base);
+            }
+        }
+        out_segment_ids.clear();
+        return -1;
+    }
+
+    {
+        std::lock_guard<std::mutex> lock(allocated_segment_records_mutex_);
+        for (size_t i = 0; i < mounted_ids.size(); ++i) {
+            allocated_segment_records_[mounted_ids[i]] = allocated_records[i];
+        }
+    }
+    if (out_allocated_size) {
+        *out_allocated_size = aligned_total_size;
+    }
+    out_segment_ids = std::move(mounted_ids);
+    return 0;
+}
+
+int RealClient::unmountAndFreeSegment(
+    const std::vector<std::string> &segment_ids,
+    uint64_t grace_period_seconds) {
+    if (!client_) {
+        LOG(ERROR) << "Client not initialized";
+        return -1;
+    }
+
+    uint64_t grace_period_ms = grace_period_seconds * 1000;
+    int first_error = 0;
+    struct SegmentToUnmount {
+        std::string segment_id;
+        UUID id;
+    };
+    std::vector<SegmentToUnmount> to_unmount;
+    std::vector<std::pair<std::string, AllocatedSegmentRecord>> to_cleanup;
+    {
+        std::lock_guard<std::mutex> lock(allocated_segment_records_mutex_);
+        for (const auto &segment_id : segment_ids) {
+            UUID id;
+            if (!StringToUuid(segment_id, id)) {
+                LOG(ERROR) << "Invalid segment_id: " << segment_id;
+                if (first_error == 0) first_error = -1;
+                continue;
+            }
+
+            auto it = allocated_segment_records_.find(segment_id);
+            if (it == allocated_segment_records_.end()) {
+                LOG(ERROR) << "segment_id not found in allocated records: "
+                           << segment_id;
+                if (first_error == 0) first_error = -1;
+                continue;
+            }
+
+            to_unmount.push_back({segment_id, id});
+        }
+    }
+
+    for (auto &entry : to_unmount) {
+        std::function<void(const UUID &)> cleanup_callback;
+        if (grace_period_ms != 0) {
+            cleanup_callback = [this](const UUID &cleanup_id) {
+                ReleaseAllocatedSegmentRecord(UuidToString(cleanup_id));
+            };
+        }
+        auto result = client_->UnmountSegmentById(entry.id, grace_period_ms,
+                                                  std::move(cleanup_callback));
+        if (!result.has_value()) {
+            LOG(ERROR) << "UnmountSegmentById failed for " << entry.segment_id;
+            if (first_error == 0) {
+                first_error = static_cast<int>(result.error());
+            }
+            continue;  // Don't release local resources on failure
+        }
+
+        if (grace_period_ms == 0) {
+            std::lock_guard<std::mutex> lock(allocated_segment_records_mutex_);
+            auto it = allocated_segment_records_.find(entry.segment_id);
+            if (it != allocated_segment_records_.end()) {
+                to_cleanup.emplace_back(entry.segment_id, it->second);
+                allocated_segment_records_.erase(it);
+            }
+        }
+    }
+
+    for (auto &p : to_cleanup) {
+        if (p.second.base) {
+            free_memory(p.second.protocol, p.second.base);
+        }
+    }
+
+    return first_error;
+}
+
 int RealClient::health_check() {
     if (closed_.load()) return HC_NOT_INITIALIZED;
     if (!client_) return HC_NOT_INITIALIZED;
@@ -1510,6 +2066,14 @@ int64_t RealClient::getSize(const std::string &key) {
 tl::expected<void, ErrorCode> RealClient::map_shm_internal(
     int fd, uint64_t dummy_base_addr, size_t shm_size, bool is_local_buffer,
     const UUID &client_id) {
+    return map_shm_internal_with_device(fd, dummy_base_addr, shm_size,
+                                        is_local_buffer,
+                                        kInvalidPhysicalDeviceId, client_id);
+}
+
+tl::expected<void, ErrorCode> RealClient::map_shm_internal_with_device(
+    int fd, uint64_t dummy_base_addr, size_t shm_size, bool is_local_buffer,
+    int32_t physical_device_id, const UUID &client_id) {
     std::stringstream addr_stream;
     addr_stream << "0x" << std::hex << dummy_base_addr;
 
@@ -1518,15 +2082,16 @@ tl::expected<void, ErrorCode> RealClient::map_shm_internal(
         "_" + std::to_string(client_id.second) + "_" + addr_stream.str();
     std::unique_lock<std::shared_mutex> lock(dummy_client_mutex_);
 
-    // Check if client context exists, create if not
-    auto &context = shm_contexts_[client_id];
-
     // Check if this shm is already mapped
-    for (const auto &shm : context.mapped_shms) {
-        if (shm.dummy_base_addr == static_cast<uintptr_t>(dummy_base_addr)) {
-            LOG(INFO) << "Segment already mapped: " << shm_name;
-            if (fd >= 0) close(fd);
-            return {};
+    auto context_it = shm_contexts_.find(client_id);
+    if (context_it != shm_contexts_.end()) {
+        for (const auto &shm : context_it->second.mapped_shms) {
+            if (shm.dummy_base_addr ==
+                static_cast<uintptr_t>(dummy_base_addr)) {
+                LOG(INFO) << "Segment already mapped: " << shm_name;
+                if (fd >= 0) close(fd);
+                return {};
+            }
         }
     }
 
@@ -1535,6 +2100,15 @@ tl::expected<void, ErrorCode> RealClient::map_shm_internal(
         return tl::make_unexpected(ErrorCode::INVALID_PARAMS);
     }
 
+#ifdef USE_ASCEND_DIRECT
+    auto context_result = set_context_if_needed(protocol, physical_device_id,
+                                                "POSIX shm registration");
+    if (!context_result) {
+        close(fd);
+        return context_result;
+    }
+#endif
+
     // Map shared memory from FD
     void *shm_buffer =
         mmap(nullptr, shm_size, PROT_READ | PROT_WRITE, MAP_SHARED, fd, 0);
@@ -1554,6 +2128,11 @@ tl::expected<void, ErrorCode> RealClient::map_shm_internal(
     shm.dummy_base_addr = static_cast<uintptr_t>(dummy_base_addr);
     shm.shm_addr_offset = reinterpret_cast<uintptr_t>(shm_buffer) -
                           reinterpret_cast<uintptr_t>(dummy_base_addr);
+    shm.is_ascend =
+        false;  // memfd/POSIX path; teardown differs from Ascend VMM/IPC
+    shm.device_id = physical_device_id;
+
+    auto &context = shm_contexts_[client_id];
 
     if (shm_size > 0) {
         auto result = client_->RegisterLocalMemory(
@@ -1829,12 +2408,59 @@ tl::expected<void, ErrorCode> RealClient::ascend_unmap_shm_internal(
         if (!shm.shm_buffer) {
             continue;
         }
-        teardown_ascend_shm_buffer(shm);
+        if (shm.is_ascend) {
+            teardown_ascend_shm_buffer(shm);
+        } else {
+#ifdef USE_ASCEND_DIRECT
+            auto context_result = set_context_if_needed(protocol, shm.device_id,
+                                                        "POSIX shm cleanup");
+            if (!context_result) {
+                LOG(WARNING) << "Skip unregisterLocalMemory for POSIX shm: "
+                             << shm.shm_name
+                             << ", error: " << toString(context_result.error());
+            } else
+#endif
+                // Host POSIX shm mapped via map_shm_internal (memfd); not
+                // Ascend VMM/IPC.
+                if (shm.shm_size > 0 && client_) {
+                    auto res =
+                        client_->unregisterLocalMemory(shm.shm_buffer, true);
+                    if (!res) {
+                        LOG(WARNING) << "Failed to unregister local memory for "
+                                        "POSIX shm: "
+                                     << shm.shm_name
+                                     << ", error: " << toString(res.error());
+                    }
+                }
+            if (munmap(shm.shm_buffer, shm.shm_size) != 0) {
+                LOG(ERROR) << "Failed to munmap POSIX shared memory: "
+                           << shm.shm_name << ", error: " << strerror(errno);
+            } else {
+                LOG(INFO) << "Unmapped POSIX shared memory: " << shm.shm_name
+                          << ", size: " << shm.shm_size;
+            }
+        }
     }
     shm_contexts_.erase(it);
     return {};
 }
 
+tl::expected<bool, ErrorCode> RealClient::is_shm_mapped_internal(
+    uint64_t dummy_base_addr, const UUID &client_id) {
+    std::shared_lock<std::shared_mutex> lock(dummy_client_mutex_);
+    auto context_it = shm_contexts_.find(client_id);
+    if (context_it == shm_contexts_.end()) {
+        return false;
+    }
+
+    const auto target_addr = static_cast<uintptr_t>(dummy_base_addr);
+    const auto &mapped_shms = context_it->second.mapped_shms;
+    return std::any_of(mapped_shms.begin(), mapped_shms.end(),
+                       [target_addr](const MappedShm &shm) {
+                           return shm.dummy_base_addr == target_addr;
+                       });
+}
+
 tl::expected<void, ErrorCode> RealClient::unregister_shm_buffer_internal(
     uint64_t dummy_base_addr, const UUID &client_id) {
     std::unique_lock<std::shared_mutex> lock(dummy_client_mutex_);
@@ -1868,6 +2494,13 @@ tl::expected<void, ErrorCode> RealClient::unregister_shm_buffer_internal(
         if (shm_it->is_ascend) {
             teardown_ascend_shm_buffer(*shm_it);
         } else {
+#ifdef USE_ASCEND_DIRECT
+            auto context_result = set_context_if_needed(
+                protocol, shm_it->device_id, "POSIX shm unregister");
+            if (!context_result) {
+                return context_result;
+            }
+#endif
             auto rc = client_->unregisterLocalMemory(shm_it->shm_buffer, true);
             if (!rc) {
                 LOG(ERROR) << "Failed to unregister memory: "
@@ -1923,43 +2556,73 @@ std::shared_ptr<BufferHandle> RealClient::get_buffer_internal(
         return nullptr;
     }
 
-    const auto &res = client_->GetPreferredReplica(replica_list);
-    if (!res) {
-        LOG(ERROR) << "Empty replica list for key: " << key;
+    // Select best replica: prefer local MEMORY, then any MEMORY,
+    // then LOCAL_DISK, then DISK.
+    // LOCAL_DISK data is on a remote node's SSD — must use offload RPC.
+    // MEMORY / DISK are handled via client_->Get below.
+    auto local_endpoints = client_->GetLocalEndpoints();
+    const auto *best_replica = SelectBestReplica(replica_list, local_endpoints);
+    if (!best_replica) {
+        LOG(ERROR) << "No usable replica for key: " << key;
         return nullptr;
     }
 
-    const auto &replica = res.value();
+    const auto &replica = *best_replica;
     uint64_t total_length = calculate_total_size(replica);
 
     if (total_length == 0) {
         return nullptr;
     }
 
-    // Normal allocation path
+    // Allocate buffer
     auto alloc_result = client_buffer_allocator->allocate(total_length);
     if (!alloc_result) {
         LOG(ERROR) << "Failed to allocate buffer for get_buffer, key: " << key;
         return nullptr;
     }
 
-    auto &buffer_handle = *alloc_result;
+    auto buffer_handle =
+        std::make_shared<BufferHandle>(std::move(*alloc_result));
+
+    if (best_replica->is_local_disk_replica()) {
+        // LOCAL_DISK: data is on remote node's SSD. Use offload RPC.
+        const auto &endpoint =
+            best_replica->get_local_disk_descriptor().transport_endpoint;
+        std::unordered_map<std::string, std::vector<Slice>> objects;
+        objects.emplace(
+            key, std::vector<Slice>{{buffer_handle->ptr(), total_length}});
+        auto read_result =
+            batch_get_into_offload_object_internal(endpoint, objects);
+        if (!read_result) {
+            LOG(ERROR) << "SSD read failed for key '" << key
+                       << "': " << toString(read_result.error());
+            return nullptr;
+        }
+        return buffer_handle;
+    }
+
+    // MEMORY / DISK: use client_->Get.  FilterQueryResult ensures
+    // Client::Get's internal FindFirstCompleteReplica can only see
+    // the replica we selected, preventing accidental LOCAL_DISK picks.
+    if (replica.is_disk_replica() &&
+        gpu_staging::IsDevicePointer(buffer_handle->ptr(), nullptr)) {
+        LOG(WARNING) << "DISK replica for key '" << key
+                     << "' received a device pointer from the allocator; "
+                     << "file I/O cannot write to GPU memory — read will fail. "
+                     << "Ensure client_buffer_allocator_ returns host memory.";
+    }
 
-    // Create slices for the allocated buffer
     std::vector<Slice> slices;
-    allocateSlices(slices, replica, buffer_handle.ptr());
-
-    // Get the object data
-    auto get_result = client_->Get(key, query_result.value(), slices);
+    allocateSlices(slices, replica, buffer_handle->ptr());
+    auto filtered_qr = FilterQueryResult(query_result.value(), replica);
+    auto get_result = client_->Get(key, filtered_qr, slices);
     if (!get_result) {
         LOG(ERROR) << "Get failed for key: " << key
                    << " with error: " << toString(get_result.error());
         return nullptr;
     }
 
-    // Create BufferHandle with the allocated memory
-    // The buffer will be managed by the BufferHandle's shared_ptr
-    return std::make_shared<BufferHandle>(std::move(buffer_handle));
+    return buffer_handle;
 }
 
 // Implementation of get_buffer method
@@ -2175,9 +2838,18 @@ RealClient::batch_get_buffer_internal(
         std::unique_ptr<BufferHandle> buffer_handle;
         std::vector<Slice> slices;
     };
+    struct DiskKeyOp {
+        size_t original_index;
+        std::string key;
+        QueryResult query_result;
+        std::unique_ptr<BufferHandle> buffer_handle;
+        uint64_t total_size;
+    };
     std::vector<KeyOp> valid_ops;
+    std::vector<DiskKeyOp> disk_ops;
     valid_ops.reserve(keys.size());
 
+    auto local_endpoints = client_->GetLocalEndpoints();
     for (size_t i = 0; i < keys.size(); ++i) {
         const auto &key = keys[i];
 
@@ -2196,7 +2868,15 @@ RealClient::batch_get_buffer_internal(
             continue;
         }
 
-        const auto &replica = query_result_values.replicas[0];
+        // Select best replica: prefer local MEMORY, then any MEMORY,
+        // then LOCAL_DISK, then DISK.
+        const auto *best_replica =
+            SelectBestReplica(query_result_values.replicas, local_endpoints);
+        if (!best_replica) {
+            LOG(ERROR) << "No usable replica for key: " << key;
+            continue;
+        }
+        const auto replica = *best_replica;
         uint64_t total_size = calculate_total_size(replica);
         if (total_size == 0) {
             continue;
@@ -2215,43 +2895,120 @@ RealClient::batch_get_buffer_internal(
         std::vector<Slice> slices;
         allocateSlices(slices, replica, buffer_handle->ptr());
 
-        valid_ops.emplace_back(
-            KeyOp{.original_index = i,
-                  .key = key,
-                  .query_result = std::move(query_result_values),
-                  .buffer_handle = std::move(buffer_handle),
-                  .slices = std::move(slices)});
+        if (replica.is_local_disk_replica()) {
+            // LOCAL_DISK: buffer is allocated and registered via
+            // client_buffer_allocator_, route to SSD RPC path below.
+            disk_ops.emplace_back(DiskKeyOp{
+                .original_index = i,
+                .key = key,
+                .query_result = FilterQueryResult(query_result_values, replica),
+                .buffer_handle = std::move(buffer_handle),
+                .total_size = total_size});
+            continue;
+        }
+
+        // DISK replicas use storage_backend::vector_read (file I/O) which
+        // can only write to CPU-addressable memory.  If the allocator ever
+        // returns device memory for DISK, the read will silently fail.
+        if (replica.is_disk_replica() &&
+            gpu_staging::IsDevicePointer(buffer_handle->ptr(), nullptr)) {
+            LOG(WARNING)
+                << "DISK replica for key '" << key
+                << "' received a device pointer from the allocator; "
+                << "file I/O cannot write to GPU memory — read will fail. "
+                << "Ensure client_buffer_allocator_ returns host memory.";
+        }
+        valid_ops.emplace_back(KeyOp{
+            .original_index = i,
+            .key = key,
+            .query_result = FilterQueryResult(query_result_values, replica),
+            .buffer_handle = std::move(buffer_handle),
+            .slices = std::move(slices)});
     }
 
-    if (valid_ops.empty()) {
+    if (valid_ops.empty() && disk_ops.empty()) {
         return final_results;
     }
 
-    // 3. Execute batch get
-    std::vector<std::string> batch_keys;
-    std::vector<QueryResult> batch_query_results;
-    std::unordered_map<std::string, std::vector<Slice>> batch_slices;
-    batch_keys.reserve(valid_ops.size());
-    batch_query_results.reserve(valid_ops.size());
+    // 3. Execute batch get for memory/disk replicas
+    if (!valid_ops.empty()) {
+        std::vector<std::string> batch_keys;
+        std::vector<QueryResult> batch_query_results;
+        std::unordered_map<std::string, std::vector<Slice>> batch_slices;
+        batch_keys.reserve(valid_ops.size());
+        batch_query_results.reserve(valid_ops.size());
 
-    for (auto &op : valid_ops) {
-        batch_keys.push_back(op.key);
-        batch_query_results.push_back(op.query_result);
-        batch_slices[op.key] = op.slices;
+        for (auto &op : valid_ops) {
+            batch_keys.push_back(op.key);
+            batch_query_results.push_back(op.query_result);
+            batch_slices[op.key] = op.slices;
+        }
+
+        auto batch_get_results =
+            client_->BatchGet(batch_keys, batch_query_results, batch_slices);
+
+        // 4. Process results and create BufferHandles
+        for (size_t i = 0; i < valid_ops.size(); ++i) {
+            if (batch_get_results[i]) {
+                auto &op = valid_ops[i];
+                final_results[op.original_index] =
+                    std::make_shared<BufferHandle>(
+                        std::move(*op.buffer_handle));
+            } else {
+                LOG(ERROR) << "BatchGet failed for key '" << valid_ops[i].key
+                           << "': " << toString(batch_get_results[i].error());
+            }
+        }
     }
 
-    auto batch_get_results =
-        client_->BatchGet(batch_keys, batch_query_results, batch_slices);
+    // 5. Execute batch get for LOCAL_DISK replicas via SSD RPC
+    if (!disk_ops.empty()) {
+        // Group by transport endpoint
+        std::unordered_map<std::string,
+                           std::unordered_map<std::string, std::vector<Slice>>>
+            offload_objects;
+        // Build key -> disk_ops index for result lookup
+        std::unordered_map<std::string, size_t> disk_key_to_idx;
 
-    // 4. Process results and create BufferHandles
-    for (size_t i = 0; i < valid_ops.size(); ++i) {
-        if (batch_get_results[i]) {
-            auto &op = valid_ops[i];
-            final_results[op.original_index] =
-                std::make_shared<BufferHandle>(std::move(*op.buffer_handle));
-        } else {
-            LOG(ERROR) << "BatchGet failed for key '" << valid_ops[i].key
-                       << "': " << toString(batch_get_results[i].error());
+        for (size_t idx = 0; idx < disk_ops.size(); ++idx) {
+            auto &op = disk_ops[idx];
+            // Find the LOCAL_DISK replica — replicas may be in any order.
+            const Replica::Descriptor *replica_ptr = nullptr;
+            for (const auto &r : op.query_result.replicas) {
+                if (r.is_local_disk_replica()) {
+                    replica_ptr = &r;
+                    break;
+                }
+            }
+            if (!replica_ptr) {
+                LOG(ERROR) << "No LOCAL_DISK replica found for key: " << op.key;
+                continue;
+            }
+            const auto &replica = *replica_ptr;
+            offload_objects[replica.get_local_disk_descriptor()
+                                .transport_endpoint]
+                .emplace(op.key, std::vector<Slice>{
+                                     {op.buffer_handle->ptr(), op.total_size}});
+            disk_key_to_idx[op.key] = idx;
+        }
+
+        for (auto &[endpoint, objects] : offload_objects) {
+            if (objects.empty()) continue;
+            auto read_result =
+                batch_get_into_offload_object_internal(endpoint, objects);
+            for (auto &[key, slices] : objects) {
+                auto idx_it = disk_key_to_idx.find(key);
+                if (idx_it == disk_key_to_idx.end()) continue;
+                auto &op = disk_ops[idx_it->second];
+                if (read_result) {
+                    final_results[op.original_index] =
+                        std::make_shared<BufferHandle>(
+                            std::move(*op.buffer_handle));
+                } else {
+                    LOG(ERROR) << "SSD read failed for key '" << key
+                               << "': " << toString(read_result.error());
+                }
+            }
         }
     }
 
@@ -2316,6 +3073,40 @@ int RealClient::unregister_buffer(void *buffer) {
     return to_py_ret(unregister_buffer_internal(buffer));
 }
 
+std::optional<RealClient::WritableBufferRegion>
+RealClient::resolve_writable_buffer_region(void *buffer) const {
+    std::shared_lock<std::shared_mutex> lock(registered_buffer_mutex_);
+    const auto target = reinterpret_cast<uintptr_t>(buffer);
+    for (const auto &[registered_buffer, registered_size] :
+         registered_buffer_sizes_) {
+        const auto base = reinterpret_cast<uintptr_t>(registered_buffer);
+        if (target < base) {
+            continue;
+        }
+        const auto offset = target - base;
+        if (offset < registered_size) {
+            return WritableBufferRegion{
+                .base = registered_buffer,
+                .size = registered_size,
+                .offset = static_cast<size_t>(offset),
+            };
+        }
+    }
+    if (local_buffer_region_.has_value()) {
+        const auto base =
+            reinterpret_cast<uintptr_t>(local_buffer_region_->base);
+        if (target >= base) {
+            const auto offset = target - base;
+            if (offset < local_buffer_region_->size) {
+                WritableBufferRegion region = *local_buffer_region_;
+                region.offset = static_cast<size_t>(offset);
+                return region;
+            }
+        }
+    }
+    return std::nullopt;
+}
+
 tl::expected<RealClient::RangedReadMetadata, ErrorCode>
 RealClient::resolve_ranged_read_metadata(const std::string &key) {
     if (!client_) {
@@ -2323,34 +3114,8 @@ RealClient::resolve_ranged_read_metadata(const std::string &key) {
         return tl::unexpected(ErrorCode::INVALID_PARAMS);
     }
 
-    auto query_result = client_->Query(key);
-    if (!query_result) {
-        if (query_result.error() == ErrorCode::OBJECT_NOT_FOUND ||
-            query_result.error() == ErrorCode::REPLICA_IS_NOT_READY) {
-            return tl::unexpected(query_result.error());
-        }
-        LOG(ERROR) << "Query failed for key: " << key
-                   << " with error: " << toString(query_result.error());
-        return tl::unexpected(query_result.error());
-    }
-
-    const auto &replica_list = query_result.value().replicas;
-    if (replica_list.empty()) {
-        LOG(ERROR) << "Internal error: replica_list is empty";
-        return tl::unexpected(ErrorCode::INVALID_PARAMS);
-    }
-
-    const auto &res = client_->GetPreferredReplica(replica_list);
-    if (!res) {
-        LOG(ERROR) << "Internal error: replica_list is empty";
-        return tl::unexpected(ErrorCode::INVALID_PARAMS);
-    }
-
-    auto query_value = std::move(query_result.value());
-    auto replica = res.value();
-    return RangedReadMetadata{.query_result = std::move(query_value),
-                              .replica = std::move(replica),
-                              .total_size = calculate_total_size(replica)};
+    return build_ranged_read_metadata_from_query_result(key,
+                                                        client_->Query(key));
 }
 
 tl::expected<int64_t, ErrorCode> RealClient::execute_ranged_read(
@@ -2375,11 +3140,57 @@ tl::expected<int64_t, ErrorCode> RealClient::execute_ranged_read(
     }
 
     if (src_offset == 0 && size == total_size) {
+        // LOCAL_DISK full-object read: destination buffer is passed
+        // directly to SSD RPC (same pattern as single-buffer batch_get_into;
+        // GPU buffers are pre-registered by vLLM via register_buffer).
+        if (replica.is_local_disk_replica()) {
+            const auto &endpoint =
+                replica.get_local_disk_descriptor().transport_endpoint;
+            std::unordered_map<std::string, std::vector<Slice>> objects;
+            objects.emplace(
+                key, std::vector<Slice>{
+                         {static_cast<char *>(buffer) + dst_offset, size}});
+            auto result =
+                batch_get_into_offload_object_internal(endpoint, objects);
+            if (!result) return tl::unexpected(result.error());
+            return static_cast<int64_t>(total_size);
+        }
+
+        if (replica.is_disk_replica()) {
+            // DISK full read: local file I/O (vector_read) cannot write to
+            // GPU memory. Use temp CPU buffer, then scatter to dst.
+            auto alloc_result = client_buffer_allocator_->allocate(total_size);
+            if (!alloc_result) {
+                LOG(ERROR) << "Failed to allocate temp buffer for DISK full "
+                           << "read, key: " << key << ", size: " << total_size;
+                return tl::unexpected(ErrorCode::NO_AVAILABLE_HANDLE);
+            }
+            BufferHandle tmp_handle(std::move(*alloc_result));
+            std::vector<mooncake::Slice> tmp_slices;
+            allocateSlices(tmp_slices, replica, tmp_handle.ptr());
+            auto filtered_qr = FilterQueryResult(query_result, replica);
+            auto get_result = client_->Get(key, filtered_qr, tmp_slices);
+            if (!get_result) {
+                LOG(ERROR) << "DISK Get failed for key: " << key
+                           << " with error: " << toString(get_result.error());
+                return tl::unexpected(get_result.error());
+            }
+            void *dst = static_cast<char *>(buffer) + dst_offset;
+            const void *src = tmp_handle.ptr();
+            if (auto r = scatter_host_to_maybe_device(
+                    dst, src, total_size, "DISK full read, key: " + key);
+                !r) {
+                return tl::unexpected(r.error());
+            }
+            return static_cast<int64_t>(total_size);
+        }
+
         std::vector<mooncake::Slice> slices;
         allocateSlices(slices, replica,
                        static_cast<char *>(buffer) + dst_offset);
 
-        auto get_result = client_->Get(key, query_result, slices);
+        auto filtered_qr = FilterQueryResult(query_result, replica);
+        auto get_result = client_->Get(key, filtered_qr, slices);
         if (!get_result) {
             LOG(ERROR) << "Get failed for key: " << key
                        << " with error: " << toString(get_result.error());
@@ -2388,8 +3199,75 @@ tl::expected<int64_t, ErrorCode> RealClient::execute_ranged_read(
         return static_cast<int64_t>(total_size);
     }
 
+    // Partial disk read: allocate temp CPU buffer, invoke read_op to
+    // fill it, then scatter [src_offset, src_offset+size) to dst.
+    //
+    // buf_size controls how much to allocate / read.  DISK must use
+    // total_size (allocateSlices requires full-object slices).
+    // LOCAL_DISK can use src_offset + size (offload RPC transfers
+    // sequentially from remote offset 0).
+    auto partial_disk_read =
+        [&](auto &&read_op,
+            size_t buf_size) -> tl::expected<int64_t, ErrorCode> {
+        auto alloc_result = client_buffer_allocator_->allocate(buf_size);
+        if (!alloc_result) {
+            LOG(ERROR) << "Failed to allocate temp buffer for ranged disk "
+                       << "read, key: " << key << ", size: " << buf_size;
+            return tl::unexpected(ErrorCode::NO_AVAILABLE_HANDLE);
+        }
+        BufferHandle tmp_handle(std::move(*alloc_result));
+        auto read_result = read_op(tmp_handle.ptr());
+        if (!read_result) return tl::unexpected(read_result.error());
+        void *dst = static_cast<char *>(buffer) + dst_offset;
+        const void *src =
+            static_cast<const char *>(tmp_handle.ptr()) + src_offset;
+        if (auto r = scatter_host_to_maybe_device(
+                dst, src, size, "ranged disk read, key: " + key);
+            !r) {
+            return tl::unexpected(r.error());
+        }
+        return static_cast<int64_t>(size);
+    };
+
+    if (replica.is_local_disk_replica()) {
+        // LOCAL_DISK: offload RPC transfers sequentially from remote offset
+        // 0, so we only need src_offset + size bytes (not total_size).
+        return partial_disk_read(
+            [&](void *tmp_buf) -> tl::expected<void, ErrorCode> {
+                const auto &endpoint =
+                    replica.get_local_disk_descriptor().transport_endpoint;
+                std::unordered_map<std::string, std::vector<Slice>> objects;
+                objects.emplace(
+                    key, std::vector<Slice>{{static_cast<char *>(tmp_buf),
+                                             src_offset + size}});
+                return batch_get_into_offload_object_internal(endpoint,
+                                                              objects);
+            },
+            src_offset + size);
+    }
+
+    if (replica.is_disk_replica()) {
+        // DISK: client_->Get + allocateSlices requires full-object slices,
+        // so we must allocate total_size.
+        return partial_disk_read(
+            [&](void *tmp_buf) -> tl::expected<void, ErrorCode> {
+                std::vector<mooncake::Slice> tmp_slices;
+                allocateSlices(tmp_slices, replica, tmp_buf);
+                auto filtered_qr = FilterQueryResult(query_result, replica);
+                auto get_result = client_->Get(key, filtered_qr, tmp_slices);
+                if (!get_result) {
+                    LOG(ERROR)
+                        << "DISK Get failed for key: " << key
+                        << " with error: " << toString(get_result.error());
+                    return tl::unexpected(get_result.error());
+                }
+                return {};
+            },
+            total_size);
+    }
+
     if (!replica.is_memory_replica()) {
-        LOG(ERROR) << "ranged reads only support memory replicas";
+        LOG(ERROR) << "ranged reads only support memory/disk replicas";
         return tl::unexpected(ErrorCode::INVALID_REPLICA);
     }
 
@@ -2446,7 +3324,15 @@ RealClient::get_into_ranges_internal(
     const std::vector<size_t> *buffer_capacities,
     std::vector<std::vector<std::vector<tl::expected<int64_t, ErrorCode>>>>
         *prepared_results,
-    const std::vector<std::vector<std::vector<bool>>> *valid_fragments) {
+    const std::vector<std::vector<std::vector<bool>>> *valid_fragments,
+    const QueryResultCache *query_result_cache) {
+    if (!client_) {
+        LOG(ERROR) << "Client is not initialized";
+        return build_ranged_read_internal_error_results(
+            buffers.size(), all_keys, all_dst_offsets,
+            ErrorCode::INVALID_PARAMS);
+    }
+
     const size_t buffer_count = buffers.size();
     PreparedRangedReadRequest prepared;
     if (prepared_results != nullptr && valid_fragments != nullptr) {
@@ -2474,21 +3360,42 @@ RealClient::get_into_ranges_internal(
         resolved_buffer_capacities = *buffer_capacities;
     } else {
         resolved_buffer_capacities.resize(buffer_count, 0);
-        std::shared_lock<std::shared_mutex> lock(registered_buffer_mutex_);
         for (size_t i = 0; i < buffer_count; ++i) {
-            auto it = registered_buffer_sizes_.find(buffers[i]);
-            if (it == registered_buffer_sizes_.end()) {
+            auto region = resolve_writable_buffer_region(buffers[i]);
+            if (!region.has_value()) {
                 LOG(ERROR)
-                    << "get_into_ranges: buffer is not registered at index "
+                    << "get_into_ranges: buffer is not Store-managed writable "
+                       "memory at index "
                     << i;
                 continue;
             }
-            resolved_buffer_capacities[i] = it->second;
+            resolved_buffer_capacities[i] = region->size - region->offset;
         }
     }
 
     std::unordered_map<std::string, tl::expected<RangedReadMetadata, ErrorCode>>
         metadata_cache;
+    size_t key_count_hint = 0;
+    for (const auto &keys : all_keys) {
+        key_count_hint += keys.size();
+    }
+    metadata_cache.reserve(key_count_hint);
+    auto now = std::chrono::steady_clock::now();
+    if (query_result_cache != nullptr) {
+        for (const auto &[key, query_result] : *query_result_cache) {
+            if (!query_result) {
+                metadata_cache.emplace(key,
+                                       tl::unexpected(query_result.error()));
+                continue;
+            }
+            if (query_result->IsLeaseExpired(now)) {
+                continue;
+            }
+            metadata_cache.emplace(key,
+                                   build_ranged_read_metadata_from_query_result(
+                                       key, query_result));
+        }
+    }
 
     for (size_t i = 0; i < buffer_count; ++i) {
         const size_t key_count = prepared.results[i].size();
@@ -2498,9 +3405,16 @@ RealClient::get_into_ranges_internal(
         }
 
         for (size_t j = 0; j < key_count; ++j) {
-            auto [metadata_it, inserted] = metadata_cache.try_emplace(
-                all_keys[i][j], resolve_ranged_read_metadata(all_keys[i][j]));
-            (void)inserted;
+            auto metadata_it = metadata_cache.find(all_keys[i][j]);
+            if (metadata_it == metadata_cache.end()) {
+                metadata_it =
+                    metadata_cache
+                        .emplace(
+                            all_keys[i][j],
+                            build_ranged_read_metadata_from_query_result(
+                                all_keys[i][j], client_->Query(all_keys[i][j])))
+                        .first;
+            }
             auto &metadata_result = metadata_it->second;
 
             for (size_t k = 0; k < prepared.results[i][j].size(); ++k) {
@@ -2552,13 +3466,14 @@ std::vector<std::vector<std::vector<int64_t>>> RealClient::get_into_ranges(
     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) {
+    const std::vector<std::vector<std::vector<size_t>>> &all_sizes,
+    const QueryResultCache *query_result_cache) {
     auto results =
         execute_timed_operation<std::vector<std::vector<std::vector<int64_t>>>>(
             [&]() {
-                return convert_ranged_read_results(
-                    get_into_ranges_internal(buffers, all_keys, all_dst_offsets,
-                                             all_src_offsets, all_sizes));
+                return convert_ranged_read_results(get_into_ranges_internal(
+                    buffers, all_keys, all_dst_offsets, all_src_offsets,
+                    all_sizes, nullptr, nullptr, nullptr, query_result_cache));
             },
             [](const auto &) { return true; },
             [&](uint64_t latency_us, const auto &ret) {
@@ -2569,6 +3484,49 @@ std::vector<std::vector<std::vector<int64_t>>> RealClient::get_into_ranges(
     return results;
 }
 
+std::vector<tl::expected<QueryResult, ErrorCode>> RealClient::batch_query(
+    const std::vector<std::string> &keys) {
+    if (!client_) {
+        LOG(ERROR) << "Client is not initialized";
+        return std::vector<tl::expected<QueryResult, ErrorCode>>(
+            keys.size(), tl::unexpected(ErrorCode::INVALID_PARAMS));
+    }
+    return client_->BatchQuery(keys);
+}
+
+tl::expected<RealClient::RangedReadMetadata, ErrorCode>
+RealClient::build_ranged_read_metadata_from_query_result(
+    const std::string &key, tl::expected<QueryResult, ErrorCode> query_result) {
+    if (!query_result) {
+        if (query_result.error() == ErrorCode::OBJECT_NOT_FOUND ||
+            query_result.error() == ErrorCode::REPLICA_IS_NOT_READY) {
+            return tl::unexpected(query_result.error());
+        }
+        LOG(ERROR) << "Query failed for key: " << key
+                   << " with error: " << toString(query_result.error());
+        return tl::unexpected(query_result.error());
+    }
+
+    const auto &replica_list = query_result->replicas;
+    if (replica_list.empty()) {
+        LOG(ERROR) << "Internal error: replica_list is empty";
+        return tl::unexpected(ErrorCode::INVALID_PARAMS);
+    }
+
+    auto local_endpoints = client_->GetLocalEndpoints();
+    const auto *best_replica = SelectBestReplica(replica_list, local_endpoints);
+    if (!best_replica) {
+        LOG(ERROR) << "No usable replica for key: " << key;
+        return tl::unexpected(ErrorCode::INVALID_REPLICA);
+    }
+
+    auto query_value = std::move(query_result.value());
+    auto replica = *best_replica;
+    return RangedReadMetadata{.query_result = std::move(query_value),
+                              .replica = std::move(replica),
+                              .total_size = calculate_total_size(replica)};
+}
+
 std::string RealClient::get_hostname() const { return local_hostname; }
 
 std::vector<int> RealClient::batch_put_from(
@@ -2695,8 +3653,8 @@ std::vector<tl::expected<void, ErrorCode>> RealClient::batch_put_from_internal(
 tl::expected<void, ErrorCode> RealClient::put_from_internal(
     const std::string &key, void *buffer, size_t size,
     const ReplicateConfig &config) {
-    // NOTE: The buffer address must be previously registered with
-    // register_buffer() for zero-copy RDMA operations to work correctly
+    // NOTE: The buffer address must resolve to Store-managed registered
+    // memory for zero-copy RDMA operations to work correctly
     if (config.prefer_alloc_in_same_node) {
         LOG(ERROR) << "prefer_alloc_in_same_node is not supported.";
         return tl::unexpected(ErrorCode::INVALID_PARAMS);
@@ -3189,26 +4147,20 @@ std::vector<int64_t> RealClient::batch_get_into(
     return results;
 }
 
-std::vector<tl::expected<int64_t, ErrorCode>>
+async_simple::coro::Lazy<std::vector<tl::expected<int64_t, ErrorCode>>>
 RealClient::batch_get_into_dummy_helper(
     const std::vector<std::string> &keys,
     const std::vector<uint64_t> &dummy_buffers,
     const std::vector<size_t> &sizes, int32_t device_id,
     const UUID &client_id) {
-#ifdef USE_ASCEND_DIRECT
-    if (!ContextManager::getInstance().setCurrentContextByPhysicalId(
-            device_id)) {
-        LOG(ERROR) << "Failed to set context for physical device " << device_id;
-        return std::vector<tl::expected<int64_t, ErrorCode>>(
-            keys.size(), tl::unexpected(ErrorCode::INVALID_PARAMS));
-    }
-#endif
-
-    std::shared_lock<std::shared_mutex> lock(dummy_client_mutex_);
+    // Hold shared_lock for the entire operation to prevent SHM from being
+    // unmapped while batch_get_into_internal is using the translated buffers.
+    auto lock = std::make_shared<std::shared_lock<std::shared_mutex>>(
+        dummy_client_mutex_);
     auto it = shm_contexts_.find(client_id);
     if (it == shm_contexts_.end()) {
         LOG(ERROR) << "client_id=" << client_id << ", error=shm_not_mapped";
-        return std::vector<tl::expected<int64_t, ErrorCode>>(
+        co_return std::vector<tl::expected<int64_t, ErrorCode>>(
             keys.size(), tl::unexpected(ErrorCode::INVALID_PARAMS));
     }
     auto &context = it->second;
@@ -3216,10 +4168,45 @@ RealClient::batch_get_into_dummy_helper(
     auto buffers_result =
         map_dummy_addrs_to_real_ptrs(context, dummy_buffers, sizes, client_id);
     if (!buffers_result) {
-        return std::vector<tl::expected<int64_t, ErrorCode>>(
+        co_return std::vector<tl::expected<int64_t, ErrorCode>>(
             keys.size(), tl::unexpected(buffers_result.error()));
     }
-    return batch_get_into_internal(keys, buffers_result.value(), sizes);
+
+    // Run batch_get_into_internal (which may block on offload RPC + SSD I/O)
+    // on a dedicated thread pool so the coro_rpc IO thread stays free for ping.
+    //
+    // Pack everything the worker-thread call needs into a single heap-owned
+    // state object so the lambda captures only a raw pointer (trivially
+    // copyable, trivially destructible). Avoids GCC 11 coroutine + coro_io
+    // double-destruction interactions when non-trivial captures end up in
+    // both the OUTER coroutine frame and the post_helper chain.
+    struct CallState {
+        std::vector<std::string> keys;
+        std::vector<size_t> sizes;
+        std::vector<void *> buffers;
+        std::shared_ptr<std::shared_lock<std::shared_mutex>> lock;
+        int32_t device_id;
+    };
+    auto state = std::make_unique<CallState>();
+    state->keys = keys;
+    state->sizes = sizes;
+    state->buffers = std::move(buffers_result.value());
+    state->lock = std::move(lock);
+    state->device_id = device_id;
+
+    auto *s = state.get();
+    auto try_result = co_await coro_io::post([this, s]() {
+#ifdef USE_ASCEND_DIRECT
+        auto context_result =
+            set_context_if_needed(protocol, s->device_id, "batch_get worker");
+        if (!context_result) {
+            return std::vector<tl::expected<int64_t, ErrorCode>>(
+                s->keys.size(), tl::unexpected(context_result.error()));
+        }
+#endif
+        return batch_get_into_internal(s->keys, s->buffers, s->sizes);
+    });
+    co_return try_result.value();
 }
 
 std::vector<tl::expected<void, ErrorCode>>
@@ -3324,6 +4311,8 @@ RealClient::get_into_ranges_shm_helper(
     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,
+    const std::map<std::string, CachedQueryResultResponse>
+        &cached_query_results,
     int32_t device_id, const UUID &client_id) {
 #ifdef USE_ASCEND_DIRECT
     if (!ContextManager::getInstance().setCurrentContextByPhysicalId(
@@ -3367,10 +4356,13 @@ RealClient::get_into_ranges_shm_helper(
         return prepared.results;
     }
 
+    auto query_result_cache =
+        build_query_result_cache_from_cached_results(cached_query_results);
     return get_into_ranges_internal(
         real_buffers_result.value(), all_keys, all_dst_offsets, all_src_offsets,
         all_sizes, &prepared.required_buffer_sizes, &prepared.results,
-        &prepared.valid_fragments);
+        &prepared.valid_fragments,
+        query_result_cache.empty() ? nullptr : &query_result_cache);
 }
 
 std::vector<tl::expected<int64_t, ErrorCode>>
@@ -3411,11 +4403,20 @@ RealClient::batch_get_into_internal(const std::vector<std::string> &keys,
         std::vector<Slice> slices;
         uint64_t total_size;
     };
+    struct DiskKeyInfo {
+        std::string key;
+        size_t original_index;
+        QueryResult query_result;
+        void *dst_buffer;
+        uint64_t total_size;
+    };
 
     std::vector<ValidKeyInfo> valid_operations;
     std::unordered_map<std::string, ValidKeyInfo> valid_local_disk_operations;
+    std::vector<DiskKeyInfo> disk_operations;
     valid_operations.reserve(num_keys);
 
+    auto local_endpoints = client_->GetLocalEndpoints();
     for (size_t i = 0; i < num_keys; ++i) {
         const auto &key = keys[i];
 
@@ -3439,8 +4440,18 @@ RealClient::batch_get_into_internal(const std::vector<std::string> &keys,
             continue;
         }
 
+        // Select best replica: prefer local MEMORY, then any MEMORY,
+        // then LOCAL_DISK, then DISK.
+        const auto *best_replica =
+            SelectBestReplica(query_result_values.replicas, local_endpoints);
+        if (!best_replica) {
+            LOG(ERROR) << "No usable replica for key: " << key;
+            results[i] = tl::unexpected(ErrorCode::INVALID_REPLICA);
+            continue;
+        }
+
         // Calculate required buffer size
-        const auto &replica = query_result_values.replicas[0];
+        const auto replica = *best_replica;
         uint64_t total_size = calculate_total_size(replica);
 
         // Validate buffer capacity
@@ -3452,12 +4463,9 @@ RealClient::batch_get_into_internal(const std::vector<std::string> &keys,
             continue;
         }
 
-        // Create slices for this key's buffer
-        std::vector<Slice> key_slices;
-        allocateSlices(key_slices, replica, buffers[i]);
-
-        if (query_result_values.replicas.size() == 1 &&
-            query_result_values.replicas.at(0).is_local_disk_replica()) {
+        if (replica.is_local_disk_replica()) {
+            std::vector<Slice> key_slices;
+            allocateSlices(key_slices, replica, buffers[i]);
             valid_local_disk_operations.emplace(
                 key,
                 ValidKeyInfo{.key = key,
@@ -3468,11 +4476,25 @@ RealClient::batch_get_into_internal(const std::vector<std::string> &keys,
             results[i] = static_cast<int64_t>(total_size);
             continue;
         }
-        // Store operation info for batch processing
+        if (replica.is_disk_replica()) {
+            // DISK: file I/O (vector_read) cannot write to user GPU buffer.
+            // Defer — allocate CPU temp buffer, BatchGet, then scatter.
+            disk_operations.emplace_back(
+                DiskKeyInfo{.key = key,
+                            .original_index = i,
+                            .query_result = std::move(query_result_values),
+                            .dst_buffer = buffers[i],
+                            .total_size = total_size});
+            results[i] = static_cast<int64_t>(total_size);
+            continue;
+        }
+        // MEMORY: RDMA directly to user buffer.
+        std::vector<Slice> key_slices;
+        allocateSlices(key_slices, replica, buffers[i]);
         valid_operations.push_back(
             {.key = key,
              .original_index = i,
-             .query_result = std::move(query_result_values),
+             .query_result = FilterQueryResult(query_result_values, replica),
              .slices = std::move(key_slices),
              .total_size = total_size});
 
@@ -3481,7 +4503,8 @@ RealClient::batch_get_into_internal(const std::vector<std::string> &keys,
     }
 
     // Early return if no valid operations
-    if (valid_operations.empty() && valid_local_disk_operations.empty()) {
+    if (valid_operations.empty() && valid_local_disk_operations.empty() &&
+        disk_operations.empty()) {
         return results;
     }
 
@@ -3516,16 +4539,103 @@ RealClient::batch_get_into_internal(const std::vector<std::string> &keys,
         }
     }
 
+    // ---- DISK replicas: BatchGet into CPU temp buffers, then scatter ----
+    if (!disk_operations.empty()) {
+        std::vector<std::string> disk_batch_keys;
+        std::vector<QueryResult> disk_batch_qrs;
+        std::unordered_map<std::string, std::vector<Slice>> disk_batch_slices;
+        std::vector<size_t> disk_batch_indices;
+        std::unordered_map<std::string, std::unique_ptr<BufferHandle>>
+            disk_temp_handles;
+
+        for (size_t di = 0; di < disk_operations.size(); ++di) {
+            auto &op = disk_operations[di];
+            // Find the DISK replica.
+            const Replica::Descriptor *replica_ptr = nullptr;
+            for (const auto &r : op.query_result.replicas) {
+                if (r.is_disk_replica()) {
+                    replica_ptr = &r;
+                    break;
+                }
+            }
+            if (!replica_ptr) {
+                LOG(ERROR) << "No DISK replica found for key: " << op.key;
+                results[op.original_index] =
+                    tl::unexpected(ErrorCode::INVALID_REPLICA);
+                continue;
+            }
+            auto alloc_result =
+                client_buffer_allocator_->allocate(op.total_size);
+            if (!alloc_result) {
+                LOG(ERROR) << "Failed to allocate temp buffer for DISK "
+                           << "read, key: " << op.key
+                           << ", size: " << op.total_size;
+                results[op.original_index] =
+                    tl::unexpected(ErrorCode::NO_AVAILABLE_HANDLE);
+                continue;
+            }
+            auto handle =
+                std::make_unique<BufferHandle>(std::move(*alloc_result));
+            std::vector<Slice> disk_slices;
+            allocateSlices(disk_slices, *replica_ptr, handle->ptr());
+            disk_batch_keys.push_back(op.key);
+            disk_batch_qrs.push_back(
+                FilterQueryResult(op.query_result, *replica_ptr));
+            disk_batch_slices[op.key] = std::move(disk_slices);
+            disk_batch_indices.push_back(di);
+            disk_temp_handles.emplace(op.key, std::move(handle));
+        }
+
+        if (!disk_batch_keys.empty()) {
+            auto disk_results = client_->BatchGet(
+                disk_batch_keys, disk_batch_qrs, disk_batch_slices);
+
+            for (size_t di = 0; di < disk_batch_indices.size(); ++di) {
+                const auto &key = disk_batch_keys[di];
+                auto &op = disk_operations[disk_batch_indices[di]];
+                auto handle_it = disk_temp_handles.find(key);
+                if (!disk_results[di]) {
+                    LOG(ERROR) << "DISK BatchGet failed for key '" << key
+                               << "': " << toString(disk_results[di].error());
+                    results[op.original_index] =
+                        tl::unexpected(disk_results[di].error());
+                    continue;
+                }
+                if (auto r = scatter_host_to_maybe_device(
+                        op.dst_buffer,
+                        static_cast<char *>(handle_it->second->ptr()),
+                        op.total_size, "DISK read, key: " + key);
+                    !r) {
+                    results[op.original_index] = tl::make_unexpected(r.error());
+                }
+            }
+        }
+    }
+
     // Prepare batch transfer data structures
-    std::unordered_map<std::string, std::unordered_map<std::string, Slice>>
+    std::unordered_map<std::string,
+                       std::unordered_map<std::string, std::vector<Slice>>>
         offload_objects;
 
     for (const auto &op_it : valid_local_disk_operations) {
-        const auto &replica = op_it.second.query_result.replicas.at(0);
+        // Find the LOCAL_DISK replica from the list — replicas may be in
+        // any order from Master.
+        const Replica::Descriptor *replica_ptr = nullptr;
+        for (const auto &r : op_it.second.query_result.replicas) {
+            if (r.is_local_disk_replica()) {
+                replica_ptr = &r;
+                break;
+            }
+        }
+        if (!replica_ptr) {
+            LOG(ERROR) << "No LOCAL_DISK replica found for key: "
+                       << op_it.first;
+            continue;
+        }
+        const auto &replica = *replica_ptr;
         auto [store_segment_it, _] = offload_objects.try_emplace(
             replica.get_local_disk_descriptor().transport_endpoint);
-        store_segment_it->second.emplace(op_it.first,
-                                         op_it.second.slices.at(0));
+        store_segment_it->second.emplace(op_it.first, op_it.second.slices);
     }
 
     size_t offload_object_count = 0;
@@ -3553,10 +4663,10 @@ RealClient::batch_get_into_internal(const std::vector<std::string> &keys,
         std::chrono::duration_cast<std::chrono::microseconds>(
             end_time - start_read_store_time)
             .count();
-    LOG(INFO) << "Time taken for batch_get_into: " << elapsed_time
-              << "us, read store: " << read_store_time
-              << "us, with memory key count: " << valid_operations.size()
-              << ", offload key count: " << offload_object_count;
+    // LOG(INFO) << "Time taken for batch_get_into: " << elapsed_time
+    //           << "us, read store: " << read_store_time
+    //           << "us, with memory key count: " << valid_operations.size()
+    //           << ", offload key count: " << offload_object_count;
 
     return results;
 }
@@ -3583,8 +4693,8 @@ int RealClient::put_from_with_metadata(const std::string &key, void *buffer,
                                        size_t metadata_size,
                                        const ReplicateConfig &config) {
     const auto start_time = std::chrono::steady_clock::now();
-    // NOTE: The buffer address must be previously registered with
-    // register_buffer() for zero-copy RDMA operations to work correctly
+    // NOTE: The buffer address must resolve to Store-managed registered
+    // memory for zero-copy RDMA operations to work correctly
     if (config.prefer_alloc_in_same_node) {
         LOG(ERROR) << "prefer_alloc_in_same_node is not supported.";
         return -1;
@@ -3772,8 +4882,21 @@ RealClient::batch_get_into_multi_buffers_internal(
         uint64_t total_size;
     };
 
+    struct DiskKeyInfo {
+        std::string key;
+        size_t original_index;
+        QueryResult query_result;
+        std::vector<void *> buffers;
+        std::vector<size_t> sizes;
+        uint64_t total_size;
+        bool is_local_disk;  // true=LOCAL_DISK (offload RPC), false=DISK
+                             // (BatchGet)
+    };
+
     std::vector<ValidKeyInfo> valid_operations;
+    std::unordered_map<std::string, DiskKeyInfo> valid_local_disk_ops;
     valid_operations.reserve(num_keys);
+    auto local_endpoints = client_->GetLocalEndpoints();
     for (size_t i = 0; i < num_keys; ++i) {
         const auto &key = keys[i];
         // Handle query failures
@@ -3793,8 +4916,17 @@ RealClient::batch_get_into_multi_buffers_internal(
             results.emplace_back(tl::unexpected(ErrorCode::INVALID_REPLICA));
             continue;
         }
-        // Calculate required buffer size
-        const auto &replica = query_result_values.replicas[0];
+        // Select best replica: prefer MEMORY (direct RDMA to GPU), then
+        // LOCAL_DISK, then DISK. Master may return multiple replicas in any
+        // order, so always scan rather than blindly taking replicas[0].
+        const auto *best_replica =
+            SelectBestReplica(query_result_values.replicas, local_endpoints);
+        if (!best_replica) {
+            LOG(ERROR) << "No usable replica for key: " << key;
+            results.emplace_back(tl::unexpected(ErrorCode::INVALID_REPLICA));
+            continue;
+        }
+        const auto replica = *best_replica;
         uint64_t total_size = calculate_total_size(replica);
         const auto &sizes = all_sizes[i];
         uint64_t dst_total_size = 0;
@@ -3813,11 +4945,28 @@ RealClient::batch_get_into_multi_buffers_internal(
         std::vector<Slice> key_slices;
         key_slices.reserve(buffers.size());
         if (replica.is_memory_replica()) {
+            // MEMORY: RDMA from remote memory directly to GPU (GPUDirect).
             for (size_t j = 0; j < buffers.size(); ++j) {
                 key_slices.emplace_back(Slice{buffers[j], sizes[j]});
             }
+        } else if (replica.is_local_disk_replica() ||
+                   replica.is_disk_replica()) {
+            // LOCAL_DISK: GPU buffers passed directly as scatter-gather slices
+            // (zero-copy). DISK: file I/O cannot write to GPU memory; temp CPU
+            // buffer used at read time.
+            valid_local_disk_ops.emplace(
+                key,
+                DiskKeyInfo{.key = key,
+                            .original_index = i,
+                            .query_result = std::move(query_result_values),
+                            .buffers = all_buffers[i],
+                            .sizes = all_sizes[i],
+                            .total_size = total_size,
+                            .is_local_disk = replica.is_local_disk_replica()});
+            results.emplace_back(static_cast<int64_t>(total_size));
+            continue;
         } else {
-            LOG(ERROR) << "Invalid replica type for key: " << key;
+            LOG(ERROR) << "Unsupported replica type for key: " << key;
             results.emplace_back(tl::unexpected(ErrorCode::INVALID_PARAMS));
             continue;
         }
@@ -3825,44 +4974,218 @@ RealClient::batch_get_into_multi_buffers_internal(
         valid_operations.push_back(
             {.key = key,
              .original_index = i,
-             .query_result = std::move(query_result_values),
+             .query_result = FilterQueryResult(query_result_values, replica),
              .slices = std::move(key_slices),
              .total_size = total_size});
         // Set success result (actual bytes transferred)
         results.emplace_back(static_cast<int64_t>(total_size));
     }
     // Early return if no valid operations
-    if (valid_operations.empty()) {
+    if (valid_operations.empty() && valid_local_disk_ops.empty()) {
         return results;
     }
 
-    // Prepare batch transfer data structures
-    std::vector<std::string> batch_keys;
-    std::vector<QueryResult> batch_query_results;
-    std::unordered_map<std::string, std::vector<Slice>> batch_slices;
-    batch_keys.reserve(valid_operations.size());
-    batch_query_results.reserve(valid_operations.size());
-    for (auto &op : valid_operations) {
-        batch_keys.push_back(op.key);
-        batch_query_results.push_back(op.query_result);
-        batch_slices[op.key] = op.slices;
-    }
+    // ---- Memory/Disk replica: existing BatchGet path ----
+    if (!valid_operations.empty()) {
+        std::vector<std::string> batch_keys;
+        std::vector<QueryResult> batch_query_results;
+        std::unordered_map<std::string, std::vector<Slice>> batch_slices;
+        batch_keys.reserve(valid_operations.size());
+        batch_query_results.reserve(valid_operations.size());
+        for (auto &op : valid_operations) {
+            batch_keys.push_back(op.key);
+            batch_query_results.push_back(op.query_result);
+            batch_slices[op.key] = op.slices;
+        }
 
-    auto batch_get_results =
-        client_->BatchGet(batch_keys, batch_query_results, batch_slices,
-                          prefer_alloc_in_same_node);
+        auto batch_get_results =
+            client_->BatchGet(batch_keys, batch_query_results, batch_slices,
+                              prefer_alloc_in_same_node);
 
-    // Process transfer results
-    for (size_t j = 0; j < batch_get_results.size(); ++j) {
-        const auto &op = valid_operations[j];
-
-        if (!batch_get_results[j]) {
-            const auto error = batch_get_results[j].error();
-            LOG(ERROR) << "BatchGet failed for key '" << op.key
-                       << "': " << toString(error);
-            results[op.original_index] = tl::unexpected(error);
+        for (size_t j = 0; j < batch_get_results.size(); ++j) {
+            const auto &op = valid_operations[j];
+            if (!batch_get_results[j]) {
+                const auto error = batch_get_results[j].error();
+                LOG(ERROR) << "BatchGet failed for key '" << op.key
+                           << "': " << toString(error);
+                results[op.original_index] = tl::unexpected(error);
+            }
         }
     }
+
+    // ---- LOCAL_DISK / DISK replica: disk read paths ----
+    if (!valid_local_disk_ops.empty()) {
+        // LOCAL_DISK: pass user GPU buffers directly as scatter-gather slices.
+        // vLLM pre-registers all GPU KV-cache memory with TransferEngine via
+        // register_buffer(), so the offload RDMA can scatter the on-disk blob
+        // into non-contiguous per-layer GPU destinations natively — no temp
+        // CPU buffer or H2D copy needed.
+        {
+            std::unordered_map<
+                std::string,
+                std::unordered_map<std::string, std::vector<Slice>>>
+                offload_objects;
+
+            for (auto &[key, op] : valid_local_disk_ops) {
+                if (!op.is_local_disk) continue;
+                // Find the correct LOCAL_DISK replica — Master may return
+                // replicas in any order (e.g. [DISK, LOCAL_DISK]).
+                const Replica::Descriptor *replica_ptr = nullptr;
+                for (const auto &r : op.query_result.replicas) {
+                    if (r.is_local_disk_replica()) {
+                        replica_ptr = &r;
+                        break;
+                    }
+                }
+                if (!replica_ptr) {
+                    LOG(ERROR)
+                        << "No LOCAL_DISK replica found for key: " << key;
+                    results[op.original_index] =
+                        tl::make_unexpected(ErrorCode::INVALID_REPLICA);
+                    continue;
+                }
+                const auto &replica = *replica_ptr;
+                std::vector<Slice> user_slices;
+                user_slices.reserve(op.buffers.size());
+                size_t slice_total = 0;
+                for (size_t j = 0; j < op.buffers.size(); ++j) {
+                    user_slices.push_back(Slice{op.buffers[j], op.sizes[j]});
+                    slice_total += op.sizes[j];
+                }
+                if (slice_total < op.total_size) {
+                    LOG(ERROR) << "Slice size too small for key " << key
+                               << ": slices=" << slice_total
+                               << ", total=" << op.total_size;
+                    results[op.original_index] =
+                        tl::make_unexpected(ErrorCode::INVALID_PARAMS);
+                    continue;
+                }
+                offload_objects[replica.get_local_disk_descriptor()
+                                    .transport_endpoint]
+                    .emplace(key, std::move(user_slices));
+            }
+
+            for (auto &[endpoint, objects] : offload_objects) {
+                if (objects.empty()) continue;
+                auto read_result =
+                    batch_get_into_offload_object_internal(endpoint, objects);
+                // On success: results[original_index] was already pre-filled
+                // with total_size when valid_local_disk_ops was built; nothing
+                // to update. Only overwrite on failure.
+                if (!read_result) {
+                    for (auto &[key, slices] : objects) {
+                        auto disk_it = valid_local_disk_ops.find(key);
+                        if (disk_it == valid_local_disk_ops.end()) continue;
+                        LOG(ERROR) << "SSD read failed for key '" << key
+                                   << "': " << toString(read_result.error());
+                        results[disk_it->second.original_index] =
+                            tl::make_unexpected(read_result.error());
+                    }
+                }
+            }
+        }
+
+        // DISK: one batched BatchGet into CPU temp buffers, then scatter.
+        // (storage_backend::vector_read cannot write directly to GPU memory)
+        {
+            // Scatter temp CPU buffer -> user multi_buffers (GPU or host).
+            // Returns false and sets results[original_index] on error.
+            auto scatter_to_buffers = [&](const std::string &key, char *src,
+                                          const DiskKeyInfo &op) -> bool {
+                size_t offset = 0;
+                for (size_t j = 0; j < op.buffers.size(); ++j) {
+                    if (offset >= op.total_size) break;
+                    size_t sz =
+                        std::min(op.sizes[j],
+                                 static_cast<size_t>(op.total_size - offset));
+                    void *dst = op.buffers[j];
+                    if (auto r = scatter_host_to_maybe_device(
+                            dst, src + offset, sz, "DISK scatter, key: " + key);
+                        !r) {
+                        results[op.original_index] =
+                            tl::make_unexpected(r.error());
+                        return false;
+                    }
+                    offset += sz;
+                }
+                return true;
+            };
+
+            std::vector<std::string> disk_batch_keys;
+            std::vector<QueryResult> disk_batch_qrs;
+            std::unordered_map<std::string, std::vector<Slice>>
+                disk_batch_slices;
+            std::vector<std::string> disk_key_order;
+            std::unordered_map<std::string, std::unique_ptr<BufferHandle>>
+                temp_handles;
+
+            for (auto &[key, op] : valid_local_disk_ops) {
+                if (op.is_local_disk) continue;
+                auto alloc_result =
+                    client_buffer_allocator_->allocate(op.total_size);
+                if (!alloc_result) {
+                    LOG(ERROR)
+                        << "Failed to allocate temp buffer for DISK "
+                        << "read, key: " << key << ", size: " << op.total_size;
+                    results[op.original_index] =
+                        tl::make_unexpected(ErrorCode::NO_AVAILABLE_HANDLE);
+                    continue;
+                }
+                auto handle =
+                    std::make_unique<BufferHandle>(std::move(*alloc_result));
+                // Find the correct DISK replica — Master may return
+                // replicas in any order (e.g. [LOCAL_DISK, DISK]).
+                const Replica::Descriptor *replica_ptr = nullptr;
+                for (const auto &r : op.query_result.replicas) {
+                    if (r.is_disk_replica()) {
+                        replica_ptr = &r;
+                        break;
+                    }
+                }
+                if (!replica_ptr) {
+                    LOG(ERROR) << "No DISK replica found for key: " << key;
+                    results[op.original_index] =
+                        tl::make_unexpected(ErrorCode::INVALID_REPLICA);
+                    continue;
+                }
+                const auto &replica = *replica_ptr;
+                std::vector<Slice> disk_slices;
+                allocateSlices(disk_slices, replica, handle->ptr());
+                disk_batch_keys.push_back(key);
+                disk_batch_qrs.push_back(
+                    FilterQueryResult(op.query_result, *replica_ptr));
+                disk_batch_slices[key] = std::move(disk_slices);
+                disk_key_order.push_back(key);
+                temp_handles.emplace(key, std::move(handle));
+            }
+
+            if (!disk_batch_keys.empty()) {
+                auto disk_results = client_->BatchGet(
+                    disk_batch_keys, disk_batch_qrs, disk_batch_slices,
+                    prefer_alloc_in_same_node);
+
+                for (size_t di = 0; di < disk_key_order.size(); ++di) {
+                    const auto &key = disk_key_order[di];
+                    auto &op = valid_local_disk_ops.at(key);
+                    auto handle_it = temp_handles.find(key);
+                    if (!disk_results[di]) {
+                        LOG(ERROR)
+                            << "DISK BatchGet failed for key '" << key
+                            << "': " << toString(disk_results[di].error());
+                        results[op.original_index] =
+                            tl::make_unexpected(disk_results[di].error());
+                        continue;
+                    }
+                    if (!scatter_to_buffers(
+                            key, static_cast<char *>(handle_it->second->ptr()),
+                            op))
+                        continue;
+                }
+            }
+            // temp_handles: BufferHandle RAII releases allocator memory
+        }
+    }
+
     return results;
 }
 
@@ -3945,6 +5268,13 @@ int RealClient::start_dummy_client_monitor() {
     return 0;
 }
 
+void RealClient::stop_dummy_client_monitor() {
+    dummy_client_monitor_running_ = false;
+    if (dummy_client_monitor_thread_.joinable()) {
+        dummy_client_monitor_thread_.join();
+        LOG(INFO) << "dummy_client_monitor_thread stopped";
+    }
+}
 int RealClient::start_ipc_server() {
     ipc_running_ = true;
     ipc_thread_ = std::jthread(&RealClient::ipc_server_func, this);
@@ -4053,8 +5383,9 @@ void RealClient::handle_ipc_shm_register(int client_sock) {
     }
 
     UUID client_id{req.client_id_first, req.client_id_second};
-    auto result = map_shm_internal(fd, req.dummy_base_addr, req.shm_size,
-                                   req.is_local_buffer, client_id);
+    auto result = map_shm_internal_with_device(
+        fd, req.dummy_base_addr, req.shm_size, req.is_local_buffer,
+        req.device_id, client_id);
 
     int status = result.has_value() ? 0 : -1;
     ::send(client_sock, &status, sizeof(status), 0);
@@ -4156,6 +5487,32 @@ RealClient::batch_get_replica_desc(const std::vector<std::string> &keys) {
     return replica_map;
 }
 
+std::vector<CachedQueryResultResponse> RealClient::batch_get_query_results(
+    const std::vector<std::string> &keys) {
+    if (!client_) {
+        LOG(ERROR) << "batch_get_query_results: client not initialized";
+        return std::vector<CachedQueryResultResponse>(
+            keys.size(), CachedQueryResultResponse(ErrorCode::INVALID_PARAMS));
+    }
+    auto query_results = client_->BatchQuery(keys);
+    std::vector<CachedQueryResultResponse> cached_results;
+    cached_results.reserve(keys.size());
+    if (query_results.size() != keys.size()) {
+        LOG(ERROR) << "Batch query response size mismatch in "
+                      "batch_get_query_results: expected "
+                   << keys.size() << ", got " << query_results.size() << ".";
+        return std::vector<CachedQueryResultResponse>(
+            keys.size(), CachedQueryResultResponse(ErrorCode::RPC_FAIL));
+    }
+
+    const auto now = std::chrono::steady_clock::now();
+    for (const auto &query_result : query_results) {
+        cached_results.push_back(
+            to_cached_query_result_response(query_result, now));
+    }
+    return cached_results;
+}
+
 std::vector<std::string> RealClient::batch_replica_clear(
     const std::vector<std::string> &keys, const std::string &segment_name) {
     if (!client_) {
@@ -4186,18 +5543,33 @@ tl::expected<QueryTaskResponse, ErrorCode> RealClient::query_task(
     const UUID &task_id) {
     return client_->QueryTask(task_id);
 }
-tl::expected<BatchGetOffloadObjectResponse, ErrorCode>
+async_simple::coro::Lazy<tl::expected<BatchGetOffloadObjectResponse, ErrorCode>>
 RealClient::batch_get_offload_object(const std::vector<std::string> &keys,
                                      const std::vector<int64_t> &sizes) {
-    auto result = file_storage_->BatchGet(keys, sizes);
+    // Run SSD I/O on a dedicated thread pool so the coro_rpc IO thread is
+    // free to handle ping and other RPCs. Same heap-owned state pattern as
+    // batch_get_into_dummy_helper: the lambda captures only a raw pointer.
+    struct CallState {
+        std::vector<std::string> keys;
+        std::vector<int64_t> sizes;
+        std::shared_ptr<FileStorage> file_storage;
+    };
+    auto state = std::make_unique<CallState>();
+    state->keys = keys;
+    state->sizes = sizes;
+    state->file_storage = file_storage_;
+    auto *s = state.get();
+    auto try_result = co_await coro_io::post(
+        [s]() { return s->file_storage->BatchGet(s->keys, s->sizes); });
+    auto result = try_result.value();
     if (!result) {
         LOG(ERROR) << "Batch get offload object failed,err_code = "
                    << result.error();
-        return tl::make_unexpected(result.error());
+        co_return tl::make_unexpected(result.error());
     }
-    return BatchGetOffloadObjectResponse(
+    co_return BatchGetOffloadObjectResponse(
         result.value().batch_id, std::move(result.value().pointers),
-        client_->GetTransportEndpoint(),
+        client_->GetSegmentEndpoint(),
         file_storage_->config_.client_buffer_gc_ttl_ms);
 }
 
@@ -4213,21 +5585,32 @@ bool RealClient::release_offload_buffer(uint64_t batch_id) {
 tl::expected<void, ErrorCode>
 RealClient::batch_get_into_offload_object_internal(
     const std::string &target_rpc_service_addr,
-    std::unordered_map<std::string, Slice> &objects) {
+    std::unordered_map<std::string, std::vector<Slice>> &objects) {
+    offload_rpc_read_count_.fetch_add(1, std::memory_order_relaxed);
     auto start_time = std::chrono::steady_clock::now();
     std::vector<std::string> keys;
+    std::vector<std::string> storage_keys;
     std::vector<int64_t> sizes;
     for (const auto &object_it : objects) {
         keys.emplace_back(object_it.first);
-        sizes.emplace_back(object_it.second.size);
+        storage_keys.emplace_back(
+            MakeTenantScopedStorageKey(client_->tenant_id(), object_it.first));
+        int64_t total = 0;
+        for (const auto &s : object_it.second) total += s.size;
+        sizes.emplace_back(total);
     }
     auto batchGetResp = client_requester_->batch_get_offload_object(
-        target_rpc_service_addr, keys, sizes);
+        target_rpc_service_addr, storage_keys, sizes);
     if (!batchGetResp) {
         LOG(ERROR) << "Batch get offload object failed with error: "
                    << batchGetResp.error();
         return tl::make_unexpected(batchGetResp.error());
     }
+    if (batchGetResp->pointers.size() != keys.size()) {
+        LOG(ERROR) << "Pointer count mismatch from owner: expected="
+                   << keys.size() << ", got=" << batchGetResp->pointers.size();
+        return tl::make_unexpected(ErrorCode::INVALID_PARAMS);
+    }
     auto result =
         client_->BatchGetOffloadObject(batchGetResp->transfer_engine_addr, keys,
                                        batchGetResp->pointers, objects);
diff --git a/mooncake-store/src/real_client_main.cpp b/mooncake-store/src/real_client_main.cpp
index 8bddca08..15c7202f 100644
--- a/mooncake-store/src/real_client_main.cpp
+++ b/mooncake-store/src/real_client_main.cpp
@@ -19,6 +19,11 @@ DEFINE_int32(port, 50052, "Real Client service port");
 DEFINE_string(global_segment_size, "4 GB", "Size of global segment");
 DEFINE_int32(threads, 1, "Number of threads for client service");
 DEFINE_bool(enable_offload, false, "Enable offload availability");
+DEFINE_bool(start_offload_rpc_server, true,
+            "Expose TCP RPC for disk-tier reads "
+            "(batch_get_offload_object / release_offload_buffer). "
+            "Effective only when --enable_offload is true. "
+            "Disable for a write-only owner.");
 DECLARE_bool(enable_http_server);
 DECLARE_int32(http_port);
 
@@ -61,6 +66,7 @@ void RegisterClientRpcService(coro_rpc::coro_rpc_server &server,
     server.register_handler<&RealClient::ascend_ipc_shm_internal>(&real_client);
     server.register_handler<&RealClient::ascend_unmap_shm_internal>(
         &real_client);
+    server.register_handler<&RealClient::is_shm_mapped_internal>(&real_client);
     server.register_handler<&RealClient::unmap_shm_internal>(&real_client);
     server.register_handler<&RealClient::unregister_shm_buffer_internal>(
         &real_client);
@@ -106,7 +112,7 @@ int main(int argc, char *argv[]) {
         FLAGS_host, FLAGS_metadata_server, global_segment_size, 0,
         FLAGS_protocol, FLAGS_device_names, FLAGS_master_server_address,
         nullptr, "@mooncake_client_" + std::to_string(FLAGS_port) + ".sock",
-        FLAGS_port, FLAGS_enable_offload);
+        FLAGS_port, FLAGS_enable_offload, FLAGS_start_offload_rpc_server);
     if (!res) {
         LOG(FATAL) << "Failed to setup client: " << toString(res.error());
         return -1;
diff --git a/mooncake-store/src/rpc_service.cpp b/mooncake-store/src/rpc_service.cpp
index 4c15549a..493a496c 100644
--- a/mooncake-store/src/rpc_service.cpp
+++ b/mooncake-store/src/rpc_service.cpp
@@ -1,21 +1,8 @@
 #include "rpc_service.h"
-#include <csignal>
 
-#include <ylt/struct_json/json_reader.h>
-#include <ylt/struct_json/json_writer.h>
-
-#include <atomic>
-#include <chrono>
-#include <cstdint>
-#include <sstream>
-#include <thread>
-#include <ylt/coro_http/coro_http_client.hpp>
-#include <ylt/coro_http/coro_http_server.hpp>
 #include <ylt/coro_rpc/coro_rpc_server.hpp>
-#include <ylt/reflection/user_reflect_macro.hpp>
 #include <ylt/util/tl/expected.hpp>
 
-#include "ha_metric_manager.h"
 #include "master_metric_manager.h"
 #include "master_service.h"
 #include "rpc_helper.h"
@@ -25,734 +12,45 @@
 
 namespace mooncake {
 
-const uint64_t kMetricReportIntervalSeconds = 10;
-
-namespace {
-
-std::string EscapeJson(std::string_view input) {
-    std::string escaped;
-    escaped.reserve(input.size());
-    for (char ch : input) {
-        switch (ch) {
-            case '\\':
-                escaped += "\\\\";
-                break;
-            case '"':
-                escaped += "\\\"";
-                break;
-            case '\n':
-                escaped += "\\n";
-                break;
-            case '\r':
-                escaped += "\\r";
-                break;
-            case '\t':
-                escaped += "\\t";
-                break;
-            default:
-                escaped.push_back(ch);
-                break;
-        }
-    }
-    return escaped;
-}
-
-std::string AppendMetricSections(std::string primary, std::string secondary) {
-    if (!primary.empty() && primary.back() != '\n') {
-        primary.push_back('\n');
-    }
-    primary += secondary;
-    return primary;
-}
-
-void SetServiceUnavailable(coro_http::coro_http_response& resp,
-                           std::string_view message) {
-    const std::string payload =
-        std::string("{\"success\":false,\"error_code\":") +
-        std::to_string(toInt(ErrorCode::UNAVAILABLE_IN_CURRENT_MODE)) +
-        ",\"error_message\":\"" + EscapeJson(message) + "\"}";
-    resp.add_header("Content-Type", "application/json; charset=utf-8");
-    resp.set_status_and_content(coro_http::status_type::service_unavailable,
-                                payload);
-}
-struct HttpErrorResponse {
-    bool success{false};
-    int32_t error_code{0};
-    std::string error_message;
-};
-YLT_REFL(HttpErrorResponse, success, error_code, error_message);
-
-struct HttpCreateDrainJobResponse {
-    bool success{false};
-    std::string job_id;
-    std::string status;
-    int32_t error_code{0};
-    std::string error_message;
-};
-YLT_REFL(HttpCreateDrainJobResponse, success, job_id, status, error_code,
-         error_message);
-
-struct HttpQueryDrainJobResponse {
-    bool success{false};
-    std::string job_id;
-    int32_t type{0};
-    std::string type_name;
-    int32_t status{0};
-    std::string status_name;
-    int64_t created_at_ms_epoch{0};
-    int64_t last_updated_at_ms_epoch{0};
-    std::vector<std::string> segments;
-    uint64_t succeeded_units{0};
-    uint64_t failed_units{0};
-    uint64_t blocked_units{0};
-    uint64_t active_units{0};
-    uint64_t migrated_bytes{0};
-    std::string message;
-    int32_t error_code{0};
-    std::string error_message;
-};
-YLT_REFL(HttpQueryDrainJobResponse, success, job_id, type, type_name, status,
-         status_name, created_at_ms_epoch, last_updated_at_ms_epoch, segments,
-         succeeded_units, failed_units, blocked_units, active_units,
-         migrated_bytes, message, error_code, error_message);
-
-struct HttpCancelDrainJobResponse {
-    bool success{false};
-    std::string job_id;
-    std::string status;
-    int32_t error_code{0};
-    std::string error_message;
-};
-YLT_REFL(HttpCancelDrainJobResponse, success, job_id, status, error_code,
-         error_message);
-
-struct HttpSegmentStatusResponse {
-    bool success{false};
-    std::string segment;
-    int32_t status{0};
-    std::string status_name;
-    int32_t error_code{0};
-    std::string error_message;
-};
-YLT_REFL(HttpSegmentStatusResponse, success, segment, status, status_name,
-         error_code, error_message);
-
-template <typename T>
-void WriteJsonResponse(coro_http::coro_http_response& resp,
-                       coro_http::status_type status, const T& payload) {
-    std::string json;
-    struct_json::to_json(payload, json);
-    resp.add_header("Content-Type", "application/json; charset=utf-8");
-    resp.set_status_and_content(status, std::move(json));
-}
-
-template <typename T>
-std::string EnumToString(const T& value) {
-    std::ostringstream oss;
-    oss << value;
-    return oss.str();
-}
-
-coro_http::status_type ErrorCodeToHttpStatus(ErrorCode error) {
-    switch (error) {
-        case ErrorCode::INVALID_PARAMS:
-            return coro_http::status_type::bad_request;
-        case ErrorCode::JOB_NOT_FOUND:
-        case ErrorCode::SEGMENT_NOT_FOUND:
-            return coro_http::status_type::not_found;
-        default:
-            return coro_http::status_type::internal_server_error;
-    }
-}
-
-void WriteErrorResponse(coro_http::coro_http_response& resp,
-                        coro_http::status_type status, ErrorCode error,
-                        std::string message = {}) {
-    HttpErrorResponse payload;
-    payload.error_code = toInt(error);
-    payload.error_message =
-        message.empty() ? toString(error) : std::move(message);
-    WriteJsonResponse(resp, status, payload);
-}
-
-tl::expected<UUID, ErrorCode> ParseJobId(std::string_view job_id_view) {
-    if (job_id_view.empty()) {
-        return tl::make_unexpected(ErrorCode::INVALID_PARAMS);
-    }
-
-    UUID job_id;
-    if (!StringToUuid(std::string(job_id_view), job_id)) {
-        return tl::make_unexpected(ErrorCode::INVALID_PARAMS);
-    }
-
-    return job_id;
-}
-
-HttpQueryDrainJobResponse ToHttpQueryDrainJobResponse(
-    const QueryJobResponse& job) {
-    HttpQueryDrainJobResponse payload;
-    payload.success = true;
-    payload.job_id = UuidToString(job.id);
-    payload.type = static_cast<int32_t>(job.type);
-    payload.type_name = EnumToString(job.type);
-    payload.status = static_cast<int32_t>(job.status);
-    payload.status_name = EnumToString(job.status);
-    payload.created_at_ms_epoch = job.created_at_ms_epoch;
-    payload.last_updated_at_ms_epoch = job.last_updated_at_ms_epoch;
-    payload.segments = job.segments;
-    payload.succeeded_units = job.succeeded_units;
-    payload.failed_units = job.failed_units;
-    payload.blocked_units = job.blocked_units;
-    payload.active_units = job.active_units;
-    payload.migrated_bytes = job.migrated_bytes;
-    payload.message = job.message;
-    return payload;
-}
-
-}  // namespace
-
 WrappedMasterService::WrappedMasterService(
-    const WrappedMasterServiceConfig& config)
-    : master_service_(MasterServiceConfig(config)) {}
+    const WrappedMasterServiceConfig& config,
+    HttpMetadataServer* http_metadata_server,
+    const std::string& http_metadata_remote_url)
+    : master_service_(MasterServiceConfig(config)) {
+    // Configure metadata cleanup on client timeout. Prefer the co-located
+    // in-process server; otherwise fall back to a separately-deployed HTTP
+    // metadata server derived from the cluster configuration.
+    if (http_metadata_server) {
+        master_service_.setHttpMetadataServer(http_metadata_server);
+    } else if (!http_metadata_remote_url.empty()) {
+        master_service_.setHttpMetadataRemoteUrl(http_metadata_remote_url);
+    }
+}
 
 WrappedMasterService::~WrappedMasterService() = default;
 
-MasterAdminServer::MasterAdminServer(uint16_t http_port,
-                                     bool enable_metric_reporting)
-    : http_port_(http_port),
-      enable_metric_reporting_(enable_metric_reporting),
-      http_server_(4, http_port) {}
-
-MasterAdminServer::~MasterAdminServer() { Stop(); }
-
-bool MasterAdminServer::Start() {
-    HAMetricManager::Init();
-    InitHttpServer();
-
-    auto ec = http_server_.async_start();
-    if (ec.hasResult()) {
-        LOG(ERROR) << "Failed to start master admin server on port "
-                   << http_port_;
-        return false;
-    }
-
-    started_.store(true);
-    if (enable_metric_reporting_) {
-        metric_report_running_.store(true);
-        metric_report_thread_ = std::thread([this]() {
-            while (metric_report_running_.load()) {
-                const auto snapshot = SnapshotState();
-                LOG(INFO) << "Master Admin Metrics: role="
-                          << ha::MasterRuntimeRoleToString(snapshot.state)
-                          << ", state="
-                          << ha::MasterRuntimeStateToString(snapshot.state)
-                          << ", summary=" << BuildMetricsSummaryText();
-                std::this_thread::sleep_for(
-                    std::chrono::seconds(kMetricReportIntervalSeconds));
-            }
-        });
-    }
-
-    LOG(INFO) << "Master admin server started on port " << http_server_.port();
-    return true;
-}
-
-void MasterAdminServer::Stop() {
-    metric_report_running_.store(false);
-    if (metric_report_thread_.joinable()) {
-        metric_report_thread_.join();
-    }
-    if (started_.exchange(false)) {
-        http_server_.stop();
-    }
-}
-
-void MasterAdminServer::SetRuntimeState(ha::MasterRuntimeState state) {
-    std::lock_guard<std::mutex> lock(state_mutex_);
-    state_ = state;
-}
-
-void MasterAdminServer::SetObservedLeader(
-    const std::optional<ha::MasterView>& leader_view) {
-    std::lock_guard<std::mutex> lock(state_mutex_);
-    leader_view_ = leader_view;
-}
-
-void MasterAdminServer::SetServiceDelegate(
-    std::shared_ptr<WrappedMasterService> service) {
-    std::lock_guard<std::mutex> lock(state_mutex_);
-    service_ = std::move(service);
-    if (!service_) {
-        service_available_ = false;
-    }
-}
-
-void MasterAdminServer::SetServiceAvailable(bool available) {
-    std::lock_guard<std::mutex> lock(state_mutex_);
-    service_available_ = available && service_ != nullptr;
-}
-
-MasterAdminServer::RuntimeSnapshot MasterAdminServer::SnapshotState() const {
-    std::lock_guard<std::mutex> lock(state_mutex_);
-    return RuntimeSnapshot{
-        .state = state_,
-        .leader_view = leader_view_,
-        .service = service_,
-        .service_available = service_available_,
-    };
-}
-
-std::string MasterAdminServer::BuildMetricsText() const {
-    return AppendMetricSections(
-        MasterMetricManager::instance().serialize_metrics(),
-        HAMetricManager::instance().serialize_metrics());
-}
-
-std::string MasterAdminServer::BuildMetricsSummaryText() const {
-    const auto snapshot = SnapshotState();
-    std::ostringstream oss;
-    oss << "role=" << ha::MasterRuntimeRoleToString(snapshot.state)
-        << ", state=" << ha::MasterRuntimeStateToString(snapshot.state)
-        << ", service_ready=" << (snapshot.service_available ? "true" : "false")
-        << ", master={" << MasterMetricManager::instance().get_summary_string()
-        << "}, ha={" << HAMetricManager::instance().get_summary_string() << "}";
-    if (snapshot.leader_view.has_value()) {
-        oss << ", leader=" << snapshot.leader_view->leader_address
-            << ", view_version=" << snapshot.leader_view->view_version;
-    }
-    return oss.str();
-}
-
-std::string MasterAdminServer::BuildHealthJson() const {
-    const auto snapshot = SnapshotState();
-    std::ostringstream oss;
-    oss << "{\"status\":\"ok\",\"role\":\""
-        << ha::MasterRuntimeRoleToString(snapshot.state) << "\",\"ha_state\":\""
-        << ha::MasterRuntimeStateToString(snapshot.state)
-        << "\",\"service_ready\":"
-        << (snapshot.service_available ? "true" : "false");
-    if (snapshot.leader_view.has_value()) {
-        oss << ",\"leader_address\":\""
-            << EscapeJson(snapshot.leader_view->leader_address)
-            << "\",\"view_version\":" << snapshot.leader_view->view_version;
-    }
-    oss << "}";
-    return oss.str();
-}
-
-std::string MasterAdminServer::BuildLeaderJson() const {
-    const auto snapshot = SnapshotState();
-    if (!snapshot.leader_view.has_value()) {
-        return "{\"present\":false}";
-    }
-
-    std::ostringstream oss;
-    oss << "{\"present\":true,\"leader_address\":\""
-        << EscapeJson(snapshot.leader_view->leader_address)
-        << "\",\"view_version\":" << snapshot.leader_view->view_version << "}";
-    return oss.str();
-}
-
-std::shared_ptr<WrappedMasterService> MasterAdminServer::GetActiveService()
-    const {
-    const auto snapshot = SnapshotState();
-    if (!snapshot.service_available) {
-        return nullptr;
-    }
-    return snapshot.service;
-}
-
-void MasterAdminServer::InitHttpServer() {
-    using namespace coro_http;
-
-    http_server_.set_http_handler<GET>(
-        "/metrics", [this](coro_http_request&, coro_http_response& resp) {
-            resp.add_header("Content-Type", "text/plain; version=0.0.4");
-            resp.set_status_and_content(status_type::ok, BuildMetricsText());
-        });
-
-    http_server_.set_http_handler<GET>(
-        "/metrics/summary",
-        [this](coro_http_request&, coro_http_response& resp) {
-            resp.add_header("Content-Type", "text/plain; version=0.0.4");
-            resp.set_status_and_content(status_type::ok,
-                                        BuildMetricsSummaryText());
-        });
-
-    http_server_.set_http_handler<GET>(
-        "/health", [this](coro_http_request&, coro_http_response& resp) {
-            resp.add_header("Content-Type", "application/json; charset=utf-8");
-            resp.set_status_and_content(status_type::ok, BuildHealthJson());
-        });
-
-    http_server_.set_http_handler<GET>(
-        "/role", [this](coro_http_request&, coro_http_response& resp) {
-            const auto snapshot = SnapshotState();
-            resp.add_header("Content-Type", "text/plain; charset=utf-8");
-            resp.set_status_and_content(
-                status_type::ok, ha::MasterRuntimeRoleToString(snapshot.state));
-        });
-
-    http_server_.set_http_handler<GET>(
-        "/ha_status", [this](coro_http_request&, coro_http_response& resp) {
-            const auto snapshot = SnapshotState();
-            resp.add_header("Content-Type", "text/plain; charset=utf-8");
-            resp.set_status_and_content(
-                status_type::ok,
-                ha::MasterRuntimeStateToString(snapshot.state));
-        });
-
-    http_server_.set_http_handler<GET>(
-        "/leader", [this](coro_http_request&, coro_http_response& resp) {
-            resp.add_header("Content-Type", "application/json; charset=utf-8");
-            resp.set_status_and_content(status_type::ok, BuildLeaderJson());
-        });
-
-    http_server_.set_http_handler<GET>(
-        "/query_key", [this](coro_http_request& req, coro_http_response& resp) {
-            auto service = GetActiveService();
-            if (!service) {
-                SetServiceUnavailable(resp, "service plane is not active");
-                return;
-            }
-
-            auto key = req.get_query_value("key");
-            auto get_result = service->GetReplicaList(std::string(key));
-            resp.add_header("Content-Type", "text/plain; version=0.0.4");
-            if (get_result) {
-                std::string ss;
-                const auto& replicas = get_result.value().replicas;
-                for (const auto& replica : replicas) {
-                    if (!replica.is_memory_replica()) {
-                        continue;
-                    }
-                    std::string tmp;
-                    struct_json::to_json(
-                        replica.get_memory_descriptor().buffer_descriptor, tmp);
-                    ss += tmp;
-                    ss += "\n";
-                }
-                resp.set_status_and_content(status_type::ok, std::move(ss));
-                return;
-            }
-
-            resp.set_status_and_content(status_type::not_found,
-                                        toString(get_result.error()));
-        });
-
-    http_server_.set_http_handler<GET>(
-        "/get_all_keys", [this](coro_http_request&, coro_http_response& resp) {
-            auto service = GetActiveService();
-            if (!service) {
-                SetServiceUnavailable(resp, "service plane is not active");
-                return;
-            }
-
-            resp.add_header("Content-Type", "text/plain; version=0.0.4");
-            auto result = service->GetAllKeysForAdmin();
-            if (!result) {
-                resp.set_status_and_content(status_type::internal_server_error,
-                                            "Failed to get all keys");
-                return;
-            }
-
-            std::string body;
-            for (const auto& key : result.value()) {
-                body += key;
-                body += "\n";
-            }
-            resp.set_status_and_content(status_type::ok, std::move(body));
-        });
-
-    http_server_.set_http_handler<GET>(
-        "/get_all_segments",
-        [this](coro_http_request&, coro_http_response& resp) {
-            auto service = GetActiveService();
-            if (!service) {
-                SetServiceUnavailable(resp, "service plane is not active");
-                return;
-            }
-
-            resp.add_header("Content-Type", "text/plain; version=0.0.4");
-            auto result = service->GetAllSegmentsForAdmin();
-            if (!result) {
-                resp.set_status_and_content(status_type::internal_server_error,
-                                            "Failed to get all segments");
-                return;
-            }
-
-            std::string body;
-            for (const auto& segment_name : result.value()) {
-                body += segment_name;
-                body += "\n";
-            }
-            resp.set_status_and_content(status_type::ok, std::move(body));
-        });
-
-    http_server_.set_http_handler<GET>(
-        "/query_segment",
-        [this](coro_http_request& req, coro_http_response& resp) {
-            auto service = GetActiveService();
-            if (!service) {
-                SetServiceUnavailable(resp, "service plane is not active");
-                return;
-            }
-
-            auto segment = req.get_query_value("segment");
-            resp.add_header("Content-Type", "text/plain; version=0.0.4");
-            auto result = service->QuerySegmentForAdmin(std::string(segment));
-
-            if (!result) {
-                resp.set_status_and_content(status_type::internal_server_error,
-                                            "Failed to query segment");
-                return;
-            }
-
-            const auto [used, capacity] = result.value();
-            std::string body(segment);
-            body += "\nUsed(bytes): ";
-            body += std::to_string(used);
-            body += "\nCapacity(bytes) : ";
-            body += std::to_string(capacity);
-            body += "\n";
-            resp.set_status_and_content(status_type::ok, std::move(body));
-        });
-
-    http_server_.set_http_handler<POST>(
-        "/api/v1/drain_jobs",
-        [&](coro_http_request& req, coro_http_response& resp) {
-            CreateDrainJobRequest request;
-            try {
-                struct_json::from_json(request, req.get_body());
-            } catch (const std::exception& e) {
-                WriteErrorResponse(
-                    resp, status_type::bad_request, ErrorCode::INVALID_PARAMS,
-                    std::string("Invalid JSON body: ") + e.what());
-                return;
-            }
-
-            auto service = GetActiveService();
-            if (!service) {
-                SetServiceUnavailable(resp, "service plane is not active");
-                return;
-            }
-
-            auto result = service->CreateDrainJob(request);
-            if (!result.has_value()) {
-                WriteErrorResponse(resp, ErrorCodeToHttpStatus(result.error()),
-                                   result.error());
-                return;
-            }
-
-            HttpCreateDrainJobResponse payload;
-            payload.success = true;
-            payload.job_id = UuidToString(result.value());
-            payload.status = "CREATED";
-            WriteJsonResponse(resp, status_type::ok, payload);
-        });
-
-    http_server_.set_http_handler<GET>(
-        "/api/v1/drain_jobs/query",
-        [&](coro_http_request& req, coro_http_response& resp) {
-            auto job_id_result =
-                ParseJobId(req.get_decode_query_value("job_id"));
-            if (!job_id_result.has_value()) {
-                WriteErrorResponse(resp, status_type::bad_request,
-                                   job_id_result.error(),
-                                   "Missing or invalid job_id");
-                return;
-            }
-
-            auto service = GetActiveService();
-            if (!service) {
-                SetServiceUnavailable(resp, "service plane is not active");
-                return;
-            }
-
-            auto result = service->QueryDrainJob(job_id_result.value());
-            if (!result.has_value()) {
-                WriteErrorResponse(resp, ErrorCodeToHttpStatus(result.error()),
-                                   result.error());
-                return;
-            }
-
-            WriteJsonResponse(resp, status_type::ok,
-                              ToHttpQueryDrainJobResponse(result.value()));
-        });
-
-    http_server_.set_http_handler<POST>(
-        "/api/v1/drain_jobs/cancel",
-        [&](coro_http_request& req, coro_http_response& resp) {
-            auto job_id_result =
-                ParseJobId(req.get_decode_query_value("job_id"));
-            if (!job_id_result.has_value()) {
-                WriteErrorResponse(resp, status_type::bad_request,
-                                   job_id_result.error(),
-                                   "Missing or invalid job_id");
-                return;
-            }
-
-            auto service = GetActiveService();
-            if (!service) {
-                SetServiceUnavailable(resp, "service plane is not active");
-                return;
-            }
-
-            auto result = service->CancelDrainJob(job_id_result.value());
-            if (!result.has_value()) {
-                WriteErrorResponse(resp, ErrorCodeToHttpStatus(result.error()),
-                                   result.error());
-                return;
-            }
-
-            HttpCancelDrainJobResponse payload;
-            payload.success = true;
-            payload.job_id = UuidToString(job_id_result.value());
-            payload.status = "CANCELED";
-            WriteJsonResponse(resp, status_type::ok, payload);
-        });
-
-    http_server_.set_http_handler<GET>(
-        "/api/v1/segments/status",
-        [&](coro_http_request& req, coro_http_response& resp) {
-            auto segment_name = req.get_decode_query_value("segment");
-            if (segment_name.empty()) {
-                WriteErrorResponse(resp, status_type::bad_request,
-                                   ErrorCode::INVALID_PARAMS,
-                                   "Missing segment query parameter");
-                return;
-            }
-
-            auto service = GetActiveService();
-            if (!service) {
-                SetServiceUnavailable(resp, "service plane is not active");
-                return;
-            }
-
-            auto result =
-                service->QuerySegmentStatus(std::string(segment_name));
-            if (!result.has_value()) {
-                WriteErrorResponse(resp, ErrorCodeToHttpStatus(result.error()),
-                                   result.error());
-                return;
-            }
-
-            HttpSegmentStatusResponse payload;
-            payload.success = true;
-            payload.segment = std::string(segment_name);
-            payload.status = static_cast<int32_t>(result.value());
-            payload.status_name = EnumToString(result.value());
-            WriteJsonResponse(resp, status_type::ok, payload);
-        });
-
-    http_server_.set_http_handler<GET>(
-        "/batch_query_keys",
-        [this](coro_http_request& req, coro_http_response& resp) {
-            auto service = GetActiveService();
-            if (!service) {
-                resp.add_header("Content-Type",
-                                "application/json; charset=utf-8");
-                resp.set_status_and_content(
-                    status_type::service_unavailable,
-                    "{\"success\":false,\"error\":\"service plane is not "
-                    "active\"}");
-                return;
-            }
-
-            auto keys_view = req.get_query_value("keys");
-            std::vector<std::string> keys;
-            if (!keys_view.empty()) {
-                std::string keys_str(keys_view);
-                std::string key;
-                std::istringstream iss(keys_str);
-                while (std::getline(iss, key, ',')) {
-                    keys.push_back(std::move(key));
-                }
-            }
-
-            resp.add_header("Content-Type", "application/json; charset=utf-8");
-            if (keys.empty()) {
-                resp.set_status_and_content(
-                    status_type::bad_request,
-                    "{\"success\":false,\"error\":\"No keys provided. Use "
-                    "?keys=key1,key2,...\"}");
-                return;
-            }
-
-            auto results = service->BatchGetReplicaList(keys);
-            const size_t n = std::min(keys.size(), results.size());
-            std::string body;
-            body.reserve(n * 512);
-            body += "{\"success\":true,\"data\":{";
-
-            for (size_t i = 0; i < n; ++i) {
-                if (i > 0) {
-                    body += ",";
-                }
-
-                body += "\"";
-                body += EscapeJson(keys[i]);
-                body += "\":";
-
-                const auto& result = results[i];
-                if (!result.has_value()) {
-                    body += "{\"ok\":false,\"error\":\"";
-                    body += EscapeJson(toString(result.error()));
-                    body += "\"}";
-                    continue;
-                }
-
-                body += "{\"ok\":true,\"values\":[";
-                bool first = true;
-                for (const auto& replica : result.value().replicas) {
-                    if (!replica.is_memory_replica()) {
-                        continue;
-                    }
-
-                    std::string tmp;
-                    struct_json::to_json(
-                        replica.get_memory_descriptor().buffer_descriptor, tmp);
-                    if (!first) {
-                        body += ",";
-                    }
-                    body += tmp;
-                    first = false;
-                }
-                body += "]}";
-            }
-
-            body += "}}";
-            if (results.size() != keys.size()) {
-                LOG(WARNING)
-                    << "BatchGetReplicaList size mismatch: keys=" << keys.size()
-                    << " results=" << results.size();
-            }
-            resp.set_status_and_content(status_type::ok, std::move(body));
-        });
-}
-
 tl::expected<MasterMetricManager::CacheHitStatDict, ErrorCode>
 WrappedMasterService::CalcCacheStats() {
     return MasterMetricManager::instance().calculate_cache_stats();
 }
 
 tl::expected<bool, ErrorCode> WrappedMasterService::ExistKey(
-    const std::string& key) {
+    const std::string& key, const std::string& tenant_id) {
     return execute_rpc(
-        "ExistKey", [&] { return master_service_.ExistKey(key); },
+        "ExistKey", [&] { return master_service_.ExistKey(key, tenant_id); },
         [&](auto& timer) { timer.LogRequest("key=", key); },
         [] { MasterMetricManager::instance().inc_exist_key_requests(); },
         [] { MasterMetricManager::instance().inc_exist_key_failures(); });
 }
 
 std::vector<tl::expected<bool, ErrorCode>> WrappedMasterService::BatchExistKey(
-    const std::vector<std::string>& keys) {
+    const std::vector<std::string>& keys, const std::string& tenant_id) {
     ScopedVLogTimer timer(1, "BatchExistKey");
     const size_t total_keys = keys.size();
     timer.LogRequest("keys_count=", total_keys);
     MasterMetricManager::instance().inc_batch_exist_key_requests(total_keys);
 
-    auto result = master_service_.BatchExistKey(keys);
+    auto result = master_service_.BatchExistKey(keys, tenant_id);
 
     size_t failure_count = 0;
     for (size_t i = 0; i < result.size(); ++i) {
@@ -859,10 +157,11 @@ WrappedMasterService::BatchReplicaClear(
 
 tl::expected<std::unordered_map<std::string, std::vector<Replica::Descriptor>>,
              ErrorCode>
-WrappedMasterService::GetReplicaListByRegex(const std::string& str) {
+WrappedMasterService::GetReplicaListByRegex(const std::string& str,
+                                            const std::string& tenant_id) {
     return execute_rpc(
         "GetReplicaListByRegex",
-        [&] { return master_service_.GetReplicaListByRegex(str); },
+        [&] { return master_service_.GetReplicaListByRegex(str, tenant_id); },
         [&](auto& timer) { timer.LogRequest("Regex=", str); },
         [] {
             MasterMetricManager::instance()
@@ -875,9 +174,11 @@ WrappedMasterService::GetReplicaListByRegex(const std::string& str) {
 }
 
 tl::expected<GetReplicaListResponse, ErrorCode>
-WrappedMasterService::GetReplicaList(const std::string& key) {
+WrappedMasterService::GetReplicaList(const std::string& key,
+                                     const std::string& tenant_id) {
     return execute_rpc(
-        "GetReplicaList", [&] { return master_service_.GetReplicaList(key); },
+        "GetReplicaList",
+        [&] { return master_service_.GetReplicaList(key, tenant_id); },
         [&](auto& timer) { timer.LogRequest("key=", key); },
         [] { MasterMetricManager::instance().inc_get_replica_list_requests(); },
         [] {
@@ -886,8 +187,8 @@ WrappedMasterService::GetReplicaList(const std::string& key) {
 }
 
 std::vector<tl::expected<GetReplicaListResponse, ErrorCode>>
-WrappedMasterService::BatchGetReplicaList(
-    const std::vector<std::string>& keys) {
+WrappedMasterService::BatchGetReplicaList(const std::vector<std::string>& keys,
+                                          const std::string& tenant_id) {
     ScopedVLogTimer timer(1, "BatchGetReplicaList");
     const size_t total_keys = keys.size();
     timer.LogRequest("keys_count=", total_keys);
@@ -897,9 +198,7 @@ WrappedMasterService::BatchGetReplicaList(
     std::vector<tl::expected<GetReplicaListResponse, ErrorCode>> results;
     results.reserve(keys.size());
 
-    for (const auto& key : keys) {
-        results.emplace_back(master_service_.GetReplicaList(key));
-    }
+    results = master_service_.BatchGetReplicaList(keys, tenant_id);
 
     size_t failure_count = 0;
     for (size_t i = 0; i < results.size(); ++i) {
@@ -934,12 +233,13 @@ WrappedMasterService::BatchGetReplicaList(
 tl::expected<std::vector<Replica::Descriptor>, ErrorCode>
 WrappedMasterService::PutStart(const UUID& client_id, const std::string& key,
                                const uint64_t slice_length,
-                               const ReplicateConfig& config) {
+                               const ReplicateConfig& config,
+                               const std::string& tenant_id) {
     return execute_rpc(
         "PutStart",
         [&] {
-            return master_service_.PutStart(client_id, key, slice_length,
-                                            config);
+            return master_service_.PutStart(client_id, key, tenant_id,
+                                            slice_length, config);
         },
         [&](auto& timer) {
             timer.LogRequest("client_id=", client_id, ", key=", key,
@@ -950,10 +250,14 @@ WrappedMasterService::PutStart(const UUID& client_id, const std::string& key,
 }
 
 tl::expected<void, ErrorCode> WrappedMasterService::PutEnd(
-    const UUID& client_id, const std::string& key, ReplicaType replica_type) {
+    const UUID& client_id, const std::string& key, ReplicaType replica_type,
+    const std::string& tenant_id) {
     return execute_rpc(
         "PutEnd",
-        [&] { return master_service_.PutEnd(client_id, key, replica_type); },
+        [&] {
+            return master_service_.PutEnd(client_id, key, tenant_id,
+                                          replica_type);
+        },
         [&](auto& timer) {
             timer.LogRequest("client_id=", client_id, ", key=", key,
                              ", replica_type=", replica_type);
@@ -963,10 +267,14 @@ tl::expected<void, ErrorCode> WrappedMasterService::PutEnd(
 }
 
 tl::expected<void, ErrorCode> WrappedMasterService::PutRevoke(
-    const UUID& client_id, const std::string& key, ReplicaType replica_type) {
+    const UUID& client_id, const std::string& key, ReplicaType replica_type,
+    const std::string& tenant_id) {
     return execute_rpc(
         "PutRevoke",
-        [&] { return master_service_.PutRevoke(client_id, key, replica_type); },
+        [&] {
+            return master_service_.PutRevoke(client_id, key, tenant_id,
+                                             replica_type);
+        },
         [&](auto& timer) {
             timer.LogRequest("client_id=", client_id, ", key=", key,
                              ", replica_type=", replica_type);
@@ -979,7 +287,8 @@ std::vector<tl::expected<std::vector<Replica::Descriptor>, ErrorCode>>
 WrappedMasterService::BatchPutStart(const UUID& client_id,
                                     const std::vector<std::string>& keys,
                                     const std::vector<uint64_t>& slice_lengths,
-                                    const ReplicateConfig& config) {
+                                    const ReplicateConfig& config,
+                                    const std::string& tenant_id) {
     ScopedVLogTimer timer(1, "BatchPutStart");
     const size_t total_keys = keys.size();
     timer.LogRequest("client_id=", client_id, ", keys_count=", total_keys);
@@ -989,11 +298,24 @@ WrappedMasterService::BatchPutStart(const UUID& client_id,
         results;
     results.reserve(keys.size());
 
-    if (config.prefer_alloc_in_same_node) {
+    if (keys.size() != slice_lengths.size()) {
+        LOG(ERROR) << "BatchPutStart: keys.size()=" << keys.size()
+                   << " != slice_lengths.size()=" << slice_lengths.size();
+        results.assign(keys.size(),
+                       tl::make_unexpected(ErrorCode::INVALID_PARAMS));
+    } else if (config.group_ids.has_value() &&
+               config.group_ids->size() != keys.size()) {
+        LOG(ERROR) << "BatchPutStart: group_ids.size()="
+                   << config.group_ids->size()
+                   << " != keys.size()=" << keys.size();
+        results.assign(keys.size(),
+                       tl::make_unexpected(ErrorCode::INVALID_PARAMS));
+    } else if (config.prefer_alloc_in_same_node) {
         ReplicateConfig new_config = config;
         for (size_t i = 0; i < keys.size(); ++i) {
+            auto key_config = new_config.ForSingleKey(i);
             auto result = master_service_.PutStart(
-                client_id, keys[i], slice_lengths[i], new_config);
+                client_id, keys[i], tenant_id, slice_lengths[i], key_config);
             results.emplace_back(result);
             if ((i == 0) && result.has_value()) {
                 std::string preferred_segment;
@@ -1013,8 +335,9 @@ WrappedMasterService::BatchPutStart(const UUID& client_id,
         }
     } else {
         for (size_t i = 0; i < keys.size(); ++i) {
+            auto key_config = config.ForSingleKey(i);
             results.emplace_back(master_service_.PutStart(
-                client_id, keys[i], slice_lengths[i], config));
+                client_id, keys[i], tenant_id, slice_lengths[i], key_config));
         }
     }
 
@@ -1056,7 +379,8 @@ WrappedMasterService::BatchPutStart(const UUID& client_id,
 }
 
 std::vector<tl::expected<void, ErrorCode>> WrappedMasterService::BatchPutEnd(
-    const UUID& client_id, const std::vector<std::string>& keys) {
+    const UUID& client_id, const std::vector<std::string>& keys,
+    ReplicaType replica_type, const std::string& tenant_id) {
     ScopedVLogTimer timer(1, "BatchPutEnd");
     const size_t total_keys = keys.size();
     timer.LogRequest("client_id=", client_id, ", keys_count=", total_keys);
@@ -1067,7 +391,7 @@ std::vector<tl::expected<void, ErrorCode>> WrappedMasterService::BatchPutEnd(
 
     for (const auto& key : keys) {
         results.emplace_back(
-            master_service_.PutEnd(client_id, key, ReplicaType::MEMORY));
+            master_service_.PutEnd(client_id, key, tenant_id, replica_type));
     }
 
     size_t failure_count = 0;
@@ -1095,7 +419,8 @@ std::vector<tl::expected<void, ErrorCode>> WrappedMasterService::BatchPutEnd(
 }
 
 std::vector<tl::expected<void, ErrorCode>> WrappedMasterService::BatchPutRevoke(
-    const UUID& client_id, const std::vector<std::string>& keys) {
+    const UUID& client_id, const std::vector<std::string>& keys,
+    ReplicaType replica_type, const std::string& tenant_id) {
     ScopedVLogTimer timer(1, "BatchPutRevoke");
     const size_t total_keys = keys.size();
     timer.LogRequest("client_id=", client_id, ", keys_count=", total_keys);
@@ -1106,7 +431,7 @@ std::vector<tl::expected<void, ErrorCode>> WrappedMasterService::BatchPutRevoke(
 
     for (const auto& key : keys) {
         results.emplace_back(
-            master_service_.PutRevoke(client_id, key, ReplicaType::MEMORY));
+            master_service_.PutRevoke(client_id, key, tenant_id, replica_type));
     }
 
     size_t failure_count = 0;
@@ -1136,12 +461,13 @@ std::vector<tl::expected<void, ErrorCode>> WrappedMasterService::BatchPutRevoke(
 tl::expected<std::vector<Replica::Descriptor>, ErrorCode>
 WrappedMasterService::UpsertStart(const UUID& client_id, const std::string& key,
                                   const uint64_t slice_length,
-                                  const ReplicateConfig& config) {
+                                  const ReplicateConfig& config,
+                                  const std::string& tenant_id) {
     return execute_rpc(
         "UpsertStart",
         [&] {
-            return master_service_.UpsertStart(client_id, key, slice_length,
-                                               config);
+            return master_service_.UpsertStart(client_id, key, tenant_id,
+                                               slice_length, config);
         },
         [&](auto& timer) {
             timer.LogRequest("client_id=", client_id, ", key=", key,
@@ -1152,10 +478,14 @@ WrappedMasterService::UpsertStart(const UUID& client_id, const std::string& key,
 }
 
 tl::expected<void, ErrorCode> WrappedMasterService::UpsertEnd(
-    const UUID& client_id, const std::string& key, ReplicaType replica_type) {
+    const UUID& client_id, const std::string& key, ReplicaType replica_type,
+    const std::string& tenant_id) {
     return execute_rpc(
         "UpsertEnd",
-        [&] { return master_service_.UpsertEnd(client_id, key, replica_type); },
+        [&] {
+            return master_service_.UpsertEnd(client_id, key, tenant_id,
+                                             replica_type);
+        },
         [&](auto& timer) {
             timer.LogRequest("client_id=", client_id, ", key=", key,
                              ", replica_type=", replica_type);
@@ -1165,11 +495,13 @@ tl::expected<void, ErrorCode> WrappedMasterService::UpsertEnd(
 }
 
 tl::expected<void, ErrorCode> WrappedMasterService::UpsertRevoke(
-    const UUID& client_id, const std::string& key, ReplicaType replica_type) {
+    const UUID& client_id, const std::string& key, ReplicaType replica_type,
+    const std::string& tenant_id) {
     return execute_rpc(
         "UpsertRevoke",
         [&] {
-            return master_service_.UpsertRevoke(client_id, key, replica_type);
+            return master_service_.UpsertRevoke(client_id, key, tenant_id,
+                                                replica_type);
         },
         [&](auto& timer) {
             timer.LogRequest("client_id=", client_id, ", key=", key,
@@ -1182,13 +514,14 @@ tl::expected<void, ErrorCode> WrappedMasterService::UpsertRevoke(
 std::vector<tl::expected<std::vector<Replica::Descriptor>, ErrorCode>>
 WrappedMasterService::BatchUpsertStart(
     const UUID& client_id, const std::vector<std::string>& keys,
-    const std::vector<uint64_t>& slice_lengths, const ReplicateConfig& config) {
+    const std::vector<uint64_t>& slice_lengths, const ReplicateConfig& config,
+    const std::string& tenant_id) {
     ScopedVLogTimer timer(1, "BatchUpsertStart");
     const size_t total_keys = keys.size();
     timer.LogRequest("client_id=", client_id, ", keys_count=", total_keys);
     MasterMetricManager::instance().inc_batch_put_start_requests(total_keys);
 
-    auto results = master_service_.BatchUpsertStart(client_id, keys,
+    auto results = master_service_.BatchUpsertStart(client_id, keys, tenant_id,
                                                     slice_lengths, config);
 
     size_t failure_count = 0;
@@ -1216,13 +549,14 @@ WrappedMasterService::BatchUpsertStart(
 }
 
 std::vector<tl::expected<void, ErrorCode>> WrappedMasterService::BatchUpsertEnd(
-    const UUID& client_id, const std::vector<std::string>& keys) {
+    const UUID& client_id, const std::vector<std::string>& keys,
+    const std::string& tenant_id) {
     ScopedVLogTimer timer(1, "BatchUpsertEnd");
     const size_t total_keys = keys.size();
     timer.LogRequest("client_id=", client_id, ", keys_count=", total_keys);
     MasterMetricManager::instance().inc_batch_put_end_requests(total_keys);
 
-    auto results = master_service_.BatchUpsertEnd(client_id, keys);
+    auto results = master_service_.BatchUpsertEnd(client_id, keys, tenant_id);
 
     size_t failure_count = 0;
     for (size_t i = 0; i < results.size(); ++i) {
@@ -1250,13 +584,15 @@ std::vector<tl::expected<void, ErrorCode>> WrappedMasterService::BatchUpsertEnd(
 
 std::vector<tl::expected<void, ErrorCode>>
 WrappedMasterService::BatchUpsertRevoke(const UUID& client_id,
-                                        const std::vector<std::string>& keys) {
+                                        const std::vector<std::string>& keys,
+                                        const std::string& tenant_id) {
     ScopedVLogTimer timer(1, "BatchUpsertRevoke");
     const size_t total_keys = keys.size();
     timer.LogRequest("client_id=", client_id, ", keys_count=", total_keys);
     MasterMetricManager::instance().inc_batch_put_revoke_requests(total_keys);
 
-    auto results = master_service_.BatchUpsertRevoke(client_id, keys);
+    auto results =
+        master_service_.BatchUpsertRevoke(client_id, keys, tenant_id);
 
     size_t failure_count = 0;
     for (size_t i = 0; i < results.size(); ++i) {
@@ -1283,19 +619,19 @@ WrappedMasterService::BatchUpsertRevoke(const UUID& client_id,
 }
 
 tl::expected<void, ErrorCode> WrappedMasterService::Remove(
-    const std::string& key, bool force) {
+    const std::string& key, bool force, const std::string& tenant_id) {
     return execute_rpc(
-        "Remove", [&] { return master_service_.Remove(key, force); },
+        "Remove", [&] { return master_service_.Remove(key, tenant_id, force); },
         [&](auto& timer) { timer.LogRequest("key=", key, ", force=", force); },
         [] { MasterMetricManager::instance().inc_remove_requests(); },
         [] { MasterMetricManager::instance().inc_remove_failures(); });
 }
 
 tl::expected<long, ErrorCode> WrappedMasterService::RemoveByRegex(
-    const std::string& str, bool force) {
+    const std::string& str, bool force, const std::string& tenant_id) {
     return execute_rpc(
         "RemoveByRegex",
-        [&] { return master_service_.RemoveByRegex(str, force); },
+        [&] { return master_service_.RemoveByRegex(str, tenant_id, force); },
         [&](auto& timer) {
             timer.LogRequest("regex=", str, ", force=", force);
         },
@@ -1303,23 +639,24 @@ tl::expected<long, ErrorCode> WrappedMasterService::RemoveByRegex(
         [] { MasterMetricManager::instance().inc_remove_by_regex_failures(); });
 }
 
-long WrappedMasterService::RemoveAll(bool force) {
+long WrappedMasterService::RemoveAll(bool force, const std::string& tenant_id) {
     ScopedVLogTimer timer(1, "RemoveAll");
     timer.LogRequest("action=remove_all_objects, force=", force);
     MasterMetricManager::instance().inc_remove_all_requests();
-    long result = master_service_.RemoveAll(force);
+    long result = master_service_.RemoveAll(tenant_id, force);
     timer.LogResponse("items_removed=", result);
     return result;
 }
 
 std::vector<tl::expected<void, ErrorCode>> WrappedMasterService::BatchRemove(
-    const std::vector<std::string>& keys, bool force) {
+    const std::vector<std::string>& keys, bool force,
+    const std::string& tenant_id) {
     ScopedVLogTimer timer(1, "BatchRemove");
     const size_t total_keys = keys.size();
     timer.LogRequest("keys_count=", total_keys, ", force=", force);
     MasterMetricManager::instance().inc_remove_requests(total_keys);
 
-    auto results = master_service_.BatchRemove(keys, force);
+    auto results = master_service_.BatchRemove(keys, tenant_id, force);
 
     size_t failure_count = 0;
     for (const auto& result : results) {
@@ -1349,6 +686,25 @@ tl::expected<void, ErrorCode> WrappedMasterService::MountSegment(
         [] { MasterMetricManager::instance().inc_mount_segment_failures(); });
 }
 
+tl::expected<void, ErrorCode> WrappedMasterService::MountNoFSegment(
+    const NoFSegment& segment, const UUID& client_id) {
+    return execute_rpc(
+        "MountNoFSegment",
+        [&] { return master_service_.MountNoFSegment(segment, client_id); },
+        [&](auto& timer) {
+            timer.LogRequest("NoF segment mount: ", "base=", segment.base,
+                             ", size=", segment.size,
+                             ", segment_name=", segment.name,
+                             ", id=", segment.id);
+        },
+        [] {
+            MasterMetricManager::instance().inc_mount_nof_segment_requests();
+        },
+        [] {
+            MasterMetricManager::instance().inc_mount_nof_segment_failures();
+        });
+}
+
 tl::expected<void, ErrorCode> WrappedMasterService::ReMountSegment(
     const std::vector<Segment>& segments, const UUID& client_id) {
     return execute_rpc(
@@ -1362,6 +718,23 @@ tl::expected<void, ErrorCode> WrappedMasterService::ReMountSegment(
         [] { MasterMetricManager::instance().inc_remount_segment_failures(); });
 }
 
+tl::expected<void, ErrorCode> WrappedMasterService::ReMountNoFSegment(
+    const std::vector<NoFSegment>& segments, const UUID& client_id) {
+    return execute_rpc(
+        "ReMountNoFSegment",
+        [&] { return master_service_.ReMountNoFSegment(segments, client_id); },
+        [&](auto& timer) {
+            timer.LogRequest("NoF segment remount: ", "segments_count=",
+                             segments.size(), ", client_id=", client_id);
+        },
+        [] {
+            MasterMetricManager::instance().inc_remount_nof_segment_requests();
+        },
+        [] {
+            MasterMetricManager::instance().inc_remount_nof_segment_failures();
+        });
+}
+
 tl::expected<void, ErrorCode> WrappedMasterService::UnmountSegment(
     const UUID& segment_id, const UUID& client_id) {
     return execute_rpc(
@@ -1375,18 +748,79 @@ tl::expected<void, ErrorCode> WrappedMasterService::UnmountSegment(
         [] { MasterMetricManager::instance().inc_unmount_segment_failures(); });
 }
 
+tl::expected<void, ErrorCode> WrappedMasterService::GracefulUnmountSegment(
+    const UUID& segment_id, const UUID& client_id, uint64_t grace_period_ms) {
+    return execute_rpc(
+        "GracefulUnmountSegment",
+        [&] {
+            return master_service_.GracefulUnmountSegment(segment_id, client_id,
+                                                          grace_period_ms);
+        },
+        [&](auto& timer) {
+            timer.LogRequest("segment_id=", segment_id,
+                             ", client_id=", client_id,
+                             ", grace_period_ms=", grace_period_ms);
+        },
+        [] {
+            MasterMetricManager::instance()
+                .inc_unmount_segment_requests();  // reuse metric or add new
+        },
+        [] {
+            MasterMetricManager::instance()
+                .inc_unmount_segment_failures();  // reuse metric or add new
+        });
+}
+
+tl::expected<void, ErrorCode> WrappedMasterService::UnmountNoFSegment(
+    const UUID& segment_id, const UUID& client_id) {
+    return execute_rpc(
+        "UnmountNoFSegment",
+        [&] {
+            return master_service_.UnmountNoFSegment(segment_id, client_id);
+        },
+        [&](auto& timer) {
+            timer.LogRequest("NoF segment unmount: ", "segment_id=", segment_id,
+                             ", client_id=", client_id);
+        },
+        [] {
+            MasterMetricManager::instance().inc_unmount_nof_segment_requests();
+        },
+        [] {
+            MasterMetricManager::instance().inc_unmount_nof_segment_failures();
+        });
+}
+
+tl::expected<std::vector<NoFSegment>, ErrorCode>
+WrappedMasterService::GetAllNoFSegments() {
+    return execute_rpc(
+        "GetAllNoFSegments",
+        [&] { return master_service_.GetAllNoFSegments(); },
+        [&](auto& timer) { timer.LogRequest("Get all NoF segments"); }, [] {},
+        [] {});
+}
+
+tl::expected<std::vector<NoFSegmentOwnerInfo>, ErrorCode>
+WrappedMasterService::GetNoFSegmentsByName(const std::string& segment_name) {
+    return execute_rpc(
+        "GetNoFSegmentsByName",
+        [&] { return master_service_.GetNoFSegmentsByName(segment_name); },
+        [&](auto& timer) { timer.LogRequest("segment_name=", segment_name); },
+        [] {}, [] {});
+}
+
 tl::expected<CopyStartResponse, ErrorCode> WrappedMasterService::CopyStart(
-    const UUID& client_id, const std::string& key,
+    const UUID& client_id, const std::string& key, const std::string& tenant_id,
     const std::string& src_segment,
     const std::vector<std::string>& tgt_segments) {
     return execute_rpc(
         "CopyStart",
         [&] {
-            return master_service_.CopyStart(client_id, key, src_segment,
-                                             tgt_segments);
+            return master_service_.CopyStart(client_id, key, tenant_id,
+                                             src_segment, tgt_segments);
         },
         [&](auto& timer) {
             timer.LogRequest("client_id=", client_id, ", key=", key,
+                             ", tenant_id=", tenant_id,
                              ", src_segment=", src_segment,
                              ", tgt_segments_count=", tgt_segments.size());
         },
@@ -1395,39 +829,45 @@ tl::expected<CopyStartResponse, ErrorCode> WrappedMasterService::CopyStart(
 }
 
 tl::expected<void, ErrorCode> WrappedMasterService::CopyEnd(
-    const UUID& client_id, const std::string& key) {
+    const UUID& client_id, const std::string& key,
+    const std::string& tenant_id) {
     return execute_rpc(
-        "CopyEnd", [&] { return master_service_.CopyEnd(client_id, key); },
+        "CopyEnd",
+        [&] { return master_service_.CopyEnd(client_id, key, tenant_id); },
         [&](auto& timer) {
-            timer.LogRequest("client_id=", client_id, ", key=", key);
+            timer.LogRequest("client_id=", client_id, ", key=", key,
+                             ", tenant_id=", tenant_id);
         },
         [] { MasterMetricManager::instance().inc_copy_end_requests(); },
         [] { MasterMetricManager::instance().inc_copy_end_failures(); });
 }
 
 tl::expected<void, ErrorCode> WrappedMasterService::CopyRevoke(
-    const UUID& client_id, const std::string& key) {
+    const UUID& client_id, const std::string& key,
+    const std::string& tenant_id) {
     return execute_rpc(
         "CopyRevoke",
-        [&] { return master_service_.CopyRevoke(client_id, key); },
+        [&] { return master_service_.CopyRevoke(client_id, key, tenant_id); },
         [&](auto& timer) {
-            timer.LogRequest("client_id=", client_id, ", key=", key);
+            timer.LogRequest("client_id=", client_id, ", key=", key,
+                             ", tenant_id=", tenant_id);
         },
         [] { MasterMetricManager::instance().inc_copy_revoke_requests(); },
         [] { MasterMetricManager::instance().inc_copy_revoke_failures(); });
 }
 
 tl::expected<MoveStartResponse, ErrorCode> WrappedMasterService::MoveStart(
-    const UUID& client_id, const std::string& key,
+    const UUID& client_id, const std::string& key, const std::string& tenant_id,
     const std::string& src_segment, const std::string& tgt_segment) {
     return execute_rpc(
         "MoveStart",
         [&] {
-            return master_service_.MoveStart(client_id, key, src_segment,
-                                             tgt_segment);
+            return master_service_.MoveStart(client_id, key, tenant_id,
+                                             src_segment, tgt_segment);
         },
         [&](auto& timer) {
             timer.LogRequest("client_id=", client_id, ", key=", key,
+                             ", tenant_id=", tenant_id,
                              ", src_segment=", src_segment,
                              ", tgt_segment=", tgt_segment);
         },
@@ -1436,38 +876,45 @@ tl::expected<MoveStartResponse, ErrorCode> WrappedMasterService::MoveStart(
 }
 
 tl::expected<void, ErrorCode> WrappedMasterService::MoveEnd(
-    const UUID& client_id, const std::string& key) {
+    const UUID& client_id, const std::string& key,
+    const std::string& tenant_id) {
     return execute_rpc(
-        "MoveEnd", [&] { return master_service_.MoveEnd(client_id, key); },
+        "MoveEnd",
+        [&] { return master_service_.MoveEnd(client_id, key, tenant_id); },
         [&](auto& timer) {
-            timer.LogRequest("client_id=", client_id, ", key=", key);
+            timer.LogRequest("client_id=", client_id, ", key=", key,
+                             ", tenant_id=", tenant_id);
         },
         [] { MasterMetricManager::instance().inc_move_end_requests(); },
         [] { MasterMetricManager::instance().inc_move_end_failures(); });
 }
 
 tl::expected<void, ErrorCode> WrappedMasterService::MoveRevoke(
-    const UUID& client_id, const std::string& key) {
+    const UUID& client_id, const std::string& key,
+    const std::string& tenant_id) {
     return execute_rpc(
         "MoveRevoke",
-        [&] { return master_service_.MoveRevoke(client_id, key); },
+        [&] { return master_service_.MoveRevoke(client_id, key, tenant_id); },
         [&](auto& timer) {
-            timer.LogRequest("client_id=", client_id, ", key=", key);
+            timer.LogRequest("client_id=", client_id, ", key=", key,
+                             ", tenant_id=", tenant_id);
         },
         [] { MasterMetricManager::instance().inc_move_revoke_requests(); },
         [] { MasterMetricManager::instance().inc_move_revoke_failures(); });
 }
 
 tl::expected<void, ErrorCode> WrappedMasterService::EvictDiskReplica(
-    const UUID& client_id, const std::string& key, ReplicaType replica_type) {
+    const UUID& client_id, const std::string& key, const std::string& tenant_id,
+    ReplicaType replica_type) {
     return execute_rpc(
         "EvictDiskReplica",
         [&] {
-            return master_service_.EvictDiskReplica(client_id, key,
+            return master_service_.EvictDiskReplica(client_id, key, tenant_id,
                                                     replica_type);
         },
         [&](auto& timer) {
             timer.LogRequest("client_id=", client_id, ", key=", key,
+                             ", tenant_id=", tenant_id,
                              ", replica_type=", replica_type);
         },
         [] {
@@ -1481,15 +928,16 @@ tl::expected<void, ErrorCode> WrappedMasterService::EvictDiskReplica(
 std::vector<tl::expected<void, ErrorCode>>
 WrappedMasterService::BatchEvictDiskReplica(
     const UUID& client_id, const std::vector<std::string>& keys,
-    ReplicaType replica_type) {
+    const std::string& tenant_id, ReplicaType replica_type) {
     ScopedVLogTimer timer(1, "BatchEvictDiskReplica");
     const size_t total_keys = keys.size();
     timer.LogRequest("client_id=", client_id, ", keys_count=", total_keys,
+                     ", tenant_id=", tenant_id,
                      ", replica_type=", replica_type);
     MasterMetricManager::instance().inc_evict_disk_replica_requests();
 
-    auto results =
-        master_service_.BatchEvictDiskReplica(client_id, keys, replica_type);
+    auto results = master_service_.BatchEvictDiskReplica(
+        client_id, keys, tenant_id, replica_type);
 
     size_t failure_count = 0;
     for (size_t i = 0; i < results.size(); ++i) {
@@ -1511,12 +959,14 @@ WrappedMasterService::BatchEvictDiskReplica(
 }
 
 tl::expected<UUID, ErrorCode> WrappedMasterService::CreateCopyTask(
-    const std::string& key, const std::vector<std::string>& targets) {
+    const std::string& key, const std::string& tenant_id,
+    const std::vector<std::string>& targets) {
     return execute_rpc(
         "CreateCopyTask",
-        [&] { return master_service_.CreateCopyTask(key, targets); },
+        [&] { return master_service_.CreateCopyTask(key, tenant_id, targets); },
         [&](auto& timer) {
-            timer.LogRequest("key=", key, ", targets_size=", targets.size());
+            timer.LogRequest("key=", key, ", tenant_id=", tenant_id,
+                             ", targets_size=", targets.size());
         },
         [] { MasterMetricManager::instance().inc_create_copy_task_requests(); },
         [] {
@@ -1525,14 +975,17 @@ tl::expected<UUID, ErrorCode> WrappedMasterService::CreateCopyTask(
 }
 
 tl::expected<UUID, ErrorCode> WrappedMasterService::CreateMoveTask(
-    const std::string& key, const std::string& source,
-    const std::string& target) {
+    const std::string& key, const std::string& tenant_id,
+    const std::string& source, const std::string& target) {
     return execute_rpc(
         "CreateMoveTask",
-        [&] { return master_service_.CreateMoveTask(key, source, target); },
+        [&] {
+            return master_service_.CreateMoveTask(key, tenant_id, source,
+                                                  target);
+        },
         [&](auto& timer) {
-            timer.LogRequest("key=", key, ", source=", source,
-                             ", target=", target);
+            timer.LogRequest("key=", key, ", tenant_id=", tenant_id,
+                             ", source=", source, ", target=", target);
         },
         [] { MasterMetricManager::instance().inc_create_move_task_requests(); },
         [] {
@@ -1612,9 +1065,57 @@ tl::expected<std::string, ErrorCode> WrappedMasterService::ServiceReady() {
     return GetMooncakeStoreVersion();
 }
 
+tl::expected<std::vector<TenantQuotaSnapshot>, ErrorCode>
+WrappedMasterService::ListTenantQuotaSnapshots() {
+    if (!master_service_.IsTenantQuotaEnabled()) {
+        return tl::make_unexpected(ErrorCode::UNAVAILABLE_IN_CURRENT_MODE);
+    }
+    return master_service_.ListTenantQuotaSnapshots();
+}
+
+tl::expected<TenantQuotaSnapshot, ErrorCode>
+WrappedMasterService::GetTenantQuotaSnapshot(const std::string& tenant_id) {
+    if (!master_service_.IsTenantQuotaEnabled()) {
+        return tl::make_unexpected(ErrorCode::UNAVAILABLE_IN_CURRENT_MODE);
+    }
+    auto snapshot = master_service_.GetTenantQuotaSnapshot(tenant_id);
+    if (!snapshot.has_value()) {
+        return tl::make_unexpected(ErrorCode::OBJECT_NOT_FOUND);
+    }
+    return snapshot.value();
+}
+
+tl::expected<TenantQuotaSnapshot, ErrorCode>
+WrappedMasterService::UpsertTenantQuotaPolicy(const std::string& tenant_id,
+                                              uint64_t requested_quota_bytes) {
+    if (!master_service_.IsTenantQuotaEnabled()) {
+        return tl::make_unexpected(ErrorCode::UNAVAILABLE_IN_CURRENT_MODE);
+    }
+    return master_service_.UpsertTenantQuotaPolicy(tenant_id,
+                                                   requested_quota_bytes);
+}
+
+tl::expected<std::optional<TenantQuotaSnapshot>, ErrorCode>
+WrappedMasterService::DeleteTenantQuotaPolicy(const std::string& tenant_id) {
+    if (!master_service_.IsTenantQuotaEnabled()) {
+        return tl::make_unexpected(ErrorCode::UNAVAILABLE_IN_CURRENT_MODE);
+    }
+    return master_service_.DeleteTenantQuotaPolicy(tenant_id);
+}
+
+tl::expected<uint64_t, ErrorCode>
+WrappedMasterService::GetTenantQuotaAllocatableCapacityBytes() {
+    if (!master_service_.IsTenantQuotaEnabled()) {
+        return tl::make_unexpected(ErrorCode::UNAVAILABLE_IN_CURRENT_MODE);
+    }
+    return master_service_.GetTenantQuotaAllocatableCapacityBytes();
+}
+
 tl::expected<std::vector<std::string>, ErrorCode>
 WrappedMasterService::GetAllKeysForAdmin() {
-    return master_service_.GetAllKeys();
+    // Compatibility endpoint: /get_all_keys historically listed only the
+    // default tenant's keys.
+    return master_service_.GetAllKeys("default");
 }
 
 tl::expected<std::vector<std::string>, ErrorCode>
@@ -1622,6 +1123,11 @@ WrappedMasterService::GetAllSegmentsForAdmin() {
     return master_service_.GetAllSegments();
 }
 
+tl::expected<std::vector<MasterService::SegmentDetailInfo>, ErrorCode>
+WrappedMasterService::GetSegmentsDetailForAdmin() {
+    return master_service_.GetSegmentsDetail();
+}
+
 tl::expected<std::pair<uint64_t, uint64_t>, ErrorCode>
 WrappedMasterService::QuerySegmentForAdmin(const std::string& segment) {
     return master_service_.QuerySegments(segment);
@@ -1640,8 +1146,7 @@ tl::expected<void, ErrorCode> WrappedMasterService::MountLocalDiskSegment(
     return result;
 }
 
-tl::expected<std::unordered_map<std::string, int64_t, std::hash<std::string>>,
-             ErrorCode>
+tl::expected<std::vector<OffloadTaskItem>, ErrorCode>
 WrappedMasterService::OffloadObjectHeartbeat(const UUID& client_id,
                                              bool enable_offloading) {
     ScopedVLogTimer timer(1, "OffloadObjectHeartbeat");
@@ -1651,14 +1156,64 @@ WrappedMasterService::OffloadObjectHeartbeat(const UUID& client_id,
     return result;
 }
 
+tl::expected<void, ErrorCode> WrappedMasterService::ReportSsdCapacity(
+    const UUID& client_id, int64_t ssd_total_capacity_bytes) {
+    ScopedVLogTimer timer(1, "ReportSsdCapacity");
+    timer.LogRequest("client_id=", client_id,
+                     ", ssd_total_capacity_bytes=", ssd_total_capacity_bytes);
+    return master_service_.ReportSsdCapacity(client_id,
+                                             ssd_total_capacity_bytes);
+}
+
 tl::expected<void, ErrorCode> WrappedMasterService::NotifyOffloadSuccess(
-    const UUID& client_id, const std::vector<std::string>& keys,
+    const UUID& client_id, const std::vector<OffloadTaskItem>& tasks,
     const std::vector<StorageObjectMetadata>& metadatas) {
     ScopedVLogTimer timer(1, "NotifyOffloadSuccess");
     timer.LogRequest("action=notify_offload_success");
 
     auto result =
-        master_service_.NotifyOffloadSuccess(client_id, keys, metadatas);
+        master_service_.NotifyOffloadSuccess(client_id, tasks, metadatas);
+    timer.LogResponseExpected(result);
+    return result;
+}
+
+tl::expected<std::vector<PromotionTaskItem>, ErrorCode>
+WrappedMasterService::PromotionObjectHeartbeat(const UUID& client_id) {
+    ScopedVLogTimer timer(1, "PromotionObjectHeartbeat");
+    timer.LogRequest("action=promotion_object_heartbeat");
+    return master_service_.PromotionObjectHeartbeat(client_id);
+}
+
+tl::expected<PromotionAllocStartResponse, ErrorCode>
+WrappedMasterService::PromotionAllocStart(
+    const UUID& client_id, const std::string& key, const std::string& tenant_id,
+    uint64_t size, const std::vector<std::string>& preferred_segments) {
+    ScopedVLogTimer timer(1, "PromotionAllocStart");
+    timer.LogRequest("action=promotion_alloc_start");
+    auto result = master_service_.PromotionAllocStart(client_id, key, tenant_id,
+                                                      size, preferred_segments);
+    timer.LogResponseExpected(result);
+    return result;
+}
+
+tl::expected<void, ErrorCode> WrappedMasterService::NotifyPromotionSuccess(
+    const UUID& client_id, const std::string& key,
+    const std::string& tenant_id) {
+    ScopedVLogTimer timer(1, "NotifyPromotionSuccess");
+    timer.LogRequest("action=notify_promotion_success");
+    auto result =
+        master_service_.NotifyPromotionSuccess(client_id, key, tenant_id);
+    timer.LogResponseExpected(result);
+    return result;
+}
+
+tl::expected<void, ErrorCode> WrappedMasterService::NotifyPromotionFailure(
+    const UUID& client_id, const std::string& key,
+    const std::string& tenant_id) {
+    ScopedVLogTimer timer(1, "NotifyPromotionFailure");
+    timer.LogRequest("action=notify_promotion_failure");
+    auto result =
+        master_service_.NotifyPromotionFailure(client_id, key, tenant_id);
     timer.LogResponseExpected(result);
     return result;
 }
@@ -1683,6 +1238,11 @@ tl::expected<SegmentStatus, ErrorCode> WrappedMasterService::QuerySegmentStatus(
     return master_service_.QuerySegmentStatus(segment_name);
 }
 
+tl::expected<SegmentStatus, ErrorCode>
+WrappedMasterService::QuerySegmentStatusById(const UUID& segment_id) {
+    return master_service_.QuerySegmentStatusById(segment_id);
+}
+
 void RegisterRpcService(
     coro_rpc::coro_rpc_server& server,
     mooncake::WrappedMasterService& wrapped_master_service) {
@@ -1734,14 +1294,34 @@ void RegisterRpcService(
         &wrapped_master_service);
     server.register_handler<&mooncake::WrappedMasterService::MountSegment>(
         &wrapped_master_service);
+    server.register_handler<&mooncake::WrappedMasterService::MountNoFSegment>(
+        &wrapped_master_service);
     server.register_handler<&mooncake::WrappedMasterService::ReMountSegment>(
         &wrapped_master_service);
+    server.register_handler<&mooncake::WrappedMasterService::ReMountNoFSegment>(
+        &wrapped_master_service);
     server.register_handler<&mooncake::WrappedMasterService::UnmountSegment>(
         &wrapped_master_service);
+    server.register_handler<
+        &mooncake::WrappedMasterService::GracefulUnmountSegment>(
+        &wrapped_master_service);
+    server.register_handler<&mooncake::WrappedMasterService::UnmountNoFSegment>(
+        &wrapped_master_service);
+    server.register_handler<&mooncake::WrappedMasterService::GetAllNoFSegments>(
+        &wrapped_master_service);
+    server.register_handler<
+        &mooncake::WrappedMasterService::GetNoFSegmentsByName>(
+        &wrapped_master_service);
     server.register_handler<&mooncake::WrappedMasterService::Ping>(
         &wrapped_master_service);
     server.register_handler<&mooncake::WrappedMasterService::GetFsdir>(
         &wrapped_master_service);
+    server
+        .register_handler<&mooncake::WrappedMasterService::QuerySegmentStatus>(
+            &wrapped_master_service);
+    server.register_handler<
+        &mooncake::WrappedMasterService::QuerySegmentStatusById>(
+        &wrapped_master_service);
     server.register_handler<&mooncake::WrappedMasterService::GetStorageConfig>(
         &wrapped_master_service);
     server.register_handler<&mooncake::WrappedMasterService::BatchExistKey>(
@@ -1754,9 +1334,23 @@ void RegisterRpcService(
     server.register_handler<
         &mooncake::WrappedMasterService::OffloadObjectHeartbeat>(
         &wrapped_master_service);
+    server.register_handler<&mooncake::WrappedMasterService::ReportSsdCapacity>(
+        &wrapped_master_service);
     server.register_handler<
         &mooncake::WrappedMasterService::NotifyOffloadSuccess>(
         &wrapped_master_service);
+    server.register_handler<
+        &mooncake::WrappedMasterService::PromotionObjectHeartbeat>(
+        &wrapped_master_service);
+    server
+        .register_handler<&mooncake::WrappedMasterService::PromotionAllocStart>(
+            &wrapped_master_service);
+    server.register_handler<
+        &mooncake::WrappedMasterService::NotifyPromotionSuccess>(
+        &wrapped_master_service);
+    server.register_handler<
+        &mooncake::WrappedMasterService::NotifyPromotionFailure>(
+        &wrapped_master_service);
     server.register_handler<&mooncake::WrappedMasterService::CopyStart>(
         &wrapped_master_service);
     server.register_handler<&mooncake::WrappedMasterService::CopyEnd>(
diff --git a/mooncake-store/src/segment.cpp b/mooncake-store/src/segment.cpp
index b96845eb..8551cedd 100644
--- a/mooncake-store/src/segment.cpp
+++ b/mooncake-store/src/segment.cpp
@@ -20,6 +20,11 @@ bool HasAllocator(const AllocatorManager& allocator_manager,
            allocators->end();
 }
 
+bool IsMsgpackInteger(const msgpack::object& object) {
+    return object.type == msgpack::type::POSITIVE_INTEGER ||
+           object.type == msgpack::type::NEGATIVE_INTEGER;
+}
+
 }  // namespace
 
 ErrorCode ScopedSegmentAccess::MountSegment(const Segment& segment,
@@ -211,6 +216,47 @@ ErrorCode ScopedSegmentAccess::PrepareUnmountSegment(
     return ErrorCode::OK;
 }
 
+ErrorCode ScopedSegmentAccess::PrepareGracefulUnmountSegment(
+    const UUID& segment_id) {
+    auto it = segment_manager_->mounted_segments_.find(segment_id);
+    if (it == segment_manager_->mounted_segments_.end()) {
+        LOG(WARNING) << "segment_id=" << segment_id
+                     << ", warn=segment_not_found";
+        return ErrorCode::SEGMENT_NOT_FOUND;
+    }
+    auto status = it->second.status;
+    if (status == SegmentStatus::UNMOUNTING) {
+        LOG(ERROR) << "segment_id=" << segment_id
+                   << ", error=segment_is_unmounting";
+        return ErrorCode::UNAVAILABLE_IN_CURRENT_STATUS;
+    }
+    if (status == SegmentStatus::GRACEFULLY_UNMOUNTING) {
+        // Idempotent: already in graceful unmount state
+        return ErrorCode::OK;
+    }
+    if (status != SegmentStatus::OK && status != SegmentStatus::DRAINING) {
+        LOG(ERROR) << "segment_id=" << segment_id
+                   << ", error=unavailable_in_current_status, status="
+                   << status;
+        return ErrorCode::UNAVAILABLE_IN_CURRENT_STATUS;
+    }
+
+    auto& mounted_segment = it->second;
+    auto& segment = mounted_segment.segment;
+
+    // Remove the allocator from the segment manager
+    std::shared_ptr<BufferAllocatorBase> allocator =
+        mounted_segment.buf_allocator;
+    if (HasAllocator(segment_manager_->allocator_manager_, segment.name,
+                     allocator)) {
+        segment_manager_->allocator_manager_.removeAllocator(segment.name,
+                                                             allocator);
+    }
+    // Set the segment status to GRACEFULLY_UNMOUNTING
+    mounted_segment.status = SegmentStatus::GRACEFULLY_UNMOUNTING;
+    return ErrorCode::OK;
+}
+
 ErrorCode ScopedSegmentAccess::CommitUnmountSegment(
     const UUID& segment_id, const UUID& client_id,
     const size_t& metrics_dec_capacity) {
@@ -240,9 +286,14 @@ ErrorCode ScopedSegmentAccess::CommitUnmountSegment(
     auto&& segment = segment_manager_->mounted_segments_.find(segment_id);
     if (segment != segment_manager_->mounted_segments_.end()) {
         segment_name = segment->second.segment.name;
-        // Also remove from segment_name_client_id_map_
-        segment_manager_->client_by_name_.erase(segment_name);
-        segment_manager_->segment_id_by_name_.erase(segment_name);
+        auto segment_id_by_name_it =
+            segment_manager_->segment_id_by_name_.find(segment_name);
+        if (segment_id_by_name_it !=
+                segment_manager_->segment_id_by_name_.end() &&
+            segment_id_by_name_it->second == segment_id) {
+            segment_manager_->segment_id_by_name_.erase(segment_id_by_name_it);
+            segment_manager_->client_by_name_.erase(segment_name);
+        }
         is_cxl = (segment->second.segment.protocol == "cxl");
     }
     // Remove from mounted_segments_
@@ -276,6 +327,20 @@ ErrorCode ScopedSegmentAccess::GetClientSegments(
 void ScopedSegmentAccess::UnmountLocalDiskSegment(const UUID& client_id) {
     auto it = segment_manager_->client_local_disk_segment_.find(client_id);
     if (it != segment_manager_->client_local_disk_segment_.end()) {
+        // Hold offloading_mutex_ while reading ssd_total_capacity_bytes to
+        // avoid a data race with OffloadObjectHeartbeat, which writes the
+        // field under the same lock.  Release the lock before erase() so we
+        // don't unlock an already-destroyed mutex (erase destroys the
+        // LocalDiskSegment, including its mutex).
+        int64_t reported_capacity = 0;
+        {
+            MutexLocker locker(&it->second->offloading_mutex_);
+            reported_capacity = it->second->ssd_total_capacity_bytes;
+        }
+        if (reported_capacity > 0) {
+            MasterMetricManager::instance().dec_total_file_capacity(
+                reported_capacity);
+        }
         segment_manager_->client_local_disk_segment_.erase(it);
         LOG(INFO) << "client_id=" << client_id
                   << ", action=unmount_local_disk_segment";
@@ -513,10 +578,10 @@ SegmentSerializer::Serialize() {
             segment_manager_->client_local_disk_segment_.at(client_uuid);
         packer.pack(UuidToString(client_uuid));
 
-        // Serialize LocalDiskSegment: [enable_offloading, count, key1, ts1,
-        // key2, ts2, ...] Sort keys to ensure determinism
+        // Serialize LocalDiskSegment: [enable_offloading, count, storage_key1,
+        // task1, storage_key2, task2, ...] Sort keys to ensure determinism.
         std::vector<std::string> sorted_keys;
-        for (const auto& [key, ts] : segment->offloading_objects) {
+        for (const auto& [key, _] : segment->offloading_objects) {
             sorted_keys.push_back(key);
         }
         std::sort(sorted_keys.begin(), sorted_keys.end());
@@ -527,7 +592,11 @@ SegmentSerializer::Serialize() {
 
         for (const auto& key : sorted_keys) {
             packer.pack(key);
-            packer.pack(segment->offloading_objects.at(key));
+            const auto& task = segment->offloading_objects.at(key);
+            packer.pack_array(3);
+            packer.pack(task.tenant_id);
+            packer.pack(task.key);
+            packer.pack(task.size);
         }
     }
 
@@ -842,8 +911,8 @@ tl::expected<void, SerializationError> SegmentSerializer::Deserialize(
                                 client_uuid_str)));
             }
 
-            // Parse LocalDiskSegment array: [enable_offloading, count, key1,
-            // ts1, ...]
+            // Parse LocalDiskSegment array: [enable_offloading, count,
+            // storage_key1, task1, ...]
             if (client_value.type != msgpack::type::ARRAY ||
                 client_value.via.array.size < 2) {
                 return tl::unexpected(
@@ -861,19 +930,62 @@ tl::expected<void, SerializationError> SegmentSerializer::Deserialize(
             // Parse offloading_objects
             for (uint64_t k = 0; k < count; ++k) {
                 size_t key_idx = 2 + k * 2;
-                size_t ts_idx = 2 + k * 2 + 1;
-                if (ts_idx >= client_value.via.array.size) {
+                size_t task_idx = 2 + k * 2 + 1;
+                if (task_idx >= client_value.via.array.size) {
                     return tl::unexpected(
                         SerializationError(ErrorCode::DESERIALIZE_FAIL,
                                            "deserialize local_disk_segments "
                                            "offloading_objects out of bounds"));
                 }
 
+                if (client_value.via.array.ptr[key_idx].type !=
+                    msgpack::type::STR) {
+                    return tl::unexpected(SerializationError(
+                        ErrorCode::DESERIALIZE_FAIL,
+                        "deserialize local_disk_segments offloading key is "
+                        "not string"));
+                }
                 std::string key(
                     client_value.via.array.ptr[key_idx].via.str.ptr,
                     client_value.via.array.ptr[key_idx].via.str.size);
-                int64_t ts = client_value.via.array.ptr[ts_idx].as<int64_t>();
-                segment->offloading_objects[key] = ts;
+                const auto& task_obj = client_value.via.array.ptr[task_idx];
+                if (task_obj.type == msgpack::type::ARRAY &&
+                    task_obj.via.array.size == 3) {
+                    if (task_obj.via.array.ptr[0].type != msgpack::type::STR ||
+                        task_obj.via.array.ptr[1].type != msgpack::type::STR) {
+                        return tl::unexpected(SerializationError(
+                            ErrorCode::DESERIALIZE_FAIL,
+                            "deserialize local_disk_segments offloading task "
+                            "fields are not strings"));
+                    }
+                    if (!IsMsgpackInteger(task_obj.via.array.ptr[2])) {
+                        return tl::unexpected(SerializationError(
+                            ErrorCode::DESERIALIZE_FAIL,
+                            "deserialize local_disk_segments offloading task "
+                            "size is not integer"));
+                    }
+                    OffloadTaskItem task;
+                    task.tenant_id =
+                        task_obj.via.array.ptr[0].as<std::string>();
+                    task.key = task_obj.via.array.ptr[1].as<std::string>();
+                    task.size = task_obj.via.array.ptr[2].as<int64_t>();
+                    segment->offloading_objects[key] = std::move(task);
+                } else {
+                    // Backward compatibility for snapshots whose
+                    // offloading_objects value was key -> size.
+                    if (!IsMsgpackInteger(task_obj)) {
+                        return tl::unexpected(SerializationError(
+                            ErrorCode::DESERIALIZE_FAIL,
+                            "deserialize local_disk_segments legacy "
+                            "offloading size is not integer"));
+                    }
+                    auto [tenant_id, user_key] =
+                        ParseTenantScopedStorageKey(key);
+                    segment->offloading_objects[key] =
+                        OffloadTaskItem{.tenant_id = std::move(tenant_id),
+                                        .key = std::move(user_key),
+                                        .size = task_obj.as<int64_t>()};
+                }
             }
 
             segment_manager_->client_local_disk_segment_[client_id] =
@@ -940,6 +1052,17 @@ ErrorCode ScopedSegmentAccess::GetSegmentStatusByName(
     return ErrorCode::OK;
 }
 
+ErrorCode ScopedSegmentAccess::GetSegmentStatusById(
+    const UUID& segment_id, SegmentStatus& status) const {
+    auto mounted_segment_it =
+        segment_manager_->mounted_segments_.find(segment_id);
+    if (mounted_segment_it == segment_manager_->mounted_segments_.end()) {
+        return ErrorCode::SEGMENT_NOT_FOUND;
+    }
+    status = mounted_segment_it->second.status;
+    return ErrorCode::OK;
+}
+
 ErrorCode ScopedSegmentAccess::SetSegmentStatusByName(
     const std::string& segment_name, SegmentStatus status) {
     auto segment_id_it =
@@ -977,6 +1100,272 @@ ErrorCode ScopedSegmentAccess::SetSegmentStatusByName(
     return ErrorCode::OK;
 }
 
+/* ScopedNoFSegmentAccess Implementation */
+ErrorCode ScopedNoFSegmentAccess::MountSegment(const NoFSegment& segment,
+                                               const UUID& client_id) {
+    const uintptr_t buffer = segment.base;
+    const size_t size = segment.size;
+
+    // NoF segment base is an NVMe namespace offset, so 0 is valid.
+    if (size == 0) {
+        LOG(ERROR) << "NoF segment mount: buffer=" << buffer
+                   << " or size=" << size << " is invalid";
+        return ErrorCode::INVALID_PARAMS;
+    }
+
+    if (nof_segment_manager_->memory_allocator_ ==
+            BufferAllocatorType::CACHELIB &&
+        (buffer % facebook::cachelib::Slab::kSize ||
+         size % facebook::cachelib::Slab::kSize)) {
+        LOG(ERROR) << "NoF segment mount: buffer=" << buffer
+                   << " or size=" << size << " is not aligned to "
+                   << facebook::cachelib::Slab::kSize
+                   << " as required by Cachelib";
+        return ErrorCode::INVALID_PARAMS;
+    }
+
+    auto exist_segment_it =
+        nof_segment_manager_->mounted_segments_.find(segment.id);
+    if (exist_segment_it != nof_segment_manager_->mounted_segments_.end()) {
+        auto& exist_segment = exist_segment_it->second;
+        if (exist_segment.status == SegmentStatus::OK) {
+            LOG(WARNING) << "NoF segment mount: segment_name=" << segment.name
+                         << ", warn=segment_already_exists_by_id";
+            return ErrorCode::SEGMENT_ALREADY_EXISTS;
+        }
+        LOG(ERROR) << "NoF segment mount: segment_name=" << segment.name
+                   << ", error=segment_already_exists_but_not_ok"
+                   << ", status=" << exist_segment.status;
+        return ErrorCode::UNAVAILABLE_IN_CURRENT_STATUS;
+    }
+
+    // Treat the same transport endpoint as the same remote SSD namespace even
+    // if a retry arrives with a different generated UUID.
+    for (const auto& [existing_id, existing_segment] :
+         nof_segment_manager_->mounted_segments_) {
+        if (existing_segment.status == SegmentStatus::OK &&
+            existing_segment.segment.te_endpoint == segment.te_endpoint) {
+            LOG(WARNING) << "NoF segment mount: segment_name=" << segment.name
+                         << ", endpoint=" << segment.te_endpoint
+                         << ", warn=segment_already_exists_with_different_id";
+            return ErrorCode::SEGMENT_ALREADY_EXISTS;
+        }
+    }
+
+    std::shared_ptr<BufferAllocatorBase> allocator;
+    try {
+        switch (nof_segment_manager_->memory_allocator_) {
+            case BufferAllocatorType::CACHELIB:
+                allocator = std::make_shared<CachelibBufferAllocator>(
+                    segment.name, buffer, size, segment.te_endpoint,
+                    ReplicaType::NOF_SSD);
+                break;
+            case BufferAllocatorType::OFFSET:
+                allocator = std::make_shared<OffsetBufferAllocator>(
+                    segment.name, buffer, size, segment.te_endpoint,
+                    ReplicaType::NOF_SSD);
+                break;
+            default:
+                LOG(ERROR) << "NoF segment mount: segment_name=" << segment.name
+                           << ", error=unknown_memory_allocator="
+                           << static_cast<int>(
+                                  nof_segment_manager_->memory_allocator_);
+                return ErrorCode::INVALID_PARAMS;
+        }
+
+        if (!allocator) {
+            LOG(ERROR) << "NoF segment mount: segment_name=" << segment.name
+                       << ", error=failed_to_create_allocator";
+            return ErrorCode::INVALID_PARAMS;
+        }
+    } catch (...) {
+        LOG(ERROR) << "NoF segment mount: segment_name=" << segment.name
+                   << ", error=exception_during_allocator_creation";
+        return ErrorCode::INVALID_PARAMS;
+    }
+
+    nof_segment_manager_->allocator_manager_.addAllocator(segment.name,
+                                                          allocator);
+    nof_segment_manager_->client_segments_[client_id].push_back(segment.id);
+    nof_segment_manager_->mounted_segments_[segment.id] = {
+        segment, client_id, SegmentStatus::OK, std::move(allocator)};
+    nof_segment_manager_->client_by_name_[segment.name] = client_id;
+    MasterMetricManager::instance().inc_total_nof_capacity(segment.name, size);
+
+    return ErrorCode::OK;
+}
+
+ErrorCode ScopedNoFSegmentAccess::ReMountSegment(
+    const std::vector<NoFSegment>& segments, const UUID& client_id) {
+    for (const auto& segment : segments) {
+        ErrorCode err = MountSegment(segment, client_id);
+        if (err == ErrorCode::UNAVAILABLE_IN_CURRENT_STATUS ||
+            err == ErrorCode::INTERNAL_ERROR) {
+            LOG(ERROR) << "NoF segment remount: segment_name=" << segment.name
+                       << ", error=fail_to_remount_segment";
+            return err;
+        }
+        if (err == ErrorCode::INVALID_PARAMS) {
+            LOG(WARNING) << "NoF segment remount: segment_name=" << segment.name
+                         << ", warn=invalid_params";
+        } else if (err == ErrorCode::SEGMENT_ALREADY_EXISTS) {
+            LOG(WARNING) << "NoF segment remount: segment_name=" << segment.name
+                         << ", warn=segment_already_exists";
+        } else if (err != ErrorCode::OK) {
+            LOG(ERROR) << "NoF segment remount: segment_name=" << segment.name
+                       << ", error=unexpected_error (" << err << ")";
+        }
+    }
+
+    return ErrorCode::OK;
+}
+
+ErrorCode ScopedNoFSegmentAccess::PrepareUnmountSegment(
+    const UUID& segment_id, size_t& metrics_dec_capacity) {
+    auto it = nof_segment_manager_->mounted_segments_.find(segment_id);
+    if (it == nof_segment_manager_->mounted_segments_.end()) {
+        LOG(WARNING) << "NoF segment unmount: segment_id=" << segment_id
+                     << ", warn=segment_not_found";
+        return ErrorCode::SEGMENT_NOT_FOUND;
+    }
+    if (it->second.status == SegmentStatus::UNMOUNTING) {
+        LOG(ERROR) << "NoF segment unmount: segment_id=" << segment_id
+                   << ", error=segment_is_unmounting";
+        return ErrorCode::UNAVAILABLE_IN_CURRENT_STATUS;
+    }
+
+    auto& mounted_segment = it->second;
+    auto& segment = mounted_segment.segment;
+    metrics_dec_capacity = segment.size;
+
+    std::shared_ptr<BufferAllocatorBase> allocator =
+        mounted_segment.buf_allocator;
+    if (HasAllocator(nof_segment_manager_->allocator_manager_, segment.name,
+                     allocator)) {
+        nof_segment_manager_->allocator_manager_.removeAllocator(segment.name,
+                                                                 allocator);
+    }
+
+    mounted_segment.buf_allocator.reset();
+    mounted_segment.status = SegmentStatus::UNMOUNTING;
+    return ErrorCode::OK;
+}
+
+ErrorCode ScopedNoFSegmentAccess::CommitUnmountSegment(
+    const UUID& segment_id, const UUID& client_id,
+    const size_t& metrics_dec_capacity) {
+    bool found_in_client_segments = false;
+    auto client_it = nof_segment_manager_->client_segments_.find(client_id);
+    if (client_it != nof_segment_manager_->client_segments_.end()) {
+        auto& segments = client_it->second;
+        auto segment_it =
+            std::find(segments.begin(), segments.end(), segment_id);
+        if (segment_it != segments.end()) {
+            segments.erase(segment_it);
+            found_in_client_segments = true;
+        }
+        if (segments.empty()) {
+            nof_segment_manager_->client_segments_.erase(client_it);
+        }
+    }
+    if (!found_in_client_segments) {
+        LOG(ERROR) << "NoF segment unmount: segment_id=" << segment_id
+                   << ", error=segment_not_found_in_client_segments";
+    }
+
+    std::string segment_name;
+    auto segment_it = nof_segment_manager_->mounted_segments_.find(segment_id);
+    if (segment_it != nof_segment_manager_->mounted_segments_.end()) {
+        segment_name = segment_it->second.segment.name;
+        nof_segment_manager_->client_by_name_.erase(segment_name);
+    }
+
+    nof_segment_manager_->mounted_segments_.erase(segment_id);
+    MasterMetricManager::instance().dec_total_nof_capacity(
+        segment_name, metrics_dec_capacity);
+
+    return ErrorCode::OK;
+}
+
+ErrorCode ScopedNoFSegmentAccess::GetClientSegments(
+    const UUID& client_id, std::vector<NoFSegment>& segments) const {
+    auto it = nof_segment_manager_->client_segments_.find(client_id);
+    if (it == nof_segment_manager_->client_segments_.end()) {
+        return ErrorCode::SEGMENT_NOT_FOUND;
+    }
+    segments.clear();
+    for (auto& segment_id : it->second) {
+        auto segment_it =
+            nof_segment_manager_->mounted_segments_.find(segment_id);
+        if (segment_it != nof_segment_manager_->mounted_segments_.end()) {
+            segments.emplace_back(segment_it->second.segment);
+        }
+    }
+    return ErrorCode::OK;
+}
+
+ErrorCode ScopedNoFSegmentAccess::GetMountedSegments(
+    std::vector<MountedNoFSegmentSnapshot>& segments) const {
+    segments.clear();
+    segments.reserve(nof_segment_manager_->mounted_segments_.size());
+    for (const auto& it : nof_segment_manager_->mounted_segments_) {
+        segments.push_back(MountedNoFSegmentSnapshot{
+            .segment_id = it.first,
+            .client_id = it.second.client_id,
+            .segment = it.second.segment,
+            .status = it.second.status,
+        });
+    }
+    return ErrorCode::OK;
+}
+
+ErrorCode ScopedNoFSegmentAccess::GetAllSegments(
+    std::vector<std::string>& all_segments) {
+    all_segments.clear();
+    for (auto& segment : nof_segment_manager_->mounted_segments_) {
+        if (segment.second.status == SegmentStatus::OK) {
+            all_segments.push_back(segment.second.segment.name);
+        }
+    }
+    return ErrorCode::OK;
+}
+
+ErrorCode ScopedNoFSegmentAccess::QuerySegments(const std::string& segment,
+                                                size_t& used,
+                                                size_t& capacity) {
+    size_t total_used = 0, total_capacity = 0;
+    const auto& allocator_manager = nof_segment_manager_->allocator_manager_;
+    const auto& allocators = allocator_manager.getAllocators(segment);
+    if (allocators != nullptr) {
+        for (const auto& allocator : *allocators) {
+            total_used += allocator->size();
+            total_capacity += allocator->capacity();
+        }
+    }
+
+    if (total_capacity == 0) {
+        VLOG(1) << "NoF segment query: segment=" << segment
+                << ", error=segment_not_found";
+        return ErrorCode::SEGMENT_NOT_FOUND;
+    }
+
+    used = total_used;
+    capacity = total_capacity;
+    return ErrorCode::OK;
+}
+
+void NoFSegmentManager::GetMountedSegmentsSnapshot(
+    std::vector<MountedNoFSegmentSnapshot>& segments) const {
+    std::shared_lock<std::shared_mutex> lock(segment_mutex_);
+    segments.clear();
+    segments.reserve(mounted_segments_.size());
+    for (const auto& [segment_id, mounted_segment] : mounted_segments_) {
+        segments.push_back(MountedNoFSegmentSnapshot{
+            segment_id, mounted_segment.client_id, mounted_segment.segment,
+            mounted_segment.status});
+    }
+}
+
 void SegmentManager::initializeCxlAllocator(const std::string& cxl_path,
                                             const size_t cxl_size) {
     LOG(INFO) << "Init CXL global allocator.";
@@ -990,4 +1379,22 @@ void SegmentManager::initializeCxlAllocator(const std::string& cxl_path,
         cxl_path, DEFAULT_CXL_BASE, cxl_size, cxl_path);
     MasterMetricManager::instance().inc_total_mem_capacity(cxl_path, cxl_size);
 }
+
+int64_t ScopedLocalDiskSegmentAccess::getSsdTotalCapacity(
+    const std::string& segment_name) const {
+    auto client_it = client_by_name_.find(segment_name);
+    if (client_it == client_by_name_.end()) return 0;
+    auto disk_it = client_local_disk_segment_.find(client_it->second);
+    if (disk_it == client_local_disk_segment_.end()) return 0;
+    return disk_it->second->ssd_total_capacity_bytes;
+}
+
+int64_t ScopedLocalDiskSegmentAccess::getSsdUsedBytes(
+    const std::string& segment_name) const {
+    auto client_it = client_by_name_.find(segment_name);
+    if (client_it == client_by_name_.end()) return 0;
+    auto disk_it = client_local_disk_segment_.find(client_it->second);
+    if (disk_it == client_local_disk_segment_.end()) return 0;
+    return disk_it->second->ssd_used_bytes.load(std::memory_order_relaxed);
+}
 }  // namespace mooncake
diff --git a/mooncake-store/src/shm_helper.cpp b/mooncake-store/src/shm_helper.cpp
index bee0021b..d1dbd424 100644
--- a/mooncake-store/src/shm_helper.cpp
+++ b/mooncake-store/src/shm_helper.cpp
@@ -54,7 +54,8 @@ bool ShmHelper::cleanup() {
         }
         if (shm->base_addr) {
 #if defined(USE_ASCEND_DIRECT)
-            if (globalConfig().ascend_agent_mode) {
+            if (globalConfig().ascend_agent_mode &&
+                globalConfig().ascend_use_fabric_mem) {
                 free_memory("ascend", shm->base_addr);
                 continue;
             }
@@ -73,22 +74,28 @@ bool ShmHelper::cleanup() {
 
 void* ShmHelper::allocate(size_t size) {
     std::lock_guard<std::mutex> lock(shm_mutex_);
+    // Dummy-real: FabricMem host uses VMM; non-Fabric host uses memfd+mmap like
+    // non-agent / GPU shm path.
 #ifdef USE_ASCEND_DIRECT
     if (globalConfig().ascend_agent_mode) {
-        void* base_addr = nullptr;
-        size_t alloc_size = size;
-        base_addr = ascend_allocate_vmm_memory_direct(alloc_size);
-        if (base_addr == nullptr) {
-            throw std::runtime_error("Failed to allocate VMM shared memory");
+        if (globalConfig().ascend_use_fabric_mem) {
+            void* base_addr = nullptr;
+            size_t alloc_size = size;
+            base_addr = ascend_allocate_vmm_memory_direct(alloc_size);
+            if (base_addr == nullptr) {
+                throw std::runtime_error(
+                    "Failed to allocate VMM shared memory");
+            }
+            auto shm = std::make_shared<ShmSegment>();
+            shm->fd = -1;
+            shm->base_addr = base_addr;
+            shm->size = alloc_size;
+            shm->name = MOONCAKE_SHM_NAME;
+            shm->registered = false;
+            shms_.push_back(shm);
+            return base_addr;
         }
-        auto shm = std::make_shared<ShmSegment>();
-        shm->fd = -1;
-        shm->base_addr = base_addr;
-        shm->size = alloc_size;
-        shm->name = MOONCAKE_SHM_NAME;
-        shm->registered = false;
-        shms_.push_back(shm);
-        return base_addr;
+        // ascend_agent_mode && !ascend_use_fabric_mem: fall through to memfd
     }
 #endif
 
@@ -142,7 +149,8 @@ int ShmHelper::free(void* addr) {
             }
             if ((*it)->base_addr) {
 #if defined(USE_ASCEND_DIRECT)
-                if (globalConfig().ascend_agent_mode) {
+                if (globalConfig().ascend_agent_mode &&
+                    globalConfig().ascend_use_fabric_mem) {
                     free_memory("ascend", (*it)->base_addr);
                 } else
 #endif
diff --git a/mooncake-store/src/spdk/CMakeLists.txt b/mooncake-store/src/spdk/CMakeLists.txt
new file mode 100644
index 00000000..31d60847
--- /dev/null
+++ b/mooncake-store/src/spdk/CMakeLists.txt
@@ -0,0 +1,5 @@
+set(SPDK_WRAPPER_SOURCES
+    ${CMAKE_CURRENT_SOURCE_DIR}/spdk_wrapper.cpp
+)
+
+set(SPDK_WRAPPER_SOURCES ${SPDK_WRAPPER_SOURCES} PARENT_SCOPE)
\ No newline at end of file
diff --git a/mooncake-store/src/spdk/spdk_wrapper.cpp b/mooncake-store/src/spdk/spdk_wrapper.cpp
new file mode 100644
index 00000000..5042c8f4
--- /dev/null
+++ b/mooncake-store/src/spdk/spdk_wrapper.cpp
@@ -0,0 +1,490 @@
+#include <glog/logging.h>
+
+#include <atomic>
+#include <cerrno>
+#include <chrono>
+#include <cstring>
+#include <cstdlib>
+#include <thread>
+#include "spdk/spdk_wrapper.h"
+
+namespace mooncake {
+namespace {
+
+bool ParseEnvU64(const char *name, uint64_t *out) {
+    const char *val = std::getenv(name);
+    if (!val || *val == '\0') {
+        return false;
+    }
+
+    errno = 0;
+    char *end = nullptr;
+    unsigned long long parsed = std::strtoull(val, &end, 10);
+    if (errno != 0 || end == val || (end && *end != '\0')) {
+        LOG(WARNING) << "Invalid value for " << name << ": " << val;
+        return false;
+    }
+
+    *out = static_cast<uint64_t>(parsed);
+    return true;
+}
+
+bool ParseEnvBool(const char *name, bool *out) {
+    uint64_t v = 0;
+    if (!ParseEnvU64(name, &v)) {
+        return false;
+    }
+    *out = (v != 0);
+    return true;
+}
+
+void ApplyCtrlrOptsFromEnv(struct spdk_nvme_ctrlr_opts *opts) {
+    uint64_t v = 0;
+    bool bv = false;
+    opts->keep_alive_timeout_ms = 0;
+
+    if (ParseEnvU64("MC_NVME_NUM_IO_QUEUES", &v)) {
+        opts->num_io_queues = static_cast<uint32_t>(v);
+    }
+    if (ParseEnvU64("MC_NVME_IO_QUEUE_SIZE", &v)) {
+        opts->io_queue_size = static_cast<uint32_t>(v);
+    }
+    if (ParseEnvU64("MC_NVME_IO_QUEUE_REQUESTS", &v)) {
+        opts->io_queue_requests = static_cast<uint32_t>(v);
+    }
+    if (ParseEnvU64("MC_NVME_TRANSPORT_ACK_TIMEOUT", &v)) {
+        opts->transport_ack_timeout = static_cast<uint8_t>(v);
+    }
+    if (ParseEnvU64("MC_NVME_ADMIN_QUEUE_SIZE", &v)) {
+        opts->admin_queue_size = static_cast<uint16_t>(v);
+    }
+    if (ParseEnvU64("MC_NVME_FABRICS_CONNECT_TIMEOUT_US", &v)) {
+        opts->fabrics_connect_timeout_us = v;
+    }
+    if (ParseEnvBool("MC_NVME_HEADER_DIGEST", &bv)) {
+        opts->header_digest = bv;
+    }
+    if (ParseEnvBool("MC_NVME_DATA_DIGEST", &bv)) {
+        opts->data_digest = bv;
+    }
+    LOG(INFO) << "NVMe ctrlr opts: num_io_queues=" << opts->num_io_queues
+              << ", io_queue_size=" << opts->io_queue_size
+              << ", io_queue_requests=" << opts->io_queue_requests
+              << ", keep_alive_timeout_ms=" << opts->keep_alive_timeout_ms
+              << ", transport_ack_timeout="
+              << static_cast<int>(opts->transport_ack_timeout)
+              << ", admin_queue_size=" << opts->admin_queue_size
+              << ", fabrics_connect_timeout_us="
+              << opts->fabrics_connect_timeout_us
+              << ", header_digest=" << opts->header_digest
+              << ", data_digest=" << opts->data_digest;
+}
+
+}  // namespace
+
+struct nof_seg_handle {
+    struct spdk_nvme_qpair *qpair;
+    struct spdk_nvme_ns *ns;
+};
+
+struct tr_info {
+    struct spdk_nvme_transport_id trid;
+    std::string ctrlr_key;
+    uint32_t ns;
+};
+
+struct ctrlr_info {
+    struct spdk_nvme_ctrlr *ctrlr;
+    std::map<uint32_t, std::unique_ptr<nof_seg_handle>> ns_seg;
+    std::mutex ns_mutex;
+};
+
+SpdkWrapper::SpdkWrapper() = default;
+
+SpdkWrapper::~SpdkWrapper() { Cleanup(); }
+
+SpdkWrapper &SpdkWrapper::GetInstance() {
+    static SpdkWrapper ins;
+    return ins;
+}
+
+bool SpdkWrapper::InitializeEnv() {
+    if (initialized.load(std::memory_order_acquire)) {
+        return true;
+    }
+
+    std::lock_guard<std::mutex> lock(init_mutex);
+    if (initialized.load(std::memory_order_acquire)) {
+        return true;
+    }
+
+    struct spdk_env_opts opts;
+    spdk_env_opts_init(&opts);
+    opts.name = "mooncake";
+
+    int rc = spdk_env_init(&opts);
+    if (rc != 0) {
+        fprintf(stderr, "SPDK init failed: %d\n", rc);
+        return false;
+    }
+
+    // Mark SPDK as initialized.
+    initialized.store(true, std::memory_order_release);
+    return true;
+}
+
+void SpdkWrapper::Cleanup() {
+    if (initialized.load(std::memory_order_acquire)) {
+        {
+            std::lock_guard<std::mutex> lock(ctrlrs_mutex);
+            for (auto &[_, info] : connected_ctrlrs) {
+                if (info) {
+                    // Free all qpairs and segment handles
+                    for (auto &[_, seg] : info->ns_seg) {
+                        if (seg && seg->qpair) {
+                            spdk_nvme_ctrlr_free_io_qpair(seg->qpair);
+                        }
+                    }
+                    // Detach controller
+                    if (info->ctrlr) {
+                        spdk_nvme_detach(info->ctrlr);
+                    }
+                }
+            }
+            connected_ctrlrs.clear();
+        }
+
+        {
+            std::lock_guard<std::mutex> lock(probe_buffers_mutex_);
+            for (auto &[_, probe_buffer] : probe_buffers_) {
+                if (probe_buffer && probe_buffer->ptr) {
+                    spdk_free(probe_buffer->ptr);
+                    probe_buffer->ptr = nullptr;
+                    probe_buffer->size = 0;
+                }
+            }
+            probe_buffers_.clear();
+        }
+        spdk_env_fini();
+        initialized.store(false, std::memory_order_release);
+    }
+}
+
+void *SpdkWrapper::Alloc(size_t size, size_t align, int socket_id) {
+    if (!InitializeEnv()) {
+        return nullptr;
+    }
+
+    return spdk_zmalloc(size, align, nullptr, socket_id, SPDK_MALLOC_DMA);
+}
+
+void SpdkWrapper::Free(void *ptr) {
+    if (ptr) {
+        spdk_free(ptr);
+    }
+}
+
+void SpdkWrapper::ProbeReadComplete(void *ctx,
+                                    const struct spdk_nvme_cpl *cpl) {
+    auto *probe_ctx = reinterpret_cast<ProbeRequestContext *>(ctx);
+    if (spdk_nvme_cpl_is_error(cpl)) {
+        {
+            std::lock_guard<std::mutex> lock(probe_ctx->error_mutex);
+            probe_ctx->error_reason =
+                std::string("completion_error:") +
+                spdk_nvme_cpl_get_status_string(&cpl->status);
+        }
+        probe_ctx->success.store(false, std::memory_order_release);
+    } else {
+        probe_ctx->success.store(true, std::memory_order_release);
+    }
+    probe_ctx->done.store(true, std::memory_order_release);
+    if (probe_ctx->owner != nullptr) {
+        probe_ctx->owner->RecycleProbeRequestContext(probe_ctx);
+    }
+}
+
+void SpdkWrapper::ReplenishProbeRequestContextPoolLocked(size_t count) {
+    for (size_t i = 0; i < count; ++i) {
+        auto probe_ctx = std::make_unique<ProbeRequestContext>();
+        probe_request_context_pool_.push(probe_ctx.get());
+        probe_request_contexts_.push_back(std::move(probe_ctx));
+    }
+}
+
+SpdkWrapper::ProbeRequestContext *SpdkWrapper::AcquireProbeRequestContext() {
+    std::lock_guard<std::mutex> lock(probe_request_context_pool_mutex_);
+    if (probe_request_context_pool_.empty()) {
+        ReplenishProbeRequestContextPoolLocked(8);
+    }
+    auto *probe_ctx = probe_request_context_pool_.top();
+    probe_request_context_pool_.pop();
+    probe_ctx->Reset(this);
+    return probe_ctx;
+}
+
+void SpdkWrapper::RecycleProbeRequestContext(ProbeRequestContext *ctx) {
+    if (ctx == nullptr) {
+        return;
+    }
+    std::lock_guard<std::mutex> lock(probe_request_context_pool_mutex_);
+    probe_request_context_pool_.push(ctx);
+}
+
+int64_t SpdkWrapper::NvmePollProcessCompletion(nof_seg_handle *seg,
+                                               uint32_t complete_per_seg) {
+    return spdk_nvme_qpair_process_completions(seg->qpair, complete_per_seg);
+}
+
+int SpdkWrapper::ParseTransPortStr(const std::string &tr_str, tr_info *info) {
+    std::memset(&info->trid, 0, sizeof(info->trid));
+    info->ns = 1;
+
+    if (spdk_nvme_transport_id_parse(&info->trid, tr_str.c_str()) != 0) {
+        LOG(ERROR) << "Error parsing transport address";
+        return -1;
+    }
+
+    std::string ns_prefix = "ns:";
+    size_t ns_pos = tr_str.find(ns_prefix);
+    if (ns_pos != std::string::npos) {
+        size_t ns_start = ns_pos + ns_prefix.length();
+        size_t ns_end = tr_str.find_first_of(" \t", ns_start);
+
+        std::string ns_str;
+        if (ns_end == std::string::npos) {
+            ns_str = tr_str.substr(ns_start);
+        } else {
+            ns_str = tr_str.substr(ns_start, ns_end - ns_start);
+        }
+
+        try {
+            info->ns = std::stoul(ns_str);
+        } catch (const std::exception &e) {
+            LOG(ERROR) << "Failed to parse ns value: " << ns_str
+                       << ", error: " << e.what();
+            return -1;
+        }
+    } else {
+        LOG(ERROR) << "No ns field found in transport string";
+    }
+
+    info->ctrlr_key = std::string(info->trid.traddr) + "|" +
+                      std::string(info->trid.trsvcid) + "|" +
+                      std::string(info->trid.subnqn) + "|" +
+                      std::to_string(static_cast<int>(info->trid.trtype));
+
+    LOG(INFO) << "traddr:" << info->trid.traddr
+              << "trsvcid:" << info->trid.trsvcid << "ns:" << info->ns
+              << "subnqn:" << info->trid.subnqn
+              << "trtype:" << info->trid.trtype;
+
+    return 0;
+}
+
+int SpdkWrapper::ConnectController(const struct spdk_nvme_transport_id *trid,
+                                   ctrlr_info *info) {
+    auto probe_cb = [](void *cb_ctx, const struct spdk_nvme_transport_id *trid,
+                       struct spdk_nvme_ctrlr_opts *opts) -> bool {
+        ApplyCtrlrOptsFromEnv(opts);
+        LOG(INFO) << "Attaching to " << trid->traddr << " " << trid->subnqn;
+        return true;
+    };
+
+    auto attach_cb = [](void *cb_ctx, const struct spdk_nvme_transport_id *trid,
+                        struct spdk_nvme_ctrlr *ctrlr,
+                        const struct spdk_nvme_ctrlr_opts *opts) {
+        LOG(INFO) << "Attached to " << trid->traddr << " " << trid->subnqn;
+        ctrlr_info *info = (ctrlr_info *)cb_ctx;
+        info->ctrlr = ctrlr;
+    };
+
+    return spdk_nvme_probe(trid, (void *)info, probe_cb, attach_cb, NULL);
+}
+
+nof_seg_handle *SpdkWrapper::OpenNofSegment(const std::string &tr_str) {
+    tr_info tr;
+    int ret = ParseTransPortStr(tr_str, &tr);
+    if (ret != 0) {
+        return nullptr;
+    }
+
+    ctrlr_info *info = nullptr;
+    {
+        std::lock_guard<std::mutex> lock(ctrlrs_mutex);
+        auto it = connected_ctrlrs.find(tr.ctrlr_key);
+        if (it == connected_ctrlrs.end()) {
+            auto new_info = std::make_unique<ctrlr_info>();
+            info = new_info.get();
+
+            ret = ConnectController(&tr.trid, info);
+            if (ret != 0) {
+                return nullptr;
+            }
+
+            connected_ctrlrs[tr.ctrlr_key] = std::move(new_info);
+        } else {
+            info = it->second.get();
+        }
+    }
+
+    nof_seg_handle *seg_handle = nullptr;
+    struct spdk_nvme_qpair *qpair = nullptr;
+    struct spdk_nvme_ns *ns = nullptr;
+    {
+        auto &ns_seg = info->ns_seg;
+        std::lock_guard<std::mutex> lock(info->ns_mutex);
+        auto ns_it = ns_seg.find(tr.ns);
+        if (ns_it != ns_seg.end()) {
+            return ns_it->second.get();
+        }
+
+        if (spdk_nvme_ctrlr_is_active_ns(info->ctrlr, tr.ns)) {
+            ns = spdk_nvme_ctrlr_get_ns(info->ctrlr, tr.ns);
+        } else {
+            LOG(ERROR) << "spdk_nvme_ctrlr_is_active_ns failed";
+            return nullptr;
+        }
+
+        qpair = spdk_nvme_ctrlr_alloc_io_qpair(info->ctrlr, nullptr, 0);
+        if (!qpair) {
+            LOG(ERROR) << "alloc spdk_nvme_qpair failed";
+            return nullptr;
+        }
+
+        auto new_seg = std::make_unique<nof_seg_handle>();
+        new_seg->qpair = qpair;
+        new_seg->ns = ns;
+        seg_handle = new_seg.get();
+        ns_seg[tr.ns] = std::move(new_seg);
+    }
+
+    return seg_handle;
+}
+
+uint32_t SpdkWrapper::GetBlockSize(const nof_seg_handle *seg_handle) {
+    if (!seg_handle || !seg_handle->ns) {
+        return INVALID_BLOCK_SIZE;
+    }
+
+    return spdk_nvme_ns_get_sector_size(seg_handle->ns);
+}
+
+int SpdkWrapper::SubmitRequest(const nof_seg_handle *seg_handle, void *ptr,
+                               uint64_t lba, uint32_t lba_count, int op,
+                               spdk_nvme_cmd_cb cb_fn, void *cb_ctx) {
+    if (!seg_handle || !ptr || !lba_count || !seg_handle->qpair ||
+        !seg_handle->ns) {
+        return -1;
+    }
+
+    struct spdk_nvme_qpair *qpair = seg_handle->qpair;
+    struct spdk_nvme_ns *ns = seg_handle->ns;
+    if (op == kSpdkNofOpRead) {
+        return spdk_nvme_ns_cmd_read(ns, qpair, ptr, lba, lba_count, cb_fn,
+                                     cb_ctx, 0);
+    } else if (op == kSpdkNofOpWrite) {
+        return spdk_nvme_ns_cmd_write(ns, qpair, ptr, lba, lba_count, cb_fn,
+                                      cb_ctx, 0);
+    }
+    return -1;
+}
+
+SpdkWrapper::ProbeBuffer *SpdkWrapper::GetOrCreateProbeBuffer(
+    const std::string &tr_str, uint32_t block_size, std::string *error_reason) {
+    std::lock_guard<std::mutex> lock(probe_buffers_mutex_);
+    auto &probe_buffer = probe_buffers_[tr_str];
+    if (!probe_buffer) {
+        probe_buffer = std::make_unique<ProbeBuffer>();
+    }
+
+    if (probe_buffer->ptr != nullptr && probe_buffer->size == block_size) {
+        return probe_buffer.get();
+    }
+
+    if (probe_buffer->ptr != nullptr) {
+        spdk_free(probe_buffer->ptr);
+        probe_buffer->ptr = nullptr;
+        probe_buffer->size = 0;
+    }
+
+    probe_buffer->ptr =
+        spdk_zmalloc(block_size, 0x1000, nullptr, -1, SPDK_MALLOC_DMA);
+    if (!probe_buffer->ptr) {
+        if (error_reason) {
+            *error_reason = "alloc_fail";
+        }
+        return nullptr;
+    }
+    probe_buffer->size = block_size;
+    return probe_buffer.get();
+}
+
+bool SpdkWrapper::ProbeNofSegment(const std::string &tr_str,
+                                  uint32_t timeout_ms,
+                                  std::string *error_reason) {
+    if (!InitializeEnv()) {
+        if (error_reason) {
+            *error_reason = "spdk_env_init_fail";
+        }
+        return false;
+    }
+
+    nof_seg_handle *seg_handle = OpenNofSegment(tr_str);
+    if (!seg_handle) {
+        if (error_reason) {
+            *error_reason = "open_fail";
+        }
+        return false;
+    }
+
+    uint32_t block_size = GetBlockSize(seg_handle);
+    if (block_size == INVALID_BLOCK_SIZE || block_size == 0) {
+        if (error_reason) {
+            *error_reason = "invalid_block_size";
+        }
+        return false;
+    }
+
+    ProbeBuffer *probe_buffer =
+        GetOrCreateProbeBuffer(tr_str, block_size, error_reason);
+    if (!probe_buffer || !probe_buffer->ptr) {
+        return false;
+    }
+
+    ProbeRequestContext *probe_ctx = AcquireProbeRequestContext();
+    int ret = SubmitRequest(seg_handle, probe_buffer->ptr, 0, 1, kSpdkNofOpRead,
+                            ProbeReadComplete, probe_ctx);
+    if (ret != 0) {
+        RecycleProbeRequestContext(probe_ctx);
+        if (error_reason) {
+            *error_reason = "submit_fail";
+        }
+        return false;
+    }
+
+    auto deadline = std::chrono::steady_clock::now() +
+                    std::chrono::milliseconds(timeout_ms);
+    while (!probe_ctx->done.load(std::memory_order_acquire) &&
+           std::chrono::steady_clock::now() < deadline) {
+        NvmePollProcessCompletion(seg_handle, 0);
+        std::this_thread::sleep_for(std::chrono::milliseconds(1));
+    }
+
+    bool ok = probe_ctx->done.load(std::memory_order_acquire) &&
+              probe_ctx->success.load(std::memory_order_acquire);
+    if (!ok && error_reason) {
+        if (!probe_ctx->done.load(std::memory_order_acquire)) {
+            *error_reason = "completion_timeout";
+        } else {
+            std::lock_guard<std::mutex> lock(probe_ctx->error_mutex);
+            *error_reason = probe_ctx->error_reason.empty()
+                                ? "completion_error"
+                                : probe_ctx->error_reason;
+        }
+    }
+
+    return ok;
+}
+
+}  // namespace mooncake
diff --git a/mooncake-store/src/ssd_register_client.cpp b/mooncake-store/src/ssd_register_client.cpp
new file mode 100644
index 00000000..39cc9f39
--- /dev/null
+++ b/mooncake-store/src/ssd_register_client.cpp
@@ -0,0 +1,138 @@
+#include "ssd_register_client.h"
+#include <algorithm>
+#include <cctype>
+#include <cstdlib>
+
+namespace mooncake {
+
+NoFRegisterClient::NoFRegisterClient()
+    : master_client_(generate_uuid(), nullptr) {}
+
+NoFRegisterClient::~NoFRegisterClient() = default;
+
+int NoFRegisterClient::set_register(const std::string &nqn, size_t nsid,
+                                    const std::string &traddr, size_t trsvcid,
+                                    uintptr_t base, size_t size,
+                                    const std::string &master_server_addr) {
+    LOG(INFO) << "Registering SSD: nqn=" << nqn << ",nsid=" << nsid
+              << ",traddr=" << traddr << ",trsvcid=" << trsvcid
+              << ",master=" << master_server_addr << ",base=" << base
+              << ",size=" << size;
+
+    auto err = master_client_.Connect(master_server_addr);
+    if (err != ErrorCode::OK) {
+        LOG(ERROR) << "Failed to connect to master";
+        return OPERATION_FAILED;
+    }
+
+    const char *trtype_env = std::getenv("MC_NOF_TRTYPE");
+    std::string trtype = trtype_env ? trtype_env : "RDMA";
+    std::transform(
+        trtype.begin(), trtype.end(), trtype.begin(),
+        [](unsigned char c) { return static_cast<char>(std::toupper(c)); });
+    if (trtype != "RDMA" && trtype != "TCP") {
+        LOG(WARNING) << "Invalid MC_NOF_TRTYPE=" << trtype
+                     << ", fallback to RDMA";
+        trtype = "RDMA";
+    }
+
+    std::string te_endpoint = "traddr:" + traddr +
+                              " trsvcid:" + std::to_string(trsvcid) +
+                              " subnqn:" + nqn + " trtype:" + trtype +
+                              " adrfam:IPv4 ns:" + std::to_string(nsid);
+
+    NoFSegment segment;
+    segment.base = base;
+    segment.size = size;
+    segment.id = generate_uuid();
+    segment.name = te_endpoint;
+    segment.te_endpoint = te_endpoint;
+    auto mount_result = master_client_.MountNoFSegment(segment);
+    if (!mount_result) {
+        LOG(ERROR) << "mount_segment_to_master_failed ";
+        return OPERATION_FAILED;
+    }
+
+    return OPERATION_OK;
+}
+
+int NoFRegisterClient::set_unregister_by_endpoint(
+    const std::string &nqn, size_t nsid, const std::string &traddr,
+    size_t trsvcid, const std::string &master_server_addr) {
+    LOG(INFO) << "Unregistering SSD by endpoint: nqn=" << nqn
+              << ",nsid=" << nsid << ",traddr=" << traddr
+              << ",trsvcid=" << trsvcid << ",master=" << master_server_addr;
+
+    // Connect to master server
+    auto err = master_client_.Connect(master_server_addr);
+    if (err != ErrorCode::OK) {
+        LOG(ERROR) << "Failed to connect to master: " << static_cast<int>(err);
+        return OPERATION_FAILED;
+    }
+
+    const char *trtype_env = std::getenv("MC_NOF_TRTYPE");
+    std::string trtype = trtype_env ? trtype_env : "RDMA";
+    std::transform(
+        trtype.begin(), trtype.end(), trtype.begin(),
+        [](unsigned char c) { return static_cast<char>(std::toupper(c)); });
+    if (trtype != "RDMA" && trtype != "TCP") {
+        LOG(WARNING) << "Invalid MC_NOF_TRTYPE=" << trtype
+                     << ", fallback to RDMA";
+        trtype = "RDMA";
+    }
+
+    // Build the te_endpoint string to match registered segments
+    std::string te_endpoint = "traddr:" + traddr +
+                              " trsvcid:" + std::to_string(trsvcid) +
+                              " subnqn:" + nqn + " trtype:" + trtype +
+                              " adrfam:IPv4 ns:" + std::to_string(nsid);
+
+    LOG(INFO) << "Built te_endpoint: " << te_endpoint;
+
+    auto matching_segments_result =
+        master_client_.GetNoFSegmentsByName(te_endpoint);
+    if (!matching_segments_result) {
+        LOG(ERROR) << "Failed to get NoF segments by name: "
+                   << static_cast<int>(matching_segments_result.error());
+        return OPERATION_FAILED;
+    }
+
+    std::vector<NoFSegmentOwnerInfo> matching_segments =
+        matching_segments_result.value();
+    LOG(INFO) << "Retrieved " << matching_segments.size()
+              << " mounted NoF segments for te_endpoint";
+    if (matching_segments.empty()) {
+        LOG(ERROR) << "No segment found for te_endpoint: " << te_endpoint;
+        return OPERATION_FAILED;
+    }
+
+    // Unmount all matching segments
+    bool all_unmounted = true;
+    for (const auto &segment : matching_segments) {
+        LOG(INFO) << "Found matching segment: id=" << segment.segment_id
+                  << ", owner_client_id=" << segment.client_id;
+        MasterClient owner_master_client(segment.client_id, nullptr);
+        err = owner_master_client.Connect(master_server_addr);
+        if (err != ErrorCode::OK) {
+            LOG(ERROR) << "Failed to connect owner master client for segment "
+                       << segment.segment_id << ": " << static_cast<int>(err);
+            all_unmounted = false;
+            continue;
+        }
+
+        auto unmount_result =
+            owner_master_client.UnmountNoFSegment(segment.segment_id);
+        if (!unmount_result) {
+            LOG(ERROR) << "Failed to unmount segment " << segment.segment_id
+                       << ": " << static_cast<int>(unmount_result.error());
+            all_unmounted = false;
+        } else {
+            LOG(INFO) << "Successfully unmounted segment " << segment.segment_id
+                      << ", owner_client_id=" << segment.client_id;
+        }
+    }
+
+    return all_unmounted ? OPERATION_OK : OPERATION_FAILED;
+}
+
+}  // namespace mooncake
diff --git a/mooncake-store/src/storage/distributed/distributed_storage_backend.cpp b/mooncake-store/src/storage/distributed/distributed_storage_backend.cpp
new file mode 100644
index 00000000..7c8e6ffb
--- /dev/null
+++ b/mooncake-store/src/storage/distributed/distributed_storage_backend.cpp
@@ -0,0 +1,320 @@
+#include "storage/distributed/distributed_storage_backend.h"
+
+#include <fmt/format.h>
+
+#include <algorithm>
+#include <cstdlib>
+#include <xxhash.h>
+#include <filesystem>
+
+#include "utils.h"
+
+namespace mooncake {
+
+// === DistributedStorageConfig ===
+
+bool DistributedStorageConfig::Validate() const {
+    if (fsdir.empty()) {
+        LOG(ERROR) << "DistributedStorageConfig: fsdir is empty";
+        return false;
+    }
+    if (!std::filesystem::path(fsdir).is_absolute()) {
+        LOG(ERROR)
+            << "DistributedStorageConfig: fsdir must be an absolute path: "
+            << fsdir;
+        return false;
+    }
+    if (fs_adapter_type.empty()) {
+        LOG(ERROR) << "DistributedStorageConfig: fs_adapter_type is empty";
+        return false;
+    }
+    if (fs_adapter_type != "hf3fs") {
+        LOG(ERROR) << "DistributedStorageConfig: unsupported fs_adapter_type: "
+                   << fs_adapter_type;
+        return false;
+    }
+    if (hash_bucket_count <= 0) {
+        LOG(ERROR) << "DistributedStorageConfig: hash_bucket_count must > 0";
+        return false;
+    }
+    return true;
+}
+
+DistributedStorageConfig DistributedStorageConfig::FromEnvironment() {
+    DistributedStorageConfig config;
+    config.fsdir =
+        GetEnvStringOr("MOONCAKE_DISTRIBUTED_ROOT_DIR", config.fsdir);
+    if (!std::filesystem::path(config.fsdir).is_absolute()) {
+        config.fsdir = std::filesystem::absolute(config.fsdir).string();
+    }
+    config.fs_adapter_type =
+        GetEnvStringOr("MOONCAKE_DISTRIBUTED_FS_TYPE", config.fs_adapter_type);
+    config.enable_health_check =
+        GetEnvOr<bool>("MOONCAKE_DISTRIBUTED_HEALTH_CHECK", false);
+    config.hash_bucket_count =
+        GetEnvOr<int>("MOONCAKE_DISTRIBUTED_HASH_BUCKET_COUNT", 256);
+    return config;
+}
+
+// === DistributedStorageBackend ===
+
+DistributedStorageBackend::DistributedStorageBackend(
+    const FileStorageConfig& file_storage_config,
+    const DistributedStorageConfig& distributed_config,
+    std::unique_ptr<FileSystemAdapter> fs_adapter)
+    : StorageBackendInterface(file_storage_config),
+      fs_adapter_(std::move(fs_adapter)),
+      distributed_config_(distributed_config),
+      root_dir_(distributed_config.fsdir),
+      hash_bucket_count_(distributed_config.hash_bucket_count) {}
+
+tl::expected<void, ErrorCode> DistributedStorageBackend::Init() {
+    if (initialized_) {
+        LOG(WARNING) << "DistributedStorageBackend is already initialized";
+        return {};
+    }
+
+    auto init_result = fs_adapter_->Init(root_dir_);
+    if (!init_result) return init_result;
+
+    // Ensure root directory exists before health check
+    std::error_code ec;
+    std::filesystem::create_directories(root_dir_, ec);
+    if (ec) {
+        LOG(ERROR) << "Failed to create root directory " << root_dir_ << ": "
+                   << ec.message();
+        return tl::make_unexpected(ErrorCode::FILE_WRITE_FAIL);
+    }
+
+    if (distributed_config_.enable_health_check) {
+        std::string probe_path =
+            fmt::format("{}/.mooncake_health_probe_{}", root_dir_,
+                        UuidToString(generate_uuid()));
+        std::string probe_data = "health_check";
+        auto write_result = fs_adapter_->WriteFile(
+            probe_path,
+            std::span<const char>(probe_data.data(), probe_data.size()));
+        if (!write_result) {
+            LOG(ERROR) << "DFS health check failed (write): "
+                       << static_cast<int>(write_result.error());
+            return tl::make_unexpected(write_result.error());
+        }
+
+        std::vector<char> read_buf(probe_data.size());
+        auto read_result =
+            fs_adapter_->ReadFile(probe_path, read_buf.data(), read_buf.size());
+        if (!read_result || *read_result != probe_data.size() ||
+            std::string(read_buf.data(), read_buf.size()) != probe_data) {
+            LOG(ERROR) << "DFS health check failed (read back mismatch)";
+            auto del_err = fs_adapter_->DeleteFile(probe_path);
+            if (!del_err) {
+                LOG(WARNING) << "Failed to delete health-check probe: "
+                             << static_cast<int>(del_err.error());
+            }
+            return tl::make_unexpected(ErrorCode::DFS_SERVICE_UNAVAILABLE);
+        }
+
+        fs_adapter_->DeleteFile(probe_path);
+        LOG(INFO) << "DFS health check passed, adapter="
+                  << fs_adapter_->GetName();
+    }
+
+    // Ensure hash bucket directories exist
+    for (int i = 0; i < hash_bucket_count_; ++i) {
+        std::string bucket_dir = fmt::format("{}/{:02x}", root_dir_, i);
+        std::filesystem::create_directories(bucket_dir, ec);
+        if (ec) {
+            LOG(ERROR) << "Failed to create bucket directory " << bucket_dir
+                       << ": " << ec.message();
+            return tl::make_unexpected(ErrorCode::FILE_WRITE_FAIL);
+        }
+    }
+
+    initialized_ = true;
+    return {};
+}
+
+tl::expected<int64_t, ErrorCode> DistributedStorageBackend::BatchOffload(
+    const std::unordered_map<std::string, std::vector<Slice>>& batch_object,
+    std::function<ErrorCode(const std::vector<std::string>& keys,
+                            std::vector<StorageObjectMetadata>& metadatas)>
+        complete_handler,
+    std::function<void(const std::vector<std::string>& evicted_keys)>
+        eviction_handler) {
+    if (!initialized_) {
+        LOG(ERROR) << "DistributedStorageBackend is not initialized";
+        return tl::make_unexpected(ErrorCode::INTERNAL_ERROR);
+    }
+
+    if (eviction_handler) {
+        LOG_FIRST_N(WARNING, 1)
+            << "DistributedStorageBackend does not support eviction, "
+               "eviction_handler ignored";
+    }
+
+    std::vector<std::string> success_keys;
+    std::vector<StorageObjectMetadata> success_metas;
+
+    for (const auto& [key, slices] : batch_object) {
+        auto path = GetObjectPath(key);
+
+        std::vector<iovec> iovs;
+        for (const auto& slice : slices) {
+            iovs.push_back({slice.ptr, slice.size});
+        }
+
+        auto result =
+            fs_adapter_->VectorWriteFile(path, iovs.data(), iovs.size(), 0);
+        if (!result) {
+            LOG(WARNING) << "Failed to offload key " << key << ": "
+                         << static_cast<int>(result.error());
+            continue;
+        }
+
+        success_keys.push_back(key);
+        StorageObjectMetadata meta{-1, 0, static_cast<int64_t>(key.size()),
+                                   static_cast<int64_t>(*result), ""};
+        success_metas.push_back(meta);
+    }
+
+    if (!success_keys.empty()) {
+        auto err = complete_handler(success_keys, success_metas);
+        if (err != ErrorCode::OK) {
+            return tl::make_unexpected(err);
+        }
+    }
+
+    return static_cast<int64_t>(success_keys.size());
+}
+
+tl::expected<void, ErrorCode> DistributedStorageBackend::BatchLoad(
+    std::unordered_map<std::string, Slice>& batched_slices) {
+    if (!initialized_) {
+        LOG(ERROR) << "DistributedStorageBackend is not initialized";
+        return tl::make_unexpected(ErrorCode::INTERNAL_ERROR);
+    }
+
+    for (auto& [key, slice] : batched_slices) {
+        auto path = GetObjectPath(key);
+
+        auto result = fs_adapter_->ReadFile(path, slice.ptr, slice.size);
+        if (!result) {
+            return tl::make_unexpected(result.error());
+        }
+        if (*result != slice.size) {
+            return tl::make_unexpected(ErrorCode::FILE_READ_FAIL);
+        }
+    }
+    return {};
+}
+
+tl::expected<bool, ErrorCode> DistributedStorageBackend::IsExist(
+    const std::string& key) {
+    if (!initialized_) {
+        LOG(ERROR) << "DistributedStorageBackend is not initialized";
+        return tl::make_unexpected(ErrorCode::INTERNAL_ERROR);
+    }
+
+    auto path = GetObjectPath(key);
+    return fs_adapter_->FileExists(path);
+}
+
+tl::expected<bool, ErrorCode> DistributedStorageBackend::IsEnableOffloading() {
+    return true;
+}
+
+tl::expected<void, ErrorCode> DistributedStorageBackend::ScanMeta(
+    const std::function<
+        ErrorCode(const std::vector<std::string>& keys,
+                  std::vector<StorageObjectMetadata>& metadatas)>& handler) {
+    if (!initialized_) {
+        LOG(ERROR) << "DistributedStorageBackend is not initialized";
+        return tl::make_unexpected(ErrorCode::INTERNAL_ERROR);
+    }
+
+    std::vector<std::string> batch_keys;
+    std::vector<StorageObjectMetadata> batch_metas;
+    const size_t batch_limit = static_cast<size_t>(std::max<int64_t>(
+        1, file_storage_config_.scanmeta_iterator_keys_limit));
+
+    for (int i = 0; i < hash_bucket_count_; ++i) {
+        std::string bucket_dir = fmt::format("{}/{:02x}", root_dir_, i);
+        auto file_infos = fs_adapter_->ListFilesWithInfo(bucket_dir);
+        if (!file_infos) {
+            if (file_infos.error() == ErrorCode::FILE_NOT_FOUND) {
+                continue;
+            }
+            LOG(ERROR) << "Failed to list files in bucket " << bucket_dir
+                       << ": " << static_cast<int>(file_infos.error());
+            return tl::make_unexpected(file_infos.error());
+        }
+
+        for (const auto& info : *file_infos) {
+            std::string key = UnescapeFilename(info.name);
+            batch_keys.push_back(key);
+            StorageObjectMetadata meta{-1, 0, static_cast<int64_t>(key.size()),
+                                       static_cast<int64_t>(info.size), ""};
+            batch_metas.push_back(meta);
+
+            if (batch_keys.size() >= batch_limit) {
+                auto err = handler(batch_keys, batch_metas);
+                if (err != ErrorCode::OK) return tl::make_unexpected(err);
+                batch_keys.clear();
+                batch_metas.clear();
+            }
+        }
+    }
+    if (!batch_keys.empty()) {
+        auto err = handler(batch_keys, batch_metas);
+        if (err != ErrorCode::OK) return tl::make_unexpected(err);
+    }
+    return {};
+}
+
+// === Key -> Path mapping ===
+
+std::string DistributedStorageBackend::GetObjectPath(
+    const std::string& key) const {
+    uint64_t hash = XXH64(key.data(), key.size(), 0);
+    std::string bucket = fmt::format("{:02x}", hash % hash_bucket_count_);
+    std::string safe_key = EscapeFilename(key);
+    return (std::filesystem::path(root_dir_) / bucket / safe_key).string();
+}
+
+std::string DistributedStorageBackend::EscapeFilename(const std::string& key) {
+    std::string result;
+    result.reserve(key.size() + 16);
+    for (unsigned char c : key) {
+        if (c == '@' || c == ':' || c == '/' || c == '\\' || c == '%' ||
+            c < 0x20 || c > 0x7e) {
+            result += fmt::format("%{:02x}", static_cast<unsigned>(c));
+        } else {
+            result += static_cast<char>(c);
+        }
+    }
+    return result;
+}
+
+std::string DistributedStorageBackend::UnescapeFilename(
+    const std::string& name) {
+    auto is_hex = [](char c) {
+        return (c >= '0' && c <= '9') || (c >= 'a' && c <= 'f') ||
+               (c >= 'A' && c <= 'F');
+    };
+    std::string result;
+    result.reserve(name.size());
+    for (size_t i = 0; i < name.size(); ++i) {
+        if (name[i] == '%' && i + 2 < name.size() && is_hex(name[i + 1]) &&
+            is_hex(name[i + 2])) {
+            char hex[3] = {name[i + 1], name[i + 2], 0};
+            unsigned long val = strtoul(hex, nullptr, 16);
+            result += static_cast<char>(val);
+            i += 2;
+        } else {
+            result += name[i];
+        }
+    }
+    return result;
+}
+
+}  // namespace mooncake
diff --git a/mooncake-store/src/storage/distributed/hf3fs_adapter.cpp b/mooncake-store/src/storage/distributed/hf3fs_adapter.cpp
new file mode 100644
index 00000000..a4ccf3b5
--- /dev/null
+++ b/mooncake-store/src/storage/distributed/hf3fs_adapter.cpp
@@ -0,0 +1,382 @@
+#include "storage/distributed/hf3fs_adapter.h"
+
+#include <dirent.h>
+#include <fcntl.h>
+#include <unistd.h>
+
+#include <cerrno>
+#include <cstring>
+#include <filesystem>
+
+#include "hf3fs/hf3fs.h"
+
+namespace mooncake {
+
+// Destructor must be defined in .cpp where USRBIOResourceManager is complete.
+Hf3fsAdapter::~Hf3fsAdapter() = default;
+
+Hf3fsAdapter::Hf3fsAdapter() = default;
+
+tl::expected<void, ErrorCode> Hf3fsAdapter::Init(
+    const std::string& mount_path) {
+    resource_manager_ = std::make_unique<USRBIOResourceManager>();
+    Hf3fsConfig config{};
+    config.mount_root = mount_path;
+    resource_manager_->setDefaultParams(config);
+    return {};
+}
+
+tl::expected<void, ErrorCode> Hf3fsAdapter::Shutdown() {
+    resource_manager_.reset();
+    return {};
+}
+
+tl::expected<size_t, ErrorCode> Hf3fsAdapter::WriteFile(
+    const std::string& path, std::span<const char> data) {
+    int fd = open(path.c_str(), O_WRONLY | O_CREAT | O_TRUNC | O_CLOEXEC, 0644);
+    if (fd < 0) return tl::make_unexpected(ErrorCode::FILE_OPEN_FAIL);
+
+    if (hf3fs_reg_fd(fd, 0) > 0) {
+        close(fd);
+        ::unlink(path.c_str());
+        return tl::make_unexpected(ErrorCode::FILE_OPEN_FAIL);
+    }
+
+    auto* resource = resource_manager_->getThreadResource();
+    if (!resource || !resource->initialized) {
+        hf3fs_dereg_fd(fd);
+        close(fd);
+        ::unlink(path.c_str());
+        return tl::make_unexpected(ErrorCode::FILE_OPEN_FAIL);
+    }
+
+    auto& threefs_iov = resource->iov_;
+    auto& ior_write = resource->ior_write_;
+    const char* data_ptr = data.data();
+    size_t length = data.size();
+    size_t total_written = 0;
+    off_t offset = 0;
+
+    while (total_written < length) {
+        size_t chunk =
+            std::min(length - total_written, resource->config_.iov_size);
+        memcpy(threefs_iov.base, data_ptr + total_written, chunk);
+
+        int ret = hf3fs_prep_io(&ior_write, &threefs_iov, false,
+                                threefs_iov.base, fd, offset, chunk, nullptr);
+        if (ret < 0) break;
+
+        ret = hf3fs_submit_ios(&ior_write);
+        if (ret < 0) break;
+
+        struct hf3fs_cqe cqe;
+        ret = hf3fs_wait_for_ios(&ior_write, &cqe, 1, 1, nullptr);
+        if (ret < 0 || cqe.result < 0) break;
+
+        size_t bytes_written = cqe.result;
+        total_written += bytes_written;
+        offset += bytes_written;
+        if (bytes_written < chunk) break;
+    }
+
+    hf3fs_dereg_fd(fd);
+    close(fd);
+
+    if (total_written != length) {
+        auto unlink_ret = ::unlink(path.c_str());
+        if (unlink_ret != 0) {
+            LOG(WARNING) << "Failed to clean up partial write: " << path;
+        }
+        return tl::make_unexpected(ErrorCode::FILE_WRITE_FAIL);
+    }
+    return total_written;
+}
+
+tl::expected<size_t, ErrorCode> Hf3fsAdapter::ReadFile(const std::string& path,
+                                                       void* buf, size_t len) {
+    if (!buf && len > 0) {
+        return tl::make_unexpected(ErrorCode::INVALID_PARAMS);
+    }
+
+    int fd = open(path.c_str(), O_RDONLY | O_CLOEXEC);
+    if (fd < 0) {
+        if (errno == ENOENT) {
+            return tl::make_unexpected(ErrorCode::FILE_NOT_FOUND);
+        }
+        return tl::make_unexpected(ErrorCode::FILE_OPEN_FAIL);
+    }
+
+    if (hf3fs_reg_fd(fd, 0) > 0) {
+        close(fd);
+        return tl::make_unexpected(ErrorCode::FILE_OPEN_FAIL);
+    }
+
+    auto* resource = resource_manager_->getThreadResource();
+    if (!resource || !resource->initialized) {
+        hf3fs_dereg_fd(fd);
+        close(fd);
+        return tl::make_unexpected(ErrorCode::FILE_OPEN_FAIL);
+    }
+
+    auto& threefs_iov = resource->iov_;
+    auto& ior_read = resource->ior_read_;
+    char* dest = static_cast<char*>(buf);
+    size_t total_read = 0;
+    off_t offset = 0;
+
+    while (total_read < len) {
+        size_t chunk = std::min(len - total_read, resource->config_.iov_size);
+
+        int ret = hf3fs_prep_io(&ior_read, &threefs_iov, true, threefs_iov.base,
+                                fd, offset, chunk, nullptr);
+        if (ret < 0) break;
+
+        ret = hf3fs_submit_ios(&ior_read);
+        if (ret < 0) break;
+
+        struct hf3fs_cqe cqe;
+        ret = hf3fs_wait_for_ios(&ior_read, &cqe, 1, 1, nullptr);
+        if (ret < 0 || cqe.result < 0) break;
+
+        size_t bytes_read = cqe.result;
+        if (bytes_read == 0) break;
+
+        memcpy(dest + total_read, threefs_iov.base, bytes_read);
+        total_read += bytes_read;
+        offset += bytes_read;
+        if (bytes_read < chunk) break;
+    }
+
+    hf3fs_dereg_fd(fd);
+    close(fd);
+
+    if (total_read != len) {
+        return tl::make_unexpected(ErrorCode::FILE_READ_FAIL);
+    }
+    return total_read;
+}
+
+tl::expected<size_t, ErrorCode> Hf3fsAdapter::VectorWriteFile(
+    const std::string& path, const iovec* iov, int iovcnt, off_t offset) {
+    for (int i = 0; i < iovcnt; ++i) {
+        if (!iov[i].iov_base && iov[i].iov_len > 0) {
+            return tl::make_unexpected(ErrorCode::INVALID_PARAMS);
+        }
+    }
+
+    int fd = open(path.c_str(), O_WRONLY | O_CREAT | O_TRUNC | O_CLOEXEC, 0644);
+    if (fd < 0) return tl::make_unexpected(ErrorCode::FILE_OPEN_FAIL);
+
+    if (hf3fs_reg_fd(fd, 0) > 0) {
+        close(fd);
+        ::unlink(path.c_str());
+        return tl::make_unexpected(ErrorCode::FILE_OPEN_FAIL);
+    }
+
+    auto* resource = resource_manager_->getThreadResource();
+    if (!resource || !resource->initialized) {
+        hf3fs_dereg_fd(fd);
+        close(fd);
+        ::unlink(path.c_str());
+        return tl::make_unexpected(ErrorCode::FILE_OPEN_FAIL);
+    }
+
+    auto& threefs_iov = resource->iov_;
+    auto& ior_write = resource->ior_write_;
+
+    size_t total_length = 0;
+    for (int i = 0; i < iovcnt; ++i) total_length += iov[i].iov_len;
+
+    size_t total_written = 0;
+    off_t current_offset = offset;
+    size_t remaining = total_length;
+    int iov_idx = 0;
+    size_t iov_off = 0;
+
+    while (remaining > 0) {
+        size_t chunk = std::min(remaining, resource->config_.iov_size);
+
+        // Copy from iovec to shared buffer
+        size_t copied = 0;
+        char* dest = reinterpret_cast<char*>(threefs_iov.base);
+        while (copied < chunk && iov_idx < iovcnt) {
+            size_t n = std::min(chunk - copied, iov[iov_idx].iov_len - iov_off);
+            memcpy(dest + copied,
+                   static_cast<char*>(iov[iov_idx].iov_base) + iov_off, n);
+            copied += n;
+            iov_off += n;
+            if (iov_off >= iov[iov_idx].iov_len) {
+                iov_idx++;
+                iov_off = 0;
+            }
+        }
+
+        int ret =
+            hf3fs_prep_io(&ior_write, &threefs_iov, false, threefs_iov.base, fd,
+                          current_offset, chunk, nullptr);
+        if (ret < 0) break;
+        ret = hf3fs_submit_ios(&ior_write);
+        if (ret < 0) break;
+        struct hf3fs_cqe cqe;
+        ret = hf3fs_wait_for_ios(&ior_write, &cqe, 1, 1, nullptr);
+        if (ret < 0 || cqe.result < 0) break;
+
+        size_t bytes_written = cqe.result;
+        total_written += bytes_written;
+        current_offset += bytes_written;
+        remaining -= bytes_written;
+        if (bytes_written < chunk) break;
+    }
+
+    hf3fs_dereg_fd(fd);
+    close(fd);
+
+    if (total_written != total_length) {
+        auto unlink_ret = ::unlink(path.c_str());
+        if (unlink_ret != 0) {
+            LOG(WARNING) << "Failed to clean up partial write: " << path;
+        }
+        return tl::make_unexpected(ErrorCode::FILE_WRITE_FAIL);
+    }
+    return total_written;
+}
+
+tl::expected<size_t, ErrorCode> Hf3fsAdapter::VectorReadFile(
+    const std::string& path, const iovec* iov, int iovcnt, off_t offset) {
+    for (int i = 0; i < iovcnt; ++i) {
+        if (!iov[i].iov_base && iov[i].iov_len > 0) {
+            return tl::make_unexpected(ErrorCode::INVALID_PARAMS);
+        }
+    }
+
+    int fd = open(path.c_str(), O_RDONLY | O_CLOEXEC);
+    if (fd < 0) {
+        if (errno == ENOENT) {
+            return tl::make_unexpected(ErrorCode::FILE_NOT_FOUND);
+        }
+        return tl::make_unexpected(ErrorCode::FILE_OPEN_FAIL);
+    }
+
+    if (hf3fs_reg_fd(fd, 0) > 0) {
+        close(fd);
+        return tl::make_unexpected(ErrorCode::FILE_OPEN_FAIL);
+    }
+
+    auto* resource = resource_manager_->getThreadResource();
+    if (!resource || !resource->initialized) {
+        hf3fs_dereg_fd(fd);
+        close(fd);
+        return tl::make_unexpected(ErrorCode::FILE_OPEN_FAIL);
+    }
+
+    auto& threefs_iov = resource->iov_;
+    auto& ior_read = resource->ior_read_;
+
+    size_t total_length = 0;
+    for (int i = 0; i < iovcnt; ++i) total_length += iov[i].iov_len;
+
+    size_t total_read = 0;
+    off_t current_offset = offset;
+    size_t remaining = total_length;
+    int iov_idx = 0;
+    size_t iov_off = 0;
+
+    while (remaining > 0) {
+        size_t chunk = std::min(remaining, resource->config_.iov_size);
+
+        int ret = hf3fs_prep_io(&ior_read, &threefs_iov, true, threefs_iov.base,
+                                fd, current_offset, chunk, nullptr);
+        if (ret < 0) break;
+        ret = hf3fs_submit_ios(&ior_read);
+        if (ret < 0) break;
+        struct hf3fs_cqe cqe;
+        ret = hf3fs_wait_for_ios(&ior_read, &cqe, 1, 1, nullptr);
+        if (ret < 0 || cqe.result < 0) break;
+
+        size_t bytes_read = cqe.result;
+        if (bytes_read == 0) break;
+
+        // Copy from shared buffer to iovec
+        size_t to_copy = bytes_read;
+        char* src = reinterpret_cast<char*>(threefs_iov.base);
+        while (to_copy > 0 && iov_idx < iovcnt) {
+            size_t n = std::min(to_copy, iov[iov_idx].iov_len - iov_off);
+            memcpy(static_cast<char*>(iov[iov_idx].iov_base) + iov_off, src, n);
+            src += n;
+            to_copy -= n;
+            total_read += n;
+            remaining -= n;
+            current_offset += n;
+            iov_off += n;
+            if (iov_off >= iov[iov_idx].iov_len) {
+                iov_idx++;
+                iov_off = 0;
+            }
+        }
+        if (bytes_read < chunk) break;
+    }
+
+    hf3fs_dereg_fd(fd);
+    close(fd);
+
+    if (total_read != total_length) {
+        return tl::make_unexpected(ErrorCode::FILE_READ_FAIL);
+    }
+    return total_read;
+}
+
+tl::expected<void, ErrorCode> Hf3fsAdapter::DeleteFile(
+    const std::string& path) {
+    if (::unlink(path.c_str()) != 0) {
+        if (errno == ENOENT) {
+            return tl::make_unexpected(ErrorCode::FILE_NOT_FOUND);
+        }
+        return tl::make_unexpected(ErrorCode::FILE_WRITE_FAIL);
+    }
+    return {};
+}
+
+// Note: FileExists and DeleteFile use POSIX access()/unlink() via the
+// kernel VFS mount. This assumes 3FS's FUSE mount namespace is consistent
+// with the USRBIO namespace used for read/write I/O.
+
+tl::expected<bool, ErrorCode> Hf3fsAdapter::FileExists(
+    const std::string& path) {
+    if (::access(path.c_str(), F_OK) == 0) {
+        return true;
+    }
+    if (errno == ENOENT) {
+        return false;
+    }
+    return tl::make_unexpected(ErrorCode::FILE_READ_FAIL);
+}
+
+tl::expected<std::vector<std::string>, ErrorCode> Hf3fsAdapter::ListFiles(
+    const std::string& dir) {
+    DIR* d = opendir(dir.c_str());
+    if (!d) {
+        if (errno == ENOENT) {
+            return tl::make_unexpected(ErrorCode::FILE_NOT_FOUND);
+        }
+        return tl::make_unexpected(ErrorCode::FILE_READ_FAIL);
+    }
+    std::vector<std::string> result;
+
+    struct dirent* entry;
+    while ((entry = readdir(d)) != nullptr) {
+        std::string name = entry->d_name;
+        if (name == "." || name == "..") continue;
+        if (entry->d_type == DT_DIR) continue;
+        if (entry->d_type == DT_UNKNOWN) {
+            struct stat st;
+            if (::stat((dir + "/" + name).c_str(), &st) == 0 &&
+                S_ISDIR(st.st_mode))
+                continue;
+        }
+        result.push_back(name);
+    }
+    closedir(d);
+    return result;
+}
+
+}  // namespace mooncake
diff --git a/mooncake-store/src/storage_backend.cpp b/mooncake-store/src/storage_backend.cpp
index 24540ab6..42f29daa 100644
--- a/mooncake-store/src/storage_backend.cpp
+++ b/mooncake-store/src/storage_backend.cpp
@@ -5,6 +5,7 @@
 #include <sys/stat.h>
 #include <sys/uio.h>
 #include <errno.h>
+#include <cstdint>
 #include <cstring>
 
 #include <regex>
@@ -21,6 +22,7 @@
 #include "utils.h"
 
 #include <ylt/util/tl/expected.hpp>
+#include "storage/distributed/distributed_storage_backend.h"
 
 namespace mooncake {
 
@@ -72,7 +74,7 @@ BucketBackendConfig BucketBackendConfig::FromEnvironment() {
 
     const auto policy_str = GetEnvStringOr(
         "MOONCAKE_OFFLOAD_BUCKET_EVICTION_POLICY",
-        GetEnvStringOr("MOONCAKE_BUCKET_EVICTION_POLICY", "none"));
+        GetEnvStringOr("MOONCAKE_BUCKET_EVICTION_POLICY", "fifo"));
     if (policy_str == "fifo") {
         config.eviction_policy = BucketEvictionPolicy::FIFO;
     } else if (policy_str == "lru") {
@@ -104,23 +106,10 @@ void StorageBackend::RecalculateAvailableSpace() {
     }
 }
 
-bool StorageBackend::IsEvictionEnabled() const {
-    // First check user configuration
-    if (!enable_eviction_) {
-        return false;
-    }
-
-#ifdef USE_3FS
-    // Eviction is only enabled for local storage, not for 3FS
-    return !is_3fs_dir_;
-#else
-    // If 3FS is not compiled in, eviction is enabled if user config allows
-    return true;
-#endif
-}
+bool StorageBackend::IsEvictionEnabled() const { return enable_eviction_; }
 
 tl::expected<void, ErrorCode> StorageBackend::Init(uint64_t quota_bytes = 0) {
-    // Skip eviction initialization for 3FS mode
+    // Skip eviction initialization if disabled
     if (!IsEvictionEnabled()) {
         initialized_.store(true, std::memory_order_release);
         return {};
@@ -208,7 +197,7 @@ tl::expected<void, ErrorCode> StorageBackend::Init(uint64_t quota_bytes = 0) {
         if (total_space_ >= used_space_) {
             RecalculateAvailableSpace();
         } else {
-            // Only enable eviction for local storage, not for 3FS
+            // Only enable eviction for local storage
             if (IsEvictionEnabled()) {
                 eviction_needed = true;
                 available_space_ = -1;
@@ -217,10 +206,11 @@ tl::expected<void, ErrorCode> StorageBackend::Init(uint64_t quota_bytes = 0) {
                     << ") exceeds the new quota (" << total_space_
                     << "). Eviction will be triggered after initial setup.";
             } else {
-                // For 3FS mode, just log a warning but don't trigger eviction
+                // Eviction disabled, just log a warning but don't trigger
+                // eviction
                 LOG(WARNING) << "Existing used space (" << used_space_
                              << ") exceeds the new quota (" << total_space_
-                             << "). Eviction is disabled for 3FS mode.";
+                             << "). Eviction is disabled.";
                 RecalculateAvailableSpace();  // Still calculate available space
             }
         }
@@ -496,7 +486,7 @@ void StorageBackend::RemoveFile(const std::string& path) {
     std::this_thread::sleep_for(
         std::chrono::microseconds(50));  // sleep for 50 us
 
-    // For 3FS mode, use original logic (no queue tracking)
+    // Eviction disabled, use simple delete (no queue tracking)
     if (!IsEvictionEnabled()) {
         if (fs::exists(path)) {
             std::error_code ec;
@@ -552,7 +542,7 @@ void StorageBackend::RemoveByRegex(const std::string& regex_pattern) {
         return;
     }
 
-    // For 3FS mode, use original logic (no queue tracking)
+    // Eviction disabled, use simple delete (no queue tracking)
     if (!IsEvictionEnabled()) {
         fs::path storage_root = fs::path(root_dir_) / fsdir_;
         if (!fs::exists(storage_root) || !fs::is_directory(storage_root)) {
@@ -625,7 +615,7 @@ void StorageBackend::RemoveByRegex(const std::string& regex_pattern) {
 void StorageBackend::RemoveAll() {
     namespace fs = std::filesystem;
 
-    // For 3FS mode, use original logic (no queue tracking)
+    // Eviction disabled, use simple delete (no queue tracking)
     if (!IsEvictionEnabled()) {
         // Iterate through the root directory and remove all files
         for (const auto& entry : fs::directory_iterator(root_dir_)) {
@@ -788,18 +778,6 @@ std::unique_ptr<StorageFile> StorageBackend::create_file(
         return nullptr;
     }
 
-#ifdef USE_3FS
-    if (is_3fs_dir_) {
-        if (hf3fs_reg_fd(fd, 0) > 0) {
-            close(fd);
-            return nullptr;
-        }
-        return resource_manager_ ? std::make_unique<ThreeFSFile>(
-                                       path, fd, resource_manager_.get())
-                                 : nullptr;
-    }
-#endif
-
 #ifdef USE_URING
     if (use_uring_) {
         // use_direct_io mirrors the O_DIRECT flag: true for reads, false for
@@ -823,6 +801,28 @@ bool StorageBackend::CheckDiskSpace(size_t required_size) {
 
     bool has_enough_space = available_space_ >= required_size;
 
+    if (has_enough_space) {
+        // Also check actual disk space to handle multiple instances
+        // sharing the same filesystem with independent quotas.
+        namespace fs = std::filesystem;
+        fs::path storage_root = fs::path(root_dir_) / GetActualFsdir();
+        std::error_code ec;
+        auto space_info = fs::space(storage_root, ec);
+        if (!ec) {
+            uint64_t actual_available = space_info.available;
+            constexpr uint64_t kMinFreeSpace = 256 * kMB;
+            if (actual_available < required_size + kMinFreeSpace) {
+                VLOG(1) << "Actual disk space low: available="
+                        << actual_available << ", required=" << required_size
+                        << ". Triggering eviction.";
+                has_enough_space = false;
+            }
+        } else {
+            LOG(WARNING) << "Failed to get disk space info for " << storage_root
+                         << ": " << ec.message();
+        }
+    }
+
     if (has_enough_space) {
         used_space_ += required_size;
         available_space_ -= required_size;
@@ -836,8 +836,7 @@ bool StorageBackend::CheckDiskSpace(size_t required_size) {
 FileRecord StorageBackend::EvictFile() {
     // Eviction is only enabled for local storage
     if (!IsEvictionEnabled()) {
-        LOG(WARNING)
-            << "Eviction is disabled for 3FS mode. Cannot evict files.";
+        LOG(WARNING) << "Eviction is disabled. Cannot evict files.";
         return {};
     }
 
@@ -909,8 +908,7 @@ FileRecord StorageBackend::SelectFileToEvictByFIFO() {
 tl::expected<std::vector<std::string>, ErrorCode>
 StorageBackend::EnsureDiskSpace(size_t required_size) {
     std::vector<std::string> evicted_keys;
-    // If eviction is disabled (3FS mode), skip space checking and eviction
-    // Let 3FS filesystem handle space management itself
+    // If eviction is disabled, skip space checking and eviction
     if (!IsEvictionEnabled()) {
         return evicted_keys;
     }
@@ -1155,7 +1153,7 @@ tl::expected<void, ErrorCode> StorageBackendAdaptor::ScanMeta(
     fs::path root = fs::path(file_storage_config_.storage_filepath) /
                     file_per_key_config_.fsdir;
     if (!fs::exists(root)) {
-        meta_scanned_.store(true, std::memory_order_acquire);
+        meta_scanned_.store(true, std::memory_order_release);
         return {};
     }
 
@@ -1228,7 +1226,7 @@ tl::expected<void, ErrorCode> StorageBackendAdaptor::ScanMeta(
         }
     }
 
-    meta_scanned_.store(true, std::memory_order_acquire);
+    meta_scanned_.store(true, std::memory_order_release);
     return {};
 }
 
@@ -1579,7 +1577,16 @@ tl::expected<void, ErrorCode> BucketStorageBackend::Init() {
             if (entry.is_regular_file() &&
                 entry.path().extension() == BUCKET_METADATA_FILE_SUFFIX) {
                 const auto& bucket_id_str = entry.path().stem();
-                int64_t bucket_id = std::stoll(bucket_id_str);
+                int64_t bucket_id = 0;
+                try {
+                    bucket_id = std::stoll(bucket_id_str);
+                } catch (const std::exception& e) {
+                    LOG(WARNING)
+                        << "Skipping metadata file with a non-numeric "
+                           "bucket id: "
+                        << entry.path().string() << " (" << e.what() << ")";
+                    continue;
+                }
                 auto [metadata_it, success] = buckets_.try_emplace(
                     bucket_id, std::make_shared<BucketMetadata>());
                 if (success) lru_index_.emplace(0LL, bucket_id);
@@ -1690,7 +1697,16 @@ tl::expected<void, ErrorCode> BucketStorageBackend::Init() {
             // Extract bucket ID from filename (e.g., "12345.bucket" ->
             // "12345")
             auto bucket_id_str = entry.path().stem();
-            int64_t bucket_id = std::stoll(bucket_id_str);
+            int64_t bucket_id = 0;
+            try {
+                bucket_id = std::stoll(bucket_id_str);
+            } catch (const std::exception& e) {
+                LOG(WARNING)
+                    << "Skipping orphan-scan file with a non-numeric "
+                       "bucket id: "
+                    << entry.path().string() << " (" << e.what() << ")";
+                continue;
+            }
 
             // Check if this bucket has valid metadata
             if (valid_bucket_ids.find(bucket_id) != valid_bucket_ids.end()) {
@@ -1764,6 +1780,13 @@ tl::expected<bool, ErrorCode> BucketStorageBackend::IsExist(
 }
 
 tl::expected<bool, ErrorCode> BucketStorageBackend::IsEnableOffloading() {
+    // When eviction is enabled, always allow offloading since PrepareEviction
+    // will manage capacity by evicting old buckets as needed.
+    if (bucket_backend_config_.eviction_policy != BucketEvictionPolicy::NONE &&
+        bucket_backend_config_.max_total_size > 0) {
+        return true;
+    }
+
     auto store_metadata_result = GetStoreMetadata();
     if (!store_metadata_result) {
         LOG(ERROR) << "Failed to get store metadata: "
@@ -2209,15 +2232,61 @@ BucketStorageBackend::PendingEviction BucketStorageBackend::PrepareEviction(
 
     SharedMutexLocker lock(&mutex_);
 
-    if (!buckets_.empty() &&
-        total_size_ + required_size > bucket_backend_config_.max_total_size) {
-        LOG(INFO) << "[Evict] triggered: total=" << total_size_ << "/"
-                  << bucket_backend_config_.max_total_size
-                  << " required=" << required_size;
+    // Check actual disk space once before the loop. PrepareEviction only
+    // removes metadata -- files are deleted later in FinalizeEviction -- so
+    // re-checking disk space inside the loop would yield the same result
+    // every iteration and is unnecessary.
+    //
+    // When disk is full, we calculate the space deficit and accumulate
+    // the estimated freed space (data_size + meta_size) per evicted
+    // bucket. Due to block alignment, actual disk usage >= data_size +
+    // meta_size, so this is a safe lower bound -- we will not under-evict.
+    bool initial_disk_full = false;
+    uint64_t deficit = 0;
+    {
+        namespace fs = std::filesystem;
+        std::error_code ec;
+        auto space_info = fs::space(storage_path_, ec);
+        if (!ec) {
+            uint64_t actual_available = space_info.available;
+            constexpr uint64_t kMinFreeSpace = 256 * kMB;
+            uint64_t req_sz =
+                required_size > 0 ? static_cast<uint64_t>(required_size) : 0;
+            initial_disk_full = actual_available < req_sz + kMinFreeSpace;
+            if (initial_disk_full) {
+                deficit = req_sz + kMinFreeSpace - actual_available;
+                LOG(WARNING)
+                    << "[Evict] Actual disk space too low: available="
+                    << actual_available << ", required=" << required_size
+                    << ", deficit=" << deficit
+                    << ". Will evict buckets to free space.";
+            }
+        } else {
+            LOG(WARNING) << "[Evict] Failed to get disk space info for "
+                         << storage_path_ << ": " << ec.message();
+        }
     }
 
-    while (!buckets_.empty() && total_size_ + required_size >
-                                    bucket_backend_config_.max_total_size) {
+    size_t evict_count = 0;
+    constexpr size_t kMaxEvictionBuckets = 1000;
+    uint64_t accumulated_freed_space = 0;
+
+    while (!buckets_.empty() && evict_count < kMaxEvictionBuckets) {
+        bool quota_exceeded =
+            total_size_ + required_size > bucket_backend_config_.max_total_size;
+
+        bool disk_still_full =
+            initial_disk_full && (accumulated_freed_space < deficit);
+
+        if (!quota_exceeded && !disk_still_full) break;
+
+        if (evict_count == 0) {
+            LOG(INFO) << "[Evict] triggered: total=" << total_size_ << "/"
+                      << bucket_backend_config_.max_total_size
+                      << " required=" << required_size
+                      << " disk_full=" << initial_disk_full;
+        }
+
         auto evict_it = SelectEvictionCandidate();
         if (evict_it == buckets_.end()) break;
 
@@ -2242,7 +2311,11 @@ BucketStorageBackend::PendingEviction BucketStorageBackend::PrepareEviction(
         for (const auto& key : evict_meta->keys) {
             result.keys.push_back(key);
         }
+        accumulated_freed_space +=
+            static_cast<uint64_t>(evict_meta->data_size) +
+            static_cast<uint64_t>(evict_meta->meta_size);
         result.buckets.emplace_back(evict_id, std::move(evict_meta));
+        evict_count++;
     }
 
     if (!result.buckets.empty()) {
@@ -2804,11 +2877,25 @@ tl::expected<int64_t, ErrorCode> OffsetAllocatorStorageBackend::BatchOffload(
             continue;  // Simulate allocation/write failure
         }
 
-        // Calculate total value size
-        uint32_t value_size = 0;
+        // Calculate total value size. RecordHeader stores value_len as a
+        // uint32_t (RecordHeader::SIZE is 8 bytes), so accumulating directly
+        // into a uint32_t would silently overflow for an object larger than
+        // 4 GiB: the record would be under-allocated and then its full slices
+        // written past the allocation. Sum in 64 bits and reject oversized
+        // objects instead of corrupting the storage arena.
+        uint64_t total_value_size = 0;
         for (const auto& slice : slices) {
-            value_size += static_cast<uint32_t>(slice.size);
+            total_value_size += slice.size;
         }
+        if (total_value_size > UINT32_MAX) {
+            LOG(ERROR)
+                << "Object too large for SSD offload (value_len must fit "
+                   "in 4 GiB) for key: "
+                << key << ", size: " << total_value_size
+                << " - skipping this key";
+            continue;  // partial-success model: keep processing other keys
+        }
+        uint32_t value_size = static_cast<uint32_t>(total_value_size);
 
         // Prepare record header
         RecordHeader header{.key_len = static_cast<uint32_t>(key.size()),
@@ -3214,9 +3301,30 @@ CreateStorageBackend(const FileStorageConfig& config) {
         case StorageBackendType::kOffsetAllocator: {
             return std::make_shared<OffsetAllocatorStorageBackend>(config);
         }
+        case StorageBackendType::kDistributed: {
+            auto distributed_config =
+                DistributedStorageConfig::FromEnvironment();
+            if (!distributed_config.Validate()) {
+                throw std::invalid_argument(
+                    "Invalid DistributedStorage configuration");
+            }
+            std::unique_ptr<FileSystemAdapter> adapter;
+            if (distributed_config.fs_adapter_type == "hf3fs") {
+#ifdef USE_3FS
+                adapter = std::make_unique<Hf3fsAdapter>();
+#else
+                return tl::make_unexpected(ErrorCode::INVALID_PARAMS);
+#endif
+            } else {
+                return tl::make_unexpected(ErrorCode::INVALID_PARAMS);
+            }
+            return std::make_shared<DistributedStorageBackend>(
+                config, distributed_config, std::move(adapter));
+        }
         default: {
-            LOG(FATAL) << "Unsupported backend type";
-            return tl::make_unexpected(ErrorCode::INTERNAL_ERROR);
+            LOG(ERROR) << "Unsupported backend type: "
+                       << static_cast<int>(config.storage_backend_type);
+            return tl::make_unexpected(ErrorCode::INVALID_PARAMS);
         }
     }
 }
diff --git a/mooncake-store/src/tenant_quota.cpp b/mooncake-store/src/tenant_quota.cpp
new file mode 100644
index 00000000..da567c48
--- /dev/null
+++ b/mooncake-store/src/tenant_quota.cpp
@@ -0,0 +1,324 @@
+#include "tenant_quota.h"
+
+#include <algorithm>
+#include <limits>
+
+#include "types.h"
+
+#include <glog/logging.h>
+
+namespace mooncake {
+namespace {
+
+struct RemainderShare {
+    std::string tenant_id;
+    uint64_t base = 0;
+    unsigned __int128 remainder = 0;
+};
+
+uint64_t SaturatingAdd(uint64_t lhs, uint64_t rhs) {
+    if (lhs > std::numeric_limits<uint64_t>::max() - rhs) {
+        return std::numeric_limits<uint64_t>::max();
+    }
+    return lhs + rhs;
+}
+
+bool IsLazyEmptyTenant(const TenantQuotaState& state) {
+    return !state.has_explicit_policy && state.used_bytes == 0 &&
+           state.reserved_bytes == 0 && state.committed_count == 0 &&
+           state.metadata_object_count == 0;
+}
+
+TenantQuotaResult AccountingMismatch(const char* operation,
+                                     const std::string& tenant_id,
+                                     uint64_t requested, uint64_t available) {
+    LOG(WARNING) << operation << " accounting mismatch for tenant " << tenant_id
+                 << ": requested=" << requested << ", available=" << available;
+    return tl::make_unexpected(TenantQuotaError::kAccountingMismatch);
+}
+
+}  // namespace
+
+std::vector<TenantQuotaAssignment> BuildEffectiveQuotaAssignments(
+    const std::map<std::string, TenantQuotaState>& tenants,
+    uint64_t allocatable_capacity_bytes) {
+    unsigned __int128 explicit_requested_sum = 0;
+    std::vector<std::string> explicit_tenants;
+
+    for (const auto& [tenant_id, state] : tenants) {
+        if (state.has_explicit_policy) {
+            explicit_tenants.push_back(tenant_id);
+            explicit_requested_sum += state.requested_quota_bytes;
+        }
+    }
+
+    std::map<std::string, uint64_t> assigned;
+    for (const auto& [tenant_id, _] : tenants) {
+        assigned[tenant_id] = 0;
+    }
+
+    auto distribute = [&](const std::vector<std::string>& tenant_ids,
+                          uint64_t capacity, unsigned __int128 denominator,
+                          bool proportional_to_requested) {
+        if (tenant_ids.empty() || capacity == 0 || denominator == 0) {
+            return;
+        }
+
+        std::vector<RemainderShare> shares;
+        shares.reserve(tenant_ids.size());
+        uint64_t base_assigned = 0;
+        for (const auto& tenant_id : tenant_ids) {
+            const auto& state = tenants.at(tenant_id);
+            unsigned __int128 numerator =
+                proportional_to_requested ? state.requested_quota_bytes : 1;
+            unsigned __int128 product =
+                static_cast<unsigned __int128>(capacity) * numerator;
+            uint64_t base = static_cast<uint64_t>(product / denominator);
+            unsigned __int128 remainder = product % denominator;
+            shares.push_back({tenant_id, base, remainder});
+            base_assigned += base;
+        }
+
+        std::sort(shares.begin(), shares.end(),
+                  [](const RemainderShare& lhs, const RemainderShare& rhs) {
+                      if (lhs.remainder != rhs.remainder) {
+                          return lhs.remainder > rhs.remainder;
+                      }
+                      return lhs.tenant_id < rhs.tenant_id;
+                  });
+
+        uint64_t remaining = capacity - base_assigned;
+        for (auto& share : shares) {
+            if (remaining > 0) {
+                ++share.base;
+                --remaining;
+            }
+            assigned[share.tenant_id] = share.base;
+        }
+    };
+
+    if (explicit_requested_sum <= allocatable_capacity_bytes) {
+        for (const auto& tenant_id : explicit_tenants) {
+            assigned[tenant_id] = tenants.at(tenant_id).requested_quota_bytes;
+        }
+    } else {
+        distribute(explicit_tenants, allocatable_capacity_bytes,
+                   explicit_requested_sum,
+                   /*proportional_to_requested=*/true);
+    }
+
+    std::vector<TenantQuotaAssignment> result;
+    result.reserve(assigned.size());
+    for (const auto& [tenant_id, effective_quota_bytes] : assigned) {
+        result.push_back({.tenant_id = tenant_id,
+                          .effective_quota_bytes = effective_quota_bytes});
+    }
+    return result;
+}
+
+TenantQuotaResult TenantQuotaTable::UpsertTenantPolicy(
+    std::string tenant_id, uint64_t requested_quota_bytes) {
+    if (requested_quota_bytes == 0) {
+        return tl::make_unexpected(TenantQuotaError::kInvalidArgument);
+    }
+
+    auto normalized_tenant_id = NormalizeTenantId(std::move(tenant_id));
+    auto& state = GetOrCreateState(normalized_tenant_id);
+    state.requested_quota_bytes = requested_quota_bytes;
+    state.has_explicit_policy = true;
+    return {};
+}
+
+void TenantQuotaTable::EraseTenantPolicy(std::string tenant_id) {
+    auto normalized_tenant_id = NormalizeTenantId(std::move(tenant_id));
+    auto it = tenants_.find(normalized_tenant_id);
+    if (it == tenants_.end()) {
+        return;
+    }
+    auto& state = it->second;
+    state.requested_quota_bytes = 0;
+    state.effective_quota_bytes = 0;
+    state.has_explicit_policy = false;
+}
+
+void TenantQuotaTable::RecomputeEffectiveQuotas(
+    uint64_t allocatable_capacity_bytes) {
+    for (auto& [tenant_id, state] : tenants_) {
+        state.effective_quota_bytes = 0;
+    }
+
+    for (const auto& assignment :
+         BuildEffectiveQuotaAssignments(tenants_, allocatable_capacity_bytes)) {
+        tenants_.at(assignment.tenant_id).effective_quota_bytes =
+            assignment.effective_quota_bytes;
+    }
+
+    for (auto& [_, state] : tenants_) {
+        RefreshOverQuota(&state);
+    }
+}
+
+std::optional<TenantQuotaSnapshot> TenantQuotaTable::GetTenantSnapshot(
+    std::string tenant_id) const {
+    auto normalized_tenant_id = NormalizeTenantId(std::move(tenant_id));
+    auto it = tenants_.find(normalized_tenant_id);
+    if (it == tenants_.end()) {
+        return std::nullopt;
+    }
+    return MakeSnapshot(it->first, it->second);
+}
+
+std::vector<TenantQuotaSnapshot> TenantQuotaTable::ListTenantSnapshots() const {
+    std::vector<TenantQuotaSnapshot> snapshots;
+    for (const auto& [tenant_id, state] : tenants_) {
+        if (IsLazyEmptyTenant(state)) {
+            continue;
+        }
+        snapshots.push_back(MakeSnapshot(tenant_id, state));
+    }
+    return snapshots;
+}
+
+TenantQuotaResult TenantQuotaTable::Reserve(std::string tenant_id,
+                                            uint64_t bytes) {
+    auto normalized_tenant_id = NormalizeTenantId(std::move(tenant_id));
+    if (bytes == 0) {
+        GetOrCreateState(normalized_tenant_id);
+        return {};
+    }
+
+    auto it = tenants_.find(normalized_tenant_id);
+    if (it == tenants_.end()) {
+        return tl::make_unexpected(TenantQuotaError::kQuotaExceeded);
+    }
+    auto& state = it->second;
+    if (!state.has_explicit_policy) {
+        return tl::make_unexpected(TenantQuotaError::kQuotaExceeded);
+    }
+
+    if (static_cast<unsigned __int128>(state.used_bytes) +
+            state.reserved_bytes + bytes >
+        state.effective_quota_bytes) {
+        return tl::make_unexpected(TenantQuotaError::kQuotaExceeded);
+    }
+
+    state.reserved_bytes += bytes;
+    RefreshOverQuota(&state);
+    return {};
+}
+
+TenantQuotaResult TenantQuotaTable::Commit(std::string tenant_id,
+                                           uint64_t bytes) {
+    auto normalized_tenant_id = NormalizeTenantId(std::move(tenant_id));
+    auto& state = GetOrCreateState(normalized_tenant_id);
+    if (bytes == 0) {
+        return {};
+    }
+
+    if (state.reserved_bytes < bytes) {
+        return AccountingMismatch("Commit", normalized_tenant_id, bytes,
+                                  state.reserved_bytes);
+    }
+
+    state.reserved_bytes -= bytes;
+    state.used_bytes = SaturatingAdd(state.used_bytes, bytes);
+    ++state.committed_count;
+    RefreshOverQuota(&state);
+    return {};
+}
+
+TenantQuotaResult TenantQuotaTable::Abort(std::string tenant_id,
+                                          uint64_t bytes) {
+    auto normalized_tenant_id = NormalizeTenantId(std::move(tenant_id));
+    auto& state = GetOrCreateState(normalized_tenant_id);
+    if (bytes == 0) {
+        return {};
+    }
+
+    if (state.reserved_bytes < bytes) {
+        return AccountingMismatch("Abort", normalized_tenant_id, bytes,
+                                  state.reserved_bytes);
+    }
+
+    state.reserved_bytes -= bytes;
+    RefreshOverQuota(&state);
+    return {};
+}
+
+TenantQuotaResult TenantQuotaTable::Release(std::string tenant_id,
+                                            uint64_t bytes) {
+    auto normalized_tenant_id = NormalizeTenantId(std::move(tenant_id));
+    auto& state = GetOrCreateState(normalized_tenant_id);
+    if (bytes == 0) {
+        return {};
+    }
+
+    if (state.used_bytes < bytes) {
+        return AccountingMismatch("Release", normalized_tenant_id, bytes,
+                                  state.used_bytes);
+    }
+
+    state.used_bytes -= bytes;
+    if (state.committed_count > 0) {
+        --state.committed_count;
+    } else {
+        LOG(WARNING) << "Release found zero committed_count for tenant "
+                     << normalized_tenant_id;
+    }
+    RefreshOverQuota(&state);
+    return {};
+}
+
+TenantQuotaResult TenantQuotaTable::ReleasePartial(std::string tenant_id,
+                                                   uint64_t bytes) {
+    auto normalized_tenant_id = NormalizeTenantId(std::move(tenant_id));
+    auto& state = GetOrCreateState(normalized_tenant_id);
+    if (bytes == 0) {
+        return {};
+    }
+
+    DCHECK_LE(bytes, state.used_bytes);
+    if (state.used_bytes < bytes) {
+        return AccountingMismatch("ReleasePartial", normalized_tenant_id, bytes,
+                                  state.used_bytes);
+    }
+
+    state.used_bytes -= bytes;
+    RefreshOverQuota(&state);
+    return {};
+}
+
+TenantQuotaState& TenantQuotaTable::GetOrCreateState(
+    const std::string& tenant_id) {
+    auto [it, inserted] = tenants_.try_emplace(tenant_id);
+    if (inserted) {
+        it->second.requested_quota_bytes = 0;
+        it->second.effective_quota_bytes = 0;
+    }
+    return it->second;
+}
+
+TenantQuotaSnapshot TenantQuotaTable::MakeSnapshot(
+    const std::string& tenant_id, const TenantQuotaState& state) const {
+    return TenantQuotaSnapshot{
+        .tenant_id = tenant_id,
+        .requested_quota_bytes = state.requested_quota_bytes,
+        .effective_quota_bytes = state.effective_quota_bytes,
+        .used_bytes = state.used_bytes,
+        .reserved_bytes = state.reserved_bytes,
+        .committed_count = state.committed_count,
+        .metadata_object_count = state.metadata_object_count,
+        .has_explicit_policy = state.has_explicit_policy,
+        .over_quota = state.over_quota,
+    };
+}
+
+void TenantQuotaTable::RefreshOverQuota(TenantQuotaState* state) const {
+    state->over_quota =
+        (!state->has_explicit_policy && state->metadata_object_count > 0) ||
+        static_cast<unsigned __int128>(state->used_bytes) +
+                state->reserved_bytes >
+            state->effective_quota_bytes;
+}
+
+}  // namespace mooncake
diff --git a/mooncake-store/src/tenant_quota_policy_store.cpp b/mooncake-store/src/tenant_quota_policy_store.cpp
new file mode 100644
index 00000000..23e10b70
--- /dev/null
+++ b/mooncake-store/src/tenant_quota_policy_store.cpp
@@ -0,0 +1,361 @@
+#include "tenant_quota_policy_store.h"
+
+#include <cerrno>
+#include <chrono>
+#include <cstring>
+#include <filesystem>
+#include <fstream>
+#include <iomanip>
+#include <limits>
+#include <sstream>
+#include <thread>
+
+#include <fcntl.h>
+#include <sys/stat.h>
+#include <sys/types.h>
+#include <unistd.h>
+
+#include <glog/logging.h>
+#include <yaml-cpp/yaml.h>
+
+#include "types.h"
+
+namespace mooncake {
+namespace {
+
+bool IsValidTenantQuotaName(const std::string& name) {
+    return IsValidTenantId(name);
+}
+
+std::string ErrnoMessage(const std::string& action, const std::string& path) {
+    return action + " '" + path + "' failed: " + std::strerror(errno);
+}
+
+tl::expected<void, std::string> WriteAll(int fd, const std::string& content,
+                                         const std::string& path) {
+    const char* data = content.data();
+    size_t remaining = content.size();
+    while (remaining > 0) {
+        ssize_t written = ::write(fd, data, remaining);
+        if (written < 0) {
+            if (errno == EINTR) {
+                continue;
+            }
+            return tl::make_unexpected(ErrnoMessage("write", path));
+        }
+        if (written == 0) {
+            return tl::make_unexpected("write '" + path + "' made no progress");
+        }
+        data += written;
+        remaining -= static_cast<size_t>(written);
+    }
+    return {};
+}
+
+tl::expected<void, std::string> FsyncDirectory(const std::string& path) {
+    std::filesystem::path dir = std::filesystem::path(path).parent_path();
+    if (dir.empty()) {
+        dir = ".";
+    }
+    int fd = ::open(dir.string().c_str(), O_RDONLY | O_DIRECTORY);
+    if (fd < 0) {
+        return tl::make_unexpected(
+            ErrnoMessage("open directory", dir.string()));
+    }
+    auto close_fd = [&] { ::close(fd); };
+    if (::fsync(fd) != 0) {
+        std::string error = ErrnoMessage("fsync directory", dir.string());
+        close_fd();
+        return tl::make_unexpected(error);
+    }
+    close_fd();
+    return {};
+}
+
+std::string MakeTempPath(const std::string& path) {
+    const auto now =
+        std::chrono::steady_clock::now().time_since_epoch().count();
+    std::ostringstream oss;
+    oss << path << ".tmp." << ::getpid() << "."
+        << std::hash<std::thread::id>{}(std::this_thread::get_id()) << "."
+        << now;
+    return oss.str();
+}
+
+std::string QuoteYamlDoubleQuotedScalar(const std::string& value) {
+    std::ostringstream out;
+    out << '"';
+    for (unsigned char c : value) {
+        switch (c) {
+            case '\\':
+                out << "\\\\";
+                break;
+            case '"':
+                out << "\\\"";
+                break;
+            case '\0':
+                out << "\\0";
+                break;
+            case '\a':
+                out << "\\a";
+                break;
+            case '\b':
+                out << "\\b";
+                break;
+            case '\t':
+                out << "\\t";
+                break;
+            case '\n':
+                out << "\\n";
+                break;
+            case '\v':
+                out << "\\v";
+                break;
+            case '\f':
+                out << "\\f";
+                break;
+            case '\r':
+                out << "\\r";
+                break;
+            default:
+                if (c < 0x20 || c == 0x7f) {
+                    out << "\\x" << std::uppercase << std::hex << std::setw(2)
+                        << std::setfill('0') << static_cast<int>(c) << std::dec
+                        << std::nouppercase << std::setfill(' ');
+                } else {
+                    out << static_cast<char>(c);
+                }
+                break;
+        }
+    }
+    out << '"';
+    return out.str();
+}
+
+}  // namespace
+
+tl::expected<uint64_t, std::string> ParseTenantQuotaBytes(
+    const std::string& value) {
+    if (value.empty()) {
+        return tl::make_unexpected("quota must not be empty");
+    }
+
+    size_t digits = 0;
+    while (digits < value.size() && value[digits] >= '0' &&
+           value[digits] <= '9') {
+        ++digits;
+    }
+    if (digits == 0) {
+        return tl::make_unexpected("quota must start with an integer");
+    }
+
+    uint64_t number = 0;
+    for (size_t i = 0; i < digits; ++i) {
+        const uint64_t digit = static_cast<uint64_t>(value[i] - '0');
+        if (number > (std::numeric_limits<uint64_t>::max() - digit) / 10) {
+            return tl::make_unexpected("quota integer overflows uint64");
+        }
+        number = number * 10 + digit;
+    }
+    if (number == 0) {
+        return tl::make_unexpected("quota must be positive");
+    }
+
+    const std::string unit = value.substr(digits);
+    uint64_t multiplier = 1;
+    if (unit.empty() || unit == "B") {
+        multiplier = 1;
+    } else if (unit == "KB") {
+        multiplier = 1024ULL;
+    } else if (unit == "MB") {
+        multiplier = 1024ULL * 1024ULL;
+    } else if (unit == "GB") {
+        multiplier = 1024ULL * 1024ULL * 1024ULL;
+    } else if (unit == "TB") {
+        multiplier = 1024ULL * 1024ULL * 1024ULL * 1024ULL;
+    } else {
+        return tl::make_unexpected("unsupported quota unit '" + unit + "'");
+    }
+
+    if (number > std::numeric_limits<uint64_t>::max() / multiplier) {
+        return tl::make_unexpected("quota byte value overflows uint64");
+    }
+    return number * multiplier;
+}
+
+tl::expected<TenantQuotaPolicySnapshot, std::string> ParseTenantQuotaPolicyYaml(
+    const std::string& yaml) {
+    YAML::Node root;
+    try {
+        root = YAML::Load(yaml);
+    } catch (const YAML::Exception& e) {
+        return tl::make_unexpected(std::string("invalid YAML: ") + e.what());
+    }
+
+    if (!root || !root.IsMap()) {
+        return tl::make_unexpected("tenant quota policy must be a YAML map");
+    }
+    const auto version_node = root["version"];
+    if (!version_node || !version_node.IsScalar()) {
+        return tl::make_unexpected("tenant quota policy version is required");
+    }
+    int version = 0;
+    try {
+        version = version_node.as<int>();
+    } catch (const YAML::Exception& e) {
+        return tl::make_unexpected(std::string("invalid version: ") + e.what());
+    }
+    if (version != 1) {
+        return tl::make_unexpected("unsupported tenant quota policy version: " +
+                                   std::to_string(version));
+    }
+
+    const auto tenants_node = root["tenants"];
+    if (!tenants_node || !tenants_node.IsSequence()) {
+        return tl::make_unexpected("tenants must be a YAML sequence");
+    }
+
+    TenantQuotaPolicySnapshot snapshot;
+    for (size_t i = 0; i < tenants_node.size(); ++i) {
+        const auto entry = tenants_node[i];
+        if (!entry || !entry.IsMap()) {
+            return tl::make_unexpected("tenant entry must be a YAML map");
+        }
+        const auto name_node = entry["name"];
+        const auto quota_node = entry["quota"];
+        if (!name_node || !name_node.IsScalar()) {
+            return tl::make_unexpected("tenant name is required");
+        }
+        if (!quota_node || !quota_node.IsScalar()) {
+            return tl::make_unexpected("tenant quota is required");
+        }
+
+        std::string name;
+        std::string quota;
+        try {
+            name = name_node.as<std::string>();
+            quota = quota_node.as<std::string>();
+        } catch (const YAML::Exception& e) {
+            return tl::make_unexpected(std::string("invalid tenant entry: ") +
+                                       e.what());
+        }
+        if (!IsValidTenantQuotaName(name)) {
+            return tl::make_unexpected("invalid tenant name '" + name + "'");
+        }
+        name = NormalizeTenantId(name);
+        if (!IsValidTenantQuotaName(name)) {
+            return tl::make_unexpected("invalid tenant name '" + name + "'");
+        }
+        if (snapshot.tenant_quotas.contains(name)) {
+            return tl::make_unexpected("duplicate tenant name '" + name + "'");
+        }
+
+        auto quota_bytes = ParseTenantQuotaBytes(quota);
+        if (!quota_bytes) {
+            return tl::make_unexpected("invalid quota for tenant '" + name +
+                                       "': " + quota_bytes.error());
+        }
+        snapshot.tenant_quotas.emplace(std::move(name), quota_bytes.value());
+    }
+
+    return snapshot;
+}
+
+std::string FormatTenantQuotaPolicyYaml(
+    const TenantQuotaPolicySnapshot& snapshot) {
+    std::ostringstream out;
+    out << "version: 1\n\n";
+    if (snapshot.tenant_quotas.empty()) {
+        out << "tenants: []\n";
+        return out.str();
+    }
+    out << "tenants:\n";
+    for (const auto& [tenant_id, quota] : snapshot.tenant_quotas) {
+        out << "  - name: " << QuoteYamlDoubleQuotedScalar(tenant_id) << "\n";
+        out << "    quota: " << quota << "\n";
+    }
+    return out.str();
+}
+
+YamlTenantQuotaPolicyStore::YamlTenantQuotaPolicyStore(std::string path)
+    : path_(std::move(path)) {}
+
+tl::expected<TenantQuotaPolicySnapshot, std::string>
+YamlTenantQuotaPolicyStore::Load() {
+    std::lock_guard<std::mutex> lock(mutex_);
+    std::ifstream input(path_);
+    if (!input.is_open()) {
+        return tl::make_unexpected("failed to open tenant quota policy file '" +
+                                   path_ + "'");
+    }
+    std::ostringstream buffer;
+    buffer << input.rdbuf();
+    if (input.bad()) {
+        return tl::make_unexpected("failed to read tenant quota policy file '" +
+                                   path_ + "'");
+    }
+    return ParseTenantQuotaPolicyYaml(buffer.str());
+}
+
+tl::expected<void, std::string> YamlTenantQuotaPolicyStore::Save(
+    const TenantQuotaPolicySnapshot& snapshot) {
+    std::lock_guard<std::mutex> lock(mutex_);
+    const std::string content = FormatTenantQuotaPolicyYaml(snapshot);
+    const std::string tmp_path = MakeTempPath(path_);
+
+    int fd = ::open(tmp_path.c_str(), O_CREAT | O_TRUNC | O_WRONLY, 0644);
+    if (fd < 0) {
+        return tl::make_unexpected(ErrnoMessage("open", tmp_path));
+    }
+
+    auto cleanup = [&] {
+        ::close(fd);
+        ::unlink(tmp_path.c_str());
+    };
+
+    auto write_result = WriteAll(fd, content, tmp_path);
+    if (!write_result) {
+        cleanup();
+        return tl::make_unexpected(write_result.error());
+    }
+    if (::fsync(fd) != 0) {
+        std::string error = ErrnoMessage("fsync", tmp_path);
+        cleanup();
+        return tl::make_unexpected(error);
+    }
+    if (::close(fd) != 0) {
+        std::string error = ErrnoMessage("close", tmp_path);
+        ::unlink(tmp_path.c_str());
+        return tl::make_unexpected(error);
+    }
+    fd = -1;
+
+    if (::rename(tmp_path.c_str(), path_.c_str()) != 0) {
+        std::string error = ErrnoMessage("rename", path_);
+        ::unlink(tmp_path.c_str());
+        return tl::make_unexpected(error);
+    }
+
+    auto fsync_result = FsyncDirectory(path_);
+    if (!fsync_result) {
+        LOG(WARNING) << "failed to fsync tenant quota policy directory after "
+                        "rename: "
+                     << fsync_result.error();
+    }
+    return {};
+}
+
+tl::expected<std::unique_ptr<TenantQuotaPolicyStore>, std::string>
+CreateTenantQuotaPolicyStore(const std::string& type, const std::string& uri) {
+    if (type != "file") {
+        return tl::make_unexpected("unsupported tenant quota connector type '" +
+                                   type + "'");
+    }
+    if (uri.empty()) {
+        return tl::make_unexpected(
+            "tenant quota file connector requires a non-empty uri");
+    }
+    return std::make_unique<YamlTenantQuotaPolicyStore>(uri);
+}
+
+}  // namespace mooncake
diff --git a/mooncake-store/src/transfer_task.cpp b/mooncake-store/src/transfer_task.cpp
index 8a2989af..5995e3c4 100644
--- a/mooncake-store/src/transfer_task.cpp
+++ b/mooncake-store/src/transfer_task.cpp
@@ -3,14 +3,172 @@
 #include <glog/logging.h>
 
 #include <algorithm>
+#include <cctype>
+#include <chrono>
+#include <cerrno>
 #include <cstdlib>
 #include <sstream>
+#include <string>
 #include <vector>
+#include "gpu_staging_utils.h"
 #include "transfer_engine.h"
 #include "transport/transport.h"
+#ifdef USE_NOF
+#include "spdk/spdk_wrapper.h"
+#endif
 
+#ifdef USE_NOF
+static bool IsTruthyEnv(const char* value) {
+    if (!value) {
+        return false;
+    }
+    std::string normalized(value);
+    std::transform(
+        normalized.begin(), normalized.end(), normalized.begin(),
+        [](unsigned char c) { return static_cast<char>(std::tolower(c)); });
+    return normalized == "1" || normalized == "true" || normalized == "yes" ||
+           normalized == "on";
+}
+
+static bool IsSpdkNofDebugEnabled() {
+    static const bool enabled = IsTruthyEnv(std::getenv("MC_NOF_DEBUG"));
+    return enabled;
+}
+
+static int GetSpdkNofDebugIntervalMs() {
+    static const int interval_ms = []() {
+        const char* raw_value = std::getenv("MC_NOF_DEBUG_INTERVAL_MS");
+        if (!raw_value) {
+            return 1000;
+        }
+        char* end_ptr = nullptr;
+        long parsed = std::strtol(raw_value, &end_ptr, 10);
+        if (end_ptr == raw_value || (end_ptr != nullptr && *end_ptr != '\0') ||
+            parsed <= 0) {
+            return 1000;
+        }
+        return static_cast<int>(parsed);
+    }();
+    return interval_ms;
+}
+
+static int GetPositiveEnvOrDefault(const char* name, int default_value) {
+    const char* raw_value = std::getenv(name);
+    if (!raw_value || raw_value[0] == '\0') {
+        return default_value;
+    }
+
+    errno = 0;
+    char* end_ptr = nullptr;
+    long parsed = std::strtol(raw_value, &end_ptr, 10);
+    if (errno != 0 || end_ptr == raw_value ||
+        (end_ptr != nullptr && *end_ptr != '\0') || parsed <= 0 ||
+        parsed > std::numeric_limits<int>::max()) {
+        LOG(WARNING) << "Invalid value for " << name << ": " << raw_value
+                     << ", using default " << default_value;
+        return default_value;
+    }
+
+    return static_cast<int>(parsed);
+}
+
+static int GetSpdkNofSubmitChunkBytes() {
+    static const int value = GetPositiveEnvOrDefault(
+        "MC_NOF_SUBMIT_CHUNK_BYTES", mooncake::kDefaultSpdkNofSubmitChunkBytes);
+    return value;
+}
+
+static int GetSpdkNofInflightBytesLimit() {
+    static const int value =
+        GetPositiveEnvOrDefault("MC_NOF_INFLIGHT_BYTES_LIMIT",
+                                mooncake::kDefaultSpdkNofInflightBytesLimit);
+    return value;
+}
+
+static int GetSpdkNofWorkerCount() {
+    static const int value = GetPositiveEnvOrDefault(
+        "MC_NOF_WORKERS", mooncake::kDefaultSpdkNofWorkers);
+    return value;
+}
+
+static int CountSpdkNofQueuedTasks(const mooncake::SpdkNofTask* head) {
+    int count = 0;
+    const mooncake::SpdkNofTask* cursor = head;
+    while (cursor != nullptr) {
+        ++count;
+        cursor = cursor->nxt;
+    }
+    return count;
+}
+
+static inline void SpdkNofTaskCompletion(mooncake::SpdkNofTask* task) {
+    if (task->remaining_lba == 0 && task->outstanding_sub_io == 0) {
+        task->state->set_completed(task->failed
+                                       ? mooncake::ErrorCode::TRANSFER_FAIL
+                                       : mooncake::ErrorCode::OK);
+        if (!task->on_chain) {
+            delete task;
+        }
+    }
+}
+
+static void nvmf_io_complete(void* ctx, const struct spdk_nvme_cpl* cpl) {
+    if (!ctx) {
+        LOG(ERROR) << "nvmf_io_complete ctx is null";
+        return;
+    }
+
+    mooncake::SpdkNofSubTask* sub_task =
+        reinterpret_cast<mooncake::SpdkNofSubTask*>(ctx);
+    mooncake::SpdkNofTask* task = sub_task->task;
+    mooncake::SpdkNofQos* nof_qos = task->nof_qos;
+    int op = task->op;
+    if (--(*task->io_count) < 0) {
+        LOG(ERROR) << "total outstanding io < 0";
+    }
+
+    if (--(task->outstanding_sub_io) < 0) {
+        LOG(ERROR) << "task outstanding io < 0";
+    }
+
+    nof_qos->inflight_blocks[op] -= sub_task->submit_lba_count;
+    if (nof_qos->inflight_blocks[op] < 0) {
+        LOG(ERROR) << "task outstanding io < 0";
+    }
+
+    if (spdk_nvme_cpl_is_error(cpl)) {
+        LOG(ERROR) << "task_complete: I/O failed"
+                   << spdk_nvme_cpl_get_status_string(&cpl->status);
+        task->remaining_lba = 0;
+        task->failed = true;
+    }
+
+    SpdkNofTaskCompletion(task);
+
+    sub_task->sub_task_pool->push(sub_task);
+}
+#endif
 namespace mooncake {
 
+#ifdef USE_NOF
+SpdkNofQos::SpdkNofQos(uint32_t block_size) {
+    int block_size_int = static_cast<int>(block_size);
+    if (block_size_int <= 0) {
+        block_size_int = 1;
+    }
+
+    blocks_per_chunk =
+        std::max(1, GetSpdkNofSubmitChunkBytes() / block_size_int);
+    inflight_blocks_limit =
+        std::max(1, GetSpdkNofInflightBytesLimit() / block_size_int);
+    for (int i = 0; i < kSpdkNofOpNum; ++i) {
+        inflight_blocks[i] = 0;
+        head[i] = nullptr;
+        tail[i] = nullptr;
+    }
+}
+#endif
+
 // ============================================================================
 // FilereadWorkerPool Implementation
 // ============================================================================
@@ -118,6 +276,304 @@ void FilereadWorkerPool::workerThread() {
     VLOG(2) << "FilereadWorkerPool worker thread exiting";
 }
 
+// ============================================================================
+// SpdkNofWorkerPool Implementation
+// ============================================================================
+// to fully utilize the available ssd bandwidth, we use a default of 4 worker
+// threads.
+
+#ifdef USE_NOF
+SpdkNofWorkerPool::SpdkNofWorkerPool(int numa_socket_id)
+    : worker_count_(GetSpdkNofWorkerCount()),
+      numa_socket_id_(numa_socket_id),
+      task_queue_(std::make_unique<std::queue<SpdkNofTask>[]>(worker_count_)),
+      queue_mutex_(std::make_unique<std::mutex[]>(worker_count_)),
+      queue_cv_(std::make_unique<std::condition_variable[]>(worker_count_)),
+      shutdown_(false) {
+    VLOG(1) << "Creating SpdkNofWorkerPool with " << worker_count_
+            << " workers";
+
+    // Start worker threads
+    workers_.reserve(worker_count_);
+    for (int i = 0; i < worker_count_; ++i) {
+        workers_.emplace_back(&SpdkNofWorkerPool::workerThread, this, i);
+    }
+}
+
+SpdkNofWorkerPool::~SpdkNofWorkerPool() {
+    if (shutdown_.exchange(true)) {
+        return;
+    }
+
+    for (int i = 0; i < worker_count_; ++i) {
+        queue_cv_[i].notify_all();
+    }
+
+    for (auto& worker : workers_) {
+        if (worker.joinable()) {
+            worker.join();
+        }
+    }
+
+    VLOG(1) << "SpdkNofWorkerPool destroyed";
+}
+
+void SpdkNofWorkerPool::submitTask(SpdkNofTask task) {
+    if (!task.state) {
+        LOG(ERROR) << "Attempting to submit spdk nof task without state";
+        return;
+    }
+
+    if (shutdown_.load()) {
+        LOG(WARNING)
+            << "Attempting to submit task to shutdown SpdkNofWorkerPool";
+        task.state->set_completed(ErrorCode::TRANSFER_FAIL);
+        return;
+    }
+
+    int worker_idx = -1;
+    {
+        std::lock_guard<std::mutex> lock(seg_mutex_);
+        nof_seg_handle* seg = task.seg_handle;
+        bool new_binding = false;
+        if (seg_to_worker_.find(seg) != seg_to_worker_.end()) {
+            worker_idx = seg_to_worker_[seg];
+        } else {
+            worker_idx = (seg_num++ % worker_count_);
+            seg_to_worker_[seg] = worker_idx;
+            new_binding = true;
+        }
+        if (new_binding && IsSpdkNofDebugEnabled()) {
+            LOG(INFO) << "nof_worker_bind seg_handle=" << seg
+                      << " worker_idx=" << worker_idx;
+        }
+    }
+    if (worker_idx < 0 || worker_idx >= worker_count_) {
+        LOG(ERROR) << "seg is not bind to invalid worker " << worker_idx;
+        task.state->set_completed(ErrorCode::TRANSFER_FAIL);
+        return;
+    }
+
+    {
+        std::lock_guard<std::mutex> lock(queue_mutex_[worker_idx]);
+        task_queue_[worker_idx].push(std::move(task));
+    }
+    queue_cv_[worker_idx].notify_one();
+}
+
+static bool HasBufferedTask(
+    const std::map<nof_seg_handle*, std::unique_ptr<SpdkNofQos>>& seg_to_qos) {
+    for (const auto& [_, nof_qos] : seg_to_qos) {
+        if (!nof_qos->Empty()) {
+            return true;
+        }
+    }
+    return false;
+}
+
+constexpr int kSpdkNofSubTaskChunkSize = 4096;
+
+static inline bool CheckSubTaskPool(
+    std::stack<SpdkNofSubTask*>& sub_task_pool,
+    std::vector<SpdkNofSubTask*>& sub_task_chunks, int work_idx) {
+    if (!sub_task_pool.empty()) {
+        return true;
+    }
+    SpdkNofSubTask* sub_tasks =
+        new (std::nothrow) SpdkNofSubTask[kSpdkNofSubTaskChunkSize];
+    if (!sub_tasks) {
+        LOG(ERROR) << "alloc SpdkNofSubTask failed, worker " << work_idx;
+        return false;
+    }
+    sub_task_chunks.push_back(sub_tasks);
+
+    for (int i = 0; i < kSpdkNofSubTaskChunkSize; ++i) {
+        sub_tasks[i].sub_task_pool = &sub_task_pool;
+        sub_task_pool.push(&sub_tasks[i]);
+    }
+
+    return true;
+}
+
+void SpdkNofWorkerPool::workerThread(int work_idx) {
+    bindToSocket(numa_socket_id_);
+    VLOG(2) << "SpdkNofWorkerPool worker thread started";
+
+    int64_t total_outstanding_io = 0;
+    // std::set<nof_seg_handle *> seg_set;
+    std::map<nof_seg_handle*, std::unique_ptr<SpdkNofQos>> seg_to_qos;
+    std::stack<SpdkNofSubTask*> sub_task_pool;
+    std::vector<SpdkNofSubTask*> sub_task_chunks;
+    auto& task_queue = task_queue_[work_idx];
+    auto& queue_cv = queue_cv_[work_idx];
+    auto& queue_mutex = queue_mutex_[work_idx];
+    auto last_debug_snapshot = std::chrono::steady_clock::now();
+
+    if (!CheckSubTaskPool(sub_task_pool, sub_task_chunks, work_idx)) {
+        return;
+    }
+
+    while (true) {
+        // Wait for task or shutdown signal
+        {
+            std::unique_lock<std::mutex> lock(queue_mutex);
+            queue_cv.wait(
+                lock, [this, &task_queue, &total_outstanding_io, &seg_to_qos] {
+                    return shutdown_.load() || !task_queue.empty() ||
+                           total_outstanding_io || HasBufferedTask(seg_to_qos);
+                });
+
+            if (shutdown_.load() && task_queue.empty() &&
+                (total_outstanding_io == 0) && !HasBufferedTask(seg_to_qos)) {
+                break;
+            }
+
+            while (!task_queue.empty()) {
+                SpdkNofTask* task = new (std::nothrow)
+                    SpdkNofTask(std::move(task_queue.front()));
+                if (task == nullptr) {
+                    LOG(ERROR)
+                        << "alloc SpdkNofTask failed, worker " << work_idx;
+                    continue;
+                }
+
+                SpdkNofQos* nof_qos = nullptr;
+                auto it = seg_to_qos.find(task->seg_handle);
+                if (it == seg_to_qos.end()) {
+                    auto qos = std::make_unique<SpdkNofQos>(
+                        SpdkWrapper::GetInstance().GetBlockSize(
+                            task->seg_handle));
+                    if (qos == nullptr) {
+                        LOG(ERROR)
+                            << "alloc SpdkNofQos failed, worker " << work_idx;
+                        delete task;
+                        continue;
+                    }
+                    nof_qos = qos.get();
+                    seg_to_qos[task->seg_handle] = std::move(qos);
+                    if (IsSpdkNofDebugEnabled()) {
+                        LOG(INFO)
+                            << "nof_qos_create worker_idx=" << work_idx
+                            << " seg_handle=" << task->seg_handle
+                            << " blocks_per_chunk="
+                            << seg_to_qos[task->seg_handle]->blocks_per_chunk
+                            << " inflight_blocks_limit="
+                            << seg_to_qos[task->seg_handle]
+                                   ->inflight_blocks_limit;
+                    }
+                } else {
+                    nof_qos = it->second.get();
+                }
+                task->io_count = &total_outstanding_io;
+                task->nof_qos = nof_qos;
+                task->on_chain = true;
+                nof_qos->PushTask(task);
+                task_queue.pop();
+            }
+        }
+
+        for (auto& [seg_handle, nof_qos] : seg_to_qos) {
+            uint32_t block_size =
+                SpdkWrapper::GetInstance().GetBlockSize(seg_handle);
+            for (int i = 0; i < kSpdkNofOpNum; ++i) {
+                int avail_blocks = nof_qos->inflight_blocks_limit -
+                                   nof_qos->inflight_blocks[i];
+                while (nof_qos->head[i] && avail_blocks > 0) {
+                    SpdkNofTask* task = nof_qos->head[i];
+                    SpdkNofSubTask* sub_task;
+                    while (task->remaining_lba > 0 && avail_blocks > 0) {
+                        uint32_t submit_lba_count = std::min(
+                            avail_blocks, std::min(task->remaining_lba,
+                                                   nof_qos->blocks_per_chunk));
+                        int lba_off = task->lba_count - task->remaining_lba;
+                        uint64_t submit_lba = task->lba + lba_off;
+                        void* submit_ptr = reinterpret_cast<void*>(
+                            reinterpret_cast<char*>(task->ptr) +
+                            lba_off * block_size);
+
+                        if (!CheckSubTaskPool(sub_task_pool, sub_task_chunks,
+                                              work_idx)) {
+                            task->failed = true;
+                            task->remaining_lba = 0;
+                            break;
+                        }
+                        sub_task = sub_task_pool.top();
+                        sub_task_pool.pop();
+                        sub_task->task = task;
+                        sub_task->submit_lba_count = submit_lba_count;
+
+                        int ret = SpdkWrapper::GetInstance().SubmitRequest(
+                            task->seg_handle, submit_ptr, submit_lba,
+                            submit_lba_count, task->op, nvmf_io_complete,
+                            sub_task);
+                        if (ret != 0) {
+                            LOG(ERROR) << "work " << work_idx << ", seg "
+                                       << task->seg_handle << " submit io fail";
+                            task->failed = true;
+                            task->remaining_lba = 0;
+                        } else {
+                            task->idx++;
+                            task->remaining_lba -= submit_lba_count;
+                            nof_qos->inflight_blocks[i] += submit_lba_count;
+                            avail_blocks -= submit_lba_count;
+                            task->outstanding_sub_io++;
+                            total_outstanding_io++;
+                        }
+                    }
+                    if (task->remaining_lba == 0) {
+                        nof_qos->PopTask(i);
+                        task->on_chain = false;
+                        SpdkNofTaskCompletion(task);
+                    }
+                }
+            }
+        }
+
+        if (total_outstanding_io > 0) {
+            int64_t ret = 0;
+            for (auto& it : seg_to_qos) {
+                ret = SpdkWrapper::GetInstance().NvmePollProcessCompletion(
+                    it.first, 0);
+                if (ret < 0) {
+                    LOG(ERROR) << "poll completion error: ret " << ret;
+                }
+            }
+        }
+
+        if (IsSpdkNofDebugEnabled()) {
+            auto now = std::chrono::steady_clock::now();
+            auto elapsed =
+                std::chrono::duration_cast<std::chrono::milliseconds>(
+                    now - last_debug_snapshot);
+            if (elapsed.count() >= GetSpdkNofDebugIntervalMs()) {
+                for (const auto& [seg_handle, nof_qos] : seg_to_qos) {
+                    LOG(INFO)
+                        << "nof_qos_state worker_idx=" << work_idx
+                        << " seg_handle=" << seg_handle
+                        << " inflight_read=" << nof_qos->inflight_blocks[0]
+                        << " inflight_write=" << nof_qos->inflight_blocks[1]
+                        << " inflight_limit=" << nof_qos->inflight_blocks_limit
+                        << " queued_read="
+                        << CountSpdkNofQueuedTasks(nof_qos->head[0])
+                        << " queued_write="
+                        << CountSpdkNofQueuedTasks(nof_qos->head[1])
+                        << " total_outstanding_io=" << total_outstanding_io;
+                }
+                last_debug_snapshot = now;
+            }
+        }
+    }
+
+    for (auto* sub_tasks : sub_task_chunks) {
+        delete[] sub_tasks;
+    }
+    // seg_to_qos will automatically clean up SpdkNofQos objects using
+    // unique_ptr
+
+    VLOG(2) << "SpdkNofWorkerPool worker thread exiting";
+}
+#endif
+
 // ============================================================================
 // MemcpyWorkerPool Implementation
 // ============================================================================
@@ -193,13 +649,35 @@ void MemcpyWorkerPool::workerThread() {
         // Execute the task if we have one
         if (task.state) {
             try {
+                bool ok = true;
                 for (const auto& op : task.operations) {
-                    std::memcpy(op.dest, op.src, op.size);
+                    int src_dev = -1, dst_dev = -1;
+                    bool src_on_gpu =
+                        gpu_staging::IsDevicePointer(op.src, &src_dev);
+                    bool dst_on_gpu =
+                        gpu_staging::IsDevicePointer(op.dest, &dst_dev);
+
+                    if (!src_on_gpu && !dst_on_gpu) {
+                        std::memcpy(op.dest, op.src, op.size);
+                    } else {
+                        int dev = src_on_gpu ? src_dev : dst_dev;
+                        gpu_staging::SetDevice(dev);
+                        if (!gpu_staging::CopyAuto(op.dest, op.src, op.size)) {
+                            LOG(ERROR)
+                                << "GPU memcpy failed: src_dev=" << src_dev
+                                << " dst_dev=" << dst_dev
+                                << " size=" << op.size;
+                            ok = false;
+                            break;
+                        }
+                    }
                 }
 
-                VLOG(2) << "Memcpy task completed successfully with "
-                        << task.operations.size() << " operations";
-                task.state->set_completed(ErrorCode::OK);
+                VLOG(2) << "Memcpy task completed with "
+                        << task.operations.size() << " operations"
+                        << (ok ? "" : " (with GPU copy failure)");
+                task.state->set_completed(ok ? ErrorCode::OK
+                                             : ErrorCode::TRANSFER_FAIL);
             } catch (const std::exception& e) {
                 LOG(ERROR) << "Exception during async memcpy: " << e.what();
                 task.state->set_completed(ErrorCode::TRANSFER_FAIL);
@@ -424,10 +902,17 @@ TransferStrategy TransferFuture::strategy() const {
 
 TransferSubmitter::TransferSubmitter(TransferEngine& engine,
                                      std::shared_ptr<StorageBackend>& backend,
-                                     TransferMetric* transfer_metric)
+                                     const std::string& local_hostname,
+                                     TransferMetric* transfer_metric,
+                                     int numa_socket_id)
     : engine_(engine),
+      local_endpoint_(engine.getLocalIpAndPort()),
       memcpy_pool_(std::make_unique<MemcpyWorkerPool>()),
+#ifdef USE_NOF
+      spdk_nvmf_pool_(std::make_unique<SpdkNofWorkerPool>(numa_socket_id)),
+#endif
       fileread_pool_(std::make_unique<FilereadWorkerPool>(backend)),
+      local_hostname_(local_hostname),
       transfer_metric_(transfer_metric) {
     // Read MC_STORE_MEMCPY environment variable.
     // When not set, auto-detect based on transport type:
@@ -444,7 +929,7 @@ TransferSubmitter::TransferSubmitter(TransferEngine& engine,
         std::string env_str(env_value);
         // Convert to lowercase for case-insensitive comparison
         std::transform(env_str.begin(), env_str.end(), env_str.begin(),
-                       ::tolower);
+                       [](unsigned char c) { return std::tolower(c); });
         if (env_str == "false" || env_str == "0" || env_str == "no" ||
             env_str == "off") {
             memcpy_enabled_ = false;
@@ -464,7 +949,7 @@ TransferSubmitter::TransferSubmitter(TransferEngine& engine,
 
 std::optional<TransferFuture> TransferSubmitter::submit(
     const Replica::Descriptor& replica, std::vector<Slice>& slices,
-    TransferRequest::OpCode op_code) {
+    TransferRequest::OpCode op_code, void* ptr, size_t size) {
     std::optional<TransferFuture> future;
 
     if (replica.is_memory_replica()) {
@@ -493,6 +978,20 @@ std::optional<TransferFuture> TransferSubmitter::submit(
                     return std::nullopt;
             }
         }
+    } else if (replica.is_nof_replica()) {
+#ifdef USE_NOF
+        auto& ssd_desc = replica.get_nof_descriptor();
+        auto& handle = ssd_desc.buffer_descriptor;
+
+        if (!ptr || (size == 0)) {
+            return std::nullopt;
+        }
+
+        future = submitSpdkNofOperation(handle, ptr, size, op_code);
+#else
+        LOG(ERROR) << "NoF transfer requested while USE_NOF is disabled";
+        return std::nullopt;
+#endif
     } else {
         future = submitFileReadOperation(replica, slices, op_code);
     }
@@ -551,25 +1050,38 @@ std::optional<TransferFuture>
 TransferSubmitter::submit_batch_get_offload_object(
     const std::string& transfer_engine_addr,
     const std::vector<std::string>& keys, const std::vector<uint64_t>& pointers,
-    const std::unordered_map<std::string, Slice>& batched_slices) {
+    const std::unordered_map<std::string, std::vector<Slice>>& batched_slices) {
     std::optional<TransferFuture> future;
     std::vector<TransferRequest> requests;
+    // Open the segment once — all keys share the same transfer_engine_addr.
+    SegmentHandle seg = engine_.openSegment(transfer_engine_addr);
+    if (seg == static_cast<uint64_t>(ERR_INVALID_ARGUMENT)) {
+        LOG(ERROR) << "Failed to open segment " << transfer_engine_addr;
+        // nullopt = failure (caller checks !future).  The function returns
+        // std::optional so tl::unexpected is not available here.
+        return std::nullopt;
+    }
     for (size_t i = 0; i < keys.size(); ++i) {
-        auto key = keys[i];
-        auto pointer = pointers[i];
-        SegmentHandle seg = engine_.openSegment(transfer_engine_addr);
-        if (seg == static_cast<uint64_t>(ERR_INVALID_ARGUMENT)) {
-            LOG(ERROR) << "Failed to open segment " << transfer_engine_addr;
-            return std::nullopt;
+        const auto& key = keys[i];
+        const uint64_t pointer = pointers[i];
+        auto it = batched_slices.find(key);
+        if (it == batched_slices.end()) {
+            LOG(ERROR) << "Key not found in batched_slices: " << key;
+            return std::nullopt;  // fail closed
+        }
+        // Emit one TransferRequest per slice: the on-disk blob is read
+        // sequentially while slices may point to non-contiguous GPU memory.
+        uint64_t offset = 0;
+        for (const auto& slice : it->second) {
+            TransferRequest request;
+            request.opcode = TransferRequest::READ;
+            request.source = static_cast<char*>(slice.ptr);
+            request.target_id = seg;
+            request.target_offset = pointer + offset;
+            request.length = slice.size;
+            requests.emplace_back(request);
+            offset += slice.size;
         }
-        const auto& slice = batched_slices.find(key)->second;
-        TransferRequest request;
-        request.opcode = TransferRequest::READ;
-        request.source = static_cast<char*>(slice.ptr);
-        request.target_id = seg;
-        request.target_offset = pointer;
-        request.length = slice.size;
-        requests.emplace_back(request);
     }
     return submitTransfer(requests);
 }
@@ -729,6 +1241,9 @@ std::optional<TransferFuture> TransferSubmitter::submitRangeRead(
         }
 
         future = submitMemoryReadOperation(handle, slices, src_offset);
+    } else if (replica.is_nof_replica()) {
+        LOG(ERROR) << "Range read not supported for NoF replicas";
+        return std::nullopt;
     } else if (replica.is_disk_replica() || replica.is_local_disk_replica()) {
         LOG(ERROR)
             << "Range read not supported for disk replicas (use full read)";
@@ -742,6 +1257,46 @@ std::optional<TransferFuture> TransferSubmitter::submitRangeRead(
     return future;
 }
 
+#ifdef USE_NOF
+std::optional<TransferFuture> TransferSubmitter::submitSpdkNofOperation(
+    const AllocatedBuffer::Descriptor& handle, void* ptr, size_t size,
+    const TransferRequest::OpCode op_code) {
+    if (handle.transport_endpoint_.empty() || handle.size_ < size) {
+        LOG(ERROR) << "Invalid NoF request endpoint="
+                   << handle.transport_endpoint_
+                   << ", buffer_size=" << handle.size_
+                   << ", request_size=" << size;
+        return std::nullopt;
+    }
+
+    nof_seg_handle* seg_handle =
+        SpdkWrapper::GetInstance().OpenNofSegment(handle.transport_endpoint_);
+    if (!seg_handle) {
+        LOG(ERROR) << "Failed to open NoF segment endpoint="
+                   << handle.transport_endpoint_;
+        return std::nullopt;
+    }
+
+    uint32_t block_size = SpdkWrapper::GetInstance().GetBlockSize(seg_handle);
+    if (block_size == INVALID_BLOCK_SIZE ||
+        handle.buffer_address_ % block_size != 0 || size % block_size != 0 ||
+        reinterpret_cast<std::uintptr_t>(ptr) % block_size != 0) {
+        LOG(ERROR) << "NoF request offset=" << handle.buffer_address_
+                   << ", ptr=" << ptr << ", size=" << size
+                   << " is not aligned to block size " << block_size;
+        return std::nullopt;
+    }
+
+    auto state = std::make_shared<SpdkNofOperationState>();
+    SpdkNofTask task(seg_handle, ptr, handle.buffer_address_ / block_size,
+                     size / block_size, op_code, state);
+    spdk_nvmf_pool_->submitTask(std::move(task));
+
+    VLOG(1) << "SPDK NoF transfer submitted to " << handle.transport_endpoint_;
+    return TransferFuture(state);
+}
+#endif
+
 std::optional<TransferFuture> TransferSubmitter::submitFileReadOperation(
     const Replica::Descriptor& replica, std::vector<Slice>& slices,
     TransferRequest::OpCode op_code) {
@@ -778,8 +1333,8 @@ TransferStrategy TransferSubmitter::selectStrategy(
 }
 
 namespace {
-// Helper function to extract IP address from endpoint string (ip:port format)
-// Supports both IPv4 (ip:port) and IPv6 ([ipv6]:port) formats
+// Helper function to extract IP address from endpoint string (ip:port format).
+// Supports both IPv4 (ip:port) and IPv6 ([ipv6]:port) formats.
 std::string extractIpAddress(const std::string& endpoint) {
     if (endpoint.empty()) {
         return "";
@@ -790,38 +1345,54 @@ std::string extractIpAddress(const std::string& endpoint) {
         size_t closing_bracket = endpoint.find(']');
         if (closing_bracket == std::string::npos) {
             LOG(WARNING) << "Invalid IPv6 endpoint format: " << endpoint;
-            // Return empty to disable local memcpy optimization
             return "";
         }
-        return endpoint.substr(1, closing_bracket - 1);  // Extract IPv6 address
+        return endpoint.substr(1, closing_bracket - 1);
     }
 
-    // Handle IPv4 or hostname:port format
-    // Find the last colon (to handle IPv6 addresses without brackets)
+    // Handle IPv4 or hostname:port format.
     size_t colon_pos = endpoint.rfind(':');
     if (colon_pos != std::string::npos) {
         return endpoint.substr(0, colon_pos);
     }
 
-    // No colon found, return the whole string (might be just IP or hostname)
+    // No colon found, return the whole string (might be just IP or hostname).
     return endpoint;
 }
 }  // namespace
 
-bool TransferSubmitter::isLocalTransfer(
-    const AllocatedBuffer::Descriptor& handle) const {
-    std::string local_ep = engine_.getLocalIpAndPort();
-    std::string local_ip = extractIpAddress(local_ep);
-
-    if (!local_ep.empty()) {
-        std::string handle_ip = extractIpAddress(handle.transport_endpoint_);
-        return !handle.transport_endpoint_.empty() && handle_ip == local_ip;
+bool TransferSubmitter::isSameProcessEndpoint(
+    const std::string& handle_endpoint, const std::string& local_endpoint) {
+    // Local memcpy requires that handle.buffer_address_ is a virtual address
+    // valid in THIS process. Same host is not enough: two processes on the
+    // same host share an IP but have distinct virtual address spaces, so a
+    // memcpy on a peer process's address would segfault. Require the full
+    // transport endpoint to match, which uniquely identifies the owning
+    // process.
+    if (handle_endpoint.empty() || local_endpoint.empty()) {
+        return false;
+    }
+    if (handle_endpoint == local_endpoint) {
+        return true;
+    }
+
+    const std::string handle_ip = extractIpAddress(handle_endpoint);
+    const std::string local_ip = extractIpAddress(local_endpoint);
+    if (!handle_ip.empty() && handle_ip == local_ip) {
+        VLOG(2) << "Disabling local memcpy for same-host endpoints with "
+                   "different process endpoints: handle="
+                << handle_endpoint << ", local=" << local_endpoint;
     }
 
-    // Without a local IP we cannot prove locality; disable memcpy.
     return false;
 }
 
+bool TransferSubmitter::isLocalTransfer(
+    const AllocatedBuffer::Descriptor& handle) const {
+    return isSameProcessEndpoint(handle.transport_endpoint_, local_hostname_) ||
+           isSameProcessEndpoint(handle.transport_endpoint_, local_endpoint_);
+}
+
 bool TransferSubmitter::validateTransferParams(
     const AllocatedBuffer::Descriptor& handle,
     const std::vector<Slice>& slices) const {
diff --git a/mooncake-store/src/types.cpp b/mooncake-store/src/types.cpp
index 2f5c4f45..68bdd7fc 100644
--- a/mooncake-store/src/types.cpp
+++ b/mooncake-store/src/types.cpp
@@ -37,10 +37,14 @@ const std::string& toString(ErrorCode errorCode) noexcept {
         {ErrorCode::OBJECT_REPLICA_BUSY, "OBJECT_REPLICA_BUSY"},
         {ErrorCode::TRANSFER_FAIL, "TRANSFER_FAIL"},
         {ErrorCode::RPC_FAIL, "RPC_FAIL"},
+        {ErrorCode::RPC_TIMEOUT, "RPC_TIMEOUT"},
         {ErrorCode::ETCD_OPERATION_ERROR, "ETCD_OPERATION_ERROR"},
         {ErrorCode::ETCD_KEY_NOT_EXIST, "ETCD_KEY_NOT_EXIST"},
         {ErrorCode::ETCD_TRANSACTION_FAIL, "ETCD_TRANSACTION_FAIL"},
         {ErrorCode::ETCD_CTX_CANCELLED, "ETCD_CTX_CANCELLED"},
+        {ErrorCode::OPLOG_ENTRY_NOT_FOUND, "OPLOG_ENTRY_NOT_FOUND"},
+        {ErrorCode::K8S_LEASE_OPERATION_ERROR, "K8S_LEASE_OPERATION_ERROR"},
+        {ErrorCode::K8S_LEASE_NOT_FOUND, "K8S_LEASE_NOT_FOUND"},
         {ErrorCode::UNAVAILABLE_IN_CURRENT_STATUS,
          "UNAVAILABLE_IN_CURRENT_STATUS"},
         {ErrorCode::UNAVAILABLE_IN_CURRENT_MODE, "UNAVAILABLE_IN_CURRENT_MODE"},
@@ -63,7 +67,16 @@ const std::string& toString(ErrorCode errorCode) noexcept {
         {ErrorCode::PERSISTENT_FAIL, "PERSISTENT_FAIL"},
         {ErrorCode::TASK_NOT_FOUND, "TASK_NOT_FOUND"},
         {ErrorCode::TASK_PENDING_LIMIT_EXCEEDED, "TASK_PENDING_LIMIT_EXCEEDED"},
-        {ErrorCode::JOB_NOT_FOUND, "JOB_NOT_FOUND"}};
+        {ErrorCode::JOB_NOT_FOUND, "JOB_NOT_FOUND"},
+        {ErrorCode::DFS_NETWORK_TIMEOUT, "DFS_NETWORK_TIMEOUT"},
+        {ErrorCode::DFS_SERVICE_UNAVAILABLE, "DFS_SERVICE_UNAVAILABLE"},
+        {ErrorCode::DFS_QUOTA_EXCEEDED, "DFS_QUOTA_EXCEEDED"},
+        {ErrorCode::DFS_PERMISSION_DENIED, "DFS_PERMISSION_DENIED"},
+        {ErrorCode::DFS_STALE_HANDLE, "DFS_STALE_HANDLE"},
+        {ErrorCode::DFS_PARTIAL_WRITE, "DFS_PARTIAL_WRITE"},
+        {ErrorCode::TENANT_QUOTA_EXCEEDED, "TENANT_QUOTA_EXCEEDED"},
+        {ErrorCode::TENANT_NOT_REGISTERED, "TENANT_NOT_REGISTERED"},
+        {ErrorCode::TENANT_NOT_EMPTY, "TENANT_NOT_EMPTY"}};
 
     auto it = errorCodeMap.find(errorCode);
     static const std::string unknownError = "UNKNOWN_ERROR";
diff --git a/mooncake-store/src/utils.cpp b/mooncake-store/src/utils.cpp
index d7820e9c..d1f03549 100644
--- a/mooncake-store/src/utils.cpp
+++ b/mooncake-store/src/utils.cpp
@@ -1,6 +1,11 @@
 #include "utils.h"
+#include "mmap_arena.h"
+#include "config.h"
+#include "common.h"
+#include "ub_allocator.h"
 
 #include <Slab.h>
+#include <gflags/gflags.h>
 #include <glog/logging.h>
 #include <ifaddrs.h>
 #include <net/if.h>
@@ -11,6 +16,7 @@
 #include <boost/algorithm/string.hpp>
 
 #include <algorithm>
+#include <atomic>
 #include <memory>
 #include <random>
 #include <cerrno>
@@ -19,12 +25,29 @@
 #include <sys/mman.h>
 #include <numa.h>
 #include <numaif.h>
-#include "config.h"
-#include "common.h"
+#include <mutex>
+
+// Feature flag to enable/disable arena allocator. Disabled by default so the
+// library does not pre-map a large pool unless the operator opts in via gflag
+// or an explicit MC_MMAP_ARENA_POOL_SIZE environment override.
+DEFINE_bool(use_mmap_arena_allocator, false,
+            "Enable the lock-free mmap arena allocator for mmap buffers");
+
+// Arena pool size (default 8GiB when explicitly enabled without an env
+// override).
+DEFINE_uint64(mmap_arena_pool_size, 8ULL * 1024 * 1024 * 1024,
+              "Arena allocator pool size in bytes");
+#ifdef USE_ASCEND_DIRECT
+#include "acl/acl.h"
+#endif
 #if defined(USE_ASCEND_DIRECT) || defined(USE_UBSHMEM)
 #include "ascend_allocator.h"
 #endif
 
+#ifdef USE_NOF
+#include "spdk/spdk_wrapper.h"
+#endif
+
 #include <ylt/coro_http/coro_http_client.hpp>
 
 namespace mooncake {
@@ -83,7 +106,7 @@ AutoPortBinder::~AutoPortBinder() {
 
 void *allocate_buffer_allocator_memory(size_t total_size,
                                        const std::string &protocol,
-                                       size_t alignment) {
+                                       size_t alignment, bool use_spdk_dma) {
     const size_t default_alignment = facebook::cachelib::Slab::kSize;
     // Ensure total_size is a multiple of alignment
     if (alignment == default_alignment && total_size < alignment) {
@@ -95,29 +118,165 @@ void *allocate_buffer_allocator_memory(size_t total_size,
         return ascend_allocate_memory(total_size, protocol);
     }
 #endif
-
+#if defined(USE_UB)
+    if (protocol == "ub") {
+        return mooncake::ub_allocate_memory(alignment, total_size);
+    }
+#endif
+#ifdef USE_NOF
+    if (use_spdk_dma && total_size > 0) {
+        return mooncake::SpdkWrapper::GetInstance().Alloc(total_size, alignment,
+                                                          -1);
+    }
+#endif
     // Allocate aligned memory
     return aligned_alloc(alignment, total_size);
 }
 
+// Global arena instance (lazy initialization)
+static std::unique_ptr<MmapArena> g_mmap_arena;
+static std::once_flag g_arena_init_flag;
+static std::atomic<uint64_t> g_arena_oom_fallback_count{0};
+static std::atomic<uint64_t> g_arena_noop_free_count{0};
+
+static void initializeGlobalArena() {
+    const std::string env_pool_size =
+        GetEnvStringOr("MC_MMAP_ARENA_POOL_SIZE", "");
+    // Allow env var to override the gflag (useful when loaded as .so from
+    // Python). An explicit pool-size env var is also treated as an opt-in,
+    // because pybind11 users cannot easily pass gflags.
+    const std::string env_disable = GetEnvStringOr("MC_DISABLE_MMAP_ARENA", "");
+    const std::optional<bool> disable_override = string_to_bool(env_disable);
+    if (!env_disable.empty() && !disable_override.has_value()) {
+        LOG(WARNING) << "Ignoring invalid MC_DISABLE_MMAP_ARENA='"
+                     << env_disable
+                     << "'; accepted values: 1/0, true/false, yes/no, on/off";
+    }
+    const bool arena_requested =
+        FLAGS_use_mmap_arena_allocator || !env_pool_size.empty();
+    const bool arena_disabled = disable_override.value_or(false);
+    if (!arena_requested || arena_disabled) {
+        LOG(INFO) << "=== ARENA ALLOCATOR DISABLED ===";
+        if (arena_disabled) {
+            LOG(INFO) << "MC_DISABLE_MMAP_ARENA=" << env_disable
+                      << " forces direct mmap()";
+        } else {
+            LOG(INFO) << "Arena is opt-in; set --use_mmap_arena_allocator or "
+                         "MC_MMAP_ARENA_POOL_SIZE to enable it";
+        }
+        return;
+    }
+
+    g_mmap_arena = std::make_unique<MmapArena>();
+    // Keep arena init consistent with the direct-mmap path:
+    //   MC_STORE_USE_HUGEPAGE=1  -> strict: hard fail if hugepage mmap fails
+    //   unset                    -> permissive: try hugepages, retry on
+    //                               regular pages if HugeTLB is unavailable
+    // This preserves both pre-existing contracts and avoids surprising
+    // operators with a silent hugepage downgrade.
+    const bool hugepages_explicitly_requested =
+        get_hugepage_size_from_env() > 0;
+
+    // Allow env var override since gflags cannot be set from Python (pybind11).
+    // Supports human-readable sizes via string_to_byte_size(): "20gb", "16GB",
+    // etc.
+    uint64_t arena_pool_size = FLAGS_mmap_arena_pool_size;
+    if (!env_pool_size.empty()) {
+        const uint64_t parsed = string_to_byte_size(env_pool_size);
+        if (parsed > 0) {
+            arena_pool_size = parsed;
+            LOG(INFO) << "MC_MMAP_ARENA_POOL_SIZE override: " << env_pool_size
+                      << " (" << byte_size_to_string(arena_pool_size) << ")";
+        } else {
+            LOG(WARNING) << "Invalid MC_MMAP_ARENA_POOL_SIZE='" << env_pool_size
+                         << "', using default "
+                         << byte_size_to_string(FLAGS_mmap_arena_pool_size);
+        }
+    }
+
+    bool success =
+        g_mmap_arena->initialize(arena_pool_size, MmapArena::kMinAlignment,
+                                 !hugepages_explicitly_requested);
+
+    if (success) {
+        auto stats = g_mmap_arena->getStats();
+        LOG(INFO) << "=== ARENA ALLOCATOR ENABLED ===";
+        LOG(INFO) << "Arena pool size: " << (stats.pool_size / BYTES_PER_GIB)
+                  << " GiB";
+        LOG(INFO) << "Using lock-free atomic bump allocation";
+    } else {
+        LOG(ERROR) << "=== ARENA INITIALIZATION FAILED ===";
+        LOG(ERROR) << "Falling back to traditional mmap()";
+        if (hugepages_explicitly_requested) {
+            LOG(ERROR) << "MC_STORE_USE_HUGEPAGE is set, so the fallback path "
+                          "will also require hugepages";
+        }
+        LOG(ERROR) << "Arena initialization is only attempted once per "
+                      "process; restart after fixing the environment if you "
+                      "want to retry arena bring-up";
+        g_mmap_arena.reset();
+    }
+}
+
+// Compute the mmap/munmap size for the fallback (non-arena) path.
+// Used by both allocate_buffer_mmap_memory and free_buffer_mmap_memory
+// so they agree on the mapping size.
+static inline size_t mmap_map_size(size_t total_size, size_t hugepage_size) {
+    const size_t page_size =
+        hugepage_size > 0 ? hugepage_size : static_cast<size_t>(getpagesize());
+    return align_up(total_size, page_size);
+}
+
 void *allocate_buffer_mmap_memory(size_t total_size, size_t alignment) {
     if (total_size == 0) {
-        LOG(ERROR) << "Total size must be greater than 0 for hugepage mmap";
+        LOG(ERROR) << "Total size must be greater than 0 for mmap";
         return nullptr;
     }
 
+    // Initialize arena on first call
+    std::call_once(g_arena_init_flag, initializeGlobalArena);
+
+    // Try arena allocation first (if enabled).
+    // Forward caller's alignment so the arena honors the contract.
+    if (g_mmap_arena && g_mmap_arena->isInitialized()) {
+        void *ptr = g_mmap_arena->allocate(total_size, alignment);
+        if (ptr != nullptr) {
+            VLOG(1) << "Allocated " << total_size << " bytes from arena at "
+                    << ptr;
+            return ptr;
+        }
+        // Arena OOM, fall through to traditional mmap
+        const uint64_t fallback_count =
+            g_arena_oom_fallback_count.fetch_add(1, std::memory_order_relaxed) +
+            1;
+        LOG_FIRST_N(WARNING, 3)
+            << "Arena OOM, falling back to mmap() for size=" << total_size
+            << " (count=" << fallback_count << ")"
+            << " (further warnings suppressed)";
+    }
+
+    // Traditional mmap allocation (fallback or arena disabled).
     unsigned int flags = MAP_PRIVATE | MAP_ANONYMOUS | MAP_POPULATE;
-    const size_t effective_alignment =
-        std::max(alignment, get_hugepage_size_from_env(&flags));
-    const size_t map_size = align_up(total_size, effective_alignment);
+    const size_t hugepage_size = get_hugepage_size_from_env(&flags);
+    const size_t map_size = mmap_map_size(total_size, hugepage_size);
+    const size_t guaranteed_alignment =
+        hugepage_size > 0 ? hugepage_size : static_cast<size_t>(getpagesize());
+    if (alignment > guaranteed_alignment) {
+        LOG_FIRST_N(WARNING, 3)
+            << "Fallback mmap cannot honor alignment=" << alignment
+            << " (guaranteed=" << guaranteed_alignment
+            << "); pointer may be under-aligned"
+            << " (further warnings suppressed)";
+    }
 
     void *ptr = mmap(nullptr, map_size, PROT_READ | PROT_WRITE, flags, -1, 0);
     if (ptr == MAP_FAILED) {
-        LOG(ERROR) << "Hugepage mmap failed, size=" << map_size
-                   << ", errno=" << errno << " (" << strerror(errno) << ")";
+        LOG(ERROR) << "mmap failed, size=" << map_size << ", errno=" << errno
+                   << " (" << strerror(errno) << ")";
         return nullptr;
     }
 
+    VLOG(1) << "Allocated " << total_size << " bytes via mmap() at " << ptr;
     return ptr;
 }
 
@@ -126,10 +285,30 @@ void free_buffer_mmap_memory(void *ptr, size_t total_size) {
         return;
     }
 
-    const size_t map_size = align_up(total_size, get_hugepage_size_from_env());
+    // Check if pointer belongs to global arena
+    std::call_once(g_arena_init_flag, initializeGlobalArena);
+
+    if (g_mmap_arena && g_mmap_arena->owns(ptr)) {
+        const uint64_t noop_free_count =
+            g_arena_noop_free_count.fetch_add(1, std::memory_order_relaxed) + 1;
+        LOG_FIRST_N(WARNING, 3)
+            << "free_buffer_mmap_memory() does not individually release "
+               "arena-owned pointer "
+            << ptr << "; the global arena releases its pool at process shutdown"
+            << " (count=" << noop_free_count << ")"
+            << " (further warnings suppressed)";
+        return;
+    }
+
+    // Direct mmap allocation - safe to unmap
+    const size_t map_size =
+        mmap_map_size(total_size, get_hugepage_size_from_env());
     if (munmap(ptr, map_size) != 0) {
         LOG(ERROR) << "munmap hugepage failed, size=" << map_size
                    << ", errno=" << errno << " (" << strerror(errno) << ")";
+    } else {
+        VLOG(1) << "Freed direct mmap allocation at " << ptr
+                << ", size=" << map_size;
     }
 }
 
@@ -198,7 +377,12 @@ void free_memory(const std::string &protocol, void *ptr) {
         return ascend_free_memory(protocol, ptr);
     }
 #endif
-
+#if defined(USE_UB)
+    if (protocol == "ub") {
+        mooncake::ub_free_memory(ptr);
+        return;
+    }
+#endif
     free(ptr);
 }
 
@@ -375,8 +559,9 @@ static std::string SanitizeKey(const std::string &key) {
 
     for (char c : key) {
         // Replace invalid characters with underscore
-        sanitized_key.push_back(
-            kInvalidChars.find(c) != std::string_view::npos ? '_' : c);
+        const bool invalid =
+            c == '\0' || kInvalidChars.find(c) != std::string_view::npos;
+        sanitized_key.push_back(invalid ? '_' : c);
     }
     return sanitized_key;
 }
diff --git a/mooncake-store/tests/CMakeLists.txt b/mooncake-store/tests/CMakeLists.txt
index 5eb6e5ad..0558cf87 100644
--- a/mooncake-store/tests/CMakeLists.txt
+++ b/mooncake-store/tests/CMakeLists.txt
@@ -36,35 +36,52 @@ endfunction()
 add_store_test(buffer_allocator_test buffer_allocator_test.cpp)
 add_store_test(allocation_strategy_test allocation_strategy_test.cpp)
 add_store_test(eviction_strategy_test eviction_strategy_test.cpp)
+add_store_test(deadline_scheduler_test deadline_scheduler_test.cpp)
 add_store_test(master_service_test master_service_test.cpp)
+add_store_test(master_service_tenant_quota_test
+               master_service_tenant_quota_test.cpp)
 add_store_test(batch_remove_test batch_remove_test.cpp)
 add_store_test(master_service_ssd_test master_service_ssd_test.cpp)
 add_store_test(offload_on_evict_test offload_on_evict_test.cpp)
+add_store_test(promotion_on_hit_test promotion_on_hit_test.cpp)
+add_store_test(file_storage_promotion_test file_storage_promotion_test.cpp)
 add_store_test(master_service_ssd_test_for_snapshot
                ha/snapshot/master_service_ssd_test_for_snapshot.cpp)
+add_store_test(master_service_promotion_test_for_snapshot
+               ha/snapshot/master_service_promotion_test_for_snapshot.cpp)
+if(USE_NOF)
+  add_store_test(nof_heartbeat_test nof_heartbeat_test.cpp)
+endif()
 add_store_test(client_integration_test client_integration_test.cpp)
+add_store_test(rpc_timeout_test rpc_timeout_test.cpp)
 if(USE_CXL)
   add_store_test(cxl_client_integration_test cxl_client_integration_test.cpp)
 endif()
 add_store_test(master_metrics_test master_metrics_test.cpp)
+add_store_test(master_admin_server_test master_admin_server_test.cpp)
 add_store_test(posix_file_test posix_file_test.cpp)
 add_store_test(thread_pool_test thread_pool_test.cpp)
 add_store_test(transfer_task_test transfer_task_test.cpp)
+add_store_test(tenant_quota_test tenant_quota_test.cpp)
 add_store_test(segment_test segment_test.cpp)
 add_store_test(offset_allocator_test offset_allocator_test.cpp)
 add_store_test(utils_test utils_test.cpp)
 add_store_test(client_buffer_test client_buffer_test.cpp)
 add_store_test(client_local_hot_cache_test client_local_hot_cache_test.cpp)
+add_store_test(client_tcp_local_memcpy_test client_tcp_local_memcpy_test.cpp)
 add_store_test(pybind_client_test pybind_client_test.cpp)
 add_store_test(ipv6_client_test ipv6_client_test.cpp)
 add_store_test(client_metrics_test client_metrics_test.cpp)
 add_store_test(ssd_metrics_test ssd_metrics_test.cpp)
 add_store_test(serializer_test serializer_test.cpp)
-add_store_test(embedded_snapshot_catalog_store_test
-               ha/snapshot/catalog/backends/embedded/embedded_snapshot_catalog_store_test.cpp)
+add_store_test(
+  embedded_snapshot_catalog_store_test
+  ha/snapshot/catalog/backends/embedded/embedded_snapshot_catalog_store_test.cpp
+)
 add_store_test(zstd_util_test zstd_util_test.cpp)
-add_store_test(local_file_snapshot_object_store_test
-               ha/snapshot/object/backends/local/local_file_snapshot_object_store_test.cpp)
+add_store_test(
+  local_file_snapshot_object_store_test
+  ha/snapshot/object/backends/local/local_file_snapshot_object_store_test.cpp)
 add_store_test(file_util_test file_util_test.cpp)
 add_store_test(snapshot_child_process_test
                ha/snapshot/snapshot_child_process_test.cpp)
@@ -79,35 +96,72 @@ add_store_test(task_executor_test task_executor_test.cpp)
 add_store_test(task_integration_test task_integration_test.cpp)
 add_store_test(dummy_client_get_buffer_test dummy_client_get_buffer_test.cpp)
 add_store_test(health_check_test health_check_test.cpp)
+add_store_test(http_metadata_server_test http_metadata_server_test.cpp)
+add_store_test(mmap_arena_test mmap_arena_test.cpp)
+add_store_test(mmap_arena_fallback_test mmap_arena_fallback_test.cpp)
+add_store_test(object_data_type_test object_data_type_test.cpp)
 add_subdirectory(e2e)
 
 add_executable(high_availability_test ha/leadership/high_availability_test.cpp)
-target_include_directories(high_availability_test PRIVATE
-                           ${CMAKE_CURRENT_SOURCE_DIR})
+target_include_directories(high_availability_test
+                           PRIVATE ${CMAKE_CURRENT_SOURCE_DIR})
 if(STORE_USE_REDIS)
-    target_sources(high_availability_test PRIVATE
-        ha/leadership/backends/redis/high_availability_redis_test.cpp)
-    target_include_directories(high_availability_test PRIVATE
-        ${MOONCAKE_STORE_HIREDIS_INCLUDE_DIR})
-    target_link_libraries(high_availability_test PRIVATE
-        ${MOONCAKE_STORE_HIREDIS_LIBRARY})
+  target_sources(
+    high_availability_test
+    PRIVATE ha/leadership/backends/redis/high_availability_redis_test.cpp)
+  target_include_directories(high_availability_test
+                             PRIVATE ${MOONCAKE_STORE_HIREDIS_INCLUDE_DIR})
+  target_link_libraries(high_availability_test
+                        PRIVATE ${MOONCAKE_STORE_HIREDIS_LIBRARY})
 
-    add_executable(redis_snapshot_catalog_store_test
-                   ha/snapshot/catalog/backends/redis/redis_snapshot_catalog_store_test.cpp)
-    target_include_directories(redis_snapshot_catalog_store_test PRIVATE
-        ${CMAKE_CURRENT_SOURCE_DIR}
-        ${MOONCAKE_STORE_HIREDIS_INCLUDE_DIR})
-    target_link_libraries(redis_snapshot_catalog_store_test
-        PUBLIC mooncake_store transfer_engine cachelib_memory_allocator
-               ${ETCD_WRAPPER_LIB} glog gflags ibverbs gtest pthread
-        PRIVATE ${MOONCAKE_STORE_HIREDIS_LIBRARY})
-    add_test(NAME redis_snapshot_catalog_store_test
-             COMMAND redis_snapshot_catalog_store_test)
+  add_executable(
+    redis_snapshot_catalog_store_test
+    ha/snapshot/catalog/backends/redis/redis_snapshot_catalog_store_test.cpp)
+  target_include_directories(
+    redis_snapshot_catalog_store_test
+    PRIVATE ${CMAKE_CURRENT_SOURCE_DIR} ${MOONCAKE_STORE_HIREDIS_INCLUDE_DIR})
+  target_link_libraries(
+    redis_snapshot_catalog_store_test
+    PUBLIC mooncake_store
+           transfer_engine
+           cachelib_memory_allocator
+           ${ETCD_WRAPPER_LIB}
+           glog
+           gflags
+           ibverbs
+           gtest
+           pthread
+    PRIVATE ${MOONCAKE_STORE_HIREDIS_LIBRARY})
+  add_test(NAME redis_snapshot_catalog_store_test
+           COMMAND redis_snapshot_catalog_store_test)
 endif()
-target_link_libraries(high_availability_test PUBLIC mooncake_store transfer_engine cachelib_memory_allocator ${ETCD_WRAPPER_LIB} glog gtest gtest_main pthread)
-if (STORE_USE_ETCD OR STORE_USE_REDIS)
-    add_test(NAME high_availability_test COMMAND high_availability_test)
+if(STORE_USE_K8S_LEASE)
+  target_sources(
+    high_availability_test
+    PRIVATE ha/leadership/backends/k8s/high_availability_k8s_test.cpp)
 endif()
+target_link_libraries(
+  high_availability_test
+  PUBLIC mooncake_store
+         transfer_engine
+         cachelib_memory_allocator
+         ${ETCD_WRAPPER_LIB}
+         glog
+         gtest
+         gtest_main
+         pthread)
+if(STORE_USE_K8S_LEASE)
+  target_link_libraries(high_availability_test PRIVATE ${K8S_LEASE_WRAPPER_LIB})
+endif()
+if(STORE_USE_ETCD
+   OR STORE_USE_REDIS
+   OR STORE_USE_K8S_LEASE)
+  add_test(NAME high_availability_test COMMAND high_availability_test)
+endif()
+add_ha_test(ha_backend_availability_test
+            ha/leadership/ha_backend_availability_test.cpp)
+add_ha_test(leader_label_reconciler_test
+            ha/leadership/leader_label_reconciler_test.cpp)
 
 add_executable(stress_workload_test stress_workload_test.cpp)
 target_link_libraries(
@@ -124,10 +178,8 @@ target_link_libraries(
 # HA tests
 add_ha_test(standby_state_machine_test
             ha/standby/standby_state_machine_test.cpp)
-add_ha_test(ha_metric_manager_test
-            ha/standby/ha_metric_manager_test.cpp)
-add_ha_test(hot_standby_service_test
-            ha/standby/hot_standby_service_test.cpp)
+add_ha_test(ha_metric_manager_test ha/standby/ha_metric_manager_test.cpp)
+add_ha_test(hot_standby_service_test ha/standby/hot_standby_service_test.cpp)
 add_ha_test(hot_standby_snapshot_bootstrap_test
             ha/standby/hot_standby_snapshot_bootstrap_test.cpp)
 add_ha_test(oplog_applier_test ha/oplog/oplog_applier_test.cpp)
@@ -144,9 +196,8 @@ add_ha_test(catalog_backed_snapshot_provider_test
 if(STORE_USE_REDIS)
   foreach(target catalog_backed_snapshot_provider_test
                  hot_standby_snapshot_bootstrap_test)
-    target_include_directories(${target} PRIVATE
-                               ${MOONCAKE_STORE_HIREDIS_INCLUDE_DIR})
-    target_link_libraries(${target}
-                          PRIVATE ${MOONCAKE_STORE_HIREDIS_LIBRARY})
+    target_include_directories(${target}
+                               PRIVATE ${MOONCAKE_STORE_HIREDIS_INCLUDE_DIR})
+    target_link_libraries(${target} PRIVATE ${MOONCAKE_STORE_HIREDIS_LIBRARY})
   endforeach()
 endif()
diff --git a/mooncake-store/tests/allocation_strategy_test.cpp b/mooncake-store/tests/allocation_strategy_test.cpp
index 293e57f0..71292444 100644
--- a/mooncake-store/tests/allocation_strategy_test.cpp
+++ b/mooncake-store/tests/allocation_strategy_test.cpp
@@ -89,6 +89,9 @@ INSTANTIATE_TEST_SUITE_P(
             case AllocationStrategyType::FREE_RATIO_FIRST:
                 strategy_str = "FreeRatioFirst";
                 break;
+            case AllocationStrategyType::SSD_FREE_RATIO_FIRST:
+                strategy_str = "SsdFreeRatioFirst";
+                break;
             default:
                 strategy_str = "Unknown";
         }
@@ -752,4 +755,255 @@ TEST_F(AllocationStrategyTest, PerformanceTest) {
 // public API. The functionality is now encapsulated within the Allocate()
 // method.
 
+// Mock SsdMetricsProvider for testing SSD-aware allocation
+class MockSsdMetricsProvider : public SsdMetricsProvider {
+   public:
+    std::unordered_map<std::string, int64_t> total_capacity;
+    std::unordered_map<std::string, int64_t> used_bytes;
+
+    int64_t getSsdTotalCapacity(const std::string& name) const override {
+        auto it = total_capacity.find(name);
+        return it != total_capacity.end() ? it->second : 0;
+    }
+
+    int64_t getSsdUsedBytes(const std::string& name) const override {
+        auto it = used_bytes.find(name);
+        return it != used_bytes.end() ? it->second : 0;
+    }
+};
+
+TEST_F(AllocationStrategyTest, SsdFreeRatioFirstChoosesHighestFreeRatio) {
+    auto ssd_strategy = std::make_unique<SsdFreeRatioFirstAllocationStrategy>();
+
+    const int kNumSegments = 3;
+    const size_t kSegmentSize = 64 * MiB;
+
+    AllocatorManager allocator_manager;
+    for (int i = 0; i < kNumSegments; i++) {
+        const auto name = std::to_string(i) + "-segment";
+        allocator_manager.addAllocator(
+            name,
+            std::make_shared<OffsetBufferAllocator>(
+                name, 0x100000000ULL + i * kSegmentSize, kSegmentSize, name));
+    }
+
+    // SSD free ratios: segment 0 = 20%, segment 1 = 60%, segment 2 = 90%
+    MockSsdMetricsProvider ssd_provider;
+    for (int i = 0; i < kNumSegments; i++) {
+        const auto name = std::to_string(i) + "-segment";
+        ssd_provider.total_capacity[name] = 1000 * MiB;
+    }
+    ssd_provider.used_bytes["0-segment"] = 800 * MiB;  // 20% free
+    ssd_provider.used_bytes["1-segment"] = 400 * MiB;  // 60% free
+    ssd_provider.used_bytes["2-segment"] = 100 * MiB;  // 90% free
+
+    auto result =
+        ssd_strategy->Allocate(allocator_manager, 64 * 1024, 1, {}, {},
+                               ReplicaType::MEMORY, &ssd_provider);
+    ASSERT_TRUE(result.has_value());
+    ASSERT_EQ(result.value().size(), 1u);
+
+    const auto& replica = result.value()[0];
+    auto descriptor = replica.get_descriptor();
+    ASSERT_TRUE(descriptor.is_memory_replica());
+    const auto& mem_desc = descriptor.get_memory_descriptor();
+    EXPECT_EQ(mem_desc.buffer_descriptor.transport_endpoint_, "2-segment");
+}
+
+// Test that SsdFreeRatioFirstAllocationStrategy works without an
+// SsdMetricsProvider (delegates to base class random allocation).
+TEST_F(AllocationStrategyTest,
+       SsdFreeRatioFirstWithoutMetricsProviderAllocates) {
+    auto ssd_strategy = std::make_unique<SsdFreeRatioFirstAllocationStrategy>();
+
+    auto allocator = std::make_shared<OffsetBufferAllocator>(
+        "segment1", 0x100000000ULL, 64 * MiB, "segment1");
+    AllocatorManager allocator_manager;
+    allocator_manager.addAllocator("segment1", allocator);
+
+    auto result = ssd_strategy->Allocate(allocator_manager, 1024, 1, {}, {});
+    ASSERT_TRUE(result.has_value());
+    EXPECT_EQ(result.value().size(), 1u);
+}
+
+// Test that SsdFreeRatioFirstAllocationStrategy skips excluded segments even
+// when they have the highest SSD free ratio.
+TEST_F(AllocationStrategyTest, SsdFreeRatioFirstExcludedSegmentsSkipped) {
+    auto ssd_strategy = std::make_unique<SsdFreeRatioFirstAllocationStrategy>();
+
+    const int kNumSegments = 3;
+    const size_t kSegmentSize = 64 * MiB;
+
+    AllocatorManager allocator_manager;
+    for (int i = 0; i < kNumSegments; i++) {
+        const auto name = std::to_string(i) + "-segment";
+        allocator_manager.addAllocator(
+            name,
+            std::make_shared<OffsetBufferAllocator>(
+                name, 0x100000000ULL + i * kSegmentSize, kSegmentSize, name));
+    }
+
+    // SSD free ratios: 0-segment=95%, 1-segment=50%, 2-segment=30%
+    MockSsdMetricsProvider ssd_provider;
+    for (int i = 0; i < kNumSegments; i++) {
+        const auto name = std::to_string(i) + "-segment";
+        ssd_provider.total_capacity[name] = 1000 * MiB;
+    }
+    ssd_provider.used_bytes["0-segment"] = 50 * MiB;   // 95% free (highest)
+    ssd_provider.used_bytes["1-segment"] = 500 * MiB;  // 50% free
+    ssd_provider.used_bytes["2-segment"] = 700 * MiB;  // 30% free
+
+    // Exclude 0-segment which has the highest SSD free ratio
+    std::set<std::string> excluded = {"0-segment"};
+
+    auto result =
+        ssd_strategy->Allocate(allocator_manager, 64 * 1024, 1, {}, excluded,
+                               ReplicaType::MEMORY, &ssd_provider);
+    ASSERT_TRUE(result.has_value());
+    ASSERT_EQ(result.value().size(), 1u);
+
+    const auto& replica = result.value()[0];
+    auto descriptor = replica.get_descriptor();
+    ASSERT_TRUE(descriptor.is_memory_replica());
+    const auto& mem_desc = descriptor.get_memory_descriptor();
+    // Must NOT be allocated to the excluded segment despite its high free ratio
+    EXPECT_NE(mem_desc.buffer_descriptor.transport_endpoint_, "0-segment");
+}
+
+// Test that ssd_used_bytes > ssd_total_capacity is clamped to 0% free,
+// so the other segment (with normal usage) is preferred.
+TEST_F(AllocationStrategyTest, SsdFreeRatioFirstUsedExceedsTotalIsClamped) {
+    auto ssd_strategy = std::make_unique<SsdFreeRatioFirstAllocationStrategy>();
+
+    const size_t kSegmentSize = 64 * MiB;
+
+    AllocatorManager allocator_manager;
+    allocator_manager.addAllocator(
+        "0-segment",
+        std::make_shared<OffsetBufferAllocator>("0-segment", 0x100000000ULL,
+                                                kSegmentSize, "0-segment"));
+    allocator_manager.addAllocator(
+        "1-segment", std::make_shared<OffsetBufferAllocator>(
+                         "1-segment", 0x100000000ULL + kSegmentSize,
+                         kSegmentSize, "1-segment"));
+
+    // 0-segment: used exceeds total (concurrent drift) → clamped to 0% free
+    // 1-segment: used=100, total=1000 → 90% free
+    MockSsdMetricsProvider ssd_provider;
+    ssd_provider.total_capacity["0-segment"] = 1000 * MiB;
+    ssd_provider.total_capacity["1-segment"] = 1000 * MiB;
+    ssd_provider.used_bytes["0-segment"] = 1500 * MiB;  // exceeds total
+    ssd_provider.used_bytes["1-segment"] = 100 * MiB;   // 90% free
+
+    auto result =
+        ssd_strategy->Allocate(allocator_manager, 64 * 1024, 1, {}, {},
+                               ReplicaType::MEMORY, &ssd_provider);
+    ASSERT_TRUE(result.has_value());
+    ASSERT_EQ(result.value().size(), 1u);
+
+    const auto& replica = result.value()[0];
+    auto descriptor = replica.get_descriptor();
+    ASSERT_TRUE(descriptor.is_memory_replica());
+    const auto& mem_desc = descriptor.get_memory_descriptor();
+    // 0-segment is clamped to 0% free; 1-segment is 90% free → pick 1-segment
+    EXPECT_EQ(mem_desc.buffer_descriptor.transport_endpoint_, "1-segment");
+}
+
+// Strategy-level performance comparison: Random vs SsdFreeRatioFirst.
+// Uses OffsetBufferAllocator (real physical memory via aligned_alloc) and
+// MockSsdMetricsProvider. Measures pure strategy overhead without MasterService
+// locking, reflecting only the allocation algorithm cost.
+TEST_F(AllocationStrategyTest, SsdFreeRatioFirstVsRandomStrategyPerformance) {
+    constexpr size_t kNumSegments = 64;
+    constexpr size_t kSegmentSize = 8 * MiB;   // 64 * 8MB = 512MB physical
+    constexpr size_t kAllocSize = 128 * 1024;  // 128KB per allocation
+    constexpr int kWarmupRounds = 200;
+    constexpr int kBenchmarkRounds = 2000;
+
+    // Build AllocatorManager: 64 segments, each 8 MiB
+    AllocatorManager allocator_manager;
+    for (size_t i = 0; i < kNumSegments; i++) {
+        const auto name = "perf_seg_" + std::to_string(i);
+        allocator_manager.addAllocator(
+            name,
+            std::make_shared<OffsetBufferAllocator>(
+                name, 0x200000000ULL + i * kSegmentSize, kSegmentSize, name));
+    }
+
+    // SSD free ratios vary uniformly from ~10% to ~90% across segments
+    MockSsdMetricsProvider ssd_provider;
+    for (size_t i = 0; i < kNumSegments; i++) {
+        const auto name = "perf_seg_" + std::to_string(i);
+        ssd_provider.total_capacity[name] = 1000 * MiB;
+        // used: 100 MiB (10% used) to 900 MiB (90% used)
+        ssd_provider.used_bytes[name] =
+            static_cast<int64_t>((100 + i * 800 / kNumSegments)) * MiB;
+    }
+
+    // -------- Random strategy --------
+    auto random_strategy = std::make_unique<RandomAllocationStrategy>();
+    {
+        // Warmup (results discarded)
+        for (int i = 0; i < kWarmupRounds; i++) {
+            (void)random_strategy->Allocate(allocator_manager, kAllocSize);
+        }
+    }
+    std::vector<std::vector<Replica>> random_replicas;
+    random_replicas.reserve(kBenchmarkRounds);
+
+    auto t_rand_start = std::chrono::steady_clock::now();
+    for (int i = 0; i < kBenchmarkRounds; i++) {
+        auto r = random_strategy->Allocate(allocator_manager, kAllocSize);
+        ASSERT_TRUE(r.has_value());
+        random_replicas.emplace_back(std::move(r.value()));
+    }
+    auto rand_us = std::chrono::duration_cast<std::chrono::microseconds>(
+        std::chrono::steady_clock::now() - t_rand_start);
+    random_replicas.clear();  // deallocate so SSD strategy starts fresh
+
+    // -------- SsdFreeRatioFirst strategy --------
+    auto ssd_strategy = std::make_unique<SsdFreeRatioFirstAllocationStrategy>();
+    {
+        // Warmup
+        for (int i = 0; i < kWarmupRounds; i++) {
+            (void)ssd_strategy->Allocate(allocator_manager, kAllocSize, 1, {},
+                                         {}, ReplicaType::MEMORY,
+                                         &ssd_provider);
+        }
+    }
+    std::vector<std::vector<Replica>> ssd_replicas;
+    ssd_replicas.reserve(kBenchmarkRounds);
+
+    auto t_ssd_start = std::chrono::steady_clock::now();
+    for (int i = 0; i < kBenchmarkRounds; i++) {
+        auto r = ssd_strategy->Allocate(allocator_manager, kAllocSize, 1, {},
+                                        {}, ReplicaType::MEMORY, &ssd_provider);
+        ASSERT_TRUE(r.has_value());
+        ssd_replicas.emplace_back(std::move(r.value()));
+    }
+    auto ssd_us = std::chrono::duration_cast<std::chrono::microseconds>(
+        std::chrono::steady_clock::now() - t_ssd_start);
+    ssd_replicas.clear();
+
+    double rand_us_per_op =
+        static_cast<double>(rand_us.count()) / kBenchmarkRounds;
+    double ssd_us_per_op =
+        static_cast<double>(ssd_us.count()) / kBenchmarkRounds;
+    double overhead_ratio =
+        static_cast<double>(ssd_us.count()) / rand_us.count();
+
+    std::cout
+        << "\n=== Strategy-Level Performance: Random vs SsdFreeRatioFirst ===\n"
+        << "Segments: " << kNumSegments
+        << " | Alloc size: " << (kAllocSize / 1024) << " KB"
+        << " | Rounds: " << kBenchmarkRounds << "\n"
+        << "Random:              " << rand_us.count() << " us total  |  "
+        << std::fixed << std::setprecision(3) << rand_us_per_op << " us/op\n"
+        << "SsdFreeRatioFirst:   " << ssd_us.count() << " us total  |  "
+        << ssd_us_per_op << " us/op\n"
+        << "Overhead ratio:      " << std::setprecision(2) << overhead_ratio
+        << "x  (" << std::setprecision(1) << (overhead_ratio - 1.0) * 100.0
+        << "% slower)\n\n";
+}
+
 }  // namespace mooncake
diff --git a/mooncake-store/tests/client_integration_test.cpp b/mooncake-store/tests/client_integration_test.cpp
index 42a13256..e30b6df4 100644
--- a/mooncake-store/tests/client_integration_test.cpp
+++ b/mooncake-store/tests/client_integration_test.cpp
@@ -4,6 +4,7 @@
 
 #include <cstdint>
 #include <filesystem>
+#include <functional>
 #include <memory>
 #include <string>
 #include <vector>
@@ -653,6 +654,82 @@ TEST_F(ClientIntegrationTest, BatchPutGetOperations) {
     }
 }
 
+TEST_F(ClientIntegrationTest, BatchPutMixedGroupIdsThroughClient) {
+    auto find_group_id_on_different_shard = [](const std::string& key) {
+        static constexpr size_t kMetadataShardCountForTest = 1024;
+        const size_t key_shard =
+            std::hash<std::string>{}(key) % kMetadataShardCountForTest;
+        for (int i = 0; i < 10000; ++i) {
+            std::string group_id = key + "_group_" + std::to_string(i);
+            if (std::hash<std::string>{}(group_id) %
+                    kMetadataShardCountForTest !=
+                key_shard) {
+                return group_id;
+            }
+        }
+        return key + "_fallback_group";
+    };
+
+    const std::vector<std::string> keys = {
+        "client_batch_grouped_a",
+        "client_batch_ungrouped",
+        "client_batch_grouped_b",
+    };
+    const std::vector<std::string> values = {
+        "grouped-value-a",
+        "ungrouped-value",
+        "grouped-value-b",
+    };
+
+    std::vector<void*> source_buffers;
+    std::vector<std::vector<Slice>> batched_slices;
+    source_buffers.reserve(values.size());
+    batched_slices.reserve(values.size());
+    for (const auto& value : values) {
+        void* buffer = client_buffer_allocator_->allocate(value.size());
+        memcpy(buffer, value.data(), value.size());
+        source_buffers.push_back(buffer);
+        batched_slices.push_back({Slice{buffer, value.size()}});
+    }
+
+    ReplicateConfig config;
+    config.replica_num = 1;
+    config.group_ids =
+        std::vector<std::string>{find_group_id_on_different_shard(keys[0]), "",
+                                 find_group_id_on_different_shard(keys[2])};
+
+    auto put_results = test_client_->BatchPut(keys, batched_slices, config);
+    ASSERT_EQ(put_results.size(), keys.size());
+    for (const auto& result : put_results) {
+        ASSERT_TRUE(result.has_value())
+            << "BatchPut failed: " << toString(result.error());
+    }
+
+    for (size_t i = 0; i < values.size(); ++i) {
+        client_buffer_allocator_->deallocate(source_buffers[i],
+                                             values[i].size());
+    }
+
+    for (size_t i = 0; i < keys.size(); ++i) {
+        void* target_buffer =
+            client_buffer_allocator_->allocate(values[i].size());
+        std::vector<Slice> slices{Slice{target_buffer, values[i].size()}};
+        auto get_result = test_client_->Get(keys[i], slices);
+        ASSERT_TRUE(get_result.has_value())
+            << "Get failed: " << toString(get_result.error());
+        EXPECT_EQ(slices[0].size, values[i].size());
+        EXPECT_EQ(memcmp(slices[0].ptr, values[i].data(), values[i].size()), 0);
+        client_buffer_allocator_->deallocate(target_buffer, values[i].size());
+    }
+
+    auto remove_results = test_client_->BatchRemove(keys, /*force=*/true);
+    ASSERT_EQ(remove_results.size(), keys.size());
+    for (const auto& result : remove_results) {
+        ASSERT_TRUE(result.has_value())
+            << "BatchRemove failed: " << toString(result.error());
+    }
+}
+
 // Test batch IsExist operations through the client
 TEST_F(ClientIntegrationTest, BatchIsExistOperations) {
     int batch_size = 50;
@@ -1734,6 +1811,40 @@ TEST_F(ClientIntegrationTest, BatchUpsertMixed) {
     }
 }
 
+TEST_F(ClientIntegrationTest, MountSegmentAndGetIdAndUnmountSegmentById) {
+    // Allocate a small buffer for this test
+    size_t test_size = 16 * 1024 * 1024;  // 16 MB
+    void* test_buffer = allocate_buffer_allocator_memory(test_size);
+    ASSERT_NE(test_buffer, nullptr);
+
+    // Test MountSegmentAndGetId returns a valid UUID
+    auto mount_result = test_client_->MountSegmentAndGetId(
+        test_buffer, test_size, FLAGS_protocol);
+    ASSERT_TRUE(mount_result.has_value())
+        << "MountSegmentAndGetId failed: " << toString(mount_result.error());
+    UUID segment_id = mount_result.value();
+    EXPECT_NE(segment_id.first, 0u);
+    EXPECT_NE(segment_id.second, 0u);
+
+    // Test UnmountSegmentById succeeds
+    auto unmount_result = test_client_->UnmountSegmentById(segment_id);
+    EXPECT_TRUE(unmount_result.has_value())
+        << "UnmountSegmentById failed: " << toString(unmount_result.error());
+
+    // Test MountSegment delegates to MountSegmentAndGetId (equivalent behavior)
+    auto mount2 =
+        test_client_->MountSegment(test_buffer, test_size, FLAGS_protocol);
+    EXPECT_TRUE(mount2.has_value())
+        << "MountSegment failed: " << toString(mount2.error());
+
+    // Test UnmountSegment delegates to UnmountSegmentImpl (equivalent behavior)
+    auto unmount2 = test_client_->UnmountSegment(test_buffer, test_size);
+    EXPECT_TRUE(unmount2.has_value())
+        << "UnmountSegment failed: " << toString(unmount2.error());
+
+    free(test_buffer);
+}
+
 }  // namespace testing
 
 }  // namespace mooncake
diff --git a/mooncake-store/tests/client_local_hot_cache_test.cpp b/mooncake-store/tests/client_local_hot_cache_test.cpp
index d16192e1..86927c28 100644
--- a/mooncake-store/tests/client_local_hot_cache_test.cpp
+++ b/mooncake-store/tests/client_local_hot_cache_test.cpp
@@ -476,11 +476,40 @@ TEST_F(LocalHotCacheTest, GetHotKeyProtectsBlockFromReuse) {
     HotMemBlock* block2 = cache.GetHotKey("key2");
     ASSERT_NE(block2, nullptr);
     VerifySliceData(block2, 1024, 'B');
+    cache.ReleaseHotKey("key2");
 
     // key1 should be evicted since we only have 1 block
     EXPECT_FALSE(cache.HasHotKey("key1"));
 }
 
+// Token check and publish must be atomic: a generation bump that races after
+// the token was captured must cancel the fill instead of resurrecting it.
+TEST_F(LocalHotCacheTest, PutHotKeyWithTokenRejectsStaleFill) {
+    const size_t cache_size = 16 * 1024 * 1024;  // 1 block
+    LocalHotCache cache(cache_size);
+
+    // Stale: token captured, then key generation bumped (as Remove would).
+    HotCachePutToken stale = cache.AcquirePutToken("k");
+    cache.BumpKeyGeneration("k");
+    HotMemBlock* b1 = cache.GetFreeBlock();
+    ASSERT_NE(b1, nullptr);
+    b1->key_ = "k";
+    b1->size = 1024;
+    EXPECT_FALSE(cache.PutHotKey(b1, stale));
+    EXPECT_FALSE(cache.HasHotKey("k"));
+    EXPECT_EQ(cache.GetCacheSize(), 1)
+        << "Stale fill must be returned to the pool, not published";
+
+    // Valid: token still current -> published.
+    HotCachePutToken fresh = cache.AcquirePutToken("k");
+    HotMemBlock* b2 = cache.GetFreeBlock();
+    ASSERT_NE(b2, nullptr);
+    b2->key_ = "k";
+    b2->size = 1024;
+    EXPECT_TRUE(cache.PutHotKey(b2, fresh));
+    EXPECT_TRUE(cache.HasHotKey("k"));
+}
+
 // Test LocalHotCacheHandler basic functionality
 TEST_F(LocalHotCacheTest, LocalHotCacheHandlerBasic) {
     const size_t cache_size = 32 * 1024 * 1024;  // 32MB = 2 blocks
@@ -518,6 +547,54 @@ TEST_F(LocalHotCacheTest, LocalHotCacheHandlerNullCache) {
     EXPECT_FALSE(handler.SubmitPutTask("key", slice));
 }
 
+// Regression test: a block recycled from the LRU must still accept objects up
+// to the full block capacity, even if it previously held a smaller object.
+// Before the fix, SubmitPutTask compared the new slice against block->size
+// (the previous object's logical length, which GetFreeBlock does not reset)
+// instead of the block capacity, so a recycled block permanently rejected any
+// object larger than the one it last held.
+TEST_F(LocalHotCacheTest, RecycledBlockAcceptsLargerObject) {
+    const size_t block_size = 64 * 1024;  // 64KB
+    // Single-block cache, so the second Put must recycle the first block.
+    auto cache = std::make_shared<LocalHotCache>(block_size, block_size);
+    ASSERT_EQ(cache->GetCacheSize(), 1u);
+    ASSERT_EQ(cache->GetBlockSize(), block_size);
+
+    LocalHotCacheHandler handler(cache, /*num_worker_threads=*/1,
+                                 /*max_queue_capacity=*/16);
+
+    // Wait until an async Put has been published into the cache.
+    auto wait_for_key = [&](const std::string& key) {
+        for (int i = 0; i < 400; ++i) {  // up to ~2s
+            if (cache->HasHotKey(key)) return true;
+            std::this_thread::sleep_for(std::chrono::milliseconds(5));
+        }
+        return cache->HasHotKey(key);
+    };
+
+    // 1) Cache a small object. The block's logical size shrinks to small_size.
+    const size_t small_size = 1024;  // 1KB
+    Slice small_slice = CreateSlice(small_size, 'S');
+    ASSERT_TRUE(handler.SubmitPutTask("small_key", small_slice));
+    ASSERT_TRUE(wait_for_key("small_key"));
+
+    // 2) Cache a larger object that still fits the block. This forces
+    //    GetFreeBlock() to recycle the block that held "small_key". Before the
+    //    fix this returned false (large_size > stale block->size of 1KB).
+    const size_t large_size = 32 * 1024;  // 32KB: > small_size, <= block_size
+    Slice large_slice = CreateSlice(large_size, 'L');
+    EXPECT_TRUE(handler.SubmitPutTask("large_key", large_slice));
+    ASSERT_TRUE(wait_for_key("large_key"));
+
+    // The larger object is fully cached with the correct length and data, and
+    // the evicted small object is gone.
+    HotMemBlock* block = cache->GetHotKey("large_key");
+    ASSERT_NE(block, nullptr);
+    VerifySliceData(block, large_size, 'L');
+    cache->ReleaseHotKey("large_key");
+    EXPECT_FALSE(cache->HasHotKey("small_key"));
+}
+
 // Test concurrent access to LocalHotCache
 TEST_F(LocalHotCacheTest, ConcurrentAccess) {
     const size_t cache_size = 128 * 1024 * 1024;  // 128MB = 8 blocks
@@ -531,8 +608,7 @@ TEST_F(LocalHotCacheTest, ConcurrentAccess) {
 
     // Each thread puts and gets keys
     for (int t = 0; t < num_threads; ++t) {
-        threads.emplace_back([&cache, t, keys_per_thread, &successful_puts,
-                              &successful_gets]() {
+        threads.emplace_back([&cache, t, &successful_puts, &successful_gets]() {
             for (int i = 0; i < keys_per_thread; ++i) {
                 std::string key =
                     "thread_" + std::to_string(t) + "_key_" + std::to_string(i);
diff --git a/mooncake-store/tests/client_tcp_local_memcpy_test.cpp b/mooncake-store/tests/client_tcp_local_memcpy_test.cpp
new file mode 100644
index 00000000..289ba85a
--- /dev/null
+++ b/mooncake-store/tests/client_tcp_local_memcpy_test.cpp
@@ -0,0 +1,508 @@
+#include <glog/logging.h>
+#include <gtest/gtest.h>
+
+#include <chrono>
+#include <cstdlib>
+#include <cstring>
+#include <memory>
+#include <optional>
+#include <string>
+#include <thread>
+#include <vector>
+
+#include "allocator.h"
+#include "client_service.h"
+#include "test_server_helpers.h"
+#include "types.h"
+#include "utils.h"
+
+namespace mooncake {
+namespace testing {
+
+enum class HandshakeMode { P2P, Metadata };
+
+namespace {
+
+class EnvGuard {
+   public:
+    explicit EnvGuard(const char* key) : key_(key) {
+        if (const char* value = std::getenv(key_)) {
+            old_value_ = value;
+        }
+    }
+
+    ~EnvGuard() {
+        if (old_value_.has_value()) {
+            setenv(key_, old_value_->c_str(), 1);
+        } else {
+            unsetenv(key_);
+        }
+    }
+
+   private:
+    const char* key_;
+    std::optional<std::string> old_value_;
+};
+
+class VLogGuard {
+   public:
+    explicit VLogGuard(int new_v) : old_v_(FLAGS_v) { FLAGS_v = new_v; }
+    ~VLogGuard() { FLAGS_v = old_v_; }
+
+   private:
+    int old_v_;
+};
+
+class StrategyCaptureSink : public google::LogSink {
+   public:
+    void send(google::LogSeverity severity, const char* full_filename,
+              const char* base_filename, int line, const struct ::tm* tm_time,
+              const char* message, size_t message_len) override {
+        (void)severity;
+        (void)full_filename;
+        (void)base_filename;
+        (void)line;
+        (void)tm_time;
+
+        std::string msg(message, message_len);
+        constexpr char kPrefix[] = "Using transfer strategy: ";
+        size_t pos = msg.find(kPrefix);
+        if (pos == std::string::npos) {
+            return;
+        }
+
+        captured_strategy_ = msg.substr(pos + std::strlen(kPrefix));
+        TrimInPlace(captured_strategy_);
+    }
+
+    const std::string& strategy() const { return captured_strategy_; }
+
+   private:
+    static void TrimInPlace(std::string& value) {
+        auto begin = value.find_first_not_of(" \t\n\r");
+        if (begin == std::string::npos) {
+            value.clear();
+            return;
+        }
+        auto end = value.find_last_not_of(" \t\n\r");
+        value = value.substr(begin, end - begin + 1);
+    }
+
+    std::string captured_strategy_;
+};
+
+struct ClientRuntime {
+    std::shared_ptr<Client> client;
+    std::unique_ptr<SimpleAllocator> io_allocator;
+    void* segment_ptr = nullptr;
+    size_t segment_size = 0;
+    std::string host_name;
+};
+
+struct PreparedObject {
+    std::string key;
+    std::string payload;
+    QueryResult query_result;
+    Replica::Descriptor replica;
+};
+
+Replica::Descriptor FindFirstCompleteMemoryReplica(const QueryResult& result) {
+    for (const auto& replica : result.replicas) {
+        if (replica.status == ReplicaStatus::COMPLETE &&
+            replica.is_memory_replica()) {
+            return replica;
+        }
+    }
+    throw std::runtime_error("No complete memory replica found");
+}
+
+bool PutHotKeyHelper(LocalHotCache& cache, const std::string& key,
+                     const Slice& slice) {
+    if (key.empty() || slice.ptr == nullptr || slice.size == 0) {
+        return false;
+    }
+    if (cache.TouchHotKey(key)) {
+        return true;
+    }
+
+    HotMemBlock* block = cache.GetFreeBlock();
+    if (block == nullptr) {
+        return false;
+    }
+    if (slice.size > block->size) {
+        block->key_.clear();
+        cache.PutHotKey(block);
+        return false;
+    }
+
+    std::memcpy(block->addr, slice.ptr, slice.size);
+    block->size = slice.size;
+    block->key_ = key;
+    return cache.PutHotKey(block);
+}
+
+const char* HandshakeModeName(HandshakeMode mode) {
+    switch (mode) {
+        case HandshakeMode::P2P:
+            return "P2P";
+        case HandshakeMode::Metadata:
+            return "Metadata";
+    }
+    return "Unknown";
+}
+
+}  // namespace
+
+class TcpLocalMemcpyAutoEnableTest : public ::testing::Test {
+   protected:
+    static void SetUpTestSuite() {
+        google::InitGoogleLogging("TcpLocalMemcpyAutoEnableTest");
+        FLAGS_logtostderr = 1;
+    }
+
+    static void TearDownTestSuite() { google::ShutdownGoogleLogging(); }
+
+    void SetUp() override {
+        memcpy_guard_ = std::make_unique<EnvGuard>("MC_STORE_MEMCPY");
+        vlog_guard_ = std::make_unique<VLogGuard>(1);
+        unsetenv("MC_STORE_MEMCPY");
+
+        InProcMasterConfig config;
+        config.http_metadata_port = getFreeTcpPort();
+        ASSERT_TRUE(master_.Start(config));
+        master_address_ = master_.master_address();
+        metadata_url_ = master_.metadata_url();
+    }
+
+    void TearDown() override {
+        CleanupRuntime(remote_runtime_);
+        CleanupRuntime(runtime_);
+        master_.Stop();
+    }
+
+    ClientRuntime CreateRuntime(const std::string& host_name,
+                                const std::string& metadata_conn) {
+        ClientRuntime runtime;
+        runtime.host_name = host_name;
+
+        auto client_opt = Client::Create(host_name, metadata_conn, "tcp",
+                                         std::nullopt, master_address_);
+        EXPECT_TRUE(client_opt.has_value())
+            << "Failed to create client for host " << host_name
+            << ", metadata=" << metadata_conn;
+        if (!client_opt.has_value()) {
+            return runtime;
+        }
+
+        runtime.client = client_opt.value();
+        runtime.io_allocator =
+            std::make_unique<SimpleAllocator>(16 * 1024 * 1024);
+        auto reg = runtime.client->RegisterLocalMemory(
+            runtime.io_allocator->getBase(), 16 * 1024 * 1024, "cpu:0", false,
+            false);
+        EXPECT_TRUE(reg.has_value())
+            << "RegisterLocalMemory failed: " << toString(reg.error());
+
+        runtime.segment_size = 64 * 1024 * 1024;
+        runtime.segment_ptr =
+            allocate_buffer_allocator_memory(runtime.segment_size);
+        EXPECT_NE(runtime.segment_ptr, nullptr);
+        if (runtime.segment_ptr == nullptr) {
+            return runtime;
+        }
+
+        auto mount = runtime.client->MountSegment(runtime.segment_ptr,
+                                                  runtime.segment_size, "tcp");
+        EXPECT_TRUE(mount.has_value())
+            << "MountSegment failed: " << toString(mount.error());
+        if (!mount.has_value()) {
+            free_memory("", runtime.segment_ptr);
+            runtime.segment_ptr = nullptr;
+        }
+
+        return runtime;
+    }
+
+    void CleanupRuntime(ClientRuntime& runtime) {
+        if (runtime.client && runtime.segment_ptr) {
+            auto unmount = runtime.client->UnmountSegment(runtime.segment_ptr,
+                                                          runtime.segment_size);
+            EXPECT_TRUE(unmount.has_value())
+                << "UnmountSegment failed: " << toString(unmount.error());
+        }
+
+        runtime.client.reset();
+        runtime.io_allocator.reset();
+
+        if (runtime.segment_ptr) {
+            free_memory("", runtime.segment_ptr);
+            runtime.segment_ptr = nullptr;
+        }
+        runtime.segment_size = 0;
+        runtime.host_name.clear();
+    }
+
+    PreparedObject PrepareLocalObject(const std::string& key,
+                                      const std::string& payload) {
+        return PrepareObjectOnRuntime(runtime_, key, payload);
+    }
+
+    PreparedObject PrepareObjectOnRuntime(ClientRuntime& runtime,
+                                          const std::string& key,
+                                          const std::string& payload) {
+        void* write_buf = runtime.io_allocator->allocate(payload.size());
+        EXPECT_NE(write_buf, nullptr);
+        std::memcpy(write_buf, payload.data(), payload.size());
+
+        std::vector<Slice> write_slices;
+        write_slices.emplace_back(Slice{write_buf, payload.size()});
+
+        ReplicateConfig config;
+        config.replica_num = 1;
+        config.preferred_segment = runtime.host_name;
+
+        auto put = runtime.client->Put(key, write_slices, config);
+        runtime.io_allocator->deallocate(write_buf, payload.size());
+        EXPECT_TRUE(put.has_value()) << "Put failed: " << toString(put.error());
+
+        auto query = runtime.client->Query(key);
+        EXPECT_TRUE(query.has_value())
+            << "Query failed: " << toString(query.error());
+
+        auto replica = FindFirstCompleteMemoryReplica(query.value());
+
+        return PreparedObject{key, payload, std::move(query.value()), replica};
+    }
+
+    void LogReplicaDiagnostics(const std::string& label,
+                               const Replica::Descriptor& replica) {
+        const auto& endpoint = replica.get_memory_descriptor()
+                                   .buffer_descriptor.transport_endpoint_;
+        LOG(INFO) << label << " replica endpoint=" << endpoint;
+        LOG(INFO) << label << " client transport endpoint="
+                  << runtime_.client->GetTransportEndpoint();
+        LOG(INFO) << label << " is local replica="
+                  << runtime_.client->IsReplicaOnLocalMemory(replica);
+    }
+
+    InProcMaster master_;
+    std::string master_address_;
+    std::string metadata_url_;
+    ClientRuntime runtime_;
+    ClientRuntime remote_runtime_;
+    std::unique_ptr<EnvGuard> memcpy_guard_;
+    std::unique_ptr<VLogGuard> vlog_guard_;
+};
+
+TEST_F(TcpLocalMemcpyAutoEnableTest, P2PLocalReplicaUsesLocalMemcpy) {
+    runtime_ = CreateRuntime("localhost", "P2PHANDSHAKE");
+    ASSERT_TRUE(runtime_.client != nullptr);
+
+    auto prepared =
+        PrepareLocalObject("p2p_local_memcpy_key", "hello-local-memcpy");
+    LogReplicaDiagnostics("[P2P]", prepared.replica);
+
+    StrategyCaptureSink sink;
+    google::AddLogSink(&sink);
+
+    std::vector<char> out(prepared.payload.size(), '\0');
+    std::vector<Slice> read_slices;
+    read_slices.emplace_back(Slice{out.data(), out.size()});
+
+    auto get =
+        runtime_.client->Get(prepared.key, prepared.query_result, read_slices);
+
+    google::RemoveLogSink(&sink);
+
+    ASSERT_TRUE(get.has_value()) << "Get failed: " << toString(get.error());
+    ASSERT_EQ(std::memcmp(out.data(), prepared.payload.data(), out.size()), 0);
+    ASSERT_TRUE(runtime_.client->IsReplicaOnLocalMemory(prepared.replica));
+
+    const auto& replica_endpoint = prepared.replica.get_memory_descriptor()
+                                       .buffer_descriptor.transport_endpoint_;
+    EXPECT_EQ(replica_endpoint, runtime_.client->GetTransportEndpoint());
+    EXPECT_EQ(sink.strategy(), "LOCAL_MEMCPY");
+}
+
+TEST_F(TcpLocalMemcpyAutoEnableTest, MetadataLocalReplicaUsesLocalMemcpy) {
+    runtime_ = CreateRuntime("localhost", metadata_url_);
+    ASSERT_TRUE(runtime_.client != nullptr);
+
+    auto prepared =
+        PrepareLocalObject("metadata_local_memcpy_key", "hello-local-memcpy");
+    LogReplicaDiagnostics("[METADATA]", prepared.replica);
+
+    const auto& replica_endpoint = prepared.replica.get_memory_descriptor()
+                                       .buffer_descriptor.transport_endpoint_;
+    const auto client_endpoint = runtime_.client->GetTransportEndpoint();
+    const bool is_local =
+        runtime_.client->IsReplicaOnLocalMemory(prepared.replica);
+
+    StrategyCaptureSink sink;
+    google::AddLogSink(&sink);
+
+    std::vector<char> out(prepared.payload.size(), '\0');
+    std::vector<Slice> read_slices;
+    read_slices.emplace_back(Slice{out.data(), out.size()});
+
+    auto get =
+        runtime_.client->Get(prepared.key, prepared.query_result, read_slices);
+
+    google::RemoveLogSink(&sink);
+
+    ASSERT_TRUE(get.has_value()) << "Get failed: " << toString(get.error());
+    ASSERT_EQ(std::memcmp(out.data(), prepared.payload.data(), out.size()), 0);
+    ASSERT_TRUE(is_local);
+    EXPECT_NE(replica_endpoint, client_endpoint);
+    EXPECT_EQ(sink.strategy(), "LOCAL_MEMCPY");
+}
+
+TEST_F(TcpLocalMemcpyAutoEnableTest,
+       P2PRemoteReplicaOnSameTcpHostUsesTransferEngine) {
+    remote_runtime_ = CreateRuntime("127.0.0.1:18001", "P2PHANDSHAKE");
+    runtime_ = CreateRuntime("127.0.0.1:18002", "P2PHANDSHAKE");
+    ASSERT_TRUE(remote_runtime_.client != nullptr);
+    ASSERT_TRUE(runtime_.client != nullptr);
+
+    auto prepared = PrepareObjectOnRuntime(
+        remote_runtime_, "p2p_remote_same_host_key", "hello-remote-transfer");
+    LogReplicaDiagnostics("[P2P_REMOTE]", prepared.replica);
+
+    const auto& replica_endpoint = prepared.replica.get_memory_descriptor()
+                                       .buffer_descriptor.transport_endpoint_;
+    const auto client_endpoint = runtime_.client->GetTransportEndpoint();
+    ASSERT_NE(replica_endpoint, client_endpoint);
+    ASSERT_FALSE(runtime_.client->IsReplicaOnLocalMemory(prepared.replica));
+
+    StrategyCaptureSink sink;
+    google::AddLogSink(&sink);
+
+    std::vector<char> out(prepared.payload.size(), '\0');
+    std::vector<Slice> read_slices;
+    read_slices.emplace_back(Slice{out.data(), out.size()});
+
+    auto get =
+        runtime_.client->Get(prepared.key, prepared.query_result, read_slices);
+
+    google::RemoveLogSink(&sink);
+
+    ASSERT_TRUE(get.has_value()) << "Get failed: " << toString(get.error());
+    ASSERT_EQ(std::memcmp(out.data(), prepared.payload.data(), out.size()), 0);
+    EXPECT_EQ(sink.strategy(), "TRANSFER_ENGINE");
+}
+
+TEST_F(TcpLocalMemcpyAutoEnableTest,
+       MetadataRemoteReplicaOnSameTcpHostUsesTransferEngine) {
+    remote_runtime_ = CreateRuntime("metadata-host-a", metadata_url_);
+    runtime_ = CreateRuntime("metadata-host-b", metadata_url_);
+    ASSERT_TRUE(remote_runtime_.client != nullptr);
+    ASSERT_TRUE(runtime_.client != nullptr);
+
+    auto prepared =
+        PrepareObjectOnRuntime(remote_runtime_, "metadata_remote_same_host_key",
+                               "hello-metadata-remote-transfer");
+    LogReplicaDiagnostics("[METADATA_REMOTE]", prepared.replica);
+
+    const auto& replica_endpoint = prepared.replica.get_memory_descriptor()
+                                       .buffer_descriptor.transport_endpoint_;
+    const auto client_endpoint = runtime_.client->GetTransportEndpoint();
+    ASSERT_EQ(replica_endpoint, remote_runtime_.host_name);
+    ASSERT_NE(replica_endpoint, runtime_.host_name);
+    ASSERT_NE(replica_endpoint, client_endpoint);
+    ASSERT_FALSE(runtime_.client->IsReplicaOnLocalMemory(prepared.replica));
+
+    StrategyCaptureSink sink;
+    google::AddLogSink(&sink);
+
+    std::vector<char> out(prepared.payload.size(), '\0');
+    std::vector<Slice> read_slices;
+    read_slices.emplace_back(Slice{out.data(), out.size()});
+
+    auto get =
+        runtime_.client->Get(prepared.key, prepared.query_result, read_slices);
+
+    google::RemoveLogSink(&sink);
+
+    ASSERT_TRUE(get.has_value()) << "Get failed: " << toString(get.error());
+    ASSERT_EQ(std::memcmp(out.data(), prepared.payload.data(), out.size()), 0);
+    EXPECT_EQ(sink.strategy(), "TRANSFER_ENGINE");
+}
+
+class HotCacheRedirectStrategyTest
+    : public TcpLocalMemcpyAutoEnableTest,
+      public ::testing::WithParamInterface<HandshakeMode> {};
+
+TEST_P(HotCacheRedirectStrategyTest, CacheHitUsesLocalMemcpy) {
+    EnvGuard cache_size_guard("MC_STORE_LOCAL_HOT_CACHE_SIZE");
+    setenv("MC_STORE_LOCAL_HOT_CACHE_SIZE", "33554432", 1);  // 32MB
+
+    const HandshakeMode mode = GetParam();
+    const bool is_p2p = mode == HandshakeMode::P2P;
+    const std::string metadata_conn = is_p2p ? "P2PHANDSHAKE" : metadata_url_;
+    const std::string host_name =
+        is_p2p ? "localhost" : "metadata-hot-cache-host";
+
+    runtime_ = CreateRuntime(host_name, metadata_conn);
+    ASSERT_TRUE(runtime_.client != nullptr) << HandshakeModeName(mode);
+    ASSERT_TRUE(runtime_.client->IsHotCacheEnabled())
+        << HandshakeModeName(mode);
+
+    const std::string key = is_p2p ? "p2p_hot_cache_redirect_key"
+                                   : "metadata_hot_cache_redirect_key";
+    const std::string payload =
+        is_p2p ? "p2p-hot-cache-data" : "metadata-hot-cache-data";
+
+    Slice cache_slice{const_cast<char*>(payload.data()), payload.size()};
+    ASSERT_TRUE(
+        PutHotKeyHelper(*runtime_.client->GetHotCache(), key, cache_slice));
+    ASSERT_TRUE(runtime_.client->GetHotCache()->HasHotKey(key));
+
+    Replica::Descriptor replica;
+    replica.id = is_p2p ? 1 : 2;
+    replica.status = ReplicaStatus::COMPLETE;
+    MemoryDescriptor mem_desc;
+    mem_desc.buffer_descriptor.transport_endpoint_ = "remote:9999";
+    mem_desc.buffer_descriptor.buffer_address_ = 0;
+    mem_desc.buffer_descriptor.size_ = payload.size();
+    replica.descriptor_variant = mem_desc;
+
+    ASSERT_FALSE(runtime_.client->IsReplicaOnLocalMemory(replica))
+        << HandshakeModeName(mode);
+
+    std::vector<Replica::Descriptor> replicas;
+    replicas.emplace_back(replica);
+    QueryResult query_result(
+        std::move(replicas),
+        std::chrono::steady_clock::now() + std::chrono::seconds(60));
+
+    StrategyCaptureSink sink;
+    google::AddLogSink(&sink);
+
+    std::vector<char> out(payload.size(), '\0');
+    std::vector<Slice> read_slices;
+    read_slices.emplace_back(Slice{out.data(), out.size()});
+
+    auto get = runtime_.client->Get(key, query_result, read_slices);
+
+    google::RemoveLogSink(&sink);
+
+    ASSERT_TRUE(get.has_value()) << "Get failed: " << toString(get.error());
+    ASSERT_EQ(std::memcmp(out.data(), payload.data(), out.size()), 0);
+    EXPECT_EQ(sink.strategy(), "LOCAL_MEMCPY");
+}
+
+INSTANTIATE_TEST_SUITE_P(
+    AllModes, HotCacheRedirectStrategyTest,
+    ::testing::Values(HandshakeMode::P2P, HandshakeMode::Metadata),
+    [](const ::testing::TestParamInfo<HandshakeMode>& info) {
+        switch (info.param) {
+            case HandshakeMode::P2P:
+                return "P2P";
+            case HandshakeMode::Metadata:
+                return "Metadata";
+        }
+        return "Unknown";
+    });
+
+}  // namespace testing
+}  // namespace mooncake
\ No newline at end of file
diff --git a/mooncake-store/tests/deadline_scheduler_test.cpp b/mooncake-store/tests/deadline_scheduler_test.cpp
new file mode 100644
index 00000000..cb6ce27b
--- /dev/null
+++ b/mooncake-store/tests/deadline_scheduler_test.cpp
@@ -0,0 +1,218 @@
+#include "deadline_scheduler.h"
+
+#include <gtest/gtest.h>
+
+#include <chrono>
+#include <condition_variable>
+#include <mutex>
+#include <thread>
+#include <vector>
+
+#include "types.h"
+
+namespace mooncake::test {
+namespace {
+
+using Scheduler = DeadlineScheduler<UUID>;
+using Clock = Scheduler::Clock;
+
+class CallbackRecorder {
+   public:
+    void Record(const UUID& id) {
+        {
+            std::lock_guard<std::mutex> lock(mutex_);
+            records_.push_back(id);
+        }
+        cv_.notify_all();
+    }
+
+    bool WaitForSize(size_t expected_size, std::chrono::milliseconds timeout) {
+        std::unique_lock<std::mutex> lock(mutex_);
+        return cv_.wait_for(lock, timeout, [this, expected_size]() {
+            return records_.size() >= expected_size;
+        });
+    }
+
+    std::vector<UUID> Records() const {
+        std::lock_guard<std::mutex> lock(mutex_);
+        return records_;
+    }
+
+    size_t Size() const {
+        std::lock_guard<std::mutex> lock(mutex_);
+        return records_.size();
+    }
+
+   private:
+    mutable std::mutex mutex_;
+    std::condition_variable cv_;
+    std::vector<UUID> records_;
+};
+
+UUID MakeUuid(uint64_t value) { return UUID{value, value + 1000}; }
+
+}  // namespace
+
+TEST(DeadlineSchedulerTest, RunsRecordsInDeadlineOrder) {
+    CallbackRecorder recorder;
+    Scheduler scheduler([&recorder](const UUID& id) { recorder.Record(id); });
+
+    const auto item1 = MakeUuid(1);
+    const auto item2 = MakeUuid(2);
+    const auto item3 = MakeUuid(3);
+    const auto now = Clock::now();
+
+    scheduler.Schedule(item1, now + std::chrono::milliseconds(90));
+    scheduler.Schedule(item2, now + std::chrono::milliseconds(30));
+    scheduler.Schedule(item3, now + std::chrono::milliseconds(60));
+
+    ASSERT_TRUE(recorder.WaitForSize(3, std::chrono::milliseconds(500)));
+
+    const auto records = recorder.Records();
+    ASSERT_EQ(records.size(), 3);
+    EXPECT_EQ(records[0], item2);
+    EXPECT_EQ(records[1], item3);
+    EXPECT_EQ(records[2], item1);
+}
+
+TEST(DeadlineSchedulerTest, EarlierRecordPreemptsCurrentWait) {
+    CallbackRecorder recorder;
+    Scheduler scheduler([&recorder](const UUID& id) { recorder.Record(id); });
+
+    const auto long_item = MakeUuid(10);
+    const auto short_item = MakeUuid(20);
+
+    scheduler.Schedule(long_item,
+                       Clock::now() + std::chrono::milliseconds(500));
+    std::this_thread::sleep_for(std::chrono::milliseconds(20));
+    scheduler.Schedule(short_item,
+                       Clock::now() + std::chrono::milliseconds(40));
+
+    ASSERT_TRUE(recorder.WaitForSize(1, std::chrono::milliseconds(250)));
+    const auto records = recorder.Records();
+    ASSERT_EQ(records.size(), 1);
+    EXPECT_EQ(records[0], short_item);
+
+    scheduler.Stop();
+    EXPECT_EQ(recorder.Size(), 1);
+}
+
+TEST(DeadlineSchedulerTest, RemoveIfDropsPendingRecords) {
+    CallbackRecorder recorder;
+    Scheduler scheduler([&recorder](const UUID& id) { recorder.Record(id); });
+
+    const auto removed_item1 = MakeUuid(31);
+    const auto kept_item = MakeUuid(41);
+    const auto removed_item2 = MakeUuid(32);
+    const auto now = Clock::now();
+
+    scheduler.Schedule(removed_item1, now + std::chrono::milliseconds(200));
+    scheduler.Schedule(kept_item, now + std::chrono::milliseconds(250));
+    scheduler.Schedule(removed_item2, now + std::chrono::milliseconds(300));
+
+    scheduler.RemoveIf([&](const UUID& id) {
+        return id == removed_item1 || id == removed_item2;
+    });
+
+    ASSERT_TRUE(recorder.WaitForSize(1, std::chrono::milliseconds(600)));
+    std::this_thread::sleep_for(std::chrono::milliseconds(150));
+
+    const auto records = recorder.Records();
+    ASSERT_EQ(records.size(), 1);
+    EXPECT_EQ(records[0], kept_item);
+}
+
+TEST(DeadlineSchedulerTest, StopCancelsPendingRecords) {
+    CallbackRecorder recorder;
+    Scheduler scheduler([&recorder](const UUID& id) { recorder.Record(id); });
+
+    scheduler.Schedule(MakeUuid(50),
+                       Clock::now() + std::chrono::milliseconds(200));
+    scheduler.Stop();
+
+    std::this_thread::sleep_for(std::chrono::milliseconds(250));
+    EXPECT_EQ(recorder.Size(), 0);
+}
+
+TEST(DeadlineSchedulerTest, StopOnEmptySchedulerIsNoOp) {
+    CallbackRecorder recorder;
+    Scheduler scheduler([&recorder](const UUID& id) { recorder.Record(id); });
+
+    scheduler.Stop();
+    EXPECT_EQ(recorder.Size(), 0);
+}
+
+TEST(DeadlineSchedulerTest, RemoveIfOnEmptySchedulerIsNoOp) {
+    CallbackRecorder recorder;
+    Scheduler scheduler([&recorder](const UUID& id) { recorder.Record(id); });
+
+    scheduler.RemoveIf([](const UUID&) { return true; });
+    EXPECT_EQ(recorder.Size(), 0);
+}
+
+TEST(DeadlineSchedulerTest, ContinuesSchedulingAfterQueueDrains) {
+    CallbackRecorder recorder;
+    Scheduler scheduler([&recorder](const UUID& id) { recorder.Record(id); });
+
+    const auto first_item = MakeUuid(60);
+    scheduler.Schedule(first_item,
+                       Clock::now() + std::chrono::milliseconds(30));
+    ASSERT_TRUE(recorder.WaitForSize(1, std::chrono::milliseconds(500)));
+
+    const auto second_item = MakeUuid(61);
+    scheduler.Schedule(second_item,
+                       Clock::now() + std::chrono::milliseconds(30));
+    ASSERT_TRUE(recorder.WaitForSize(2, std::chrono::milliseconds(500)));
+
+    const auto records = recorder.Records();
+    ASSERT_EQ(records.size(), 2);
+    EXPECT_EQ(records[0], first_item);
+    EXPECT_EQ(records[1], second_item);
+}
+
+TEST(DeadlineSchedulerTest, RemoveAllDrainsQueueWithoutRunning) {
+    CallbackRecorder recorder;
+    Scheduler scheduler([&recorder](const UUID& id) { recorder.Record(id); });
+
+    scheduler.Schedule(MakeUuid(70),
+                       Clock::now() + std::chrono::milliseconds(200));
+    scheduler.Schedule(MakeUuid(71),
+                       Clock::now() + std::chrono::milliseconds(250));
+    scheduler.RemoveIf([](const UUID&) { return true; });
+
+    std::this_thread::sleep_for(std::chrono::milliseconds(350));
+    EXPECT_EQ(recorder.Size(), 0);
+}
+
+TEST(DeadlineSchedulerTest, RunsAlreadyExpiredRecordImmediately) {
+    CallbackRecorder recorder;
+    Scheduler scheduler([&recorder](const UUID& id) { recorder.Record(id); });
+
+    const auto past_item = MakeUuid(80);
+    scheduler.Schedule(past_item, Clock::now() - std::chrono::milliseconds(10));
+
+    ASSERT_TRUE(recorder.WaitForSize(1, std::chrono::milliseconds(500)));
+    const auto records = recorder.Records();
+    ASSERT_EQ(records.size(), 1);
+    EXPECT_EQ(records[0], past_item);
+}
+
+TEST(DeadlineSchedulerTest, RunsExpiredBeforePendingRecord) {
+    CallbackRecorder recorder;
+    Scheduler scheduler([&recorder](const UUID& id) { recorder.Record(id); });
+
+    const auto past_item = MakeUuid(90);
+    const auto future_item = MakeUuid(91);
+    const auto now = Clock::now();
+
+    scheduler.Schedule(future_item, now + std::chrono::milliseconds(120));
+    scheduler.Schedule(past_item, now - std::chrono::milliseconds(10));
+
+    ASSERT_TRUE(recorder.WaitForSize(2, std::chrono::milliseconds(500)));
+    const auto records = recorder.Records();
+    ASSERT_EQ(records.size(), 2);
+    EXPECT_EQ(records[0], past_item);
+    EXPECT_EQ(records[1], future_item);
+}
+
+}  // namespace mooncake::test
diff --git a/mooncake-store/tests/dummy_client_get_buffer_test.cpp b/mooncake-store/tests/dummy_client_get_buffer_test.cpp
index 3635d3de..b1b1dc53 100644
--- a/mooncake-store/tests/dummy_client_get_buffer_test.cpp
+++ b/mooncake-store/tests/dummy_client_get_buffer_test.cpp
@@ -9,6 +9,7 @@
 #include <chrono>
 #include <cstdlib>
 #include <cstring>
+#include <iostream>
 #include <memory>
 #include <string>
 #include <thread>
@@ -184,6 +185,16 @@ class DummyClientGetBufferTest : public ::testing::Test {
         std::this_thread::sleep_for(std::chrono::milliseconds(2000));
     }
 
+    // Trigger hot-cache admission via DummyClient fallback reads (CMS threshold
+    // default 2) but do NOT wait for async PutHotKey to publish.
+    void TriggerAdmissionWithoutWaiting(const std::string &key) {
+        for (int i = 0; i < 2; ++i) {
+            auto buf = dummy_client_->get_buffer(key);
+            ASSERT_NE(buf, nullptr)
+                << "get_buffer failed while triggering admission";
+        }
+    }
+
     mooncake::testing::InProcMaster master_;
     std::shared_ptr<RealClient> real_client_;
     std::shared_ptr<DummyClient> dummy_client_;
@@ -194,6 +205,156 @@ class DummyClientGetBufferTest : public ::testing::Test {
     std::optional<std::string> saved_hot_block_env_;
     std::optional<std::string> saved_hot_shm_env_;
 };
+// ---- Regression: a stable hot-cache entry must be invalidated by remove ----
+TEST_F(DummyClientGetBufferTest,
+       StableHotCacheEntryShouldBeInvalidatedByRemove) {
+    ASSERT_TRUE(SetupStack()) << "Failed to bring up real+dummy stack";
+
+    const std::string key = "stable_hot_cache_remove_invalidation";
+
+    const char fill = static_cast<char>(7);
+    const std::string data(kPayloadSize, fill);
+
+    PutData(key, data);
+    WarmHotCache(key);
+
+    auto hot_buf_1 = dummy_client_->get_buffer(key);
+    ASSERT_NE(hot_buf_1, nullptr)
+        << "dummy get_buffer should succeed after WarmHotCache";
+    ASSERT_EQ(hot_buf_1->size(), data.size())
+        << "hot cache buffer size mismatch before remove";
+    ASSERT_TRUE(dummy_client_->is_hot_cache_ptr(hot_buf_1->ptr()))
+        << "buffer should be in hot cache shm before remove";
+
+    {
+        std::string got(static_cast<char *>(hot_buf_1->ptr()),
+                        hot_buf_1->size());
+        ASSERT_EQ(got.size(), data.size());
+        ASSERT_TRUE(std::equal(got.begin(), got.end(), data.begin()))
+            << "hot cache data mismatch before remove";
+    }
+
+    auto hot_buf_2 = dummy_client_->get_buffer(key);
+    ASSERT_NE(hot_buf_2, nullptr)
+        << "second dummy get_buffer should still succeed before remove";
+    ASSERT_TRUE(dummy_client_->is_hot_cache_ptr(hot_buf_2->ptr()))
+        << "second buffer should still be in hot cache shm before remove";
+    ASSERT_EQ(hot_buf_2->size(), data.size())
+        << "second hot cache buffer size mismatch before remove";
+
+    {
+        std::string got(static_cast<char *>(hot_buf_2->ptr()),
+                        hot_buf_2->size());
+        ASSERT_EQ(got.size(), data.size());
+        ASSERT_TRUE(std::equal(got.begin(), got.end(), data.begin()))
+            << "second hot cache data mismatch before remove";
+    }
+
+    int remove_rc = real_client_->remove(key, true);
+    ASSERT_EQ(remove_rc, 0) << "remove failed";
+
+    auto after_remove_buf = dummy_client_->get_buffer(key);
+    if (after_remove_buf != nullptr) {
+        const bool after_remove_is_hot =
+            dummy_client_->is_hot_cache_ptr(after_remove_buf->ptr());
+        std::cerr << "[StableHotCacheEntryShouldBeInvalidatedByRemove] "
+                  << "after remove get_buffer returned non-null; "
+                  << "size=" << after_remove_buf->size()
+                  << ", is_hot_cache_ptr=" << after_remove_is_hot << std::endl;
+        FAIL() << "remove(key, true) succeeded, but dummy get_buffer(key) "
+               << "still returned a non-null buffer";
+    }
+}
+
+// ---- Diagnosis: removed key must not be resurrected by async hot-cache fill
+// ----
+//
+// Compare against pre-fix behavior:
+//   mooncake-store/tests/scripts/compare_hot_cache_fix.sh all
+TEST_F(DummyClientGetBufferTest,
+       RemoveShouldNotBeResurrectedByAsyncHotCacheFill) {
+    ASSERT_TRUE(SetupStack()) << "Failed to bring up real+dummy stack";
+
+    const std::string key = "async_fill_resurrect_removed_key";
+    constexpr int kRounds = 8;
+
+    for (int round = 0; round < kRounds; ++round) {
+        SCOPED_TRACE("round=" + std::to_string(round));
+
+        const char fill = static_cast<char>(1 + (round % 250));
+        const std::string data(kPayloadSize, fill);
+
+        PutData(key, data);
+
+        // Two fallback reads reach CMS admission threshold (default 2) and
+        // submit async fill, but we do not wait for worker PutHotKey.
+        TriggerAdmissionWithoutWaiting(key);
+
+        int remove_rc = real_client_->remove(key, true);
+        ASSERT_EQ(remove_rc, 0) << "remove failed at round=" << round;
+
+        auto immediately_after_remove = dummy_client_->get_buffer(key);
+        const bool immediate_non_null = (immediately_after_remove != nullptr);
+        bool immediate_is_hot = false;
+        size_t immediate_size = 0;
+        if (immediate_non_null) {
+            immediate_is_hot = dummy_client_->is_hot_cache_ptr(
+                immediately_after_remove->ptr());
+            immediate_size = immediately_after_remove->size();
+            std::cerr << "[round=" << round
+                      << "] immediately after remove: NON_NULL"
+                      << ", size=" << immediate_size
+                      << ", is_hot_cache_ptr=" << immediate_is_hot << std::endl;
+        } else {
+            std::cerr << "[round=" << round
+                      << "] immediately after remove: nullptr" << std::endl;
+        }
+
+        std::this_thread::sleep_for(std::chrono::milliseconds(2000));
+
+        auto delayed_after_remove = dummy_client_->get_buffer(key);
+        const bool delayed_non_null = (delayed_after_remove != nullptr);
+        bool delayed_is_hot = false;
+        size_t delayed_size = 0;
+        if (delayed_non_null) {
+            delayed_is_hot =
+                dummy_client_->is_hot_cache_ptr(delayed_after_remove->ptr());
+            delayed_size = delayed_after_remove->size();
+            std::cerr << "[round=" << round
+                      << "] delayed after remove: NON_NULL"
+                      << ", size=" << delayed_size
+                      << ", is_hot_cache_ptr=" << delayed_is_hot << std::endl;
+        } else {
+            std::cerr << "[round=" << round << "] delayed after remove: nullptr"
+                      << std::endl;
+        }
+
+        if (immediate_non_null || delayed_non_null) {
+            if (!immediate_non_null && delayed_non_null) {
+                FAIL() << "FAILED at round=" << round
+                       << ": key was nullptr immediately after remove, "
+                       << "but became non-null after waiting. "
+                       << "This strongly suggests async hot-cache fill "
+                       << "resurrected a removed key. "
+                       << "delayed_is_hot_cache_ptr=" << delayed_is_hot
+                       << ", delayed_size=" << delayed_size;
+            }
+            if (immediate_non_null && delayed_non_null) {
+                FAIL() << "FAILED at round=" << round
+                       << ": key was non-null immediately after remove "
+                       << "and remained non-null after waiting. "
+                       << "immediate_is_hot_cache_ptr=" << immediate_is_hot
+                       << ", delayed_is_hot_cache_ptr=" << delayed_is_hot;
+            }
+            if (immediate_non_null && !delayed_non_null) {
+                FAIL() << "FAILED at round=" << round
+                       << ": key was non-null immediately after remove, "
+                       << "but became nullptr after waiting. "
+                       << "immediate_is_hot_cache_ptr=" << immediate_is_hot;
+            }
+        }
+    }
+}
 
 // ---- Test: get_buffer via allocator fallback (no hot cache hit) ----
 TEST_F(DummyClientGetBufferTest, GetBuffer_AllocatorFallback) {
@@ -409,4 +570,4 @@ int main(int argc, char **argv) {
     ::testing::InitGoogleTest(&argc, argv);
     gflags::ParseCommandLineFlags(&argc, &argv, false);
     return RUN_ALL_TESTS();
-}
\ No newline at end of file
+}
diff --git a/mooncake-store/tests/e2e/readme.md b/mooncake-store/tests/e2e/readme.md
index 0f7a4b39..e8ad2aed 100644
--- a/mooncake-store/tests/e2e/readme.md
+++ b/mooncake-store/tests/e2e/readme.md
@@ -11,6 +11,8 @@ The E2E test suite includes several executable programs designed to test differe
 - **e2e_rand_test**: Long-term randomized end-to-end testing.
 - **chaos_test**: Short-term chaos testing with predefined scenarios.
 - **chaos_rand_test**: Long-term randomized chaos testing with configurable parameters.
+- **store_client_e2e.py**: Python `MooncakeDistributedStore` client that continuously issues `put/get` operations.
+- **run_nof_heartbeat_tcp_e2e.sh**: Scripted NoF heartbeat end-to-end test using a TCP SPDK target.
 
 ## Parameters
 
@@ -72,4 +74,78 @@ Currently it only has few test cases. Will add more in the future.
 2. Run the test.
 
 **[WIP]**:
-Currently it only has few test cases. Will add more in the future.
\ No newline at end of file
+Currently it only has few test cases. Will add more in the future.
+
+### run_nof_heartbeat_tcp_e2e.sh
+
+**Brief**: Launches a real four-component path for NoF heartbeat validation:
+
+- `mooncake_master`
+- standalone Python HTTP metadata server
+- SPDK `nvmf_tgt` with TCP transport
+- Python client built on `MooncakeDistributedStore`
+
+The script first verifies steady-state `put/get` success with `memory + nof` replicas. It then kills the SPDK target, waits for the master heartbeat thread to emit `action=unmount_nof_segment_by_heartbeat`, and finally verifies that the client still observes successful I/O after the NoF segment is removed.
+
+**Prerequisites**:
+
+- `BUILD_DIR` points to a build tree that already contains:
+  - `mooncake-store/src/mooncake_master`
+  - `mooncake-integration/store*.so`
+- SPDK has already been built under `extern/spdk`
+- Python environment contains `aiohttp` because the script launches a standalone metadata process with `mooncake-wheel/mooncake/http_metadata_server.py`
+- The script uses `sudo -n` to set hugepages and mount `/dev/hugepages`, so the current user must have passwordless sudo
+
+**Usage**:
+
+```bash
+cd mooncake-store/tests/e2e
+BUILD_DIR=/path/to/build ./run_nof_heartbeat_tcp_e2e.sh
+```
+
+To run in **NoF-only** mode (do not mount a local memory segment), set:
+
+```bash
+CLIENT_GLOBAL_SEGMENT_SIZE=0 BUILD_DIR=/path/to/build ./run_nof_heartbeat_tcp_e2e.sh
+```
+
+To increase the amount of steady-state traffic before killing the target, set:
+
+```bash
+PRE_FAULT_SUCCESS_TARGET=10 BUILD_DIR=/path/to/build ./run_nof_heartbeat_tcp_e2e.sh
+```
+
+**Notes**:
+
+- The client payload size defaults to `4096` bytes because the current NoF path requires 4K-aligned I/O.
+- The script uses a standalone metadata server process (`mooncake-wheel/mooncake/http_metadata_server.py`) instead of the embedded master metadata server so all four components remain explicit during the test.
+- In default mode, the script verifies **service continuity** after NoF unmount by checking that post-fault I/O still succeeds.
+- In `CLIENT_GLOBAL_SEGMENT_SIZE=0` mode, the script verifies **NoF-only failure behavior** by checking that post-fault I/O starts failing after the NoF segment is removed.
+- Logs are written under `LOG_DIR` (default `/tmp/mooncake_nof_heartbeat_e2e`) and the final pass/fail summary is printed from `summary.log`.
+
+### store_client_e2e.py
+
+**Brief**: A standalone Python workload generator built on `MooncakeDistributedStore`. It continuously issues `put/get` against the configured master/metadata pair and prints `put_ok/get_ok/put_fail/get_fail` lines that can be consumed by shell scripts.
+
+**Standalone Usage**:
+
+```bash
+PYTHONPATH=/path/to/build/mooncake-integration \
+python3 store_client_e2e.py \
+  --local-hostname 127.0.0.1:50071 \
+  --metadata-server http://127.0.0.1:8080/metadata \
+  --master-server 127.0.0.1:50051 \
+  --global-segment-size 67108864 \
+  --local-buffer-size 33554432 \
+  --payload-size 4096 \
+  --duration-sec 20 \
+  --sleep-ms 200 \
+  --key-prefix demo
+```
+
+**Key Parameters**:
+
+- `--global-segment-size 0`: run in NoF-only mode
+- `--payload-size 4096`: keep NoF writes 4K aligned
+- `--duration-sec`: total workload duration
+- `--sleep-ms`: interval between operations
diff --git a/mooncake-store/tests/e2e/run_nof_heartbeat_tcp_e2e.sh b/mooncake-store/tests/e2e/run_nof_heartbeat_tcp_e2e.sh
new file mode 100644
index 00000000..4e748824
--- /dev/null
+++ b/mooncake-store/tests/e2e/run_nof_heartbeat_tcp_e2e.sh
@@ -0,0 +1,245 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+SCRIPT_DIR=$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)
+REPO_ROOT=$(cd -- "$SCRIPT_DIR/../../.." && pwd)
+BUILD_DIR=${BUILD_DIR:-"$REPO_ROOT/build"}
+LOG_DIR=${LOG_DIR:-/tmp/mooncake_nof_heartbeat_e2e}
+MASTER_RPC=${MASTER_RPC:-127.0.0.1:50051}
+MASTER_HOST=${MASTER_RPC%:*}
+MASTER_PORT=${MASTER_RPC##*:}
+METADATA_HOST=${METADATA_HOST:-127.0.0.1}
+METADATA_PORT=${METADATA_PORT:-8080}
+METADATA_SERVER="http://$METADATA_HOST:$METADATA_PORT/metadata"
+TARGET_HOST=${TARGET_HOST:-127.0.0.1}
+TARGET_PORT=${TARGET_PORT:-4420}
+TARGET_NQN=${TARGET_NQN:-nqn.2016-06.io.spdk:cnode1}
+NOF_SIZE=${NOF_SIZE:-67108864}
+PAYLOAD_SIZE=${PAYLOAD_SIZE:-4096}
+CLIENT_DURATION=${CLIENT_DURATION:-0}
+CLIENT_SLEEP_MS=${CLIENT_SLEEP_MS:-200}
+HEARTBEAT_INTERVAL=${HEARTBEAT_INTERVAL:-2}
+HEARTBEAT_TIMEOUT_MS=${HEARTBEAT_TIMEOUT_MS:-500}
+HEARTBEAT_FAILURES=${HEARTBEAT_FAILURES:-3}
+CLIENT_GLOBAL_SEGMENT_SIZE=${CLIENT_GLOBAL_SEGMENT_SIZE:-67108864}
+CLIENT_LOCAL_BUFFER_SIZE=${CLIENT_LOCAL_BUFFER_SIZE:-33554432}
+CLIENT_MEMORY_REPLICA_NUM=${CLIENT_MEMORY_REPLICA_NUM:-1}
+CLIENT_NOF_REPLICA_NUM=${CLIENT_NOF_REPLICA_NUM:-1}
+PRE_FAULT_SUCCESS_TARGET=${PRE_FAULT_SUCCESS_TARGET:-3}
+
+TARGET_PID=""
+MASTER_PID=""
+META_PID=""
+CLIENT_PID=""
+
+cleanup() {
+  [[ -n "$CLIENT_PID" ]] && kill "$CLIENT_PID" >/dev/null 2>&1 || true
+  [[ -n "$MASTER_PID" ]] && kill "$MASTER_PID" >/dev/null 2>&1 || true
+  [[ -n "$TARGET_PID" ]] && kill "$TARGET_PID" >/dev/null 2>&1 || true
+  [[ -n "$META_PID" ]] && kill "$META_PID" >/dev/null 2>&1 || true
+}
+trap cleanup EXIT
+
+wait_for_pattern() {
+  local file=$1
+  local pattern=$2
+  local timeout=$3
+  local waited=0
+  while (( waited < timeout )); do
+    if grep -q "$pattern" "$file" 2>/dev/null; then
+      return 0
+    fi
+    sleep 1
+    waited=$((waited + 1))
+  done
+  return 1
+}
+
+count_pattern() {
+  local file=$1
+  local pattern=$2
+  grep -c "$pattern" "$file" 2>/dev/null || true
+}
+
+rm -rf "$LOG_DIR"
+mkdir -p "$LOG_DIR"
+
+sudo -n sh -c 'echo 512 > /proc/sys/vm/nr_hugepages'
+sudo -n umount /dev/hugepages >/dev/null 2>&1 || true
+sudo -n mkdir -p /dev/hugepages
+sudo -n mount -t hugetlbfs -o pagesize=2M,mode=1777 none /dev/hugepages
+grep -E 'HugePages_Total|HugePages_Free|Hugepagesize' /proc/meminfo >"$LOG_DIR/hugepages.log"
+
+pkill -f '/nvmf_tgt' >/dev/null 2>&1 || true
+pkill -f 'mooncake_master' >/dev/null 2>&1 || true
+pkill -f 'http_metadata_server.py' >/dev/null 2>&1 || true
+sleep 1
+
+"$REPO_ROOT/extern/spdk/build/bin/nvmf_tgt" -m 0x1 -u --iova-mode=va --wait-for-rpc >"$LOG_DIR/target.log" 2>&1 &
+TARGET_PID=$!
+sleep 3
+
+python3 "$REPO_ROOT/extern/spdk/scripts/rpc.py" framework_start_init >/dev/null
+python3 "$REPO_ROOT/extern/spdk/scripts/rpc.py" framework_wait_init >/dev/null
+python3 "$REPO_ROOT/extern/spdk/scripts/rpc.py" bdev_malloc_create -b Malloc0 64 4096 >/dev/null
+python3 "$REPO_ROOT/extern/spdk/scripts/rpc.py" nvmf_create_transport -t TCP >/dev/null || true
+python3 "$REPO_ROOT/extern/spdk/scripts/rpc.py" nvmf_create_subsystem "$TARGET_NQN" -a -s SPDK00000000000001 >/dev/null || true
+python3 "$REPO_ROOT/extern/spdk/scripts/rpc.py" nvmf_subsystem_add_ns "$TARGET_NQN" Malloc0 >/dev/null || true
+python3 "$REPO_ROOT/extern/spdk/scripts/rpc.py" nvmf_subsystem_add_listener "$TARGET_NQN" -t tcp -a "$TARGET_HOST" -s "$TARGET_PORT" >/dev/null || true
+
+python3 "$REPO_ROOT/mooncake-wheel/mooncake/http_metadata_server.py" --host "$METADATA_HOST" --port "$METADATA_PORT" >"$LOG_DIR/metadata.log" 2>&1 &
+META_PID=$!
+sleep 2
+
+"$BUILD_DIR/mooncake-store/src/mooncake_master" \
+  --rpc_address="$MASTER_HOST" \
+  --rpc_port="$MASTER_PORT" \
+  --enable_http_metadata_server=false \
+  --logtostderr=true \
+  --nof_heartbeat_interval_sec="$HEARTBEAT_INTERVAL" \
+  --nof_heartbeat_probe_timeout_ms="$HEARTBEAT_TIMEOUT_MS" \
+  --nof_heartbeat_failures_threshold="$HEARTBEAT_FAILURES" >"$LOG_DIR/master.log" 2>&1 &
+MASTER_PID=$!
+sleep 3
+
+PYTHONPATH="$BUILD_DIR/mooncake-integration" MC_NOF_TRTYPE=TCP python3 - <<PY >"$LOG_DIR/register.log" 2>&1
+import store
+ret = store.MooncakeDistributedNoFRegister().real_register(
+    "$TARGET_NQN",
+    1,
+    "$TARGET_HOST",
+    int("$TARGET_PORT"),
+    0,
+    int("$NOF_SIZE"),
+    "$MASTER_RPC",
+)
+print(f"register_ret {ret}", flush=True)
+raise SystemExit(0 if ret == 0 else 1)
+PY
+
+PYTHONPATH="$BUILD_DIR/mooncake-integration" python3 "$SCRIPT_DIR/store_client_e2e.py" \
+  --local-hostname "127.0.0.1:50071" \
+  --metadata-server "$METADATA_SERVER" \
+  --master-server "$MASTER_RPC" \
+  --global-segment-size "$CLIENT_GLOBAL_SEGMENT_SIZE" \
+  --local-buffer-size "$CLIENT_LOCAL_BUFFER_SIZE" \
+  --memory-replica-num "$CLIENT_MEMORY_REPLICA_NUM" \
+  --nof-replica-num "$CLIENT_NOF_REPLICA_NUM" \
+  --payload-size "$PAYLOAD_SIZE" \
+  --duration-sec "$CLIENT_DURATION" \
+  --sleep-ms "$CLIENT_SLEEP_MS" \
+  --key-prefix "nof-heartbeat" >"$LOG_DIR/client.log" 2>&1 &
+CLIENT_PID=$!
+
+deadline=$((SECONDS + 30))
+while (( SECONDS < deadline )); do
+  put_ok_count=$(count_pattern "$LOG_DIR/client.log" 'put_ok')
+  get_ok_count=$(count_pattern "$LOG_DIR/client.log" 'get_ok')
+  if (( put_ok_count >= PRE_FAULT_SUCCESS_TARGET && get_ok_count >= PRE_FAULT_SUCCESS_TARGET )); then
+    break
+  fi
+  sleep 1
+done
+
+put_ok_count=$(count_pattern "$LOG_DIR/client.log" 'put_ok')
+get_ok_count=$(count_pattern "$LOG_DIR/client.log" 'get_ok')
+if (( put_ok_count < PRE_FAULT_SUCCESS_TARGET || get_ok_count < PRE_FAULT_SUCCESS_TARGET )); then
+  echo "client did not reach enough initial success: put_ok=$put_ok_count get_ok=$get_ok_count"
+  exit 1
+fi
+
+pre_fault_line_count=$(wc -l <"$LOG_DIR/client.log")
+
+kill "$TARGET_PID" >/dev/null 2>&1 || true
+TARGET_PID=""
+
+if ! wait_for_pattern "$LOG_DIR/master.log" 'action=unmount_nof_segment_by_heartbeat' $((HEARTBEAT_INTERVAL * HEARTBEAT_FAILURES + 20)); then
+  echo "master did not unmount nof segment by heartbeat"
+  exit 1
+fi
+
+post_unmount_line_count=$(wc -l <"$LOG_DIR/client.log")
+
+observation_deadline=$((SECONDS + 10))
+while (( SECONDS < observation_deadline )); do
+  post_unmount_successes=$(awk -v start="$post_unmount_line_count" '
+    NR > start && ($0 ~ /put_ok/ || $0 ~ /get_ok/) { count++ }
+    END { print count + 0 }
+  ' "$LOG_DIR/client.log")
+  post_unmount_failures=$(awk -v start="$post_unmount_line_count" '
+    NR > start && ($0 ~ /put_fail/ || $0 ~ /get_fail/) { count++ }
+    END { print count + 0 }
+  ' "$LOG_DIR/client.log")
+
+  if [[ "$CLIENT_GLOBAL_SEGMENT_SIZE" -eq 0 ]]; then
+    (( post_unmount_failures > 0 )) && break
+  else
+    (( post_unmount_successes > 0 )) && break
+  fi
+  sleep 1
+done
+
+kill "$CLIENT_PID" >/dev/null 2>&1 || true
+wait "$CLIENT_PID" >/dev/null 2>&1 || true
+CLIENT_PID=""
+
+post_fault_successes=$(awk -v start="$pre_fault_line_count" '
+  NR > start && ($0 ~ /put_ok/ || $0 ~ /get_ok/) { count++ }
+  END { print count + 0 }
+' "$LOG_DIR/client.log")
+
+post_fault_failures=$(awk -v start="$pre_fault_line_count" '
+  NR > start && ($0 ~ /put_fail/ || $0 ~ /get_fail/) { count++ }
+  END { print count + 0 }
+' "$LOG_DIR/client.log")
+
+post_unmount_successes=$(awk -v start="$post_unmount_line_count" '
+  NR > start && ($0 ~ /put_ok/ || $0 ~ /get_ok/) { count++ }
+  END { print count + 0 }
+' "$LOG_DIR/client.log")
+
+post_unmount_failures=$(awk -v start="$post_unmount_line_count" '
+  NR > start && ($0 ~ /put_fail/ || $0 ~ /get_fail/) { count++ }
+  END { print count + 0 }
+' "$LOG_DIR/client.log")
+
+if [[ "$CLIENT_GLOBAL_SEGMENT_SIZE" -eq 0 ]]; then
+  if [[ "$post_unmount_failures" -le 0 ]]; then
+    echo "expected IO failures after target down in nof-only mode"
+    exit 1
+  fi
+else
+  if [[ "$post_unmount_successes" -le 0 ]]; then
+    echo "no successful IO observed after target down"
+    exit 1
+  fi
+fi
+
+{
+  echo "=== hugepages ==="
+  cat "$LOG_DIR/hugepages.log"
+  echo "=== register ==="
+  cat "$LOG_DIR/register.log"
+  echo "=== client ==="
+  cat "$LOG_DIR/client.log"
+  echo "=== master tail ==="
+  tail -n 200 "$LOG_DIR/master.log"
+  echo "=== metadata tail ==="
+  tail -n 80 "$LOG_DIR/metadata.log"
+  echo "=== target tail ==="
+  tail -n 120 "$LOG_DIR/target.log"
+  echo "=== verdict ==="
+  echo "post_fault_successes=$post_fault_successes"
+  echo "post_fault_failures=$post_fault_failures"
+  echo "post_unmount_successes=$post_unmount_successes"
+  echo "post_unmount_failures=$post_unmount_failures"
+  echo "pre_fault_put_ok=$put_ok_count"
+  echo "pre_fault_get_ok=$get_ok_count"
+  echo "pre_fault_line_count=$pre_fault_line_count"
+  echo "post_unmount_line_count=$post_unmount_line_count"
+  echo "client_global_segment_size=$CLIENT_GLOBAL_SEGMENT_SIZE"
+  echo "client_memory_replica_num=$CLIENT_MEMORY_REPLICA_NUM"
+  echo "client_nof_replica_num=$CLIENT_NOF_REPLICA_NUM"
+} >"$LOG_DIR/summary.log"
+
+cat "$LOG_DIR/summary.log"
diff --git a/mooncake-store/tests/e2e/store_client_e2e.py b/mooncake-store/tests/e2e/store_client_e2e.py
new file mode 100644
index 00000000..9e061677
--- /dev/null
+++ b/mooncake-store/tests/e2e/store_client_e2e.py
@@ -0,0 +1,108 @@
+#!/usr/bin/env python3
+import argparse
+import os
+import sys
+import time
+
+
+def make_payload(seq: int, size: int) -> bytes:
+    seed = f"value{seq}".encode()
+    repeat = (size // len(seed)) + 1
+    return (seed * repeat)[:size]
+
+
+def main() -> int:
+    parser = argparse.ArgumentParser(
+        description="Continuous MooncakeDistributedStore put/get workload"
+    )
+    parser.add_argument("--local-hostname", default="127.0.0.1:50071")
+    parser.add_argument("--metadata-server", default="http://127.0.0.1:8080/metadata")
+    parser.add_argument("--master-server", default="127.0.0.1:50051")
+    parser.add_argument("--protocol", default="tcp")
+    parser.add_argument("--device-name", default="")
+    parser.add_argument("--global-segment-size", type=int, default=64 * 1024 * 1024)
+    parser.add_argument("--local-buffer-size", type=int, default=32 * 1024 * 1024)
+    parser.add_argument("--payload-size", type=int, default=4096)
+    parser.add_argument("--duration-sec", type=int, default=20)
+    parser.add_argument("--sleep-ms", type=int, default=200)
+    parser.add_argument("--key-prefix", default="nof-e2e")
+    parser.add_argument("--memory-replica-num", type=int, default=1)
+    parser.add_argument("--nof-replica-num", type=int, default=1)
+    args = parser.parse_args()
+
+    try:
+        import store  # type: ignore
+    except Exception as exc:
+        print(f"import_fail {exc}", flush=True)
+        return 10
+
+    mc = store.MooncakeDistributedStore()
+    setup_ret = mc.setup(
+        args.local_hostname,
+        args.metadata_server,
+        args.global_segment_size,
+        args.local_buffer_size,
+        args.protocol,
+        args.device_name,
+        args.master_server,
+    )
+    print(f"setup_ret {setup_ret}", flush=True)
+    if setup_ret != 0:
+        return setup_ret
+
+    replicate_config = store.ReplicateConfig()
+    replicate_config.replica_num = args.memory_replica_num
+    replicate_config.nof_replica_num = args.nof_replica_num
+
+    deadline = time.time() + args.duration_sec if args.duration_sec > 0 else None
+    seq = 0
+    put_ok = 0
+    get_ok = 0
+    put_fail = 0
+    get_fail = 0
+    mismatch = 0
+
+    while deadline is None or time.time() < deadline:
+        seq += 1
+        key = f"{args.key_prefix}-{seq}"
+        expected = make_payload(seq, args.payload_size)
+
+        put_ret = mc.put(key, expected, replicate_config)
+        if put_ret == 0:
+            put_ok += 1
+            print(f"put_ok seq={seq} key={key} len={len(expected)}", flush=True)
+        else:
+            put_fail += 1
+            print(f"put_fail seq={seq} key={key} ret={put_ret}", flush=True)
+            time.sleep(args.sleep_ms / 1000.0)
+            continue
+
+        actual = mc.get(key)
+        if actual == expected:
+            get_ok += 1
+            print(f"get_ok seq={seq} key={key} len={len(actual)}", flush=True)
+        elif actual in (b"", None):
+            get_fail += 1
+            actual_len = 0 if actual in (b"", None) else len(actual)
+            print(f"get_fail seq={seq} key={key} len={actual_len}", flush=True)
+        else:
+            mismatch += 1
+            print(
+                f"data_mismatch seq={seq} key={key} expected_len={len(expected)} actual_len={len(actual)}",
+                flush=True,
+            )
+            return 20
+
+        time.sleep(args.sleep_ms / 1000.0)
+
+    print(
+        "summary "
+        f"put_ok={put_ok} put_fail={put_fail} "
+        f"get_ok={get_ok} get_fail={get_fail} mismatch={mismatch}",
+        flush=True,
+    )
+    return 0 if mismatch == 0 else 20
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())
diff --git a/mooncake-store/tests/file_storage_promotion_test.cpp b/mooncake-store/tests/file_storage_promotion_test.cpp
new file mode 100644
index 00000000..a553d92d
--- /dev/null
+++ b/mooncake-store/tests/file_storage_promotion_test.cpp
@@ -0,0 +1,402 @@
+// Unit tests for FileStorage::ProcessPromotionTasks. Uses a stub Client
+// that overrides the 4 promotion-RPC entry points so we can drive the
+// orchestration loop deterministically without standing up a master.
+
+#include <glog/logging.h>
+#include <gtest/gtest.h>
+
+#include <atomic>
+#include <filesystem>
+#include <set>
+#include <thread>
+
+#include "client_service.h"
+#include "file_storage.h"
+#include "storage_backend.h"
+
+namespace mooncake {
+
+namespace fs_test {
+
+// Programmable fake. Each Set* method records what the next call should
+// return; the call itself records that it happened so tests can assert on
+// invocation count and ordering.
+class FakeClient : public Client {
+   public:
+    FakeClient()
+        : Client(/*local_hostname=*/"localhost:9003",
+                 /*metadata_connstring=*/"",
+                 /*protocol=*/"tcp",
+                 /*labels=*/{}) {}
+
+    // Drives the queue returned to the heartbeat caller.
+    std::vector<PromotionTaskItem> heartbeat_queue;
+    tl::expected<void, ErrorCode> heartbeat_result =
+        tl::expected<void, ErrorCode>{};
+
+    tl::expected<void, ErrorCode> PromotionObjectHeartbeat(
+        std::vector<PromotionTaskItem>& promotion_objects) override {
+        heartbeat_calls.fetch_add(1);
+        if (!heartbeat_result.has_value()) {
+            return tl::make_unexpected(heartbeat_result.error());
+        }
+        // Mirror production master: PromotionObjectHeartbeat returns at
+        // most kMaxPerHeartbeat keys per call (see
+        // MasterService::PromotionObjectHeartbeat) and leaves the rest
+        // queued for subsequent calls. Tests iterate by calling
+        // ProcessPromotionTasks multiple times until heartbeat_queue is
+        // empty.
+        constexpr size_t kMaxPerHeartbeat = 1;
+        promotion_objects.clear();
+        while (promotion_objects.size() < kMaxPerHeartbeat &&
+               !heartbeat_queue.empty()) {
+            promotion_objects.push_back(std::move(heartbeat_queue.back()));
+            heartbeat_queue.pop_back();
+        }
+        return {};
+    }
+
+    // PromotionAllocStart: per-key dispatch table. Keys absent from
+    // alloc_overrides return the default response.
+    std::unordered_map<std::string, ErrorCode> alloc_overrides;
+    PromotionAllocStartResponse default_alloc_response;
+    bool default_alloc_succeeds = true;
+
+    tl::expected<PromotionAllocStartResponse, ErrorCode> PromotionAllocStart(
+        const std::string& key, uint64_t size,
+        const std::vector<std::string>& preferred_segments) override {
+        return PromotionAllocStart(key, "default", size, preferred_segments);
+    }
+
+    tl::expected<PromotionAllocStartResponse, ErrorCode> PromotionAllocStart(
+        const std::string& key, const std::string& tenant_id, uint64_t size,
+        const std::vector<std::string>& preferred_segments) override {
+        (void)tenant_id;
+        (void)size;
+        (void)preferred_segments;
+        alloc_calls.fetch_add(1);
+        last_alloc_key = key;
+        auto it = alloc_overrides.find(key);
+        if (it != alloc_overrides.end()) {
+            return tl::make_unexpected(it->second);
+        }
+        if (!default_alloc_succeeds) {
+            return tl::make_unexpected(ErrorCode::NO_AVAILABLE_HANDLE);
+        }
+        return default_alloc_response;
+    }
+
+    // PromotionWrite: per-key dispatch.
+    std::unordered_map<std::string, ErrorCode> write_overrides;
+    ErrorCode default_write_result = ErrorCode::OK;
+
+    ErrorCode PromotionWrite(const Replica::Descriptor&,
+                             std::vector<Slice>&) override {
+        write_calls.fetch_add(1);
+        auto it = write_overrides.find(last_alloc_key);
+        if (it != write_overrides.end()) {
+            return it->second;
+        }
+        return default_write_result;
+    }
+
+    std::unordered_map<std::string, ErrorCode> notify_overrides;
+    tl::expected<void, ErrorCode> default_notify_result =
+        tl::expected<void, ErrorCode>{};
+
+    tl::expected<void, ErrorCode> NotifyPromotionSuccess(
+        const std::string& key) override {
+        return NotifyPromotionSuccess(key, "default");
+    }
+
+    tl::expected<void, ErrorCode> NotifyPromotionSuccess(
+        const std::string& key, const std::string& tenant_id) override {
+        (void)tenant_id;
+        notify_calls.fetch_add(1);
+        notify_keys.push_back(key);
+        auto it = notify_overrides.find(key);
+        if (it != notify_overrides.end()) {
+            return tl::make_unexpected(it->second);
+        }
+        if (!default_notify_result.has_value()) {
+            return tl::make_unexpected(default_notify_result.error());
+        }
+        return {};
+    }
+
+    // NotifyPromotionFailure: records calls so tests can assert that
+    // post-AllocStart failure paths in ProcessPromotionTasks actually
+    // notify the master.
+    tl::expected<void, ErrorCode> NotifyPromotionFailure(
+        const std::string& key) override {
+        return NotifyPromotionFailure(key, "default");
+    }
+
+    tl::expected<void, ErrorCode> NotifyPromotionFailure(
+        const std::string& key, const std::string& tenant_id) override {
+        (void)tenant_id;
+        notify_failure_calls.fetch_add(1);
+        notify_failure_keys.push_back(key);
+        return {};
+    }
+
+    std::atomic<int> heartbeat_calls{0};
+    std::atomic<int> alloc_calls{0};
+    std::atomic<int> write_calls{0};
+    std::atomic<int> notify_calls{0};
+    std::atomic<int> notify_failure_calls{0};
+    std::vector<std::string> notify_keys;
+    std::vector<std::string> notify_failure_keys;
+    std::string last_alloc_key;
+};
+
+}  // namespace fs_test
+
+class FileStoragePromotionTest : public ::testing::Test {
+   protected:
+    std::string data_path;
+    std::shared_ptr<fs_test::FakeClient> fake;
+    std::unique_ptr<FileStorage> file_storage;
+
+    void SetUp() override {
+        google::InitGoogleLogging("FileStoragePromotionTest");
+        FLAGS_logtostderr = true;
+        unsetenv("MOONCAKE_OFFLOAD_FILE_STORAGE_PATH");
+
+        data_path = std::filesystem::current_path().string() + "/data_prom";
+        std::filesystem::create_directories(data_path);
+        for (const auto& entry :
+             std::filesystem::directory_iterator(data_path)) {
+            if (entry.is_regular_file()) std::filesystem::remove(entry.path());
+        }
+
+        FileStorageConfig cfg = FileStorageConfig::FromEnvironment();
+        cfg.storage_filepath = data_path;
+        cfg.local_buffer_size = 4 * 1024 * 1024;
+        fake = std::make_shared<fs_test::FakeClient>();
+        file_storage =
+            std::make_unique<FileStorage>(cfg, fake, "localhost:9003");
+    }
+
+    void TearDown() override {
+        file_storage.reset();
+        fake.reset();
+        google::ShutdownGoogleLogging();
+        for (const auto& entry :
+             std::filesystem::directory_iterator(data_path)) {
+            if (entry.is_regular_file()) std::filesystem::remove(entry.path());
+        }
+    }
+
+    tl::expected<void, ErrorCode> CallProcessPromotionTasks() {
+        return file_storage->ProcessPromotionTasks();
+    }
+
+    // Drain a multi-key queue across multiple ticks. ProcessPromotionTasks
+    // caps work at 1 task/tick (heartbeat-safety) and the FakeClient's
+    // PromotionObjectHeartbeat mirrors production by clearing the queue on
+    // drain. So the test fixture has to push the remaining keys back
+    // between ticks, the same way the master would re-push if the gate
+    // re-fires. We track which key was processed via fake->last_alloc_key
+    // and remove it from the working set.
+    tl::expected<void, ErrorCode> DrainAllPromotionTasks(
+        std::unordered_map<std::string, int64_t> remaining) {
+        tl::expected<void, ErrorCode> last_res{};
+        while (!remaining.empty()) {
+            std::string before = fake->last_alloc_key;
+            fake->heartbeat_queue.clear();
+            fake->heartbeat_queue.reserve(remaining.size());
+            for (const auto& [key, size] : remaining) {
+                fake->heartbeat_queue.push_back(PromotionTaskItem{
+                    .tenant_id = "default", .key = key, .size = size});
+            }
+            last_res = CallProcessPromotionTasks();
+            if (!last_res.has_value()) return last_res;
+            if (fake->last_alloc_key == before ||
+                !remaining.contains(fake->last_alloc_key)) {
+                // No forward progress (e.g., empty effective queue / size<=0
+                // skip on every key); avoid infinite loop.
+                break;
+            }
+            remaining.erase(fake->last_alloc_key);
+        }
+        return last_res;
+    }
+};
+
+// Empty heartbeat queue: do nothing past the heartbeat call.
+TEST_F(FileStoragePromotionTest, EmptyQueueIsNoOp) {
+    fake->heartbeat_queue = {};
+    auto res = CallProcessPromotionTasks();
+    EXPECT_TRUE(res.has_value());
+    EXPECT_EQ(fake->heartbeat_calls.load(), 1);
+    EXPECT_EQ(fake->alloc_calls.load(), 0);
+    EXPECT_EQ(fake->write_calls.load(), 0);
+    EXPECT_EQ(fake->notify_calls.load(), 0);
+}
+
+// Heartbeat returns SEGMENT_NOT_FOUND (transient post-restart):
+// swallow silently, no error to caller.
+TEST_F(FileStoragePromotionTest, HeartbeatSegmentNotFoundIsBenign) {
+    fake->heartbeat_result = tl::make_unexpected(ErrorCode::SEGMENT_NOT_FOUND);
+    auto res = CallProcessPromotionTasks();
+    EXPECT_TRUE(res.has_value());
+    EXPECT_EQ(fake->alloc_calls.load(), 0);
+}
+
+// Heartbeat returns a non-benign error: propagate it.
+TEST_F(FileStoragePromotionTest, HeartbeatHardErrorPropagates) {
+    fake->heartbeat_result = tl::make_unexpected(ErrorCode::INTERNAL_ERROR);
+    auto res = CallProcessPromotionTasks();
+    ASSERT_FALSE(res.has_value());
+    EXPECT_EQ(res.error(), ErrorCode::INTERNAL_ERROR);
+    EXPECT_EQ(fake->alloc_calls.load(), 0);
+}
+
+// Non-positive size in queue: skip that key, continue.
+TEST_F(FileStoragePromotionTest, NonPositiveSizeSkipped) {
+    fake->heartbeat_queue = {
+        {.tenant_id = "default", .key = "k_bad", .size = 0},
+        {.tenant_id = "default", .key = "k_good", .size = 1024}};
+    auto res = CallProcessPromotionTasks();
+    EXPECT_TRUE(res.has_value());
+    // Only k_good should reach AllocStart.
+    EXPECT_EQ(fake->alloc_calls.load(), 1);
+    EXPECT_EQ(fake->last_alloc_key, "k_good");
+}
+
+// PromotionAllocStart fails (e.g. master out of DRAM): skip key, no
+// write, no notify, advance to next.
+TEST_F(FileStoragePromotionTest, AllocStartFailureSkipsKey) {
+    fake->alloc_overrides["k1"] = ErrorCode::NO_AVAILABLE_HANDLE;
+    auto res = DrainAllPromotionTasks({{"k1", 1024}, {"k2", 1024}});
+    EXPECT_TRUE(res.has_value());
+    EXPECT_EQ(fake->alloc_calls.load(), 2);
+    // k1's failure prevents its write+notify; k2 succeeds (then fails at
+    // BatchLoad because the SSD file doesn't exist).
+    EXPECT_LE(fake->write_calls.load(), 1);
+    EXPECT_LE(fake->notify_calls.load(), 1);
+}
+
+// BatchLoad failure (SSD file missing): no PromotionWrite, no Notify.
+// Master-side reaper handles the orphaned PROCESSING replica.
+TEST_F(FileStoragePromotionTest, BatchLoadFailureLeavesNoNotify) {
+    fake->heartbeat_queue = {
+        {.tenant_id = "default", .key = "k_missing", .size = 1024}};
+    // Default alloc succeeds; BatchLoad will fail because there's no file
+    // at data_path/k_missing for the storage backend to read.
+    auto res = CallProcessPromotionTasks();
+    EXPECT_TRUE(res.has_value());  // ProcessPromotionTasks itself is best-
+                                   // effort; per-key failures are warnings.
+    EXPECT_EQ(fake->alloc_calls.load(), 1);
+    EXPECT_EQ(fake->write_calls.load(), 0)
+        << "TransferWrite must not run if BatchLoad failed";
+    EXPECT_EQ(fake->notify_calls.load(), 0)
+        << "NotifyPromotionSuccess must not run if BatchLoad failed";
+}
+
+// PromotionWrite failure: no Notify.
+TEST_F(FileStoragePromotionTest, TransferWriteFailureLeavesNoNotify) {
+    fake->heartbeat_queue = {
+        {.tenant_id = "default", .key = "k_te_fail", .size = 1024}};
+    fake->default_write_result = ErrorCode::TRANSFER_FAIL;
+    auto res = CallProcessPromotionTasks();
+    EXPECT_TRUE(res.has_value());
+    EXPECT_EQ(fake->alloc_calls.load(), 1);
+    // Whether write_calls is 0 or 1 depends on whether BatchLoad succeeded
+    // first; either way Notify must not fire.
+    EXPECT_EQ(fake->notify_calls.load(), 0)
+        << "NotifyPromotionSuccess must not run on TransferWrite failure";
+}
+
+// NotifyPromotionSuccess failure: logged, but processing of remaining
+// keys continues.
+TEST_F(FileStoragePromotionTest, NotifyFailureDoesNotAbortBatch) {
+    fake->notify_overrides["k1"] = ErrorCode::OBJECT_NOT_FOUND;
+    auto res = DrainAllPromotionTasks({{"k1", 1024}, {"k2", 1024}});
+    EXPECT_TRUE(res.has_value());
+    EXPECT_EQ(fake->alloc_calls.load(), 2)
+        << "Both keys must be alloc-attempted regardless of k1's notify "
+        << "failure";
+}
+
+// Per-key independence: failures on one key don't prevent attempts on
+// others.
+TEST_F(FileStoragePromotionTest, PerKeyFailuresAreIndependent) {
+    fake->alloc_overrides["k_alloc_fail"] = ErrorCode::NO_AVAILABLE_HANDLE;
+    fake->write_overrides["k_write_fail"] = ErrorCode::TRANSFER_FAIL;
+    auto res = DrainAllPromotionTasks({
+        {"k_alloc_fail", 1024},
+        {"k_write_fail", 1024},
+        {"k_normal", 1024},
+    });
+    EXPECT_TRUE(res.has_value());
+    // All three reach AllocStart.
+    EXPECT_EQ(fake->alloc_calls.load(), 3);
+}
+
+// Every post-admission failure path (including PromotionAllocStart's
+// own failure) must call NotifyPromotionFailure so the master releases
+// the task slot immediately. Without this, the slot stays pinned until
+// put_start_release_timeout_sec_ (~10 min default), turning a transient
+// DRAM-pressure spike into a sustained outage of promotion_queue_limit_.
+TEST_F(FileStoragePromotionTest, AllocStartFailureNotifiesMaster) {
+    fake->heartbeat_queue = {
+        {.tenant_id = "default", .key = "k_alloc_fail", .size = 1024}};
+    fake->alloc_overrides["k_alloc_fail"] = ErrorCode::NO_AVAILABLE_HANDLE;
+    auto res = CallProcessPromotionTasks();
+    EXPECT_TRUE(res.has_value());
+    EXPECT_EQ(fake->alloc_calls.load(), 1);
+    EXPECT_EQ(fake->notify_calls.load(), 0)
+        << "Success-notify must not fire on AllocStart failure.";
+    EXPECT_EQ(fake->notify_failure_calls.load(), 1)
+        << "AllocStart failure must invoke NotifyPromotionFailure so the "
+        << "master can release the task slot immediately. Without this "
+        << "the slot is pinned for ~10 min and transient DRAM pressure "
+        << "saturates promotion_queue_limit_ for the same window.";
+    ASSERT_EQ(fake->notify_failure_keys.size(), 1u);
+    EXPECT_EQ(fake->notify_failure_keys[0], "k_alloc_fail");
+}
+
+// Same invariant on the post-AllocStart paths (BatchLoad /
+// TransferWrite / Notify-Success). Each failure mode must release the
+// master slot. Exercises three modes (AllocStart itself, missing-file
+// BatchLoad, override on Notify) by draining across successive
+// heartbeats (the FakeClient mirrors the master's kMaxPerHeartbeat = 1
+// cap) and verifies each one releases.
+TEST_F(FileStoragePromotionTest, PostAllocFailuresAllNotifyMaster) {
+    fake->alloc_overrides["k_alloc_fail"] = ErrorCode::NO_AVAILABLE_HANDLE;
+    // k_load_fail: no override -> AllocStart succeeds, but BatchLoad
+    // will fail because no SSD file exists for this key in data_path.
+    // k_notify_fail: AllocStart and BatchLoad and TransferWrite all
+    // succeed; Notify is overridden to fail.
+    fake->notify_overrides["k_notify_fail"] = ErrorCode::OBJECT_NOT_FOUND;
+
+    auto res = DrainAllPromotionTasks({
+        {"k_alloc_fail", 1024},
+        {"k_load_fail", 1024},
+        {"k_notify_fail", 1024},
+    });
+    EXPECT_TRUE(res.has_value());
+    EXPECT_EQ(fake->alloc_calls.load(), 3);
+
+    // Every failed key must have its slot released via Notify-Failure.
+    // k_notify_fail also counts: Notify-Success failed, so we still
+    // need to free the slot.
+    EXPECT_EQ(fake->notify_failure_calls.load(), 3)
+        << "All three failed promotions must each invoke "
+        << "NotifyPromotionFailure. If this fires with <3, one of the "
+        << "failure paths is still leaking the master slot.";
+    std::set<std::string> got(fake->notify_failure_keys.begin(),
+                              fake->notify_failure_keys.end());
+    EXPECT_EQ(got.count("k_alloc_fail"), 1u);
+    EXPECT_EQ(got.count("k_load_fail"), 1u);
+    EXPECT_EQ(got.count("k_notify_fail"), 1u);
+}
+
+}  // namespace mooncake
+
+int main(int argc, char** argv) {
+    ::testing::InitGoogleTest(&argc, argv);
+    return RUN_ALL_TESTS();
+}
diff --git a/mooncake-store/tests/ha/leadership/backends/k8s/high_availability_k8s_test.cpp b/mooncake-store/tests/ha/leadership/backends/k8s/high_availability_k8s_test.cpp
new file mode 100644
index 00000000..7978258a
--- /dev/null
+++ b/mooncake-store/tests/ha/leadership/backends/k8s/high_availability_k8s_test.cpp
@@ -0,0 +1,347 @@
+#include <gflags/gflags.h>
+
+#include <atomic>
+#include <chrono>
+#include <cstdlib>
+#include <future>
+#include <memory>
+#include <optional>
+#include <string>
+#include <thread>
+
+#ifdef STORE_USE_K8S_LEASE
+#include "k8s_lease_helper.h"
+#endif
+#include "ha/leadership/leader_coordinator_factory.h"
+#include "ha/leadership/high_availability_test_fixture.h"
+#include "types.h"
+
+namespace mooncake {
+namespace testing {
+
+DEFINE_string(k8s_namespace, "default",
+              "K8s namespace for HA integration tests");
+DEFINE_string(k8s_lease_name, "mooncake-ha-test",
+              "K8s Lease name for HA integration tests");
+
+namespace {
+
+std::optional<std::string> GetK8sSkipReason() {
+#ifdef STORE_USE_K8S_LEASE
+    // Probe: try to init K8s client and read a lease.
+    auto err = K8sLeaseHelper::Init();
+    if (err != ErrorCode::OK) {
+        return "K8s API not reachable (Init failed)";
+    }
+    std::string holder;
+    int64_t transitions = 0;
+    err = K8sLeaseHelper::GetHolder(FLAGS_k8s_namespace, FLAGS_k8s_lease_name,
+                                    holder, transitions);
+    if (err != ErrorCode::OK && err != ErrorCode::K8S_LEASE_NOT_FOUND) {
+        return "K8s API not reachable (GetHolder probe failed)";
+    }
+    return std::nullopt;
+#else
+    return "K8s Lease HA backend is not enabled in this build";
+#endif
+}
+
+ha::HABackendSpec MakeK8sBackendSpec(const std::string& ns,
+                                     const std::string& lease_name) {
+    return ha::HABackendSpec{
+        .type = ha::HABackendType::K8S,
+        .connstring = ns + "/" + lease_name,
+        .cluster_namespace = "",
+    };
+}
+
+std::unique_ptr<ha::LeaderCoordinator> CreateK8sCoordinatorOrNull(
+    const std::string& ns, const std::string& lease_name) {
+    auto coordinator =
+        ha::CreateLeaderCoordinator(MakeK8sBackendSpec(ns, lease_name));
+    if (!coordinator) {
+        return nullptr;
+    }
+    return std::move(coordinator.value());
+}
+
+std::string MakeK8sTestLeaseName(const std::string& suffix) {
+    return FLAGS_k8s_lease_name + "-" + suffix;
+}
+
+}  // namespace
+
+TEST_F(HighAvailabilityTest, K8sBasicMasterViewOperations) {
+    if (auto skip_reason = GetK8sSkipReason(); skip_reason.has_value()) {
+        GTEST_SKIP() << *skip_reason;
+    }
+
+    const auto lease_name = MakeK8sTestLeaseName("basic");
+    auto coordinator =
+        CreateK8sCoordinatorOrNull(FLAGS_k8s_namespace, lease_name);
+    ASSERT_NE(coordinator, nullptr);
+
+    // Initially, the master view should be empty (no leader)
+    auto initial_view = coordinator->ReadCurrentView();
+    ASSERT_TRUE(initial_view.has_value());
+    // Note: may or may not have value depending on prior test state
+
+    // Acquire leadership
+    auto acquire = coordinator->TryAcquireLeadership("127.0.0.1:8899");
+    ASSERT_TRUE(acquire.has_value());
+    ASSERT_EQ(ha::AcquireLeadershipStatus::ACQUIRED, acquire->status);
+    ASSERT_TRUE(acquire->session.has_value());
+
+    // Read current view — should show our address
+    auto current_view = coordinator->ReadCurrentView();
+    ASSERT_TRUE(current_view.has_value());
+    ASSERT_TRUE(current_view->has_value());
+    EXPECT_EQ("127.0.0.1:8899", current_view->value().leader_address);
+
+    // Renew
+    auto renewed = coordinator->RenewLeadership(*acquire->session);
+    ASSERT_TRUE(renewed.has_value());
+    EXPECT_TRUE(renewed.value());
+
+    // WaitForViewChange — should time out since nothing changed
+    auto no_change = coordinator->WaitForViewChange(
+        acquire->session->view.view_version, std::chrono::milliseconds(200));
+    ASSERT_TRUE(no_change.has_value());
+    ASSERT_FALSE(no_change->changed);
+    ASSERT_TRUE(no_change->timed_out);
+
+    // Release
+    ASSERT_EQ(ErrorCode::OK, coordinator->ReleaseLeadership(*acquire->session));
+
+    // After release, view should eventually show no leader or a different
+    // transitions count
+    auto released = coordinator->WaitForViewChange(
+        acquire->session->view.view_version, std::chrono::seconds(10));
+    ASSERT_TRUE(released.has_value());
+    ASSERT_TRUE(released->changed);
+}
+
+TEST_F(HighAvailabilityTest, K8sLeadershipMonitorIgnoresExplicitRelease) {
+    if (auto skip_reason = GetK8sSkipReason(); skip_reason.has_value()) {
+        GTEST_SKIP() << *skip_reason;
+    }
+
+    const auto lease_name = MakeK8sTestLeaseName("monitor-release");
+    auto coordinator =
+        CreateK8sCoordinatorOrNull(FLAGS_k8s_namespace, lease_name);
+    ASSERT_NE(coordinator, nullptr);
+
+    auto acquire = coordinator->TryAcquireLeadership("127.0.0.1:9922");
+    ASSERT_TRUE(acquire.has_value());
+    ASSERT_EQ(ha::AcquireLeadershipStatus::ACQUIRED, acquire->status);
+    ASSERT_TRUE(acquire->session.has_value());
+    const auto session = *acquire->session;
+
+    auto renew = coordinator->RenewLeadership(session);
+    ASSERT_TRUE(renew.has_value());
+    ASSERT_TRUE(renew.value());
+
+    auto callback_fired = std::make_shared<std::atomic<bool>>(false);
+    auto monitor = coordinator->StartLeadershipMonitor(
+        session, [callback_fired](ha::LeadershipLossReason) {
+            callback_fired->store(true);
+        });
+    ASSERT_TRUE(monitor.has_value());
+
+    // Explicit release should NOT fire the monitor callback
+    ASSERT_EQ(ErrorCode::OK, coordinator->ReleaseLeadership(session));
+    std::this_thread::sleep_for(std::chrono::seconds(2));
+    EXPECT_FALSE(callback_fired->load());
+}
+
+TEST_F(HighAvailabilityTest, K8sLeadershipMonitorReportsElectionCancellation) {
+    if (auto skip_reason = GetK8sSkipReason(); skip_reason.has_value()) {
+        GTEST_SKIP() << *skip_reason;
+    }
+
+    const auto lease_name = MakeK8sTestLeaseName("monitor-loss");
+    auto coordinator =
+        CreateK8sCoordinatorOrNull(FLAGS_k8s_namespace, lease_name);
+    ASSERT_NE(coordinator, nullptr);
+
+    auto acquire = coordinator->TryAcquireLeadership("127.0.0.1:9955");
+    ASSERT_TRUE(acquire.has_value());
+    ASSERT_EQ(ha::AcquireLeadershipStatus::ACQUIRED, acquire->status);
+    ASSERT_TRUE(acquire->session.has_value());
+    const auto session = *acquire->session;
+
+    auto renew = coordinator->RenewLeadership(session);
+    ASSERT_TRUE(renew.has_value());
+    ASSERT_TRUE(renew.value());
+
+    auto callback_fired = std::make_shared<std::atomic<bool>>(false);
+    auto monitor = coordinator->StartLeadershipMonitor(
+        session, [callback_fired](ha::LeadershipLossReason) {
+            callback_fired->store(true);
+        });
+    ASSERT_TRUE(monitor.has_value());
+
+    // Cancel the election goroutine directly, bypassing ReleaseLeadership.
+    // This simulates leadership loss: WaitLost returns while
+    // election_shutdown_requested_ is still false, so the monitor fires.
+    K8sLeaseHelper::CancelElection(FLAGS_k8s_namespace, lease_name);
+
+    // Wait for the monitor thread to detect the loss and fire the callback
+    for (int i = 0; i < 50 && !callback_fired->load(); ++i) {
+        std::this_thread::sleep_for(std::chrono::milliseconds(200));
+    }
+    EXPECT_TRUE(callback_fired->load());
+}
+
+TEST_F(HighAvailabilityTest, K8sCanReacquireAfterRelease) {
+    if (auto skip_reason = GetK8sSkipReason(); skip_reason.has_value()) {
+        GTEST_SKIP() << *skip_reason;
+    }
+
+    const auto lease_name = MakeK8sTestLeaseName("reacquire");
+    auto coordinator =
+        CreateK8sCoordinatorOrNull(FLAGS_k8s_namespace, lease_name);
+    ASSERT_NE(coordinator, nullptr);
+
+    // First acquisition
+    auto first_acquire = coordinator->TryAcquireLeadership("127.0.0.1:9933");
+    ASSERT_TRUE(first_acquire.has_value());
+    ASSERT_EQ(ha::AcquireLeadershipStatus::ACQUIRED, first_acquire->status);
+    ASSERT_TRUE(first_acquire->session.has_value());
+
+    auto first_renew = coordinator->RenewLeadership(*first_acquire->session);
+    ASSERT_TRUE(first_renew.has_value());
+    ASSERT_TRUE(first_renew.value());
+
+    ASSERT_EQ(ErrorCode::OK,
+              coordinator->ReleaseLeadership(*first_acquire->session));
+
+    // Wait for lease to expire
+    std::this_thread::sleep_for(std::chrono::seconds(2));
+
+    // Second acquisition
+    auto second_acquire = coordinator->TryAcquireLeadership("127.0.0.1:9944");
+    ASSERT_TRUE(second_acquire.has_value());
+    ASSERT_EQ(ha::AcquireLeadershipStatus::ACQUIRED, second_acquire->status);
+    ASSERT_TRUE(second_acquire->session.has_value());
+
+    auto second_renew = coordinator->RenewLeadership(*second_acquire->session);
+    ASSERT_TRUE(second_renew.has_value());
+    ASSERT_TRUE(second_renew.value());
+
+    ASSERT_EQ(ErrorCode::OK,
+              coordinator->ReleaseLeadership(*second_acquire->session));
+}
+
+TEST_F(HighAvailabilityTest, K8sContendedLeadershipAndHandover) {
+    if (auto skip_reason = GetK8sSkipReason(); skip_reason.has_value()) {
+        GTEST_SKIP() << *skip_reason;
+    }
+
+    const auto lease_name = MakeK8sTestLeaseName("contention");
+
+    // Coordinator A acquires leadership
+    auto coord_a = CreateK8sCoordinatorOrNull(FLAGS_k8s_namespace, lease_name);
+    ASSERT_NE(coord_a, nullptr);
+
+    auto acquire_a = coord_a->TryAcquireLeadership("127.0.0.1:7701");
+    ASSERT_TRUE(acquire_a.has_value());
+    ASSERT_EQ(ha::AcquireLeadershipStatus::ACQUIRED, acquire_a->status);
+    ASSERT_TRUE(acquire_a->session.has_value());
+
+    auto renew_a = coord_a->RenewLeadership(*acquire_a->session);
+    ASSERT_TRUE(renew_a.has_value());
+    ASSERT_TRUE(renew_a.value());
+
+    // Coordinator B attempts acquisition on the same lease — should be
+    // contended
+    auto coord_b = CreateK8sCoordinatorOrNull(FLAGS_k8s_namespace, lease_name);
+    ASSERT_NE(coord_b, nullptr);
+
+    auto acquire_b = coord_b->TryAcquireLeadership("127.0.0.1:7702");
+    ASSERT_TRUE(acquire_b.has_value());
+    EXPECT_EQ(ha::AcquireLeadershipStatus::CONTENDED, acquire_b->status);
+    ASSERT_TRUE(acquire_b->observed_view.has_value());
+    EXPECT_EQ(acquire_a->session->view.view_version,
+              acquire_b->observed_view->view_version);
+
+    // A releases leadership
+    ASSERT_EQ(ErrorCode::OK, coord_a->ReleaseLeadership(*acquire_a->session));
+
+    // Wait for lease to expire so B can take over
+    auto view_changed = coord_b->WaitForViewChange(
+        acquire_a->session->view.view_version, std::chrono::seconds(10));
+    ASSERT_TRUE(view_changed.has_value());
+    ASSERT_TRUE(view_changed->changed);
+
+    // B acquires leadership
+    auto acquire_b2 = coord_b->TryAcquireLeadership("127.0.0.1:7702");
+    ASSERT_TRUE(acquire_b2.has_value());
+    ASSERT_EQ(ha::AcquireLeadershipStatus::ACQUIRED, acquire_b2->status);
+    ASSERT_TRUE(acquire_b2->session.has_value());
+
+    ASSERT_EQ(ErrorCode::OK, coord_b->ReleaseLeadership(*acquire_b2->session));
+}
+
+// --- Connstring validation ---
+
+TEST_F(HighAvailabilityTest, K8sConnstringEmptyIsRejected) {
+    auto coordinator = CreateK8sCoordinatorOrNull("", "");
+    ASSERT_EQ(coordinator, nullptr);
+}
+
+TEST_F(HighAvailabilityTest, K8sConnstringEmptyNamespaceIsRejected) {
+    auto spec = ha::HABackendSpec{
+        .type = ha::HABackendType::K8S,
+        .connstring = "/lease-name",
+    };
+    auto coordinator = ha::CreateLeaderCoordinator(spec);
+    ASSERT_FALSE(coordinator.has_value());
+    EXPECT_EQ(ErrorCode::INVALID_PARAMS, coordinator.error());
+}
+
+TEST_F(HighAvailabilityTest, K8sConnstringEmptyLeaseNameIsRejected) {
+    auto spec = ha::HABackendSpec{
+        .type = ha::HABackendType::K8S,
+        .connstring = "default/",
+    };
+    auto coordinator = ha::CreateLeaderCoordinator(spec);
+    ASSERT_FALSE(coordinator.has_value());
+    EXPECT_EQ(ErrorCode::INVALID_PARAMS, coordinator.error());
+}
+
+TEST_F(HighAvailabilityTest, K8sConnstringMultipleSlashesIsRejected) {
+    auto spec = ha::HABackendSpec{
+        .type = ha::HABackendType::K8S,
+        .connstring = "a/b/c",
+    };
+    auto coordinator = ha::CreateLeaderCoordinator(spec);
+    ASSERT_FALSE(coordinator.has_value());
+    EXPECT_EQ(ErrorCode::INVALID_PARAMS, coordinator.error());
+}
+
+TEST_F(HighAvailabilityTest, K8sConnstringValidFormatPassesParsing) {
+    if (auto skip_reason = GetK8sSkipReason(); skip_reason.has_value()) {
+        GTEST_SKIP() << *skip_reason;
+    }
+
+    auto coordinator =
+        CreateK8sCoordinatorOrNull(FLAGS_k8s_namespace, "parse-valid");
+    ASSERT_NE(coordinator, nullptr);
+}
+
+TEST_F(HighAvailabilityTest, K8sConnstringNoSlashDefaultsNamespace) {
+    if (auto skip_reason = GetK8sSkipReason(); skip_reason.has_value()) {
+        GTEST_SKIP() << *skip_reason;
+    }
+
+    auto spec = ha::HABackendSpec{
+        .type = ha::HABackendType::K8S,
+        .connstring = FLAGS_k8s_lease_name + "-parse-noslash",
+    };
+    auto coordinator = ha::CreateLeaderCoordinator(spec);
+    ASSERT_TRUE(coordinator.has_value());
+}
+
+}  // namespace testing
+}  // namespace mooncake
diff --git a/mooncake-store/tests/ha/leadership/etcd_leader_hang_e2e.sh b/mooncake-store/tests/ha/leadership/etcd_leader_hang_e2e.sh
new file mode 100755
index 00000000..c0913e90
--- /dev/null
+++ b/mooncake-store/tests/ha/leadership/etcd_leader_hang_e2e.sh
@@ -0,0 +1,173 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+ROOT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/../../../../.." && pwd)"
+WORK_DIR="${WORK_DIR:-/tmp/mooncake-etcd-leader-hang-e2e}"
+ETCD_BIN="${ETCD_BIN:-etcd}"
+ETCDCTL_BIN="${ETCDCTL_BIN:-etcdctl}"
+MASTER_BIN="${MASTER_BIN:-${ROOT_DIR}/build/mooncake-store/src/mooncake_master}"
+CLUSTER_ID="${CLUSTER_ID:-mooncake_cluster}"
+MASTER_VIEW_KEY="${MASTER_VIEW_KEY:-mooncake-store/${CLUSTER_ID}/master_view}"
+MASTER_ADDR="${MASTER_ADDR:-127.0.0.1}"
+MASTER_PORT="${MASTER_PORT:-50051}"
+MASTER_METRICS_PORT="${MASTER_METRICS_PORT:-9003}"
+ENDPOINTS="127.0.0.1:23791,127.0.0.1:23792,127.0.0.1:23793"
+ETCD_QUORUM_SIZE="${ETCD_QUORUM_SIZE:-2}"
+
+mkdir -p "${WORK_DIR}"
+
+ETCD_PIDS=()
+MASTER_PID=""
+
+cleanup() {
+  set +e
+  if [[ -n "${MASTER_PID}" ]]; then
+    kill "${MASTER_PID}" 2>/dev/null || true
+  fi
+  for pid in "${ETCD_PIDS[@]}"; do
+    kill -CONT "${pid}" 2>/dev/null || true
+    kill "${pid}" 2>/dev/null || true
+  done
+  wait 2>/dev/null || true
+  rm -rf "${WORK_DIR}"
+}
+trap cleanup EXIT
+
+start_etcd_node() {
+  local name="$1"
+  local client_port="$2"
+  local peer_port="$3"
+  local data_dir="${WORK_DIR}/${name}"
+
+  "${ETCD_BIN}" \
+    --name "${name}" \
+    --data-dir "${data_dir}" \
+    --listen-client-urls "http://127.0.0.1:${client_port}" \
+    --advertise-client-urls "http://127.0.0.1:${client_port}" \
+    --listen-peer-urls "http://127.0.0.1:${peer_port}" \
+    --initial-advertise-peer-urls "http://127.0.0.1:${peer_port}" \
+    --initial-cluster "etcd1=http://127.0.0.1:23891,etcd2=http://127.0.0.1:23892,etcd3=http://127.0.0.1:23893" \
+    --initial-cluster-state new \
+    >"${WORK_DIR}/${name}.log" 2>&1 &
+  ETCD_PIDS+=("$!")
+}
+
+healthy_endpoints_csv() {
+  local endpoints=()
+  IFS=',' read -r -a all_endpoints <<<"${ENDPOINTS}"
+  for endpoint in "${all_endpoints[@]}"; do
+    if "${ETCDCTL_BIN}" --endpoints="${endpoint}" endpoint health >/dev/null 2>&1; then
+      endpoints+=("${endpoint}")
+    fi
+  done
+  local IFS=,
+  echo "${endpoints[*]}"
+}
+
+healthy_endpoint_count() {
+  local healthy
+  healthy="$(healthy_endpoints_csv)"
+  if [[ -z "${healthy}" ]]; then
+    echo 0
+    return
+  fi
+  awk -F',' '{print NF}' <<<"${healthy}"
+}
+
+dump_logs() {
+  set +e
+  echo "---- mooncake_master.log ----" >&2
+  tail -n 80 "${WORK_DIR}/mooncake_master.log" 2>/dev/null >&2 || true
+  for name in etcd1 etcd2 etcd3; do
+    echo "---- ${name}.log ----" >&2
+    tail -n 40 "${WORK_DIR}/${name}.log" 2>/dev/null >&2 || true
+  done
+}
+
+wait_for_etcd() {
+  for _ in $(seq 1 60); do
+    if (( $(healthy_endpoint_count) >= ETCD_QUORUM_SIZE )); then
+      return 0
+    fi
+    sleep 1
+  done
+  echo "etcd quorum did not become healthy" >&2
+  dump_logs
+  return 1
+}
+
+leader_pid() {
+  local live_endpoints
+  live_endpoints="$(healthy_endpoints_csv)"
+  if [[ -z "${live_endpoints}" ]]; then
+    echo "no healthy etcd endpoints" >&2
+    return 1
+  fi
+
+  local leader_client_url
+  leader_client_url="$("${ETCDCTL_BIN}" --endpoints="${live_endpoints}" endpoint status -w json \
+    | python3 -c 'import json,sys; data=json.load(sys.stdin); leaders=[x["Endpoint"] for x in data if x["Status"].get("leader")==x["Status"].get("header",{}).get("member_id")]; print(leaders[0] if leaders else "")')"
+
+  case "${leader_client_url}" in
+    *23791*) echo "${ETCD_PIDS[0]}" ;;
+    *23792*) echo "${ETCD_PIDS[1]}" ;;
+    *23793*) echo "${ETCD_PIDS[2]}" ;;
+    *) echo "unknown leader endpoint ${leader_client_url}" >&2; return 1 ;;
+  esac
+}
+
+wait_for_master_view() {
+  local expected="$1"
+  for _ in $(seq 1 90); do
+    local live_endpoints
+    live_endpoints="$(healthy_endpoints_csv)"
+    if [[ -z "${live_endpoints}" ]]; then
+      sleep 1
+      continue
+    fi
+
+    local value
+    value="$("${ETCDCTL_BIN}" --endpoints="${live_endpoints}" get "${MASTER_VIEW_KEY}" --print-value-only 2>/dev/null || true)"
+    if [[ "${value}" == "${expected}" ]]; then
+      return 0
+    fi
+    sleep 1
+  done
+  echo "master_view ${MASTER_VIEW_KEY} did not recover to ${expected}" >&2
+  dump_logs
+  return 1
+}
+
+start_master() {
+  "${MASTER_BIN}" \
+    --enable_ha=true \
+    --ha_backend_type=etcd \
+    --etcd_endpoints="${ENDPOINTS}" \
+    --cluster_id="${CLUSTER_ID}" \
+    --rpc_address="${MASTER_ADDR}" \
+    --rpc_port="${MASTER_PORT}" \
+    --metrics_port="${MASTER_METRICS_PORT}" \
+    >"${WORK_DIR}/mooncake_master.log" 2>&1 &
+  MASTER_PID="$!"
+}
+
+start_etcd_node etcd1 23791 23891
+start_etcd_node etcd2 23792 23892
+start_etcd_node etcd3 23793 23893
+wait_for_etcd
+
+start_master
+wait_for_master_view "${MASTER_ADDR}:${MASTER_PORT}"
+
+STOPPED_LEADER_PID="$(leader_pid)"
+kill -STOP "${STOPPED_LEADER_PID}"
+wait_for_etcd
+wait_for_master_view "${MASTER_ADDR}:${MASTER_PORT}"
+kill -CONT "${STOPPED_LEADER_PID}"
+
+KILLED_LEADER_PID="$(leader_pid)"
+kill -KILL "${KILLED_LEADER_PID}"
+wait_for_etcd
+wait_for_master_view "${MASTER_ADDR}:${MASTER_PORT}"
+
+echo "etcd leader hang e2e passed"
diff --git a/mooncake-store/tests/ha/leadership/ha_backend_availability_test.cpp b/mooncake-store/tests/ha/leadership/ha_backend_availability_test.cpp
new file mode 100644
index 00000000..e301a05f
--- /dev/null
+++ b/mooncake-store/tests/ha/leadership/ha_backend_availability_test.cpp
@@ -0,0 +1,63 @@
+#include "ha/ha_types.h"
+
+#include <gtest/gtest.h>
+#include <ylt/util/tl/expected.hpp>
+
+namespace mooncake {
+
+tl::expected<std::optional<ha::HABackendSpec>, ErrorCode> ParseHABackendSpec(
+    const std::string& master_server_entry);
+
+namespace ha {
+namespace {
+
+TEST(HABackendAvailabilityTest, UnknownBackendIsInvalid) {
+    EXPECT_EQ(ErrorCode::INVALID_PARAMS,
+              ValidateHABackendAvailability(HABackendType::UNKNOWN));
+}
+
+TEST(HABackendAvailabilityTest, EtcdAvailabilityMatchesBuildFlag) {
+#ifdef STORE_USE_ETCD
+    EXPECT_EQ(ErrorCode::OK,
+              ValidateHABackendAvailability(HABackendType::ETCD));
+#else
+    EXPECT_EQ(ErrorCode::UNAVAILABLE_IN_CURRENT_MODE,
+              ValidateHABackendAvailability(HABackendType::ETCD));
+#endif
+}
+
+TEST(HABackendAvailabilityTest, RedisAvailabilityMatchesBuildFlag) {
+#ifdef STORE_USE_REDIS
+    EXPECT_EQ(ErrorCode::OK,
+              ValidateHABackendAvailability(HABackendType::REDIS));
+#else
+    EXPECT_EQ(ErrorCode::UNAVAILABLE_IN_CURRENT_MODE,
+              ValidateHABackendAvailability(HABackendType::REDIS));
+#endif
+}
+
+TEST(HABackendAvailabilityTest, K8sAvailabilityMatchesBuildFlag) {
+#ifdef STORE_USE_K8S_LEASE
+    EXPECT_EQ(ErrorCode::OK, ValidateHABackendAvailability(HABackendType::K8S));
+#else
+    EXPECT_EQ(ErrorCode::UNAVAILABLE_IN_CURRENT_MODE,
+              ValidateHABackendAvailability(HABackendType::K8S));
+#endif
+}
+
+TEST(HABackendAvailabilityTest, ClientSpecParsingMatchesK8sBuildFlag) {
+    auto spec = ParseHABackendSpec("k8s://default/master");
+#ifdef STORE_USE_K8S_LEASE
+    ASSERT_TRUE(spec.has_value());
+    ASSERT_TRUE(spec->has_value());
+    EXPECT_EQ(HABackendType::K8S, (*spec)->type);
+    EXPECT_EQ("default/master", (*spec)->connstring);
+#else
+    ASSERT_FALSE(spec.has_value());
+    EXPECT_EQ(ErrorCode::UNAVAILABLE_IN_CURRENT_MODE, spec.error());
+#endif
+}
+
+}  // namespace
+}  // namespace ha
+}  // namespace mooncake
diff --git a/mooncake-store/tests/ha/leadership/high_availability_test.cpp b/mooncake-store/tests/ha/leadership/high_availability_test.cpp
index 4b91ecda..dbd2c5fe 100644
--- a/mooncake-store/tests/ha/leadership/high_availability_test.cpp
+++ b/mooncake-store/tests/ha/leadership/high_availability_test.cpp
@@ -7,6 +7,7 @@
 #include <mutex>
 #include <optional>
 #include <string>
+#include <thread>
 
 #ifdef STORE_USE_ETCD
 #include "etcd_helper.h"
@@ -317,6 +318,120 @@ TEST_F(HighAvailabilityTest, BasicMasterViewOperations) {
 
 #ifdef STORE_USE_ETCD
 
+// WaitForViewChange must return promptly when the leader's master_view key is
+// deleted, driven by the etcd watch rather than the timeout. We give it a long
+// (5s) timeout but delete the key after ~300ms; a watch-based implementation
+// returns shortly after the deletion, well before the timeout would fire.
+TEST_F(HighAvailabilityTest, WaitForViewChangeReturnsPromptlyOnLeaderLoss) {
+    if (auto skip_reason = GetEtcdSkipReason(); skip_reason.has_value()) {
+        GTEST_SKIP() << *skip_reason;
+    }
+
+    auto coordinator = CreateEtcdCoordinatorOrNull(FLAGS_etcd_endpoints);
+    ASSERT_NE(coordinator, nullptr);
+
+    auto acquire = coordinator->TryAcquireLeadership("0.0.0.0:5555");
+    ASSERT_TRUE(acquire.has_value());
+    ASSERT_EQ(ha::AcquireLeadershipStatus::ACQUIRED, acquire->status);
+    ASSERT_TRUE(acquire->session.has_value());
+    const auto session = *acquire->session;
+
+    auto renew = coordinator->RenewLeadership(session);
+    ASSERT_TRUE(renew.has_value());
+    ASSERT_TRUE(renew.value());
+
+    // Release leadership from another thread after a short delay; this revokes
+    // the lease and deletes the master_view key, which the watch observes.
+    std::thread releaser([&]() {
+        std::this_thread::sleep_for(std::chrono::milliseconds(300));
+        coordinator->ReleaseLeadership(session);
+    });
+
+    const auto start = std::chrono::steady_clock::now();
+    auto changed = coordinator->WaitForViewChange(session.view.view_version,
+                                                  std::chrono::seconds(5));
+    const auto elapsed = std::chrono::steady_clock::now() - start;
+    releaser.join();
+
+    ASSERT_TRUE(changed.has_value());
+    EXPECT_TRUE(changed->changed);
+    EXPECT_FALSE(changed->current_view.has_value());
+    // Must be driven by the watch (key deleted at ~300ms), not the 5s timeout.
+    EXPECT_LT(elapsed, std::chrono::seconds(3));
+}
+
+// WaitForViewChange must honor its timeout when the leader is stable: the watch
+// blocks with no deletion event, and the timer cancels it so the call returns
+// timed_out at roughly the requested deadline -- neither returning early nor
+// hanging past it.
+TEST_F(HighAvailabilityTest, WaitForViewChangeTimesOutWhenStable) {
+    if (auto skip_reason = GetEtcdSkipReason(); skip_reason.has_value()) {
+        GTEST_SKIP() << *skip_reason;
+    }
+
+    auto coordinator = CreateEtcdCoordinatorOrNull(FLAGS_etcd_endpoints);
+    ASSERT_NE(coordinator, nullptr);
+
+    auto acquire = coordinator->TryAcquireLeadership("0.0.0.0:4444");
+    ASSERT_TRUE(acquire.has_value());
+    ASSERT_EQ(ha::AcquireLeadershipStatus::ACQUIRED, acquire->status);
+    ASSERT_TRUE(acquire->session.has_value());
+    const auto session = *acquire->session;
+
+    auto renew = coordinator->RenewLeadership(session);
+    ASSERT_TRUE(renew.has_value());
+    ASSERT_TRUE(renew.value());
+
+    const auto timeout = std::chrono::milliseconds(500);
+    const auto start = std::chrono::steady_clock::now();
+    auto result =
+        coordinator->WaitForViewChange(session.view.view_version, timeout);
+    const auto elapsed = std::chrono::steady_clock::now() - start;
+
+    ASSERT_TRUE(result.has_value());
+    EXPECT_FALSE(result->changed);
+    EXPECT_TRUE(result->timed_out);
+    // Did not return early (watch did not spuriously fire) ...
+    EXPECT_GE(elapsed, std::chrono::milliseconds(400));
+    // ... and did not hang past the deadline (timer cancelled the watch).
+    EXPECT_LT(elapsed, std::chrono::seconds(3));
+
+    ASSERT_EQ(ErrorCode::OK, coordinator->ReleaseLeadership(session));
+}
+
+// When the observed view already differs from the caller's known version,
+// WaitForViewChange returns promptly: it arms the watch (now established
+// synchronously via WithCreatedNotify) and then the initial read short-circuits
+// on the version mismatch before any event is awaited. Passing no known version
+// while a leader exists is one such case.
+TEST_F(HighAvailabilityTest, WaitForViewChangeReturnsCurrentViewImmediately) {
+    if (auto skip_reason = GetEtcdSkipReason(); skip_reason.has_value()) {
+        GTEST_SKIP() << *skip_reason;
+    }
+
+    auto coordinator = CreateEtcdCoordinatorOrNull(FLAGS_etcd_endpoints);
+    ASSERT_NE(coordinator, nullptr);
+
+    auto acquire = coordinator->TryAcquireLeadership("0.0.0.0:3333");
+    ASSERT_TRUE(acquire.has_value());
+    ASSERT_EQ(ha::AcquireLeadershipStatus::ACQUIRED, acquire->status);
+    ASSERT_TRUE(acquire->session.has_value());
+    const auto session = *acquire->session;
+
+    const auto start = std::chrono::steady_clock::now();
+    auto result =
+        coordinator->WaitForViewChange(std::nullopt, std::chrono::seconds(5));
+    const auto elapsed = std::chrono::steady_clock::now() - start;
+
+    ASSERT_TRUE(result.has_value());
+    EXPECT_TRUE(result->changed);
+    ASSERT_TRUE(result->current_view.has_value());
+    EXPECT_EQ(result->current_view->view_version, session.view.view_version);
+    EXPECT_LT(elapsed, std::chrono::seconds(1));
+
+    ASSERT_EQ(ErrorCode::OK, coordinator->ReleaseLeadership(session));
+}
+
 TEST_F(HighAvailabilityTest, LeadershipMonitorReportsKeepAliveLoss) {
     if (auto skip_reason = GetEtcdSkipReason(); skip_reason.has_value()) {
         GTEST_SKIP() << *skip_reason;
@@ -354,6 +469,229 @@ TEST_F(HighAvailabilityTest, LeadershipMonitorReportsKeepAliveLoss) {
     ASSERT_EQ(ErrorCode::OK, coordinator->ReleaseLeadership(session));
 }
 
+TEST_F(HighAvailabilityTest, EtcdStoreClientResetKeepsBasicOperationsWorking) {
+    if (auto skip_reason = GetEtcdSkipReason(); skip_reason.has_value()) {
+        GTEST_SKIP() << *skip_reason;
+    }
+
+    const std::string key =
+        FLAGS_etcd_test_key_prefix + "reset_basic_operation";
+    const std::string value_before = "before_reset";
+    const std::string value_after = "after_reset";
+
+    ASSERT_EQ(ErrorCode::OK,
+              EtcdHelper::Put(key.c_str(), key.size(), value_before.c_str(),
+                              value_before.size()));
+
+    std::string got;
+    EtcdRevisionId rev = 0;
+    ASSERT_EQ(ErrorCode::OK,
+              EtcdHelper::Get(key.c_str(), key.size(), got, rev));
+    ASSERT_EQ(value_before, got);
+
+    ASSERT_EQ(ErrorCode::OK,
+              EtcdHelper::ResetEtcdStoreClient(FLAGS_etcd_endpoints));
+
+    ASSERT_EQ(ErrorCode::OK,
+              EtcdHelper::Put(key.c_str(), key.size(), value_after.c_str(),
+                              value_after.size()));
+    ASSERT_EQ(ErrorCode::OK,
+              EtcdHelper::Get(key.c_str(), key.size(), got, rev));
+    ASSERT_EQ(value_after, got);
+
+    EtcdLeaseId lease_id = 0;
+    ASSERT_EQ(ErrorCode::OK, EtcdHelper::GrantLease(10, lease_id));
+    ASSERT_EQ(ErrorCode::OK, EtcdHelper::RevokeLease(lease_id));
+}
+
+TEST_F(HighAvailabilityTest, EtcdStoreClientResetStopsOldKeepAlive) {
+    if (auto skip_reason = GetEtcdSkipReason(); skip_reason.has_value()) {
+        GTEST_SKIP() << *skip_reason;
+    }
+
+    EtcdLeaseId lease_id = 0;
+    ASSERT_EQ(ErrorCode::OK, EtcdHelper::GrantLease(10, lease_id));
+
+    std::promise<ErrorCode> promise;
+    auto future = promise.get_future();
+    std::thread keep_alive_thread(
+        [&]() { promise.set_value(EtcdHelper::KeepAlive(lease_id)); });
+
+    auto cleanup_keep_alive = [&]() {
+        if (keep_alive_thread.joinable()) {
+            (void)EtcdHelper::CancelKeepAlive(lease_id);
+            keep_alive_thread.join();
+        }
+    };
+
+    auto ready_err = EtcdHelper::WaitKeepAliveReady(lease_id, 1000);
+    if (ready_err != ErrorCode::OK) {
+        cleanup_keep_alive();
+    }
+    ASSERT_EQ(ErrorCode::OK, ready_err);
+
+    auto reset_err = EtcdHelper::ResetEtcdStoreClient(FLAGS_etcd_endpoints);
+    if (reset_err != ErrorCode::OK) {
+        cleanup_keep_alive();
+    }
+    ASSERT_EQ(ErrorCode::OK, reset_err);
+
+    auto keep_alive_status = future.wait_for(std::chrono::seconds(5));
+    if (keep_alive_status != std::future_status::ready) {
+        cleanup_keep_alive();
+    }
+    ASSERT_EQ(keep_alive_status, std::future_status::ready);
+    EXPECT_NE(ErrorCode::OK, future.get());
+    keep_alive_thread.join();
+
+    EtcdLeaseId new_lease_id = 0;
+    ASSERT_EQ(ErrorCode::OK, EtcdHelper::GrantLease(10, new_lease_id));
+    ASSERT_EQ(ErrorCode::OK, EtcdHelper::RevokeLease(new_lease_id));
+}
+
+// Static data for prefix watch callback (C-linkage required for Go callback)
+static std::atomic<int> g_prefix_watch_event_count{0};
+static std::atomic<int> g_prefix_watch_broken_count{0};
+
+extern "C" void PrefixWatchTestCallback(void* /*ctx*/, const char* /*key*/,
+                                        size_t /*key_size*/,
+                                        const char* /*value*/,
+                                        size_t /*value_size*/, int event_type,
+                                        int64_t /*mod_revision*/) {
+    if (event_type == 2) {
+        g_prefix_watch_broken_count.fetch_add(1);
+    } else {
+        g_prefix_watch_event_count.fetch_add(1);
+    }
+}
+
+bool WaitForAtomicAtLeast(const std::atomic<int>& value, int expected,
+                          std::chrono::milliseconds timeout) {
+    const auto deadline = std::chrono::steady_clock::now() + timeout;
+    while (std::chrono::steady_clock::now() < deadline) {
+        if (value.load() >= expected) {
+            return true;
+        }
+        std::this_thread::sleep_for(std::chrono::milliseconds(50));
+    }
+    return value.load() >= expected;
+}
+
+TEST_F(HighAvailabilityTest,
+       EtcdStoreClientResetTriggersPrefixWatchBrokenAndReconnect) {
+    if (auto skip_reason = GetEtcdSkipReason(); skip_reason.has_value()) {
+        GTEST_SKIP() << *skip_reason;
+    }
+
+    const std::string prefix =
+        FLAGS_etcd_test_key_prefix + "reset_prefix_watch/";
+    const std::string key1 = prefix + "key1";
+    const std::string value1 = "value1";
+    const std::string value2 = "value2";
+
+    // Clean up any leftover keys
+    std::string prefix_end = prefix;
+    if (!prefix_end.empty()) prefix_end.back()++;
+    (void)EtcdHelper::DeleteRange(prefix.c_str(), prefix.size(),
+                                  prefix_end.c_str(), prefix_end.size());
+
+    // Reset counters and start prefix watch
+    g_prefix_watch_event_count.store(0);
+    g_prefix_watch_broken_count.store(0);
+
+    ASSERT_EQ(ErrorCode::OK,
+              EtcdHelper::WatchWithPrefixFromRevision(
+                  prefix.c_str(), prefix.size(), /*start_revision=*/0,
+                  /*callback_context=*/nullptr, PrefixWatchTestCallback));
+
+    // Allow watch to establish
+    std::this_thread::sleep_for(std::chrono::milliseconds(500));
+
+    // Write first value -- watch should receive it
+    ASSERT_EQ(ErrorCode::OK, EtcdHelper::Put(key1.c_str(), key1.size(),
+                                             value1.c_str(), value1.size()));
+
+    std::this_thread::sleep_for(std::chrono::milliseconds(500));
+    EXPECT_GE(g_prefix_watch_event_count.load(), 1);
+
+    // Trigger reset
+    ASSERT_EQ(ErrorCode::OK,
+              EtcdHelper::ResetEtcdStoreClient(FLAGS_etcd_endpoints));
+
+    // Wait for WATCH_BROKEN callback
+    ASSERT_TRUE(WaitForAtomicAtLeast(g_prefix_watch_broken_count, 1,
+                                     std::chrono::seconds(5)))
+        << "Expected WATCH_BROKEN after reset";
+    EXPECT_EQ(1, g_prefix_watch_broken_count.load());
+
+    // Wait for old goroutine to fully exit
+    ASSERT_EQ(ErrorCode::OK, EtcdHelper::WaitWatchWithPrefixStopped(
+                                 prefix.c_str(), prefix.size(), 5000));
+
+    // Start a new watch -- should succeed (old entry cleaned up by defer)
+    g_prefix_watch_event_count.store(0);
+    g_prefix_watch_broken_count.store(0);
+
+    ASSERT_EQ(ErrorCode::OK,
+              EtcdHelper::WatchWithPrefixFromRevision(
+                  prefix.c_str(), prefix.size(), /*start_revision=*/0,
+                  /*callback_context=*/nullptr, PrefixWatchTestCallback));
+
+    std::this_thread::sleep_for(std::chrono::milliseconds(500));
+
+    // Write second value -- new watch should receive it
+    ASSERT_EQ(ErrorCode::OK, EtcdHelper::Put(key1.c_str(), key1.size(),
+                                             value2.c_str(), value2.size()));
+
+    std::this_thread::sleep_for(std::chrono::milliseconds(500));
+    EXPECT_GE(g_prefix_watch_event_count.load(), 1)
+        << "New watch should receive events after reset";
+
+    // Clean up
+    ASSERT_EQ(ErrorCode::OK,
+              EtcdHelper::CancelWatchWithPrefix(prefix.c_str(), prefix.size()));
+    ASSERT_EQ(ErrorCode::OK, EtcdHelper::WaitWatchWithPrefixStopped(
+                                 prefix.c_str(), prefix.size(), 5000));
+}
+
+TEST_F(HighAvailabilityTest, EtcdStorePrefixWatchCancelDoesNotReportBroken) {
+    if (auto skip_reason = GetEtcdSkipReason(); skip_reason.has_value()) {
+        GTEST_SKIP() << *skip_reason;
+    }
+
+    const std::string prefix =
+        FLAGS_etcd_test_key_prefix + "cancel_prefix_watch/";
+    const std::string key = prefix + "key";
+    const std::string value = "value";
+
+    std::string prefix_end = prefix;
+    if (!prefix_end.empty()) prefix_end.back()++;
+    (void)EtcdHelper::DeleteRange(prefix.c_str(), prefix.size(),
+                                  prefix_end.c_str(), prefix_end.size());
+
+    g_prefix_watch_event_count.store(0);
+    g_prefix_watch_broken_count.store(0);
+
+    ASSERT_EQ(ErrorCode::OK,
+              EtcdHelper::WatchWithPrefixFromRevision(
+                  prefix.c_str(), prefix.size(), /*start_revision=*/0,
+                  /*callback_context=*/nullptr, PrefixWatchTestCallback));
+
+    ASSERT_EQ(ErrorCode::OK, EtcdHelper::Put(key.c_str(), key.size(),
+                                             value.c_str(), value.size()));
+    ASSERT_TRUE(WaitForAtomicAtLeast(g_prefix_watch_event_count, 1,
+                                     std::chrono::seconds(5)));
+
+    ASSERT_EQ(ErrorCode::OK,
+              EtcdHelper::CancelWatchWithPrefix(prefix.c_str(), prefix.size()));
+    ASSERT_EQ(ErrorCode::OK, EtcdHelper::WaitWatchWithPrefixStopped(
+                                 prefix.c_str(), prefix.size(), 5000));
+
+    std::this_thread::sleep_for(std::chrono::milliseconds(200));
+    EXPECT_EQ(0, g_prefix_watch_broken_count.load())
+        << "Explicit cancel should not report WATCH_BROKEN";
+}
+
 #endif
 
 TEST_F(HighAvailabilityTest, LeadershipMonitorIgnoresExplicitRelease) {
diff --git a/mooncake-store/tests/ha/leadership/leader_label_reconciler_test.cpp b/mooncake-store/tests/ha/leadership/leader_label_reconciler_test.cpp
new file mode 100644
index 00000000..799fd74b
--- /dev/null
+++ b/mooncake-store/tests/ha/leadership/leader_label_reconciler_test.cpp
@@ -0,0 +1,186 @@
+#include "ha/leadership/leader_label_reconciler.h"
+
+#include <gtest/gtest.h>
+
+#include <chrono>
+#include <condition_variable>
+#include <mutex>
+#include <optional>
+#include <thread>
+
+namespace mooncake {
+namespace ha {
+namespace {
+
+using namespace std::chrono_literals;
+
+// Records every apply call and can be told to fail the first N attempts to
+// simulate a flaky / unreachable K8s API server.
+class FakeLabelBackend {
+   public:
+    LeaderLabelReconciler::ApplyFn Fn() {
+        return [this](bool desired) { return Apply(desired); };
+    }
+
+    void FailNext(int count) {
+        std::lock_guard<std::mutex> lock(mutex_);
+        fail_remaining_ = count;
+    }
+
+    void CommitThenFailNext(int count) {
+        std::lock_guard<std::mutex> lock(mutex_);
+        commit_then_fail_remaining_ = count;
+    }
+
+    // Waits until the last successful apply matches `want`, or times out.
+    bool WaitForApplied(bool want, std::chrono::milliseconds timeout) {
+        std::unique_lock<std::mutex> lock(mutex_);
+        return cv_.wait_for(lock, timeout, [&] { return applied_ == want; });
+    }
+
+    bool WaitForCommitted(bool want, std::chrono::milliseconds timeout) {
+        std::unique_lock<std::mutex> lock(mutex_);
+        return cv_.wait_for(lock, timeout, [&] { return committed_ == want; });
+    }
+
+    int CallCount() {
+        std::lock_guard<std::mutex> lock(mutex_);
+        return call_count_;
+    }
+
+    std::optional<bool> Applied() {
+        std::lock_guard<std::mutex> lock(mutex_);
+        return applied_;
+    }
+
+    std::optional<bool> Committed() {
+        std::lock_guard<std::mutex> lock(mutex_);
+        return committed_;
+    }
+
+   private:
+    ErrorCode Apply(bool desired) {
+        std::lock_guard<std::mutex> lock(mutex_);
+        ++call_count_;
+        if (commit_then_fail_remaining_ > 0) {
+            --commit_then_fail_remaining_;
+            committed_ = desired;
+            cv_.notify_all();
+            return ErrorCode::K8S_LEASE_OPERATION_ERROR;
+        }
+        if (fail_remaining_ > 0) {
+            --fail_remaining_;
+            cv_.notify_all();
+            return ErrorCode::K8S_LEASE_OPERATION_ERROR;
+        }
+        committed_ = desired;
+        applied_ = desired;
+        cv_.notify_all();
+        return ErrorCode::OK;
+    }
+
+    std::mutex mutex_;
+    std::condition_variable cv_;
+    int call_count_ = 0;
+    int fail_remaining_ = 0;
+    int commit_then_fail_remaining_ = 0;
+    std::optional<bool> applied_;
+    std::optional<bool> committed_;
+};
+
+constexpr auto kRetry = 5ms;
+constexpr auto kTimeout = 2s;
+constexpr auto kSlowRetry = 10s;
+
+TEST(LeaderLabelReconcilerTest, ConvergesToLeaderAfterTransientFailures) {
+    FakeLabelBackend backend;
+    backend.FailNext(3);
+    LeaderLabelReconciler reconciler(/*enabled=*/true, backend.Fn(), kRetry);
+
+    reconciler.SetLeader(true);
+
+    ASSERT_TRUE(backend.WaitForApplied(true, kTimeout));
+    EXPECT_GT(backend.CallCount(), 3);
+}
+
+TEST(LeaderLabelReconcilerTest, ClearsLeaderAfterTransientFailures) {
+    FakeLabelBackend backend;
+    LeaderLabelReconciler reconciler(/*enabled=*/true, backend.Fn(), kRetry);
+
+    reconciler.SetLeader(true);
+    ASSERT_TRUE(backend.WaitForApplied(true, kTimeout));
+
+    backend.FailNext(3);
+    reconciler.SetLeader(false);
+
+    ASSERT_TRUE(backend.WaitForApplied(false, kTimeout));
+}
+
+TEST(LeaderLabelReconcilerTest, LatestDesiredStateWins) {
+    FakeLabelBackend backend;
+    LeaderLabelReconciler reconciler(/*enabled=*/true, backend.Fn(), kRetry);
+
+    reconciler.SetLeader(true);
+    reconciler.SetLeader(false);
+
+    ASSERT_TRUE(backend.WaitForApplied(false, kTimeout));
+    EXPECT_EQ(backend.Applied(), std::make_optional(false));
+}
+
+TEST(LeaderLabelReconcilerTest, StopsCallingOnceConverged) {
+    FakeLabelBackend backend;
+    LeaderLabelReconciler reconciler(/*enabled=*/true, backend.Fn(), kRetry);
+
+    reconciler.SetLeader(true);
+    ASSERT_TRUE(backend.WaitForApplied(true, kTimeout));
+
+    int settled = backend.CallCount();
+    std::this_thread::sleep_for(kRetry * 20);
+    EXPECT_EQ(backend.CallCount(), settled);
+}
+
+TEST(LeaderLabelReconcilerTest, DisabledNeverApplies) {
+    FakeLabelBackend backend;
+    LeaderLabelReconciler reconciler(/*enabled=*/false, backend.Fn(), kRetry);
+
+    reconciler.SetLeader(true);
+    std::this_thread::sleep_for(kRetry * 20);
+
+    EXPECT_EQ(backend.CallCount(), 0);
+    EXPECT_EQ(backend.Applied(), std::nullopt);
+}
+
+TEST(LeaderLabelReconcilerTest, AmbiguousSetFailureDoesNotLeaveStaleLeader) {
+    FakeLabelBackend backend;
+    backend.CommitThenFailNext(1);
+    LeaderLabelReconciler reconciler(/*enabled=*/true, backend.Fn(),
+                                     kSlowRetry);
+
+    reconciler.SetLeader(true);
+    ASSERT_TRUE(backend.WaitForCommitted(true, kTimeout));
+
+    reconciler.SetLeader(false);
+
+    ASSERT_TRUE(backend.WaitForCommitted(false, kTimeout));
+}
+
+TEST(LeaderLabelReconcilerTest, AmbiguousClearFailureDoesNotLeaveMissingLabel) {
+    FakeLabelBackend backend;
+    LeaderLabelReconciler reconciler(/*enabled=*/true, backend.Fn(),
+                                     kSlowRetry);
+
+    reconciler.SetLeader(true);
+    ASSERT_TRUE(backend.WaitForCommitted(true, kTimeout));
+
+    backend.CommitThenFailNext(1);
+    reconciler.SetLeader(false);
+    ASSERT_TRUE(backend.WaitForCommitted(false, kTimeout));
+
+    reconciler.SetLeader(true);
+
+    ASSERT_TRUE(backend.WaitForCommitted(true, kTimeout));
+}
+
+}  // namespace
+}  // namespace ha
+}  // namespace mooncake
diff --git a/mooncake-store/tests/ha/oplog/ha_recovery_test.cpp b/mooncake-store/tests/ha/oplog/ha_recovery_test.cpp
index 7775016a..d9fb1a8f 100644
--- a/mooncake-store/tests/ha/oplog/ha_recovery_test.cpp
+++ b/mooncake-store/tests/ha/oplog/ha_recovery_test.cpp
@@ -259,6 +259,7 @@ TEST_F(HaRecoveryTest, GC_NoSnapshot_IncompleteRecovery) {
     // This demonstrates data incompleteness without snapshot protection:
     // the applier sees a gap it cannot fill (seq 1~9 are GC'd).
     size_t applied = applier_->ApplyOpLogEntries(entries);
+    EXPECT_EQ(applied, 0);
 
     // Wait for gap timeout so pending entries eventually drain
     std::this_thread::sleep_for(std::chrono::seconds(4));
diff --git a/mooncake-store/tests/ha/snapshot/catalog/backends/embedded/embedded_snapshot_catalog_store_test.cpp b/mooncake-store/tests/ha/snapshot/catalog/backends/embedded/embedded_snapshot_catalog_store_test.cpp
index 528a6a42..4304732f 100644
--- a/mooncake-store/tests/ha/snapshot/catalog/backends/embedded/embedded_snapshot_catalog_store_test.cpp
+++ b/mooncake-store/tests/ha/snapshot/catalog/backends/embedded/embedded_snapshot_catalog_store_test.cpp
@@ -109,7 +109,7 @@ class EmbeddedSnapshotCatalogStoreTest : public ::testing::Test {
     }
 
     FakeSnapshotObjectStore backend_;
-    ha::backends::embedded::EmbeddedSnapshotCatalogStore store_{&backend_};
+    ha::backends::embedded::EmbeddedSnapshotCatalogStore store_{&backend_, ""};
 };
 
 TEST_F(EmbeddedSnapshotCatalogStoreTest, PublishAndGetLatestRoundTrip) {
@@ -205,7 +205,7 @@ TEST_F(EmbeddedSnapshotCatalogStoreTest,
 
 TEST(EmbeddedSnapshotCatalogStoreStandaloneTest,
      ListReturnsInvalidParamsWhenObjectStoreMissing) {
-    ha::backends::embedded::EmbeddedSnapshotCatalogStore store(nullptr);
+    ha::backends::embedded::EmbeddedSnapshotCatalogStore store(nullptr, "");
 
     auto snapshots = store.List(0);
     ASSERT_FALSE(snapshots.has_value());
diff --git a/mooncake-store/tests/ha/snapshot/catalog/backends/redis/redis_snapshot_catalog_store_test.cpp b/mooncake-store/tests/ha/snapshot/catalog/backends/redis/redis_snapshot_catalog_store_test.cpp
index fe21ac02..f22712d6 100644
--- a/mooncake-store/tests/ha/snapshot/catalog/backends/redis/redis_snapshot_catalog_store_test.cpp
+++ b/mooncake-store/tests/ha/snapshot/catalog/backends/redis/redis_snapshot_catalog_store_test.cpp
@@ -78,12 +78,12 @@ class FakeObjectStore final : public SnapshotObjectStore {
     std::unordered_map<std::string, std::string> objects_;
 };
 
-ha::SnapshotDescriptor MakeDescriptor(const std::string& snapshot_id) {
+ha::SnapshotDescriptor MakeDescriptor(const std::string& snapshot_root,
+                                      const std::string& snapshot_id) {
     ha::SnapshotDescriptor descriptor;
     descriptor.snapshot_id = snapshot_id;
-    descriptor.manifest_key =
-        "mooncake_master_snapshot/" + snapshot_id + "/manifest.txt";
-    descriptor.object_prefix = "mooncake_master_snapshot/" + snapshot_id + "/";
+    descriptor.manifest_key = snapshot_root + snapshot_id + "/manifest.txt";
+    descriptor.object_prefix = snapshot_root + snapshot_id + "/";
     descriptor.last_included_seq = 42;
     descriptor.producer_view_version = 7;
     descriptor.created_at_ms = 1700000000000;
@@ -122,6 +122,15 @@ class RedisSnapshotCatalogStoreTest : public ::testing::Test {
         ASSERT_NE(reply->type, REDIS_REPLY_ERROR);
     }
 
+    ha::SnapshotDescriptor MakeDesc(const std::string& snapshot_id) {
+        return MakeDescriptor(store_->GetSnapshotRoot(), snapshot_id);
+    }
+
+    std::string DescriptorKey(const std::string& snapshot_id) {
+        return ha::snapshot_catalog_store_detail::BuildDescriptorKey(
+            store_->GetSnapshotRoot(), snapshot_id);
+    }
+
     FakeObjectStore object_store_;
     std::string cluster_namespace_;
     std::unique_ptr<ha::backends::redis::RedisSnapshotCatalogStore> store_;
@@ -134,17 +143,15 @@ TEST_F(RedisSnapshotCatalogStoreTest, GetLatestReturnsEmptyWhenCatalogMissing) {
 }
 
 TEST_F(RedisSnapshotCatalogStoreTest, PublishListAndGetLatestRoundTrip) {
-    ASSERT_EQ(store_->Publish(MakeDescriptor("20240301_120000_001")),
-              ErrorCode::OK);
-    ASSERT_EQ(store_->Publish(MakeDescriptor("20240302_120000_001")),
-              ErrorCode::OK);
+    ASSERT_EQ(store_->Publish(MakeDesc("20240301_120000_001")), ErrorCode::OK);
+    ASSERT_EQ(store_->Publish(MakeDesc("20240302_120000_001")), ErrorCode::OK);
 
     auto latest = store_->GetLatest();
     ASSERT_TRUE(latest.has_value());
     ASSERT_TRUE(latest->has_value());
     EXPECT_EQ(latest->value().snapshot_id, "20240302_120000_001");
     EXPECT_EQ(latest->value().manifest_key,
-              "mooncake_master_snapshot/20240302_120000_001/manifest.txt");
+              store_->GetSnapshotRoot() + "20240302_120000_001/manifest.txt");
     EXPECT_EQ(latest->value().last_included_seq, 42u);
     EXPECT_EQ(latest->value().producer_view_version, 7u);
     EXPECT_EQ(latest->value().created_at_ms, 1700000000000);
@@ -161,10 +168,9 @@ TEST_F(RedisSnapshotCatalogStoreTest, PublishListAndGetLatestRoundTrip) {
 TEST_F(RedisSnapshotCatalogStoreTest,
        GetLatestReturnsErrorWhenDescriptorMissing) {
     const std::string snapshot_id = "20240302_120000_001";
-    ASSERT_EQ(store_->Publish(MakeDescriptor(snapshot_id)), ErrorCode::OK);
+    ASSERT_EQ(store_->Publish(MakeDesc(snapshot_id)), ErrorCode::OK);
 
-    object_store_.objects_.erase(
-        ha::snapshot_catalog_store_detail::BuildDescriptorKey(snapshot_id));
+    object_store_.objects_.erase(DescriptorKey(snapshot_id));
 
     auto latest = store_->GetLatest();
     ASSERT_FALSE(latest.has_value());
@@ -173,10 +179,9 @@ TEST_F(RedisSnapshotCatalogStoreTest,
 
 TEST_F(RedisSnapshotCatalogStoreTest, ListSkipsSnapshotsWhenDescriptorMissing) {
     const std::string snapshot_id = "20240302_120000_001";
-    ASSERT_EQ(store_->Publish(MakeDescriptor(snapshot_id)), ErrorCode::OK);
+    ASSERT_EQ(store_->Publish(MakeDesc(snapshot_id)), ErrorCode::OK);
 
-    object_store_.objects_.erase(
-        ha::snapshot_catalog_store_detail::BuildDescriptorKey(snapshot_id));
+    object_store_.objects_.erase(DescriptorKey(snapshot_id));
 
     auto snapshots = store_->List(0);
     ASSERT_TRUE(snapshots.has_value());
@@ -195,14 +200,10 @@ TEST(RedisSnapshotCatalogStoreStandaloneTest,
 
 TEST_F(RedisSnapshotCatalogStoreTest,
        ListSkipsUnreadableSnapshotsAndKeepsHealthyEntries) {
-    ASSERT_EQ(store_->Publish(MakeDescriptor("20240301_120000_001")),
-              ErrorCode::OK);
-    ASSERT_EQ(store_->Publish(MakeDescriptor("20240302_120000_001")),
-              ErrorCode::OK);
+    ASSERT_EQ(store_->Publish(MakeDesc("20240301_120000_001")), ErrorCode::OK);
+    ASSERT_EQ(store_->Publish(MakeDesc("20240302_120000_001")), ErrorCode::OK);
 
-    object_store_.objects_.erase(
-        ha::snapshot_catalog_store_detail::BuildDescriptorKey(
-            "20240302_120000_001"));
+    object_store_.objects_.erase(DescriptorKey("20240302_120000_001"));
 
     auto snapshots = store_->List(0);
     ASSERT_TRUE(snapshots.has_value());
@@ -212,10 +213,8 @@ TEST_F(RedisSnapshotCatalogStoreTest,
 
 TEST_F(RedisSnapshotCatalogStoreTest,
        DeleteUpdatesLatestAndDeletesPayloadPrefix) {
-    ASSERT_EQ(store_->Publish(MakeDescriptor("20240301_120000_001")),
-              ErrorCode::OK);
-    ASSERT_EQ(store_->Publish(MakeDescriptor("20240302_120000_001")),
-              ErrorCode::OK);
+    ASSERT_EQ(store_->Publish(MakeDesc("20240301_120000_001")), ErrorCode::OK);
+    ASSERT_EQ(store_->Publish(MakeDesc("20240302_120000_001")), ErrorCode::OK);
 
     ASSERT_EQ(store_->Delete("20240302_120000_001"), ErrorCode::OK);
 
diff --git a/mooncake-store/tests/ha/snapshot/catalog_backed_snapshot_provider_test.cpp b/mooncake-store/tests/ha/snapshot/catalog_backed_snapshot_provider_test.cpp
index 7068bdbf..071940c8 100644
--- a/mooncake-store/tests/ha/snapshot/catalog_backed_snapshot_provider_test.cpp
+++ b/mooncake-store/tests/ha/snapshot/catalog_backed_snapshot_provider_test.cpp
@@ -2,7 +2,9 @@
 #include <glog/logging.h>
 #include <gtest/gtest.h>
 
+#include <cstdint>
 #include <filesystem>
+#include <limits>
 #include <memory>
 #include <optional>
 #include <string>
@@ -42,7 +44,7 @@ class CatalogBackedSnapshotProviderTest
             GetParam(), object_store_.get(), cluster_id_, FLAGS_redis_endpoint);
         ASSERT_NE(catalog_store_, nullptr);
 
-        descriptor_ = MakeTestSnapshotDescriptor();
+        descriptor_ = MakeTestSnapshotDescriptor(cluster_id_);
     }
 
     void TearDown() override {
@@ -57,13 +59,42 @@ class CatalogBackedSnapshotProviderTest
         }
     }
 
-    void PublishSnapshotPayload() {
+    void PublishSnapshotPayload(
+        SnapshotMetadataFormat format = SnapshotMetadataFormat::kLegacy) {
         auto result = mooncake::test::PublishSnapshotPayload(
-            *object_store_, *catalog_store_, descriptor_);
+            *object_store_, *catalog_store_, descriptor_, UUID{1, 2},
+            kDefaultTestObjectKey, kDefaultTestDiskFilePath,
+            kDefaultTestObjectSize, format);
         ASSERT_TRUE(result.has_value()) << result.error();
         snapshot_published_ = true;
     }
 
+    // Loads the published snapshot and asserts the single default object
+    // round-trips intact, regardless of the metadata on-wire format.
+    void ExpectLoadsDefaultObject() {
+        auto provider = CreateProvider();
+        ASSERT_TRUE(provider.has_value()) << toString(provider.error());
+
+        auto snapshot = provider.value()->LoadLatestSnapshot(cluster_id_);
+        ASSERT_TRUE(snapshot.has_value()) << toString(snapshot.error());
+        ASSERT_TRUE(snapshot->has_value());
+        ASSERT_EQ(snapshot->value().metadata.size(), 1u);
+
+        const auto& [key, metadata] = snapshot->value().metadata.front();
+        EXPECT_EQ(key, kDefaultTestObjectKey);
+        EXPECT_EQ(metadata.client_id, (UUID{1, 2}));
+        EXPECT_EQ(metadata.size, kDefaultTestObjectSize);
+        ASSERT_EQ(metadata.replicas.size(), 1u);
+
+        const auto& replica = metadata.replicas.front();
+        EXPECT_EQ(replica.status, ReplicaStatus::COMPLETE);
+        ASSERT_TRUE(replica.is_disk_replica());
+        EXPECT_EQ(replica.get_disk_descriptor().file_path,
+                  kDefaultTestDiskFilePath);
+        EXPECT_EQ(replica.get_disk_descriptor().object_size,
+                  kDefaultTestObjectSize);
+    }
+
     tl::expected<std::unique_ptr<SnapshotProvider>, ErrorCode> CreateProvider()
         const {
         return CreateCatalogBackedSnapshotProvider(MakeSnapshotProviderConfig(
@@ -119,6 +150,60 @@ TEST_P(CatalogBackedSnapshotProviderTest, LoadLatestSnapshotRoundTrip) {
               kDefaultTestObjectSize);
 }
 
+// The master snapshot writer evolved its per-object metadata layout over time
+// (data_type field, trailing hard_pinned flag, trailing group_id). The standby
+// restore reader must accept every shape; otherwise it rejects the snapshot and
+// falls back to OpLog-only bootstrap. These tests pin each on-wire format.
+
+TEST_P(CatalogBackedSnapshotProviderTest, LoadLatestSnapshotWithDataTypeField) {
+    // 8 + replica_count: data_type packed right after replica_count.
+    PublishSnapshotPayload(SnapshotMetadataFormat::kDataTypeOnly);
+    ExpectLoadsDefaultObject();
+}
+
+TEST_P(CatalogBackedSnapshotProviderTest,
+       LoadLatestSnapshotWithHardPinnedField) {
+    // 8 + replica_count: trailing hard_pinned flag, no data_type. Exercises the
+    // type-based disambiguation (first replica is not a positive integer).
+    PublishSnapshotPayload(SnapshotMetadataFormat::kHardPinnedOnly);
+    ExpectLoadsDefaultObject();
+}
+
+TEST_P(CatalogBackedSnapshotProviderTest,
+       LoadLatestSnapshotWithDataTypeAndHardPinned) {
+    // 9 + replica_count: data_type + trailing hard_pinned.
+    PublishSnapshotPayload(SnapshotMetadataFormat::kDataTypeAndHardPinned);
+    ExpectLoadsDefaultObject();
+}
+
+TEST_P(CatalogBackedSnapshotProviderTest, LoadLatestSnapshotWithGroupId) {
+    // 10 + replica_count: current writer format (data_type + hard_pinned +
+    // trailing group_id). Regression test for the live snapshot restore
+    // failure against the latest metadata layout.
+    PublishSnapshotPayload(SnapshotMetadataFormat::kWithGroupId);
+    ExpectLoadsDefaultObject();
+}
+
+TEST_P(CatalogBackedSnapshotProviderTest, RejectsOverflowingReplicaCount) {
+    // A near-UINT32_MAX replica_count must not wrap the format-detection
+    // arithmetic into a valid-looking total and slip an out-of-bounds index
+    // past the size check. The entry packs zero replicas (array size 7) but
+    // declares UINT32_MAX, so every format must be rejected, not parsed.
+    auto published = mooncake::test::PublishSnapshotPayloadBytes(
+        *object_store_, *catalog_store_, descriptor_,
+        BuildMetadataPayloadWithDeclaredReplicaCount(
+            std::numeric_limits<uint32_t>::max()));
+    ASSERT_TRUE(published.has_value()) << published.error();
+    snapshot_published_ = true;
+
+    auto provider = CreateProvider();
+    ASSERT_TRUE(provider.has_value()) << toString(provider.error());
+
+    auto snapshot = provider.value()->LoadLatestSnapshot(cluster_id_);
+    ASSERT_FALSE(snapshot.has_value());
+    EXPECT_EQ(snapshot.error(), ErrorCode::DESERIALIZE_FAIL);
+}
+
 TEST_P(CatalogBackedSnapshotProviderTest, RejectsClusterMismatch) {
     PublishSnapshotPayload();
 
diff --git a/mooncake-store/tests/ha/snapshot/master_service_promotion_test_for_snapshot.cpp b/mooncake-store/tests/ha/snapshot/master_service_promotion_test_for_snapshot.cpp
new file mode 100644
index 00000000..20ca542b
--- /dev/null
+++ b/mooncake-store/tests/ha/snapshot/master_service_promotion_test_for_snapshot.cpp
@@ -0,0 +1,200 @@
+// Snapshot round-trip tests for the L2->L1 promotion-on-hit feature.
+//
+// The base fixture's TearDown automatically calls TestSnapshotAndRestore
+// against the service_ member, which:
+//   1. Persists the master state to local snapshot storage.
+//   2. Boots a fresh MasterService in restore mode against that snapshot.
+//   3. Asserts CaptureServiceState() matches before vs. after.
+//
+// LOCAL_DISK replica content is captured via CaptureServiceState's
+// CompareReplicaDescriptor (which compares LocalDiskDescriptor fields), and
+// LocalDiskSegmentState (offloading_objects map) is also captured. Anything
+// these tests put into those fields will round-trip; anything not captured
+// (e.g., per-shard PromotionTask map, per-segment promotion_objects map)
+// is, by design, transient — the master is allowed to drop it on restart
+// and let clients re-trigger via heartbeat.
+
+#include "master_service_test_for_snapshot_base.h"
+
+namespace mooncake::test {
+
+class MasterServicePromotionSnapshotTest
+    : public MasterServiceSnapshotTestBase {
+   protected:
+    static bool glog_initialized_;
+
+    void SetUp() override {
+        MasterServiceSnapshotTestBase::SetUp();
+        if (!glog_initialized_) {
+            google::InitGoogleLogging("MasterServicePromotionSnapshotTest");
+            FLAGS_logtostderr = true;
+            glog_initialized_ = true;
+        }
+    }
+
+    // Build a master with promotion-on-hit enabled. enable_offload is
+    // required (promotion gates on it). admission_threshold=1 keeps the
+    // gate first-touch so tests don't need many reads.
+    void CreateMasterServiceWithPromotion() {
+        auto config = MasterServiceConfig::builder().build();
+        // Builder doesn't expose promotion knobs. Use direct field
+        // assignment.
+        config.enable_offload = true;
+        config.promotion_on_hit = true;
+        config.promotion_admission_threshold = 1;
+        config.put_start_discard_timeout_sec = 0;
+        config.put_start_release_timeout_sec = 1;
+        service_ = std::make_unique<MasterService>(config);
+    }
+
+    Segment MakeMemSegment(const std::string& name, size_t base) const {
+        Segment s;
+        s.id = generate_uuid();
+        s.name = name;
+        s.base = base;
+        s.size = kDefaultSegmentSize;
+        s.te_endpoint = name;
+        return s;
+    }
+
+    UUID MountMemAndDisk(const std::string& seg_name, size_t base) {
+        Segment seg = MakeMemSegment(seg_name, base);
+        UUID client_id = generate_uuid();
+        auto mount = service_->MountSegment(seg, client_id);
+        EXPECT_TRUE(mount.has_value());
+        auto mount_ld = service_->MountLocalDiskSegment(client_id, true);
+        EXPECT_TRUE(mount_ld.has_value());
+        return client_id;
+    }
+
+    bool InjectLocalDiskReplica(const UUID& client_id, const std::string& key,
+                                int64_t size,
+                                const std::string& transport_endpoint) {
+        std::vector<OffloadTaskItem> tasks{
+            OffloadTaskItem{.tenant_id = "default", .key = key, .size = size}};
+        StorageObjectMetadata sm;
+        sm.bucket_id = 0;
+        sm.offset = 0;
+        sm.key_size = static_cast<int64_t>(key.size());
+        sm.data_size = size;
+        sm.transport_endpoint = transport_endpoint;
+        std::vector<StorageObjectMetadata> metas{sm};
+        return service_->NotifyOffloadSuccess(client_id, tasks, metas)
+            .has_value();
+    }
+};
+
+bool MasterServicePromotionSnapshotTest::glog_initialized_ = false;
+
+// A LOCAL_DISK replica created via the offload path round-trips intact:
+// after snapshot+restore, the descriptor type, client_id, object_size, and
+// transport_endpoint must all match.
+TEST_F(MasterServicePromotionSnapshotTest, LocalDiskReplicaRoundTrip) {
+    CreateMasterServiceWithPromotion();
+    UUID client_id = MountMemAndDisk("seg_a", kDefaultSegmentBase);
+
+    ASSERT_TRUE(
+        InjectLocalDiskReplica(client_id, "k_cold", 1024, "seg_a_endpoint"));
+
+    auto before = service_->GetReplicaList("k_cold", "default");
+    ASSERT_TRUE(before.has_value());
+    ASSERT_EQ(before->replicas.size(), 1u);
+    EXPECT_TRUE(before->replicas[0].is_local_disk_replica());
+}
+
+// A key with both LOCAL_DISK and a separate MEMORY replica round-trips.
+// Tests that mixed-replica metadata serialization handles both descriptor
+// variants in the same object.
+TEST_F(MasterServicePromotionSnapshotTest, MixedMemoryAndLocalDiskRoundTrip) {
+    CreateMasterServiceWithPromotion();
+    UUID client_id = MountMemAndDisk("seg_a", kDefaultSegmentBase);
+
+    // Put a MEMORY replica.
+    ReplicateConfig rc;
+    rc.replica_num = 1;
+    ASSERT_TRUE(service_->PutStart(client_id, "k_mixed", "default", 1024, rc)
+                    .has_value());
+    ASSERT_TRUE(
+        service_->PutEnd(client_id, "k_mixed", "default", ReplicaType::MEMORY)
+            .has_value());
+
+    // Add LOCAL_DISK alongside.
+    ASSERT_TRUE(
+        InjectLocalDiskReplica(client_id, "k_mixed", 1024, "seg_a_endpoint"));
+
+    auto descs = service_->GetReplicaList("k_mixed", "default");
+    ASSERT_TRUE(descs.has_value());
+    EXPECT_EQ(descs->replicas.size(), 2u);
+}
+
+// LocalDiskSegment.enable_offloading flag is preserved across snapshot.
+// (CaptureServiceState compares this explicitly.)
+TEST_F(MasterServicePromotionSnapshotTest,
+       LocalDiskSegmentEnableOffloadingPreserved) {
+    CreateMasterServiceWithPromotion();
+    UUID client_id = MountMemAndDisk("seg_a", kDefaultSegmentBase);
+    (void)client_id;
+    // No further action — the LocalDiskSegment exists with enable_offloading
+    // = true. TearDown's TestSnapshotAndRestore round-trips and compares.
+}
+
+// Multiple LOCAL_DISK holders, one DRAM segment. Each holder's segment
+// metadata must round-trip independently.
+TEST_F(MasterServicePromotionSnapshotTest, MultipleLocalDiskHoldersRoundTrip) {
+    CreateMasterServiceWithPromotion();
+
+    // Mount two DRAM+LOCAL_DISK pairs.
+    UUID client_a = MountMemAndDisk("seg_a", kDefaultSegmentBase);
+    UUID client_b =
+        MountMemAndDisk("seg_b", kDefaultSegmentBase + kDefaultSegmentSize);
+
+    ASSERT_TRUE(
+        InjectLocalDiskReplica(client_a, "k_a", 1024, "seg_a_endpoint"));
+    ASSERT_TRUE(
+        InjectLocalDiskReplica(client_b, "k_b", 2048, "seg_b_endpoint"));
+
+    auto a = service_->GetReplicaList("k_a", "default");
+    auto b = service_->GetReplicaList("k_b", "default");
+    ASSERT_TRUE(a.has_value());
+    ASSERT_TRUE(b.has_value());
+    EXPECT_EQ(a->replicas[0].get_local_disk_descriptor().client_id, client_a);
+    EXPECT_EQ(b->replicas[0].get_local_disk_descriptor().client_id, client_b);
+    EXPECT_EQ(a->replicas[0].get_local_disk_descriptor().object_size, 1024u);
+    EXPECT_EQ(b->replicas[0].get_local_disk_descriptor().object_size, 2048u);
+}
+
+// Persisting state with an in-flight PromotionTask must not crash and must
+// not mutate the visible replica state. The PromotionTask itself is
+// allowed to be dropped on restore (transient), but the LOCAL_DISK source
+// it references must come back unchanged.
+TEST_F(MasterServicePromotionSnapshotTest, InFlightPromotionTaskSnapshotSafe) {
+    CreateMasterServiceWithPromotion();
+    UUID client_id = MountMemAndDisk("seg_a", kDefaultSegmentBase);
+
+    ASSERT_TRUE(
+        InjectLocalDiskReplica(client_id, "k_cold", 1024, "seg_a_endpoint"));
+
+    // Trigger promotion gate to enqueue a task. This pins the source
+    // replica's refcnt and adds a per-shard PromotionTask plus a per-
+    // segment promotion_objects entry.
+    auto get = service_->GetReplicaList("k_cold", "default");
+    ASSERT_TRUE(get.has_value());
+
+    // The visible replica list should still expose exactly one LOCAL_DISK
+    // replica with correct content. Snapshot+restore happens in TearDown
+    // and must preserve this view.
+    EXPECT_EQ(get->replicas.size(), 1u);
+    EXPECT_TRUE(get->replicas[0].is_local_disk_replica());
+}
+
+// Promotion-on-hit can be configured on, off, or implicitly off (when
+// enable_offload is false). The persisted master config must not affect
+// replica state round-trip in any of these modes — only the runtime gate
+// behavior changes. This test asserts the off-by-feature-flag case:
+// LOCAL_DISK keys still round-trip even when promotion is disabled.
+}  // namespace mooncake::test
+
+int main(int argc, char** argv) {
+    ::testing::InitGoogleTest(&argc, argv);
+    return RUN_ALL_TESTS();
+}
diff --git a/mooncake-store/tests/ha/snapshot/master_service_ssd_test_for_snapshot.cpp b/mooncake-store/tests/ha/snapshot/master_service_ssd_test_for_snapshot.cpp
index 83887256..73c6a173 100644
--- a/mooncake-store/tests/ha/snapshot/master_service_ssd_test_for_snapshot.cpp
+++ b/mooncake-store/tests/ha/snapshot/master_service_ssd_test_for_snapshot.cpp
@@ -57,7 +57,7 @@ TEST_F(MasterServiceSSDSnapshotTest, PutEndBothReplica) {
     config.replica_num = 1;
 
     auto put_start_result =
-        service_->PutStart(client_id, key, slice_length, config);
+        service_->PutStart(client_id, key, "default", slice_length, config);
     ASSERT_TRUE(put_start_result.has_value());
     auto replicas = put_start_result.value();
     ASSERT_EQ(2, replicas.size());
@@ -70,17 +70,17 @@ TEST_F(MasterServiceSSDSnapshotTest, PutEndBothReplica) {
     EXPECT_TRUE(has_mem);
     EXPECT_TRUE(has_disk);
 
-    auto get_result = service_->GetReplicaList(key);
+    auto get_result = service_->GetReplicaList(key, "default");
     ASSERT_FALSE(get_result.has_value());
     EXPECT_EQ(ErrorCode::REPLICA_IS_NOT_READY, get_result.error());
 
     // PutEnd for both memory and disk
-    EXPECT_TRUE(
-        service_->PutEnd(client_id, key, ReplicaType::MEMORY).has_value());
-    EXPECT_TRUE(
-        service_->PutEnd(client_id, key, ReplicaType::DISK).has_value());
+    EXPECT_TRUE(service_->PutEnd(client_id, key, "default", ReplicaType::MEMORY)
+                    .has_value());
+    EXPECT_TRUE(service_->PutEnd(client_id, key, "default", ReplicaType::DISK)
+                    .has_value());
 
-    get_result = service_->GetReplicaList(key);
+    get_result = service_->GetReplicaList(key, "default");
     ASSERT_TRUE(get_result.has_value());
     EXPECT_EQ(2, get_result.value().replicas.size());
 
@@ -89,6 +89,65 @@ TEST_F(MasterServiceSSDSnapshotTest, PutEndBothReplica) {
     }
 }
 
+TEST_F(MasterServiceSSDSnapshotTest, RestorePreservesCacheTotalMetrics) {
+    MasterMetricManager::instance().reset_cache_total_nums();
+
+    CreateMasterServiceWithSSDFeat("/mnt/ssd");
+    EnsureSnapshotStores(service_.get());
+
+    constexpr size_t buffer = 0x330000000;
+    constexpr size_t size = 1024 * 1024 * 64;
+    Segment segment;
+    segment.id = generate_uuid();
+    segment.name = "test_segment_restore_cache_totals";
+    segment.base = buffer;
+    segment.size = size;
+    segment.te_endpoint = segment.name;
+    UUID client_id = generate_uuid();
+
+    ASSERT_TRUE(service_->MountSegment(segment, client_id).has_value());
+
+    std::string key = "restore_cache_total_metric_key";
+    ASSERT_TRUE(
+        service_->PutStart(client_id, key, "default", 1024, {.replica_num = 1})
+            .has_value());
+    EXPECT_TRUE(service_->PutEnd(client_id, key, "default", ReplicaType::MEMORY)
+                    .has_value());
+    EXPECT_TRUE(service_->PutEnd(client_id, key, "default", ReplicaType::DISK)
+                    .has_value());
+
+    auto& metrics = MasterMetricManager::instance();
+    using CacheHitStat = MasterMetricManager::CacheHitStat;
+    auto stats = metrics.calculate_cache_stats();
+    ASSERT_EQ(stats[CacheHitStat::MEMORY_TOTAL], 1);
+    ASSERT_EQ(stats[CacheHitStat::SSD_TOTAL], 1);
+
+    std::string snapshot_id = GenerateSnapshotId();
+    auto persist_result = CallPersistState(service_.get(), snapshot_id);
+    ASSERT_TRUE(persist_result.has_value())
+        << "Failed to persist state: " << persist_result.error().message;
+
+    metrics.reset_cache_total_nums();
+
+    ::setenv("MOONCAKE_MASTER_SERVICE_SNAPSHOT_TEST_SKIP_CLEANUP", "1", 1);
+    auto restore_config = MasterServiceConfig::builder()
+                              .set_enable_snapshot_restore(true)
+                              .set_snapshot_object_store_type("local")
+                              .set_root_fs_dir("/mnt/ssd")
+                              .build();
+    std::unique_ptr<MasterService> restored_service(
+        new MasterService(restore_config));
+    ::unsetenv("MOONCAKE_MASTER_SERVICE_SNAPSHOT_TEST_SKIP_CLEANUP");
+
+    stats = metrics.calculate_cache_stats();
+    EXPECT_EQ(stats[CacheHitStat::MEMORY_TOTAL], 1);
+    EXPECT_EQ(stats[CacheHitStat::SSD_TOTAL], 1);
+
+    auto get_result = restored_service->GetReplicaList(key, "default");
+    ASSERT_TRUE(get_result.has_value());
+    EXPECT_EQ(2, get_result.value().replicas.size());
+}
+
 TEST_F(MasterServiceSSDSnapshotTest, PutRevokeDiskReplica) {
     CreateMasterServiceWithSSDFeat("/mnt/ssd");
 
@@ -112,19 +171,21 @@ TEST_F(MasterServiceSSDSnapshotTest, PutRevokeDiskReplica) {
     config.replica_num = 1;
 
     ASSERT_TRUE(
-        service_->PutStart(client_id, key, slice_length, config).has_value());
-    EXPECT_TRUE(
-        service_->PutEnd(client_id, key, ReplicaType::MEMORY).has_value());
+        service_->PutStart(client_id, key, "default", slice_length, config)
+            .has_value());
+    EXPECT_TRUE(service_->PutEnd(client_id, key, "default", ReplicaType::MEMORY)
+                    .has_value());
 
-    auto get_result = service_->GetReplicaList(key);
+    auto get_result = service_->GetReplicaList(key, "default");
     ASSERT_TRUE(get_result.has_value());
     EXPECT_EQ(1, get_result.value().replicas.size());
     ASSERT_TRUE(get_result.value().replicas[0].is_memory_replica());
 
     EXPECT_TRUE(
-        service_->PutRevoke(client_id, key, ReplicaType::DISK).has_value());
+        service_->PutRevoke(client_id, key, "default", ReplicaType::DISK)
+            .has_value());
 
-    get_result = service_->GetReplicaList(key);
+    get_result = service_->GetReplicaList(key, "default");
     ASSERT_TRUE(get_result.has_value());
     EXPECT_EQ(1, get_result.value().replicas.size());
     ASSERT_TRUE(get_result.value().replicas[0].is_memory_replica());
@@ -153,17 +214,19 @@ TEST_F(MasterServiceSSDSnapshotTest, PutRevokeMemoryReplica) {
     config.replica_num = 1;
 
     ASSERT_TRUE(
-        service_->PutStart(client_id, key, slice_length, config).has_value());
+        service_->PutStart(client_id, key, "default", slice_length, config)
+            .has_value());
     EXPECT_TRUE(
-        service_->PutRevoke(client_id, key, ReplicaType::MEMORY).has_value());
+        service_->PutRevoke(client_id, key, "default", ReplicaType::MEMORY)
+            .has_value());
 
-    auto get_result = service_->GetReplicaList(key);
+    auto get_result = service_->GetReplicaList(key, "default");
     ASSERT_FALSE(get_result.has_value());
     EXPECT_EQ(ErrorCode::REPLICA_IS_NOT_READY, get_result.error());
 
-    EXPECT_TRUE(
-        service_->PutEnd(client_id, key, ReplicaType::DISK).has_value());
-    get_result = service_->GetReplicaList(key);
+    EXPECT_TRUE(service_->PutEnd(client_id, key, "default", ReplicaType::DISK)
+                    .has_value());
+    get_result = service_->GetReplicaList(key, "default");
     ASSERT_TRUE(get_result.has_value());
     EXPECT_EQ(1, get_result.value().replicas.size());
     ASSERT_TRUE(get_result.value().replicas[0].is_disk_replica());
@@ -192,17 +255,20 @@ TEST_F(MasterServiceSSDSnapshotTest, PutRevokeBothReplica) {
     config.replica_num = 1;
 
     ASSERT_TRUE(
-        service_->PutStart(client_id, key, slice_length, config).has_value());
+        service_->PutStart(client_id, key, "default", slice_length, config)
+            .has_value());
     EXPECT_TRUE(
-        service_->PutRevoke(client_id, key, ReplicaType::DISK).has_value());
+        service_->PutRevoke(client_id, key, "default", ReplicaType::DISK)
+            .has_value());
 
-    auto get_result = service_->GetReplicaList(key);
+    auto get_result = service_->GetReplicaList(key, "default");
     ASSERT_FALSE(get_result.has_value());
     EXPECT_EQ(ErrorCode::REPLICA_IS_NOT_READY, get_result.error());
 
     EXPECT_TRUE(
-        service_->PutRevoke(client_id, key, ReplicaType::MEMORY).has_value());
-    get_result = service_->GetReplicaList(key);
+        service_->PutRevoke(client_id, key, "default", ReplicaType::MEMORY)
+            .has_value());
+    get_result = service_->GetReplicaList(key, "default");
     ASSERT_FALSE(get_result.has_value());
     EXPECT_EQ(ErrorCode::OBJECT_NOT_FOUND, get_result.error());
 }
@@ -230,15 +296,16 @@ TEST_F(MasterServiceSSDSnapshotTest, RemoveKey) {
     config.replica_num = 1;
 
     ASSERT_TRUE(
-        service_->PutStart(client_id, key, slice_length, config).has_value());
-    EXPECT_TRUE(
-        service_->PutEnd(client_id, key, ReplicaType::MEMORY).has_value());
-    EXPECT_TRUE(
-        service_->PutEnd(client_id, key, ReplicaType::DISK).has_value());
+        service_->PutStart(client_id, key, "default", slice_length, config)
+            .has_value());
+    EXPECT_TRUE(service_->PutEnd(client_id, key, "default", ReplicaType::MEMORY)
+                    .has_value());
+    EXPECT_TRUE(service_->PutEnd(client_id, key, "default", ReplicaType::DISK)
+                    .has_value());
 
-    EXPECT_TRUE(service_->Remove(key).has_value());
+    EXPECT_TRUE(service_->Remove(key, "default").has_value());
 
-    auto get_result = service_->GetReplicaList(key);
+    auto get_result = service_->GetReplicaList(key, "default");
     EXPECT_FALSE(get_result.has_value());
     EXPECT_EQ(ErrorCode::OBJECT_NOT_FOUND, get_result.error());
 }
@@ -271,12 +338,12 @@ TEST_F(MasterServiceSSDSnapshotTest, EvictObject) {
         ReplicateConfig config;
         config.replica_num = 1;
         auto put_start_result =
-            service_->PutStart(client_id, key, slice_length, config);
+            service_->PutStart(client_id, key, "default", slice_length, config);
         if (put_start_result.has_value()) {
-            auto put_end_mem_result =
-                service_->PutEnd(client_id, key, ReplicaType::MEMORY);
+            auto put_end_mem_result = service_->PutEnd(
+                client_id, key, "default", ReplicaType::MEMORY);
             auto put_end_disk_result =
-                service_->PutEnd(client_id, key, ReplicaType::DISK);
+                service_->PutEnd(client_id, key, "default", ReplicaType::DISK);
             ASSERT_TRUE(put_end_mem_result.has_value());
             ASSERT_TRUE(put_end_disk_result.has_value());
             success_puts++;
@@ -291,7 +358,7 @@ TEST_F(MasterServiceSSDSnapshotTest, EvictObject) {
     int success_gets = 0;
     for (int i = 0; i < 1024 * 16 + 50; ++i) {
         std::string key = "test_key" + std::to_string(i);
-        auto get_result = service_->GetReplicaList(key);
+        auto get_result = service_->GetReplicaList(key, "default");
         if (get_result.has_value()) {
             success_gets++;
         }
@@ -342,7 +409,7 @@ TEST_F(MasterServiceSSDSnapshotTest, PutStartExpires) {
 
         // Put key, should success.
         auto put_start_result =
-            service_->PutStart(client_id, key, slice_length, config);
+            service_->PutStart(client_id, key, "default", slice_length, config);
         EXPECT_TRUE(put_start_result.has_value());
         auto replica_list = put_start_result.value();
         EXPECT_EQ(replica_list.size(), kReplicaCnt);
@@ -351,7 +418,8 @@ TEST_F(MasterServiceSSDSnapshotTest, PutStartExpires) {
         }
 
         // Complete the reserved replica.
-        auto put_end_result = service_->PutEnd(client_id, key, reserve_type);
+        auto put_end_result =
+            service_->PutEnd(client_id, key, "default", reserve_type);
         EXPECT_TRUE(put_end_result.has_value());
 
         // Wait for a while until the put-start expired.
@@ -361,7 +429,7 @@ TEST_F(MasterServiceSSDSnapshotTest, PutStartExpires) {
             auto result = service_->Ping(client_id);
             EXPECT_TRUE(result.has_value());
             // Protect the key from eviction.
-            auto get_result = service_->GetReplicaList(key);
+            auto get_result = service_->GetReplicaList(key, "default");
             EXPECT_TRUE(get_result.has_value());
             std::this_thread::sleep_for(std::chrono::seconds(1));
         }
@@ -369,7 +437,7 @@ TEST_F(MasterServiceSSDSnapshotTest, PutStartExpires) {
         // Put key again, should fail because the object has had an completed
         // replica.
         put_start_result =
-            service_->PutStart(client_id, key, slice_length, config);
+            service_->PutStart(client_id, key, "default", slice_length, config);
         EXPECT_FALSE(put_start_result.has_value());
         EXPECT_EQ(put_start_result.error(), ErrorCode::OBJECT_ALREADY_EXISTS);
 
@@ -380,17 +448,18 @@ TEST_F(MasterServiceSSDSnapshotTest, PutStartExpires) {
             auto result = service_->Ping(client_id);
             EXPECT_TRUE(result.has_value());
             // Protect the key from eviction.
-            auto get_result = service_->GetReplicaList(key);
+            auto get_result = service_->GetReplicaList(key, "default");
             EXPECT_TRUE(get_result.has_value());
             std::this_thread::sleep_for(std::chrono::seconds(1));
         }
 
         // Try PutEnd the discarded replica.
-        put_end_result = service_->PutEnd(client_id, key, discard_type);
+        put_end_result =
+            service_->PutEnd(client_id, key, "default", discard_type);
         EXPECT_TRUE(put_end_result.has_value());
 
         // Check that the key has only one replica.
-        auto get_result = service_->GetReplicaList(key);
+        auto get_result = service_->GetReplicaList(key, "default");
         EXPECT_TRUE(get_result.has_value());
         EXPECT_EQ(get_result.value().replicas.size(), 1);
         if (reserve_type == ReplicaType::MEMORY) {
diff --git a/mooncake-store/tests/ha/snapshot/master_service_test_for_snapshot.cpp b/mooncake-store/tests/ha/snapshot/master_service_test_for_snapshot.cpp
index ec9f290b..9ca98ab6 100644
--- a/mooncake-store/tests/ha/snapshot/master_service_test_for_snapshot.cpp
+++ b/mooncake-store/tests/ha/snapshot/master_service_test_for_snapshot.cpp
@@ -1,5 +1,6 @@
 #include "master_service_test_for_snapshot_base.h"
 
+#include <algorithm>
 #include <atomic>
 #include <random>
 #include <unordered_set>
@@ -35,14 +36,14 @@ std::string GenerateKeyForSegment(const UUID& client_id,
         std::vector<Replica::Descriptor> replica_list;
 
         // Check if the key already exists.
-        auto exist_result = service->ExistKey(key);
+        auto exist_result = service->ExistKey(key, "default");
         if (exist_result.has_value() && exist_result.value()) {
             continue;  // Retry if the key already exists
         }
 
         // Attempt to put the key.
-        auto put_result =
-            service->PutStart(client_id, key, {1024}, {.replica_num = 1});
+        auto put_result = service->PutStart(client_id, key, "default", {1024},
+                                            {.replica_num = 1});
         if (put_result.has_value()) {
             replica_list = std::move(put_result.value());
         }
@@ -57,7 +58,7 @@ std::string GenerateKeyForSegment(const UUID& client_id,
                                      std::to_string(static_cast<int>(code)));
         }
         auto put_end_result =
-            service->PutEnd(client_id, key, ReplicaType::MEMORY);
+            service->PutEnd(client_id, key, "default", ReplicaType::MEMORY);
         if (!put_end_result.has_value()) {
             throw std::runtime_error("PutEnd failed");
         }
@@ -67,7 +68,7 @@ std::string GenerateKeyForSegment(const UUID& client_id,
             return key;
         }
         // Clean up failed attempt
-        auto remove_result = service->Remove(key);
+        auto remove_result = service->Remove(key, "default");
         if (!remove_result.has_value()) {
             // Ignore cleanup failure
         }
@@ -205,13 +206,14 @@ TEST_F(MasterServiceSnapshotTest, PutStartInvalidParams) {
 
     // Test invalid replica_num
     config.replica_num = 0;
-    auto put_result1 = service_->PutStart(client_id, key, 1024, config);
+    auto put_result1 =
+        service_->PutStart(client_id, key, "default", 1024, config);
     EXPECT_FALSE(put_result1.has_value());
     EXPECT_EQ(ErrorCode::INVALID_PARAMS, put_result1.error());
 
     // Test zero slice_length
     config.replica_num = 1;
-    auto put_result2 = service_->PutStart(client_id, key, 0, config);
+    auto put_result2 = service_->PutStart(client_id, key, "default", 0, config);
     EXPECT_FALSE(put_result2.has_value());
     EXPECT_EQ(ErrorCode::INVALID_PARAMS, put_result2.error());
 }
@@ -230,38 +232,39 @@ TEST_F(MasterServiceSnapshotTest, PutStartEndFlow) {
     config.replica_num = 1;
 
     auto put_start_result =
-        service_->PutStart(client_id, key, value_length, config);
+        service_->PutStart(client_id, key, "default", value_length, config);
     EXPECT_TRUE(put_start_result.has_value());
     replica_list = put_start_result.value();
     EXPECT_FALSE(replica_list.empty());
     EXPECT_EQ(ReplicaStatus::PROCESSING, replica_list[0].status);
 
     // During put, Get/Remove should fail
-    auto get_replica_result = service_->GetReplicaList(key);
+    auto get_replica_result = service_->GetReplicaList(key, "default");
     EXPECT_FALSE(get_replica_result.has_value());
     EXPECT_EQ(ErrorCode::REPLICA_IS_NOT_READY, get_replica_result.error());
-    auto remove_result = service_->Remove(key);
+    auto remove_result = service_->Remove(key, "default");
     EXPECT_FALSE(remove_result.has_value());
     EXPECT_EQ(ErrorCode::REPLICA_IS_NOT_READY, remove_result.error());
 
     // PutEnd should fail if the client_id does not match.
-    auto put_end_fail_result =
-        service_->PutEnd(invalid_client_id, key, ReplicaType::MEMORY);
+    auto put_end_fail_result = service_->PutEnd(invalid_client_id, key,
+                                                "default", ReplicaType::MEMORY);
     EXPECT_FALSE(put_end_fail_result.has_value());
     EXPECT_EQ(put_end_fail_result.error(), ErrorCode::ILLEGAL_CLIENT);
 
     // PutRevoke should fail if the client_id does not match.
-    auto put_revoke_fail_result =
-        service_->PutRevoke(invalid_client_id, key, ReplicaType::MEMORY);
+    auto put_revoke_fail_result = service_->PutRevoke(
+        invalid_client_id, key, "default", ReplicaType::MEMORY);
     EXPECT_FALSE(put_revoke_fail_result.has_value());
     EXPECT_EQ(put_revoke_fail_result.error(), ErrorCode::ILLEGAL_CLIENT);
 
     // Test PutEnd
-    auto put_end_result = service_->PutEnd(client_id, key, ReplicaType::MEMORY);
+    auto put_end_result =
+        service_->PutEnd(client_id, key, "default", ReplicaType::MEMORY);
     EXPECT_TRUE(put_end_result.has_value());
 
     // Verify replica list after PutEnd
-    auto final_get_result = service_->GetReplicaList(key);
+    auto final_get_result = service_->GetReplicaList(key, "default");
     EXPECT_TRUE(final_get_result.has_value());
     replica_list = final_get_result.value().replicas;
     EXPECT_EQ(1, replica_list.size());
@@ -291,23 +294,24 @@ TEST_F(MasterServiceSnapshotTest, RandomPutStartEndFlow) {
     int random_number = dis(gen);
     config.replica_num = random_number;
     auto put_start_result =
-        service_->PutStart(client_id, key, value_length, config);
+        service_->PutStart(client_id, key, "default", value_length, config);
     EXPECT_TRUE(put_start_result.has_value());
     replica_list = put_start_result.value();
     EXPECT_FALSE(replica_list.empty());
     EXPECT_EQ(ReplicaStatus::PROCESSING, replica_list[0].status);
     // During put, Get/Remove should fail
-    auto get_result = service_->GetReplicaList(key);
+    auto get_result = service_->GetReplicaList(key, "default");
     EXPECT_FALSE(get_result.has_value());
     EXPECT_EQ(ErrorCode::REPLICA_IS_NOT_READY, get_result.error());
-    auto remove_result = service_->Remove(key);
+    auto remove_result = service_->Remove(key, "default");
     EXPECT_FALSE(remove_result.has_value());
     EXPECT_EQ(ErrorCode::REPLICA_IS_NOT_READY, remove_result.error());
     // Test PutEnd
-    auto put_end_result = service_->PutEnd(client_id, key, ReplicaType::MEMORY);
+    auto put_end_result =
+        service_->PutEnd(client_id, key, "default", ReplicaType::MEMORY);
     EXPECT_TRUE(put_end_result.has_value());
     // Verify replica list after PutEnd
-    auto get_result2 = service_->GetReplicaList(key);
+    auto get_result2 = service_->GetReplicaList(key, "default");
     EXPECT_TRUE(get_result2.has_value());
     replica_list = get_result2.value().replicas;
     EXPECT_EQ(random_number, replica_list.size());
@@ -324,7 +328,7 @@ TEST_F(MasterServiceSnapshotTest, GetReplicaListByRegex) {
     service_.reset(new MasterService(service_config));
     const UUID client_id = generate_uuid();
     // Test getting non-existent key
-    auto get_result = service_->GetReplicaList(".*non_existent.*");
+    auto get_result = service_->GetReplicaList(".*non_existent.*", "default");
     EXPECT_FALSE(get_result.has_value());
     EXPECT_EQ(ErrorCode::OBJECT_NOT_FOUND, get_result.error());
 
@@ -337,19 +341,19 @@ TEST_F(MasterServiceSnapshotTest, GetReplicaListByRegex) {
         ReplicateConfig config;
         config.replica_num = 1;
         auto put_start_result =
-            service_->PutStart(client_id, key, value_length, config);
+            service_->PutStart(client_id, key, "default", value_length, config);
         ASSERT_TRUE(put_start_result.has_value());
         auto put_end_result =
-            service_->PutEnd(client_id, key, ReplicaType::MEMORY);
+            service_->PutEnd(client_id, key, "default", ReplicaType::MEMORY);
         ASSERT_TRUE(put_end_result.has_value());
-        auto exist_result = service_->ExistKey(key);
+        auto exist_result = service_->ExistKey(key, "default");
         ASSERT_TRUE(exist_result.has_value());
     }
     // wait for all the lease to expire
     std::this_thread::sleep_for(std::chrono::milliseconds(kv_lease_ttl));
 
     // Test getting existing key
-    auto get_result2 = service_->GetReplicaListByRegex("^test_key");
+    auto get_result2 = service_->GetReplicaListByRegex("^test_key", "default");
     EXPECT_TRUE(get_result2.has_value());
     auto replica_list_local = get_result2.value();
     EXPECT_EQ(10, replica_list_local.size());
@@ -362,13 +366,14 @@ void put_object(MasterService& service, const UUID& client_id,
     ReplicateConfig config;
     config.replica_num = 1;
     auto put_start_result =
-        service.PutStart(client_id, key, value_length, config);
+        service.PutStart(client_id, key, "default", value_length, config);
     ASSERT_TRUE(put_start_result.has_value())
         << "Failed to PutStart for key: " << key;
-    auto put_end_result = service.PutEnd(client_id, key, ReplicaType::MEMORY);
+    auto put_end_result =
+        service.PutEnd(client_id, key, "default", ReplicaType::MEMORY);
     ASSERT_TRUE(put_end_result.has_value())
         << "Failed to PutEnd for key: " << key;
-    auto exist_result = service.ExistKey(key);
+    auto exist_result = service.ExistKey(key, "default");
     ASSERT_TRUE(exist_result.has_value())
         << "Key does not exist after put: " << key;
 }
@@ -412,7 +417,7 @@ TEST_F(MasterServiceSnapshotTest, GetReplicaListByRegexComplex) {
 
     // Test 3.1: Simple prefix matching
     {
-        auto result = service_->GetReplicaListByRegex("^test_key_");
+        auto result = service_->GetReplicaListByRegex("^test_key_", "default");
         ASSERT_TRUE(result.has_value());
         EXPECT_EQ(result.value().size(),
                   3);  // Matches test_key_01, test_key_02, test_key_10
@@ -420,7 +425,8 @@ TEST_F(MasterServiceSnapshotTest, GetReplicaListByRegexComplex) {
 
     // Test 3.2: Matching with a wildcard for any number
     {
-        auto result = service_->GetReplicaListByRegex("^test_key_\\d+$");
+        auto result =
+            service_->GetReplicaListByRegex("^test_key_\\d+$", "default");
         ASSERT_TRUE(result.has_value());
         EXPECT_EQ(result.value().size(), 3);
     }
@@ -428,8 +434,8 @@ TEST_F(MasterServiceSnapshotTest, GetReplicaListByRegexComplex) {
     // Test 3.3: Matching a specific pattern with wildcards
     {
         // Matches "data_part_1_chunk_a" and "data_part_2_chunk_b"
-        auto result =
-            service_->GetReplicaListByRegex("^data_part_\\d_chunk_.$");
+        auto result = service_->GetReplicaListByRegex("^data_part_\\d_chunk_.$",
+                                                      "default");
         ASSERT_TRUE(result.has_value());
         EXPECT_EQ(result.value().size(), 2);
     }
@@ -437,7 +443,7 @@ TEST_F(MasterServiceSnapshotTest, GetReplicaListByRegexComplex) {
     // Test 3.4: Matching keys containing a specific substring
     {
         // Matches all keys with "key" in them
-        auto result = service_->GetReplicaListByRegex("key");
+        auto result = service_->GetReplicaListByRegex("key", "default");
         ASSERT_TRUE(result.has_value());
         // Expected: test_key_01, test_key_02, test_key_10,
         //           prod_key_alpha, prod_key_beta,
@@ -449,7 +455,7 @@ TEST_F(MasterServiceSnapshotTest, GetReplicaListByRegexComplex) {
     // Test 3.5: Matching based on file-like paths
     {
         // Match all .log files
-        auto result = service_->GetReplicaListByRegex("\\.log$");
+        auto result = service_->GetReplicaListByRegex("\\.log$", "default");
         ASSERT_TRUE(result.has_value());
         EXPECT_EQ(result.value().size(), 1);
         EXPECT_EQ(result.value().begin()->first, "logs/app-2025-08-13.log");
@@ -458,7 +464,8 @@ TEST_F(MasterServiceSnapshotTest, GetReplicaListByRegexComplex) {
     // Test 3.6: OR condition using |
     {
         // Match keys starting with "prod" OR ending with "json"
-        auto result = service_->GetReplicaListByRegex("^prod|\\.json$");
+        auto result =
+            service_->GetReplicaListByRegex("^prod|\\.json$", "default");
         ASSERT_TRUE(result.has_value());
         // Expected: prod_key_alpha, prod_key_beta, config/user/settings.json
         EXPECT_EQ(result.value().size(), 3);
@@ -466,7 +473,8 @@ TEST_F(MasterServiceSnapshotTest, GetReplicaListByRegexComplex) {
 
     // Test 3.7: Regex that should not match anything
     {
-        auto result = service_->GetReplicaListByRegex("^non_existent_prefix_");
+        auto result =
+            service_->GetReplicaListByRegex("^non_existent_prefix_", "default");
         // This should succeed but return an empty map.
         ASSERT_TRUE(result.has_value());
         EXPECT_TRUE(result.value().empty());
@@ -474,7 +482,7 @@ TEST_F(MasterServiceSnapshotTest, GetReplicaListByRegexComplex) {
 
     // Test 3.8: Exact match regex
     {
-        auto result = service_->GetReplicaListByRegex("^short$");
+        auto result = service_->GetReplicaListByRegex("^short$", "default");
         ASSERT_TRUE(result.has_value());
         EXPECT_EQ(result.value().size(), 1);
         EXPECT_EQ(result.value().begin()->first, "short");
@@ -482,8 +490,8 @@ TEST_F(MasterServiceSnapshotTest, GetReplicaListByRegexComplex) {
 
     // Test 3.9: Initial test for non-existent key (as a sanity check)
     {
-        auto get_result =
-            service_->GetReplicaListByRegex(".*absolutely_non_existent.*");
+        auto get_result = service_->GetReplicaListByRegex(
+            ".*absolutely_non_existent.*", "default");
         // Depending on implementation, this could return an empty map or an
         // error. Let's assume it returns an empty map for a valid regex with no
         // matches.
@@ -496,7 +504,7 @@ TEST_F(MasterServiceSnapshotTest, GetReplicaList) {
     service_.reset(new MasterService());
     const UUID client_id = generate_uuid();
     // Test getting non-existent key
-    auto get_result = service_->GetReplicaList("non_existent");
+    auto get_result = service_->GetReplicaList("non_existent", "default");
     EXPECT_FALSE(get_result.has_value());
     EXPECT_EQ(ErrorCode::OBJECT_NOT_FOUND, get_result.error());
 
@@ -507,13 +515,14 @@ TEST_F(MasterServiceSnapshotTest, GetReplicaList) {
     ReplicateConfig config;
     config.replica_num = 1;
     auto put_start_result =
-        service_->PutStart(client_id, key, value_length, config);
+        service_->PutStart(client_id, key, "default", value_length, config);
     ASSERT_TRUE(put_start_result.has_value());
-    auto put_end_result = service_->PutEnd(client_id, key, ReplicaType::MEMORY);
+    auto put_end_result =
+        service_->PutEnd(client_id, key, "default", ReplicaType::MEMORY);
     ASSERT_TRUE(put_end_result.has_value());
 
     // Test getting existing key
-    auto get_result2 = service_->GetReplicaList(key);
+    auto get_result2 = service_->GetReplicaList(key, "default");
     EXPECT_TRUE(get_result2.has_value());
     auto replica_list_local = get_result2.value().replicas;
     EXPECT_FALSE(replica_list_local.empty());
@@ -529,22 +538,23 @@ TEST_F(MasterServiceSnapshotTest, RemoveObject) {
     ReplicateConfig config;
     config.replica_num = 1;
     auto put_start_result =
-        service_->PutStart(client_id, key, value_length, config);
+        service_->PutStart(client_id, key, "default", value_length, config);
     ASSERT_TRUE(put_start_result.has_value());
-    auto put_end_result = service_->PutEnd(client_id, key, ReplicaType::MEMORY);
+    auto put_end_result =
+        service_->PutEnd(client_id, key, "default", ReplicaType::MEMORY);
     ASSERT_TRUE(put_end_result.has_value());
 
     // Test removing the object
-    auto remove_result = service_->Remove(key);
+    auto remove_result = service_->Remove(key, "default");
     EXPECT_TRUE(remove_result.has_value());
 
     // Verify object is removed
-    auto get_result = service_->GetReplicaList(key);
+    auto get_result = service_->GetReplicaList(key, "default");
     EXPECT_FALSE(get_result.has_value());
     EXPECT_EQ(ErrorCode::OBJECT_NOT_FOUND, get_result.error());
 
     // Test removing non-existent object
-    auto remove_result2 = service_->Remove("non_existent");
+    auto remove_result2 = service_->Remove("non_existent", "default");
     EXPECT_FALSE(remove_result2.has_value());
     EXPECT_EQ(ErrorCode::OBJECT_NOT_FOUND, remove_result2.error());
 }
@@ -563,18 +573,18 @@ TEST_F(MasterServiceSnapshotTest, RandomRemoveObject) {
         ReplicateConfig config;
         config.replica_num = 1;
         auto put_start_result =
-            service_->PutStart(client_id, key, value_length, config);
+            service_->PutStart(client_id, key, "default", value_length, config);
         ASSERT_TRUE(put_start_result.has_value());
         auto put_end_result =
-            service_->PutEnd(client_id, key, ReplicaType::MEMORY);
+            service_->PutEnd(client_id, key, "default", ReplicaType::MEMORY);
         ASSERT_TRUE(put_end_result.has_value());
 
         // Test removing the object
-        auto remove_result = service_->Remove(key);
+        auto remove_result = service_->Remove(key, "default");
         EXPECT_TRUE(remove_result.has_value());
 
         // Verify object is removed
-        auto get_result = service_->GetReplicaList(key);
+        auto get_result = service_->GetReplicaList(key, "default");
         EXPECT_FALSE(get_result.has_value());
         EXPECT_EQ(ErrorCode::OBJECT_NOT_FOUND, get_result.error());
     }
@@ -595,23 +605,23 @@ TEST_F(MasterServiceSnapshotTest, RemoveByRegex) {
         ReplicateConfig config;
         config.replica_num = 1;
         auto put_start_result =
-            service_->PutStart(client_id, key, value_length, config);
+            service_->PutStart(client_id, key, "default", value_length, config);
         ASSERT_TRUE(put_start_result.has_value());
         auto put_end_result =
-            service_->PutEnd(client_id, key, ReplicaType::MEMORY);
+            service_->PutEnd(client_id, key, "default", ReplicaType::MEMORY);
         ASSERT_TRUE(put_end_result.has_value());
-        auto exist_result = service_->ExistKey(key);
+        auto exist_result = service_->ExistKey(key, "default");
         ASSERT_TRUE(exist_result.has_value());
     }
     // wait for all the lease to expire
     std::this_thread::sleep_for(std::chrono::milliseconds(kv_lease_ttl));
-    auto res = service_->RemoveByRegex("^test_key");
+    auto res = service_->RemoveByRegex("^test_key", "default");
     ASSERT_TRUE(res.has_value());
     ASSERT_EQ(10, res.value());
     times = 10;
     while (times--) {
         std::string key = "test_key" + std::to_string(times);
-        auto exist_result = service_->ExistKey(key);
+        auto exist_result = service_->ExistKey(key, "default");
         ASSERT_TRUE(exist_result.has_value());
         ASSERT_FALSE(exist_result.value());
     }
@@ -658,7 +668,7 @@ TEST_F(MasterServiceSnapshotTest, RemoveByRegexComplex) {
         populate_store();
 
         // Action: Remove keys starting with "test_key_"
-        auto remove_result = service_->RemoveByRegex("^test_key_");
+        auto remove_result = service_->RemoveByRegex("^test_key_", "default");
         ASSERT_TRUE(remove_result.has_value());
         EXPECT_EQ(remove_result.value(), 3);  // Should remove 3 keys
 
@@ -666,7 +676,7 @@ TEST_F(MasterServiceSnapshotTest, RemoveByRegexComplex) {
         std::vector<std::string> deleted_keys = {"test_key_01", "test_key_02",
                                                  "test_key_10"};
         for (const auto& key : deleted_keys) {
-            auto exist_result = service_->ExistKey(key);
+            auto exist_result = service_->ExistKey(key, "default");
             ASSERT_TRUE(exist_result.has_value());
             EXPECT_FALSE(exist_result.value())
                 << "Key " << key << " should have been deleted.";
@@ -675,7 +685,7 @@ TEST_F(MasterServiceSnapshotTest, RemoveByRegexComplex) {
         std::vector<std::string> remaining_keys = {
             "prod_key_alpha", "short", "test-key-extra"};  // Sample a few
         for (const auto& key : remaining_keys) {
-            auto exist_result = service_->ExistKey(key);
+            auto exist_result = service_->ExistKey(key, "default");
             ASSERT_TRUE(exist_result.has_value());
             EXPECT_TRUE(exist_result.value())
                 << "Key " << key << " should NOT have been deleted.";
@@ -698,12 +708,12 @@ TEST_F(MasterServiceSnapshotTest, RemoveByRegexComplex) {
         size_t total_keys = 13;  // Count from the keys_to_put vector
 
         // Action: Remove all keys
-        auto remove_result = service_->RemoveByRegex(".*");
+        auto remove_result = service_->RemoveByRegex(".*", "default");
         ASSERT_TRUE(remove_result.has_value());
         EXPECT_EQ(remove_result.value(), total_keys);
 
         // Verification: Check that no keys remain
-        auto get_all_result = service_->GetReplicaListByRegex(".*");
+        auto get_all_result = service_->GetReplicaListByRegex(".*", "default");
         ASSERT_TRUE(get_all_result.has_value());
         EXPECT_TRUE(get_all_result.value().empty());
     }
@@ -723,12 +733,13 @@ TEST_F(MasterServiceSnapshotTest, RemoveByRegexComplex) {
         size_t total_keys_before_remove = 13;
 
         // Action: Attempt to remove using a pattern that matches nothing
-        auto remove_result = service_->RemoveByRegex("^nonexistent-pattern-");
+        auto remove_result =
+            service_->RemoveByRegex("^nonexistent-pattern-", "default");
         ASSERT_TRUE(remove_result.has_value());
         EXPECT_EQ(remove_result.value(), 0);  // Should remove 0 keys
 
         // Verification: Check that all keys still exist
-        auto get_all_result = service_->GetReplicaListByRegex(".*");
+        auto get_all_result = service_->GetReplicaListByRegex(".*", "default");
         ASSERT_TRUE(get_all_result.has_value());
         EXPECT_EQ(get_all_result.value().size(), total_keys_before_remove);
     }
@@ -746,7 +757,7 @@ TEST_F(MasterServiceSnapshotTest, RemoveByRegexComplex) {
         populate_store();
 
         // Action: Remove all keys that contain a slash '/' OR end with a number
-        auto remove_result = service_->RemoveByRegex("/|\\d$");
+        auto remove_result = service_->RemoveByRegex("/|\\d$", "default");
         ASSERT_TRUE(remove_result.has_value());
         // Matches: "config/user/settings.json", "logs/app-2025-08-13.log",
         //          "test_key_01", "test_key_02", "test_key_10"
@@ -771,23 +782,25 @@ TEST_F(MasterServiceSnapshotTest, RemoveByRegexComplex) {
         populate_store();
 
         // Action: Remove all keys that contain "chunk" OR "config"
-        auto remove_result = service_->RemoveByRegex("chunk|config");
+        auto remove_result = service_->RemoveByRegex("chunk|config", "default");
         ASSERT_TRUE(remove_result.has_value());
         // Matches: "data_part_1_chunk_a", "data_part_2_chunk_b",
         // "config/user/settings.json"
         EXPECT_EQ(remove_result.value(), 3);
 
         // Verification
-        auto exist_result_chunk = service_->ExistKey("data_part_1_chunk_a");
+        auto exist_result_chunk =
+            service_->ExistKey("data_part_1_chunk_a", "default");
         ASSERT_TRUE(exist_result_chunk.has_value());
         EXPECT_FALSE(exist_result_chunk.value());
 
         auto exist_result_config =
-            service_->ExistKey("config/user/settings.json");
+            service_->ExistKey("config/user/settings.json", "default");
         ASSERT_TRUE(exist_result_config.has_value());
         EXPECT_FALSE(exist_result_config.value());
 
-        auto exist_result_untouched = service_->ExistKey("prod_key_alpha");
+        auto exist_result_untouched =
+            service_->ExistKey("prod_key_alpha", "default");
         ASSERT_TRUE(exist_result_untouched.has_value());
         EXPECT_TRUE(exist_result_untouched.value());
     }
@@ -808,12 +821,12 @@ TEST_F(MasterServiceSnapshotTest, RemoveAll) {
         ReplicateConfig config;
         config.replica_num = 1;
         auto put_start_result =
-            service_->PutStart(client_id, key, value_length, config);
+            service_->PutStart(client_id, key, "default", value_length, config);
         ASSERT_TRUE(put_start_result.has_value());
         auto put_end_result =
-            service_->PutEnd(client_id, key, ReplicaType::MEMORY);
+            service_->PutEnd(client_id, key, "default", ReplicaType::MEMORY);
         ASSERT_TRUE(put_end_result.has_value());
-        auto exist_result = service_->ExistKey(key);
+        auto exist_result = service_->ExistKey(key, "default");
         ASSERT_TRUE(exist_result.has_value());
     }
     // wait for all the lease to expire
@@ -822,7 +835,7 @@ TEST_F(MasterServiceSnapshotTest, RemoveAll) {
     // before TearDown snapshot verification ASSERT_EQ(10,
     // service_->RemoveAll()); times = 10; while (times--) {
     //     std::string key = "test_key" + std::to_string(times);
-    //     auto exist_result = service_->ExistKey(key);
+    //     auto exist_result = service_->ExistKey(key, "default");
     //     ASSERT_TRUE(exist_result.has_value());
     //     ASSERT_FALSE(exist_result.value());
     // }
@@ -857,7 +870,7 @@ TEST_F(MasterServiceSnapshotTest, SingleSliceMultiReplicaFlow) {
 
     // Test PutStart with multiple slices and replicas
     auto put_start_result =
-        service_->PutStart(client_id, key, slice_length, config);
+        service_->PutStart(client_id, key, "default", slice_length, config);
     ASSERT_TRUE(put_start_result.has_value());
     replica_list = put_start_result.value();
 
@@ -873,16 +886,17 @@ TEST_F(MasterServiceSnapshotTest, SingleSliceMultiReplicaFlow) {
     }
 
     // Test GetReplicaList during processing (should fail)
-    auto get_result = service_->GetReplicaList(key);
+    auto get_result = service_->GetReplicaList(key, "default");
     EXPECT_FALSE(get_result.has_value());
     EXPECT_EQ(ErrorCode::REPLICA_IS_NOT_READY, get_result.error());
 
     // Complete the put operation
-    auto put_end_result = service_->PutEnd(client_id, key, ReplicaType::MEMORY);
+    auto put_end_result =
+        service_->PutEnd(client_id, key, "default", ReplicaType::MEMORY);
     ASSERT_TRUE(put_end_result.has_value());
 
     // Test GetReplicaList after completion
-    auto get_result2 = service_->GetReplicaList(key);
+    auto get_result2 = service_->GetReplicaList(key, "default");
     ASSERT_TRUE(get_result2.has_value());
     auto retrieved_replicas = get_result2.value().replicas;
     ASSERT_EQ(num_replicas, retrieved_replicas.size());
@@ -916,13 +930,14 @@ TEST_F(MasterServiceSnapshotTest, CleanupStaleHandlesTest) {
 
     // Create the object
     auto put_start_result =
-        service_->PutStart(client_id, key, slice_length, config);
+        service_->PutStart(client_id, key, "default", slice_length, config);
     ASSERT_TRUE(put_start_result.has_value());
-    auto put_end_result = service_->PutEnd(client_id, key, ReplicaType::MEMORY);
+    auto put_end_result =
+        service_->PutEnd(client_id, key, "default", ReplicaType::MEMORY);
     ASSERT_TRUE(put_end_result.has_value());
 
     // Verify object exists
-    auto get_result = service_->GetReplicaList(key);
+    auto get_result = service_->GetReplicaList(key, "default");
     ASSERT_TRUE(get_result.has_value());
     auto retrieved_replicas = get_result.value().replicas;
     ASSERT_EQ(1, retrieved_replicas.size());
@@ -933,7 +948,7 @@ TEST_F(MasterServiceSnapshotTest, CleanupStaleHandlesTest) {
 
     // Try to get the object - it should be automatically removed since the
     // replica is invalid
-    auto get_result2 = service_->GetReplicaList(key);
+    auto get_result2 = service_->GetReplicaList(key, "default");
     EXPECT_FALSE(get_result2.has_value());
     EXPECT_EQ(ErrorCode::OBJECT_NOT_FOUND, get_result2.error());
 
@@ -944,14 +959,14 @@ TEST_F(MasterServiceSnapshotTest, CleanupStaleHandlesTest) {
     // Create another object
     std::string key2 = "another_segment_object";
     auto put_start_result2 =
-        service_->PutStart(client_id, key2, slice_length, config);
+        service_->PutStart(client_id, key2, "default", slice_length, config);
     ASSERT_TRUE(put_start_result2.has_value());
     auto put_end_result2 =
-        service_->PutEnd(client_id, key2, ReplicaType::MEMORY);
+        service_->PutEnd(client_id, key2, "default", ReplicaType::MEMORY);
     ASSERT_TRUE(put_end_result2.has_value());
 
     // Verify we can get it
-    auto get_result3 = service_->GetReplicaList(key2);
+    auto get_result3 = service_->GetReplicaList(key2, "default");
     ASSERT_TRUE(get_result3.has_value());
 
     // Unmount the segment
@@ -959,7 +974,7 @@ TEST_F(MasterServiceSnapshotTest, CleanupStaleHandlesTest) {
     ASSERT_TRUE(unmount_result2.has_value());
 
     // Try to remove the object that should already be cleaned up
-    auto remove_result = service_->Remove(key2);
+    auto remove_result = service_->Remove(key2, "default");
     EXPECT_FALSE(remove_result.has_value());
     EXPECT_EQ(ErrorCode::OBJECT_NOT_FOUND, remove_result.error());
 }
@@ -991,11 +1006,11 @@ TEST_F(MasterServiceSnapshotTest, ConcurrentWriteAndRemoveAll) {
                 config.replica_num = 1;
                 std::vector<Replica::Descriptor> replica_list;
 
-                auto put_start_result =
-                    service_->PutStart(client_id, key, slice_length, config);
+                auto put_start_result = service_->PutStart(
+                    client_id, key, "default", slice_length, config);
                 if (put_start_result.has_value()) {
-                    auto put_end_result =
-                        service_->PutEnd(client_id, key, ReplicaType::MEMORY);
+                    auto put_end_result = service_->PutEnd(
+                        client_id, key, "default", ReplicaType::MEMORY);
                     if (put_end_result.has_value()) {
                         success_writes++;
                     }
@@ -1062,10 +1077,10 @@ TEST_F(MasterServiceSnapshotTest, ConcurrentReadAndRemoveAll) {
         config.replica_num = 1;
 
         auto put_start_result =
-            service_->PutStart(client_id, key, slice_length, config);
+            service_->PutStart(client_id, key, "default", slice_length, config);
         ASSERT_TRUE(put_start_result.has_value());
         auto put_end_result =
-            service_->PutEnd(client_id, key, ReplicaType::MEMORY);
+            service_->PutEnd(client_id, key, "default", ReplicaType::MEMORY);
         ASSERT_TRUE(put_end_result.has_value());
     }
 
@@ -1078,7 +1093,7 @@ TEST_F(MasterServiceSnapshotTest, ConcurrentReadAndRemoveAll) {
         readers.emplace_back([&]() {
             for (int j = 0; j < num_objects; ++j) {
                 std::string key = "pre_key_" + std::to_string(j);
-                auto get_result = service_->GetReplicaList(key);
+                auto get_result = service_->GetReplicaList(key, "default");
                 if (get_result.has_value()) {
                     success_reads++;
                 }
@@ -1123,7 +1138,7 @@ TEST_F(MasterServiceSnapshotTest, ConcurrentReadAndRemoveAll) {
     // // Verify all objects were removed
     // for (int i = 0; i < num_objects; ++i) {
     //     std::string key = "pre_key_" + std::to_string(i);
-    //     auto get_result = service_->GetReplicaList(key);
+    //     auto get_result = service_->GetReplicaList(key, "default");
     //     EXPECT_FALSE(get_result.has_value());
     //     EXPECT_EQ(ErrorCode::OBJECT_NOT_FOUND, get_result.error());
     // }
@@ -1147,10 +1162,10 @@ TEST_F(MasterServiceSnapshotTest, ConcurrentRemoveAllOperations) {
         config.replica_num = 1;
 
         auto put_start_result =
-            service_->PutStart(client_id, key, slice_length, config);
+            service_->PutStart(client_id, key, "default", slice_length, config);
         ASSERT_TRUE(put_start_result.has_value());
         auto put_end_result =
-            service_->PutEnd(client_id, key, ReplicaType::MEMORY);
+            service_->PutEnd(client_id, key, "default", ReplicaType::MEMORY);
         ASSERT_TRUE(put_end_result.has_value());
     }
 
@@ -1178,7 +1193,7 @@ TEST_F(MasterServiceSnapshotTest, ConcurrentRemoveAllOperations) {
     // // Verify all objects were removed
     // for (int i = 0; i < num_objects; ++i) {
     //     std::string key = "pre_key_" + std::to_string(i);
-    //     auto get_result = service_->GetReplicaList(key);
+    //     auto get_result = service_->GetReplicaList(key, "default");
     //     EXPECT_FALSE(get_result.has_value());
     //     EXPECT_EQ(ErrorCode::OBJECT_NOT_FOUND, get_result.error());
     // }
@@ -1215,23 +1230,23 @@ TEST_F(MasterServiceSnapshotTest, UnmountSegmentImmediateCleanup) {
     // Umount will remove all objects in the segment, include the key1
     ASSERT_EQ(1, service_->GetKeyCount());
     // Verify objects in segment1 is gone
-    auto get_result1 = service_->GetReplicaList(key1);
+    auto get_result1 = service_->GetReplicaList(key1, "default");
     ASSERT_FALSE(get_result1.has_value());
     ASSERT_EQ(ErrorCode::OBJECT_NOT_FOUND, get_result1.error());
 
     // Verify objects in segment2 is still there
-    auto get_result2 = service_->GetReplicaList(key2);
+    auto get_result2 = service_->GetReplicaList(key2, "default");
     ASSERT_TRUE(get_result2.has_value());
 
     // Verify put key1 will put into segment2 rather than segment1
     auto put_start_result =
-        service_->PutStart(client_id, key1, slice_length, config);
+        service_->PutStart(client_id, key1, "default", slice_length, config);
     ASSERT_TRUE(put_start_result.has_value());
     replica_list = put_start_result.value();
     auto put_end_result =
-        service_->PutEnd(client_id, key1, ReplicaType::MEMORY);
+        service_->PutEnd(client_id, key1, "default", ReplicaType::MEMORY);
     ASSERT_TRUE(put_end_result.has_value());
-    auto get_result3 = service_->GetReplicaList(key1);
+    auto get_result3 = service_->GetReplicaList(key1, "default");
     ASSERT_TRUE(get_result3.has_value());
     auto retrieved = get_result3.value();
     ASSERT_EQ(replica_list[0]
@@ -1264,14 +1279,14 @@ TEST_F(MasterServiceSnapshotTest, ReadableAfterPartialUnmountWithReplication) {
     config.replica_num = 2;
 
     auto put_start_result =
-        service_->PutStart(client_id, key, slice_length, config);
+        service_->PutStart(client_id, key, "default", slice_length, config);
     ASSERT_TRUE(put_start_result.has_value());
     ASSERT_EQ(2u, put_start_result->size());
-    ASSERT_TRUE(
-        service_->PutEnd(client_id, key, ReplicaType::MEMORY).has_value());
+    ASSERT_TRUE(service_->PutEnd(client_id, key, "default", ReplicaType::MEMORY)
+                    .has_value());
 
     // Verify two replicas exist and they are on distinct segments
-    auto get_result = service_->GetReplicaList(key);
+    auto get_result = service_->GetReplicaList(key, "default");
     ASSERT_TRUE(get_result.has_value());
     auto replicas = get_result.value().replicas;
     ASSERT_EQ(2u, replicas.size());
@@ -1289,7 +1304,7 @@ TEST_F(MasterServiceSnapshotTest, ReadableAfterPartialUnmountWithReplication) {
     ASSERT_TRUE(service_->UnmountSegment(segment1.id, client_id).has_value());
 
     // Key should still be readable via the remaining replica
-    auto get_after_unmount = service_->GetReplicaList(key);
+    auto get_after_unmount = service_->GetReplicaList(key, "default");
     ASSERT_TRUE(get_after_unmount.has_value())
         << "Object should remain accessible with surviving replica";
 }
@@ -1339,7 +1354,7 @@ TEST_F(MasterServiceSnapshotTest, UnmountSegmentPerformance) {
 
     // Verify all keys are gone
     for (const auto& key : keys) {
-        auto get_result = service_->GetReplicaList(key);
+        auto get_result = service_->GetReplicaList(key, "default");
         EXPECT_FALSE(get_result.has_value());
         EXPECT_EQ(ErrorCode::OBJECT_NOT_FOUND, get_result.error());
     }
@@ -1370,75 +1385,76 @@ TEST_F(MasterServiceSnapshotTest, RemoveLeasedObject) {
 
     // Verify lease is granted on ExistsKey
     auto put_start_result =
-        service_->PutStart(client_id, key, slice_length, config);
+        service_->PutStart(client_id, key, "default", slice_length, config);
     ASSERT_TRUE(put_start_result.has_value());
-    auto put_end_result = service_->PutEnd(client_id, key, ReplicaType::MEMORY);
+    auto put_end_result =
+        service_->PutEnd(client_id, key, "default", ReplicaType::MEMORY);
     ASSERT_TRUE(put_end_result.has_value());
-    auto exist_result = service_->ExistKey(key);
+    auto exist_result = service_->ExistKey(key, "default");
     ASSERT_TRUE(exist_result.has_value());
-    auto remove_result = service_->Remove(key);
+    auto remove_result = service_->Remove(key, "default");
     EXPECT_FALSE(remove_result.has_value());
     EXPECT_EQ(ErrorCode::OBJECT_HAS_LEASE, remove_result.error());
     std::this_thread::sleep_for(std::chrono::milliseconds(kv_lease_ttl));
-    auto remove_result2 = service_->Remove(key);
+    auto remove_result2 = service_->Remove(key, "default");
     EXPECT_TRUE(remove_result2.has_value());
 
     // Verify lease is extended on successive ExistsKey
     auto put_start_result2 =
-        service_->PutStart(client_id, key, slice_length, config);
+        service_->PutStart(client_id, key, "default", slice_length, config);
     ASSERT_TRUE(put_start_result2.has_value());
     auto put_end_result2 =
-        service_->PutEnd(client_id, key, ReplicaType::MEMORY);
+        service_->PutEnd(client_id, key, "default", ReplicaType::MEMORY);
     ASSERT_TRUE(put_end_result2.has_value());
-    auto exist_result2 = service_->ExistKey(key);
+    auto exist_result2 = service_->ExistKey(key, "default");
     ASSERT_TRUE(exist_result2.has_value());
     std::this_thread::sleep_for(std::chrono::milliseconds(kv_lease_ttl));
-    auto exist_result3 = service_->ExistKey(key);
+    auto exist_result3 = service_->ExistKey(key, "default");
     ASSERT_TRUE(exist_result3.has_value());
-    auto remove_result3 = service_->Remove(key);
+    auto remove_result3 = service_->Remove(key, "default");
     EXPECT_FALSE(remove_result3.has_value());
     EXPECT_EQ(ErrorCode::OBJECT_HAS_LEASE, remove_result3.error());
     std::this_thread::sleep_for(std::chrono::milliseconds(kv_lease_ttl));
-    auto remove_result4 = service_->Remove(key);
+    auto remove_result4 = service_->Remove(key, "default");
     EXPECT_TRUE(remove_result4.has_value());
 
     // Verify lease is granted on GetReplicaList
     auto put_start_result3 =
-        service_->PutStart(client_id, key, slice_length, config);
+        service_->PutStart(client_id, key, "default", slice_length, config);
     ASSERT_TRUE(put_start_result3.has_value());
     auto put_end_result3 =
-        service_->PutEnd(client_id, key, ReplicaType::MEMORY);
+        service_->PutEnd(client_id, key, "default", ReplicaType::MEMORY);
     ASSERT_TRUE(put_end_result3.has_value());
-    auto get_result = service_->GetReplicaList(key);
+    auto get_result = service_->GetReplicaList(key, "default");
     ASSERT_TRUE(get_result.has_value());
-    auto remove_result5 = service_->Remove(key);
+    auto remove_result5 = service_->Remove(key, "default");
     EXPECT_FALSE(remove_result5.has_value());
     EXPECT_EQ(ErrorCode::OBJECT_HAS_LEASE, remove_result5.error());
     std::this_thread::sleep_for(std::chrono::milliseconds(kv_lease_ttl));
-    auto remove_result6 = service_->Remove(key);
+    auto remove_result6 = service_->Remove(key, "default");
     EXPECT_TRUE(remove_result6.has_value());
 
     // Verify lease is extended on successive GetReplicaList
     auto put_start_result4 =
-        service_->PutStart(client_id, key, slice_length, config);
+        service_->PutStart(client_id, key, "default", slice_length, config);
     ASSERT_TRUE(put_start_result4.has_value());
     auto put_end_result4 =
-        service_->PutEnd(client_id, key, ReplicaType::MEMORY);
+        service_->PutEnd(client_id, key, "default", ReplicaType::MEMORY);
     ASSERT_TRUE(put_end_result4.has_value());
-    auto get_result2 = service_->GetReplicaList(key);
+    auto get_result2 = service_->GetReplicaList(key, "default");
     ASSERT_TRUE(get_result2.has_value());
     std::this_thread::sleep_for(std::chrono::milliseconds(kv_lease_ttl));
-    auto get_result3 = service_->GetReplicaList(key);
+    auto get_result3 = service_->GetReplicaList(key, "default");
     ASSERT_TRUE(get_result3.has_value());
-    auto remove_result7 = service_->Remove(key);
+    auto remove_result7 = service_->Remove(key, "default");
     EXPECT_FALSE(remove_result7.has_value());
     EXPECT_EQ(ErrorCode::OBJECT_HAS_LEASE, remove_result7.error());
     std::this_thread::sleep_for(std::chrono::milliseconds(kv_lease_ttl));
-    auto remove_result8 = service_->Remove(key);
+    auto remove_result8 = service_->Remove(key, "default");
     EXPECT_TRUE(remove_result8.has_value());
 
     // Verify object is removed
-    auto get_result4 = service_->GetReplicaList(key);
+    auto get_result4 = service_->GetReplicaList(key, "default");
     EXPECT_FALSE(get_result4.has_value());
     EXPECT_EQ(ErrorCode::OBJECT_NOT_FOUND, get_result4.error());
 }
@@ -1457,13 +1473,13 @@ TEST_F(MasterServiceSnapshotTest, RemoveAllLeasedObject) {
         ReplicateConfig config;
         config.replica_num = 1;
         auto put_start_result =
-            service_->PutStart(client_id, key, slice_length, config);
+            service_->PutStart(client_id, key, "default", slice_length, config);
         ASSERT_TRUE(put_start_result.has_value());
         auto put_end_result =
-            service_->PutEnd(client_id, key, ReplicaType::MEMORY);
+            service_->PutEnd(client_id, key, "default", ReplicaType::MEMORY);
         ASSERT_TRUE(put_end_result.has_value());
         if (i >= 5) {
-            auto exist_result = service_->ExistKey(key);
+            auto exist_result = service_->ExistKey(key, "default");
             ASSERT_TRUE(exist_result.has_value());
         }
     }
@@ -1471,7 +1487,7 @@ TEST_F(MasterServiceSnapshotTest, RemoveAllLeasedObject) {
     // before TearDown snapshot verification ASSERT_EQ(5,
     // service_->RemoveAll()); for (int i = 0; i < 5; ++i) {
     //     std::string key = "test_key" + std::to_string(i);
-    //     auto exist_result = service_->ExistKey(key);
+    //     auto exist_result = service_->ExistKey(key, "default");
     //     ASSERT_FALSE(exist_result.value());
     // }
     // // wait for all the lease to expire
@@ -1479,7 +1495,7 @@ TEST_F(MasterServiceSnapshotTest, RemoveAllLeasedObject) {
     // ASSERT_EQ(5, service_->RemoveAll());
     // for (int i = 5; i < 10; ++i) {
     //     std::string key = "test_key" + std::to_string(i);
-    //     auto exist_result = service_->ExistKey(key);
+    //     auto exist_result = service_->ExistKey(key, "default");
     //     ASSERT_FALSE(exist_result.value());
     // }
 }
@@ -1510,10 +1526,10 @@ TEST_F(MasterServiceSnapshotTest, EvictObject) {
         ReplicateConfig config;
         config.replica_num = 1;
         auto put_start_result =
-            service_->PutStart(client_id, key, slice_length, config);
+            service_->PutStart(client_id, key, "default", slice_length, config);
         if (put_start_result.has_value()) {
-            auto put_end_result =
-                service_->PutEnd(client_id, key, ReplicaType::MEMORY);
+            auto put_end_result = service_->PutEnd(client_id, key, "default",
+                                                   ReplicaType::MEMORY);
             ASSERT_TRUE(put_end_result.has_value());
             success_puts++;
         } else {
@@ -1555,13 +1571,13 @@ TEST_F(MasterServiceSnapshotTest, TryEvictLeasedObject) {
         ReplicateConfig config;
         config.replica_num = 1;
         auto put_start_result =
-            service_->PutStart(client_id, key, slice_length, config);
+            service_->PutStart(client_id, key, "default", slice_length, config);
         if (put_start_result.has_value()) {
-            auto put_end_result =
-                service_->PutEnd(client_id, key, ReplicaType::MEMORY);
+            auto put_end_result = service_->PutEnd(client_id, key, "default",
+                                                   ReplicaType::MEMORY);
             ASSERT_TRUE(put_end_result.has_value());
             // the object is leased
-            auto get_result = service_->GetReplicaList(key);
+            auto get_result = service_->GetReplicaList(key, "default");
             ASSERT_TRUE(get_result.has_value());
             leased_keys.push_back(key);
             success_puts++;
@@ -1575,7 +1591,7 @@ TEST_F(MasterServiceSnapshotTest, TryEvictLeasedObject) {
     std::this_thread::sleep_for(std::chrono::milliseconds(50));
     // All leased objects should be accessible
     for (const auto& key : leased_keys) {
-        auto get_result = service_->GetReplicaList(key);
+        auto get_result = service_->GetReplicaList(key, "default");
         ASSERT_TRUE(get_result.has_value());
     }
     // [Commented for snapshot test] The following RemoveAll would clear data
@@ -1611,19 +1627,21 @@ TEST_F(MasterServiceSnapshotTest, RemoveSoftPinObject) {
 
     // Verify soft pin does not block remove
     ASSERT_TRUE(
-        service_->PutStart(client_id, key, slice_length, config).has_value());
-    ASSERT_TRUE(
-        service_->PutEnd(client_id, key, ReplicaType::MEMORY).has_value());
-    EXPECT_TRUE(service_->Remove(key).has_value());
+        service_->PutStart(client_id, key, "default", slice_length, config)
+            .has_value());
+    ASSERT_TRUE(service_->PutEnd(client_id, key, "default", ReplicaType::MEMORY)
+                    .has_value());
+    EXPECT_TRUE(service_->Remove(key, "default").has_value());
 
     // [Commented for snapshot test] The following RemoveAll would clear data
     // before TearDown snapshot verification
     // // Verify soft pin does not block RemoveAll
     // ASSERT_TRUE(
-    //     service_->PutStart(client_id, key, slice_length,
+    //     service_->PutStart(client_id, key, "default", slice_length, //
     //     config).has_value());
     // ASSERT_TRUE(
-    //     service_->PutEnd(client_id, key, ReplicaType::MEMORY).has_value());
+    //     service_->PutEnd(client_id, key, "default",
+    //     ReplicaType::MEMORY).has_value());
     // EXPECT_EQ(1, service_->RemoveAll());
 }
 
@@ -1661,11 +1679,12 @@ TEST_F(MasterServiceSnapshotTest, SoftPinObjectsNotEvictedBeforeOtherObjects) {
             soft_pin_config.with_soft_pin = true;
 
             ASSERT_TRUE(service_
-                            ->PutStart(client_id, pin_key, slice_length,
-                                       soft_pin_config)
+                            ->PutStart(client_id, pin_key, "default",
+                                       slice_length, soft_pin_config)
                             .has_value());
             ASSERT_TRUE(
-                service_->PutEnd(client_id, pin_key, ReplicaType::MEMORY)
+                service_
+                    ->PutEnd(client_id, pin_key, "default", ReplicaType::MEMORY)
                     .has_value());
         }
 
@@ -1676,10 +1695,12 @@ TEST_F(MasterServiceSnapshotTest, SoftPinObjectsNotEvictedBeforeOtherObjects) {
             uint64_t slice_length = value_size;
             ReplicateConfig config;
             config.replica_num = 1;
-            if (service_->PutStart(client_id, key, slice_length, config)
+            if (service_
+                    ->PutStart(client_id, key, "default", slice_length, config)
                     .has_value()) {
                 ASSERT_TRUE(
-                    service_->PutEnd(client_id, key, ReplicaType::MEMORY)
+                    service_
+                        ->PutEnd(client_id, key, "default", ReplicaType::MEMORY)
                         .has_value());
             } else {
                 failed_puts++;
@@ -1692,7 +1713,8 @@ TEST_F(MasterServiceSnapshotTest, SoftPinObjectsNotEvictedBeforeOtherObjects) {
         // pin_key should still be accessible
         for (int i = 0; i < 2; i++) {
             std::string pin_key = "pin_key" + std::to_string(i);
-            ASSERT_TRUE(service_->GetReplicaList(pin_key).has_value());
+            ASSERT_TRUE(
+                service_->GetReplicaList(pin_key, "default").has_value());
         }
 
         // wait for the lease to expire
@@ -1735,10 +1757,11 @@ TEST_F(MasterServiceSnapshotTest, SoftPinObjectsCanBeEvicted) {
         ReplicateConfig config;
         config.replica_num = 1;
         config.with_soft_pin = true;
-        if (service_->PutStart(client_id, key, slice_length, config)
+        if (service_->PutStart(client_id, key, "default", slice_length, config)
                 .has_value()) {
-            ASSERT_TRUE(service_->PutEnd(client_id, key, ReplicaType::MEMORY)
-                            .has_value());
+            ASSERT_TRUE(
+                service_->PutEnd(client_id, key, "default", ReplicaType::MEMORY)
+                    .has_value());
             success_puts++;
         } else {
             // wait for eviction to work
@@ -1795,10 +1818,11 @@ TEST_F(MasterServiceSnapshotTest, SoftPinExtendedOnGet) {
             soft_pin_config.replica_num = 1;
             soft_pin_config.with_soft_pin = true;
 
-            ASSERT_TRUE(service_->PutStart(client_id, pin_key, slice_length,
-                                           soft_pin_config));
+            ASSERT_TRUE(service_->PutStart(client_id, pin_key, "default",
+                                           slice_length, soft_pin_config));
             ASSERT_TRUE(
-                service_->PutEnd(client_id, pin_key, ReplicaType::MEMORY)
+                service_
+                    ->PutEnd(client_id, pin_key, "default", ReplicaType::MEMORY)
                     .has_value());
         }
 
@@ -1808,7 +1832,8 @@ TEST_F(MasterServiceSnapshotTest, SoftPinExtendedOnGet) {
         // Get the pin_key to extend the soft pin
         for (int i = 0; i < 2; i++) {
             std::string pin_key = "pin_key" + std::to_string(i);
-            ASSERT_TRUE(service_->GetReplicaList(pin_key).has_value());
+            ASSERT_TRUE(
+                service_->GetReplicaList(pin_key, "default").has_value());
         }
 
         // Fill the segment to trigger eviction
@@ -1818,10 +1843,12 @@ TEST_F(MasterServiceSnapshotTest, SoftPinExtendedOnGet) {
             uint64_t slice_length = value_size;
             ReplicateConfig config;
             config.replica_num = 1;
-            if (service_->PutStart(client_id, key, slice_length, config)
+            if (service_
+                    ->PutStart(client_id, key, "default", slice_length, config)
                     .has_value()) {
                 ASSERT_TRUE(
-                    service_->PutEnd(client_id, key, ReplicaType::MEMORY)
+                    service_
+                        ->PutEnd(client_id, key, "default", ReplicaType::MEMORY)
                         .has_value());
             } else {
                 failed_puts++;
@@ -1835,7 +1862,8 @@ TEST_F(MasterServiceSnapshotTest, SoftPinExtendedOnGet) {
         // pin_key should still be accessible
         for (int i = 0; i < 2; i++) {
             std::string pin_key = "pin_key" + std::to_string(i);
-            ASSERT_TRUE(service_->GetReplicaList(pin_key).has_value());
+            ASSERT_TRUE(
+                service_->GetReplicaList(pin_key, "default").has_value());
         }
         // [Commented for snapshot test] The following RemoveAll would clear
         // data before TearDown snapshot verification Only remove all objects
@@ -1879,10 +1907,11 @@ TEST_F(MasterServiceSnapshotTest, SoftPinObjectsNotAllowEvict) {
         ReplicateConfig config;
         config.replica_num = 1;
         config.with_soft_pin = true;
-        if (service_->PutStart(client_id, key, slice_length, config)
+        if (service_->PutStart(client_id, key, "default", slice_length, config)
                 .has_value()) {
-            ASSERT_TRUE(service_->PutEnd(client_id, key, ReplicaType::MEMORY)
-                            .has_value());
+            ASSERT_TRUE(
+                service_->PutEnd(client_id, key, "default", ReplicaType::MEMORY)
+                    .has_value());
             success_keys.push_back(key);
         } else {
             // wait for eviction to work
@@ -1892,7 +1921,7 @@ TEST_F(MasterServiceSnapshotTest, SoftPinObjectsNotAllowEvict) {
     ASSERT_LE(success_keys.size(), 17);
     // All soft pinned objects should be accessible
     for (const auto& key : success_keys) {
-        ASSERT_TRUE(service_->GetReplicaList(key).has_value());
+        ASSERT_TRUE(service_->GetReplicaList(key, "default").has_value());
     }
     // [Commented for snapshot test] The following RemoveAll would clear data
     // before TearDown snapshot verification
@@ -1920,7 +1949,7 @@ TEST_F(MasterServiceSnapshotTest, ReplicaSegmentsAreUnique) {
     config.replica_num = 10;
 
     auto put_start_result =
-        service_->PutStart(client_id, key, slice_length, config);
+        service_->PutStart(client_id, key, "default", slice_length, config);
     ASSERT_TRUE(put_start_result.has_value());
     auto replica_list_local = put_start_result.value();
     ASSERT_EQ(config.replica_num, replica_list_local.size());
@@ -1936,8 +1965,8 @@ TEST_F(MasterServiceSnapshotTest, ReplicaSegmentsAreUnique) {
     EXPECT_EQ(segment_names.size(), config.replica_num)
         << "Duplicate segment found";
 
-    ASSERT_TRUE(
-        service_->PutEnd(client_id, key, ReplicaType::MEMORY).has_value());
+    ASSERT_TRUE(service_->PutEnd(client_id, key, "default", ReplicaType::MEMORY)
+                    .has_value());
 }
 
 TEST_F(MasterServiceSnapshotTest, ReplicationFactorTwoWithSingleSegment) {
@@ -1958,7 +1987,7 @@ TEST_F(MasterServiceSnapshotTest, ReplicationFactorTwoWithSingleSegment) {
     config.replica_num = 2;
 
     auto put_start_result =
-        service_->PutStart(client_id, key, slice_length, config);
+        service_->PutStart(client_id, key, "default", slice_length, config);
     ASSERT_TRUE(put_start_result.has_value());
     auto replicas = put_start_result.value();
 
@@ -1990,23 +2019,23 @@ TEST_F(MasterServiceSnapshotTest, BatchExistKeyTest) {
         ReplicateConfig config;
         config.replica_num = 1;
         uint64_t slice_length = value_size;
-        auto put_start_result =
-            service_->PutStart(client_id, test_keys[i], slice_length, config);
+        auto put_start_result = service_->PutStart(
+            client_id, test_keys[i], "default", slice_length, config);
         ASSERT_TRUE(put_start_result.has_value());
-        auto put_end_result =
-            service_->PutEnd(client_id, test_keys[i], ReplicaType::MEMORY);
+        auto put_end_result = service_->PutEnd(client_id, test_keys[i],
+                                               "default", ReplicaType::MEMORY);
         ASSERT_TRUE(put_end_result.has_value());
     }
 
     // Test individual ExistKey calls to verify the underlying functionality
     for (int i = 0; i < test_object_num; ++i) {
-        auto exist_result = service_->ExistKey(test_keys[i]);
+        auto exist_result = service_->ExistKey(test_keys[i], "default");
         EXPECT_TRUE(exist_result.value());
     }
 
     // Tets batch
     test_keys.push_back("non_existent_key");
-    auto exist_resp = service_->BatchExistKey(test_keys);
+    auto exist_resp = service_->BatchExistKey(test_keys, "default");
     for (int i = 0; i < test_object_num; ++i) {
         ASSERT_TRUE(exist_resp[i].value());
     }
@@ -2156,6 +2185,117 @@ TEST_F(MasterServiceSnapshotTest,
         << "Client with all empty te_endpoints should have empty IP vector";
 }
 
+TEST_F(MasterServiceSnapshotTest, BatchQueryIpBracketedIpv6Test) {
+    service_.reset(new MasterService());
+    const UUID client_id = generate_uuid();
+
+    // Mount a segment with a bracketed IPv6 endpoint
+    constexpr size_t buffer = 0x300000000;
+    constexpr size_t size = 1024 * 1024 * 16;
+    Segment segment = MakeSegment("test_segment", buffer, size);
+    segment.te_endpoint = "[::1]:17813";
+    auto mount_result = service_->MountSegment(segment, client_id);
+    ASSERT_TRUE(mount_result.has_value());
+
+    std::vector<UUID> client_ids = {client_id};
+    auto query_result = service_->BatchQueryIp(client_ids);
+    ASSERT_TRUE(query_result.has_value());
+
+    const auto& results = query_result.value();
+    auto it = results.find(client_id);
+    ASSERT_NE(it, results.end());
+
+    const auto& ip_addresses = it->second;
+    ASSERT_EQ(1u, ip_addresses.size());
+    EXPECT_EQ("::1", ip_addresses[0]);
+}
+
+TEST_F(MasterServiceSnapshotTest, BatchQueryIpLinkLocalIpv6WithScopeTest) {
+    service_.reset(new MasterService());
+    const UUID client_id = generate_uuid();
+
+    // Mount a segment with a link-local IPv6 address with scope ID
+    constexpr size_t buffer = 0x300000000;
+    constexpr size_t size = 1024 * 1024 * 16;
+    Segment segment = MakeSegment("test_segment", buffer, size);
+    segment.te_endpoint = "fe80::a236:bcff:fecb:a1be%eno2:15773";
+    auto mount_result = service_->MountSegment(segment, client_id);
+    ASSERT_TRUE(mount_result.has_value());
+
+    std::vector<UUID> client_ids = {client_id};
+    auto query_result = service_->BatchQueryIp(client_ids);
+    ASSERT_TRUE(query_result.has_value());
+
+    const auto& results = query_result.value();
+    auto it = results.find(client_id);
+    ASSERT_NE(it, results.end());
+
+    const auto& ip_addresses = it->second;
+    ASSERT_EQ(1u, ip_addresses.size());
+    EXPECT_EQ("fe80::a236:bcff:fecb:a1be%eno2", ip_addresses[0]);
+}
+
+TEST_F(MasterServiceSnapshotTest, BatchQueryIpIpv6NoPortTest) {
+    service_.reset(new MasterService());
+    const UUID client_id = generate_uuid();
+
+    // Mount a segment with an IPv6 address without port
+    constexpr size_t buffer = 0x300000000;
+    constexpr size_t size = 1024 * 1024 * 16;
+    Segment segment = MakeSegment("test_segment", buffer, size);
+    segment.te_endpoint = "::1";
+    auto mount_result = service_->MountSegment(segment, client_id);
+    ASSERT_TRUE(mount_result.has_value());
+
+    std::vector<UUID> client_ids = {client_id};
+    auto query_result = service_->BatchQueryIp(client_ids);
+    ASSERT_TRUE(query_result.has_value());
+
+    const auto& results = query_result.value();
+    auto it = results.find(client_id);
+    ASSERT_NE(it, results.end());
+
+    const auto& ip_addresses = it->second;
+    ASSERT_EQ(1u, ip_addresses.size());
+    EXPECT_EQ("::1", ip_addresses[0]);
+}
+
+TEST_F(MasterServiceSnapshotTest, BatchQueryIpMixedIpv4AndIpv6Test) {
+    service_.reset(new MasterService());
+    const UUID client_id = generate_uuid();
+
+    // Mount segments with IPv4 and IPv6 endpoints for the same client
+    constexpr size_t buffer1 = 0x300000000;
+    constexpr size_t buffer2 = 0x400000000;
+    constexpr size_t size = 1024 * 1024 * 16;
+
+    Segment segment1 = MakeSegment("segment1", buffer1, size);
+    segment1.te_endpoint = "192.168.1.1:12345";
+    auto mount_result1 = service_->MountSegment(segment1, client_id);
+    ASSERT_TRUE(mount_result1.has_value());
+
+    Segment segment2 = MakeSegment("segment2", buffer2, size);
+    segment2.te_endpoint = "[::1]:17813";
+    auto mount_result2 = service_->MountSegment(segment2, client_id);
+    ASSERT_TRUE(mount_result2.has_value());
+
+    std::vector<UUID> client_ids = {client_id};
+    auto query_result = service_->BatchQueryIp(client_ids);
+    ASSERT_TRUE(query_result.has_value());
+
+    const auto& results = query_result.value();
+    auto it = results.find(client_id);
+    ASSERT_NE(it, results.end());
+
+    const auto& ip_addresses = it->second;
+    ASSERT_EQ(2u, ip_addresses.size());
+
+    std::unordered_set<std::string> ip_set(ip_addresses.begin(),
+                                           ip_addresses.end());
+    EXPECT_NE(ip_set.find("192.168.1.1"), ip_set.end());
+    EXPECT_NE(ip_set.find("::1"), ip_set.end());
+}
+
 TEST_F(MasterServiceSnapshotTest, PutStartExpiringTest) {
     // Reset storage space metrics.
     MasterMetricManager::instance().reset_allocated_mem_size();
@@ -2190,7 +2330,7 @@ TEST_F(MasterServiceSnapshotTest, PutStartExpiringTest) {
 
     // Put key_1, should success.
     auto put_start_result =
-        service_->PutStart(client_id, key_1, slice_length, config);
+        service_->PutStart(client_id, key_1, "default", slice_length, config);
     EXPECT_TRUE(put_start_result.has_value());
     replica_list = put_start_result.value();
     EXPECT_EQ(replica_list.size(), kReplicaCnt);
@@ -2200,7 +2340,7 @@ TEST_F(MasterServiceSnapshotTest, PutStartExpiringTest) {
 
     // Put key_1 again, should fail because the key exists.
     put_start_result =
-        service_->PutStart(client_id, key_1, slice_length, config);
+        service_->PutStart(client_id, key_1, "default", slice_length, config);
     EXPECT_FALSE(put_start_result.has_value());
     EXPECT_EQ(put_start_result.error(), ErrorCode::OBJECT_ALREADY_EXISTS);
 
@@ -2216,7 +2356,7 @@ TEST_F(MasterServiceSnapshotTest, PutStartExpiringTest) {
     // Put key_1 again, should success because the old one has expired and will
     // be discarded by this put.
     put_start_result =
-        service_->PutStart(client_id, key_1, slice_length, config);
+        service_->PutStart(client_id, key_1, "default", slice_length, config);
     EXPECT_TRUE(put_start_result.has_value());
     replica_list = put_start_result.value();
     EXPECT_EQ(replica_list.size(), kReplicaCnt);
@@ -2226,17 +2366,17 @@ TEST_F(MasterServiceSnapshotTest, PutStartExpiringTest) {
 
     // Complete key_1.
     auto put_end_result =
-        service_->PutEnd(client_id, key_1, ReplicaType::MEMORY);
+        service_->PutEnd(client_id, key_1, "default", ReplicaType::MEMORY);
     EXPECT_TRUE(put_end_result.has_value());
 
     // Protect key_1 from eviction.
-    auto get_result = service_->GetReplicaList(key_1);
+    auto get_result = service_->GetReplicaList(key_1, "default");
     EXPECT_TRUE(get_result.has_value());
 
     // Put key_2, should fail because the key_1 occupied 12MB (6MB processing,
     // 6MB discarded but not yet released) on each segment.
     put_start_result =
-        service_->PutStart(client_id, key_2, slice_length, config);
+        service_->PutStart(client_id, key_2, "default", slice_length, config);
     EXPECT_FALSE(put_start_result.has_value());
     EXPECT_EQ(put_start_result.error(), ErrorCode::NO_AVAILABLE_HANDLE);
 
@@ -2249,7 +2389,7 @@ TEST_F(MasterServiceSnapshotTest, PutStartExpiringTest) {
             EXPECT_TRUE(result.has_value());
         }
         // Protect key_1 from eviction.
-        auto get_result = service_->GetReplicaList(key_1);
+        auto get_result = service_->GetReplicaList(key_1, "default");
         EXPECT_TRUE(get_result.has_value());
         std::this_thread::sleep_for(std::chrono::seconds(1));
     }
@@ -2257,7 +2397,7 @@ TEST_F(MasterServiceSnapshotTest, PutStartExpiringTest) {
     // Put key_2 again, should success because the discarded replica has been
     // released.
     put_start_result =
-        service_->PutStart(client_id, key_2, slice_length, config);
+        service_->PutStart(client_id, key_2, "default", slice_length, config);
     EXPECT_TRUE(put_start_result.has_value());
     replica_list = put_start_result.value();
     EXPECT_EQ(replica_list.size(), kReplicaCnt);
@@ -2272,7 +2412,7 @@ TEST_F(MasterServiceSnapshotTest, PutStartExpiringTest) {
             EXPECT_TRUE(result.has_value());
         }
         // Protect key_1 from eviction.
-        auto get_result = service_->GetReplicaList(key_1);
+        auto get_result = service_->GetReplicaList(key_1, "default");
         EXPECT_TRUE(get_result.has_value());
         std::this_thread::sleep_for(std::chrono::seconds(1));
     }
@@ -2280,7 +2420,7 @@ TEST_F(MasterServiceSnapshotTest, PutStartExpiringTest) {
     // Put key_2 again, should fail because eviction has not been triggered. And
     // this PutStart should trigger the eviction.
     put_start_result =
-        service_->PutStart(client_id, key_2, slice_length, config);
+        service_->PutStart(client_id, key_2, "default", slice_length, config);
     EXPECT_FALSE(put_start_result.has_value());
     EXPECT_EQ(put_start_result.error(), ErrorCode::NO_AVAILABLE_HANDLE);
 
@@ -2290,7 +2430,7 @@ TEST_F(MasterServiceSnapshotTest, PutStartExpiringTest) {
     // Put key_2 again, should success because the previous one has been
     // discarded and released.
     put_start_result =
-        service_->PutStart(client_id, key_2, slice_length, config);
+        service_->PutStart(client_id, key_2, "default", slice_length, config);
     EXPECT_TRUE(put_start_result.has_value());
     replica_list = put_start_result.value();
     EXPECT_EQ(replica_list.size(), kReplicaCnt);
@@ -2299,7 +2439,8 @@ TEST_F(MasterServiceSnapshotTest, PutStartExpiringTest) {
     }
 
     // Complete key_2.
-    put_end_result = service_->PutEnd(client_id, key_2, ReplicaType::MEMORY);
+    put_end_result =
+        service_->PutEnd(client_id, key_2, "default", ReplicaType::MEMORY);
     EXPECT_TRUE(put_end_result.has_value());
 }
 
@@ -2369,8 +2510,11 @@ TEST_F(MasterServiceSnapshotTest, OffloadObjectHeartbeat) {
     }
     ASSERT_EQ(res->size(), keys.size());
     for (auto& key : keys) {
-        ASSERT_TRUE(res.value().find(key) != res.value().end());
-        ASSERT_EQ(res.value().find(key)->second, 1024);
+        auto it = std::find_if(
+            res->begin(), res->end(),
+            [&key](const OffloadTaskItem& task) { return task.key == key; });
+        ASSERT_TRUE(it != res->end());
+        ASSERT_EQ(it->size, 1024);
     }
 
     keys.clear();
@@ -2386,8 +2530,11 @@ TEST_F(MasterServiceSnapshotTest, OffloadObjectHeartbeat) {
     }
     ASSERT_EQ(res->size(), keys.size());
     for (auto& key : keys) {
-        ASSERT_TRUE(res.value().find(key) != res.value().end());
-        ASSERT_EQ(res.value().find(key)->second, 1024);
+        auto it = std::find_if(
+            res->begin(), res->end(),
+            [&key](const OffloadTaskItem& task) { return task.key == key; });
+        ASSERT_TRUE(it != res->end());
+        ASSERT_EQ(it->size, 1024);
     }
 }
 
@@ -2410,16 +2557,16 @@ TEST_F(MasterServiceSnapshotTest, BatchReplicaClearAllSegments) {
         ReplicateConfig config;
         config.replica_num = 1;
         auto put_start_result =
-            service_->PutStart(client_id, key, value_length, config);
+            service_->PutStart(client_id, key, "default", value_length, config);
         ASSERT_TRUE(put_start_result.has_value());
         auto put_end_result =
-            service_->PutEnd(client_id, key, ReplicaType::MEMORY);
+            service_->PutEnd(client_id, key, "default", ReplicaType::MEMORY);
         ASSERT_TRUE(put_end_result.has_value());
     }
 
     // Verify objects exist
     for (const auto& key : keys) {
-        auto exist_result = service_->ExistKey(key);
+        auto exist_result = service_->ExistKey(key, "default");
         ASSERT_TRUE(exist_result.has_value());
         ASSERT_TRUE(exist_result.value());
     }
@@ -2436,7 +2583,7 @@ TEST_F(MasterServiceSnapshotTest, BatchReplicaClearAllSegments) {
 
     // Verify objects are removed
     for (const auto& key : keys) {
-        auto exist_result = service_->ExistKey(key);
+        auto exist_result = service_->ExistKey(key, "default");
         ASSERT_TRUE(exist_result.has_value());
         ASSERT_FALSE(exist_result.value())
             << "Key " << key << " should be removed";
@@ -2467,9 +2614,10 @@ TEST_F(MasterServiceSnapshotTest, BatchReplicaClearSpecificSegment) {
     config.preferred_segment =
         segment_name;  // Ensure object is placed on segment1
     auto put_start_result =
-        service_->PutStart(client_id, key, value_length, config);
+        service_->PutStart(client_id, key, "default", value_length, config);
     ASSERT_TRUE(put_start_result.has_value());
-    auto put_end_result = service_->PutEnd(client_id, key, ReplicaType::MEMORY);
+    auto put_end_result =
+        service_->PutEnd(client_id, key, "default", ReplicaType::MEMORY);
     ASSERT_TRUE(put_end_result.has_value());
 
     // 4. Wait for lease to expire and verify it's actually expired
@@ -2507,7 +2655,7 @@ TEST_F(MasterServiceSnapshotTest, BatchReplicaClearSpecificSegment) {
     const auto& cleared_keys = clear_result.value();
     ASSERT_EQ(1u, cleared_keys.size()) << "Key should be cleared";
 
-    auto exist_result = service_->ExistKey(key);
+    auto exist_result = service_->ExistKey(key, "default");
     ASSERT_TRUE(exist_result.has_value());
     ASSERT_FALSE(exist_result.value())
         << "Key should be removed after being cleared.";
@@ -2528,13 +2676,14 @@ TEST_F(MasterServiceSnapshotTest, BatchReplicaClearWithLeaseActive) {
     ReplicateConfig config;
     config.replica_num = 1;
     auto put_start_result =
-        service_->PutStart(client_id, key, value_length, config);
+        service_->PutStart(client_id, key, "default", value_length, config);
     ASSERT_TRUE(put_start_result.has_value());
-    auto put_end_result = service_->PutEnd(client_id, key, ReplicaType::MEMORY);
+    auto put_end_result =
+        service_->PutEnd(client_id, key, "default", ReplicaType::MEMORY);
     ASSERT_TRUE(put_end_result.has_value());
 
     // Grant a lease by calling GetReplicaList (similar to normal usage)
-    auto get_result = service_->GetReplicaList(key);
+    auto get_result = service_->GetReplicaList(key, "default");
     ASSERT_TRUE(get_result.has_value());
 
     // Try to clear immediately (lease should still be active)
@@ -2548,7 +2697,7 @@ TEST_F(MasterServiceSnapshotTest, BatchReplicaClearWithLeaseActive) {
         << "No keys should be cleared when lease is active";
 
     // Verify object still exists
-    auto exist_result = service_->ExistKey(key);
+    auto exist_result = service_->ExistKey(key, "default");
     ASSERT_TRUE(exist_result.has_value());
     ASSERT_TRUE(exist_result.value()) << "Key should still exist";
 }
@@ -2569,10 +2718,10 @@ TEST_F(MasterServiceSnapshotTest, BatchReplicaClearWithDifferentClientId) {
     ReplicateConfig config;
     config.replica_num = 1;
     auto put_start_result =
-        service_->PutStart(client_id1, key, value_length, config);
+        service_->PutStart(client_id1, key, "default", value_length, config);
     ASSERT_TRUE(put_start_result.has_value());
     auto put_end_result =
-        service_->PutEnd(client_id1, key, ReplicaType::MEMORY);
+        service_->PutEnd(client_id1, key, "default", ReplicaType::MEMORY);
     ASSERT_TRUE(put_end_result.has_value());
 
     // Wait for lease to expire
@@ -2589,7 +2738,7 @@ TEST_F(MasterServiceSnapshotTest, BatchReplicaClearWithDifferentClientId) {
         << "No keys should be cleared for different client_id";
 
     // Verify object still exists
-    auto exist_result = service_->ExistKey(key);
+    auto exist_result = service_->ExistKey(key, "default");
     ASSERT_TRUE(exist_result.has_value());
     ASSERT_TRUE(exist_result.value()) << "Key should still exist";
 }
@@ -2644,11 +2793,11 @@ TEST_F(MasterServiceSnapshotTest, BatchReplicaClearWithEmptyStringKeys) {
     uint64_t value_length = 1024;
     ReplicateConfig config;
     config.replica_num = 1;
-    auto put_start_result =
-        service_->PutStart(client_id, valid_key, value_length, config);
+    auto put_start_result = service_->PutStart(client_id, valid_key, "default",
+                                               value_length, config);
     ASSERT_TRUE(put_start_result.has_value());
     auto put_end_result =
-        service_->PutEnd(client_id, valid_key, ReplicaType::MEMORY);
+        service_->PutEnd(client_id, valid_key, "default", ReplicaType::MEMORY);
     ASSERT_TRUE(put_end_result.has_value());
 
     // Wait for lease to expire
@@ -2686,22 +2835,25 @@ TEST_F(MasterServiceSnapshotTest, BatchReplicaClearMixedScenario) {
 
     // Create key1 and key2 with client_id1
     auto put_start1 =
-        service_->PutStart(client_id1, key1, value_length, config);
+        service_->PutStart(client_id1, key1, "default", value_length, config);
     ASSERT_TRUE(put_start1.has_value());
-    auto put_end1 = service_->PutEnd(client_id1, key1, ReplicaType::MEMORY);
+    auto put_end1 =
+        service_->PutEnd(client_id1, key1, "default", ReplicaType::MEMORY);
     ASSERT_TRUE(put_end1.has_value());
 
     auto put_start2 =
-        service_->PutStart(client_id1, key2, value_length, config);
+        service_->PutStart(client_id1, key2, "default", value_length, config);
     ASSERT_TRUE(put_start2.has_value());
-    auto put_end2 = service_->PutEnd(client_id1, key2, ReplicaType::MEMORY);
+    auto put_end2 =
+        service_->PutEnd(client_id1, key2, "default", ReplicaType::MEMORY);
     ASSERT_TRUE(put_end2.has_value());
 
     // Create key3 with client_id2
     auto put_start3 =
-        service_->PutStart(client_id2, key3, value_length, config);
+        service_->PutStart(client_id2, key3, "default", value_length, config);
     ASSERT_TRUE(put_start3.has_value());
-    auto put_end3 = service_->PutEnd(client_id2, key3, ReplicaType::MEMORY);
+    auto put_end3 =
+        service_->PutEnd(client_id2, key3, "default", ReplicaType::MEMORY);
     ASSERT_TRUE(put_end3.has_value());
 
     // Wait for lease to expire
@@ -2719,16 +2871,16 @@ TEST_F(MasterServiceSnapshotTest, BatchReplicaClearMixedScenario) {
         << "Only keys belonging to client_id1 should be cleared";
 
     // Verify key1 and key2 are cleared
-    auto exist1 = service_->ExistKey(key1);
+    auto exist1 = service_->ExistKey(key1, "default");
     ASSERT_TRUE(exist1.has_value());
     ASSERT_FALSE(exist1.value()) << "key1 should be cleared";
 
-    auto exist2 = service_->ExistKey(key2);
+    auto exist2 = service_->ExistKey(key2, "default");
     ASSERT_TRUE(exist2.has_value());
     ASSERT_FALSE(exist2.value()) << "key2 should be cleared";
 
     // Verify key3 still exists (different client_id)
-    auto exist3 = service_->ExistKey(key3);
+    auto exist3 = service_->ExistKey(key3, "default");
     ASSERT_TRUE(exist3.has_value());
     ASSERT_TRUE(exist3.value())
         << "key3 should still exist (different client_id)";
@@ -2764,15 +2916,15 @@ TEST_F(MasterServiceSnapshotTest, CreateCopyTaskTest) {
     config.replica_num = 1;
     config.preferred_segment = "segment_0";
     auto put_start_result =
-        service_->PutStart(client_id, key1, slice_length, config);
+        service_->PutStart(client_id, key1, "default", slice_length, config);
     EXPECT_TRUE(put_start_result.has_value());
     auto put_end_result =
-        service_->PutEnd(client_id, key1, ReplicaType::MEMORY);
+        service_->PutEnd(client_id, key1, "default", ReplicaType::MEMORY);
     EXPECT_TRUE(put_end_result.has_value());
 
     // Copy key1 to "segment_1" and "segment_2"
     auto copy_result =
-        service_->CreateCopyTask(key1, {"segment_1", "segment_2"});
+        service_->CreateCopyTask(key1, "default", {"segment_1", "segment_2"});
     EXPECT_TRUE(copy_result.has_value());
 
     // verify the copy task is created and assigned to the client who executed
@@ -2783,18 +2935,19 @@ TEST_F(MasterServiceSnapshotTest, CreateCopyTaskTest) {
     EXPECT_EQ(contexts[0].client_id, task.value().assigned_client);
 
     // Copy with empty targets should fail
-    auto copy_result1 = service_->CreateCopyTask(key1, {});
+    auto copy_result1 = service_->CreateCopyTask(key1, "default", {});
     EXPECT_FALSE(copy_result1.has_value());
     EXPECT_EQ(ErrorCode::INVALID_PARAMS, copy_result1.error());
 
     // Copy not exist key should fail
     auto copy_result2 =
-        service_->CreateCopyTask("not_exist_key", {"segment_1"});
+        service_->CreateCopyTask("not_exist_key", "default", {"segment_1"});
     EXPECT_FALSE(copy_result2.has_value());
     EXPECT_EQ(ErrorCode::OBJECT_NOT_FOUND, copy_result2.error());
 
     // Copy to segment that not mounted should fail
-    auto copy_result3 = service_->CreateCopyTask(key1, {"not_mounted_segment"});
+    auto copy_result3 =
+        service_->CreateCopyTask(key1, "default", {"not_mounted_segment"});
     EXPECT_FALSE(copy_result3.has_value());
     EXPECT_EQ(ErrorCode::INVALID_PARAMS, copy_result3.error());
 }
@@ -2829,14 +2982,15 @@ TEST_F(MasterServiceSnapshotTest, CreateMoveTaskTest) {
     config.replica_num = 1;
     config.preferred_segment = "segment_0";
     auto put_start_result =
-        service_->PutStart(client_id, key1, slice_length, config);
+        service_->PutStart(client_id, key1, "default", slice_length, config);
     EXPECT_TRUE(put_start_result.has_value());
     auto put_end_result =
-        service_->PutEnd(client_id, key1, ReplicaType::MEMORY);
+        service_->PutEnd(client_id, key1, "default", ReplicaType::MEMORY);
     EXPECT_TRUE(put_end_result.has_value());
 
     // Move key1 from "segment_0" to "segment_1"
-    auto move_result = service_->CreateMoveTask(key1, "segment_0", "segment_1");
+    auto move_result =
+        service_->CreateMoveTask(key1, "default", "segment_0", "segment_1");
     EXPECT_TRUE(move_result.has_value());
 
     // Verify the move task is created and assigned to the client owning the
@@ -2847,32 +3001,32 @@ TEST_F(MasterServiceSnapshotTest, CreateMoveTaskTest) {
     EXPECT_EQ(contexts[0].client_id, task.value().assigned_client);
 
     // Move non-existent key should fail
-    auto move_result1 =
-        service_->CreateMoveTask("not_exist_key", "segment_0", "segment_1");
+    auto move_result1 = service_->CreateMoveTask("not_exist_key", "default",
+                                                 "segment_0", "segment_1");
     EXPECT_FALSE(move_result1.has_value());
     EXPECT_EQ(ErrorCode::OBJECT_NOT_FOUND, move_result1.error());
 
     // Move to segment that is same as source should fail
     auto move_result_same =
-        service_->CreateMoveTask(key1, "segment_1", "segment_1");
+        service_->CreateMoveTask(key1, "default", "segment_1", "segment_1");
     EXPECT_FALSE(move_result_same.has_value());
     EXPECT_EQ(ErrorCode::INVALID_PARAMS, move_result_same.error());
 
     // Move to segment that is not mounted should fail
-    auto move_result2 =
-        service_->CreateMoveTask(key1, "segment_0", "not_mounted_segment");
+    auto move_result2 = service_->CreateMoveTask(key1, "default", "segment_0",
+                                                 "not_mounted_segment");
     EXPECT_FALSE(move_result2.has_value());
     EXPECT_EQ(ErrorCode::INVALID_PARAMS, move_result2.error());
 
     // Move from segment that does not have the replica should fail
     auto move_result3 =
-        service_->CreateMoveTask(key1, "segment_2", "segment_1");
+        service_->CreateMoveTask(key1, "default", "segment_2", "segment_1");
     EXPECT_FALSE(move_result3.has_value());
     EXPECT_EQ(ErrorCode::INVALID_PARAMS, move_result3.error());
 
     // Move from segment that is not mounted should fail
-    auto move_result4 =
-        service_->CreateMoveTask(key1, "not_mounted_segment", "segment_1");
+    auto move_result4 = service_->CreateMoveTask(
+        key1, "default", "not_mounted_segment", "segment_1");
     EXPECT_FALSE(move_result4.has_value());
     EXPECT_EQ(ErrorCode::INVALID_PARAMS, move_result4.error());
 }
@@ -2907,14 +3061,15 @@ TEST_F(MasterServiceSnapshotTest, QueryTaskTest) {
     config.replica_num = 1;
     config.preferred_segment = "segment_0";
     auto put_start_result =
-        service_->PutStart(client_id, key1, slice_length, config);
+        service_->PutStart(client_id, key1, "default", slice_length, config);
     EXPECT_TRUE(put_start_result.has_value());
     auto put_end_result =
-        service_->PutEnd(client_id, key1, ReplicaType::MEMORY);
+        service_->PutEnd(client_id, key1, "default", ReplicaType::MEMORY);
     EXPECT_TRUE(put_end_result.has_value());
 
     // Move key1 from "segment_0" to "segment_1"
-    auto move_result = service_->CreateMoveTask(key1, "segment_0", "segment_1");
+    auto move_result =
+        service_->CreateMoveTask(key1, "default", "segment_0", "segment_1");
     EXPECT_TRUE(move_result.has_value());
 
     // Query non-existent task should fail
@@ -2955,18 +3110,21 @@ TEST_F(MasterServiceSnapshotTest,
     config.replica_num = 1;
     config.preferred_segment = "segment_0";
 
+    ASSERT_TRUE(service_
+                    ->PutStart(put_client_id, key, "default",
+                               /*slice_length=*/1024, config)
+                    .has_value());
     ASSERT_TRUE(
-        service_->PutStart(put_client_id, key, /*slice_length=*/1024, config)
+        service_->PutEnd(put_client_id, key, "default", ReplicaType::MEMORY)
             .has_value());
-    ASSERT_TRUE(
-        service_->PutEnd(put_client_id, key, ReplicaType::MEMORY).has_value());
 
     // Create two tasks; both should be assigned to the client owning source
     // segment_0.
-    auto copy_task_id = service_->CreateCopyTask(key, {"segment_1"});
+    auto copy_task_id = service_->CreateCopyTask(key, "default", {"segment_1"});
     ASSERT_TRUE(copy_task_id.has_value());
 
-    auto move_task_id = service_->CreateMoveTask(key, "segment_0", "segment_1");
+    auto move_task_id =
+        service_->CreateMoveTask(key, "default", "segment_0", "segment_1");
     ASSERT_TRUE(move_task_id.has_value());
 
     // Fetch from client_0 should get both tasks (order not guaranteed).
@@ -3015,15 +3173,18 @@ TEST_F(MasterServiceSnapshotTest, FetchTasksRespectsBatchSize) {
     config.replica_num = 1;
     config.preferred_segment = "segment_0";
 
+    ASSERT_TRUE(service_
+                    ->PutStart(put_client_id, key, "default",
+                               /*slice_length=*/1024, config)
+                    .has_value());
     ASSERT_TRUE(
-        service_->PutStart(put_client_id, key, /*slice_length=*/1024, config)
+        service_->PutEnd(put_client_id, key, "default", ReplicaType::MEMORY)
             .has_value());
-    ASSERT_TRUE(
-        service_->PutEnd(put_client_id, key, ReplicaType::MEMORY).has_value());
 
-    auto t1 = service_->CreateCopyTask(key, {"segment_1"});
+    auto t1 = service_->CreateCopyTask(key, "default", {"segment_1"});
     ASSERT_TRUE(t1.has_value());
-    auto t2 = service_->CreateMoveTask(key, "segment_0", "segment_1");
+    auto t2 =
+        service_->CreateMoveTask(key, "default", "segment_0", "segment_1");
     ASSERT_TRUE(t2.has_value());
 
     auto fetch_first = service_->FetchTasks(ctx0.client_id, /*batch_size=*/1);
@@ -3064,14 +3225,16 @@ TEST_F(MasterServiceSnapshotTest, UpdateTaskSuccessFlow) {
     config.replica_num = 1;
     config.preferred_segment = "segment_0";
 
+    ASSERT_TRUE(service_
+                    ->PutStart(put_client_id, key, "default",
+                               /*slice_length=*/1024, config)
+                    .has_value());
     ASSERT_TRUE(
-        service_->PutStart(put_client_id, key, /*slice_length=*/1024, config)
+        service_->PutEnd(put_client_id, key, "default", ReplicaType::MEMORY)
             .has_value());
-    ASSERT_TRUE(
-        service_->PutEnd(put_client_id, key, ReplicaType::MEMORY).has_value());
 
     // Create a task assigned to client owning segment_0.
-    auto task_id_res = service_->CreateCopyTask(key, {"segment_1"});
+    auto task_id_res = service_->CreateCopyTask(key, "default", {"segment_1"});
     ASSERT_TRUE(task_id_res.has_value());
     const UUID task_id = task_id_res.value();
 
@@ -3120,13 +3283,16 @@ TEST_F(MasterServiceSnapshotTest, UpdateTaskRejectsWrongClient) {
     config.replica_num = 1;
     config.preferred_segment = "segment_0";
 
+    ASSERT_TRUE(service_
+                    ->PutStart(put_client_id, key, "default",
+                               /*slice_length=*/1024, config)
+                    .has_value());
     ASSERT_TRUE(
-        service_->PutStart(put_client_id, key, /*slice_length=*/1024, config)
+        service_->PutEnd(put_client_id, key, "default", ReplicaType::MEMORY)
             .has_value());
-    ASSERT_TRUE(
-        service_->PutEnd(put_client_id, key, ReplicaType::MEMORY).has_value());
 
-    auto task_id_res = service_->CreateMoveTask(key, "segment_0", "segment_1");
+    auto task_id_res =
+        service_->CreateMoveTask(key, "default", "segment_0", "segment_1");
     ASSERT_TRUE(task_id_res.has_value());
     const UUID task_id = task_id_res.value();
 
@@ -3183,8 +3349,8 @@ TEST_F(MasterServiceSnapshotTest, CopyStart) {
     UUID client_id = generate_uuid();
 
     // Test Case 1: CopyStart a non-existent key, should fail.
-    auto copy_result = service_->CopyStart(client_id, "non_existent_key",
-                                           "segment_1", {"segment_2"});
+    auto copy_result = service_->CopyStart(
+        client_id, "non_existent_key", "default", "segment_1", {"segment_2"});
     EXPECT_FALSE(copy_result.has_value());
     EXPECT_EQ(ErrorCode::OBJECT_NOT_FOUND, copy_result.error());
 
@@ -3197,22 +3363,23 @@ TEST_F(MasterServiceSnapshotTest, CopyStart) {
     config.preferred_segment = "segment_1";
 
     auto put_start_result =
-        service_->PutStart(client_id, key, slice_length, config);
+        service_->PutStart(client_id, key, "default", slice_length, config);
     ASSERT_TRUE(put_start_result.has_value());
 
     // Test Case 2: CopyStart to segment_2 and segment_3, should fail because
     // the only replica is not completed.
-    copy_result = service_->CopyStart(client_id, key, "segment_1",
+    copy_result = service_->CopyStart(client_id, key, "default", "segment_1",
                                       {"segment_2", "segment_3"});
     EXPECT_FALSE(copy_result.has_value());
     EXPECT_EQ(ErrorCode::REPLICA_NOT_FOUND, copy_result.error());
 
     // PutEnd the object.
-    auto put_end_result = service_->PutEnd(client_id, key, ReplicaType::MEMORY);
+    auto put_end_result =
+        service_->PutEnd(client_id, key, "default", ReplicaType::MEMORY);
     ASSERT_TRUE(put_end_result.has_value());
 
     // Test Case 3: CopyStart to segment_2 and segment_3, should success.
-    copy_result = service_->CopyStart(client_id, key, "segment_1",
+    copy_result = service_->CopyStart(client_id, key, "default", "segment_1",
                                       {"segment_2", "segment_3"});
     EXPECT_TRUE(copy_result.has_value());
     auto copy_response = copy_result.value();
@@ -3221,40 +3388,41 @@ TEST_F(MasterServiceSnapshotTest, CopyStart) {
     EXPECT_EQ(2, copy_response.targets.size());
 
     // Test Case 4: Try remove the object, should fail because it is copying.
-    auto remove_result = service_->Remove(key);
+    auto remove_result = service_->Remove(key, "default");
     EXPECT_FALSE(remove_result.has_value());
     EXPECT_EQ(ErrorCode::REPLICA_IS_NOT_READY, remove_result.error());
 
     // Test Case 5: CopyStart to segment_4, should fail because there is an
     // ongoing copy task.
-    copy_result =
-        service_->CopyStart(client_id, key, "segment_1", {"segment_4"});
+    copy_result = service_->CopyStart(client_id, key, "default", "segment_1",
+                                      {"segment_4"});
     EXPECT_FALSE(copy_result.has_value());
     EXPECT_EQ(ErrorCode::OBJECT_HAS_REPLICATION_TASK, copy_result.error());
 
     // Test Case 6: CopyEnd, should success and the object now has 3 replicas.
-    auto copy_end_result = service_->CopyEnd(client_id, key);
+    auto copy_end_result = service_->CopyEnd(client_id, key, "default");
     EXPECT_TRUE(copy_end_result.has_value());
-    auto get_result = service_->GetReplicaList(key);
+    auto get_result = service_->GetReplicaList(key, "default");
     EXPECT_TRUE(get_result.has_value());
     EXPECT_EQ(3, get_result.value().replicas.size());
 
     // Test Case 7: Copy from a non-existent replica to segment_3 and
     // segment_4, should fail.
-    copy_result = service_->CopyStart(client_id, key, "non_existent_segment",
-                                      {"segment_3", "segment_4"});
+    copy_result =
+        service_->CopyStart(client_id, key, "default", "non_existent_segment",
+                            {"segment_3", "segment_4"});
     EXPECT_FALSE(copy_result.has_value());
     EXPECT_EQ(ErrorCode::REPLICA_NOT_FOUND, copy_result.error());
 
     // Test Case 8: Copy to segment_4 and a non-existent segment, should fail.
-    copy_result = service_->CopyStart(client_id, key, "segment_1",
+    copy_result = service_->CopyStart(client_id, key, "default", "segment_1",
                                       {"segment_4", "non_existent_segment"});
     EXPECT_FALSE(copy_result.has_value());
     EXPECT_EQ(ErrorCode::SEGMENT_NOT_FOUND, copy_result.error());
 
     // Test Case 9: Copy to segment_3 and segment_4, should skip segment_3 and
     // successfully copy to segment_4.
-    copy_result = service_->CopyStart(client_id, key, "segment_1",
+    copy_result = service_->CopyStart(client_id, key, "default", "segment_1",
                                       {"segment_3", "segment_4"});
     EXPECT_TRUE(copy_result.has_value());
     copy_response = copy_result.value();
@@ -3267,16 +3435,16 @@ TEST_F(MasterServiceSnapshotTest, CopyStart) {
                                .buffer_descriptor.transport_endpoint_);
 
     // End the copy operation to clean up state
-    copy_end_result = service_->CopyEnd(client_id, key);
+    copy_end_result = service_->CopyEnd(client_id, key, "default");
     EXPECT_TRUE(copy_end_result.has_value());
-    get_result = service_->GetReplicaList(key);
+    get_result = service_->GetReplicaList(key, "default");
     EXPECT_TRUE(get_result.has_value());
     EXPECT_EQ(4, get_result.value().replicas.size());
 
     // Test Case 10: Copy to segment_4 again, should skip because it's already
     // used.
-    copy_result =
-        service_->CopyStart(client_id, key, "segment_1", {"segment_4"});
+    copy_result = service_->CopyStart(client_id, key, "default", "segment_1",
+                                      {"segment_4"});
     EXPECT_TRUE(copy_result.has_value());
     copy_response = copy_result.value();
     EXPECT_EQ("segment_1", copy_response.source.get_memory_descriptor()
@@ -3289,19 +3457,19 @@ TEST_F(MasterServiceSnapshotTest, CopyStart) {
     std::this_thread::sleep_for(std::chrono::milliseconds(kv_lease_ttl * 2));
 
     // Test Case 11: Try remove the object, should fail because it is copying.
-    remove_result = service_->Remove(key);
+    remove_result = service_->Remove(key, "default");
     EXPECT_FALSE(remove_result.has_value());
     EXPECT_EQ(ErrorCode::OBJECT_HAS_REPLICATION_TASK, remove_result.error());
 
     // Clean up the copy operation
-    // copy_end_result = service_->CopyEnd(client_id, key);
+    // copy_end_result = service_->CopyEnd(client_id, key, "default");
     // EXPECT_TRUE(copy_end_result.has_value());
 
     // Wait for the lease to expire
     // std::this_thread::sleep_for(std::chrono::milliseconds(kv_lease_ttl * 2));
 
     // Test Case 12: Try remove the object, should success.
-    // remove_result = service_->Remove(key);
+    // remove_result = service_->Remove(key, "default");
     // EXPECT_TRUE(remove_result.has_value());
 }
 
@@ -3321,7 +3489,8 @@ TEST_F(MasterServiceSnapshotTest, CopyEnd) {
     UUID invalid_client_id = generate_uuid();
 
     // Test Case 1: CopyEnd a non-existent key, should fail.
-    auto copy_end_result = service_->CopyEnd(client_id, "non_existent_key");
+    auto copy_end_result =
+        service_->CopyEnd(client_id, "non_existent_key", "default");
     EXPECT_FALSE(copy_end_result.has_value());
     EXPECT_EQ(ErrorCode::OBJECT_NOT_FOUND, copy_end_result.error());
 
@@ -3333,45 +3502,46 @@ TEST_F(MasterServiceSnapshotTest, CopyEnd) {
     config.preferred_segment = "segment_1";
 
     auto put_start_result =
-        service_->PutStart(client_id, key, slice_length, config);
+        service_->PutStart(client_id, key, "default", slice_length, config);
     ASSERT_TRUE(put_start_result.has_value());
-    auto put_end_result = service_->PutEnd(client_id, key, ReplicaType::MEMORY);
+    auto put_end_result =
+        service_->PutEnd(client_id, key, "default", ReplicaType::MEMORY);
     ASSERT_TRUE(put_end_result.has_value());
 
     // Test Case 2: CopyEnd the object, should fail because there is no ongoing
     // copy task.
-    copy_end_result = service_->CopyEnd(client_id, key);
+    copy_end_result = service_->CopyEnd(client_id, key, "default");
     EXPECT_FALSE(copy_end_result.has_value());
     EXPECT_EQ(ErrorCode::OBJECT_NO_REPLICATION_TASK, copy_end_result.error());
 
     // CopyStart the object to segment_2
-    auto copy_start_result =
-        service_->CopyStart(client_id, key, "segment_1", {"segment_2"});
+    auto copy_start_result = service_->CopyStart(client_id, key, "default",
+                                                 "segment_1", {"segment_2"});
     ASSERT_TRUE(copy_start_result.has_value());
 
     // Test Case 3: CopyEnd with an invalid client id, should fail.
-    copy_end_result = service_->CopyEnd(invalid_client_id, key);
+    copy_end_result = service_->CopyEnd(invalid_client_id, key, "default");
     EXPECT_FALSE(copy_end_result.has_value());
     EXPECT_EQ(ErrorCode::ILLEGAL_CLIENT, copy_end_result.error());
 
     // Test Case 4: MoveEnd the object, should fail because the ongoing task is
     // Copy.
-    auto move_end_result = service_->MoveEnd(client_id, key);
+    auto move_end_result = service_->MoveEnd(client_id, key, "default");
     EXPECT_FALSE(move_end_result.has_value());
     EXPECT_EQ(ErrorCode::INVALID_PARAMS, move_end_result.error());
 
     // Test Case 5: CopyEnd, should success.
-    copy_end_result = service_->CopyEnd(client_id, key);
+    copy_end_result = service_->CopyEnd(client_id, key, "default");
     EXPECT_TRUE(copy_end_result.has_value());
 
     // Verify we now have 2 replicas
-    auto get_result = service_->GetReplicaList(key);
+    auto get_result = service_->GetReplicaList(key, "default");
     EXPECT_TRUE(get_result.has_value());
     EXPECT_EQ(2, get_result.value().replicas.size());
 
     // CopyStart the object from segment_1 to segment_3, then unmount segment_1
-    copy_start_result =
-        service_->CopyStart(client_id, key, "segment_1", {"segment_3"});
+    copy_start_result = service_->CopyStart(client_id, key, "default",
+                                            "segment_1", {"segment_3"});
     ASSERT_TRUE(copy_start_result.has_value());
 
     // Unmount segment_1 to simulate source gone
@@ -3381,10 +3551,10 @@ TEST_F(MasterServiceSnapshotTest, CopyEnd) {
 
     // Test Case 6: CopyEnd, should fail because the source is gone, the object
     // should have only 1 replica from segment_2.
-    copy_end_result = service_->CopyEnd(client_id, key);
+    copy_end_result = service_->CopyEnd(client_id, key, "default");
     EXPECT_FALSE(copy_end_result.has_value());
     EXPECT_EQ(ErrorCode::REPLICA_IS_GONE, copy_end_result.error());
-    get_result = service_->GetReplicaList(key);
+    get_result = service_->GetReplicaList(key, "default");
     EXPECT_TRUE(get_result.has_value());
     auto& replicas = get_result.value().replicas;
     EXPECT_EQ(1, replicas.size());
@@ -3393,8 +3563,8 @@ TEST_F(MasterServiceSnapshotTest, CopyEnd) {
                                .buffer_descriptor.transport_endpoint_);
 
     // CopyStart the object from segment_2 to segment_3, then unmount segment_3
-    copy_start_result =
-        service_->CopyStart(client_id, key, "segment_2", {"segment_3"});
+    copy_start_result = service_->CopyStart(client_id, key, "default",
+                                            "segment_2", {"segment_3"});
     ASSERT_TRUE(copy_start_result.has_value());
 
     // Unmount segment_3 to simulate target gone
@@ -3404,10 +3574,10 @@ TEST_F(MasterServiceSnapshotTest, CopyEnd) {
 
     // Test Case 7: CopyEnd, should fail because the target is gone, the object
     // should have only 1 replica from segment_2.
-    copy_end_result = service_->CopyEnd(client_id, key);
+    copy_end_result = service_->CopyEnd(client_id, key, "default");
     EXPECT_FALSE(copy_end_result.has_value());
     EXPECT_EQ(ErrorCode::REPLICA_IS_GONE, copy_end_result.error());
-    get_result = service_->GetReplicaList(key);
+    get_result = service_->GetReplicaList(key, "default");
     EXPECT_TRUE(get_result.has_value());
     replicas = get_result.value().replicas;
     EXPECT_EQ(1, replicas.size());
@@ -3431,7 +3601,7 @@ TEST_F(MasterServiceSnapshotTest, CopyRevoke) {
 
     // Test Case 1: CopyRevoke a non-existent key, should fail.
     auto copy_revoke_result =
-        service_->CopyRevoke(client_id, "non_existent_key");
+        service_->CopyRevoke(client_id, "non_existent_key", "default");
     EXPECT_FALSE(copy_revoke_result.has_value());
     EXPECT_EQ(ErrorCode::OBJECT_NOT_FOUND, copy_revoke_result.error());
 
@@ -3443,47 +3613,49 @@ TEST_F(MasterServiceSnapshotTest, CopyRevoke) {
     config.preferred_segment = "segment_1";
 
     auto put_start_result =
-        service_->PutStart(client_id, key, slice_length, config);
+        service_->PutStart(client_id, key, "default", slice_length, config);
     ASSERT_TRUE(put_start_result.has_value());
-    auto put_end_result = service_->PutEnd(client_id, key, ReplicaType::MEMORY);
+    auto put_end_result =
+        service_->PutEnd(client_id, key, "default", ReplicaType::MEMORY);
     ASSERT_TRUE(put_end_result.has_value());
 
     // Test Case 2: CopyRevoke the object, should fail because there is no
     // ongoing copy task.
-    copy_revoke_result = service_->CopyRevoke(client_id, key);
+    copy_revoke_result = service_->CopyRevoke(client_id, key, "default");
     EXPECT_FALSE(copy_revoke_result.has_value());
     EXPECT_EQ(ErrorCode::OBJECT_NO_REPLICATION_TASK,
               copy_revoke_result.error());
 
     // CopyStart the object to segment_2
-    auto copy_start_result =
-        service_->CopyStart(client_id, key, "segment_1", {"segment_2"});
+    auto copy_start_result = service_->CopyStart(client_id, key, "default",
+                                                 "segment_1", {"segment_2"});
     ASSERT_TRUE(copy_start_result.has_value());
 
     // Test Case 3: CopyRevoke with an invalid client id, should fail.
-    copy_revoke_result = service_->CopyRevoke(invalid_client_id, key);
+    copy_revoke_result =
+        service_->CopyRevoke(invalid_client_id, key, "default");
     EXPECT_FALSE(copy_revoke_result.has_value());
     EXPECT_EQ(ErrorCode::ILLEGAL_CLIENT, copy_revoke_result.error());
 
     // Test Case 4: MoveRevoke the object, should fail because the ongoing task
     // is Copy.
-    auto move_revoke_result = service_->MoveRevoke(client_id, key);
+    auto move_revoke_result = service_->MoveRevoke(client_id, key, "default");
     EXPECT_FALSE(move_revoke_result.has_value());
     EXPECT_EQ(ErrorCode::INVALID_PARAMS, move_revoke_result.error());
 
     // Test Case 5: CopyRevoke, should success.
-    copy_revoke_result = service_->CopyRevoke(client_id, key);
+    copy_revoke_result = service_->CopyRevoke(client_id, key, "default");
     EXPECT_TRUE(copy_revoke_result.has_value());
 
     // Verify we still have 1 replica (the copy was revoked)
-    auto get_result = service_->GetReplicaList(key);
+    auto get_result = service_->GetReplicaList(key, "default");
     EXPECT_TRUE(get_result.has_value());
     EXPECT_EQ(1, get_result.value().replicas.size());
 
     // CopyStart the object from segment_1 to segment_2 again, then unmount
     // segment_1
-    copy_start_result =
-        service_->CopyStart(client_id, key, "segment_1", {"segment_2"});
+    copy_start_result = service_->CopyStart(client_id, key, "default",
+                                            "segment_1", {"segment_2"});
     ASSERT_TRUE(copy_start_result.has_value());
 
     // Unmount segment_1 to simulate source gone
@@ -3493,11 +3665,11 @@ TEST_F(MasterServiceSnapshotTest, CopyRevoke) {
 
     // Test Case 6: CopyRevoke, should success even though the source is gone,
     // the object should be erased too.
-    copy_revoke_result = service_->CopyRevoke(client_id, key);
+    copy_revoke_result = service_->CopyRevoke(client_id, key, "default");
     EXPECT_TRUE(copy_revoke_result.has_value());
 
     // Verify the object has been removed.
-    get_result = service_->GetReplicaList(key);
+    get_result = service_->GetReplicaList(key, "default");
     EXPECT_FALSE(get_result.has_value());
 }
 
@@ -3515,7 +3687,8 @@ TEST_F(MasterServiceSnapshotTest, MoveEnd) {
     UUID invalid_client_id = generate_uuid();
 
     // Test Case 1: MoveEnd a non-existent key, should fail.
-    auto move_end_result = service_->MoveEnd(client_id, "non_existent_key");
+    auto move_end_result =
+        service_->MoveEnd(client_id, "non_existent_key", "default");
     EXPECT_FALSE(move_end_result.has_value());
     EXPECT_EQ(ErrorCode::OBJECT_NOT_FOUND, move_end_result.error());
 
@@ -3527,46 +3700,47 @@ TEST_F(MasterServiceSnapshotTest, MoveEnd) {
     config.preferred_segment = "segment_1";
 
     auto put_start_result =
-        service_->PutStart(client_id, key, slice_length, config);
+        service_->PutStart(client_id, key, "default", slice_length, config);
     ASSERT_TRUE(put_start_result.has_value());
-    auto put_end_result = service_->PutEnd(client_id, key, ReplicaType::MEMORY);
+    auto put_end_result =
+        service_->PutEnd(client_id, key, "default", ReplicaType::MEMORY);
     ASSERT_TRUE(put_end_result.has_value());
 
     // Test Case 2: MoveEnd the object, should fail because there is no ongoing
     // move task.
-    move_end_result = service_->MoveEnd(client_id, key);
+    move_end_result = service_->MoveEnd(client_id, key, "default");
     EXPECT_FALSE(move_end_result.has_value());
     EXPECT_EQ(ErrorCode::OBJECT_NO_REPLICATION_TASK, move_end_result.error());
 
     // MoveStart the object to segment_2
-    auto move_start_result =
-        service_->MoveStart(client_id, key, "segment_1", "segment_2");
+    auto move_start_result = service_->MoveStart(client_id, key, "default",
+                                                 "segment_1", "segment_2");
     ASSERT_TRUE(move_start_result.has_value());
 
     // Test Case 3: MoveEnd with an invalid client id, should fail.
-    move_end_result = service_->MoveEnd(invalid_client_id, key);
+    move_end_result = service_->MoveEnd(invalid_client_id, key, "default");
     EXPECT_FALSE(move_end_result.has_value());
     EXPECT_EQ(ErrorCode::ILLEGAL_CLIENT, move_end_result.error());
 
     // Test Case 4: CopyEnd the object, should fail because the ongoing task is
     // Move.
-    auto copy_end_result = service_->CopyEnd(client_id, key);
+    auto copy_end_result = service_->CopyEnd(client_id, key, "default");
     EXPECT_FALSE(copy_end_result.has_value());
     EXPECT_EQ(ErrorCode::INVALID_PARAMS, copy_end_result.error());
 
     // Test Case 5: MoveEnd, should success.
-    move_end_result = service_->MoveEnd(client_id, key);
+    move_end_result = service_->MoveEnd(client_id, key, "default");
     EXPECT_TRUE(move_end_result.has_value());
 
     // Verify we still have 1 replica (the move was successful)
-    auto get_result = service_->GetReplicaList(key);
+    auto get_result = service_->GetReplicaList(key, "default");
     EXPECT_TRUE(get_result.has_value());
     EXPECT_EQ(1, get_result.value().replicas.size());
 
     // MoveStart the object from segment_2 to segment_1 again, then unmount
     // segment_2
-    move_start_result =
-        service_->MoveStart(client_id, key, "segment_2", "segment_1");
+    move_start_result = service_->MoveStart(client_id, key, "default",
+                                            "segment_2", "segment_1");
     ASSERT_TRUE(move_start_result.has_value());
 
     // Unmount segment_2 to simulate source gone
@@ -3575,7 +3749,7 @@ TEST_F(MasterServiceSnapshotTest, MoveEnd) {
     ASSERT_TRUE(unmount_result.has_value());
 
     // Test Case 6: MoveEnd, should fail because the source is gone.
-    move_end_result = service_->MoveEnd(client_id, key);
+    move_end_result = service_->MoveEnd(client_id, key, "default");
     EXPECT_FALSE(move_end_result.has_value());
     EXPECT_EQ(ErrorCode::REPLICA_IS_GONE, move_end_result.error());
 }
@@ -3594,7 +3768,7 @@ TEST_F(MasterServiceSnapshotTest, MoveRevoke) {
 
     // Test Case 1: MoveRevoke a non-existent key, should fail.
     auto move_revoke_result =
-        service_->MoveRevoke(client_id, "non_existent_key");
+        service_->MoveRevoke(client_id, "non_existent_key", "default");
     EXPECT_FALSE(move_revoke_result.has_value());
     EXPECT_EQ(ErrorCode::OBJECT_NOT_FOUND, move_revoke_result.error());
 
@@ -3606,40 +3780,42 @@ TEST_F(MasterServiceSnapshotTest, MoveRevoke) {
     config.preferred_segment = "segment_1";
 
     auto put_start_result =
-        service_->PutStart(client_id, key, slice_length, config);
+        service_->PutStart(client_id, key, "default", slice_length, config);
     ASSERT_TRUE(put_start_result.has_value());
-    auto put_end_result = service_->PutEnd(client_id, key, ReplicaType::MEMORY);
+    auto put_end_result =
+        service_->PutEnd(client_id, key, "default", ReplicaType::MEMORY);
     ASSERT_TRUE(put_end_result.has_value());
 
     // Test Case 2: MoveRevoke the object, should fail because there is no
     // ongoing move task.
-    move_revoke_result = service_->MoveRevoke(client_id, key);
+    move_revoke_result = service_->MoveRevoke(client_id, key, "default");
     EXPECT_FALSE(move_revoke_result.has_value());
     EXPECT_EQ(ErrorCode::OBJECT_NO_REPLICATION_TASK,
               move_revoke_result.error());
 
     // MoveStart the object from segment_1 to segment_2
-    auto move_start_result =
-        service_->MoveStart(client_id, key, "segment_1", "segment_2");
+    auto move_start_result = service_->MoveStart(client_id, key, "default",
+                                                 "segment_1", "segment_2");
     ASSERT_TRUE(move_start_result.has_value());
 
     // Test Case 3: MoveRevoke with an invalid client id, should fail.
-    move_revoke_result = service_->MoveRevoke(invalid_client_id, key);
+    move_revoke_result =
+        service_->MoveRevoke(invalid_client_id, key, "default");
     EXPECT_FALSE(move_revoke_result.has_value());
     EXPECT_EQ(ErrorCode::ILLEGAL_CLIENT, move_revoke_result.error());
 
     // Test Case 4: CopyRevoke the object, should fail because the ongoing task
     // is Move.
-    auto copy_revoke_result = service_->CopyRevoke(client_id, key);
+    auto copy_revoke_result = service_->CopyRevoke(client_id, key, "default");
     EXPECT_FALSE(copy_revoke_result.has_value());
     EXPECT_EQ(ErrorCode::INVALID_PARAMS, copy_revoke_result.error());
 
     // Test Case 5: MoveRevoke, should succeed.
-    move_revoke_result = service_->MoveRevoke(client_id, key);
+    move_revoke_result = service_->MoveRevoke(client_id, key, "default");
     EXPECT_TRUE(move_revoke_result.has_value());
 
     // Verify we still have 1 replica (the move was revoked)
-    auto get_result = service_->GetReplicaList(key);
+    auto get_result = service_->GetReplicaList(key, "default");
     EXPECT_TRUE(get_result.has_value());
     auto& replicas = get_result.value().replicas;
     EXPECT_EQ(1, replicas.size());
@@ -3649,8 +3825,8 @@ TEST_F(MasterServiceSnapshotTest, MoveRevoke) {
 
     // MoveStart the object from segment_1 to segment_2 again, then unmount
     // segment_1
-    move_start_result =
-        service_->MoveStart(client_id, key, "segment_1", "segment_2");
+    move_start_result = service_->MoveStart(client_id, key, "default",
+                                            "segment_1", "segment_2");
     ASSERT_TRUE(move_start_result.has_value());
 
     // Unmount segment_1 to simulate source gone
@@ -3659,11 +3835,11 @@ TEST_F(MasterServiceSnapshotTest, MoveRevoke) {
     ASSERT_TRUE(unmount_result.has_value());
 
     // Test Case 6: MoveRevoke, should succeed even though the source is gone.
-    move_revoke_result = service_->MoveRevoke(client_id, key);
+    move_revoke_result = service_->MoveRevoke(client_id, key, "default");
     EXPECT_TRUE(move_revoke_result.has_value());
 
     // The object should be erased as there is no replica left.
-    get_result = service_->GetReplicaList(key);
+    get_result = service_->GetReplicaList(key, "default");
     EXPECT_FALSE(get_result.has_value());
 }
 
@@ -3686,8 +3862,8 @@ TEST_F(MasterServiceSnapshotTest, MoveStart) {
     UUID client_id = generate_uuid();
 
     // Test Case 1: MoveStart a non-existent key, should fail.
-    auto move_start_result = service_->MoveStart(client_id, "non_existent_key",
-                                                 "segment_1", "segment_2");
+    auto move_start_result = service_->MoveStart(
+        client_id, "non_existent_key", "default", "segment_1", "segment_2");
     EXPECT_FALSE(move_start_result.has_value());
     EXPECT_EQ(ErrorCode::OBJECT_NOT_FOUND, move_start_result.error());
 
@@ -3699,36 +3875,37 @@ TEST_F(MasterServiceSnapshotTest, MoveStart) {
     config.preferred_segment = "segment_1";
 
     auto put_start_result =
-        service_->PutStart(client_id, key, slice_length, config);
+        service_->PutStart(client_id, key, "default", slice_length, config);
     ASSERT_TRUE(put_start_result.has_value());
 
     // Test Case 2: MoveStart the object, should fail because the only replica
     // is not completed.
-    move_start_result =
-        service_->MoveStart(client_id, key, "segment_1", "segment_2");
+    move_start_result = service_->MoveStart(client_id, key, "default",
+                                            "segment_1", "segment_2");
     EXPECT_FALSE(move_start_result.has_value());
     EXPECT_EQ(ErrorCode::REPLICA_NOT_FOUND, move_start_result.error());
 
     // PutEnd the object.
-    auto put_end_result = service_->PutEnd(client_id, key, ReplicaType::MEMORY);
+    auto put_end_result =
+        service_->PutEnd(client_id, key, "default", ReplicaType::MEMORY);
     ASSERT_TRUE(put_end_result.has_value());
 
     // Copy the object to segment_3.
-    auto copy_start_result =
-        service_->CopyStart(client_id, key, "segment_1", {"segment_3"});
+    auto copy_start_result = service_->CopyStart(client_id, key, "default",
+                                                 "segment_1", {"segment_3"});
     ASSERT_TRUE(copy_start_result.has_value());
-    auto copy_end_result = service_->CopyEnd(client_id, key);
+    auto copy_end_result = service_->CopyEnd(client_id, key, "default");
     ASSERT_TRUE(copy_end_result.has_value());
 
     // Test Case 3: MoveStart with source and target be the same, should fail.
-    move_start_result =
-        service_->MoveStart(client_id, key, "segment_1", "segment_1");
+    move_start_result = service_->MoveStart(client_id, key, "default",
+                                            "segment_1", "segment_1");
     EXPECT_FALSE(move_start_result.has_value());
     EXPECT_EQ(move_start_result.error(), ErrorCode::INVALID_PARAMS);
 
     // Test Case 4: MoveStart to segment_2, should succeed.
-    move_start_result =
-        service_->MoveStart(client_id, key, "segment_1", "segment_2");
+    move_start_result = service_->MoveStart(client_id, key, "default",
+                                            "segment_1", "segment_2");
     EXPECT_TRUE(move_start_result.has_value());
     auto move_response = move_start_result.value();
     EXPECT_EQ("segment_1", move_response.source.get_memory_descriptor()
@@ -3739,44 +3916,44 @@ TEST_F(MasterServiceSnapshotTest, MoveStart) {
                                .buffer_descriptor.transport_endpoint_);
 
     // Test Case 5: Try remove the object, should fail because it is moving.
-    auto remove_result = service_->Remove(key);
+    auto remove_result = service_->Remove(key, "default");
     EXPECT_FALSE(remove_result.has_value());
     EXPECT_EQ(ErrorCode::REPLICA_IS_NOT_READY, remove_result.error());
 
     // Test Case 6: MoveStart again, should fail because there is an ongoing
     // move task.
-    move_start_result =
-        service_->MoveStart(client_id, key, "segment_1", "segment_3");
+    move_start_result = service_->MoveStart(client_id, key, "default",
+                                            "segment_1", "segment_3");
     EXPECT_FALSE(move_start_result.has_value());
     EXPECT_EQ(ErrorCode::OBJECT_HAS_REPLICATION_TASK,
               move_start_result.error());
 
     // Test Case 7: MoveEnd, should succeed and the object now has 2 replicas
     // from segment_2 and segment_3
-    auto move_end_result = service_->MoveEnd(client_id, key);
+    auto move_end_result = service_->MoveEnd(client_id, key, "default");
     EXPECT_TRUE(move_end_result.has_value());
 
-    auto get_result = service_->GetReplicaList(key);
+    auto get_result = service_->GetReplicaList(key, "default");
     EXPECT_TRUE(get_result.has_value());
     auto& replicas = get_result.value().replicas;
     EXPECT_EQ(2, replicas.size());
 
     // Test Case 8: Move from a non-existent replica to segment_1, should fail.
     move_start_result = service_->MoveStart(
-        client_id, key, "non_existent_segment", "segment_1");
+        client_id, key, "default", "non_existent_segment", "segment_1");
     EXPECT_FALSE(move_start_result.has_value());
     EXPECT_EQ(ErrorCode::REPLICA_NOT_FOUND, move_start_result.error());
 
     // Test Case 8.5: Move to a non-existent target segment, should fail.
-    move_start_result = service_->MoveStart(client_id, key, "segment_2",
-                                            "non_existent_segment");
+    move_start_result = service_->MoveStart(
+        client_id, key, "default", "segment_2", "non_existent_segment");
     EXPECT_FALSE(move_start_result.has_value());
     EXPECT_EQ(ErrorCode::SEGMENT_NOT_FOUND, move_start_result.error());
 
     // Test Case 9: Move to an already existing segment, should succeed but
     // return nullopt.
-    move_start_result =
-        service_->MoveStart(client_id, key, "segment_2", "segment_3");
+    move_start_result = service_->MoveStart(client_id, key, "default",
+                                            "segment_2", "segment_3");
     EXPECT_TRUE(move_start_result.has_value());
     move_response = move_start_result.value();
     EXPECT_EQ("segment_2", move_response.source.get_memory_descriptor()
@@ -3785,16 +3962,16 @@ TEST_F(MasterServiceSnapshotTest, MoveStart) {
 
     // Test Case 10: Try remove the object, should fail because it is moving.
     std::this_thread::sleep_for(std::chrono::milliseconds(kv_lease_ttl * 2));
-    remove_result = service_->Remove(key);
+    remove_result = service_->Remove(key, "default");
     EXPECT_FALSE(remove_result.has_value());
     EXPECT_EQ(ErrorCode::OBJECT_HAS_REPLICATION_TASK, remove_result.error());
 
     // End the move.
-    move_end_result = service_->MoveEnd(client_id, key);
+    move_end_result = service_->MoveEnd(client_id, key, "default");
     EXPECT_TRUE(move_end_result.has_value());
 
     // Now the object should have only 1 replica on segment_3.
-    get_result = service_->GetReplicaList(key);
+    get_result = service_->GetReplicaList(key, "default");
     EXPECT_TRUE(get_result.has_value());
     replicas = get_result.value().replicas;
     EXPECT_EQ(1, replicas.size());
@@ -3804,7 +3981,7 @@ TEST_F(MasterServiceSnapshotTest, MoveStart) {
 
     // Test Case 11: Try remove the object, should succeed after lease expires.
     std::this_thread::sleep_for(std::chrono::milliseconds(kv_lease_ttl * 2));
-    remove_result = service_->Remove(key);
+    remove_result = service_->Remove(key, "default");
     EXPECT_TRUE(remove_result.has_value());
 }
 
@@ -3836,26 +4013,27 @@ TEST_F(MasterServiceSnapshotTest, ProtectCopyMoveSourceFromEviction) {
     config.preferred_segment = "segment_1";
 
     // Put two objects for move and copy tests.
-    auto put_start_result =
-        service_->PutStart(client_id, copy_key, slice_length, config);
+    auto put_start_result = service_->PutStart(client_id, copy_key, "default",
+                                               slice_length, config);
     ASSERT_TRUE(put_start_result.has_value());
     auto put_end_result =
-        service_->PutEnd(client_id, copy_key, ReplicaType::MEMORY);
+        service_->PutEnd(client_id, copy_key, "default", ReplicaType::MEMORY);
     ASSERT_TRUE(put_end_result.has_value());
 
-    put_start_result =
-        service_->PutStart(client_id, move_key, slice_length, config);
+    put_start_result = service_->PutStart(client_id, move_key, "default",
+                                          slice_length, config);
     ASSERT_TRUE(put_start_result.has_value());
-    put_end_result = service_->PutEnd(client_id, move_key, ReplicaType::MEMORY);
+    put_end_result =
+        service_->PutEnd(client_id, move_key, "default", ReplicaType::MEMORY);
     ASSERT_TRUE(put_end_result.has_value());
 
     // Start copy and move operations.
-    auto copy_start_result =
-        service_->CopyStart(client_id, copy_key, "segment_1", {"segment_2"});
+    auto copy_start_result = service_->CopyStart(client_id, copy_key, "default",
+                                                 "segment_1", {"segment_2"});
     ASSERT_TRUE(copy_start_result.has_value());
 
-    auto move_start_result =
-        service_->MoveStart(client_id, move_key, "segment_1", "segment_2");
+    auto move_start_result = service_->MoveStart(client_id, move_key, "default",
+                                                 "segment_1", "segment_2");
     ASSERT_TRUE(move_start_result.has_value());
 
     // Put more objects to trigger eviction. Do not prefer any segments.
@@ -3863,10 +4041,10 @@ TEST_F(MasterServiceSnapshotTest, ProtectCopyMoveSourceFromEviction) {
     for (size_t i = 0; i < 128 * (kSegmentSize * 2 / slice_length); ++i) {
         std::string key = "test_key_" + std::to_string(i);
         auto put_start_result =
-            service_->PutStart(client_id, key, slice_length, config);
+            service_->PutStart(client_id, key, "default", slice_length, config);
         if (put_start_result.has_value()) {
-            auto put_end_result =
-                service_->PutEnd(client_id, key, ReplicaType::MEMORY);
+            auto put_end_result = service_->PutEnd(client_id, key, "default",
+                                                   ReplicaType::MEMORY);
             ASSERT_TRUE(put_end_result.has_value());
         } else {
             // wait for eviction to work
@@ -3880,10 +4058,10 @@ TEST_F(MasterServiceSnapshotTest, ProtectCopyMoveSourceFromEviction) {
     ASSERT_TRUE(remove_all_result > 0);
 
     // Try end copy and move operations, should success.
-    auto copy_end_result = service_->CopyEnd(client_id, copy_key);
+    auto copy_end_result = service_->CopyEnd(client_id, copy_key, "default");
     EXPECT_TRUE(copy_end_result.has_value());
 
-    auto move_end_result = service_->MoveEnd(client_id, move_key);
+    auto move_end_result = service_->MoveEnd(client_id, move_key, "default");
     EXPECT_TRUE(move_end_result.has_value());
 }
 
diff --git a/mooncake-store/tests/ha/snapshot/master_service_test_for_snapshot_base.h b/mooncake-store/tests/ha/snapshot/master_service_test_for_snapshot_base.h
index d2edb675..5d48eb25 100644
--- a/mooncake-store/tests/ha/snapshot/master_service_test_for_snapshot_base.h
+++ b/mooncake-store/tests/ha/snapshot/master_service_test_for_snapshot_base.h
@@ -54,6 +54,7 @@ class MasterServiceSnapshotTestBase : public ::testing::Test {
         // eviction behavior inconsistency and snapshot comparison failures
         MasterMetricManager::instance().reset_allocated_mem_size();
         MasterMetricManager::instance().reset_total_mem_capacity();
+        MasterMetricManager::instance().reset_cache_total_nums();
 
         // Create a unique temporary directory for this test
         namespace fs = std::filesystem;
@@ -72,8 +73,8 @@ class MasterServiceSnapshotTestBase : public ::testing::Test {
     // LocalDiskSegment state for comparison
     struct LocalDiskSegmentState {
         bool enable_offloading = false;
-        std::map<std::string, int64_t>
-            offloading_objects;  // key -> timestamp (sorted)
+        std::map<std::string, OffloadTaskItem>
+            offloading_objects;  // storage key -> task (sorted)
     };
 
     // Task state for comparison
@@ -165,6 +166,18 @@ class MasterServiceSnapshotTestBase : public ::testing::Test {
         return service->PersistState(snapshot_id);
     }
 
+    static void EnsureSnapshotStores(MasterService* service) {
+        if (!service->snapshot_object_store_) {
+            service->snapshot_object_store_ = SnapshotObjectStore::Create(
+                SnapshotObjectStoreType::LOCAL_FILE);
+        }
+        if (!service->snapshot_catalog_store_ &&
+            service->snapshot_object_store_) {
+            service->snapshot_catalog_store_ =
+                service->CreateSnapshotCatalogStore();
+        }
+    }
+
     // Generate unique snapshot ID (timestamp format)
     std::string GenerateSnapshotId() const {
         auto now = std::chrono::system_clock::now();
@@ -188,7 +201,10 @@ class MasterServiceSnapshotTestBase : public ::testing::Test {
 
     // Get msgpack snapshot directory path
     std::string GetSnapshotDir(const std::string& snapshot_id) const {
-        return tmp_dir() + "/mooncake_master_snapshot/" + snapshot_id + "/";
+        return tmp_dir() + "/" +
+               ha::snapshot_catalog_store_detail::BuildSnapshotRoot(
+                   DEFAULT_CLUSTER_ID) +
+               snapshot_id + "/";
     }
 
     // Get backup directory path
@@ -204,7 +220,7 @@ class MasterServiceSnapshotTestBase : public ::testing::Test {
         ServiceStateSnapshot state;
 
         // === Basic State ===
-        auto keys_result = service->GetAllKeys();
+        auto keys_result = service->GetAllKeys("default");
         if (keys_result.has_value()) {
             state.all_keys = std::move(keys_result.value());
             std::sort(state.all_keys.begin(), state.all_keys.end());
@@ -217,7 +233,7 @@ class MasterServiceSnapshotTestBase : public ::testing::Test {
         }
 
         for (const auto& key : state.all_keys) {
-            auto replica_result = service->GetReplicaList(key);
+            auto replica_result = service->GetReplicaList(key, "default");
             if (replica_result.has_value()) {
                 state.replica_lists[key] = std::move(replica_result.value());
             }
@@ -247,9 +263,8 @@ class MasterServiceSnapshotTestBase : public ::testing::Test {
                 seg_state.enable_offloading = segment->enable_offloading;
                 // Copy offloading_objects to sorted map
                 std::lock_guard<Mutex> lock(segment->offloading_mutex_);
-                for (const auto& [key, timestamp] :
-                     segment->offloading_objects) {
-                    seg_state.offloading_objects[key] = timestamp;
+                for (const auto& [key, task] : segment->offloading_objects) {
+                    seg_state.offloading_objects[key] = task;
                 }
                 state.local_disk_segments[client_id] = std::move(seg_state);
             }
@@ -530,8 +545,10 @@ class MasterServiceSnapshotTestBase : public ::testing::Test {
         const std::string key = "snapshot_putstart_consistency_key";
         const uint64_t slice_length = 1024;
 
-        auto before = original->PutStart(client_id, key, slice_length, config);
-        auto after = restored->PutStart(client_id, key, slice_length, config);
+        auto before =
+            original->PutStart(client_id, key, "default", slice_length, config);
+        auto after =
+            restored->PutStart(client_id, key, "default", slice_length, config);
 
         ASSERT_EQ(before.has_value(), after.has_value())
             << "PutStart has_value mismatch between original and restored "
@@ -788,19 +805,10 @@ class MasterServiceSnapshotTestBase : public ::testing::Test {
             << "Use 'service_.reset(new MasterService(...))' instead of "
                "'std::unique_ptr<MasterService> service_(...)'";
 
-        // Ensure snapshot_object_store_ is initialized for PersistState
         // Some test configs may not enable snapshot/restore, so the backend
         // is not created in the constructor. We create it here for TearDown
         // validation.
-        if (!service_->snapshot_object_store_) {
-            service_->snapshot_object_store_ = SnapshotObjectStore::Create(
-                SnapshotObjectStoreType::LOCAL_FILE);
-        }
-        if (!service_->snapshot_catalog_store_ &&
-            service_->snapshot_object_store_) {
-            service_->snapshot_catalog_store_ =
-                service_->CreateSnapshotCatalogStore();
-        }
+        EnsureSnapshotStores(service_.get());
 
         // Test snapshot and restore functionality for all test cases
         TestSnapshotAndRestore(service_);
diff --git a/mooncake-store/tests/ha/snapshot/snapshot_child_process_test.cpp b/mooncake-store/tests/ha/snapshot/snapshot_child_process_test.cpp
index 0f2a4493..8902329a 100644
--- a/mooncake-store/tests/ha/snapshot/snapshot_child_process_test.cpp
+++ b/mooncake-store/tests/ha/snapshot/snapshot_child_process_test.cpp
@@ -2,6 +2,7 @@
 #include "master_metric_manager.h"
 #include "ha/snapshot/catalog/snapshot_catalog_store.h"
 #include "ha/snapshot/object/snapshot_object_store.h"
+#include "ha/snapshot/snapshot_test_utils.h"
 #ifdef STORE_USE_ETCD
 #include "etcd_helper.h"
 #include "ha/oplog/etcd_oplog_store.h"
@@ -33,6 +34,13 @@ class SnapshotChildProcessTest : public ::testing::Test {
    protected:
     const std::string& tmp_dir() const { return tmp_dir_; }
 
+    // Returns the snapshot root path used by the default service (with default
+    // cluster_id).
+    std::string default_snapshot_root() const {
+        return ha::snapshot_catalog_store_detail::BuildSnapshotRoot(
+            DEFAULT_CLUSTER_ID);
+    }
+
     std::unique_ptr<MasterService> service_;
 
     static constexpr const char* kEnvSnapshotLocalPath =
@@ -158,7 +166,44 @@ class SnapshotChildProcessTest : public ::testing::Test {
         size_t shard_idx = svc->getShardIndex(key);
         auto& shard = svc->metadata_shards_[shard_idx];
         SharedMutexLocker lock(&shard.mutex, shared_lock_t{});
-        return shard.metadata.find(key) != shard.metadata.end();
+        auto tenant_it = shard.tenants.find("default");
+        return tenant_it != shard.tenants.end() &&
+               tenant_it->second.metadata.find(key) !=
+                   tenant_it->second.metadata.end();
+    }
+
+    uint32_t GetShardIndexForTest(const std::string& key) {
+        return static_cast<uint32_t>(service_->getShardIndex(key));
+    }
+
+    tl::expected<void, SerializationError> DeserializeMetadataForTest(
+        const std::vector<uint8_t>& data) {
+        MasterService::MetadataSerializer serializer(service_.get());
+        return serializer.Deserialize(data);
+    }
+
+    bool ObjectIsGroupedInMetadata(const std::string& key, size_t shard_idx) {
+        auto& shard = service_->metadata_shards_[shard_idx];
+        SharedMutexLocker lock(&shard.mutex, shared_lock_t{});
+        for (const auto& [tenant_id, tenant_state] : shard.tenants) {
+            auto it = tenant_state.metadata.find(key);
+            if (it != tenant_state.metadata.end()) {
+                return it->second.IsGrouped();
+            }
+        }
+        return false;
+    }
+
+    std::string FindGroupIdOnDifferentShard(MasterService* svc,
+                                            const std::string& key) {
+        const size_t key_shard = svc->getShardIndex(key);
+        for (int i = 0; i < 1024; ++i) {
+            std::string group_id = key + "_group_" + std::to_string(i);
+            if (svc->getShardIndex(group_id) != key_shard) {
+                return group_id;
+            }
+        }
+        return key + "_group";
     }
 
    private:
@@ -260,20 +305,22 @@ TEST_F(SnapshotChildProcessTest, CleanupOldSnapshot_KeepsRecentDeletesOld) {
         "20240101_000000_000", "20240102_000000_000", "20240103_000000_000",
         "20240104_000000_000", "20240105_000000_000"};
 
+    const std::string snapshot_root = default_snapshot_root();
     for (const auto& id : snapshot_ids) {
-        std::string key = "mooncake_master_snapshot/" + id + "/metadata";
+        std::string key = snapshot_root + id + "/metadata";
         backend->UploadString(key, "dummy_metadata");
-        std::string manifest_key =
-            "mooncake_master_snapshot/" + id + "/manifest.txt";
+        std::string manifest_key = snapshot_root + id + "/manifest.txt";
         backend->UploadString(manifest_key, "messagepack|1.0.0|" + id);
 
         auto descriptor =
-            ha::snapshot_catalog_store_detail::MakeSnapshotDescriptor(id);
+            ha::snapshot_catalog_store_detail::MakeSnapshotDescriptor(
+                snapshot_root, id);
         auto descriptor_payload =
             ha::snapshot_catalog_store_detail::SerializeSnapshotDescriptor(
                 descriptor);
         auto descriptor_key =
-            ha::snapshot_catalog_store_detail::BuildDescriptorKey(id);
+            ha::snapshot_catalog_store_detail::BuildDescriptorKey(snapshot_root,
+                                                                  id);
         backend->UploadString(descriptor_key, descriptor_payload);
     }
 
@@ -282,7 +329,7 @@ TEST_F(SnapshotChildProcessTest, CleanupOldSnapshot_KeepsRecentDeletesOld) {
 
     // Verify: list remaining objects
     std::vector<std::string> remaining;
-    backend->ListObjectsWithPrefix("mooncake_master_snapshot/", remaining);
+    backend->ListObjectsWithPrefix(snapshot_root, remaining);
 
     // Only the 2 newest (20240104, 20240105) should remain
     for (const auto& key : remaining) {
@@ -313,6 +360,33 @@ TEST_F(SnapshotChildProcessTest, UploadSnapshotPayloadFile_Success) {
 
 // ========== Auto Snapshot Thread ==========
 
+TEST_F(SnapshotChildProcessTest, DestructorInterruptsSnapshotThreadSleep) {
+    auto config = MasterServiceConfigBuilder()
+                      .set_enable_snapshot(true)
+                      .set_enable_snapshot_restore(false)
+                      .set_snapshot_backup_dir(tmp_dir() + "/backup")
+                      .set_snapshot_interval_seconds(30)
+                      .set_snapshot_child_timeout_seconds(60)
+                      .set_snapshot_retention_count(3)
+                      .set_snapshot_object_store_type("local")
+                      .build();
+    auto auto_service = std::make_unique<MasterService>(config);
+
+    // Let the background thread enter its snapshot interval wait.
+    std::this_thread::sleep_for(std::chrono::milliseconds(100));
+
+    const auto destroy_start = std::chrono::steady_clock::now();
+    auto_service.reset();
+    const auto elapsed_ms =
+        std::chrono::duration_cast<std::chrono::milliseconds>(
+            std::chrono::steady_clock::now() - destroy_start)
+            .count();
+
+    EXPECT_LT(elapsed_ms, 5000)
+        << "MasterService shutdown should not wait for the full snapshot "
+           "interval";
+}
+
 TEST_F(SnapshotChildProcessTest, AutoSnapshot_GeneratesFiles) {
     // Create a service with snapshot enabled and short interval
     auto config = MasterServiceConfigBuilder()
@@ -332,7 +406,7 @@ TEST_F(SnapshotChildProcessTest, AutoSnapshot_GeneratesFiles) {
 
     // Check if latest.txt was generated
     std::string latest_path =
-        tmp_dir() + "/mooncake_master_snapshot/latest.txt";
+        tmp_dir() + "/" + default_snapshot_root() + "latest.txt";
     bool found = fs::exists(latest_path);
 
     // Destroy service to stop snapshot thread before assertions
@@ -365,9 +439,9 @@ TEST_F(SnapshotChildProcessTest, PersistState_PublishesSnapshotDescriptor) {
 
     EXPECT_EQ(latest->value().snapshot_id, snapshot_id);
     EXPECT_EQ(latest->value().manifest_key,
-              "mooncake_master_snapshot/" + snapshot_id + "/manifest.txt");
+              default_snapshot_root() + snapshot_id + "/manifest.txt");
     EXPECT_EQ(latest->value().object_prefix,
-              "mooncake_master_snapshot/" + snapshot_id + "/");
+              default_snapshot_root() + snapshot_id + "/");
     EXPECT_EQ(latest->value().last_included_seq, 0u);
     EXPECT_EQ(latest->value().producer_view_version, kViewVersion);
     const auto created_at_ms = latest->value().created_at_ms;
@@ -379,13 +453,11 @@ TEST_F(SnapshotChildProcessTest, PersistState_UsesFrozenSnapshotDescriptor) {
     const std::string snapshot_id = "20240601_120000_124";
     CreateDefaultService();
 
-    auto descriptor =
-        ha::snapshot_catalog_store_detail::MakeSnapshotDescriptor(snapshot_id);
+    const std::string snapshot_root = default_snapshot_root();
+    auto descriptor = ha::snapshot_catalog_store_detail::MakeSnapshotDescriptor(
+        snapshot_root, snapshot_id);
     descriptor.last_included_seq = 123;
     descriptor.producer_view_version = 41;
-    descriptor.manifest_key =
-        "mooncake_master_snapshot/" + snapshot_id + "/manifest.txt";
-    descriptor.object_prefix = "mooncake_master_snapshot/" + snapshot_id + "/";
     descriptor.created_at_ms = 1717243200123;
 
     auto persist_result = CallPersistState(descriptor);
@@ -407,6 +479,110 @@ TEST_F(SnapshotChildProcessTest, PersistState_UsesFrozenSnapshotDescriptor) {
     EXPECT_EQ(latest->value().created_at_ms, descriptor.created_at_ms);
 }
 
+TEST_F(SnapshotChildProcessTest, RestoreRebuildsGroupedObjectRouting) {
+    auto make_config = [this]() {
+        return MasterServiceConfigBuilder()
+            .set_enable_snapshot(false)
+            .set_enable_snapshot_restore(true)
+            .set_snapshot_backup_dir(tmp_dir() + "/backup")
+            .set_snapshot_interval_seconds(100)
+            .set_snapshot_child_timeout_seconds(60)
+            .set_snapshot_retention_count(3)
+            .set_snapshot_object_store_type("local")
+            .set_default_kv_lease_ttl(600000)
+            .build();
+    };
+    service_ = std::make_unique<MasterService>(make_config());
+
+    Segment segment;
+    segment.id = generate_uuid();
+    segment.name = "grouped_snapshot_segment";
+    segment.base = 0x310000000;
+    segment.size = 1024 * 1024 * 16;
+    segment.te_endpoint = segment.name;
+    const UUID client_id = generate_uuid();
+    ASSERT_TRUE(service_->MountSegment(segment, client_id).has_value());
+
+    const std::string key = "snapshot_grouped_route_key";
+    ReplicateConfig replicate_config;
+    replicate_config.replica_num = 1;
+    replicate_config.group_ids = std::vector<std::string>{
+        FindGroupIdOnDifferentShard(service_.get(), key)};
+
+    auto put_start =
+        service_->PutStart(client_id, key, "default", 1024, replicate_config);
+    ASSERT_TRUE(put_start.has_value()) << toString(put_start.error());
+    ASSERT_TRUE(service_->PutEnd(client_id, key, "default", ReplicaType::MEMORY)
+                    .has_value());
+    ASSERT_TRUE(service_->ExistKey(key, "default").value_or(false));
+
+    auto persist_result = CallPersistState("20240701_130000_000");
+    ASSERT_TRUE(persist_result.has_value())
+        << "PersistState failed: " << persist_result.error().message;
+
+    service_.reset();
+    service_ = std::make_unique<MasterService>(make_config());
+
+    auto restored_replicas = service_->GetReplicaList(key, "default");
+    ASSERT_TRUE(restored_replicas.has_value())
+        << "Grouped key should remain reachable by key after restore";
+    ASSERT_TRUE(service_->Remove(key, "default", /*force=*/true).has_value());
+    EXPECT_FALSE(service_->ExistKey(key, "default").value_or(true));
+}
+
+TEST_F(SnapshotChildProcessTest,
+       DeserializeLegacyMetadataWithoutGroupIdRestoresUngroupedObject) {
+    CreateDefaultService();
+    const std::string key = "legacy_snapshot_no_group_id_key";
+    const uint32_t shard_idx = GetShardIndexForTest(key);
+    const UUID client_id = generate_uuid();
+
+    msgpack::sbuffer shard_buffer;
+    MsgpackPacker shard_packer(&shard_buffer);
+    shard_packer.pack_map(1);
+    shard_packer.pack(std::string("metadata"));
+    shard_packer.pack_array(1);
+    shard_packer.pack_array(2);
+    shard_packer.pack(key);
+
+    shard_packer.pack_array(8);
+    shard_packer.pack(UuidToString(client_id));
+    shard_packer.pack(kDefaultTestPutStartTimeMs);
+    shard_packer.pack(kDefaultTestObjectSize);
+    shard_packer.pack(kDefaultTestLeaseTimeoutMs);
+    shard_packer.pack(false);
+    shard_packer.pack(uint64_t{0});
+    shard_packer.pack(uint32_t{1});
+    PackDiskReplica(shard_packer, kDefaultTestDiskFilePath,
+                    kDefaultTestObjectSize);
+
+    auto compressed_shard =
+        zstd_compress(reinterpret_cast<const uint8_t*>(shard_buffer.data()),
+                      shard_buffer.size(), 3);
+
+    msgpack::sbuffer root_buffer;
+    MsgpackPacker root_packer(&root_buffer);
+    root_packer.pack_map(3);
+    root_packer.pack(std::string("shards"));
+    root_packer.pack_map(1);
+    root_packer.pack(shard_idx);
+    root_packer.pack_bin(compressed_shard.size());
+    root_packer.pack_bin_body(
+        reinterpret_cast<const char*>(compressed_shard.data()),
+        compressed_shard.size());
+    root_packer.pack(std::string("discarded_replicas"));
+    root_packer.pack_array(0);
+    root_packer.pack(std::string("replica_next_id"));
+    root_packer.pack(uint64_t{10});
+
+    auto deserialize_result =
+        DeserializeMetadataForTest(ToByteVector(root_buffer));
+    ASSERT_TRUE(deserialize_result.has_value())
+        << deserialize_result.error().message;
+
+    EXPECT_FALSE(ObjectIsGroupedInMetadata(key, shard_idx));
+}
+
 TEST_F(SnapshotChildProcessTest, LegacyEtcdConnstringFallbackIsPreserved) {
     MasterConfig legacy_config;
     legacy_config.enable_ha = true;
@@ -563,12 +739,14 @@ TEST_F(SnapshotChildProcessTest,
         << "MountSegment failed";
 
     const std::string key1 = "restore_fallback_key_1";
-    auto put1 = service_->PutStart(client_id, key1, {1024}, {.replica_num = 1});
+    auto put1 = service_->PutStart(client_id, key1, "default", {1024},
+                                   {.replica_num = 1});
     ASSERT_TRUE(put1.has_value()) << "PutStart for key1 failed";
     ASSERT_TRUE(
-        service_->PutEnd(client_id, key1, ReplicaType::MEMORY).has_value())
+        service_->PutEnd(client_id, key1, "default", ReplicaType::MEMORY)
+            .has_value())
         << "PutEnd for key1 failed";
-    EXPECT_TRUE(service_->ExistKey(key1).value_or(false))
+    EXPECT_TRUE(service_->ExistKey(key1, "default").value_or(false))
         << "ExistKey should refresh lease for key1 before snapshot1";
 
     const std::string snapshot_id1 = "20240702_120000_000";
@@ -578,12 +756,14 @@ TEST_F(SnapshotChildProcessTest,
         << persist_result.error().message;
 
     const std::string key2 = "restore_fallback_key_2";
-    auto put2 = service_->PutStart(client_id, key2, {1024}, {.replica_num = 1});
+    auto put2 = service_->PutStart(client_id, key2, "default", {1024},
+                                   {.replica_num = 1});
     ASSERT_TRUE(put2.has_value()) << "PutStart for key2 failed";
     ASSERT_TRUE(
-        service_->PutEnd(client_id, key2, ReplicaType::MEMORY).has_value())
+        service_->PutEnd(client_id, key2, "default", ReplicaType::MEMORY)
+            .has_value())
         << "PutEnd for key2 failed";
-    EXPECT_TRUE(service_->ExistKey(key2).value_or(false))
+    EXPECT_TRUE(service_->ExistKey(key2, "default").value_or(false))
         << "ExistKey should refresh lease for key2 before snapshot2";
 
     const std::string snapshot_id2 = "20240702_120500_000";
@@ -593,8 +773,8 @@ TEST_F(SnapshotChildProcessTest,
         << persist_result.error().message;
 
     const fs::path corrupted_metadata = fs::path(tmp_dir()) /
-                                        "mooncake_master_snapshot" /
-                                        snapshot_id2 / "metadata";
+                                        default_snapshot_root() / snapshot_id2 /
+                                        "metadata";
     std::ofstream corrupt_stream(corrupted_metadata, std::ios::binary);
     ASSERT_TRUE(corrupt_stream.is_open())
         << "Failed to corrupt latest snapshot metadata";
@@ -614,9 +794,9 @@ TEST_F(SnapshotChildProcessTest,
                               .build();
     auto restored_service = std::make_unique<MasterService>(restore_config);
 
-    EXPECT_TRUE(restored_service->ExistKey(key1).value_or(false))
+    EXPECT_TRUE(restored_service->ExistKey(key1, "default").value_or(false))
         << "Restore should fall back to the previous healthy snapshot";
-    EXPECT_FALSE(restored_service->ExistKey(key2).value_or(false))
+    EXPECT_FALSE(restored_service->ExistKey(key2, "default").value_or(false))
         << "Corrupted latest snapshot must not be partially restored";
 
     restored_service.reset();
@@ -687,22 +867,22 @@ TEST_F(SnapshotChildProcessTest, RestoreCleansNonCompleteReplica) {
 
     // Add a complete object (clean data)
     std::string clean_key = "clean_object";
-    auto put_result =
-        service_->PutStart(client_id, clean_key, {1024}, {.replica_num = 1});
+    auto put_result = service_->PutStart(client_id, clean_key, "default",
+                                         {1024}, {.replica_num = 1});
     ASSERT_TRUE(put_result.has_value()) << "PutStart clean failed";
     auto put_end_result =
-        service_->PutEnd(client_id, clean_key, ReplicaType::MEMORY);
+        service_->PutEnd(client_id, clean_key, "default", ReplicaType::MEMORY);
     ASSERT_TRUE(put_end_result.has_value()) << "PutEnd clean failed";
 
     // Add an incomplete object (PutStart without PutEnd -> non-COMPLETE)
     std::string dirty_key = "dirty_incomplete";
-    auto put_dirty_result =
-        service_->PutStart(client_id, dirty_key, {1024}, {.replica_num = 1});
+    auto put_dirty_result = service_->PutStart(client_id, dirty_key, "default",
+                                               {1024}, {.replica_num = 1});
     ASSERT_TRUE(put_dirty_result.has_value()) << "PutStart dirty failed";
     // Intentionally NO PutEnd -> replica stays in PENDING status
 
     // Verify both keys exist in metadata before snapshot
-    EXPECT_TRUE(service_->ExistKey(clean_key).value_or(false));
+    EXPECT_TRUE(service_->ExistKey(clean_key, "default").value_or(false));
     EXPECT_TRUE(KeyExistsInMetadata(service_.get(), dirty_key))
         << "Dirty key should exist in raw metadata after PutStart";
 
@@ -725,7 +905,8 @@ TEST_F(SnapshotChildProcessTest, RestoreCleansNonCompleteReplica) {
     auto restored_service = std::make_unique<MasterService>(restore_config);
 
     // Step 4: Verify non-COMPLETE replica was cleaned, complete one remains
-    EXPECT_TRUE(restored_service->ExistKey(clean_key).value_or(false))
+    EXPECT_TRUE(
+        restored_service->ExistKey(clean_key, "default").value_or(false))
         << "Complete object should survive restore";
     EXPECT_FALSE(KeyExistsInMetadata(restored_service.get(), dirty_key))
         << "Non-COMPLETE object should be cleaned from metadata during restore";
@@ -761,23 +942,25 @@ TEST_F(SnapshotChildProcessTest, RestoreCleansExpiredLease) {
     // Add two complete objects via PutStart + PutEnd
     // Note: PutEnd calls GrantLease(0, ...) so lease is immediately expired
     std::string expired_key = "expired_lease_object";
-    auto put_exp =
-        service_->PutStart(client_id, expired_key, {1024}, {.replica_num = 1});
+    auto put_exp = service_->PutStart(client_id, expired_key, "default", {1024},
+                                      {.replica_num = 1});
     ASSERT_TRUE(put_exp.has_value()) << "PutStart expired failed";
-    ASSERT_TRUE(service_->PutEnd(client_id, expired_key, ReplicaType::MEMORY)
-                    .has_value())
+    ASSERT_TRUE(
+        service_->PutEnd(client_id, expired_key, "default", ReplicaType::MEMORY)
+            .has_value())
         << "PutEnd expired failed";
 
     std::string normal_key = "normal_lease_object";
-    auto put_norm =
-        service_->PutStart(client_id, normal_key, {1024}, {.replica_num = 1});
+    auto put_norm = service_->PutStart(client_id, normal_key, "default", {1024},
+                                       {.replica_num = 1});
     ASSERT_TRUE(put_norm.has_value()) << "PutStart normal failed";
-    ASSERT_TRUE(service_->PutEnd(client_id, normal_key, ReplicaType::MEMORY)
-                    .has_value())
+    ASSERT_TRUE(
+        service_->PutEnd(client_id, normal_key, "default", ReplicaType::MEMORY)
+            .has_value())
         << "PutEnd normal failed";
 
     // ExistKey grants a fresh lease (now + 600s) to normal_key
-    EXPECT_TRUE(service_->ExistKey(normal_key).value_or(false));
+    EXPECT_TRUE(service_->ExistKey(normal_key, "default").value_or(false));
     // Do NOT call ExistKey on expired_key, its lease stays expired from PutEnd
 
     // Step 2: Persist state
@@ -799,7 +982,8 @@ TEST_F(SnapshotChildProcessTest, RestoreCleansExpiredLease) {
     auto restored_service = std::make_unique<MasterService>(restore_config);
 
     // Step 4: Verify normal data retained, expired-lease data cleaned
-    EXPECT_TRUE(restored_service->ExistKey(normal_key).value_or(false))
+    EXPECT_TRUE(
+        restored_service->ExistKey(normal_key, "default").value_or(false))
         << "Normal object with valid lease should survive restore";
     EXPECT_FALSE(KeyExistsInMetadata(restored_service.get(), expired_key))
         << "Lease-expired object should be cleaned during restore";
diff --git a/mooncake-store/tests/ha/snapshot/snapshot_test_utils.h b/mooncake-store/tests/ha/snapshot/snapshot_test_utils.h
index 8ff7af27..a160c143 100644
--- a/mooncake-store/tests/ha/snapshot/snapshot_test_utils.h
+++ b/mooncake-store/tests/ha/snapshot/snapshot_test_utils.h
@@ -44,6 +44,24 @@ struct CatalogBackendParam {
     bool requires_redis{false};
 };
 
+// On-wire shapes the master snapshot writer has emitted over time. The standby
+// restore reader must tolerate all of them. See
+// MasterService::MetadataSerializer in master_service.cpp.
+//   kLegacy:         7 + replica_count, no data_type, no trailing hard_pinned
+//   kDataTypeOnly:   8 + replica_count, data_type after replica_count
+//   kHardPinnedOnly: 8 + replica_count, trailing hard_pinned, no data_type
+//   kDataTypeAndHardPinned:
+//                    9 + replica_count, data_type plus trailing hard_pinned
+//   kWithGroupId:    10 + replica_count, data_type + hard_pinned + group_id
+//                    (the current writer format)
+enum class SnapshotMetadataFormat {
+    kLegacy,
+    kDataTypeOnly,
+    kHardPinnedOnly,
+    kDataTypeAndHardPinned,
+    kWithGroupId,
+};
+
 class ScopedEnvVar {
    public:
     ScopedEnvVar(std::string name, std::string value) : name_(std::move(name)) {
@@ -128,30 +146,10 @@ inline std::vector<uint8_t> BuildSegmentsPayload() {
     return serialized.value();
 }
 
-inline std::vector<uint8_t> BuildMetadataPayload(
-    const UUID& client_id, std::string_view object_key = kDefaultTestObjectKey,
-    std::string_view disk_file_path = kDefaultTestDiskFilePath,
-    uint64_t object_size = kDefaultTestObjectSize,
-    uint64_t put_start_time_ms = kDefaultTestPutStartTimeMs,
-    uint64_t lease_timeout_ms = kDefaultTestLeaseTimeoutMs) {
-    msgpack::sbuffer shard_buffer;
-    MsgpackPacker shard_packer(&shard_buffer);
-    shard_packer.pack_map(1);
-    shard_packer.pack(std::string("metadata"));
-    shard_packer.pack_array(1);
-    shard_packer.pack_array(2);
-    shard_packer.pack(std::string(object_key));
-
-    shard_packer.pack_array(8);
-    shard_packer.pack(UuidToString(client_id));
-    shard_packer.pack(put_start_time_ms);
-    shard_packer.pack(object_size);
-    shard_packer.pack(lease_timeout_ms);
-    shard_packer.pack(false);
-    shard_packer.pack(uint64_t{0});
-    shard_packer.pack(uint32_t{1});
-    PackDiskReplica(shard_packer, disk_file_path, object_size);
-
+// Compresses a packed shard buffer and wraps it in the {"shards": {"0": ...}}
+// root map that the snapshot metadata payload expects.
+inline std::vector<uint8_t> WrapShardIntoMetadataRoot(
+    const msgpack::sbuffer& shard_buffer) {
     auto compressed_shard =
         zstd_compress(reinterpret_cast<const uint8_t*>(shard_buffer.data()),
                       shard_buffer.size(), 3);
@@ -169,12 +167,95 @@ inline std::vector<uint8_t> BuildMetadataPayload(
     return ToByteVector(root_buffer);
 }
 
+inline std::vector<uint8_t> BuildMetadataPayload(
+    const UUID& client_id, std::string_view object_key = kDefaultTestObjectKey,
+    std::string_view disk_file_path = kDefaultTestDiskFilePath,
+    uint64_t object_size = kDefaultTestObjectSize,
+    uint64_t put_start_time_ms = kDefaultTestPutStartTimeMs,
+    uint64_t lease_timeout_ms = kDefaultTestLeaseTimeoutMs,
+    SnapshotMetadataFormat format = SnapshotMetadataFormat::kLegacy) {
+    const bool include_data_type =
+        format == SnapshotMetadataFormat::kDataTypeOnly ||
+        format == SnapshotMetadataFormat::kDataTypeAndHardPinned ||
+        format == SnapshotMetadataFormat::kWithGroupId;
+    const bool include_hard_pinned =
+        format == SnapshotMetadataFormat::kHardPinnedOnly ||
+        format == SnapshotMetadataFormat::kDataTypeAndHardPinned ||
+        format == SnapshotMetadataFormat::kWithGroupId;
+    const bool include_group_id =
+        format == SnapshotMetadataFormat::kWithGroupId;
+    constexpr uint32_t kReplicaCount = 1;
+    // 7 leading fields + replicas + optional data_type/hard_pinned/group_id.
+    const size_t array_size = 7 + kReplicaCount + (include_data_type ? 1 : 0) +
+                              (include_hard_pinned ? 1 : 0) +
+                              (include_group_id ? 1 : 0);
+
+    msgpack::sbuffer shard_buffer;
+    MsgpackPacker shard_packer(&shard_buffer);
+    shard_packer.pack_map(1);
+    shard_packer.pack(std::string("metadata"));
+    shard_packer.pack_array(1);
+    shard_packer.pack_array(2);
+    shard_packer.pack(std::string(object_key));
+
+    shard_packer.pack_array(array_size);
+    shard_packer.pack(UuidToString(client_id));
+    shard_packer.pack(put_start_time_ms);
+    shard_packer.pack(object_size);
+    shard_packer.pack(lease_timeout_ms);
+    shard_packer.pack(false);
+    shard_packer.pack(uint64_t{0});
+    shard_packer.pack(kReplicaCount);
+    if (include_data_type) {
+        shard_packer.pack(static_cast<uint8_t>(ObjectDataType::TENSOR));
+    }
+    PackDiskReplica(shard_packer, disk_file_path, object_size);
+    if (include_hard_pinned) {
+        shard_packer.pack(true);
+    }
+    if (include_group_id) {
+        shard_packer.pack(std::string("test-group"));
+    }
+
+    return WrapShardIntoMetadataRoot(shard_buffer);
+}
+
+// Builds a metadata payload whose declared replica_count field is set to
+// `declared_replica_count` while no replicas are actually packed (the entry
+// array stays at the 7 leading fields). Used to verify the deserializer
+// rejects a hostile count instead of overflowing into an out-of-bounds read.
+inline std::vector<uint8_t> BuildMetadataPayloadWithDeclaredReplicaCount(
+    uint32_t declared_replica_count,
+    std::string_view object_key = kDefaultTestObjectKey) {
+    msgpack::sbuffer shard_buffer;
+    MsgpackPacker shard_packer(&shard_buffer);
+    shard_packer.pack_map(1);
+    shard_packer.pack(std::string("metadata"));
+    shard_packer.pack_array(1);
+    shard_packer.pack_array(2);
+    shard_packer.pack(std::string(object_key));
+
+    // Exactly the 7 leading fields, zero trailing replicas.
+    shard_packer.pack_array(7);
+    shard_packer.pack(UuidToString(UUID{1, 2}));
+    shard_packer.pack(kDefaultTestPutStartTimeMs);
+    shard_packer.pack(kDefaultTestObjectSize);
+    shard_packer.pack(kDefaultTestLeaseTimeoutMs);
+    shard_packer.pack(false);
+    shard_packer.pack(uint64_t{0});
+    shard_packer.pack(declared_replica_count);
+
+    return WrapShardIntoMetadataRoot(shard_buffer);
+}
+
 inline ha::SnapshotDescriptor MakeTestSnapshotDescriptor(
+    const std::string& cluster_id = "",
     std::string_view snapshot_id = kDefaultTestSnapshotId,
     uint64_t last_included_seq = kDefaultTestSnapshotSeq,
     uint64_t producer_view_version = kDefaultTestProducerViewVersion,
     int64_t created_at_ms = kDefaultTestCreatedAtMs) {
     auto descriptor = ha::snapshot_catalog_store_detail::MakeSnapshotDescriptor(
+        ha::snapshot_catalog_store_detail::BuildSnapshotRoot(cluster_id),
         std::string(snapshot_id));
     descriptor.last_included_seq = last_included_seq;
     descriptor.producer_view_version = producer_view_version;
@@ -201,7 +282,8 @@ inline std::unique_ptr<ha::SnapshotCatalogStore> CreateCatalogStoreForTest(
     const std::string& cluster_id, const std::string& redis_endpoint) {
     if (param.catalog_store_type == "embedded") {
         return std::make_unique<
-            ha::backends::embedded::EmbeddedSnapshotCatalogStore>(object_store);
+            ha::backends::embedded::EmbeddedSnapshotCatalogStore>(object_store,
+                                                                  cluster_id);
     }
 
 #ifdef STORE_USE_REDIS
@@ -216,13 +298,12 @@ inline std::unique_ptr<ha::SnapshotCatalogStore> CreateCatalogStoreForTest(
     return nullptr;
 }
 
-inline tl::expected<void, std::string> PublishSnapshotPayload(
+// Uploads the manifest/segments/metadata objects (using the caller-provided
+// metadata payload) and publishes the descriptor.
+inline tl::expected<void, std::string> PublishSnapshotPayloadBytes(
     SnapshotObjectStore& object_store, ha::SnapshotCatalogStore& catalog_store,
     const ha::SnapshotDescriptor& descriptor,
-    const UUID& client_id = UUID{1, 2},
-    std::string_view object_key = kDefaultTestObjectKey,
-    std::string_view disk_file_path = kDefaultTestDiskFilePath,
-    uint64_t object_size = kDefaultTestObjectSize) {
+    const std::vector<uint8_t>& metadata_payload) {
     auto manifest = object_store.UploadString(descriptor.manifest_key,
                                               "messagepack|1.0.0|standby-test");
     if (!manifest) {
@@ -236,9 +317,7 @@ inline tl::expected<void, std::string> PublishSnapshotPayload(
     }
 
     auto metadata = object_store.UploadBuffer(
-        descriptor.object_prefix + "metadata",
-        BuildMetadataPayload(client_id, object_key, disk_file_path,
-                             object_size));
+        descriptor.object_prefix + "metadata", metadata_payload);
     if (!metadata) {
         return tl::make_unexpected(metadata.error());
     }
@@ -251,4 +330,19 @@ inline tl::expected<void, std::string> PublishSnapshotPayload(
     return {};
 }
 
+inline tl::expected<void, std::string> PublishSnapshotPayload(
+    SnapshotObjectStore& object_store, ha::SnapshotCatalogStore& catalog_store,
+    const ha::SnapshotDescriptor& descriptor,
+    const UUID& client_id = UUID{1, 2},
+    std::string_view object_key = kDefaultTestObjectKey,
+    std::string_view disk_file_path = kDefaultTestDiskFilePath,
+    uint64_t object_size = kDefaultTestObjectSize,
+    SnapshotMetadataFormat format = SnapshotMetadataFormat::kLegacy) {
+    return PublishSnapshotPayloadBytes(
+        object_store, catalog_store, descriptor,
+        BuildMetadataPayload(client_id, object_key, disk_file_path, object_size,
+                             kDefaultTestPutStartTimeMs,
+                             kDefaultTestLeaseTimeoutMs, format));
+}
+
 }  // namespace mooncake::test
diff --git a/mooncake-store/tests/ha/standby/hot_standby_snapshot_bootstrap_test.cpp b/mooncake-store/tests/ha/standby/hot_standby_snapshot_bootstrap_test.cpp
index ab0bf4b3..091d8846 100644
--- a/mooncake-store/tests/ha/standby/hot_standby_snapshot_bootstrap_test.cpp
+++ b/mooncake-store/tests/ha/standby/hot_standby_snapshot_bootstrap_test.cpp
@@ -47,7 +47,7 @@ class HotStandbySnapshotBootstrapTest
             GetParam(), object_store_.get(), cluster_id_, FLAGS_redis_endpoint);
         ASSERT_NE(catalog_store_, nullptr);
 
-        descriptor_ = MakeTestSnapshotDescriptor();
+        descriptor_ = MakeTestSnapshotDescriptor(cluster_id_);
     }
 
     void TearDown() override {
diff --git a/mooncake-store/tests/http_metadata_server_test.cpp b/mooncake-store/tests/http_metadata_server_test.cpp
new file mode 100644
index 00000000..5ff426a5
--- /dev/null
+++ b/mooncake-store/tests/http_metadata_server_test.cpp
@@ -0,0 +1,99 @@
+#include <gtest/gtest.h>
+
+#include <chrono>
+#include <string>
+#include <thread>
+
+#include <async_simple/coro/SyncAwait.h>
+#include <ylt/coro_http/coro_http_client.hpp>
+
+#include "http_metadata_server.h"
+#include "utils.h"
+
+namespace mooncake::testing {
+
+class HttpMetadataServerTest : public ::testing::Test {
+   protected:
+    struct HttpResponse {
+        int status;
+        std::string body;
+    };
+
+    HttpResponse Get(int port, const std::string& path) {
+        coro_http::coro_http_client client;
+        auto response = async_simple::coro::syncAwait(client.async_get(
+            "http://127.0.0.1:" + std::to_string(port) + path));
+        return {response.status, std::string(response.resp_body)};
+    }
+
+    HttpResponse Put(int port, const std::string& path,
+                     const std::string& body) {
+        coro_http::coro_http_client client;
+        auto response = async_simple::coro::syncAwait(
+            client.async_put("http://127.0.0.1:" + std::to_string(port) + path,
+                             body, coro_http::req_content_type::json));
+        return {response.status, std::string(response.resp_body)};
+    }
+
+    void WaitUntilReady(int port) {
+        for (int i = 0; i < 50; ++i) {
+            if (Get(port, "/health").status == 200) {
+                return;
+            }
+            std::this_thread::sleep_for(std::chrono::milliseconds(20));
+        }
+        FAIL() << "HTTP metadata server did not become ready";
+    }
+};
+
+TEST_F(HttpMetadataServerTest, AllowsIdempotentRpcMetaRepublish) {
+    int port = getFreeTcpPort();
+    HttpMetadataServer server(static_cast<uint16_t>(port), "127.0.0.1");
+    ASSERT_TRUE(server.start());
+    WaitUntilReady(port);
+
+    const std::string path =
+        "/metadata?key=mooncake%2Frpc_meta%2F10.0.0.1%3A12384";
+    const std::string body =
+        R"({"ip_or_host_name":"10.0.0.1","rpc_port":15228})";
+
+    EXPECT_EQ(Put(port, path, body).status, 200);
+
+    auto second = Put(port, path, body);
+    EXPECT_EQ(second.status, 200);
+    EXPECT_EQ(second.body, "metadata unchanged");
+
+    auto stored = Get(port, path);
+    EXPECT_EQ(stored.status, 200);
+    EXPECT_EQ(stored.body, body);
+
+    server.stop();
+}
+
+TEST_F(HttpMetadataServerTest, RejectsChangedRpcMetaRepublish) {
+    int port = getFreeTcpPort();
+    HttpMetadataServer server(static_cast<uint16_t>(port), "127.0.0.1");
+    ASSERT_TRUE(server.start());
+    WaitUntilReady(port);
+
+    const std::string path =
+        "/metadata?key=mooncake%2Frpc_meta%2F10.0.0.1%3A12384";
+    const std::string original =
+        R"({"ip_or_host_name":"10.0.0.1","rpc_port":15228})";
+    const std::string changed =
+        R"({"ip_or_host_name":"10.0.0.1","rpc_port":16000})";
+
+    EXPECT_EQ(Put(port, path, original).status, 200);
+
+    auto second = Put(port, path, changed);
+    EXPECT_EQ(second.status, 400);
+    EXPECT_EQ(second.body, "Duplicate rpc_meta key not allowed");
+
+    auto stored = Get(port, path);
+    EXPECT_EQ(stored.status, 200);
+    EXPECT_EQ(stored.body, original);
+
+    server.stop();
+}
+
+}  // namespace mooncake::testing
diff --git a/mooncake-store/tests/master_admin_server_test.cpp b/mooncake-store/tests/master_admin_server_test.cpp
new file mode 100644
index 00000000..5ef0cd62
--- /dev/null
+++ b/mooncake-store/tests/master_admin_server_test.cpp
@@ -0,0 +1,1219 @@
+#include <glog/logging.h>
+#include <gtest/gtest.h>
+
+#include <chrono>
+#include <csignal>
+#include <filesystem>
+#include <fstream>
+#include <map>
+#include <memory>
+#include <string>
+#include <thread>
+
+#include <ylt/coro_http/coro_http_client.hpp>
+#include <ylt/struct_json/json_reader.h>
+#include <ylt/struct_json/json_writer.h>
+
+#include "ha/ha_types.h"
+#include "master_admin_service.h"
+#include "master_config.h"
+#include "rpc_service.h"
+#include "tenant_quota_policy_store.h"
+#include "types.h"
+#include "utils.h"
+
+#include <ylt/reflection/user_reflect_macro.hpp>
+
+namespace mooncake {
+namespace test {
+
+namespace {
+
+struct HttpCreateDrainJobResponse {
+    bool success{false};
+    std::string job_id;
+    std::string status;
+    int32_t error_code{0};
+    std::string error_message;
+};
+YLT_REFL(HttpCreateDrainJobResponse, success, job_id, status, error_code,
+         error_message);
+
+struct HttpQueryDrainJobResponse {
+    bool success{false};
+    std::string job_id;
+    int32_t type{0};
+    std::string type_name;
+    int32_t status{0};
+    std::string status_name;
+    std::string message;
+    int32_t error_code{0};
+    std::string error_message;
+};
+YLT_REFL(HttpQueryDrainJobResponse, success, job_id, type, type_name, status,
+         status_name, message, error_code, error_message);
+
+struct HttpCancelDrainJobResponse {
+    bool success{false};
+    std::string job_id;
+    std::string status;
+    int32_t error_code{0};
+    std::string error_message;
+};
+YLT_REFL(HttpCancelDrainJobResponse, success, job_id, status, error_code,
+         error_message);
+
+struct HttpSegmentStatusResponse {
+    bool success{false};
+    std::string segment;
+    int32_t status{0};
+    std::string status_name;
+    int32_t error_code{0};
+    std::string error_message;
+};
+YLT_REFL(HttpSegmentStatusResponse, success, segment, status, status_name,
+         error_code, error_message);
+
+struct HttpErrorResponse {
+    bool success{false};
+    int32_t error_code{0};
+    std::string error_message;
+};
+YLT_REFL(HttpErrorResponse, success, error_code, error_message);
+
+struct HttpSegmentsDetailResponse {
+    uint64_t total_segments{0};
+};
+YLT_REFL(HttpSegmentsDetailResponse, total_segments);
+
+std::string WriteTenantQuotaPolicyForTest(
+    const std::map<std::string, uint64_t>& tenant_quotas) {
+    TenantQuotaPolicySnapshot snapshot;
+    snapshot.tenant_quotas = tenant_quotas;
+    auto path = std::filesystem::temp_directory_path() /
+                ("mooncake_admin_tenant_quota_" +
+                 UuidToString(generate_uuid()) + ".yaml");
+    std::ofstream out(path);
+    out << FormatTenantQuotaPolicyYaml(snapshot);
+    return path.string();
+}
+
+}  // namespace
+
+// =========================================================================
+// MasterAdminServerTest — lightweight tests without a shared service.
+// Each test that needs a server creates and destroys it within the test.
+// =========================================================================
+
+class MasterAdminServerTest : public ::testing::Test {
+   protected:
+    struct HttpResponse {
+        int http_status;
+        std::string body;
+    };
+
+    void SetUp() override {}
+
+    void TearDown() override {}
+
+    static std::string BaseUrl(int port) {
+        return "http://127.0.0.1:" + std::to_string(port);
+    }
+
+    HttpResponse HttpGet(int port, const std::string& path) {
+        coro_http::coro_http_client client;
+        auto result = client.get(BaseUrl(port) + path);
+        return {result.status, std::string(result.resp_body)};
+    }
+
+    HttpResponse HttpPostJson(int port, const std::string& path,
+                              const std::string& body) {
+        coro_http::coro_http_client client;
+        auto result = client.post(BaseUrl(port) + path, body,
+                                  coro_http::req_content_type::json);
+        return {result.status, std::string(result.resp_body)};
+    }
+
+    HttpResponse HttpPutJson(int port, const std::string& path,
+                             const std::string& body) {
+        coro_http::coro_http_client client;
+        auto result = async_simple::coro::syncAwait(client.async_put(
+            BaseUrl(port) + path, body, coro_http::req_content_type::json));
+        return {result.status, std::string(result.resp_body)};
+    }
+
+    HttpResponse HttpDelete(int port, const std::string& path) {
+        coro_http::coro_http_client client;
+        auto result = async_simple::coro::syncAwait(client.async_delete(
+            BaseUrl(port) + path, "", coro_http::req_content_type::json));
+        return {result.status, std::string(result.resp_body)};
+    }
+};
+
+// =========================================================================
+// Always-available endpoint tests
+// =========================================================================
+
+TEST_F(MasterAdminServerTest, MetricsEndpointReturns200) {
+    int port = getFreeTcpPort();
+    MasterAdminServer admin(static_cast<uint16_t>(port), false);
+    ASSERT_TRUE(admin.Start());
+    admin.SetRuntimeState(ha::MasterRuntimeState::kServing);
+
+    auto resp = HttpGet(port, "/metrics");
+    EXPECT_EQ(resp.http_status, 200);
+    EXPECT_NE(resp.body.find("master_"), std::string::npos);
+
+    admin.Stop();
+}
+
+TEST_F(MasterAdminServerTest, MetricsSummaryEndpointReturns200) {
+    int port = getFreeTcpPort();
+    MasterAdminServer admin(static_cast<uint16_t>(port), false);
+    ASSERT_TRUE(admin.Start());
+    admin.SetRuntimeState(ha::MasterRuntimeState::kServing);
+
+    auto resp = HttpGet(port, "/metrics/summary");
+    EXPECT_EQ(resp.http_status, 200);
+    EXPECT_NE(resp.body.find("role="), std::string::npos);
+    EXPECT_NE(resp.body.find("state="), std::string::npos);
+
+    admin.Stop();
+}
+
+TEST_F(MasterAdminServerTest, HealthEndpointReturns200InStandby) {
+    int port = getFreeTcpPort();
+    MasterAdminServer admin(static_cast<uint16_t>(port), false);
+    ASSERT_TRUE(admin.Start());
+    admin.SetRuntimeState(ha::MasterRuntimeState::kStandby);
+
+    auto resp = HttpGet(port, "/health");
+    EXPECT_EQ(resp.http_status, 200);
+    EXPECT_NE(resp.body.find("\"status\":\"ok\""), std::string::npos);
+    EXPECT_NE(resp.body.find("\"role\":\"standby\""), std::string::npos);
+
+    admin.Stop();
+}
+
+TEST_F(MasterAdminServerTest, HealthEndpointReturns200InServing) {
+    int port = getFreeTcpPort();
+    MasterAdminServer admin(static_cast<uint16_t>(port), false);
+    ASSERT_TRUE(admin.Start());
+    admin.SetRuntimeState(ha::MasterRuntimeState::kServing);
+
+    auto resp = HttpGet(port, "/health");
+    EXPECT_EQ(resp.http_status, 200);
+    EXPECT_NE(resp.body.find("\"role\":\"leader\""), std::string::npos);
+
+    admin.Stop();
+}
+
+TEST_F(MasterAdminServerTest, HealthEndpointIncludesLeaderInfoWhenSet) {
+    int port = getFreeTcpPort();
+    MasterAdminServer admin(static_cast<uint16_t>(port), false);
+    ASSERT_TRUE(admin.Start());
+    admin.SetRuntimeState(ha::MasterRuntimeState::kStandby);
+    admin.SetObservedLeader(ha::MasterView{
+        .leader_address = "10.0.0.1:19000",
+        .view_version = 42,
+    });
+
+    auto resp = HttpGet(port, "/health");
+    EXPECT_EQ(resp.http_status, 200);
+    EXPECT_NE(resp.body.find("\"leader_address\":\"10.0.0.1:19000\""),
+              std::string::npos);
+    EXPECT_NE(resp.body.find("\"view_version\":42"), std::string::npos);
+
+    admin.Stop();
+}
+
+TEST_F(MasterAdminServerTest, RoleEndpointReturnsCorrectRoleForServing) {
+    int port = getFreeTcpPort();
+    MasterAdminServer admin(static_cast<uint16_t>(port), false);
+    ASSERT_TRUE(admin.Start());
+    admin.SetRuntimeState(ha::MasterRuntimeState::kServing);
+
+    auto resp = HttpGet(port, "/role");
+    EXPECT_EQ(resp.http_status, 200);
+    EXPECT_EQ(resp.body, "leader");
+
+    admin.Stop();
+}
+
+TEST_F(MasterAdminServerTest, RoleEndpointReturnsCorrectRoleForStandby) {
+    int port = getFreeTcpPort();
+    MasterAdminServer admin(static_cast<uint16_t>(port), false);
+    ASSERT_TRUE(admin.Start());
+    admin.SetRuntimeState(ha::MasterRuntimeState::kStandby);
+
+    auto resp = HttpGet(port, "/role");
+    EXPECT_EQ(resp.http_status, 200);
+    EXPECT_EQ(resp.body, "standby");
+
+    admin.Stop();
+}
+
+TEST_F(MasterAdminServerTest, RoleEndpointReturnsCorrectRoleForCandidate) {
+    int port = getFreeTcpPort();
+    MasterAdminServer admin(static_cast<uint16_t>(port), false);
+    ASSERT_TRUE(admin.Start());
+    admin.SetRuntimeState(ha::MasterRuntimeState::kCandidate);
+
+    auto resp = HttpGet(port, "/role");
+    EXPECT_EQ(resp.http_status, 200);
+    EXPECT_EQ(resp.body, "standby");
+
+    admin.Stop();
+}
+
+TEST_F(MasterAdminServerTest, RoleEndpointReturnsCorrectRoleForRecovering) {
+    int port = getFreeTcpPort();
+    MasterAdminServer admin(static_cast<uint16_t>(port), false);
+    ASSERT_TRUE(admin.Start());
+    admin.SetRuntimeState(ha::MasterRuntimeState::kRecovering);
+
+    auto resp = HttpGet(port, "/role");
+    EXPECT_EQ(resp.http_status, 200);
+    EXPECT_EQ(resp.body, "standby");
+
+    admin.Stop();
+}
+
+TEST_F(MasterAdminServerTest, HaStatusEndpointReturnsCorrectState) {
+    int port = getFreeTcpPort();
+    MasterAdminServer admin(static_cast<uint16_t>(port), false);
+    ASSERT_TRUE(admin.Start());
+
+    admin.SetRuntimeState(ha::MasterRuntimeState::kStandby);
+    auto resp = HttpGet(port, "/ha_status");
+    EXPECT_EQ(resp.http_status, 200);
+    EXPECT_EQ(resp.body, "standby");
+
+    admin.SetRuntimeState(ha::MasterRuntimeState::kServing);
+    resp = HttpGet(port, "/ha_status");
+    EXPECT_EQ(resp.http_status, 200);
+    EXPECT_EQ(resp.body, "serving");
+
+    admin.SetRuntimeState(ha::MasterRuntimeState::kRecovering);
+    resp = HttpGet(port, "/ha_status");
+    EXPECT_EQ(resp.http_status, 200);
+    EXPECT_EQ(resp.body, "recovering");
+
+    admin.SetRuntimeState(ha::MasterRuntimeState::kCatchingUp);
+    resp = HttpGet(port, "/ha_status");
+    EXPECT_EQ(resp.http_status, 200);
+    EXPECT_EQ(resp.body, "catching_up");
+
+    admin.Stop();
+}
+
+TEST_F(MasterAdminServerTest, LeaderEndpointWithoutLeader) {
+    int port = getFreeTcpPort();
+    MasterAdminServer admin(static_cast<uint16_t>(port), false);
+    ASSERT_TRUE(admin.Start());
+
+    auto resp = HttpGet(port, "/leader");
+    EXPECT_EQ(resp.http_status, 200);
+    EXPECT_NE(resp.body.find("\"present\":false"), std::string::npos);
+
+    admin.Stop();
+}
+
+TEST_F(MasterAdminServerTest, LeaderEndpointWithLeader) {
+    int port = getFreeTcpPort();
+    MasterAdminServer admin(static_cast<uint16_t>(port), false);
+    ASSERT_TRUE(admin.Start());
+    admin.SetObservedLeader(ha::MasterView{
+        .leader_address = "192.168.1.1:19000",
+        .view_version = 5,
+    });
+
+    auto resp = HttpGet(port, "/leader");
+    EXPECT_EQ(resp.http_status, 200);
+    EXPECT_NE(resp.body.find("\"present\":true"), std::string::npos);
+    EXPECT_NE(resp.body.find("192.168.1.1:19000"), std::string::npos);
+    EXPECT_NE(resp.body.find("\"view_version\":5"), std::string::npos);
+
+    admin.Stop();
+}
+
+TEST_F(MasterAdminServerTest, LeaderEndpointAfterClearingLeader) {
+    int port = getFreeTcpPort();
+    MasterAdminServer admin(static_cast<uint16_t>(port), false);
+    ASSERT_TRUE(admin.Start());
+    admin.SetObservedLeader(ha::MasterView{
+        .leader_address = "192.168.1.1:19000",
+        .view_version = 5,
+    });
+
+    auto resp = HttpGet(port, "/leader");
+    EXPECT_NE(resp.body.find("\"present\":true"), std::string::npos);
+
+    admin.SetObservedLeader(std::nullopt);
+    resp = HttpGet(port, "/leader");
+    EXPECT_NE(resp.body.find("\"present\":false"), std::string::npos);
+
+    admin.Stop();
+}
+
+TEST_F(MasterAdminServerTest, AllAlwaysAvailableEndpointsInStartingState) {
+    int port = getFreeTcpPort();
+    MasterAdminServer admin(static_cast<uint16_t>(port), false);
+    ASSERT_TRUE(admin.Start());
+    admin.SetRuntimeState(ha::MasterRuntimeState::kStarting);
+
+    auto metrics = HttpGet(port, "/metrics");
+    EXPECT_EQ(metrics.http_status, 200);
+
+    auto summary = HttpGet(port, "/metrics/summary");
+    EXPECT_EQ(summary.http_status, 200);
+
+    auto health = HttpGet(port, "/health");
+    EXPECT_EQ(health.http_status, 200);
+    EXPECT_NE(health.body.find("\"ha_state\":\"starting\""), std::string::npos);
+
+    auto role = HttpGet(port, "/role");
+    EXPECT_EQ(role.http_status, 200);
+    EXPECT_EQ(role.body, "standby");
+
+    auto ha_status = HttpGet(port, "/ha_status");
+    EXPECT_EQ(ha_status.http_status, 200);
+    EXPECT_EQ(ha_status.body, "starting");
+
+    auto leader = HttpGet(port, "/leader");
+    EXPECT_EQ(leader.http_status, 200);
+    EXPECT_NE(leader.body.find("\"present\":false"), std::string::npos);
+
+    admin.Stop();
+}
+
+TEST_F(MasterAdminServerTest, AlwaysAvailableEndpointsInLeaderWarmup) {
+    int port = getFreeTcpPort();
+    MasterAdminServer admin(static_cast<uint16_t>(port), false);
+    ASSERT_TRUE(admin.Start());
+    admin.SetRuntimeState(ha::MasterRuntimeState::kLeaderWarmup);
+
+    auto role = HttpGet(port, "/role");
+    EXPECT_EQ(role.body, "leader");
+
+    auto health = HttpGet(port, "/health");
+    EXPECT_NE(health.body.find("\"role\":\"leader\""), std::string::npos);
+    EXPECT_NE(health.body.find("\"ha_state\":\"leader_warmup\""),
+              std::string::npos);
+
+    admin.Stop();
+}
+
+TEST_F(MasterAdminServerTest, RoleWithAllStates) {
+    int port = getFreeTcpPort();
+    MasterAdminServer admin(static_cast<uint16_t>(port), false);
+    ASSERT_TRUE(admin.Start());
+
+    // MasterRuntimeRoleToString groups kStarting/kStandby/kCandidate/
+    // kRecovering/kCatchingUp as "standby"; only leader states are "leader".
+    struct {
+        ha::MasterRuntimeState state;
+        std::string expected_role;
+    } cases[] = {
+        {ha::MasterRuntimeState::kStarting, "standby"},
+        {ha::MasterRuntimeState::kStandby, "standby"},
+        {ha::MasterRuntimeState::kCandidate, "standby"},
+        {ha::MasterRuntimeState::kRecovering, "standby"},
+        {ha::MasterRuntimeState::kCatchingUp, "standby"},
+        {ha::MasterRuntimeState::kLeaderWarmup, "leader"},
+        {ha::MasterRuntimeState::kServing, "leader"},
+    };
+
+    for (const auto& tc : cases) {
+        admin.SetRuntimeState(tc.state);
+        auto resp = HttpGet(port, "/role");
+        EXPECT_EQ(resp.body, tc.expected_role)
+            << "Unexpected role for state "
+            << ha::MasterRuntimeStateToString(tc.state);
+    }
+
+    admin.Stop();
+}
+
+TEST_F(MasterAdminServerTest, ServiceEndpointsReturn503WhenServiceUnavailable) {
+    int port = getFreeTcpPort();
+    MasterAdminServer admin(static_cast<uint16_t>(port), false);
+    ASSERT_TRUE(admin.Start());
+    admin.SetRuntimeState(ha::MasterRuntimeState::kStandby);
+
+    const std::string unavailable_msg = "service plane is not active";
+
+    auto keys = HttpGet(port, "/get_all_keys");
+    EXPECT_EQ(keys.http_status, 503);
+    EXPECT_NE(keys.body.find(unavailable_msg), std::string::npos);
+
+    auto segments = HttpGet(port, "/get_all_segments");
+    EXPECT_EQ(segments.http_status, 503);
+    EXPECT_NE(segments.body.find(unavailable_msg), std::string::npos);
+
+    auto detail = HttpGet(port, "/get_segments_detail");
+    EXPECT_EQ(detail.http_status, 503);
+    EXPECT_NE(detail.body.find(unavailable_msg), std::string::npos);
+
+    auto query_seg = HttpGet(port, "/query_segment?segment=foo");
+    EXPECT_EQ(query_seg.http_status, 503);
+    EXPECT_NE(query_seg.body.find(unavailable_msg), std::string::npos);
+
+    auto query_key = HttpGet(port, "/query_key?key=foo");
+    EXPECT_EQ(query_key.http_status, 503);
+    EXPECT_NE(query_key.body.find(unavailable_msg), std::string::npos);
+
+    auto batch = HttpGet(port, "/batch_query_keys?keys=foo");
+    EXPECT_EQ(batch.http_status, 503);
+    EXPECT_NE(batch.body.find(unavailable_msg), std::string::npos);
+
+    auto seg_status = HttpGet(port, "/api/v1/segments/status?segment=foo");
+    EXPECT_EQ(seg_status.http_status, 503);
+    EXPECT_NE(seg_status.body.find(unavailable_msg), std::string::npos);
+
+    auto tenant_quotas = HttpGet(port, "/api/v1/tenant_quotas");
+    EXPECT_EQ(tenant_quotas.http_status, 503);
+    EXPECT_NE(tenant_quotas.body.find(unavailable_msg), std::string::npos);
+
+    auto drain_create = HttpPostJson(port, "/api/v1/drain_jobs", "{}");
+    EXPECT_EQ(drain_create.http_status, 503);
+
+    // drain_jobs/query and drain_jobs/cancel validate job_id before
+    // checking service availability, so we must use a valid UUID.
+    std::string valid_uuid = UuidToString(generate_uuid());
+    auto drain_query =
+        HttpGet(port, "/api/v1/drain_jobs/query?job_id=" + valid_uuid);
+    EXPECT_EQ(drain_query.http_status, 503);
+
+    auto drain_cancel = HttpPostJson(
+        port, "/api/v1/drain_jobs/cancel?job_id=" + valid_uuid, "");
+    EXPECT_EQ(drain_cancel.http_status, 503);
+
+    admin.Stop();
+}
+
+TEST_F(MasterAdminServerTest, TenantQuotaAdminLifecycleEndpoints) {
+    const std::string policy_path = WriteTenantQuotaPolicyForTest({});
+    WrappedMasterServiceConfig svc_config;
+    svc_config.default_kv_lease_ttl = 5000;
+    svc_config.enable_metric_reporting = false;
+    svc_config.enable_multi_tenants = true;
+    svc_config.tenant_quota_connector_type = "file";
+    svc_config.tenant_quota_connector_uri = policy_path;
+    auto service = std::make_shared<WrappedMasterService>(svc_config);
+
+    Segment segment;
+    segment.id = generate_uuid();
+    segment.name = "quota_admin_segment";
+    segment.base = 0x600000000;
+    segment.size = 2000;
+    UUID client_id = generate_uuid();
+    ASSERT_TRUE(service->MountSegment(segment, client_id).has_value());
+
+    int port = getFreeTcpPort();
+    MasterAdminServer admin(static_cast<uint16_t>(port), false);
+    ASSERT_TRUE(admin.Start());
+    admin.SetRuntimeState(ha::MasterRuntimeState::kServing);
+    admin.SetServiceDelegate(service);
+    admin.SetServiceAvailable(true);
+
+    auto upsert = HttpPutJson(port, "/api/v1/tenant_quotas?tenant_id=tenant-a",
+                              "{\"requested_quota_bytes\":800}");
+    EXPECT_EQ(upsert.http_status, 200);
+    EXPECT_NE(upsert.body.find("\"tenant_id\":\"tenant-a\""),
+              std::string::npos);
+    EXPECT_NE(upsert.body.find("\"requested_quota_bytes\":800"),
+              std::string::npos);
+    EXPECT_NE(upsert.body.find("\"effective_quota_bytes\":800"),
+              std::string::npos);
+    EXPECT_NE(upsert.body.find("\"has_explicit_policy\":true"),
+              std::string::npos);
+
+    auto list = HttpGet(port, "/api/v1/tenant_quotas");
+    EXPECT_EQ(list.http_status, 200);
+    EXPECT_NE(list.body.find("\"tenant_id\":\"tenant-a\""), std::string::npos);
+
+    auto one = HttpGet(port, "/api/v1/tenant_quotas?tenant_id=tenant-a");
+    EXPECT_EQ(one.http_status, 200);
+    EXPECT_NE(one.body.find("\"committed_count\":0"), std::string::npos);
+    EXPECT_NE(one.body.find("\"over_quota\":false"), std::string::npos);
+
+    ReplicateConfig cfg;
+    cfg.replica_num = 1;
+    auto put =
+        service->PutStart(client_id, "quota_admin_key", 100, cfg, "tenant-a");
+    ASSERT_TRUE(put.has_value()) << toString(put.error());
+    ASSERT_TRUE(service
+                    ->PutEnd(client_id, "quota_admin_key", ReplicaType::MEMORY,
+                             "tenant-a")
+                    .has_value());
+
+    auto delete_non_empty =
+        HttpDelete(port, "/api/v1/tenant_quotas?tenant_id=tenant-a");
+    EXPECT_EQ(delete_non_empty.http_status, 409);
+    EXPECT_NE(delete_non_empty.body.find("TENANT_NOT_EMPTY"),
+              std::string::npos);
+
+    ASSERT_TRUE(service->Remove("quota_admin_key", /*force=*/true, "tenant-a")
+                    .has_value());
+
+    auto deleted = HttpDelete(port, "/api/v1/tenant_quotas?tenant_id=tenant-a");
+    EXPECT_EQ(deleted.http_status, 200);
+
+    auto missing = HttpGet(port, "/api/v1/tenant_quotas?tenant_id=tenant-a");
+    EXPECT_EQ(missing.http_status, 404);
+
+    admin.Stop();
+    std::filesystem::remove(policy_path);
+}
+
+TEST_F(MasterAdminServerTest, TenantQuotaAdminValidationErrors) {
+    const std::string policy_path = WriteTenantQuotaPolicyForTest({});
+    WrappedMasterServiceConfig svc_config;
+    svc_config.default_kv_lease_ttl = 5000;
+    svc_config.enable_metric_reporting = false;
+    svc_config.enable_multi_tenants = true;
+    svc_config.tenant_quota_connector_type = "file";
+    svc_config.tenant_quota_connector_uri = policy_path;
+    auto service = std::make_shared<WrappedMasterService>(svc_config);
+
+    int port = getFreeTcpPort();
+    MasterAdminServer admin(static_cast<uint16_t>(port), false);
+    ASSERT_TRUE(admin.Start());
+    admin.SetRuntimeState(ha::MasterRuntimeState::kServing);
+    admin.SetServiceDelegate(service);
+    admin.SetServiceAvailable(true);
+
+    auto missing_tenant = HttpPutJson(port, "/api/v1/tenant_quotas",
+                                      "{\"requested_quota_bytes\":100}");
+    EXPECT_EQ(missing_tenant.http_status, 400);
+
+    auto empty_tenant = HttpPutJson(port, "/api/v1/tenant_quotas?tenant_id=",
+                                    "{\"requested_quota_bytes\":100}");
+    EXPECT_EQ(empty_tenant.http_status, 400);
+
+    auto zero_explicit =
+        HttpPutJson(port, "/api/v1/tenant_quotas?tenant_id=tenant-a",
+                    "{\"requested_quota_bytes\":0}");
+    EXPECT_EQ(zero_explicit.http_status, 400);
+
+    auto reserved_tenant =
+        HttpGet(port, "/api/v1/tenant_quotas?tenant_id=_system");
+    EXPECT_EQ(reserved_tenant.http_status, 400);
+
+    auto missing_query =
+        HttpGet(port, "/api/v1/tenant_quotas?tenant_id=missing");
+    EXPECT_EQ(missing_query.http_status, 404);
+
+    admin.Stop();
+    std::filesystem::remove(policy_path);
+}
+
+TEST_F(MasterAdminServerTest, TenantQuotaAdminDisabledModeReturns409) {
+    WrappedMasterServiceConfig svc_config;
+    svc_config.default_kv_lease_ttl = 5000;
+    svc_config.enable_metric_reporting = false;
+    svc_config.enable_multi_tenants = false;
+    auto service = std::make_shared<WrappedMasterService>(svc_config);
+
+    int port = getFreeTcpPort();
+    MasterAdminServer admin(static_cast<uint16_t>(port), false);
+    ASSERT_TRUE(admin.Start());
+    admin.SetRuntimeState(ha::MasterRuntimeState::kServing);
+    admin.SetServiceDelegate(service);
+    admin.SetServiceAvailable(true);
+
+    auto list = HttpGet(port, "/api/v1/tenant_quotas");
+    EXPECT_EQ(list.http_status, 409);
+    EXPECT_NE(list.body.find("UNAVAILABLE_IN_CURRENT_MODE"), std::string::npos);
+
+    admin.Stop();
+}
+
+// =========================================================================
+// MasterAdminServerWithServiceTest — reuses a single server+service across
+// all tests via SetUpTestSuite / TearDownTestSuite for fast execution.
+// =========================================================================
+
+class MasterAdminServerWithServiceTest : public ::testing::Test {
+   protected:
+    struct HttpResponse {
+        int http_status;
+        std::string body;
+    };
+
+    static void SetUpTestSuite() {
+        WrappedMasterServiceConfig svc_config;
+        svc_config.default_kv_lease_ttl = 5000;
+        svc_config.enable_metric_reporting = false;
+        svc_config.client_live_ttl_sec =
+            3600;  // prevent client expiry in slow CI
+        service_ = std::make_shared<WrappedMasterService>(svc_config);
+
+        segment_.id = generate_uuid();
+        segment_.name = "admin_test_segment";
+        segment_.base = 0x300000000;
+        segment_.size = 8 * 1024 * 1024;
+        UUID client_id = generate_uuid();
+        (void)service_->MountSegment(segment_, client_id);
+
+        ReplicateConfig cfg;
+        cfg.replica_num = 1;
+        auto ps = service_->PutStart(client_id, kDefaultKey, 1024, cfg);
+        if (ps.has_value()) {
+            (void)service_->PutEnd(client_id, kDefaultKey, ReplicaType::MEMORY);
+        }
+
+        port_ = getFreeTcpPort();
+        admin_ = std::make_unique<MasterAdminServer>(
+            static_cast<uint16_t>(port_), false);
+        ASSERT_TRUE(admin_->Start());
+        admin_->SetRuntimeState(ha::MasterRuntimeState::kServing);
+        admin_->SetServiceDelegate(service_);
+        admin_->SetServiceAvailable(true);
+    }
+
+    static void TearDownTestSuite() {
+        if (admin_) {
+            admin_->Stop();
+            admin_.reset();
+        }
+        service_.reset();
+    }
+
+    HttpResponse HttpGet(const std::string& path) {
+        coro_http::coro_http_client client;
+        auto result = client.get(BaseUrl() + path);
+        return {result.status, std::string(result.resp_body)};
+    }
+
+    HttpResponse HttpPostJson(const std::string& path,
+                              const std::string& body) {
+        coro_http::coro_http_client client;
+        auto result = client.post(BaseUrl() + path, body,
+                                  coro_http::req_content_type::json);
+        return {result.status, std::string(result.resp_body)};
+    }
+
+    static std::string BaseUrl() {
+        return "http://127.0.0.1:" + std::to_string(port_);
+    }
+
+    static std::shared_ptr<WrappedMasterService> service_;
+    static std::unique_ptr<MasterAdminServer> admin_;
+    static int port_;
+    static Segment segment_;
+    static constexpr const char* kDefaultKey = "admin_test_key";
+};
+
+std::shared_ptr<WrappedMasterService>
+    MasterAdminServerWithServiceTest::service_;
+std::unique_ptr<MasterAdminServer> MasterAdminServerWithServiceTest::admin_;
+int MasterAdminServerWithServiceTest::port_ = 0;
+Segment MasterAdminServerWithServiceTest::segment_;
+
+// -----------------------------------------------------------------------
+// GET /get_all_keys
+// -----------------------------------------------------------------------
+
+TEST_F(MasterAdminServerWithServiceTest, GetAllKeysReturnsKeys) {
+    auto resp = HttpGet("/get_all_keys");
+    EXPECT_EQ(resp.http_status, 200);
+    EXPECT_NE(resp.body.find(kDefaultKey), std::string::npos);
+}
+
+TEST_F(MasterAdminServerWithServiceTest, GetAllKeysExcludesRemovedKey) {
+    // Use a unique key so removal doesn't affect other tests.
+    const std::string key = "ephemeral_empty_test_key";
+    UUID client_id = generate_uuid();
+    ReplicateConfig cfg;
+    cfg.replica_num = 1;
+    auto ps = service_->PutStart(client_id, key, 1024, cfg);
+    if (ps.has_value()) {
+        (void)service_->PutEnd(client_id, key, ReplicaType::MEMORY);
+    }
+    (void)service_->Remove(key, "default");
+
+    auto resp = HttpGet("/get_all_keys");
+    EXPECT_EQ(resp.http_status, 200);
+    // The default key is still present; only the ephemeral key was removed.
+    EXPECT_NE(resp.body.find(kDefaultKey), std::string::npos);
+    EXPECT_EQ(resp.body.find(key), std::string::npos);
+}
+
+// -----------------------------------------------------------------------
+// GET /query_key
+// -----------------------------------------------------------------------
+
+TEST_F(MasterAdminServerWithServiceTest, QueryKeyReturnsDataForExistingKey) {
+    auto resp = HttpGet("/query_key?key=" + std::string(kDefaultKey));
+    EXPECT_EQ(resp.http_status, 200);
+    EXPECT_NE(resp.body.find("\"buffer_address_\""), std::string::npos);
+}
+
+TEST_F(MasterAdminServerWithServiceTest, QueryKeyReturns404ForNonexistentKey) {
+    auto resp = HttpGet("/query_key?key=nonexistent_key_xyz");
+    EXPECT_EQ(resp.http_status, 404);
+}
+
+TEST_F(MasterAdminServerWithServiceTest, QueryKeyWithoutKeyParamReturns404) {
+    auto resp = HttpGet("/query_key");
+    EXPECT_EQ(resp.http_status, 404);
+}
+
+// -----------------------------------------------------------------------
+// GET /query_segment
+// -----------------------------------------------------------------------
+
+TEST_F(MasterAdminServerWithServiceTest,
+       QuerySegmentReturnsDataForExistingSegment) {
+    auto resp = HttpGet("/query_segment?segment=" + segment_.name);
+    EXPECT_EQ(resp.http_status, 200);
+    EXPECT_NE(resp.body.find(segment_.name), std::string::npos);
+    EXPECT_NE(resp.body.find("Used(bytes)"), std::string::npos);
+    EXPECT_NE(resp.body.find("Capacity(bytes)"), std::string::npos);
+}
+
+TEST_F(MasterAdminServerWithServiceTest,
+       QuerySegmentReturns500ForNonexistentSegment) {
+    auto resp = HttpGet("/query_segment?segment=nonexistent_seg");
+    EXPECT_EQ(resp.http_status, 500);
+}
+
+// -----------------------------------------------------------------------
+// GET /get_all_segments
+// -----------------------------------------------------------------------
+
+TEST_F(MasterAdminServerWithServiceTest, GetAllSegmentsReturnsSegments) {
+    auto resp = HttpGet("/get_all_segments");
+    EXPECT_EQ(resp.http_status, 200);
+    EXPECT_NE(resp.body.find(segment_.name), std::string::npos);
+}
+
+// -----------------------------------------------------------------------
+// GET /get_segments_detail
+// -----------------------------------------------------------------------
+
+TEST_F(MasterAdminServerWithServiceTest, GetSegmentsDetailReturnsDetailedInfo) {
+    auto resp = HttpGet("/get_segments_detail");
+    ASSERT_EQ(resp.http_status, 200);
+
+    HttpSegmentsDetailResponse parsed;
+    struct_json::from_json(parsed, resp.body);
+    EXPECT_GT(parsed.total_segments, 0u);
+    EXPECT_NE(resp.body.find(segment_.name), std::string::npos);
+    EXPECT_NE(resp.body.find("\"allocator_used_bytes\""), std::string::npos);
+    EXPECT_NE(resp.body.find("\"allocator_capacity_bytes\""),
+              std::string::npos);
+}
+
+// -----------------------------------------------------------------------
+// POST /api/v1/drain_jobs
+// -----------------------------------------------------------------------
+
+TEST_F(MasterAdminServerWithServiceTest,
+       CreateDrainJobSucceedsWithValidRequest) {
+    // Use a unique segment so subsequent drain tests are not blocked.
+    std::string seg = "drain_create_seg_" + UuidToString(generate_uuid());
+    Segment s;
+    s.id = generate_uuid();
+    s.name = seg;
+    s.base = 0x600000000;
+    s.size = 4 * 1024 * 1024;
+    (void)service_->MountSegment(s, generate_uuid());
+
+    std::string body = R"({"segments":[")" + seg + R"("]})";
+    auto resp = HttpPostJson("/api/v1/drain_jobs", body);
+    EXPECT_EQ(resp.http_status, 200);
+
+    HttpCreateDrainJobResponse parsed;
+    struct_json::from_json(parsed, resp.body);
+    EXPECT_TRUE(parsed.success);
+    EXPECT_FALSE(parsed.job_id.empty());
+    EXPECT_EQ(parsed.status, "CREATED");
+}
+
+TEST_F(MasterAdminServerWithServiceTest, CreateDrainJobFailsWithInvalidJson) {
+    auto resp = HttpPostJson("/api/v1/drain_jobs", "not json");
+    ASSERT_EQ(resp.http_status, 400);
+
+    HttpErrorResponse parsed;
+    struct_json::from_json(parsed, resp.body);
+    EXPECT_FALSE(parsed.success);
+}
+
+TEST_F(MasterAdminServerWithServiceTest, CreateDrainJobFailsWithEmptyBody) {
+    auto resp = HttpPostJson("/api/v1/drain_jobs", "{}");
+    EXPECT_EQ(resp.http_status, 400);
+}
+
+// -----------------------------------------------------------------------
+// GET /api/v1/drain_jobs/query
+// -----------------------------------------------------------------------
+
+TEST_F(MasterAdminServerWithServiceTest, QueryDrainJobReturnsCreatedJob) {
+    std::string seg = "drain_query_seg_" + UuidToString(generate_uuid());
+    Segment s;
+    s.id = generate_uuid();
+    s.name = seg;
+    s.base = 0x700000000;
+    s.size = 4 * 1024 * 1024;
+    (void)service_->MountSegment(s, generate_uuid());
+
+    std::string body = R"({"segments":[")" + seg + R"("]})";
+    auto create_resp = HttpPostJson("/api/v1/drain_jobs", body);
+    ASSERT_EQ(create_resp.http_status, 200);
+
+    HttpCreateDrainJobResponse create_parsed;
+    struct_json::from_json(create_parsed, create_resp.body);
+    ASSERT_FALSE(create_parsed.job_id.empty());
+
+    auto query_resp =
+        HttpGet("/api/v1/drain_jobs/query?job_id=" + create_parsed.job_id);
+    EXPECT_EQ(query_resp.http_status, 200);
+
+    HttpQueryDrainJobResponse query_parsed;
+    struct_json::from_json(query_parsed, query_resp.body);
+    EXPECT_TRUE(query_parsed.success);
+    EXPECT_EQ(query_parsed.job_id, create_parsed.job_id);
+}
+
+TEST_F(MasterAdminServerWithServiceTest, QueryDrainJobFailsWithInvalidJobId) {
+    auto resp = HttpGet("/api/v1/drain_jobs/query?job_id=not-a-uuid");
+    EXPECT_EQ(resp.http_status, 400);
+}
+
+TEST_F(MasterAdminServerWithServiceTest, QueryDrainJobFailsWithMissingJobId) {
+    auto resp = HttpGet("/api/v1/drain_jobs/query");
+    EXPECT_EQ(resp.http_status, 400);
+}
+
+TEST_F(MasterAdminServerWithServiceTest, QueryDrainJobFailsForNonexistentJob) {
+    auto resp = HttpGet("/api/v1/drain_jobs/query?job_id=" +
+                        UuidToString(generate_uuid()));
+    EXPECT_EQ(resp.http_status, 404);
+}
+
+// -----------------------------------------------------------------------
+// POST /api/v1/drain_jobs/cancel
+// -----------------------------------------------------------------------
+
+TEST_F(MasterAdminServerWithServiceTest, CancelDrainJobSucceeds) {
+    std::string seg = "drain_cancel_seg_" + UuidToString(generate_uuid());
+    Segment s;
+    s.id = generate_uuid();
+    s.name = seg;
+    s.base = 0x800000000;
+    s.size = 4 * 1024 * 1024;
+    (void)service_->MountSegment(s, generate_uuid());
+
+    std::string body = R"({"segments":[")" + seg + R"("]})";
+    auto create_resp = HttpPostJson("/api/v1/drain_jobs", body);
+    EXPECT_EQ(create_resp.http_status, 200);
+
+    HttpCreateDrainJobResponse create_parsed;
+    struct_json::from_json(create_parsed, create_resp.body);
+    ASSERT_FALSE(create_parsed.job_id.empty());
+
+    auto cancel_resp = HttpPostJson(
+        "/api/v1/drain_jobs/cancel?job_id=" + create_parsed.job_id, "");
+    EXPECT_EQ(cancel_resp.http_status, 200);
+
+    HttpCancelDrainJobResponse cancel_parsed;
+    struct_json::from_json(cancel_parsed, cancel_resp.body);
+    EXPECT_TRUE(cancel_parsed.success);
+    EXPECT_EQ(cancel_parsed.job_id, create_parsed.job_id);
+    EXPECT_EQ(cancel_parsed.status, "CANCELED");
+}
+
+TEST_F(MasterAdminServerWithServiceTest, CancelDrainJobFailsWithInvalidJobId) {
+    auto resp = HttpPostJson("/api/v1/drain_jobs/cancel?job_id=not-a-uuid", "");
+    EXPECT_EQ(resp.http_status, 400);
+}
+
+TEST_F(MasterAdminServerWithServiceTest, CancelDrainJobFailsWithMissingJobId) {
+    auto resp = HttpPostJson("/api/v1/drain_jobs/cancel", "");
+    EXPECT_EQ(resp.http_status, 400);
+}
+
+// -----------------------------------------------------------------------
+// GET /api/v1/segments/status
+// -----------------------------------------------------------------------
+
+TEST_F(MasterAdminServerWithServiceTest,
+       SegmentStatusReturnsDataForExistingSegment) {
+    auto resp = HttpGet("/api/v1/segments/status?segment=" + segment_.name);
+    EXPECT_EQ(resp.http_status, 200);
+
+    HttpSegmentStatusResponse parsed;
+    struct_json::from_json(parsed, resp.body);
+    EXPECT_TRUE(parsed.success);
+    EXPECT_EQ(parsed.segment, segment_.name);
+    EXPECT_FALSE(parsed.status_name.empty());
+}
+
+TEST_F(MasterAdminServerWithServiceTest, SegmentStatusFailsWithMissingSegment) {
+    auto resp = HttpGet("/api/v1/segments/status");
+    EXPECT_EQ(resp.http_status, 400);
+}
+
+TEST_F(MasterAdminServerWithServiceTest,
+       SegmentStatusReturnsErrorForNonexistentSegment) {
+    auto resp = HttpGet("/api/v1/segments/status?segment=no_such_segment");
+    EXPECT_EQ(resp.http_status, 404);
+}
+
+// -----------------------------------------------------------------------
+// GET /batch_query_keys
+// -----------------------------------------------------------------------
+
+TEST_F(MasterAdminServerWithServiceTest,
+       BatchQueryKeysWithNoKeysParamReturns400) {
+    auto resp = HttpGet("/batch_query_keys");
+    EXPECT_EQ(resp.http_status, 400);
+}
+
+TEST_F(MasterAdminServerWithServiceTest,
+       BatchQueryKeysWithEmptyKeysParamReturns400) {
+    auto resp = HttpGet("/batch_query_keys?keys=");
+    EXPECT_EQ(resp.http_status, 400);
+}
+
+TEST_F(MasterAdminServerWithServiceTest,
+       BatchQueryKeysReturnsDataForExistingKey) {
+    auto resp = HttpGet("/batch_query_keys?keys=" + std::string(kDefaultKey));
+    EXPECT_EQ(resp.http_status, 200);
+    EXPECT_NE(resp.body.find("\"success\":true"), std::string::npos);
+    EXPECT_NE(resp.body.find("\"data\":{"), std::string::npos);
+    EXPECT_NE(resp.body.find(kDefaultKey), std::string::npos);
+}
+
+TEST_F(MasterAdminServerWithServiceTest,
+       BatchQueryKeysReturnsErrorForNonexistentKey) {
+    auto resp = HttpGet("/batch_query_keys?keys=nonexistent_key");
+    EXPECT_EQ(resp.http_status, 200);
+    EXPECT_NE(resp.body.find("\"success\":true"), std::string::npos);
+    EXPECT_NE(resp.body.find("\"ok\":false"), std::string::npos);
+}
+
+TEST_F(MasterAdminServerWithServiceTest, BatchQueryKeysMultipleKeys) {
+    UUID client_id = generate_uuid();
+    ReplicateConfig cfg;
+    cfg.replica_num = 1;
+    auto ps = service_->PutStart(client_id, "second_key", 512, cfg);
+    if (ps.has_value()) {
+        (void)service_->PutEnd(client_id, "second_key", ReplicaType::MEMORY);
+    }
+
+    auto resp = HttpGet("/batch_query_keys?keys=" + std::string(kDefaultKey) +
+                        ",second_key");
+    EXPECT_EQ(resp.http_status, 200);
+    EXPECT_NE(resp.body.find(kDefaultKey), std::string::npos);
+    EXPECT_NE(resp.body.find("second_key"), std::string::npos);
+
+    (void)service_->Remove("second_key", "default");
+}
+
+TEST_F(MasterAdminServerWithServiceTest,
+       BatchQueryKeysWithExistingAndNonexistentKeys) {
+    auto resp = HttpGet("/batch_query_keys?keys=" + std::string(kDefaultKey) +
+                        ",nonexistent");
+    EXPECT_EQ(resp.http_status, 200);
+    EXPECT_NE(resp.body.find("\"success\":true"), std::string::npos);
+    EXPECT_NE(resp.body.find("\"ok\":true"), std::string::npos);
+    EXPECT_NE(resp.body.find("\"ok\":false"), std::string::npos);
+}
+
+TEST_F(MasterAdminServerWithServiceTest, DrainJobFullLifecycle) {
+    std::string seg = "drain_lifecycle_seg_" + UuidToString(generate_uuid());
+    Segment s;
+    s.id = generate_uuid();
+    s.name = seg;
+    s.base = 0x900000000;
+    s.size = 4 * 1024 * 1024;
+    (void)service_->MountSegment(s, generate_uuid());
+
+    std::string body = R"({"segments":[")" + seg + R"("]})";
+    auto create_resp = HttpPostJson("/api/v1/drain_jobs", body);
+    ASSERT_EQ(create_resp.http_status, 200);
+    HttpCreateDrainJobResponse create_parsed;
+    struct_json::from_json(create_parsed, create_resp.body);
+    ASSERT_TRUE(create_parsed.success);
+    std::string job_id = create_parsed.job_id;
+
+    auto query_resp = HttpGet("/api/v1/drain_jobs/query?job_id=" + job_id);
+    EXPECT_EQ(query_resp.http_status, 200);
+    HttpQueryDrainJobResponse query_parsed;
+    struct_json::from_json(query_parsed, query_resp.body);
+    EXPECT_TRUE(query_parsed.success);
+    EXPECT_EQ(query_parsed.job_id, job_id);
+
+    auto cancel_resp =
+        HttpPostJson("/api/v1/drain_jobs/cancel?job_id=" + job_id, "");
+    EXPECT_EQ(cancel_resp.http_status, 200);
+    HttpCancelDrainJobResponse cancel_parsed;
+    struct_json::from_json(cancel_parsed, cancel_resp.body);
+    EXPECT_TRUE(cancel_parsed.success);
+    EXPECT_EQ(cancel_parsed.job_id, job_id);
+    EXPECT_EQ(cancel_parsed.status, "CANCELED");
+
+    query_resp = HttpGet("/api/v1/drain_jobs/query?job_id=" + job_id);
+    EXPECT_EQ(query_resp.http_status, 200);
+    struct_json::from_json(query_parsed, query_resp.body);
+    EXPECT_TRUE(query_parsed.success);
+}
+
+// =========================================================================
+// Destructive / isolated tests that must run on their own server instance.
+// =========================================================================
+
+TEST_F(MasterAdminServerTest, ServiceUnavailableAfterDelegateCleared) {
+    WrappedMasterServiceConfig svc_config;
+    svc_config.default_kv_lease_ttl = 5000;
+    svc_config.enable_metric_reporting = false;
+    auto service = std::make_shared<WrappedMasterService>(svc_config);
+
+    int port = getFreeTcpPort();
+    MasterAdminServer admin(static_cast<uint16_t>(port), false);
+    ASSERT_TRUE(admin.Start());
+    admin.SetRuntimeState(ha::MasterRuntimeState::kServing);
+    admin.SetServiceDelegate(service);
+    admin.SetServiceAvailable(true);
+
+    auto resp = HttpGet(port, "/get_all_segments");
+    EXPECT_EQ(resp.http_status, 200);
+
+    admin.SetServiceDelegate(nullptr);
+
+    resp = HttpGet(port, "/get_all_segments");
+    EXPECT_EQ(resp.http_status, 503);
+    EXPECT_NE(resp.body.find("service plane is not active"), std::string::npos);
+
+    admin.Stop();
+}
+
+TEST_F(MasterAdminServerTest, ServiceUnavailableAfterSetServiceAvailableFalse) {
+    WrappedMasterServiceConfig svc_config;
+    svc_config.default_kv_lease_ttl = 5000;
+    svc_config.enable_metric_reporting = false;
+    auto service = std::make_shared<WrappedMasterService>(svc_config);
+
+    int port = getFreeTcpPort();
+    MasterAdminServer admin(static_cast<uint16_t>(port), false);
+    ASSERT_TRUE(admin.Start());
+    admin.SetRuntimeState(ha::MasterRuntimeState::kServing);
+    admin.SetServiceDelegate(service);
+    admin.SetServiceAvailable(true);
+
+    auto resp = HttpGet(port, "/get_all_segments");
+    EXPECT_EQ(resp.http_status, 200);
+
+    admin.SetServiceAvailable(false);
+
+    resp = HttpGet(port, "/get_all_segments");
+    EXPECT_EQ(resp.http_status, 503);
+
+    admin.SetServiceAvailable(true);
+    resp = HttpGet(port, "/get_all_segments");
+    EXPECT_EQ(resp.http_status, 200);
+
+    admin.Stop();
+}
+
+TEST_F(MasterAdminServerTest, MultipleSegmentsAndKeys) {
+    WrappedMasterServiceConfig svc_config;
+    svc_config.default_kv_lease_ttl = 5000;
+    svc_config.enable_metric_reporting = false;
+    auto service = std::make_shared<WrappedMasterService>(svc_config);
+
+    UUID client_id = generate_uuid();
+
+    Segment seg1;
+    seg1.id = generate_uuid();
+    seg1.name = "seg_alpha";
+    seg1.base = 0x400000000;
+    seg1.size = 8 * 1024 * 1024;
+    ASSERT_TRUE(service->MountSegment(seg1, client_id).has_value());
+
+    Segment seg2;
+    seg2.id = generate_uuid();
+    seg2.name = "seg_beta";
+    seg2.base = 0x500000000;
+    seg2.size = 4 * 1024 * 1024;
+    ASSERT_TRUE(service->MountSegment(seg2, client_id).has_value());
+
+    ReplicateConfig cfg;
+    cfg.replica_num = 1;
+    auto ps1 = service->PutStart(client_id, "key_one", 1024, cfg);
+    if (ps1.has_value()) {
+        (void)service->PutEnd(client_id, "key_one", ReplicaType::MEMORY);
+    }
+    auto ps2 = service->PutStart(client_id, "key_two", 2048, cfg);
+    if (ps2.has_value()) {
+        (void)service->PutEnd(client_id, "key_two", ReplicaType::MEMORY);
+    }
+
+    int port = getFreeTcpPort();
+    MasterAdminServer admin(static_cast<uint16_t>(port), false);
+    ASSERT_TRUE(admin.Start());
+    admin.SetRuntimeState(ha::MasterRuntimeState::kServing);
+    admin.SetServiceDelegate(service);
+    admin.SetServiceAvailable(true);
+
+    auto seg_resp = HttpGet(port, "/get_all_segments");
+    EXPECT_EQ(seg_resp.http_status, 200);
+    EXPECT_NE(seg_resp.body.find("seg_alpha"), std::string::npos);
+    EXPECT_NE(seg_resp.body.find("seg_beta"), std::string::npos);
+
+    auto keys_resp = HttpGet(port, "/get_all_keys");
+    EXPECT_EQ(keys_resp.http_status, 200);
+    EXPECT_NE(keys_resp.body.find("key_one"), std::string::npos);
+    EXPECT_NE(keys_resp.body.find("key_two"), std::string::npos);
+
+    auto q1 = HttpGet(port, "/query_segment?segment=seg_alpha");
+    EXPECT_EQ(q1.http_status, 200);
+    EXPECT_NE(q1.body.find("seg_alpha"), std::string::npos);
+
+    auto q2 = HttpGet(port, "/query_segment?segment=seg_beta");
+    EXPECT_EQ(q2.http_status, 200);
+    EXPECT_NE(q2.body.find("seg_beta"), std::string::npos);
+
+    auto s1 = HttpGet(port, "/api/v1/segments/status?segment=seg_alpha");
+    EXPECT_EQ(s1.http_status, 200);
+    HttpSegmentStatusResponse parsed1;
+    struct_json::from_json(parsed1, s1.body);
+    EXPECT_TRUE(parsed1.success);
+    EXPECT_EQ(parsed1.segment, "seg_alpha");
+
+    auto s2 = HttpGet(port, "/api/v1/segments/status?segment=seg_beta");
+    EXPECT_EQ(s2.http_status, 200);
+    HttpSegmentStatusResponse parsed2;
+    struct_json::from_json(parsed2, s2.body);
+    EXPECT_TRUE(parsed2.success);
+    EXPECT_EQ(parsed2.segment, "seg_beta");
+
+    auto detail_resp = HttpGet(port, "/get_segments_detail");
+    EXPECT_EQ(detail_resp.http_status, 200);
+    HttpSegmentsDetailResponse detail_parsed;
+    struct_json::from_json(detail_parsed, detail_resp.body);
+    EXPECT_EQ(detail_parsed.total_segments, 2u);
+
+    auto batch_resp = HttpGet(port, "/batch_query_keys?keys=key_one,key_two");
+    EXPECT_EQ(batch_resp.http_status, 200);
+    EXPECT_NE(batch_resp.body.find("key_one"), std::string::npos);
+    EXPECT_NE(batch_resp.body.find("key_two"), std::string::npos);
+
+    admin.Stop();
+}
+
+}  // namespace test
+}  // namespace mooncake
+
+int main(int argc, char** argv) {
+    google::InitGoogleLogging(argv[0]);
+    FLAGS_logtostderr = true;
+    ::testing::InitGoogleTest(&argc, argv);
+    int result = RUN_ALL_TESTS();
+    google::ShutdownGoogleLogging();
+    return result;
+}
diff --git a/mooncake-store/tests/master_metrics_test.cpp b/mooncake-store/tests/master_metrics_test.cpp
index 6e110fd8..812d0ff3 100644
--- a/mooncake-store/tests/master_metrics_test.cpp
+++ b/mooncake-store/tests/master_metrics_test.cpp
@@ -1,6 +1,9 @@
 #include <glog/logging.h>
 #include <gtest/gtest.h>
 
+#include <algorithm>
+#include <chrono>
+#include <cmath>
 #include <csignal>
 #include <thread>
 #include <vector>
@@ -8,6 +11,8 @@
 #include <ylt/coro_http/coro_http_client.hpp>
 
 #include "utils.h"
+#include "master_admin_service.h"
+#include "master_service.h"
 #include "rpc_service.h"
 #include "types.h"
 #include "master_config.h"
@@ -58,6 +63,7 @@ TEST_F(MasterMetricsTest, InitialStatusTest) {
     // Operation Statistics
     ASSERT_EQ(metrics.get_put_start_requests(), 0);
     ASSERT_EQ(metrics.get_put_start_failures(), 0);
+    ASSERT_EQ(metrics.get_put_start_alloc_failures(), 0);
     ASSERT_EQ(metrics.get_put_end_requests(), 0);
     ASSERT_EQ(metrics.get_put_end_failures(), 0);
     ASSERT_EQ(metrics.get_put_revoke_requests(), 0);
@@ -205,13 +211,13 @@ TEST_F(MasterMetricsTest, BasicRequestTest) {
     ASSERT_EQ(metrics.get_put_end_failures(), 0);
 
     // Test ExistKey request
-    auto exist_result = service_.ExistKey(key);
+    auto exist_result = service_.ExistKey(key, "default");
     ASSERT_TRUE(exist_result.has_value() && exist_result.value());
     ASSERT_EQ(metrics.get_exist_key_requests(), 1);
     ASSERT_EQ(metrics.get_exist_key_failures(), 0);
 
     // Test GetReplicaList request
-    auto get_replica_result = service_.GetReplicaList(key);
+    auto get_replica_result = service_.GetReplicaList(key, "default");
     ASSERT_TRUE(get_replica_result.has_value());
     ASSERT_EQ(metrics.get_get_replica_list_requests(), 1);
     ASSERT_EQ(metrics.get_get_replica_list_failures(), 0);
@@ -219,7 +225,7 @@ TEST_F(MasterMetricsTest, BasicRequestTest) {
     // Test Remove request
     std::this_thread::sleep_for(
         std::chrono::milliseconds(default_kv_lease_ttl));
-    auto remove_result = service_.Remove(key);
+    auto remove_result = service_.Remove(key, "default");
     ASSERT_TRUE(remove_result.has_value());
     ASSERT_EQ(metrics.get_remove_requests(), 1);
     ASSERT_EQ(metrics.get_remove_failures(), 0);
@@ -267,6 +273,55 @@ TEST_F(MasterMetricsTest, BasicRequestTest) {
 TEST_F(MasterMetricsTest, CalcCacheStatsTest) {
     const uint64_t default_kv_lease_ttl = 100;
     auto& metrics = MasterMetricManager::instance();
+    using CacheHitStat = MasterMetricManager::CacheHitStat;
+
+    // These values are part of the RPC/API contract. New enum entries should
+    // be appended instead of renumbering existing values.
+    ASSERT_EQ(static_cast<int>(CacheHitStat::MEMORY_HITS), 0);
+    ASSERT_EQ(static_cast<int>(CacheHitStat::SSD_HITS), 1);
+    ASSERT_EQ(static_cast<int>(CacheHitStat::MEMORY_TOTAL), 2);
+    ASSERT_EQ(static_cast<int>(CacheHitStat::SSD_TOTAL), 3);
+    ASSERT_EQ(static_cast<int>(CacheHitStat::MEMORY_HIT_RATE), 4);
+    ASSERT_EQ(static_cast<int>(CacheHitStat::SSD_HIT_RATE), 5);
+    ASSERT_EQ(static_cast<int>(CacheHitStat::OVERALL_HIT_RATE), 6);
+    ASSERT_EQ(static_cast<int>(CacheHitStat::VALID_GET_RATE), 7);
+
+    auto round_to_2 = [](double value) {
+        return std::round(value * 100.0) / 100.0;
+    };
+    auto expected_ratio = [&round_to_2](double hits, double total) {
+        return total > 0.0 ? round_to_2(hits / total) : 0.0;
+    };
+    auto expect_aliases =
+        [](const MasterMetricManager::CacheHitStatDict& stats) {
+            ASSERT_EQ(stats.at(CacheHitStat::MEMORY_CURRENT_CACHED_OBJECTS),
+                      stats.at(CacheHitStat::MEMORY_TOTAL));
+            ASSERT_EQ(stats.at(CacheHitStat::SSD_CURRENT_CACHED_OBJECTS),
+                      stats.at(CacheHitStat::SSD_TOTAL));
+            ASSERT_EQ(
+                stats.at(CacheHitStat::MEMORY_HITS_PER_CURRENT_CACHED_OBJECT),
+                stats.at(CacheHitStat::MEMORY_HIT_RATE));
+            ASSERT_EQ(
+                stats.at(CacheHitStat::SSD_HITS_PER_CURRENT_CACHED_OBJECT),
+                stats.at(CacheHitStat::SSD_HIT_RATE));
+            ASSERT_EQ(
+                stats.at(CacheHitStat::OVERALL_HITS_PER_CURRENT_CACHED_OBJECT),
+                stats.at(CacheHitStat::OVERALL_HIT_RATE));
+        };
+    auto expect_reuse_ratios =
+        [&expected_ratio](const MasterMetricManager::CacheHitStatDict& stats) {
+            const double memory_hits = stats.at(CacheHitStat::MEMORY_HITS);
+            const double ssd_hits = stats.at(CacheHitStat::SSD_HITS);
+            const double memory_total = stats.at(CacheHitStat::MEMORY_TOTAL);
+            const double ssd_total = stats.at(CacheHitStat::SSD_TOTAL);
+            ASSERT_EQ(stats.at(CacheHitStat::MEMORY_HIT_RATE),
+                      expected_ratio(memory_hits, memory_total));
+            ASSERT_EQ(stats.at(CacheHitStat::SSD_HIT_RATE),
+                      expected_ratio(ssd_hits, ssd_total));
+            ASSERT_EQ(stats.at(CacheHitStat::OVERALL_HIT_RATE),
+                      expected_ratio(memory_hits + ssd_hits,
+                                     memory_total + ssd_total));
+        };
     // Use a wrapped master service to test the metrics manager
     WrappedMasterServiceConfig service_config;
     service_config.default_kv_lease_ttl = default_kv_lease_ttl;
@@ -289,17 +344,16 @@ TEST_F(MasterMetricsTest, CalcCacheStatsTest) {
     ReplicateConfig config;
     config.replica_num = 1;
 
-    auto stats_dict = metrics.calculate_cache_stats();
-    ASSERT_EQ(stats_dict[MasterMetricManager::CacheHitStat::MEMORY_HITS], 1);
-    ASSERT_EQ(stats_dict[MasterMetricManager::CacheHitStat::SSD_HITS], 0);
-    ASSERT_EQ(stats_dict[MasterMetricManager::CacheHitStat::MEMORY_TOTAL], 2);
-    ASSERT_EQ(stats_dict[MasterMetricManager::CacheHitStat::SSD_TOTAL], 0);
-    ASSERT_EQ(stats_dict[MasterMetricManager::CacheHitStat::MEMORY_HIT_RATE],
-              0.5);
-    ASSERT_EQ(stats_dict[MasterMetricManager::CacheHitStat::SSD_HIT_RATE], 0);
-    ASSERT_EQ(stats_dict[MasterMetricManager::CacheHitStat::OVERALL_HIT_RATE],
-              0.5);
-    ASSERT_EQ(stats_dict[MasterMetricManager::CacheHitStat::VALID_GET_RATE], 1);
+    // MasterMetricManager is a process-wide singleton and these counters do
+    // not have reset APIs, so assert deltas from the current baseline.
+    const auto base_stats = metrics.calculate_cache_stats();
+    expect_aliases(base_stats);
+    expect_reuse_ratios(base_stats);
+
+    const double base_memory_hits = base_stats.at(CacheHitStat::MEMORY_HITS);
+    const double base_memory_total = base_stats.at(CacheHitStat::MEMORY_TOTAL);
+    const double base_valid_get_rate =
+        base_stats.at(CacheHitStat::VALID_GET_RATE);
 
     auto mount_result = service_.MountSegment(segment, client_id);
     ASSERT_TRUE(mount_result.has_value());
@@ -308,26 +362,43 @@ TEST_F(MasterMetricsTest, CalcCacheStatsTest) {
     ASSERT_TRUE(put_start_result1.has_value());
     auto put_end_result1 = service_.PutEnd(client_id, key, ReplicaType::MEMORY);
     ASSERT_TRUE(put_end_result1.has_value());
-    stats_dict = metrics.calculate_cache_stats();
+    auto stats_dict = metrics.calculate_cache_stats();
 
-    ASSERT_EQ(stats_dict[MasterMetricManager::CacheHitStat::MEMORY_TOTAL], 3);
+    expect_aliases(stats_dict);
+    expect_reuse_ratios(stats_dict);
+    ASSERT_EQ(stats_dict[CacheHitStat::MEMORY_HITS], base_memory_hits);
+    ASSERT_EQ(stats_dict[CacheHitStat::MEMORY_TOTAL], base_memory_total + 1);
 
-    auto get_replica_result = service_.GetReplicaList(key);
+    auto get_replica_result = service_.GetReplicaList(key, "default");
+    ASSERT_TRUE(get_replica_result.has_value());
     stats_dict = metrics.calculate_cache_stats();
-    ASSERT_EQ(stats_dict[MasterMetricManager::CacheHitStat::MEMORY_HITS], 2);
-    ASSERT_EQ(stats_dict[MasterMetricManager::CacheHitStat::SSD_HITS], 0);
-    ASSERT_EQ(stats_dict[MasterMetricManager::CacheHitStat::MEMORY_TOTAL], 3);
-    ASSERT_EQ(stats_dict[MasterMetricManager::CacheHitStat::SSD_TOTAL], 0);
-    ASSERT_NEAR(stats_dict[MasterMetricManager::CacheHitStat::MEMORY_HIT_RATE],
-                0.67, 0.01);
-    ASSERT_EQ(stats_dict[MasterMetricManager::CacheHitStat::SSD_HIT_RATE], 0);
-    ASSERT_NEAR(stats_dict[MasterMetricManager::CacheHitStat::OVERALL_HIT_RATE],
-                0.67, 0.01);
-    ASSERT_EQ(stats_dict[MasterMetricManager::CacheHitStat::VALID_GET_RATE], 1);
+    expect_aliases(stats_dict);
+    expect_reuse_ratios(stats_dict);
+    ASSERT_EQ(stats_dict[CacheHitStat::MEMORY_HITS], base_memory_hits + 1);
+    ASSERT_EQ(stats_dict[CacheHitStat::MEMORY_TOTAL], base_memory_total + 1);
+    ASSERT_GE(stats_dict[CacheHitStat::VALID_GET_RATE], base_valid_get_rate);
+    ASSERT_LE(stats_dict[CacheHitStat::VALID_GET_RATE], 1.0);
+
+    // This value is not a bounded hit ratio: hits are cumulative while cached
+    // objects are a current gauge.
+    // Keep fetching until cumulative hits exceed current cached objects.
+    const auto extra_gets = std::max<int64_t>(
+        1, static_cast<int64_t>(stats_dict[CacheHitStat::MEMORY_TOTAL] -
+                                stats_dict[CacheHitStat::MEMORY_HITS]) +
+               1);
+    for (int64_t i = 0; i < extra_gets; ++i) {
+        get_replica_result = service_.GetReplicaList(key, "default");
+        ASSERT_TRUE(get_replica_result.has_value());
+    }
+    stats_dict = metrics.calculate_cache_stats();
+    expect_aliases(stats_dict);
+    expect_reuse_ratios(stats_dict);
+    ASSERT_GT(stats_dict[CacheHitStat::MEMORY_HITS_PER_CURRENT_CACHED_OBJECT],
+              1.0);
 
     std::this_thread::sleep_for(
         std::chrono::milliseconds(default_kv_lease_ttl));
-    auto remove_result = service_.Remove(key);
+    auto remove_result = service_.Remove(key, "default");
     ASSERT_TRUE(remove_result.has_value());
 }
 
@@ -361,6 +432,11 @@ TEST_F(MasterMetricsTest, AdminServerExposesStandbyStateWithoutService) {
     EXPECT_NE(segments_resp.body.find("service plane is not active"),
               std::string::npos);
 
+    auto detail_resp = FetchUrl(http_port, "/get_segments_detail");
+    EXPECT_EQ(detail_resp.http_status, 503);
+    EXPECT_NE(detail_resp.body.find("service plane is not active"),
+              std::string::npos);
+
     admin_server.Stop();
 }
 
@@ -396,6 +472,16 @@ TEST_F(MasterMetricsTest, AdminServerRoutesServiceEndpointsWhenAvailable) {
     EXPECT_NE(query_resp.body.find(segment.name), std::string::npos);
     EXPECT_NE(query_resp.body.find("Capacity(bytes)"), std::string::npos);
 
+    auto detail_resp = FetchUrl(http_port, "/get_segments_detail");
+    EXPECT_EQ(detail_resp.http_status, 200);
+    EXPECT_NE(detail_resp.body.find("\"total_segments\""), std::string::npos);
+    EXPECT_NE(detail_resp.body.find("\"segments\""), std::string::npos);
+    EXPECT_NE(detail_resp.body.find(segment.name), std::string::npos);
+    EXPECT_NE(detail_resp.body.find("\"allocator_used_bytes\""),
+              std::string::npos);
+    EXPECT_NE(detail_resp.body.find("\"allocator_capacity_bytes\""),
+              std::string::npos);
+
     admin_server.Stop();
 }
 
@@ -427,7 +513,7 @@ TEST_F(MasterMetricsTest, BatchRequestTest) {
     ASSERT_TRUE(mount_result.has_value());
 
     // Test BatchExistKey request (should all return false initially)
-    auto batch_exist_result = service_.BatchExistKey(keys);
+    auto batch_exist_result = service_.BatchExistKey(keys, "default");
     ASSERT_EQ(batch_exist_result.size(), 3);
     ASSERT_EQ(metrics.get_batch_exist_key_requests(), 1);
     ASSERT_EQ(metrics.get_batch_exist_key_partial_successes(), 0);
@@ -464,7 +550,7 @@ TEST_F(MasterMetricsTest, BatchRequestTest) {
     ASSERT_EQ(metrics.get_batch_put_end_failed_items(), 0);
 
     // Test BatchExistKey again (should all return true now)
-    auto batch_exist_result2 = service_.BatchExistKey(keys);
+    auto batch_exist_result2 = service_.BatchExistKey(keys, "default");
     ASSERT_EQ(batch_exist_result2.size(), 3);
     ASSERT_EQ(metrics.get_batch_exist_key_requests(), 2);
     ASSERT_EQ(metrics.get_batch_exist_key_partial_successes(), 0);
@@ -511,6 +597,194 @@ TEST_F(MasterMetricsTest, BatchRequestTest) {
     ASSERT_EQ(metrics.get_batch_put_start_failed_items(), 3);
 }
 
+// ---------------------------------------------------------------------------
+// Tests for local SSD metrics (Bug 1 fix)
+// ---------------------------------------------------------------------------
+
+// Helper: put a key into a mem segment and immediately notify offload success,
+// simulating the client-side SSD write completing.
+static std::string PutKeyAndOffload(MasterService& svc, const UUID& client_id,
+                                    const std::string& segment_name,
+                                    uint64_t value_size,
+                                    const std::string& key) {
+    ReplicateConfig cfg;
+    cfg.replica_num = 1;
+    auto put_start = svc.PutStart(client_id, key, "default", value_size, cfg);
+    if (!put_start) return "";
+    svc.PutEnd(client_id, key, "default", ReplicaType::MEMORY);
+
+    StorageObjectMetadata meta;
+    meta.data_size = static_cast<int64_t>(value_size);
+    meta.transport_endpoint = "127.0.0.1:9999";
+    std::vector<OffloadTaskItem> tasks{
+        OffloadTaskItem{.tenant_id = "default",
+                        .key = key,
+                        .size = static_cast<int64_t>(value_size)}};
+    svc.NotifyOffloadSuccess(client_id, tasks, {meta});
+    return key;
+}
+
+// Verify that creating a LocalDiskReplica (via NotifyOffloadSuccess) increments
+// file_allocated_size, and that removing the key decrements it back to zero.
+TEST_F(MasterMetricsTest, LocalDiskReplicaAllocatedSize) {
+    auto& metrics = MasterMetricManager::instance();
+    MasterServiceConfig config;
+    config.enable_offload = true;
+    MasterService svc(config);
+
+    constexpr size_t kBuf = 0x400000000;
+    constexpr size_t kSegSize = 64 * 1024 * 1024;
+    constexpr uint64_t kValueSize = 4096;
+
+    UUID client_id = generate_uuid();
+    Segment seg;
+    seg.id = generate_uuid();
+    seg.name = "ssd_alloc_test_segment";
+    seg.base = kBuf;
+    seg.size = kSegSize;
+    seg.te_endpoint = seg.name;
+
+    ASSERT_TRUE(svc.MountSegment(seg, client_id).has_value());
+    ASSERT_TRUE(svc.MountLocalDiskSegment(client_id, true).has_value());
+
+    const int64_t baseline = metrics.get_allocated_file_size();
+
+    // After NotifyOffloadSuccess, a LocalDiskReplica is created.
+    std::string key = PutKeyAndOffload(svc, client_id, seg.name, kValueSize,
+                                       "ssd_alloc_test_key");
+    ASSERT_FALSE(key.empty());
+    EXPECT_EQ(metrics.get_allocated_file_size(), baseline + kValueSize);
+
+    // After removing the key the LocalDiskReplica is destroyed; gauge resets.
+    ASSERT_TRUE(svc.Remove(key, "default").has_value());
+    EXPECT_EQ(metrics.get_allocated_file_size(), baseline);
+}
+
+// Verify that OffloadObjectHeartbeat updates total_file_capacity correctly,
+// including when a client reports a changed capacity on a subsequent heartbeat.
+TEST_F(MasterMetricsTest, LocalDiskSegmentCapacityHeartbeat) {
+    auto& metrics = MasterMetricManager::instance();
+    MasterServiceConfig config;
+    config.enable_offload = true;
+    MasterService svc(config);
+
+    constexpr size_t kBuf = 0x500000000;
+    constexpr size_t kSegSize = 64 * 1024 * 1024;
+    constexpr int64_t kCap1 = 800LL * 1024 * 1024 * 1024;  // 800 GB
+    constexpr int64_t kCap2 = 400LL * 1024 * 1024 * 1024;  // 400 GB
+
+    UUID client_id = generate_uuid();
+    Segment seg;
+    seg.id = generate_uuid();
+    seg.name = "ssd_capacity_test_segment";
+    seg.base = kBuf;
+    seg.size = kSegSize;
+    seg.te_endpoint = seg.name;
+
+    ASSERT_TRUE(svc.MountSegment(seg, client_id).has_value());
+    ASSERT_TRUE(svc.MountLocalDiskSegment(client_id, true).has_value());
+
+    const int64_t baseline = metrics.get_total_file_capacity();
+
+    // ReportSsdCapacity: client reports 800 GB.
+    ASSERT_TRUE(svc.ReportSsdCapacity(client_id, kCap1).has_value());
+    EXPECT_EQ(metrics.get_total_file_capacity(), baseline + kCap1);
+
+    // Client reports 400 GB (e.g. config changed).
+    // Gauge must be updated to reflect the new value, not double-counted.
+    ASSERT_TRUE(svc.ReportSsdCapacity(client_id, kCap2).has_value());
+    EXPECT_EQ(metrics.get_total_file_capacity(), baseline + kCap2);
+
+    // Idempotent: same capacity reported again — gauge must not change.
+    ASSERT_TRUE(svc.ReportSsdCapacity(client_id, kCap2).has_value());
+    EXPECT_EQ(metrics.get_total_file_capacity(), baseline + kCap2);
+}
+
+TEST_F(MasterMetricsTest, PutStartReplicaAllocationFailureMetric) {
+    const uint64_t default_kv_lease_ttl = 100;
+    auto& metrics = MasterMetricManager::instance();
+    WrappedMasterServiceConfig service_config;
+    service_config.default_kv_lease_ttl = default_kv_lease_ttl;
+    WrappedMasterService service_(service_config);
+
+    const int64_t allocation_failures_before =
+        metrics.get_put_start_alloc_failures();
+    const int64_t put_start_failures_before = metrics.get_put_start_failures();
+
+    ReplicateConfig config;
+    config.replica_num = 1;
+    auto put_start_result = service_.PutStart(
+        generate_uuid(), "allocation_failure_key", 1024, config);
+
+    ASSERT_FALSE(put_start_result.has_value());
+    ASSERT_EQ(put_start_result.error(), ErrorCode::NO_AVAILABLE_HANDLE);
+    ASSERT_EQ(metrics.get_put_start_alloc_failures(),
+              allocation_failures_before + 1);
+    ASSERT_EQ(metrics.get_put_start_failures(), put_start_failures_before + 1);
+}
+
+TEST_F(MasterMetricsTest, SummaryUsesWindowRatesAndCumulativeEviction) {
+    auto& metrics = MasterMetricManager::instance();
+
+    const std::string baseline_summary =
+        metrics.get_summary_string_and_update_snapshot();
+    EXPECT_NE(baseline_summary.find("Requests (Success/Total per sec):"),
+              std::string::npos);
+    EXPECT_NE(baseline_summary.find("PutStart=0.00/0.00"), std::string::npos);
+
+    metrics.inc_put_start_requests(4);
+    metrics.inc_put_start_failures(1);
+    metrics.inc_batch_put_start_requests(5);
+    metrics.inc_batch_put_start_partial_success(2);
+    metrics.inc_eviction_success(3, 4096);
+    metrics.inc_eviction_fail();
+    metrics.inc_mem_eviction_success(3, 4096);
+    metrics.inc_mem_eviction_fail();
+    metrics.inc_nof_eviction_success(1, 2048);
+    metrics.inc_nof_eviction_fail();
+
+    std::this_thread::sleep_for(std::chrono::milliseconds(20));
+    const std::string window_summary = metrics.get_summary_string();
+    EXPECT_NE(window_summary.find("Requests (Success/Total per sec):"),
+              std::string::npos);
+    EXPECT_NE(window_summary.find("PutStart="), std::string::npos);
+    EXPECT_EQ(window_summary.find("/s"), std::string::npos);
+    EXPECT_EQ(window_summary.find("PutStart=3/4"), std::string::npos);
+    EXPECT_NE(window_summary.find("Batch Requests (per sec"),
+              std::string::npos);
+    EXPECT_NE(
+        window_summary.find("Eviction: Success/Attempts=1/2, AllocFail=0, "
+                            "keys=3, size=4.00 KB"),
+        std::string::npos);
+    EXPECT_NE(window_summary.find("Mem Eviction: Success/Attempts=1/2, "
+                                  "keys=3, size=4.00 KB"),
+              std::string::npos);
+    EXPECT_NE(window_summary.find("NoF Eviction: Success/Attempts=1/2, "
+                                  "keys=1, size=2.00 KB"),
+              std::string::npos);
+
+    const std::string reported_summary =
+        metrics.get_summary_string_and_update_snapshot();
+    EXPECT_NE(
+        reported_summary.find("Eviction: Success/Attempts=1/2, AllocFail=0, "
+                              "keys=3, size=4.00 KB"),
+        std::string::npos);
+
+    std::this_thread::sleep_for(std::chrono::milliseconds(20));
+    const std::string idle_summary =
+        metrics.get_summary_string_and_update_snapshot();
+    EXPECT_NE(idle_summary.find("PutStart=0.00/0.00"), std::string::npos);
+    EXPECT_NE(idle_summary.find("Eviction: Success/Attempts=1/2, "
+                                "AllocFail=0, keys=3, size=4.00 KB"),
+              std::string::npos);
+    EXPECT_NE(idle_summary.find("Mem Eviction: Success/Attempts=1/2, "
+                                "keys=3, size=4.00 KB"),
+              std::string::npos);
+    EXPECT_NE(idle_summary.find("NoF Eviction: Success/Attempts=1/2, "
+                                "keys=1, size=2.00 KB"),
+              std::string::npos);
+}
+
 }  // namespace mooncake::test
 
 int main(int argc, char** argv) {
diff --git a/mooncake-store/tests/master_service_ssd_test.cpp b/mooncake-store/tests/master_service_ssd_test.cpp
index 0a2f6033..b24e7db6 100644
--- a/mooncake-store/tests/master_service_ssd_test.cpp
+++ b/mooncake-store/tests/master_service_ssd_test.cpp
@@ -7,6 +7,7 @@
 #include <thread>
 #include <vector>
 
+#include "master_metric_manager.h"
 #include "types.h"
 
 namespace mooncake::test {
@@ -27,6 +28,66 @@ class MasterServiceSSDTest : public ::testing::Test {
     void TearDown() override { google::ShutdownGoogleLogging(); }
 };
 
+std::unique_ptr<MasterService> CreateSsdAwareOffloadService() {
+    MasterServiceConfig config;
+    config.enable_offload = true;
+    config.default_kv_lease_ttl = 0;
+    config.allocation_strategy_type =
+        AllocationStrategyType::SSD_FREE_RATIO_FIRST;
+    return std::make_unique<MasterService>(config);
+}
+
+void MountMemoryAndLocalDisk(MasterService& service, const UUID& client_id,
+                             const std::string& segment_name,
+                             size_t base_addr) {
+    Segment segment;
+    segment.id = generate_uuid();
+    segment.name = segment_name;
+    segment.base = base_addr;
+    segment.size = 64 * 1024 * 1024;
+    segment.te_endpoint = segment.name;
+
+    ASSERT_TRUE(service.MountSegment(segment, client_id).has_value());
+    ASSERT_TRUE(service.MountLocalDiskSegment(client_id, true).has_value());
+    ASSERT_TRUE(service.ReportSsdCapacity(client_id, 1000).has_value());
+}
+
+void PutAndOffload(MasterService& service, const UUID& client_id,
+                   const std::string& key, int64_t object_size,
+                   const std::string& local_disk_endpoint) {
+    ReplicateConfig config;
+    config.replica_num = 1;
+
+    ASSERT_TRUE(service.PutStart(client_id, key, "default", object_size, config)
+                    .has_value());
+    ASSERT_TRUE(service.PutEnd(client_id, key, "default", ReplicaType::MEMORY)
+                    .has_value());
+
+    StorageObjectMetadata metadata;
+    metadata.data_size = object_size;
+    metadata.transport_endpoint = local_disk_endpoint;
+    OffloadTaskItem task{
+        .tenant_id = "default", .key = key, .size = object_size};
+    ASSERT_TRUE(service.NotifyOffloadSuccess(client_id, {task}, {metadata})
+                    .has_value());
+}
+
+void ExpectNextAllocationOnSegment(MasterService& service,
+                                   const UUID& client_id,
+                                   const std::string& key,
+                                   const std::string& expected_segment) {
+    ReplicateConfig config;
+    config.replica_num = 1;
+    auto result = service.PutStart(client_id, key, "default", 64, config);
+    ASSERT_TRUE(result.has_value());
+    ASSERT_EQ(result->size(), 1u);
+    ASSERT_TRUE((*result)[0].is_memory_replica());
+    EXPECT_EQ((*result)[0]
+                  .get_memory_descriptor()
+                  .buffer_descriptor.transport_endpoint_,
+              expected_segment);
+}
+
 TEST_F(MasterServiceSSDTest, PutEndBothReplica) {
     auto service_ = CreateMasterServiceWithSSDFeat("/mnt/ssd");
 
@@ -50,7 +111,7 @@ TEST_F(MasterServiceSSDTest, PutEndBothReplica) {
     config.replica_num = 1;
 
     auto put_start_result =
-        service_->PutStart(client_id, key, slice_length, config);
+        service_->PutStart(client_id, key, "default", slice_length, config);
     ASSERT_TRUE(put_start_result.has_value());
     auto replicas = put_start_result.value();
     ASSERT_EQ(2, replicas.size());
@@ -63,17 +124,17 @@ TEST_F(MasterServiceSSDTest, PutEndBothReplica) {
     EXPECT_TRUE(has_mem);
     EXPECT_TRUE(has_disk);
 
-    auto get_result = service_->GetReplicaList(key);
+    auto get_result = service_->GetReplicaList(key, "default");
     ASSERT_FALSE(get_result.has_value());
     EXPECT_EQ(ErrorCode::REPLICA_IS_NOT_READY, get_result.error());
 
     // PutEnd for both memory and disk
-    EXPECT_TRUE(
-        service_->PutEnd(client_id, key, ReplicaType::MEMORY).has_value());
-    EXPECT_TRUE(
-        service_->PutEnd(client_id, key, ReplicaType::DISK).has_value());
+    EXPECT_TRUE(service_->PutEnd(client_id, key, "default", ReplicaType::MEMORY)
+                    .has_value());
+    EXPECT_TRUE(service_->PutEnd(client_id, key, "default", ReplicaType::DISK)
+                    .has_value());
 
-    get_result = service_->GetReplicaList(key);
+    get_result = service_->GetReplicaList(key, "default");
     ASSERT_TRUE(get_result.has_value());
     EXPECT_EQ(2, get_result.value().replicas.size());
 
@@ -105,24 +166,71 @@ TEST_F(MasterServiceSSDTest, PutRevokeDiskReplica) {
     config.replica_num = 1;
 
     ASSERT_TRUE(
-        service_->PutStart(client_id, key, slice_length, config).has_value());
-    EXPECT_TRUE(
-        service_->PutEnd(client_id, key, ReplicaType::MEMORY).has_value());
+        service_->PutStart(client_id, key, "default", slice_length, config)
+            .has_value());
+    EXPECT_TRUE(service_->PutEnd(client_id, key, "default", ReplicaType::MEMORY)
+                    .has_value());
 
-    auto get_result = service_->GetReplicaList(key);
+    auto get_result = service_->GetReplicaList(key, "default");
     ASSERT_TRUE(get_result.has_value());
     EXPECT_EQ(1, get_result.value().replicas.size());
     ASSERT_TRUE(get_result.value().replicas[0].is_memory_replica());
 
     EXPECT_TRUE(
-        service_->PutRevoke(client_id, key, ReplicaType::DISK).has_value());
+        service_->PutRevoke(client_id, key, "default", ReplicaType::DISK)
+            .has_value());
 
-    get_result = service_->GetReplicaList(key);
+    get_result = service_->GetReplicaList(key, "default");
     ASSERT_TRUE(get_result.has_value());
     EXPECT_EQ(1, get_result.value().replicas.size());
     ASSERT_TRUE(get_result.value().replicas[0].is_memory_replica());
 }
 
+TEST_F(MasterServiceSSDTest, PutRevokeProcessingDiskKeepsSsdTotal) {
+    auto service_ = CreateMasterServiceWithSSDFeat("/mnt/ssd");
+    auto& metrics = MasterMetricManager::instance();
+    using CacheHitStat = MasterMetricManager::CacheHitStat;
+    const auto base_stats = metrics.calculate_cache_stats();
+    const double base_memory_total = base_stats.at(CacheHitStat::MEMORY_TOTAL);
+    const double base_ssd_total = base_stats.at(CacheHitStat::SSD_TOTAL);
+
+    constexpr size_t buffer = 0x310000000;
+    constexpr size_t size = 1024 * 1024 * 64;
+    Segment segment;
+    segment.id = generate_uuid();
+    segment.name = "test_segment_revoke_processing_disk";
+    segment.base = buffer;
+    segment.size = size;
+    segment.te_endpoint = segment.name;
+    UUID client_id = generate_uuid();
+
+    ASSERT_TRUE(service_->MountSegment(segment, client_id).has_value());
+
+    std::string key = "revoke_processing_disk_metric_key";
+    ASSERT_TRUE(
+        service_->PutStart(client_id, key, "default", 1024, {.replica_num = 1})
+            .has_value());
+    EXPECT_TRUE(service_->PutEnd(client_id, key, "default", ReplicaType::MEMORY)
+                    .has_value());
+
+    auto stats = metrics.calculate_cache_stats();
+    EXPECT_EQ(stats[CacheHitStat::MEMORY_TOTAL], base_memory_total + 1);
+    EXPECT_EQ(stats[CacheHitStat::SSD_TOTAL], base_ssd_total);
+
+    EXPECT_TRUE(
+        service_->PutRevoke(client_id, key, "default", ReplicaType::DISK)
+            .has_value());
+
+    stats = metrics.calculate_cache_stats();
+    EXPECT_EQ(stats[CacheHitStat::MEMORY_TOTAL], base_memory_total + 1);
+    EXPECT_EQ(stats[CacheHitStat::SSD_TOTAL], base_ssd_total);
+
+    ASSERT_TRUE(service_->Remove(key, "default", /*force=*/true).has_value());
+    stats = metrics.calculate_cache_stats();
+    EXPECT_EQ(stats[CacheHitStat::MEMORY_TOTAL], base_memory_total);
+    EXPECT_EQ(stats[CacheHitStat::SSD_TOTAL], base_ssd_total);
+}
+
 TEST_F(MasterServiceSSDTest, PutRevokeMemoryReplica) {
     auto service_ = CreateMasterServiceWithSSDFeat("/mnt/ssd");
 
@@ -146,17 +254,19 @@ TEST_F(MasterServiceSSDTest, PutRevokeMemoryReplica) {
     config.replica_num = 1;
 
     ASSERT_TRUE(
-        service_->PutStart(client_id, key, slice_length, config).has_value());
+        service_->PutStart(client_id, key, "default", slice_length, config)
+            .has_value());
     EXPECT_TRUE(
-        service_->PutRevoke(client_id, key, ReplicaType::MEMORY).has_value());
+        service_->PutRevoke(client_id, key, "default", ReplicaType::MEMORY)
+            .has_value());
 
-    auto get_result = service_->GetReplicaList(key);
+    auto get_result = service_->GetReplicaList(key, "default");
     ASSERT_FALSE(get_result.has_value());
     EXPECT_EQ(ErrorCode::REPLICA_IS_NOT_READY, get_result.error());
 
-    EXPECT_TRUE(
-        service_->PutEnd(client_id, key, ReplicaType::DISK).has_value());
-    get_result = service_->GetReplicaList(key);
+    EXPECT_TRUE(service_->PutEnd(client_id, key, "default", ReplicaType::DISK)
+                    .has_value());
+    get_result = service_->GetReplicaList(key, "default");
     ASSERT_TRUE(get_result.has_value());
     EXPECT_EQ(1, get_result.value().replicas.size());
     ASSERT_TRUE(get_result.value().replicas[0].is_disk_replica());
@@ -185,17 +295,20 @@ TEST_F(MasterServiceSSDTest, PutRevokeBothReplica) {
     config.replica_num = 1;
 
     ASSERT_TRUE(
-        service_->PutStart(client_id, key, slice_length, config).has_value());
+        service_->PutStart(client_id, key, "default", slice_length, config)
+            .has_value());
     EXPECT_TRUE(
-        service_->PutRevoke(client_id, key, ReplicaType::DISK).has_value());
+        service_->PutRevoke(client_id, key, "default", ReplicaType::DISK)
+            .has_value());
 
-    auto get_result = service_->GetReplicaList(key);
+    auto get_result = service_->GetReplicaList(key, "default");
     ASSERT_FALSE(get_result.has_value());
     EXPECT_EQ(ErrorCode::REPLICA_IS_NOT_READY, get_result.error());
 
     EXPECT_TRUE(
-        service_->PutRevoke(client_id, key, ReplicaType::MEMORY).has_value());
-    get_result = service_->GetReplicaList(key);
+        service_->PutRevoke(client_id, key, "default", ReplicaType::MEMORY)
+            .has_value());
+    get_result = service_->GetReplicaList(key, "default");
     ASSERT_FALSE(get_result.has_value());
     EXPECT_EQ(ErrorCode::OBJECT_NOT_FOUND, get_result.error());
 }
@@ -223,15 +336,16 @@ TEST_F(MasterServiceSSDTest, RemoveKey) {
     config.replica_num = 1;
 
     ASSERT_TRUE(
-        service_->PutStart(client_id, key, slice_length, config).has_value());
-    EXPECT_TRUE(
-        service_->PutEnd(client_id, key, ReplicaType::MEMORY).has_value());
-    EXPECT_TRUE(
-        service_->PutEnd(client_id, key, ReplicaType::DISK).has_value());
+        service_->PutStart(client_id, key, "default", slice_length, config)
+            .has_value());
+    EXPECT_TRUE(service_->PutEnd(client_id, key, "default", ReplicaType::MEMORY)
+                    .has_value());
+    EXPECT_TRUE(service_->PutEnd(client_id, key, "default", ReplicaType::DISK)
+                    .has_value());
 
-    EXPECT_TRUE(service_->Remove(key).has_value());
+    EXPECT_TRUE(service_->Remove(key, "default").has_value());
 
-    auto get_result = service_->GetReplicaList(key);
+    auto get_result = service_->GetReplicaList(key, "default");
     EXPECT_FALSE(get_result.has_value());
     EXPECT_EQ(ErrorCode::OBJECT_NOT_FOUND, get_result.error());
 }
@@ -264,12 +378,12 @@ TEST_F(MasterServiceSSDTest, EvictObject) {
         ReplicateConfig config;
         config.replica_num = 1;
         auto put_start_result =
-            service_->PutStart(client_id, key, slice_length, config);
+            service_->PutStart(client_id, key, "default", slice_length, config);
         if (put_start_result.has_value()) {
-            auto put_end_mem_result =
-                service_->PutEnd(client_id, key, ReplicaType::MEMORY);
+            auto put_end_mem_result = service_->PutEnd(
+                client_id, key, "default", ReplicaType::MEMORY);
             auto put_end_disk_result =
-                service_->PutEnd(client_id, key, ReplicaType::DISK);
+                service_->PutEnd(client_id, key, "default", ReplicaType::DISK);
             ASSERT_TRUE(put_end_mem_result.has_value());
             ASSERT_TRUE(put_end_disk_result.has_value());
             success_puts++;
@@ -284,7 +398,7 @@ TEST_F(MasterServiceSSDTest, EvictObject) {
     int success_gets = 0;
     for (int i = 0; i < 1024 * 16 + 50; ++i) {
         std::string key = "test_key" + std::to_string(i);
-        auto get_result = service_->GetReplicaList(key);
+        auto get_result = service_->GetReplicaList(key, "default");
         if (get_result.has_value()) {
             success_gets++;
         }
@@ -335,7 +449,7 @@ TEST_F(MasterServiceSSDTest, PutStartExpires) {
 
         // Put key, should success.
         auto put_start_result =
-            service_->PutStart(client_id, key, slice_length, config);
+            service_->PutStart(client_id, key, "default", slice_length, config);
         EXPECT_TRUE(put_start_result.has_value());
         auto replica_list = put_start_result.value();
         EXPECT_EQ(replica_list.size(), kReplicaCnt);
@@ -344,7 +458,8 @@ TEST_F(MasterServiceSSDTest, PutStartExpires) {
         }
 
         // Complete the reserved replica.
-        auto put_end_result = service_->PutEnd(client_id, key, reserve_type);
+        auto put_end_result =
+            service_->PutEnd(client_id, key, "default", reserve_type);
         EXPECT_TRUE(put_end_result.has_value());
 
         // Wait for a while until the put-start expired.
@@ -354,7 +469,7 @@ TEST_F(MasterServiceSSDTest, PutStartExpires) {
             auto result = service_->Ping(client_id);
             EXPECT_TRUE(result.has_value());
             // Protect the key from eviction.
-            auto get_result = service_->GetReplicaList(key);
+            auto get_result = service_->GetReplicaList(key, "default");
             EXPECT_TRUE(get_result.has_value());
             std::this_thread::sleep_for(std::chrono::seconds(1));
         }
@@ -362,7 +477,7 @@ TEST_F(MasterServiceSSDTest, PutStartExpires) {
         // Put key again, should fail because the object has had an completed
         // replica.
         put_start_result =
-            service_->PutStart(client_id, key, slice_length, config);
+            service_->PutStart(client_id, key, "default", slice_length, config);
         EXPECT_FALSE(put_start_result.has_value());
         EXPECT_EQ(put_start_result.error(), ErrorCode::OBJECT_ALREADY_EXISTS);
 
@@ -373,17 +488,18 @@ TEST_F(MasterServiceSSDTest, PutStartExpires) {
             auto result = service_->Ping(client_id);
             EXPECT_TRUE(result.has_value());
             // Protect the key from eviction.
-            auto get_result = service_->GetReplicaList(key);
+            auto get_result = service_->GetReplicaList(key, "default");
             EXPECT_TRUE(get_result.has_value());
             std::this_thread::sleep_for(std::chrono::seconds(1));
         }
 
         // Try PutEnd the discarded replica.
-        put_end_result = service_->PutEnd(client_id, key, discard_type);
+        put_end_result =
+            service_->PutEnd(client_id, key, "default", discard_type);
         EXPECT_TRUE(put_end_result.has_value());
 
         // Check that the key has only one replica.
-        auto get_result = service_->GetReplicaList(key);
+        auto get_result = service_->GetReplicaList(key, "default");
         EXPECT_TRUE(get_result.has_value());
         EXPECT_EQ(get_result.value().replicas.size(), 1);
         if (reserve_type == ReplicaType::MEMORY) {
@@ -423,38 +539,78 @@ TEST_F(MasterServiceSSDTest, EvictDiskReplica_RemovesDiskReplica) {
 
     std::string key = "evict_disk_key";
     auto put_result =
-        service_->PutStart(client_id, key, 1024, {.replica_num = 1});
+        service_->PutStart(client_id, key, "default", 1024, {.replica_num = 1});
     ASSERT_TRUE(put_result.has_value());
 
     // Complete both replicas
-    EXPECT_TRUE(
-        service_->PutEnd(client_id, key, ReplicaType::MEMORY).has_value());
-    EXPECT_TRUE(
-        service_->PutEnd(client_id, key, ReplicaType::DISK).has_value());
+    EXPECT_TRUE(service_->PutEnd(client_id, key, "default", ReplicaType::MEMORY)
+                    .has_value());
+    EXPECT_TRUE(service_->PutEnd(client_id, key, "default", ReplicaType::DISK)
+                    .has_value());
 
     // Verify we have 2 replicas (MEM + DISK)
-    auto get_result = service_->GetReplicaList(key);
+    auto get_result = service_->GetReplicaList(key, "default");
     ASSERT_TRUE(get_result.has_value());
     EXPECT_EQ(2, get_result.value().replicas.size());
 
     // Evict disk replica
-    auto evict_result =
-        service_->EvictDiskReplica(client_id, key, ReplicaType::DISK);
+    auto evict_result = service_->EvictDiskReplica(client_id, key, "default",
+                                                   ReplicaType::DISK);
     ASSERT_TRUE(evict_result.has_value());
 
     // Verify only memory replica remains
-    get_result = service_->GetReplicaList(key);
+    get_result = service_->GetReplicaList(key, "default");
     ASSERT_TRUE(get_result.has_value());
     EXPECT_EQ(1, get_result.value().replicas.size());
     EXPECT_TRUE(get_result.value().replicas[0].is_memory_replica());
 }
 
+TEST_F(MasterServiceSSDTest, RemoveDecrementsCacheTotalMetrics) {
+    auto service_ = CreateMasterServiceWithSSDFeat("/mnt/ssd");
+    auto& metrics = MasterMetricManager::instance();
+    using CacheHitStat = MasterMetricManager::CacheHitStat;
+    const auto base_stats = metrics.calculate_cache_stats();
+    const double base_memory_total = base_stats.at(CacheHitStat::MEMORY_TOTAL);
+    const double base_ssd_total = base_stats.at(CacheHitStat::SSD_TOTAL);
+
+    constexpr size_t buffer = 0x320000000;
+    constexpr size_t size = 1024 * 1024 * 64;
+    Segment segment;
+    segment.id = generate_uuid();
+    segment.name = "test_segment_remove_metrics";
+    segment.base = buffer;
+    segment.size = size;
+    segment.te_endpoint = segment.name;
+    UUID client_id = generate_uuid();
+
+    ASSERT_TRUE(service_->MountSegment(segment, client_id).has_value());
+
+    std::string key = "remove_cache_total_metric_key";
+    ASSERT_TRUE(
+        service_->PutStart(client_id, key, "default", 1024, {.replica_num = 1})
+            .has_value());
+    EXPECT_TRUE(service_->PutEnd(client_id, key, "default", ReplicaType::MEMORY)
+                    .has_value());
+    EXPECT_TRUE(service_->PutEnd(client_id, key, "default", ReplicaType::DISK)
+                    .has_value());
+
+    auto stats = metrics.calculate_cache_stats();
+    EXPECT_EQ(stats[CacheHitStat::MEMORY_TOTAL], base_memory_total + 1);
+    EXPECT_EQ(stats[CacheHitStat::SSD_TOTAL], base_ssd_total + 1);
+
+    ASSERT_TRUE(service_->Remove(key, "default", /*force=*/true).has_value());
+
+    stats = metrics.calculate_cache_stats();
+    EXPECT_EQ(stats[CacheHitStat::MEMORY_TOTAL], base_memory_total);
+    EXPECT_EQ(stats[CacheHitStat::SSD_TOTAL], base_ssd_total);
+}
+
 TEST_F(MasterServiceSSDTest, EvictDiskReplica_NonExistentKeyReturnsError) {
     auto service_ = CreateMasterServiceWithSSDFeat("/mnt/ssd");
 
     UUID client_id = generate_uuid();
-    auto evict_result = service_->EvictDiskReplica(client_id, "nonexistent_key",
-                                                   ReplicaType::DISK);
+    auto evict_result = service_->EvictDiskReplica(
+        client_id, "nonexistent_key", "default", ReplicaType::DISK);
     EXPECT_FALSE(evict_result.has_value());
     EXPECT_EQ(evict_result.error(), ErrorCode::OBJECT_NOT_FOUND);
 }
@@ -477,20 +633,234 @@ TEST_F(MasterServiceSSDTest, EvictDiskReplica_InvalidReplicaTypeReturnsError) {
 
     std::string key = "evict_invalid_type_key";
     auto put_result =
-        service_->PutStart(client_id, key, 1024, {.replica_num = 1});
+        service_->PutStart(client_id, key, "default", 1024, {.replica_num = 1});
     ASSERT_TRUE(put_result.has_value());
-    EXPECT_TRUE(
-        service_->PutEnd(client_id, key, ReplicaType::MEMORY).has_value());
-    EXPECT_TRUE(
-        service_->PutEnd(client_id, key, ReplicaType::DISK).has_value());
+    EXPECT_TRUE(service_->PutEnd(client_id, key, "default", ReplicaType::MEMORY)
+                    .has_value());
+    EXPECT_TRUE(service_->PutEnd(client_id, key, "default", ReplicaType::DISK)
+                    .has_value());
 
     // Attempting to evict with MEMORY type should fail
-    auto evict_result =
-        service_->EvictDiskReplica(client_id, key, ReplicaType::MEMORY);
+    auto evict_result = service_->EvictDiskReplica(client_id, key, "default",
+                                                   ReplicaType::MEMORY);
     EXPECT_FALSE(evict_result.has_value());
     EXPECT_EQ(evict_result.error(), ErrorCode::INVALID_PARAMS);
 }
 
+TEST_F(MasterServiceSSDTest, RemoveReleasesLocalDiskUsageTracking) {
+    auto service = CreateSsdAwareOffloadService();
+    UUID client1 = generate_uuid();
+    UUID client2 = generate_uuid();
+    const std::string segment1 = "ssd_remove_segment_1";
+    const std::string segment2 = "ssd_remove_segment_2";
+    MountMemoryAndLocalDisk(*service, client1, segment1, 0x400000000);
+    MountMemoryAndLocalDisk(*service, client2, segment2, 0x500000000);
+
+    PutAndOffload(*service, client1, "ssd_remove_released", 800, segment1);
+    PutAndOffload(*service, client2, "ssd_remove_baseline", 100, segment2);
+
+    ASSERT_TRUE(service->Remove("ssd_remove_released", "default").has_value());
+
+    ExpectNextAllocationOnSegment(*service, client1, "ssd_remove_probe",
+                                  segment1);
+}
+
+TEST_F(MasterServiceSSDTest,
+       BatchReplicaClearAllSegmentsReleasesLocalDiskUsageTracking) {
+    auto service = CreateSsdAwareOffloadService();
+    UUID client1 = generate_uuid();
+    UUID client2 = generate_uuid();
+    const std::string segment1 = "ssd_clear_segment_1";
+    const std::string segment2 = "ssd_clear_segment_2";
+    MountMemoryAndLocalDisk(*service, client1, segment1, 0x600000000);
+    MountMemoryAndLocalDisk(*service, client2, segment2, 0x700000000);
+
+    PutAndOffload(*service, client1, "ssd_clear_released", 800, segment1);
+    PutAndOffload(*service, client2, "ssd_clear_baseline", 100, segment2);
+
+    auto clear_result =
+        service->BatchReplicaClear({"ssd_clear_released"}, client1, "");
+    ASSERT_TRUE(clear_result.has_value());
+    ASSERT_EQ(clear_result->size(), 1u);
+    EXPECT_EQ((*clear_result)[0], "ssd_clear_released");
+
+    ExpectNextAllocationOnSegment(*service, client1, "ssd_clear_probe",
+                                  segment1);
+}
+
+// Test that after offloading more data to segment1, the next allocation prefers
+// segment2 which has more SSD free space.
+TEST_F(MasterServiceSSDTest, SsdFreeRatioFirstPrefersFresherSsdAfterOffload) {
+    auto service = CreateSsdAwareOffloadService();
+    UUID client1 = generate_uuid();
+    UUID client2 = generate_uuid();
+    const std::string segment1 = "ssd_fresher_seg_1";
+    const std::string segment2 = "ssd_fresher_seg_2";
+    // Each segment reports total SSD capacity = 1000 bytes
+    MountMemoryAndLocalDisk(*service, client1, segment1, 0x800000000);
+    MountMemoryAndLocalDisk(*service, client2, segment2, 0x900000000);
+
+    // Offload 800 bytes to segment1 → ssd_used[seg1]=800, free=20%
+    PutAndOffload(*service, client1, "ssd_fresher_heavy", 800, segment1);
+    // Offload 100 bytes to segment2 → ssd_used[seg2]=100, free=90%
+    PutAndOffload(*service, client2, "ssd_fresher_light", 100, segment2);
+
+    // segment2 has higher SSD free ratio → allocation should prefer segment2
+    ExpectNextAllocationOnSegment(*service, client2, "ssd_fresher_probe",
+                                  segment2);
+}
+
+// Test that EvictDiskReplica decrements ssd_used_bytes so that the evicted
+// segment becomes preferred again for the next allocation.
+TEST_F(MasterServiceSSDTest, EvictDiskReplicaDecrementsLocalDiskUsageTracking) {
+    auto service = CreateSsdAwareOffloadService();
+    UUID client1 = generate_uuid();
+    UUID client2 = generate_uuid();
+    const std::string segment1 = "ssd_evict_dec_seg_1";
+    const std::string segment2 = "ssd_evict_dec_seg_2";
+    MountMemoryAndLocalDisk(*service, client1, segment1, 0xa00000000);
+    MountMemoryAndLocalDisk(*service, client2, segment2, 0xb00000000);
+
+    // Offload 800 bytes to segment1 → ssd_used[seg1]=800 (20% free)
+    PutAndOffload(*service, client1, "ssd_evict_dec_heavy", 800, segment1);
+    // Offload 100 bytes to segment2 → ssd_used[seg2]=100 (90% free)
+    PutAndOffload(*service, client2, "ssd_evict_dec_light", 100, segment2);
+
+    // segment2 has more SSD free space → should be preferred
+    ExpectNextAllocationOnSegment(*service, client2, "ssd_evict_dec_probe1",
+                                  segment2);
+
+    // Evict the LOCAL_DISK replica of the heavy object from segment1.
+    // NotifyOffloadSuccess creates a LOCAL_DISK replica (not DISK).
+    // This decrements ssd_used[seg1] by 800 → ssd_used[seg1]=0 (100% free)
+    auto evict_result = service->EvictDiskReplica(
+        client1, "ssd_evict_dec_heavy", "default", ReplicaType::LOCAL_DISK);
+    ASSERT_TRUE(evict_result.has_value());
+
+    // After eviction: segment1 has 100% free, segment2 has 90% free
+    // → segment1 should now be preferred
+    ExpectNextAllocationOnSegment(*service, client1, "ssd_evict_dec_probe2",
+                                  segment1);
+}
+
+// Real-path performance comparison: MasterService PutStart throughput for
+// three configurations:
+//   (A) RANDOM, no offload        — baseline, original behavior
+//   (B) RANDOM, with offload      — isolates disk-replica creation overhead
+//   (C) SSD_FREE_RATIO_FIRST, with offload — adds SSD metrics lock + sorting
+//
+// Comparing A→B separates the cost of mounting LocalDisk segments.
+// Comparing B→C isolates the pure SSD-ranking strategy overhead.
+//
+// Each round: PutStart → PutEnd(MEMORY) (timed) → Remove (not timed).
+TEST_F(MasterServiceSSDTest,
+       SsdFreeRatioFirstVsRandomMasterServicePerformance) {
+    constexpr int kNumNodes = 32;
+    constexpr size_t kSegmentSize = 8 * 1024 * 1024;  // 8 MiB each
+    constexpr size_t kSliceSize = 512;  // 512 B – focus on strategy cost
+    constexpr int kWarmupRounds = 50;
+    constexpr int kBenchmarkRounds = 300;
+
+    // Build a MasterService with kNumNodes segments. with_ssd=true also
+    // mounts LocalDisk and reports varied SSD capacity per node.
+    auto buildAndMount =
+        [&](AllocationStrategyType strategy, bool with_ssd, size_t base_start,
+            const std::string& tag) -> std::unique_ptr<MasterService> {
+        MasterServiceConfig config;
+        config.enable_offload = with_ssd;
+        config.default_kv_lease_ttl = 10000;
+        config.allocation_strategy_type = strategy;
+        auto svc = std::make_unique<MasterService>(config);
+
+        for (int i = 0; i < kNumNodes; i++) {
+            UUID cid = generate_uuid();
+            Segment seg;
+            seg.id = generate_uuid();
+            seg.name = "ms_perf_" + std::to_string(i) + "_" + tag;
+            seg.base = base_start + static_cast<size_t>(i) * kSegmentSize;
+            seg.size = kSegmentSize;
+            seg.te_endpoint = seg.name;
+            (void)svc->MountSegment(seg, cid);
+            if (with_ssd) {
+                (void)svc->MountLocalDiskSegment(cid, true);
+                // Vary total SSD capacity so nodes have distinct free ratios
+                (void)svc->ReportSsdCapacity(
+                    cid, static_cast<int64_t>(1024 * 1024) * (i + 1));
+            }
+        }
+        return svc;
+    };
+
+    // Measure kRounds of PutStart + PutEnd(MEMORY). Remove is called after
+    // timing to free allocator space without inflating the measurement.
+    auto runBenchmark = [&](MasterService& svc, const std::string& key_pfx,
+                            int rounds) -> std::chrono::microseconds {
+        const UUID writer = generate_uuid();
+        ReplicateConfig cfg;
+        cfg.replica_num = 1;
+        std::chrono::microseconds total{0};
+
+        for (int i = 0; i < rounds; i++) {
+            const std::string key = key_pfx + std::to_string(i);
+            auto t0 = std::chrono::steady_clock::now();
+            (void)svc.PutStart(writer, key, "default", kSliceSize, cfg);
+            (void)svc.PutEnd(writer, key, "default", ReplicaType::MEMORY);
+            total += std::chrono::duration_cast<std::chrono::microseconds>(
+                std::chrono::steady_clock::now() - t0);
+            (void)svc.Remove(key, "default", /*force=*/true);
+        }
+        return total;
+    };
+
+    // (A) RANDOM, no offload – baseline
+    auto svc_a = buildAndMount(AllocationStrategyType::RANDOM, false,
+                               0xc00000000ULL, "A");
+    (void)runBenchmark(*svc_a, "ms_A_wu_", kWarmupRounds);
+    auto elapsed_a = runBenchmark(*svc_a, "ms_A_bm_", kBenchmarkRounds);
+
+    // (B) RANDOM, with offload – quantify disk-replica overhead alone
+    auto svc_b = buildAndMount(AllocationStrategyType::RANDOM, true,
+                               0xd00000000ULL, "B");
+    (void)runBenchmark(*svc_b, "ms_B_wu_", kWarmupRounds);
+    auto elapsed_b = runBenchmark(*svc_b, "ms_B_bm_", kBenchmarkRounds);
+
+    // (C) SSD_FREE_RATIO_FIRST, with offload – full new feature
+    auto svc_c = buildAndMount(AllocationStrategyType::SSD_FREE_RATIO_FIRST,
+                               true, 0xe00000000ULL, "C");
+    (void)runBenchmark(*svc_c, "ms_C_wu_", kWarmupRounds);
+    auto elapsed_c = runBenchmark(*svc_c, "ms_C_bm_", kBenchmarkRounds);
+
+    auto us_per_op = [&](std::chrono::microseconds us) {
+        return static_cast<double>(us.count()) / kBenchmarkRounds;
+    };
+    double ratio_b_a =
+        static_cast<double>(elapsed_b.count()) / elapsed_a.count();
+    double ratio_c_b =
+        static_cast<double>(elapsed_c.count()) / elapsed_b.count();
+    double ratio_c_a =
+        static_cast<double>(elapsed_c.count()) / elapsed_a.count();
+
+    std::cout
+        << "\n=== MasterService Real-Path Performance (PutStart+PutEnd) ===\n"
+        << "Nodes: " << kNumNodes << " | Slice: " << kSliceSize
+        << " B | Rounds: " << kBenchmarkRounds << "\n\n"
+        << "  (A) RANDOM, offload=OFF (baseline):         " << elapsed_a.count()
+        << " us  |  " << std::fixed << std::setprecision(3)
+        << us_per_op(elapsed_a) << " us/op\n"
+        << "  (B) RANDOM, offload=ON  (disk replica cost):" << elapsed_b.count()
+        << " us  |  " << us_per_op(elapsed_b) << " us/op  ["
+        << std::setprecision(2) << ratio_b_a << "x vs A]\n"
+        << "  (C) SSD_FREE_RATIO_FIRST, offload=ON:       " << elapsed_c.count()
+        << " us  |  " << us_per_op(elapsed_c) << " us/op  [" << ratio_c_b
+        << "x vs B]\n\n"
+        << "  A→B  disk-replica overhead:   " << std::setprecision(1)
+        << (ratio_b_a - 1.0) * 100.0 << "%\n"
+        << "  B→C  SSD-ranking overhead:    " << (ratio_c_b - 1.0) * 100.0
+        << "%\n"
+        << "  A→C  total overhead vs origin:" << (ratio_c_a - 1.0) * 100.0
+        << "%\n\n";
+}
+
 }  // namespace mooncake::test
 
 int main(int argc, char** argv) {
diff --git a/mooncake-store/tests/master_service_tenant_quota_test.cpp b/mooncake-store/tests/master_service_tenant_quota_test.cpp
new file mode 100644
index 00000000..7a53b845
--- /dev/null
+++ b/mooncake-store/tests/master_service_tenant_quota_test.cpp
@@ -0,0 +1,851 @@
+#include "master_service.h"
+
+#include <atomic>
+#include <chrono>
+#include <filesystem>
+#include <future>
+#include <fstream>
+#include <map>
+#include <memory>
+#include <mutex>
+#include <optional>
+#include <shared_mutex>
+#include <string>
+#include <thread>
+#include <vector>
+
+#include <gtest/gtest.h>
+#include <unistd.h>
+
+#include "allocation_strategy.h"
+#include "tenant_quota_policy_store.h"
+#include "types.h"
+
+namespace mooncake::test {
+
+class BlockingTenantQuotaPolicyStore final : public TenantQuotaPolicyStore {
+   public:
+    explicit BlockingTenantQuotaPolicyStore(TenantQuotaPolicySnapshot snapshot)
+        : snapshot_(std::move(snapshot)),
+          allow_save_(allow_save_promise_.get_future()) {}
+
+    std::future<void> SaveStarted() {
+        return save_started_promise_.get_future();
+    }
+
+    void AllowSave() { allow_save_promise_.set_value(); }
+
+    tl::expected<TenantQuotaPolicySnapshot, std::string> Load() override {
+        return snapshot_;
+    }
+
+    tl::expected<void, std::string> Save(
+        const TenantQuotaPolicySnapshot& snapshot) override {
+        snapshot_ = snapshot;
+        save_started_promise_.set_value();
+        allow_save_.wait();
+        return {};
+    }
+
+   private:
+    TenantQuotaPolicySnapshot snapshot_;
+    std::promise<void> save_started_promise_;
+    std::promise<void> allow_save_promise_;
+    std::future<void> allow_save_;
+};
+
+#ifdef USE_NOF
+class BlockingAllocationStrategy final : public AllocationStrategy {
+   public:
+    BlockingAllocationStrategy()
+        : allow_allocation_(allow_allocation_promise_.get_future()) {}
+
+    std::future<void> AllocationStarted() {
+        return allocation_started_promise_.get_future();
+    }
+
+    void AllowAllocation() { allow_allocation_promise_.set_value(); }
+
+    tl::expected<std::vector<Replica>, ErrorCode> Allocate(
+        const AllocatorManager& allocator_manager, const size_t slice_length,
+        const size_t replica_num,
+        const std::vector<std::string>& preferred_segments,
+        const std::set<std::string>& excluded_segments,
+        const ReplicaType replica_type) override {
+        BlockOnce();
+        return delegate_.Allocate(allocator_manager, slice_length, replica_num,
+                                  preferred_segments, excluded_segments,
+                                  replica_type);
+    }
+
+    tl::expected<std::vector<Replica>, ErrorCode> Allocate(
+        const AllocatorManager& allocator_manager, const size_t slice_length,
+        const size_t replica_num,
+        const std::vector<std::string>& preferred_segments,
+        const std::set<std::string>& excluded_segments,
+        const ReplicaType replica_type,
+        const SsdMetricsProvider* ssd_provider) override {
+        (void)ssd_provider;
+        return Allocate(allocator_manager, slice_length, replica_num,
+                        preferred_segments, excluded_segments, replica_type);
+    }
+
+    tl::expected<Replica, ErrorCode> AllocateFrom(
+        const AllocatorManager& allocator_manager, const size_t slice_length,
+        const std::string& segment_name) override {
+        return delegate_.AllocateFrom(allocator_manager, slice_length,
+                                      segment_name);
+    }
+
+   private:
+    void BlockOnce() {
+        bool expected = true;
+        if (block_next_allocation_.compare_exchange_strong(expected, false)) {
+            allocation_started_promise_.set_value();
+            allow_allocation_.wait();
+        }
+    }
+
+    RandomAllocationStrategy delegate_;
+    std::atomic<bool> block_next_allocation_{true};
+    std::promise<void> allocation_started_promise_;
+    std::promise<void> allow_allocation_promise_;
+    std::future<void> allow_allocation_;
+};
+#endif
+
+class MasterServiceTenantQuotaTest : public ::testing::Test {
+   protected:
+    static constexpr size_t kSegmentBase = 0x500000000;
+
+    std::string WritePolicyFile(
+        const std::map<std::string, uint64_t>& tenant_quotas) {
+        TenantQuotaPolicySnapshot snapshot;
+        snapshot.tenant_quotas = tenant_quotas;
+        auto path =
+            std::filesystem::temp_directory_path() /
+            ("mooncake_tenant_quota_test_" + std::to_string(::getpid()) + "_" +
+             std::to_string(next_policy_file_++) + ".yaml");
+        std::ofstream out(path);
+        out << FormatTenantQuotaPolicyYaml(snapshot);
+        out.close();
+        policy_files_.push_back(path.string());
+        return path.string();
+    }
+
+    MasterServiceConfig MakeConfig(
+        const std::map<std::string, uint64_t>& tenant_quotas,
+        bool enable_multi_tenants = true) {
+        auto builder = MasterServiceConfig::builder().set_enable_multi_tenants(
+            enable_multi_tenants);
+        if (enable_multi_tenants) {
+            builder.set_tenant_quota_connector_type("file")
+                .set_tenant_quota_connector_uri(WritePolicyFile(tenant_quotas));
+        }
+        return builder.build();
+    }
+
+    UUID MountSegment(MasterService& service, size_t size = 4096,
+                      std::string name = "quota_segment") {
+        Segment segment;
+        segment.id = generate_uuid();
+        segment.name = std::move(name);
+        segment.base = kSegmentBase + next_segment_offset_;
+        segment.size = size;
+        segment.te_endpoint = segment.name;
+        next_segment_offset_ += size + 4096;
+
+        UUID client_id = generate_uuid();
+        auto result = service.MountSegment(segment, client_id);
+        EXPECT_TRUE(result.has_value()) << toString(result.error());
+        return client_id;
+    }
+
+#ifdef USE_NOF
+    UUID MountNoFSegment(MasterService& service, size_t size = 4096,
+                         std::string name = "quota_nof_segment") {
+        NoFSegment segment;
+        segment.id = generate_uuid();
+        segment.name = std::move(name);
+        segment.base = kSegmentBase + next_segment_offset_;
+        segment.size = size;
+        segment.te_endpoint = segment.name;
+        next_segment_offset_ += size + 4096;
+
+        UUID client_id = generate_uuid();
+        auto result = service.MountNoFSegment(segment, client_id);
+        EXPECT_TRUE(result.has_value()) << toString(result.error());
+        return client_id;
+    }
+#endif
+
+    ReplicateConfig MemoryConfig() {
+        ReplicateConfig config;
+        config.replica_num = 1;
+        return config;
+    }
+
+    void PutComplete(MasterService& service, const UUID& client_id,
+                     const std::string& key, const std::string& tenant_id,
+                     uint64_t size) {
+        auto start =
+            service.PutStart(client_id, key, tenant_id, size, MemoryConfig());
+        ASSERT_TRUE(start.has_value()) << toString(start.error());
+        auto end =
+            service.PutEnd(client_id, key, tenant_id, ReplicaType::MEMORY);
+        ASSERT_TRUE(end.has_value()) << toString(end.error());
+    }
+
+    TenantQuotaSnapshot Snapshot(MasterService& service,
+                                 const std::string& tenant_id) {
+        auto snapshot = service.GetTenantQuotaSnapshotForTesting(tenant_id);
+        EXPECT_TRUE(snapshot.has_value());
+        return *snapshot;
+    }
+
+    void ReloadTenantQuotaPolicyFromStore(MasterService& service) {
+        service.LoadTenantQuotaPoliciesFromStoreOrThrow();
+        service.RebuildTenantQuotaUsageFromMetadata();
+    }
+
+    void ReplaceTenantQuotaPolicyStore(
+        MasterService& service, std::unique_ptr<TenantQuotaPolicyStore> store) {
+        service.tenant_quota_policy_store_ = std::move(store);
+    }
+
+    int64_t LocalDiskUsedBytes(MasterService& service, const UUID& client_id) {
+        auto access = service.segment_manager_.getLocalDiskSegmentAccess();
+        auto& segments = access.getClientLocalDiskSegment();
+        auto it = segments.find(client_id);
+        EXPECT_TRUE(it != segments.end());
+        if (it == segments.end()) {
+            return -1;
+        }
+        return it->second->ssd_used_bytes.load(std::memory_order_relaxed);
+    }
+
+#ifdef USE_NOF
+    void ReplaceAllocationStrategy(
+        MasterService& service, std::shared_ptr<AllocationStrategy> strategy) {
+        service.allocation_strategy_ = std::move(strategy);
+    }
+#endif
+
+    tl::expected<void, ErrorCode> ReserveTenantQuotaForTest(
+        MasterService& service, const std::string& tenant_id, uint64_t bytes) {
+        return service.ReserveTenantQuota(tenant_id, bytes);
+    }
+
+    std::unique_lock<std::shared_mutex> LockSnapshotForTest(
+        MasterService& service) {
+        return std::unique_lock<std::shared_mutex>(service.snapshot_mutex_);
+    }
+
+    bool WaitForTenantQuotaPolicyMutexContention(MasterService& service) {
+        for (int i = 0; i < 500; ++i) {
+            if (!service.tenant_quota_policy_mutex_.try_lock()) {
+                return true;
+            }
+            service.tenant_quota_policy_mutex_.unlock();
+            std::this_thread::sleep_for(std::chrono::milliseconds(10));
+        }
+        return false;
+    }
+
+    void TearDown() override {
+        for (const auto& path : policy_files_) {
+            std::error_code ec;
+            std::filesystem::remove(path, ec);
+        }
+    }
+
+    size_t next_segment_offset_ = 0;
+    size_t next_policy_file_ = 0;
+    std::vector<std::string> policy_files_;
+};
+
+TEST_F(MasterServiceTenantQuotaTest,
+       SingleTenantModeCollapsesTenantsAndDisablesQuota) {
+    MasterService service(MakeConfig({}, /*enable_multi_tenants=*/false));
+    UUID client_id = MountSegment(service, /*size=*/1024);
+
+    PutComplete(service, client_id, "shared-key", "tenant-a", 800);
+
+    EXPECT_TRUE(service.ExistKey("shared-key", "tenant-b").value());
+    auto duplicate = service.PutStart(client_id, "shared-key", "tenant-b", 1,
+                                      MemoryConfig());
+    ASSERT_FALSE(duplicate.has_value());
+    EXPECT_EQ(duplicate.error(), ErrorCode::OBJECT_ALREADY_EXISTS);
+    EXPECT_TRUE(service
+                    .Remove("shared-key", "tenant-b",
+                            /*force=*/true)
+                    .has_value());
+    EXPECT_FALSE(
+        service.GetTenantQuotaSnapshotForTesting("tenant-a").has_value());
+}
+
+TEST_F(MasterServiceTenantQuotaTest,
+       MultiTenantModeRejectsEmptyUnregisteredAndImplicitDefaultWrites) {
+    MasterService service(MakeConfig({{"tenant-a", 1000}}));
+    UUID client_id = MountSegment(service);
+
+    auto empty = service.PutStart(client_id, "empty", "", 10, MemoryConfig());
+    ASSERT_FALSE(empty.has_value());
+    EXPECT_EQ(empty.error(), ErrorCode::TENANT_NOT_REGISTERED);
+
+    auto missing =
+        service.PutStart(client_id, "missing", "tenant-b", 10, MemoryConfig());
+    ASSERT_FALSE(missing.has_value());
+    EXPECT_EQ(missing.error(), ErrorCode::TENANT_NOT_REGISTERED);
+
+    auto implicit_default = service.PutStart(client_id, "default-key",
+                                             "default", 10, MemoryConfig());
+    ASSERT_FALSE(implicit_default.has_value());
+    EXPECT_EQ(implicit_default.error(), ErrorCode::TENANT_NOT_REGISTERED);
+
+    const std::string control_tenant("tenant\0bad", 10);
+    auto control = service.PutStart(client_id, "control-key", control_tenant,
+                                    10, MemoryConfig());
+    ASSERT_FALSE(control.has_value());
+    EXPECT_EQ(control.error(), ErrorCode::TENANT_NOT_REGISTERED);
+
+    auto control_policy = service.UpsertTenantQuotaPolicy(control_tenant, 100);
+    ASSERT_FALSE(control_policy.has_value());
+    EXPECT_EQ(control_policy.error(), ErrorCode::INVALID_PARAMS);
+
+    auto register_default = service.UpsertTenantQuotaPolicy("default", 100);
+    ASSERT_TRUE(register_default.has_value())
+        << toString(register_default.error());
+    PutComplete(service, client_id, "registered-default", "default", 10);
+
+    PutComplete(service, client_id, "ok", "tenant-a", 10);
+}
+
+TEST_F(MasterServiceTenantQuotaTest,
+       MultiTenantModeRejectsUnregisteredOffloadSuccess) {
+    MasterService service(MakeConfig({{"tenant-a", 1000}}));
+    UUID client_id = MountSegment(service);
+
+    StorageObjectMetadata metadata;
+    metadata.data_size = 128;
+    metadata.transport_endpoint = "disk-endpoint";
+    std::vector<OffloadTaskItem> tasks{
+        OffloadTaskItem{.tenant_id = "tenant-b", .key = "ghost", .size = 128}};
+
+    auto result = service.NotifyOffloadSuccess(client_id, tasks, {metadata});
+
+    ASSERT_FALSE(result.has_value());
+    EXPECT_EQ(result.error(), ErrorCode::TENANT_NOT_REGISTERED);
+    auto missing = service.ExistKey("ghost", "tenant-b");
+    ASSERT_TRUE(missing.has_value()) << toString(missing.error());
+    EXPECT_FALSE(missing.value());
+}
+
+TEST_F(MasterServiceTenantQuotaTest,
+       MultiTenantModeAllowsRegisteredOffloadSuccess) {
+    MasterService service(MakeConfig({{"tenant-a", 1000}}));
+    UUID client_id = MountSegment(service);
+
+    StorageObjectMetadata metadata;
+    metadata.data_size = 128;
+    metadata.transport_endpoint = "disk-endpoint";
+    std::vector<OffloadTaskItem> tasks{
+        OffloadTaskItem{.tenant_id = "tenant-a", .key = "cold", .size = 128}};
+
+    auto result = service.NotifyOffloadSuccess(client_id, tasks, {metadata});
+
+    ASSERT_TRUE(result.has_value()) << toString(result.error());
+    auto exists = service.ExistKey("cold", "tenant-a");
+    ASSERT_TRUE(exists.has_value()) << toString(exists.error());
+    EXPECT_TRUE(exists.value());
+    EXPECT_EQ(Snapshot(service, "tenant-a").used_bytes, 0);
+}
+
+TEST_F(MasterServiceTenantQuotaTest,
+       ConnectorPolicyReloadKeepsLocalDiskOnlyOrphanVisible) {
+    const std::string initial_policy =
+        WritePolicyFile({{"tenant-a", 1000}, {"tenant-b", 1000}});
+    auto config = MasterServiceConfig::builder()
+                      .set_enable_multi_tenants(true)
+                      .set_tenant_quota_connector_type("file")
+                      .set_tenant_quota_connector_uri(initial_policy)
+                      .build();
+    MasterService service(config);
+    UUID client_id = MountSegment(service);
+
+    StorageObjectMetadata metadata;
+    metadata.data_size = 128;
+    metadata.transport_endpoint = "disk-endpoint";
+    std::vector<OffloadTaskItem> tasks{
+        OffloadTaskItem{.tenant_id = "tenant-b", .key = "cold", .size = 128}};
+    ASSERT_TRUE(
+        service.NotifyOffloadSuccess(client_id, tasks, {metadata}).has_value());
+
+    {
+        std::ofstream out(initial_policy);
+        TenantQuotaPolicySnapshot replacement;
+        replacement.tenant_quotas = {{"tenant-a", 1000}};
+        out << FormatTenantQuotaPolicyYaml(replacement);
+    }
+    ReloadTenantQuotaPolicyFromStore(service);
+
+    auto orphan = Snapshot(service, "tenant-b");
+    EXPECT_FALSE(orphan.has_explicit_policy);
+    EXPECT_EQ(orphan.used_bytes, 0);
+    EXPECT_EQ(orphan.committed_count, 0);
+    EXPECT_EQ(orphan.metadata_object_count, 1);
+    EXPECT_TRUE(orphan.over_quota);
+
+    EXPECT_TRUE(service.Remove("cold", "tenant-b", /*force=*/true).has_value());
+    EXPECT_FALSE(
+        service.GetTenantQuotaSnapshotForTesting("tenant-b").has_value());
+}
+
+TEST_F(MasterServiceTenantQuotaTest,
+       NotifyOffloadSuccessCompletesExistingOrphanObject) {
+    const std::string initial_policy =
+        WritePolicyFile({{"tenant-a", 1000}, {"tenant-b", 1000}});
+    auto config = MasterServiceConfig::builder()
+                      .set_enable_multi_tenants(true)
+                      .set_enable_offload(true)
+                      .set_tenant_quota_connector_type("file")
+                      .set_tenant_quota_connector_uri(initial_policy)
+                      .build();
+    MasterService service(config);
+    UUID client_id = MountSegment(service);
+    ASSERT_TRUE(service.MountLocalDiskSegment(client_id, true).has_value());
+    PutComplete(service, client_id, "warming", "tenant-b", 128);
+
+    {
+        std::ofstream out(initial_policy);
+        TenantQuotaPolicySnapshot replacement;
+        replacement.tenant_quotas = {{"tenant-a", 1000}};
+        out << FormatTenantQuotaPolicyYaml(replacement);
+    }
+    ReloadTenantQuotaPolicyFromStore(service);
+    EXPECT_FALSE(Snapshot(service, "tenant-b").has_explicit_policy);
+
+    StorageObjectMetadata metadata;
+    metadata.data_size = 128;
+    metadata.transport_endpoint = "disk-endpoint";
+    std::vector<OffloadTaskItem> tasks{OffloadTaskItem{
+        .tenant_id = "tenant-b", .key = "warming", .size = 128}};
+
+    auto result = service.NotifyOffloadSuccess(client_id, tasks, {metadata});
+
+    ASSERT_TRUE(result.has_value()) << toString(result.error());
+    auto replicas = service.GetReplicaList("warming", "tenant-b");
+    ASSERT_TRUE(replicas.has_value()) << toString(replicas.error());
+    EXPECT_TRUE(std::any_of(replicas->replicas.begin(),
+                            replicas->replicas.end(),
+                            [](const Replica::Descriptor& replica) {
+                                return replica.is_local_disk_replica();
+                            }));
+}
+
+TEST_F(MasterServiceTenantQuotaTest,
+       NotifyOffloadSuccessRejectsOrphanObjectWithoutOffloadTask) {
+    const std::string initial_policy =
+        WritePolicyFile({{"tenant-a", 1000}, {"tenant-b", 1000}});
+    auto config = MasterServiceConfig::builder()
+                      .set_enable_multi_tenants(true)
+                      .set_tenant_quota_connector_type("file")
+                      .set_tenant_quota_connector_uri(initial_policy)
+                      .build();
+    MasterService service(config);
+    UUID client_id = MountSegment(service);
+    PutComplete(service, client_id, "warming", "tenant-b", 128);
+
+    {
+        std::ofstream out(initial_policy);
+        TenantQuotaPolicySnapshot replacement;
+        replacement.tenant_quotas = {{"tenant-a", 1000}};
+        out << FormatTenantQuotaPolicyYaml(replacement);
+    }
+    ReloadTenantQuotaPolicyFromStore(service);
+    EXPECT_FALSE(Snapshot(service, "tenant-b").has_explicit_policy);
+
+    StorageObjectMetadata metadata;
+    metadata.data_size = 128;
+    metadata.transport_endpoint = "disk-endpoint";
+    std::vector<OffloadTaskItem> tasks{OffloadTaskItem{
+        .tenant_id = "tenant-b", .key = "warming", .size = 128}};
+
+    auto result = service.NotifyOffloadSuccess(client_id, tasks, {metadata});
+
+    ASSERT_FALSE(result.has_value());
+    EXPECT_EQ(result.error(), ErrorCode::TENANT_NOT_REGISTERED);
+}
+
+TEST_F(MasterServiceTenantQuotaTest,
+       NotifyOffloadSuccessDoesNotCountAddReplicaUpdateAsNewDiskUsage) {
+    const std::string policy = WritePolicyFile({{"tenant-a", 1000}});
+    auto config = MasterServiceConfig::builder()
+                      .set_enable_multi_tenants(true)
+                      .set_enable_offload(true)
+                      .set_tenant_quota_connector_type("file")
+                      .set_tenant_quota_connector_uri(policy)
+                      .build();
+    MasterService service(config);
+    UUID client_a = MountSegment(service, 4096, "quota_segment_a");
+    UUID client_b = MountSegment(service, 4096, "quota_segment_b");
+    ASSERT_TRUE(service.MountLocalDiskSegment(client_a, true).has_value());
+    ASSERT_TRUE(service.MountLocalDiskSegment(client_b, true).has_value());
+
+    StorageObjectMetadata first_metadata;
+    first_metadata.data_size = 128;
+    first_metadata.transport_endpoint = "disk-endpoint-a";
+    std::vector<OffloadTaskItem> tasks{
+        OffloadTaskItem{.tenant_id = "tenant-a", .key = "cold", .size = 128}};
+    ASSERT_TRUE(service.NotifyOffloadSuccess(client_a, tasks, {first_metadata})
+                    .has_value());
+    EXPECT_EQ(LocalDiskUsedBytes(service, client_a), 128);
+    EXPECT_EQ(LocalDiskUsedBytes(service, client_b), 0);
+
+    StorageObjectMetadata second_metadata;
+    second_metadata.data_size = 128;
+    second_metadata.transport_endpoint = "disk-endpoint-b";
+    auto result =
+        service.NotifyOffloadSuccess(client_b, tasks, {second_metadata});
+
+    ASSERT_TRUE(result.has_value()) << toString(result.error());
+    EXPECT_EQ(LocalDiskUsedBytes(service, client_a), 128);
+    EXPECT_EQ(LocalDiskUsedBytes(service, client_b), 0);
+}
+
+TEST_F(MasterServiceTenantQuotaTest,
+       RegisteredTenantQuotaAdmissionDoesNotCreateImplicitTenants) {
+    MasterService service(MakeConfig({{"tenant-a", 100}}));
+    UUID client_id = MountSegment(service);
+
+    auto hard_pinned = MemoryConfig();
+    hard_pinned.with_hard_pin = true;
+    auto first =
+        service.PutStart(client_id, "key-a", "tenant-a", 80, hard_pinned);
+    ASSERT_TRUE(first.has_value()) << toString(first.error());
+    ASSERT_TRUE(
+        service.PutEnd(client_id, "key-a", "tenant-a", ReplicaType::MEMORY)
+            .has_value());
+
+    auto over =
+        service.PutStart(client_id, "key-b", "tenant-a", 30, MemoryConfig());
+
+    ASSERT_FALSE(over.has_value());
+    EXPECT_EQ(over.error(), ErrorCode::TENANT_QUOTA_EXCEEDED);
+    EXPECT_EQ(Snapshot(service, "tenant-a").used_bytes, 80);
+    EXPECT_FALSE(
+        service.GetTenantQuotaSnapshotForTesting("tenant-b").has_value());
+}
+
+TEST_F(MasterServiceTenantQuotaTest, CopyStartRequiresQuotaForNewReplica) {
+    MasterService service(MakeConfig({{"tenant-a", 150}}));
+    UUID client_id = MountSegment(service, /*size=*/1024, "segment-a");
+    MountSegment(service, /*size=*/1024, "segment-b");
+
+    ReplicateConfig config = MemoryConfig();
+    config.preferred_segment = "segment-a";
+    auto put_start =
+        service.PutStart(client_id, "key", "tenant-a", 100, config);
+    ASSERT_TRUE(put_start.has_value()) << toString(put_start.error());
+    ASSERT_TRUE(
+        service.PutEnd(client_id, "key", "tenant-a", ReplicaType::MEMORY)
+            .has_value());
+
+    auto copy = service.CopyStart(client_id, "key", "tenant-a", "segment-a",
+                                  {"segment-b"});
+
+    ASSERT_FALSE(copy.has_value());
+    EXPECT_EQ(copy.error(), ErrorCode::TENANT_QUOTA_EXCEEDED);
+    auto snapshot = Snapshot(service, "tenant-a");
+    EXPECT_EQ(snapshot.used_bytes, 100);
+    EXPECT_EQ(snapshot.reserved_bytes, 0);
+    EXPECT_EQ(snapshot.committed_count, 1);
+}
+
+TEST_F(MasterServiceTenantQuotaTest,
+       CopyEndCommitsAdditionalReplicaWithoutExtraObjectCount) {
+    MasterService service(MakeConfig({{"tenant-a", 300}}));
+    UUID client_id = MountSegment(service, /*size=*/1024, "segment-a");
+    MountSegment(service, /*size=*/1024, "segment-b");
+
+    ReplicateConfig config = MemoryConfig();
+    config.preferred_segment = "segment-a";
+    auto put_start =
+        service.PutStart(client_id, "key", "tenant-a", 100, config);
+    ASSERT_TRUE(put_start.has_value()) << toString(put_start.error());
+    ASSERT_TRUE(
+        service.PutEnd(client_id, "key", "tenant-a", ReplicaType::MEMORY)
+            .has_value());
+
+    auto copy = service.CopyStart(client_id, "key", "tenant-a", "segment-a",
+                                  {"segment-b"});
+    ASSERT_TRUE(copy.has_value()) << toString(copy.error());
+    auto in_flight = Snapshot(service, "tenant-a");
+    EXPECT_EQ(in_flight.used_bytes, 100);
+    EXPECT_EQ(in_flight.reserved_bytes, 100);
+
+    ASSERT_TRUE(service.CopyEnd(client_id, "key", "tenant-a").has_value());
+    auto completed = Snapshot(service, "tenant-a");
+    EXPECT_EQ(completed.used_bytes, 200);
+    EXPECT_EQ(completed.reserved_bytes, 0);
+    EXPECT_EQ(completed.committed_count, 1);
+    EXPECT_EQ(completed.metadata_object_count, 1);
+}
+
+TEST_F(MasterServiceTenantQuotaTest,
+       MoveStartRequiresQuotaForTemporaryReplica) {
+    MasterService service(MakeConfig({{"tenant-a", 150}}));
+    UUID client_id = MountSegment(service, /*size=*/1024, "segment-a");
+    MountSegment(service, /*size=*/1024, "segment-b");
+
+    ReplicateConfig config = MemoryConfig();
+    config.preferred_segment = "segment-a";
+    auto put_start =
+        service.PutStart(client_id, "key", "tenant-a", 100, config);
+    ASSERT_TRUE(put_start.has_value()) << toString(put_start.error());
+    ASSERT_TRUE(
+        service.PutEnd(client_id, "key", "tenant-a", ReplicaType::MEMORY)
+            .has_value());
+
+    auto move = service.MoveStart(client_id, "key", "tenant-a", "segment-a",
+                                  "segment-b");
+
+    ASSERT_FALSE(move.has_value());
+    EXPECT_EQ(move.error(), ErrorCode::TENANT_QUOTA_EXCEEDED);
+    auto snapshot = Snapshot(service, "tenant-a");
+    EXPECT_EQ(snapshot.used_bytes, 100);
+    EXPECT_EQ(snapshot.reserved_bytes, 0);
+}
+
+TEST_F(MasterServiceTenantQuotaTest, AdminDeleteRequiresEmptyTenant) {
+    MasterService service(MakeConfig({{"tenant-a", 1000}}));
+    UUID client_id = MountSegment(service);
+    PutComplete(service, client_id, "key", "tenant-a", 100);
+
+    auto empty_upsert = service.UpsertTenantQuotaPolicy("", 100);
+    ASSERT_FALSE(empty_upsert.has_value());
+    EXPECT_EQ(empty_upsert.error(), ErrorCode::INVALID_PARAMS);
+
+    auto empty_delete = service.DeleteTenantQuotaPolicy("");
+    ASSERT_FALSE(empty_delete.has_value());
+    EXPECT_EQ(empty_delete.error(), ErrorCode::INVALID_PARAMS);
+
+    auto delete_non_empty = service.DeleteTenantQuotaPolicy("tenant-a");
+    ASSERT_FALSE(delete_non_empty.has_value());
+    EXPECT_EQ(delete_non_empty.error(), ErrorCode::TENANT_NOT_EMPTY);
+
+    auto upsert = service.UpsertTenantQuotaPolicy("tenant-b", 100);
+    ASSERT_TRUE(upsert.has_value()) << toString(upsert.error());
+    auto delete_empty = service.DeleteTenantQuotaPolicy("tenant-b");
+    ASSERT_TRUE(delete_empty.has_value()) << toString(delete_empty.error());
+    EXPECT_FALSE(delete_empty.value().has_value());
+}
+
+TEST_F(MasterServiceTenantQuotaTest,
+       DeletePolicyBlocksValidatedReservationsBeforeConnectorSave) {
+    MasterService service(MakeConfig({{"tenant-a", 1000}}));
+    MountSegment(service);
+
+    TenantQuotaPolicySnapshot current_policy;
+    current_policy.tenant_quotas = {{"tenant-a", 1000}};
+    auto blocking_store =
+        std::make_unique<BlockingTenantQuotaPolicyStore>(current_policy);
+    auto* blocking_store_ptr = blocking_store.get();
+    auto save_started = blocking_store_ptr->SaveStarted();
+    ReplaceTenantQuotaPolicyStore(service, std::move(blocking_store));
+
+    using DeleteResult =
+        tl::expected<std::optional<TenantQuotaSnapshot>, ErrorCode>;
+    std::optional<DeleteResult> delete_result;
+    std::thread delete_thread([&] {
+        delete_result.emplace(service.DeleteTenantQuotaPolicy("tenant-a"));
+    });
+
+    if (save_started.wait_for(std::chrono::seconds(5)) !=
+        std::future_status::ready) {
+        blocking_store_ptr->AllowSave();
+        delete_thread.join();
+        FAIL() << "timed out waiting for connector save";
+    }
+
+    auto reserve = ReserveTenantQuotaForTest(service, "tenant-a", 1);
+    EXPECT_FALSE(reserve.has_value());
+    EXPECT_EQ(reserve.error(), ErrorCode::TENANT_NOT_REGISTERED);
+
+    auto zero_byte_reserve = ReserveTenantQuotaForTest(service, "tenant-a", 0);
+    EXPECT_FALSE(zero_byte_reserve.has_value());
+    EXPECT_EQ(zero_byte_reserve.error(), ErrorCode::TENANT_NOT_REGISTERED);
+
+    blocking_store_ptr->AllowSave();
+    delete_thread.join();
+
+    ASSERT_TRUE(delete_result.has_value());
+    ASSERT_TRUE(delete_result->has_value()) << toString(delete_result->error());
+    EXPECT_FALSE(delete_result->value().has_value());
+    EXPECT_FALSE(
+        service.GetTenantQuotaSnapshotForTesting("tenant-a").has_value());
+}
+
+TEST_F(MasterServiceTenantQuotaTest,
+       DeletePolicyWaitsForInFlightAddReplicaBeforeEmptyCheck) {
+    MasterService service(MakeConfig({{"tenant-a", 1000}}));
+    UUID client_id = MountSegment(service);
+
+    TenantQuotaPolicySnapshot current_policy;
+    current_policy.tenant_quotas = {{"tenant-a", 1000}};
+    auto blocking_store =
+        std::make_unique<BlockingTenantQuotaPolicyStore>(current_policy);
+    auto* blocking_store_ptr = blocking_store.get();
+    auto save_started = blocking_store_ptr->SaveStarted();
+    ReplaceTenantQuotaPolicyStore(service, std::move(blocking_store));
+
+    auto snapshot_lock = LockSnapshotForTest(service);
+    std::optional<tl::expected<bool, ErrorCode>> add_result;
+    std::thread add_thread([&] {
+        Replica replica(client_id, 128, "disk-endpoint",
+                        ReplicaStatus::COMPLETE);
+        add_result.emplace(
+            service.AddReplica(client_id, "cold", "tenant-a", replica));
+    });
+
+    if (!WaitForTenantQuotaPolicyMutexContention(service)) {
+        snapshot_lock.unlock();
+        add_thread.join();
+        FAIL() << "timed out waiting for AddReplica to enter tenant policy "
+                  "critical section";
+    }
+
+    using DeleteResult =
+        tl::expected<std::optional<TenantQuotaSnapshot>, ErrorCode>;
+    std::optional<DeleteResult> delete_result;
+    std::thread delete_thread([&] {
+        delete_result.emplace(service.DeleteTenantQuotaPolicy("tenant-a"));
+    });
+
+    const auto premature_save =
+        save_started.wait_for(std::chrono::milliseconds(200));
+    if (premature_save == std::future_status::ready) {
+        blocking_store_ptr->AllowSave();
+    }
+    snapshot_lock.unlock();
+    add_thread.join();
+    delete_thread.join();
+
+    ASSERT_EQ(premature_save, std::future_status::timeout)
+        << "tenant deletion reached connector save before in-flight "
+           "AddReplica completed";
+    ASSERT_TRUE(add_result.has_value());
+    ASSERT_TRUE(add_result->has_value()) << toString(add_result->error());
+    ASSERT_TRUE(delete_result.has_value());
+    ASSERT_FALSE(delete_result->has_value());
+    EXPECT_EQ(delete_result->error(), ErrorCode::TENANT_NOT_EMPTY);
+    auto exists = service.ExistKey("cold", "tenant-a");
+    ASSERT_TRUE(exists.has_value()) << toString(exists.error());
+    EXPECT_TRUE(exists.value());
+}
+
+#ifdef USE_NOF
+TEST_F(MasterServiceTenantQuotaTest,
+       DeletePolicyWaitsForZeroChargePutStartMetadataCreate) {
+    MasterService service(MakeConfig({{"tenant-a", 1000}}));
+    UUID client_id = MountNoFSegment(service);
+
+    auto blocking_strategy = std::make_shared<BlockingAllocationStrategy>();
+    auto* blocking_strategy_ptr = blocking_strategy.get();
+    auto allocation_started = blocking_strategy_ptr->AllocationStarted();
+    ReplaceAllocationStrategy(service, std::move(blocking_strategy));
+
+    ReplicateConfig config;
+    config.replica_num = 0;
+    config.nof_replica_num = 1;
+
+    std::optional<tl::expected<std::vector<Replica::Descriptor>, ErrorCode>>
+        put_result;
+    std::thread put_thread([&] {
+        put_result.emplace(
+            service.PutStart(client_id, "nof-key", "tenant-a", 128, config));
+    });
+
+    if (allocation_started.wait_for(std::chrono::seconds(5)) !=
+        std::future_status::ready) {
+        blocking_strategy_ptr->AllowAllocation();
+        put_thread.join();
+        FAIL() << "timed out waiting for PutStart allocation";
+    }
+
+    using DeleteResult =
+        tl::expected<std::optional<TenantQuotaSnapshot>, ErrorCode>;
+    std::optional<DeleteResult> delete_result;
+    std::thread delete_thread([&] {
+        delete_result.emplace(service.DeleteTenantQuotaPolicy("tenant-a"));
+    });
+
+    ASSERT_TRUE(WaitForTenantQuotaPolicyMutexContention(service))
+        << "DeleteTenantQuotaPolicy did not wait for zero-charge PutStart";
+
+    blocking_strategy_ptr->AllowAllocation();
+    put_thread.join();
+    delete_thread.join();
+
+    ASSERT_TRUE(put_result.has_value());
+    ASSERT_TRUE(put_result->has_value()) << toString(put_result->error());
+    ASSERT_TRUE(delete_result.has_value());
+    ASSERT_FALSE(delete_result->has_value());
+    EXPECT_EQ(delete_result->error(), ErrorCode::TENANT_NOT_EMPTY);
+
+    auto snapshot = Snapshot(service, "tenant-a");
+    EXPECT_EQ(snapshot.used_bytes, 0);
+    EXPECT_EQ(snapshot.reserved_bytes, 0);
+    EXPECT_EQ(snapshot.metadata_object_count, 1);
+}
+#endif
+
+TEST_F(MasterServiceTenantQuotaTest,
+       EffectiveQuotaUsesOnlyExplicitPolicyAndScalesProportionally) {
+    MasterService service(MakeConfig({{"tenant-a", 200}, {"tenant-b", 400}}));
+    MountSegment(service, /*size=*/300);
+
+    EXPECT_EQ(Snapshot(service, "tenant-a").effective_quota_bytes, 100);
+    EXPECT_EQ(Snapshot(service, "tenant-b").effective_quota_bytes, 200);
+}
+
+TEST_F(MasterServiceTenantQuotaTest,
+       ConnectorPolicyReloadCreatesOrphanStateAndAllowsCleanup) {
+    const std::string initial_policy =
+        WritePolicyFile({{"tenant-a", 1000}, {"tenant-b", 1000}});
+    auto config = MasterServiceConfig::builder()
+                      .set_enable_multi_tenants(true)
+                      .set_tenant_quota_connector_type("file")
+                      .set_tenant_quota_connector_uri(initial_policy)
+                      .build();
+    MasterService service(config);
+    UUID client_id = MountSegment(service);
+    PutComplete(service, client_id, "orphan-key", "tenant-b", 100);
+
+    {
+        std::ofstream out(initial_policy);
+        TenantQuotaPolicySnapshot replacement;
+        replacement.tenant_quotas = {{"tenant-a", 1000}};
+        out << FormatTenantQuotaPolicyYaml(replacement);
+    }
+    ReloadTenantQuotaPolicyFromStore(service);
+
+    auto orphan = Snapshot(service, "tenant-b");
+    EXPECT_FALSE(orphan.has_explicit_policy);
+    EXPECT_EQ(orphan.requested_quota_bytes, 0);
+    EXPECT_EQ(orphan.effective_quota_bytes, 0);
+    EXPECT_TRUE(orphan.over_quota);
+
+    EXPECT_TRUE(service.GetReplicaList("orphan-key", "tenant-b").has_value());
+    auto write =
+        service.PutStart(client_id, "new-key", "tenant-b", 1, MemoryConfig());
+    ASSERT_FALSE(write.has_value());
+    EXPECT_EQ(write.error(), ErrorCode::TENANT_NOT_REGISTERED);
+
+    EXPECT_TRUE(service
+                    .Remove("orphan-key", "tenant-b",
+                            /*force=*/true)
+                    .has_value());
+}
+
+}  // namespace mooncake::test
diff --git a/mooncake-store/tests/master_service_test.cpp b/mooncake-store/tests/master_service_test.cpp
index 91cb5144..30740f63 100644
--- a/mooncake-store/tests/master_service_test.cpp
+++ b/mooncake-store/tests/master_service_test.cpp
@@ -1,18 +1,29 @@
 #include "master_service.h"
+#include "rpc_service.h"
 
 #include <glog/logging.h>
 #include <gtest/gtest.h>
 #include <ylt/struct_json/json_reader.h>
 
+#include <algorithm>
 #include <atomic>
 #include <chrono>
+#include <filesystem>
+#include <fstream>
+#include <functional>
+#include <limits>
+#include <map>
 #include <memory>
 #include <random>
+#include <string>
 #include <thread>
 #include <vector>
 #include <unordered_set>
 #include <utility>
 
+#include <unistd.h>
+
+#include "tenant_quota_policy_store.h"
 #include "types.h"
 
 namespace mooncake::test {
@@ -31,6 +42,52 @@ class MasterServiceTest : public ::testing::Test {
 
     static constexpr size_t kDefaultSegmentBase = 0x300000000;
     static constexpr size_t kDefaultSegmentSize = 1024 * 1024 * 16;
+    static constexpr uint64_t kStrictTenantQuotaBytes = 4 * 1024 * 1024;
+
+    std::string WriteTenantPolicyFile(
+        const std::map<std::string, uint64_t>& tenant_quotas) {
+        TenantQuotaPolicySnapshot snapshot;
+        snapshot.tenant_quotas = tenant_quotas;
+        auto path =
+            std::filesystem::temp_directory_path() /
+            ("mooncake_master_service_test_" + std::to_string(::getpid()) +
+             "_" + std::to_string(next_policy_file_++) + ".yaml");
+        std::ofstream out(path);
+        out << FormatTenantQuotaPolicyYaml(snapshot);
+        out.close();
+        policy_files_.push_back(path.string());
+        return path.string();
+    }
+
+    MasterServiceConfig MakeStrictTenantConfig(
+        const std::vector<std::string>& tenants) {
+        std::map<std::string, uint64_t> tenant_quotas;
+        for (const auto& tenant : tenants) {
+            tenant_quotas.emplace(tenant, kStrictTenantQuotaBytes);
+        }
+        return MasterServiceConfig::builder()
+            .set_enable_multi_tenants(true)
+            .set_tenant_quota_connector_type("file")
+            .set_tenant_quota_connector_uri(
+                WriteTenantPolicyFile(tenant_quotas))
+            .build();
+    }
+
+    WrappedMasterServiceConfig MakeStrictWrappedConfig(
+        const std::vector<std::string>& tenants) {
+        WrappedMasterServiceConfig config;
+        config.default_kv_lease_ttl = 100;
+        config.enable_metric_reporting = false;
+        config.enable_multi_tenants = true;
+        config.tenant_quota_connector_type = "file";
+        std::map<std::string, uint64_t> tenant_quotas;
+        for (const auto& tenant : tenants) {
+            tenant_quotas.emplace(tenant, kStrictTenantQuotaBytes);
+        }
+        config.tenant_quota_connector_uri =
+            WriteTenantPolicyFile(tenant_quotas);
+        return config;
+    }
 
     Segment MakeSegment(std::string name = "test_segment",
                         size_t base = kDefaultSegmentBase,
@@ -44,6 +101,22 @@ class MasterServiceTest : public ::testing::Test {
         return segment;
     }
 
+#ifdef USE_NOF
+    NoFSegment MakeNoFSegment(
+        std::string name = "test_nof_segment",
+        std::string endpoint = "test_nof_segment_endpoint",
+        size_t base = kDefaultSegmentBase + kDefaultSegmentSize,
+        size_t size = kDefaultSegmentSize) const {
+        NoFSegment segment;
+        segment.id = generate_uuid();
+        segment.name = std::move(name);
+        segment.base = base;
+        segment.size = size;
+        segment.te_endpoint = std::move(endpoint);
+        return segment;
+    }
+#endif
+
     MountedSegmentContext PrepareSimpleSegment(
         MasterService& service, std::string name = "test_segment",
         size_t base = kDefaultSegmentBase,
@@ -67,13 +140,59 @@ class MasterServiceTest : public ::testing::Test {
         config.replica_num = 1;
         config.preferred_segment = segment_name;
 
-        auto put_start = service.PutStart(client_id, key, slice_length, config);
+        auto put_start =
+            service.PutStart(client_id, key, "default", slice_length, config);
         EXPECT_TRUE(put_start.has_value());
         EXPECT_TRUE(
-            service.PutEnd(client_id, key, ReplicaType::MEMORY).has_value());
+            service.PutEnd(client_id, key, "default", ReplicaType::MEMORY)
+                .has_value());
         return key;
     }
 
+    std::string FindGroupIdOnDifferentShard(const std::string& key) const {
+        static constexpr size_t kMetadataShardCountForTest = 1024;
+        const size_t key_shard =
+            std::hash<std::string>{}(key) % kMetadataShardCountForTest;
+        for (int i = 0; i < 10000; ++i) {
+            std::string group_id = key + "_group_" + std::to_string(i);
+            if (std::hash<std::string>{}(group_id) %
+                    kMetadataShardCountForTest !=
+                key_shard) {
+                return group_id;
+            }
+        }
+        return key + "_fallback_group";
+    }
+
+    void PutCompletedObject(MasterService& service, const UUID& client_id,
+                            const std::string& key,
+                            const ReplicateConfig& config,
+                            uint64_t slice_length = 1024) const {
+        auto put_start =
+            service.PutStart(client_id, key, "default", slice_length, config);
+        ASSERT_TRUE(put_start.has_value())
+            << "PutStart failed for key=" << key
+            << ", error=" << toString(put_start.error());
+        ASSERT_TRUE(
+            service.PutEnd(client_id, key, "default", ReplicaType::MEMORY)
+                .has_value());
+    }
+
+    void PutCompletedObject(MasterService& service, const UUID& client_id,
+                            const std::string& key,
+                            const std::string& tenant_id,
+                            const ReplicateConfig& config,
+                            uint64_t slice_length = 1024) const {
+        auto put_start =
+            service.PutStart(client_id, key, tenant_id, slice_length, config);
+        ASSERT_TRUE(put_start.has_value())
+            << "PutStart failed for key=" << key << ", tenant_id=" << tenant_id
+            << ", error=" << toString(put_start.error());
+        ASSERT_TRUE(
+            service.PutEnd(client_id, key, tenant_id, ReplicaType::MEMORY)
+                .has_value());
+    }
+
     bool ExecutePendingMoveTasks(MasterService& service,
                                  const UUID& client_id) const {
         auto fetched = service.FetchTasks(client_id, /*batch_size=*/16);
@@ -91,10 +210,13 @@ class MasterServiceTest : public ::testing::Test {
             ReplicaMovePayload payload;
             struct_json::from_json(payload, assignment.payload);
 
-            auto move_start = service.MoveStart(client_id, payload.key,
-                                                payload.source, payload.target);
+            auto move_start =
+                service.MoveStart(client_id, payload.key, payload.tenant_id,
+                                  payload.source, payload.target);
             EXPECT_TRUE(move_start.has_value());
-            EXPECT_TRUE(service.MoveEnd(client_id, payload.key).has_value());
+            EXPECT_TRUE(
+                service.MoveEnd(client_id, payload.key, payload.tenant_id)
+                    .has_value());
 
             TaskCompleteRequest complete_request;
             complete_request.id = assignment.id;
@@ -149,10 +271,39 @@ class MasterServiceTest : public ::testing::Test {
     }
 
     std::vector<Replica::Descriptor> replica_list;
+    std::vector<std::string> policy_files_;
+    size_t next_policy_file_ = 0;
 
-    void TearDown() override { google::ShutdownGoogleLogging(); }
+    void TearDown() override {
+        for (const auto& path : policy_files_) {
+            std::error_code ec;
+            std::filesystem::remove(path, ec);
+        }
+        google::ShutdownGoogleLogging();
+    }
 };
 
+TEST(TenantScopedStorageKeyTest, RoundTripsAndParsesLegacyKeys) {
+    const auto scoped =
+        MakeTenantScopedStorageKey("tenant:with:colon", "path/key:with:colon");
+    EXPECT_NE(scoped.find('\0'), std::string::npos);
+
+    auto [tenant_id, key] = ParseTenantScopedStorageKey(scoped);
+    EXPECT_EQ(tenant_id, "tenant:with:colon");
+    EXPECT_EQ(key, "path/key:with:colon");
+
+    auto [default_tenant, default_key] = ParseTenantScopedStorageKey("raw_key");
+    EXPECT_EQ(default_tenant, "default");
+    EXPECT_EQ(default_key, "raw_key");
+
+    std::string legacy = "legacy_tenant";
+    legacy.push_back('\0');
+    legacy.append("legacy_key");
+    auto [legacy_tenant, legacy_key] = ParseTenantScopedStorageKey(legacy);
+    EXPECT_EQ(legacy_tenant, "legacy_tenant");
+    EXPECT_EQ(legacy_key, "legacy_key");
+}
+
 std::string GenerateKeyForSegment(const UUID& client_id,
                                   const std::unique_ptr<MasterService>& service,
                                   const std::string& segment_name) {
@@ -163,14 +314,14 @@ std::string GenerateKeyForSegment(const UUID& client_id,
         std::vector<Replica::Descriptor> replica_list;
 
         // Check if the key already exists.
-        auto exist_result = service->ExistKey(key);
+        auto exist_result = service->ExistKey(key, "default");
         if (exist_result.has_value() && exist_result.value()) {
             continue;  // Retry if the key already exists
         }
 
         // Attempt to put the key.
-        auto put_result =
-            service->PutStart(client_id, key, {1024}, {.replica_num = 1});
+        auto put_result = service->PutStart(client_id, key, "default", {1024},
+                                            {.replica_num = 1});
         if (put_result.has_value()) {
             replica_list = std::move(put_result.value());
         }
@@ -185,7 +336,7 @@ std::string GenerateKeyForSegment(const UUID& client_id,
                                      std::to_string(static_cast<int>(code)));
         }
         auto put_end_result =
-            service->PutEnd(client_id, key, ReplicaType::MEMORY);
+            service->PutEnd(client_id, key, "default", ReplicaType::MEMORY);
         if (!put_end_result.has_value()) {
             throw std::runtime_error("PutEnd failed");
         }
@@ -195,7 +346,7 @@ std::string GenerateKeyForSegment(const UUID& client_id,
             return key;
         }
         // Clean up failed attempt
-        auto remove_result = service->Remove(key);
+        auto remove_result = service->Remove(key, "default");
         if (!remove_result.has_value()) {
             // Ignore cleanup failure
         }
@@ -406,17 +557,996 @@ TEST_F(MasterServiceTest, PutStartInvalidParams) {
     std::string key = "test_key";
     ReplicateConfig config;
 
-    // Test invalid replica_num
+    // Test invalid replica config
     config.replica_num = 0;
-    auto put_result1 = service_->PutStart(client_id, key, 1024, config);
+    config.nof_replica_num = 0;
+    auto put_result1 =
+        service_->PutStart(client_id, key, "default", 1024, config);
     EXPECT_FALSE(put_result1.has_value());
     EXPECT_EQ(ErrorCode::INVALID_PARAMS, put_result1.error());
 
     // Test zero slice_length
     config.replica_num = 1;
-    auto put_result2 = service_->PutStart(client_id, key, 0, config);
+    auto put_result2 = service_->PutStart(client_id, key, "default", 0, config);
     EXPECT_FALSE(put_result2.has_value());
     EXPECT_EQ(ErrorCode::INVALID_PARAMS, put_result2.error());
+
+    // Test prefer_alloc_in_same_node with nof replicas
+    config.nof_replica_num = 1;
+    config.prefer_alloc_in_same_node = true;
+    auto put_result3 =
+        service_->PutStart(client_id, key, "default", 1024, config);
+    EXPECT_FALSE(put_result3.has_value());
+    EXPECT_EQ(ErrorCode::INVALID_PARAMS, put_result3.error());
+}
+
+#ifdef USE_NOF
+TEST_F(MasterServiceTest, PutEndAllCompletesMemoryAndNoFReplicas) {
+    std::unique_ptr<MasterService> service_(new MasterService());
+    [[maybe_unused]] const auto mem_context = PrepareSimpleSegment(*service_);
+    NoFSegment nof_segment = MakeNoFSegment();
+    const UUID client_id = generate_uuid();
+    ASSERT_TRUE(service_->MountNoFSegment(nof_segment, client_id).has_value());
+
+    ReplicateConfig config;
+    config.replica_num = 1;
+    config.nof_replica_num = 1;
+    auto put_start_result =
+        service_->PutStart(client_id, "test_key_all", "default", 1024, config);
+    ASSERT_TRUE(put_start_result.has_value());
+
+    auto put_end_result = service_->PutEnd(client_id, "test_key_all", "default",
+                                           ReplicaType::ALL);
+    ASSERT_TRUE(put_end_result.has_value());
+
+    auto get_replica_result =
+        service_->GetReplicaList("test_key_all", "default");
+    ASSERT_TRUE(get_replica_result.has_value());
+
+    bool has_complete_memory = false;
+    bool has_complete_nof = false;
+    for (const auto& replica : get_replica_result->replicas) {
+        if (replica.is_memory_replica() &&
+            replica.status == ReplicaStatus::COMPLETE) {
+            has_complete_memory = true;
+        }
+        if (replica.is_nof_replica() &&
+            replica.status == ReplicaStatus::COMPLETE) {
+            has_complete_nof = true;
+        }
+    }
+    EXPECT_TRUE(has_complete_memory);
+    EXPECT_TRUE(has_complete_nof);
+}
+
+TEST_F(MasterServiceTest, PutEndMemoryDoesNotCompleteNoFReplica) {
+    std::unique_ptr<MasterService> service_(new MasterService());
+    [[maybe_unused]] const auto mem_context = PrepareSimpleSegment(*service_);
+    NoFSegment nof_segment =
+        MakeNoFSegment("test_nof_segment_2", "test_nof_segment_endpoint_2");
+    const UUID client_id = generate_uuid();
+    ASSERT_TRUE(service_->MountNoFSegment(nof_segment, client_id).has_value());
+
+    ReplicateConfig config;
+    config.replica_num = 1;
+    config.nof_replica_num = 1;
+    auto put_start_result = service_->PutStart(client_id, "test_key_split",
+                                               "default", 1024, config);
+    ASSERT_TRUE(put_start_result.has_value());
+
+    auto put_end_result = service_->PutEnd(client_id, "test_key_split",
+                                           "default", ReplicaType::MEMORY);
+    ASSERT_TRUE(put_end_result.has_value());
+
+    auto get_replica_result =
+        service_->GetReplicaList("test_key_split", "default");
+    ASSERT_TRUE(get_replica_result.has_value());
+    ASSERT_EQ(get_replica_result->replicas.size(), 1u);
+    EXPECT_TRUE(get_replica_result->replicas[0].is_memory_replica());
+    EXPECT_EQ(get_replica_result->replicas[0].status, ReplicaStatus::COMPLETE);
+
+    auto put_revoke_result = service_->PutRevoke(
+        client_id, "test_key_split", "default", ReplicaType::NOF_SSD);
+    ASSERT_TRUE(put_revoke_result.has_value());
+
+    auto final_replica_result =
+        service_->GetReplicaList("test_key_split", "default");
+    ASSERT_TRUE(final_replica_result.has_value());
+    ASSERT_EQ(final_replica_result->replicas.size(), 1u);
+    EXPECT_TRUE(final_replica_result->replicas[0].is_memory_replica());
+    EXPECT_EQ(final_replica_result->replicas[0].status,
+              ReplicaStatus::COMPLETE);
+}
+
+TEST_F(MasterServiceTest, PutStartOnePlusOneAllowsSingleAllocatedReplica) {
+    std::unique_ptr<MasterService> service_(new MasterService());
+    [[maybe_unused]] const auto mem_context = PrepareSimpleSegment(*service_);
+    const UUID client_id = generate_uuid();
+
+    ReplicateConfig config;
+    config.replica_num = 1;
+    config.nof_replica_num = 1;
+    auto put_start_result = service_->PutStart(
+        client_id, "test_key_one_plus_one", "default", 1024, config);
+    ASSERT_TRUE(put_start_result.has_value());
+    ASSERT_EQ(put_start_result->size(), 1u);
+    EXPECT_TRUE(put_start_result->front().is_memory_replica());
+}
+#endif
+
+TEST_F(MasterServiceTest, PutStartGroupIdsValidation) {
+    std::unique_ptr<MasterService> service_(new MasterService());
+    [[maybe_unused]] const auto context = PrepareSimpleSegment(*service_);
+    const UUID client_id = generate_uuid();
+
+    ReplicateConfig config;
+    config.replica_num = 1;
+
+    config.group_ids = std::vector<std::string>{};
+    auto empty_group_ids = service_->PutStart(client_id, "empty_group_ids",
+                                              "default", 1024, config);
+    EXPECT_FALSE(empty_group_ids.has_value());
+    EXPECT_EQ(ErrorCode::INVALID_PARAMS, empty_group_ids.error());
+
+    config.group_ids = std::vector<std::string>{"g0", "g1"};
+    auto too_many_group_ids = service_->PutStart(
+        client_id, "too_many_group_ids", "default", 1024, config);
+    EXPECT_FALSE(too_many_group_ids.has_value());
+    EXPECT_EQ(ErrorCode::INVALID_PARAMS, too_many_group_ids.error());
+
+    config.group_ids = std::vector<std::string>{""};
+    auto ungrouped = service_->PutStart(client_id, "explicit_ungrouped",
+                                        "default", 1024, config);
+    ASSERT_TRUE(ungrouped.has_value());
+    ASSERT_TRUE(service_
+                    ->PutEnd(client_id, "explicit_ungrouped", "default",
+                             ReplicaType::MEMORY)
+                    .has_value());
+    auto exists = service_->ExistKey("explicit_ungrouped", "default");
+    ASSERT_TRUE(exists.has_value());
+    EXPECT_TRUE(exists.value());
+}
+
+TEST_F(MasterServiceTest, GroupedObjectRoutesKeyLevelLookupAndRemove) {
+    std::unique_ptr<MasterService> service_(new MasterService());
+    [[maybe_unused]] const auto context = PrepareSimpleSegment(*service_);
+    const UUID client_id = generate_uuid();
+
+    const std::string key = "grouped_route_key";
+    ReplicateConfig config;
+    config.replica_num = 1;
+    config.group_ids =
+        std::vector<std::string>{FindGroupIdOnDifferentShard(key)};
+
+    PutCompletedObject(*service_, client_id, key, config);
+
+    auto exists = service_->ExistKey(key, "default");
+    ASSERT_TRUE(exists.has_value());
+    EXPECT_TRUE(exists.value());
+    EXPECT_TRUE(service_->GetReplicaList(key, "default").has_value());
+
+    ASSERT_TRUE(service_->Remove(key, "default", /*force=*/true).has_value());
+    auto exists_after_remove = service_->ExistKey(key, "default");
+    ASSERT_TRUE(exists_after_remove.has_value());
+    EXPECT_FALSE(exists_after_remove.value());
+}
+
+TEST_F(MasterServiceTest, GroupRoutingIsTenantScopedForSameUserKey) {
+    const std::string key = "tenant_grouped_shared_user_key";
+    const std::string tenant_a = "tenant_group_route_a";
+    const std::string tenant_b = "tenant_group_route_b";
+    auto service_ = std::make_unique<MasterService>(
+        MakeStrictTenantConfig({tenant_a, tenant_b}));
+    [[maybe_unused]] const auto context = PrepareSimpleSegment(*service_);
+    const UUID client_id = generate_uuid();
+
+    const std::string group_a = FindGroupIdOnDifferentShard(key);
+    std::string group_b;
+    for (int i = 0; i < 10000; ++i) {
+        group_b = key + "_tenant_b_group_" + std::to_string(i);
+        if (std::hash<std::string>{}(group_b) % 1024 !=
+            std::hash<std::string>{}(group_a) % 1024) {
+            break;
+        }
+    }
+
+    ReplicateConfig config_a;
+    config_a.replica_num = 1;
+    config_a.group_ids = std::vector<std::string>{group_a};
+    ReplicateConfig config_b;
+    config_b.replica_num = 1;
+    config_b.group_ids = std::vector<std::string>{group_b};
+
+    ASSERT_TRUE(service_->PutStart(client_id, key, tenant_a, 1024, config_a)
+                    .has_value());
+    ASSERT_TRUE(service_->PutEnd(client_id, key, tenant_a, ReplicaType::MEMORY)
+                    .has_value());
+    ASSERT_TRUE(service_->PutStart(client_id, key, tenant_b, 2048, config_b)
+                    .has_value());
+    ASSERT_TRUE(service_->PutEnd(client_id, key, tenant_b, ReplicaType::MEMORY)
+                    .has_value());
+
+    EXPECT_TRUE(service_->ExistKey(key, tenant_a).value_or(false));
+    EXPECT_TRUE(service_->ExistKey(key, tenant_b).value_or(false));
+    EXPECT_TRUE(service_->GetReplicaList(key, tenant_a).has_value());
+    EXPECT_TRUE(service_->GetReplicaList(key, tenant_b).has_value());
+
+    ASSERT_TRUE(service_->Remove(key, tenant_a, /*force=*/true).has_value());
+    EXPECT_FALSE(service_->GetReplicaList(key, tenant_a).has_value());
+    EXPECT_TRUE(service_->GetReplicaList(key, tenant_b).has_value());
+}
+
+TEST_F(MasterServiceTest, BatchGetReplicaListPreservesOrderWithGroupedKeys) {
+    std::unique_ptr<MasterService> service_(new MasterService());
+    [[maybe_unused]] const auto context = PrepareSimpleSegment(*service_);
+    const UUID client_id = generate_uuid();
+
+    const std::string grouped_key_a = "batch_get_grouped_a";
+    const std::string missing_key = "batch_get_missing";
+    const std::string ungrouped_key = "batch_get_ungrouped";
+    const std::string grouped_key_b = "batch_get_grouped_b";
+    const std::string pending_key = "batch_get_pending";
+
+    ReplicateConfig grouped_config_a;
+    grouped_config_a.replica_num = 1;
+    grouped_config_a.group_ids =
+        std::vector<std::string>{FindGroupIdOnDifferentShard(grouped_key_a)};
+    PutCompletedObject(*service_, client_id, grouped_key_a, grouped_config_a);
+
+    ReplicateConfig ungrouped_config;
+    ungrouped_config.replica_num = 1;
+    PutCompletedObject(*service_, client_id, ungrouped_key, ungrouped_config);
+
+    ReplicateConfig grouped_config_b;
+    grouped_config_b.replica_num = 1;
+    grouped_config_b.group_ids =
+        std::vector<std::string>{FindGroupIdOnDifferentShard(grouped_key_b)};
+    PutCompletedObject(*service_, client_id, grouped_key_b, grouped_config_b);
+
+    ASSERT_TRUE(service_
+                    ->PutStart(client_id, pending_key, "default", 1024,
+                               ungrouped_config)
+                    .has_value());
+
+    const std::vector<std::string> keys = {
+        grouped_key_a, missing_key, ungrouped_key, grouped_key_b, pending_key};
+    auto results = service_->BatchGetReplicaList(keys, "default");
+
+    ASSERT_EQ(results.size(), keys.size());
+    ASSERT_TRUE(results[0].has_value());
+    EXPECT_FALSE(results[0]->replicas.empty());
+    ASSERT_FALSE(results[1].has_value());
+    EXPECT_EQ(ErrorCode::OBJECT_NOT_FOUND, results[1].error());
+    ASSERT_TRUE(results[2].has_value());
+    EXPECT_FALSE(results[2]->replicas.empty());
+    ASSERT_TRUE(results[3].has_value());
+    EXPECT_FALSE(results[3]->replicas.empty());
+    ASSERT_FALSE(results[4].has_value());
+    EXPECT_EQ(ErrorCode::REPLICA_IS_NOT_READY, results[4].error());
+}
+
+TEST_F(MasterServiceTest, BatchGetReplicaListKeepsTenantIsolation) {
+    const std::string key = "batch_get_tenant_shared_key";
+    const std::string tenant_a = "batch_get_tenant_a";
+    const std::string tenant_b = "batch_get_tenant_b";
+    auto service_ = std::make_unique<MasterService>(
+        MakeStrictTenantConfig({"default", tenant_a, tenant_b}));
+    [[maybe_unused]] const auto context = PrepareSimpleSegment(*service_);
+    const UUID client_id = generate_uuid();
+
+    ReplicateConfig config_a;
+    config_a.replica_num = 1;
+    config_a.group_ids =
+        std::vector<std::string>{FindGroupIdOnDifferentShard(key)};
+    ASSERT_TRUE(service_->PutStart(client_id, key, tenant_a, 1024, config_a)
+                    .has_value());
+    ASSERT_TRUE(service_->PutEnd(client_id, key, tenant_a, ReplicaType::MEMORY)
+                    .has_value());
+
+    ReplicateConfig config_b;
+    config_b.replica_num = 1;
+    ASSERT_TRUE(service_->PutStart(client_id, key, tenant_b, 2048, config_b)
+                    .has_value());
+    ASSERT_TRUE(service_->PutEnd(client_id, key, tenant_b, ReplicaType::MEMORY)
+                    .has_value());
+
+    auto tenant_a_results = service_->BatchGetReplicaList({key}, tenant_a);
+    auto tenant_b_results = service_->BatchGetReplicaList({key}, tenant_b);
+    auto default_results = service_->BatchGetReplicaList({key}, "default");
+
+    ASSERT_EQ(tenant_a_results.size(), 1u);
+    ASSERT_EQ(tenant_b_results.size(), 1u);
+    ASSERT_EQ(default_results.size(), 1u);
+    EXPECT_TRUE(tenant_a_results[0].has_value());
+    EXPECT_TRUE(tenant_b_results[0].has_value());
+    ASSERT_FALSE(default_results[0].has_value());
+    EXPECT_EQ(ErrorCode::OBJECT_NOT_FOUND, default_results[0].error());
+}
+
+TEST_F(MasterServiceTest, GetAllKeysListsOnlyRequestedTenant) {
+    const std::string tenant_a = "tenant_get_all_keys_a";
+    auto service_ = std::make_unique<MasterService>(
+        MakeStrictTenantConfig({"default", tenant_a}));
+    [[maybe_unused]] const auto context = PrepareSimpleSegment(*service_);
+    const UUID client_id = generate_uuid();
+
+    const std::string shared_key = "shared_listing_key";
+    const std::string default_only_key = "default_listing_key";
+    const std::string tenant_only_key = "tenant_listing_key";
+
+    ReplicateConfig config;
+    config.replica_num = 1;
+    ASSERT_TRUE(
+        service_->PutStart(client_id, shared_key, "default", 1024, config)
+            .has_value());
+    ASSERT_TRUE(
+        service_->PutEnd(client_id, shared_key, "default", ReplicaType::MEMORY)
+            .has_value());
+    ASSERT_TRUE(
+        service_->PutStart(client_id, default_only_key, "default", 1024, config)
+            .has_value());
+    ASSERT_TRUE(service_
+                    ->PutEnd(client_id, default_only_key, "default",
+                             ReplicaType::MEMORY)
+                    .has_value());
+    ASSERT_TRUE(
+        service_->PutStart(client_id, shared_key, tenant_a, 1024, config)
+            .has_value());
+    ASSERT_TRUE(
+        service_->PutEnd(client_id, shared_key, tenant_a, ReplicaType::MEMORY)
+            .has_value());
+    ASSERT_TRUE(
+        service_->PutStart(client_id, tenant_only_key, tenant_a, 1024, config)
+            .has_value());
+    ASSERT_TRUE(
+        service_
+            ->PutEnd(client_id, tenant_only_key, tenant_a, ReplicaType::MEMORY)
+            .has_value());
+
+    auto default_keys = service_->GetAllKeys("default");
+    ASSERT_TRUE(default_keys.has_value());
+    EXPECT_NE(std::find(default_keys->begin(), default_keys->end(), shared_key),
+              default_keys->end());
+    EXPECT_NE(
+        std::find(default_keys->begin(), default_keys->end(), default_only_key),
+        default_keys->end());
+    EXPECT_EQ(
+        std::find(default_keys->begin(), default_keys->end(), tenant_only_key),
+        default_keys->end());
+
+    auto tenant_keys = service_->GetAllKeys(tenant_a);
+    ASSERT_TRUE(tenant_keys.has_value());
+    EXPECT_NE(std::find(tenant_keys->begin(), tenant_keys->end(), shared_key),
+              tenant_keys->end());
+    EXPECT_NE(
+        std::find(tenant_keys->begin(), tenant_keys->end(), tenant_only_key),
+        tenant_keys->end());
+    EXPECT_EQ(
+        std::find(tenant_keys->begin(), tenant_keys->end(), default_only_key),
+        tenant_keys->end());
+}
+
+TEST_F(MasterServiceTest,
+       ConcurrentGroupedAndUngroupedFirstCreateDoesNotDuplicateMetadata) {
+    std::unique_ptr<MasterService> service_(new MasterService());
+    [[maybe_unused]] const auto context = PrepareSimpleSegment(*service_);
+    const UUID client_id = generate_uuid();
+
+    const std::string key = "concurrent_grouped_ungrouped_first_create";
+    const std::string tenant_id = "tenant_concurrent_first_create";
+    ReplicateConfig ungrouped_config;
+    ungrouped_config.replica_num = 1;
+    ReplicateConfig grouped_config;
+    grouped_config.replica_num = 1;
+    grouped_config.group_ids =
+        std::vector<std::string>{FindGroupIdOnDifferentShard(key)};
+
+    static constexpr size_t kThreadCount = 16;
+    std::atomic<size_t> ready{0};
+    std::atomic<bool> start{false};
+    std::vector<int> put_start_success(kThreadCount, 0);
+    std::vector<int> put_end_success(kThreadCount, 0);
+    std::vector<std::thread> threads;
+    threads.reserve(kThreadCount);
+
+    for (size_t i = 0; i < kThreadCount; ++i) {
+        threads.emplace_back([&, i]() {
+            ready.fetch_add(1, std::memory_order_acq_rel);
+            while (!start.load(std::memory_order_acquire)) {
+                std::this_thread::yield();
+            }
+            const auto& config =
+                (i % 2 == 0) ? grouped_config : ungrouped_config;
+            auto put_start =
+                service_->PutStart(client_id, key, tenant_id, 1024, config);
+            put_start_success[i] = put_start.has_value() ? 1 : 0;
+            if (put_start.has_value()) {
+                put_end_success[i] = service_->PutEnd(client_id, key, tenant_id,
+                                                      ReplicaType::MEMORY)
+                                             .has_value()
+                                         ? 1
+                                         : 0;
+            } else {
+                EXPECT_EQ(ErrorCode::OBJECT_ALREADY_EXISTS, put_start.error());
+            }
+        });
+    }
+
+    while (ready.load(std::memory_order_acquire) < kThreadCount) {
+        std::this_thread::yield();
+    }
+    start.store(true, std::memory_order_release);
+    for (auto& thread : threads) {
+        thread.join();
+    }
+
+    EXPECT_EQ(std::count(put_start_success.begin(), put_start_success.end(), 1),
+              1);
+    EXPECT_EQ(std::count(put_end_success.begin(), put_end_success.end(), 1), 1);
+    EXPECT_EQ(service_->GetKeyCount(), 1u);
+    EXPECT_TRUE(service_->GetReplicaList(key, tenant_id).has_value());
+}
+
+TEST_F(MasterServiceTest,
+       ConcurrentDifferentGroupedFirstCreateDoesNotDuplicateMetadata) {
+    std::unique_ptr<MasterService> service_(new MasterService());
+    [[maybe_unused]] const auto context = PrepareSimpleSegment(*service_);
+    const UUID client_id = generate_uuid();
+
+    const std::string key = "concurrent_different_grouped_first_create";
+    const std::string tenant_id = "tenant_concurrent_grouped_first_create";
+    const std::string group_a = FindGroupIdOnDifferentShard(key);
+    std::string group_b;
+    for (int i = 0; i < 10000; ++i) {
+        group_b = key + "_other_group_" + std::to_string(i);
+        if (std::hash<std::string>{}(group_b) % 1024 !=
+            std::hash<std::string>{}(group_a) % 1024) {
+            break;
+        }
+    }
+    ReplicateConfig config_a;
+    config_a.replica_num = 1;
+    config_a.group_ids = std::vector<std::string>{group_a};
+    ReplicateConfig config_b;
+    config_b.replica_num = 1;
+    config_b.group_ids = std::vector<std::string>{group_b};
+
+    static constexpr size_t kThreadCount = 16;
+    std::atomic<size_t> ready{0};
+    std::atomic<bool> start{false};
+    std::vector<int> put_start_success(kThreadCount, 0);
+    std::vector<int> put_end_success(kThreadCount, 0);
+    std::vector<std::thread> threads;
+    threads.reserve(kThreadCount);
+
+    for (size_t i = 0; i < kThreadCount; ++i) {
+        threads.emplace_back([&, i]() {
+            ready.fetch_add(1, std::memory_order_acq_rel);
+            while (!start.load(std::memory_order_acquire)) {
+                std::this_thread::yield();
+            }
+            const auto& config = (i % 2 == 0) ? config_a : config_b;
+            auto put_start =
+                service_->PutStart(client_id, key, tenant_id, 1024, config);
+            put_start_success[i] = put_start.has_value() ? 1 : 0;
+            if (put_start.has_value()) {
+                put_end_success[i] = service_->PutEnd(client_id, key, tenant_id,
+                                                      ReplicaType::MEMORY)
+                                             .has_value()
+                                         ? 1
+                                         : 0;
+            } else {
+                EXPECT_EQ(ErrorCode::OBJECT_ALREADY_EXISTS, put_start.error());
+            }
+        });
+    }
+
+    while (ready.load(std::memory_order_acquire) < kThreadCount) {
+        std::this_thread::yield();
+    }
+    start.store(true, std::memory_order_release);
+    for (auto& thread : threads) {
+        thread.join();
+    }
+
+    EXPECT_EQ(std::count(put_start_success.begin(), put_start_success.end(), 1),
+              1);
+    EXPECT_EQ(std::count(put_end_success.begin(), put_end_success.end(), 1), 1);
+    EXPECT_EQ(service_->GetKeyCount(), 1u);
+    EXPECT_TRUE(service_->GetReplicaList(key, tenant_id).has_value());
+}
+
+TEST_F(MasterServiceTest, ExpiredGroupedPutCanBeReplacedByUngroupedPut) {
+    auto service_config = MasterServiceConfig::builder()
+                              .set_put_start_discard_timeout_sec(0)
+                              .set_put_start_release_timeout_sec(1)
+                              .build();
+    std::unique_ptr<MasterService> service_(new MasterService(service_config));
+    const auto context = PrepareSimpleSegment(*service_);
+    const UUID client_id = context.client_id;
+
+    const std::string key = "expired_grouped_put_to_ungrouped";
+    ReplicateConfig grouped_config;
+    grouped_config.replica_num = 1;
+    grouped_config.group_ids =
+        std::vector<std::string>{FindGroupIdOnDifferentShard(key)};
+
+    ASSERT_TRUE(
+        service_->PutStart(client_id, key, "default", 1024, grouped_config)
+            .has_value());
+    std::this_thread::sleep_for(std::chrono::milliseconds(2));
+
+    ReplicateConfig ungrouped_config;
+    ungrouped_config.replica_num = 1;
+    auto put_start =
+        service_->PutStart(client_id, key, "default", 1024, ungrouped_config);
+    ASSERT_TRUE(put_start.has_value()) << toString(put_start.error());
+    ASSERT_TRUE(service_->PutEnd(client_id, key, "default", ReplicaType::MEMORY)
+                    .has_value());
+    EXPECT_TRUE(service_->ExistKey(key, "default").value_or(false));
+}
+
+TEST_F(MasterServiceTest, BatchRemoveUnregistersGroupedRoute) {
+    std::unique_ptr<MasterService> service_(new MasterService());
+    [[maybe_unused]] const auto context = PrepareSimpleSegment(*service_);
+    const UUID client_id = generate_uuid();
+
+    const std::string key = "batch_remove_grouped_route";
+    ReplicateConfig grouped_config;
+    grouped_config.replica_num = 1;
+    grouped_config.group_ids =
+        std::vector<std::string>{FindGroupIdOnDifferentShard(key)};
+    PutCompletedObject(*service_, client_id, key, grouped_config);
+
+    auto remove_results = service_->BatchRemove(std::vector<std::string>{key},
+                                                "default", /*force=*/true);
+    ASSERT_EQ(remove_results.size(), 1u);
+    ASSERT_TRUE(remove_results[0].has_value());
+
+    ReplicateConfig ungrouped_config;
+    ungrouped_config.replica_num = 1;
+    PutCompletedObject(*service_, client_id, key, ungrouped_config);
+    EXPECT_TRUE(service_->GetReplicaList(key, "default").has_value());
+}
+
+TEST_F(MasterServiceTest, RemoveByRegexUnregistersGroupedRoute) {
+    std::unique_ptr<MasterService> service_(new MasterService());
+    [[maybe_unused]] const auto context = PrepareSimpleSegment(*service_);
+    const UUID client_id = generate_uuid();
+
+    const std::string key = "regex_remove_grouped_route";
+    ReplicateConfig grouped_config;
+    grouped_config.replica_num = 1;
+    grouped_config.group_ids =
+        std::vector<std::string>{FindGroupIdOnDifferentShard(key)};
+    PutCompletedObject(*service_, client_id, key, grouped_config);
+
+    auto removed =
+        service_->RemoveByRegex("^regex_remove_grouped_route$", "default",
+                                /*force=*/true);
+    ASSERT_TRUE(removed.has_value());
+    EXPECT_EQ(removed.value(), 1);
+
+    ReplicateConfig ungrouped_config;
+    ungrouped_config.replica_num = 1;
+    PutCompletedObject(*service_, client_id, key, ungrouped_config);
+    EXPECT_TRUE(service_->GetReplicaList(key, "default").has_value());
+}
+
+TEST_F(MasterServiceTest, GroupedLeaseRefreshNearExpiryProtectsCurrentMembers) {
+    auto service_config =
+        MasterServiceConfig::builder().set_default_kv_lease_ttl(200).build();
+    std::unique_ptr<MasterService> service_(new MasterService(service_config));
+    [[maybe_unused]] const auto context = PrepareSimpleSegment(*service_);
+    const UUID client_id = generate_uuid();
+
+    const std::string key_a = "lease_group_key_a";
+    const std::string key_b = "lease_group_key_b";
+    const std::string group_id = FindGroupIdOnDifferentShard(key_a);
+
+    ReplicateConfig config_a;
+    config_a.replica_num = 1;
+    config_a.group_ids = std::vector<std::string>{group_id};
+    ReplicateConfig config_b = config_a;
+
+    PutCompletedObject(*service_, client_id, key_a, config_a);
+    PutCompletedObject(*service_, client_id, key_b, config_b);
+
+    auto exists = service_->ExistKey(key_a, "default");
+    ASSERT_TRUE(exists.has_value());
+    ASSERT_TRUE(exists.value());
+
+    std::this_thread::sleep_for(std::chrono::milliseconds(120));
+    exists = service_->ExistKey(key_a, "default");
+    ASSERT_TRUE(exists.has_value());
+    ASSERT_TRUE(exists.value());
+    std::this_thread::sleep_for(std::chrono::milliseconds(100));
+
+    auto remove_group_peer = service_->Remove(key_b, "default");
+    ASSERT_FALSE(remove_group_peer.has_value());
+    EXPECT_EQ(ErrorCode::OBJECT_HAS_LEASE, remove_group_peer.error());
+
+    EXPECT_TRUE(service_->Remove(key_a, "default", /*force=*/true).has_value());
+    EXPECT_TRUE(service_->Remove(key_b, "default", /*force=*/true).has_value());
+}
+
+TEST_F(MasterServiceTest,
+       GroupedLeaseRefreshAfterMembershipChangeDoesNotWaitForTriggerExpiry) {
+    auto service_config =
+        MasterServiceConfig::builder().set_default_kv_lease_ttl(500).build();
+    std::unique_ptr<MasterService> service_(new MasterService(service_config));
+    [[maybe_unused]] const auto context = PrepareSimpleSegment(*service_);
+    const UUID client_id = generate_uuid();
+
+    const std::string key_a = "lease_group_dirty_key_a";
+    const std::string key_b = "lease_group_dirty_key_b";
+    const std::string group_id = FindGroupIdOnDifferentShard(key_a);
+
+    ReplicateConfig config;
+    config.replica_num = 1;
+    config.group_ids = std::vector<std::string>{group_id};
+
+    PutCompletedObject(*service_, client_id, key_a, config);
+    ASSERT_TRUE(service_->ExistKey(key_a, "default").value_or(false));
+
+    PutCompletedObject(*service_, client_id, key_b, config);
+    std::this_thread::sleep_for(std::chrono::milliseconds(150));
+
+    auto exists = service_->ExistKey(key_a, "default");
+    ASSERT_TRUE(exists.has_value());
+    ASSERT_TRUE(exists.value());
+    std::this_thread::sleep_for(std::chrono::milliseconds(390));
+
+    auto remove_group_peer = service_->Remove(key_b, "default");
+    ASSERT_FALSE(remove_group_peer.has_value());
+    EXPECT_EQ(ErrorCode::OBJECT_HAS_LEASE, remove_group_peer.error());
+
+    EXPECT_TRUE(service_->Remove(key_a, "default", /*force=*/true).has_value());
+    EXPECT_TRUE(service_->Remove(key_b, "default", /*force=*/true).has_value());
+}
+
+TEST_F(MasterServiceTest, RemoveGroupedMemberPreservesOtherMembers) {
+    std::unique_ptr<MasterService> service_(new MasterService());
+    [[maybe_unused]] const auto context = PrepareSimpleSegment(*service_);
+    const UUID client_id = generate_uuid();
+
+    const std::string key_a = "remove_group_key_a";
+    const std::string key_b = "remove_group_key_b";
+    const std::string group_id = FindGroupIdOnDifferentShard(key_a);
+
+    ReplicateConfig config;
+    config.replica_num = 1;
+    config.group_ids = std::vector<std::string>{group_id};
+    PutCompletedObject(*service_, client_id, key_a, config);
+    PutCompletedObject(*service_, client_id, key_b, config);
+
+    ASSERT_TRUE(service_->Remove(key_a, "default", /*force=*/true).has_value());
+
+    auto removed_exists = service_->ExistKey(key_a, "default");
+    ASSERT_TRUE(removed_exists.has_value());
+    EXPECT_FALSE(removed_exists.value());
+    EXPECT_TRUE(service_->GetReplicaList(key_b, "default").has_value());
+
+    ASSERT_TRUE(service_->Remove(key_b, "default", /*force=*/true).has_value());
+    auto group_empty = service_->ExistKey(key_b, "default");
+    ASSERT_TRUE(group_empty.has_value());
+    EXPECT_FALSE(group_empty.value());
+}
+
+TEST_F(MasterServiceTest, UpsertPreservesGroupMembership) {
+    std::unique_ptr<MasterService> service_(new MasterService());
+    [[maybe_unused]] const auto context = PrepareSimpleSegment(*service_);
+    const UUID client_id = generate_uuid();
+
+    const std::string key = "upsert_group_key";
+    const std::string group_id = FindGroupIdOnDifferentShard(key);
+
+    ReplicateConfig grouped_config;
+    grouped_config.replica_num = 1;
+    grouped_config.group_ids = std::vector<std::string>{group_id};
+    PutCompletedObject(*service_, client_id, key, grouped_config);
+
+    ReplicateConfig unset_group_config;
+    unset_group_config.replica_num = 1;
+    auto preserve_result = service_->UpsertStart(client_id, key, "default",
+                                                 1024, unset_group_config);
+    ASSERT_TRUE(preserve_result.has_value())
+        << "Unset group_ids should preserve existing group membership";
+    ASSERT_TRUE(
+        service_->UpsertEnd(client_id, key, "default", ReplicaType::MEMORY)
+            .has_value());
+    EXPECT_TRUE(service_->GetReplicaList(key, "default").has_value());
+
+    ReplicateConfig different_group_config;
+    different_group_config.replica_num = 1;
+    different_group_config.group_ids =
+        std::vector<std::string>{FindGroupIdOnDifferentShard(key + "_other")};
+    auto different_group_result = service_->UpsertStart(
+        client_id, key, "default", 1024, different_group_config);
+    ASSERT_FALSE(different_group_result.has_value());
+    EXPECT_EQ(ErrorCode::INVALID_PARAMS, different_group_result.error());
+
+    ReplicateConfig explicit_ungrouped_config;
+    explicit_ungrouped_config.replica_num = 1;
+    explicit_ungrouped_config.group_ids = std::vector<std::string>{""};
+    auto explicit_ungrouped_result = service_->UpsertStart(
+        client_id, key, "default", 1024, explicit_ungrouped_config);
+    ASSERT_FALSE(explicit_ungrouped_result.has_value());
+    EXPECT_EQ(ErrorCode::INVALID_PARAMS, explicit_ungrouped_result.error());
+}
+
+TEST_F(MasterServiceTest, IncompleteGroupedUpsertCanBecomeUngrouped) {
+    std::unique_ptr<MasterService> service_(new MasterService());
+    const auto context = PrepareSimpleSegment(*service_);
+    const UUID client_id = context.client_id;
+
+    const std::string key = "incomplete_grouped_upsert_to_ungrouped";
+    ReplicateConfig grouped_config;
+    grouped_config.replica_num = 1;
+    grouped_config.group_ids =
+        std::vector<std::string>{FindGroupIdOnDifferentShard(key)};
+
+    ASSERT_TRUE(
+        service_->PutStart(client_id, key, "default", 1024, grouped_config)
+            .has_value());
+    std::this_thread::sleep_for(std::chrono::milliseconds(2));
+
+    ReplicateConfig ungrouped_config;
+    ungrouped_config.replica_num = 1;
+    auto upsert_start = service_->UpsertStart(client_id, key, "default", 1024,
+                                              ungrouped_config);
+    ASSERT_TRUE(upsert_start.has_value()) << toString(upsert_start.error());
+    ASSERT_TRUE(
+        service_->UpsertEnd(client_id, key, "default", ReplicaType::MEMORY)
+            .has_value());
+    EXPECT_TRUE(service_->ExistKey(key, "default").value_or(false));
+}
+
+TEST_F(MasterServiceTest, UpsertRejectsExistingUngroupedToGrouped) {
+    std::unique_ptr<MasterService> service_(new MasterService());
+    [[maybe_unused]] const auto context = PrepareSimpleSegment(*service_);
+    const UUID client_id = generate_uuid();
+
+    const std::string key = "upsert_ungrouped_to_grouped";
+    ReplicateConfig ungrouped_config;
+    ungrouped_config.replica_num = 1;
+    PutCompletedObject(*service_, client_id, key, ungrouped_config);
+
+    ReplicateConfig grouped_config;
+    grouped_config.replica_num = 1;
+    grouped_config.group_ids =
+        std::vector<std::string>{FindGroupIdOnDifferentShard(key)};
+    auto upsert_start =
+        service_->UpsertStart(client_id, key, "default", 2048, grouped_config);
+    ASSERT_FALSE(upsert_start.has_value());
+    EXPECT_EQ(ErrorCode::INVALID_PARAMS, upsert_start.error());
+
+    EXPECT_TRUE(service_->GetReplicaList(key, "default").has_value());
+}
+
+TEST_F(MasterServiceTest,
+       GroupedEvictionExpandsSafeMembersAndSkipsLeasedGroup) {
+    auto service_config =
+        MasterServiceConfig::builder().set_default_kv_lease_ttl(1000).build();
+    constexpr size_t kSegmentSize = 4 * 1024 * 1024;
+    constexpr size_t kObjectSize = 2 * 1024 * 1024;
+
+    {
+        std::unique_ptr<MasterService> service_(
+            new MasterService(service_config));
+        [[maybe_unused]] const auto context =
+            PrepareSimpleSegment(*service_, "grouped_evict_segment",
+                                 kDefaultSegmentBase, kSegmentSize);
+        const UUID client_id = generate_uuid();
+
+        const std::string evict_key_a = "grouped_evict_key_a";
+        const std::string evict_key_b = "grouped_evict_key_b";
+        ReplicateConfig evict_config;
+        evict_config.replica_num = 1;
+        evict_config.group_ids =
+            std::vector<std::string>{FindGroupIdOnDifferentShard(evict_key_a)};
+        PutCompletedObject(*service_, client_id, evict_key_a, evict_config,
+                           kObjectSize);
+        PutCompletedObject(*service_, client_id, evict_key_b, evict_config,
+                           kObjectSize);
+
+        ReplicateConfig trigger_config;
+        trigger_config.replica_num = 1;
+        auto trigger_result =
+            service_->PutStart(client_id, "trigger_grouped_eviction", "default",
+                               kObjectSize, trigger_config);
+        ASSERT_FALSE(trigger_result.has_value());
+        EXPECT_EQ(ErrorCode::NO_AVAILABLE_HANDLE, trigger_result.error());
+
+        std::this_thread::sleep_for(std::chrono::milliseconds(200));
+
+        EXPECT_FALSE(service_->ExistKey(evict_key_a, "default").value_or(true));
+        EXPECT_FALSE(service_->ExistKey(evict_key_b, "default").value_or(true));
+    }
+
+    {
+        std::unique_ptr<MasterService> service_(
+            new MasterService(service_config));
+        [[maybe_unused]] const auto context =
+            PrepareSimpleSegment(*service_, "grouped_lease_segment",
+                                 kDefaultSegmentBase, kSegmentSize);
+        const UUID client_id = generate_uuid();
+
+        const std::string leased_key_a = "grouped_leased_key_a";
+        const std::string leased_key_b = "grouped_leased_key_b";
+        ReplicateConfig leased_config;
+        leased_config.replica_num = 1;
+        leased_config.group_ids =
+            std::vector<std::string>{FindGroupIdOnDifferentShard(leased_key_a)};
+        PutCompletedObject(*service_, client_id, leased_key_a, leased_config,
+                           kObjectSize);
+        PutCompletedObject(*service_, client_id, leased_key_b, leased_config,
+                           kObjectSize);
+
+        auto exists = service_->ExistKey(leased_key_a, "default");
+        ASSERT_TRUE(exists.has_value());
+        ASSERT_TRUE(exists.value());
+
+        ReplicateConfig trigger_config;
+        trigger_config.replica_num = 1;
+        auto trigger_result =
+            service_->PutStart(client_id, "trigger_leased_group_eviction",
+                               "default", kObjectSize, trigger_config);
+        ASSERT_FALSE(trigger_result.has_value());
+        EXPECT_EQ(ErrorCode::NO_AVAILABLE_HANDLE, trigger_result.error());
+
+        std::this_thread::sleep_for(std::chrono::milliseconds(200));
+
+        EXPECT_TRUE(
+            service_->GetReplicaList(leased_key_a, "default").has_value());
+        EXPECT_TRUE(
+            service_->GetReplicaList(leased_key_b, "default").has_value());
+    }
+}
+
+TEST_F(MasterServiceTest, GroupedEvictionSkipsUnsafeMembersAndEvictsSafePeers) {
+    constexpr size_t kSegmentSize = 4 * 1024 * 1024;
+    constexpr size_t kObjectSize = 2 * 1024 * 1024;
+    std::unique_ptr<MasterService> service_(new MasterService());
+    [[maybe_unused]] const auto context =
+        PrepareSimpleSegment(*service_, "grouped_mixed_safety_segment",
+                             kDefaultSegmentBase, kSegmentSize);
+    const UUID client_id = generate_uuid();
+
+    const std::string safe_key = "grouped_mixed_safe_key";
+    const std::string hard_pinned_key = "grouped_mixed_hard_pinned_key";
+    const std::string group_id = FindGroupIdOnDifferentShard(safe_key);
+
+    ReplicateConfig safe_config;
+    safe_config.replica_num = 1;
+    safe_config.group_ids = std::vector<std::string>{group_id};
+    PutCompletedObject(*service_, client_id, safe_key, safe_config,
+                       kObjectSize);
+
+    ReplicateConfig hard_pinned_config = safe_config;
+    hard_pinned_config.with_hard_pin = true;
+    PutCompletedObject(*service_, client_id, hard_pinned_key,
+                       hard_pinned_config, kObjectSize);
+
+    ReplicateConfig trigger_config;
+    trigger_config.replica_num = 1;
+    auto trigger_result =
+        service_->PutStart(client_id, "trigger_mixed_safety_group_eviction",
+                           "default", kObjectSize, trigger_config);
+    ASSERT_FALSE(trigger_result.has_value());
+    EXPECT_EQ(ErrorCode::NO_AVAILABLE_HANDLE, trigger_result.error());
+
+    std::this_thread::sleep_for(std::chrono::milliseconds(200));
+
+    EXPECT_FALSE(service_->ExistKey(safe_key, "default").value_or(true));
+    EXPECT_TRUE(
+        service_->GetReplicaList(hard_pinned_key, "default").has_value());
+    EXPECT_TRUE(service_->Remove(hard_pinned_key, "default", /*force=*/true)
+                    .has_value());
+}
+
+TEST_F(MasterServiceTest, BatchUpsertStartMixedGroupIdsPreservesOrder) {
+    std::unique_ptr<MasterService> service_(new MasterService());
+    [[maybe_unused]] const auto context = PrepareSimpleSegment(*service_);
+    const UUID client_id = generate_uuid();
+
+    const std::vector<std::string> keys = {
+        "batch_grouped_a",
+        "batch_ungrouped",
+        "batch_grouped_b",
+    };
+    const std::vector<uint64_t> sizes = {1024, 2048, 4096};
+
+    ReplicateConfig config;
+    config.replica_num = 1;
+    config.group_ids =
+        std::vector<std::string>{FindGroupIdOnDifferentShard(keys[0]), "",
+                                 FindGroupIdOnDifferentShard(keys[2])};
+
+    auto results =
+        service_->BatchUpsertStart(client_id, keys, "default", sizes, config);
+    ASSERT_EQ(results.size(), keys.size());
+    for (const auto& result : results) {
+        ASSERT_TRUE(result.has_value());
+    }
+
+    auto end_results = service_->BatchUpsertEnd(client_id, keys, "default");
+    ASSERT_EQ(end_results.size(), keys.size());
+    for (const auto& result : end_results) {
+        ASSERT_TRUE(result.has_value());
+    }
+
+    for (const auto& key : keys) {
+        EXPECT_TRUE(service_->GetReplicaList(key, "default").has_value());
+    }
+
+    ReplicateConfig invalid_config = config;
+    invalid_config.group_ids = std::vector<std::string>{"only_one"};
+    auto invalid_results = service_->BatchUpsertStart(
+        client_id, keys, "default", sizes, invalid_config);
+    ASSERT_EQ(invalid_results.size(), keys.size());
+    for (const auto& result : invalid_results) {
+        ASSERT_FALSE(result.has_value());
+        EXPECT_EQ(ErrorCode::INVALID_PARAMS, result.error());
+    }
+}
+
+TEST_F(MasterServiceTest, WrappedBatchPutStartMixedGroupIdsPreservesOrder) {
+    WrappedMasterServiceConfig service_config;
+    service_config.default_kv_lease_ttl = 100;
+    service_config.enable_metric_reporting = false;
+    WrappedMasterService service_(service_config);
+
+    Segment segment = MakeSegment("wrapped_batch_group_segment");
+    const UUID client_id = generate_uuid();
+    ASSERT_TRUE(service_.MountSegment(segment, client_id).has_value());
+
+    const std::vector<std::string> keys = {
+        "wrapped_batch_grouped_a",
+        "wrapped_batch_ungrouped",
+        "wrapped_batch_grouped_b",
+    };
+    const std::vector<uint64_t> sizes = {1024, 2048, 4096};
+
+    ReplicateConfig config;
+    config.replica_num = 1;
+    config.group_ids =
+        std::vector<std::string>{FindGroupIdOnDifferentShard(keys[0]), "",
+                                 FindGroupIdOnDifferentShard(keys[2])};
+
+    auto results = service_.BatchPutStart(client_id, keys, sizes, config);
+    ASSERT_EQ(results.size(), keys.size());
+    for (const auto& result : results) {
+        ASSERT_TRUE(result.has_value()) << toString(result.error());
+    }
+
+    auto end_results = service_.BatchPutEnd(client_id, keys);
+    ASSERT_EQ(end_results.size(), keys.size());
+    for (const auto& result : end_results) {
+        ASSERT_TRUE(result.has_value());
+    }
+
+    for (const auto& key : keys) {
+        EXPECT_TRUE(service_.GetReplicaList(key, "default").has_value());
+    }
+
+    ReplicateConfig invalid_config = config;
+    invalid_config.group_ids = std::vector<std::string>{"only_one"};
+    auto invalid_group_results =
+        service_.BatchPutStart(client_id, keys, sizes, invalid_config);
+    ASSERT_EQ(invalid_group_results.size(), keys.size());
+    for (const auto& result : invalid_group_results) {
+        ASSERT_FALSE(result.has_value());
+        EXPECT_EQ(ErrorCode::INVALID_PARAMS, result.error());
+    }
+
+    auto invalid_size_results =
+        service_.BatchPutStart(client_id, keys, {1024}, config);
+    ASSERT_EQ(invalid_size_results.size(), keys.size());
+    for (const auto& result : invalid_size_results) {
+        ASSERT_FALSE(result.has_value());
+        EXPECT_EQ(ErrorCode::INVALID_PARAMS, result.error());
+    }
 }
 
 TEST_F(MasterServiceTest, PutStartEndFlow) {
@@ -433,44 +1563,255 @@ TEST_F(MasterServiceTest, PutStartEndFlow) {
     config.replica_num = 1;
 
     auto put_start_result =
-        service_->PutStart(client_id, key, value_length, config);
+        service_->PutStart(client_id, key, "default", value_length, config);
     EXPECT_TRUE(put_start_result.has_value());
     replica_list = put_start_result.value();
     EXPECT_FALSE(replica_list.empty());
     EXPECT_EQ(ReplicaStatus::PROCESSING, replica_list[0].status);
 
     // During put, Get/Remove should fail
-    auto get_replica_result = service_->GetReplicaList(key);
+    auto get_replica_result = service_->GetReplicaList(key, "default");
     EXPECT_FALSE(get_replica_result.has_value());
     EXPECT_EQ(ErrorCode::REPLICA_IS_NOT_READY, get_replica_result.error());
-    auto remove_result = service_->Remove(key);
+    auto remove_result = service_->Remove(key, "default");
     EXPECT_FALSE(remove_result.has_value());
     EXPECT_EQ(ErrorCode::REPLICA_IS_NOT_READY, remove_result.error());
 
     // PutEnd should fail if the client_id does not match.
-    auto put_end_fail_result =
-        service_->PutEnd(invalid_client_id, key, ReplicaType::MEMORY);
+    auto put_end_fail_result = service_->PutEnd(invalid_client_id, key,
+                                                "default", ReplicaType::MEMORY);
     EXPECT_FALSE(put_end_fail_result.has_value());
     EXPECT_EQ(put_end_fail_result.error(), ErrorCode::ILLEGAL_CLIENT);
 
     // PutRevoke should fail if the client_id does not match.
-    auto put_revoke_fail_result =
-        service_->PutRevoke(invalid_client_id, key, ReplicaType::MEMORY);
+    auto put_revoke_fail_result = service_->PutRevoke(
+        invalid_client_id, key, "default", ReplicaType::MEMORY);
     EXPECT_FALSE(put_revoke_fail_result.has_value());
     EXPECT_EQ(put_revoke_fail_result.error(), ErrorCode::ILLEGAL_CLIENT);
 
     // Test PutEnd
-    auto put_end_result = service_->PutEnd(client_id, key, ReplicaType::MEMORY);
+    auto put_end_result =
+        service_->PutEnd(client_id, key, "default", ReplicaType::MEMORY);
     EXPECT_TRUE(put_end_result.has_value());
 
     // Verify replica list after PutEnd
-    auto final_get_result = service_->GetReplicaList(key);
+    auto final_get_result = service_->GetReplicaList(key, "default");
     EXPECT_TRUE(final_get_result.has_value());
     replica_list = final_get_result.value().replicas;
     EXPECT_EQ(1, replica_list.size());
     EXPECT_EQ(ReplicaStatus::COMPLETE, replica_list[0].status);
 }
 
+TEST_F(MasterServiceTest, TenantPutGetRemoveIsolatesSameUserKey) {
+    const std::string key = "shared_user_key";
+    const std::string tenant_a = "tenant_a";
+    const std::string tenant_b = "tenant_b";
+    auto service_ = std::make_unique<MasterService>(
+        MakeStrictTenantConfig({"default", tenant_a, tenant_b}));
+    [[maybe_unused]] const auto context = PrepareSimpleSegment(*service_);
+    const UUID client_id = generate_uuid();
+
+    ReplicateConfig config;
+    config.replica_num = 1;
+
+    ASSERT_TRUE(
+        service_->PutStart(client_id, key, tenant_a, 1024, config).has_value());
+    ASSERT_TRUE(service_->PutEnd(client_id, key, tenant_a, ReplicaType::MEMORY)
+                    .has_value());
+    ASSERT_TRUE(
+        service_->PutStart(client_id, key, tenant_b, 2048, config).has_value());
+    ASSERT_TRUE(service_->PutEnd(client_id, key, tenant_b, ReplicaType::MEMORY)
+                    .has_value());
+
+    EXPECT_FALSE(service_->GetReplicaList(key, "default").has_value());
+    EXPECT_FALSE(service_->ExistKey(key, "default").value());
+    EXPECT_TRUE(service_->ExistKey(key, tenant_a).value());
+    EXPECT_TRUE(service_->ExistKey(key, tenant_b).value());
+    EXPECT_TRUE(service_->GetReplicaList(key, tenant_a).has_value());
+    EXPECT_TRUE(service_->GetReplicaList(key, tenant_b).has_value());
+    EXPECT_EQ(service_->GetKeyCount(), 2u);
+
+    ASSERT_TRUE(service_->Remove(key, tenant_a, /*force=*/true).has_value());
+    EXPECT_FALSE(service_->GetReplicaList(key, tenant_a).has_value());
+    EXPECT_TRUE(service_->GetReplicaList(key, tenant_b).has_value());
+    EXPECT_EQ(service_->GetKeyCount(), 1u);
+}
+
+TEST_F(MasterServiceTest, RegexOperationsAreTenantScoped) {
+    const std::string key = "regex_shared_key";
+    const std::string tenant_a = "tenant_regex_a";
+    const std::string tenant_b = "tenant_regex_b";
+    auto service_ = std::make_unique<MasterService>(
+        MakeStrictTenantConfig({"default", tenant_a, tenant_b}));
+    [[maybe_unused]] const auto context = PrepareSimpleSegment(*service_);
+    const UUID client_id = generate_uuid();
+
+    ReplicateConfig config;
+    config.replica_num = 1;
+
+    ASSERT_TRUE(service_->PutStart(client_id, key, "default", 1024, config)
+                    .has_value());
+    ASSERT_TRUE(service_->PutEnd(client_id, key, "default", ReplicaType::MEMORY)
+                    .has_value());
+    ASSERT_TRUE(
+        service_->PutStart(client_id, key, tenant_a, 1024, config).has_value());
+    ASSERT_TRUE(service_->PutEnd(client_id, key, tenant_a, ReplicaType::MEMORY)
+                    .has_value());
+    ASSERT_TRUE(
+        service_->PutStart(client_id, key, tenant_b, 1024, config).has_value());
+    ASSERT_TRUE(service_->PutEnd(client_id, key, tenant_b, ReplicaType::MEMORY)
+                    .has_value());
+
+    auto default_matches =
+        service_->GetReplicaListByRegex("^regex_shared", "default");
+    ASSERT_TRUE(default_matches.has_value());
+    EXPECT_EQ(default_matches->size(), 1);
+
+    auto remove_default =
+        service_->RemoveByRegex("^regex_shared", "default", /*force=*/true);
+    ASSERT_TRUE(remove_default.has_value());
+    EXPECT_EQ(remove_default.value(), 1);
+    EXPECT_FALSE(service_->GetReplicaList(key, "default").has_value());
+    EXPECT_TRUE(service_->GetReplicaList(key, tenant_a).has_value());
+    EXPECT_TRUE(service_->GetReplicaList(key, tenant_b).has_value());
+
+    auto remove_tenant_a =
+        service_->RemoveByRegex("^regex_shared", tenant_a, /*force=*/true);
+    ASSERT_TRUE(remove_tenant_a.has_value());
+    EXPECT_EQ(remove_tenant_a.value(), 1);
+    EXPECT_FALSE(service_->GetReplicaList(key, tenant_a).has_value());
+    EXPECT_TRUE(service_->GetReplicaList(key, tenant_b).has_value());
+}
+
+TEST_F(MasterServiceTest, TenantBatchUpsertAndRevokeAreScoped) {
+    const std::vector<std::string> keys = {"tenant_batch_upsert_key_a",
+                                           "tenant_batch_upsert_key_b"};
+    const std::vector<uint64_t> sizes = {1024, 2048};
+    const std::string tenant_a = "tenant_batch_upsert_a";
+    const std::string tenant_b = "tenant_batch_upsert_b";
+    auto svc = std::make_unique<MasterService>(
+        MakeStrictTenantConfig({"default", tenant_a, tenant_b}));
+    [[maybe_unused]] const auto context = PrepareSimpleSegment(*svc);
+    const UUID client_id = generate_uuid();
+
+    ReplicateConfig config;
+    config.replica_num = 1;
+
+    auto tenant_a_results =
+        svc->BatchUpsertStart(client_id, keys, tenant_a, sizes, config);
+    ASSERT_EQ(tenant_a_results.size(), keys.size());
+    for (const auto& result : tenant_a_results) {
+        ASSERT_TRUE(result.has_value());
+    }
+    auto tenant_a_end = svc->BatchUpsertEnd(client_id, keys, tenant_a);
+    ASSERT_EQ(tenant_a_end.size(), keys.size());
+    for (const auto& result : tenant_a_end) {
+        ASSERT_TRUE(result.has_value());
+    }
+
+    auto tenant_b_results =
+        svc->BatchUpsertStart(client_id, keys, tenant_b, sizes, config);
+    ASSERT_EQ(tenant_b_results.size(), keys.size());
+    for (const auto& result : tenant_b_results) {
+        ASSERT_TRUE(result.has_value());
+    }
+    auto tenant_b_end = svc->BatchUpsertEnd(client_id, keys, tenant_b);
+    ASSERT_EQ(tenant_b_end.size(), keys.size());
+    for (const auto& result : tenant_b_end) {
+        ASSERT_TRUE(result.has_value());
+    }
+
+    for (const auto& key : keys) {
+        EXPECT_FALSE(svc->GetReplicaList(key, "default").has_value());
+        EXPECT_TRUE(svc->GetReplicaList(key, tenant_a).has_value());
+        EXPECT_TRUE(svc->GetReplicaList(key, tenant_b).has_value());
+    }
+
+    const std::string revoke_key = "tenant_batch_upsert_revoke_key";
+    auto revoke_start =
+        svc->UpsertStart(client_id, revoke_key, tenant_a, 1024, config);
+    ASSERT_TRUE(revoke_start.has_value());
+    ASSERT_TRUE(
+        svc->UpsertRevoke(client_id, revoke_key, tenant_a, ReplicaType::MEMORY)
+            .has_value());
+    EXPECT_FALSE(svc->GetReplicaList(revoke_key, tenant_a).has_value());
+}
+
+TEST_F(MasterServiceTest, TenantBatchRemoveAndRemoveAllAreScoped) {
+    const std::string shared_key = "tenant_batch_remove_shared_key";
+    const std::string tenant_a = "tenant_batch_remove_a";
+    const std::string tenant_b = "tenant_batch_remove_b";
+    auto svc = std::make_unique<MasterService>(
+        MakeStrictTenantConfig({"default", tenant_a, tenant_b}));
+    [[maybe_unused]] const auto context = PrepareSimpleSegment(*svc);
+    const UUID client_id = generate_uuid();
+
+    ReplicateConfig config;
+    config.replica_num = 1;
+
+    ASSERT_TRUE(svc->PutStart(client_id, shared_key, "default", 1024, config)
+                    .has_value());
+    ASSERT_TRUE(
+        svc->PutEnd(client_id, shared_key, "default", ReplicaType::MEMORY)
+            .has_value());
+    ASSERT_TRUE(svc->PutStart(client_id, shared_key, tenant_a, 1024, config)
+                    .has_value());
+    ASSERT_TRUE(
+        svc->PutEnd(client_id, shared_key, tenant_a, ReplicaType::MEMORY)
+            .has_value());
+    ASSERT_TRUE(svc->PutStart(client_id, shared_key, tenant_b, 1024, config)
+                    .has_value());
+    ASSERT_TRUE(
+        svc->PutEnd(client_id, shared_key, tenant_b, ReplicaType::MEMORY)
+            .has_value());
+
+    auto remove_a = svc->BatchRemove({shared_key}, tenant_a, /*force=*/true);
+    ASSERT_EQ(remove_a.size(), 1u);
+    ASSERT_TRUE(remove_a[0].has_value());
+    EXPECT_FALSE(svc->GetReplicaList(shared_key, tenant_a).has_value());
+    EXPECT_TRUE(svc->GetReplicaList(shared_key, "default").has_value());
+    EXPECT_TRUE(svc->GetReplicaList(shared_key, tenant_b).has_value());
+
+    EXPECT_EQ(svc->RemoveAll(tenant_b, /*force=*/true), 1);
+    EXPECT_FALSE(svc->GetReplicaList(shared_key, tenant_b).has_value());
+    EXPECT_TRUE(svc->GetReplicaList(shared_key, "default").has_value());
+
+    EXPECT_EQ(svc->RemoveAll(/*force=*/true), 1);
+    EXPECT_FALSE(svc->GetReplicaList(shared_key, "default").has_value());
+}
+
+TEST_F(MasterServiceTest, LegacyRemoveAllRemovesAllTenants) {
+    const std::string key = "legacy_remove_all_shared_key";
+    const std::string tenant_a = "legacy_remove_all_a";
+    const std::string tenant_b = "legacy_remove_all_b";
+    auto svc = std::make_unique<MasterService>(
+        MakeStrictTenantConfig({"default", tenant_a, tenant_b}));
+    [[maybe_unused]] const auto context = PrepareSimpleSegment(*svc);
+    const UUID client_id = generate_uuid();
+
+    ReplicateConfig config;
+    config.replica_num = 1;
+
+    ASSERT_TRUE(
+        svc->PutStart(client_id, key, "default", 1024, config).has_value());
+    ASSERT_TRUE(svc->PutEnd(client_id, key, "default", ReplicaType::MEMORY)
+                    .has_value());
+    ASSERT_TRUE(
+        svc->PutStart(client_id, key, tenant_a, 1024, config).has_value());
+    ASSERT_TRUE(
+        svc->PutEnd(client_id, key, tenant_a, ReplicaType::MEMORY).has_value());
+    ASSERT_TRUE(
+        svc->PutStart(client_id, key, tenant_b, 1024, config).has_value());
+    ASSERT_TRUE(
+        svc->PutEnd(client_id, key, tenant_b, ReplicaType::MEMORY).has_value());
+
+    EXPECT_EQ(svc->RemoveAll(/*force=*/true), 3);
+    EXPECT_FALSE(svc->GetReplicaList(key, "default").has_value());
+    EXPECT_FALSE(svc->GetReplicaList(key, tenant_a).has_value());
+    EXPECT_FALSE(svc->GetReplicaList(key, tenant_b).has_value());
+    EXPECT_EQ(svc->RemoveAll(/*force=*/true), 0);
+}
+
 TEST_F(MasterServiceTest, PutWithPreferredSegment) {
     // For backward compatibility, test the deprecated single preferred_segment
     std::unique_ptr<MasterService> service_(new MasterService());
@@ -496,7 +1837,7 @@ TEST_F(MasterServiceTest, PutWithPreferredSegment) {
     config.preferred_segment = preferred_segment;
 
     auto put_start_result =
-        service_->PutStart(client_id, key, value_length, config);
+        service_->PutStart(client_id, key, "default", value_length, config);
     EXPECT_TRUE(put_start_result.has_value());
     replica_list = put_start_result.value();
     EXPECT_EQ(1, replica_list.size());
@@ -506,7 +1847,8 @@ TEST_F(MasterServiceTest, PutWithPreferredSegment) {
               mem_desc.buffer_descriptor.transport_endpoint_);
 
     // Complete the Put operation
-    auto put_end_result = service_->PutEnd(client_id, key, ReplicaType::MEMORY);
+    auto put_end_result =
+        service_->PutEnd(client_id, key, "default", ReplicaType::MEMORY);
     EXPECT_TRUE(put_end_result.has_value());
 }
 
@@ -534,7 +1876,7 @@ TEST_F(MasterServiceTest, PutWithPreferredSegments) {
     config.preferred_segments = preferred_segments;
 
     auto put_start_result =
-        service_->PutStart(client_id, key, value_length, config);
+        service_->PutStart(client_id, key, "default", value_length, config);
     EXPECT_TRUE(put_start_result.has_value());
     replica_list = put_start_result.value();
     EXPECT_EQ(2, replica_list.size());
@@ -549,7 +1891,8 @@ TEST_F(MasterServiceTest, PutWithPreferredSegments) {
     EXPECT_TRUE(used_segments.find("segment_1") != used_segments.end());
 
     // Complete the Put operation
-    auto put_end_result = service_->PutEnd(client_id, key, ReplicaType::MEMORY);
+    auto put_end_result =
+        service_->PutEnd(client_id, key, "default", ReplicaType::MEMORY);
     EXPECT_TRUE(put_end_result.has_value());
 }
 
@@ -576,23 +1919,24 @@ TEST_F(MasterServiceTest, RandomPutStartEndFlow) {
     int random_number = dis(gen);
     config.replica_num = random_number;
     auto put_start_result =
-        service_->PutStart(client_id, key, value_length, config);
+        service_->PutStart(client_id, key, "default", value_length, config);
     EXPECT_TRUE(put_start_result.has_value());
     replica_list = put_start_result.value();
     EXPECT_FALSE(replica_list.empty());
     EXPECT_EQ(ReplicaStatus::PROCESSING, replica_list[0].status);
     // During put, Get/Remove should fail
-    auto get_result = service_->GetReplicaList(key);
+    auto get_result = service_->GetReplicaList(key, "default");
     EXPECT_FALSE(get_result.has_value());
     EXPECT_EQ(ErrorCode::REPLICA_IS_NOT_READY, get_result.error());
-    auto remove_result = service_->Remove(key);
+    auto remove_result = service_->Remove(key, "default");
     EXPECT_FALSE(remove_result.has_value());
     EXPECT_EQ(ErrorCode::REPLICA_IS_NOT_READY, remove_result.error());
     // Test PutEnd
-    auto put_end_result = service_->PutEnd(client_id, key, ReplicaType::MEMORY);
+    auto put_end_result =
+        service_->PutEnd(client_id, key, "default", ReplicaType::MEMORY);
     EXPECT_TRUE(put_end_result.has_value());
     // Verify replica list after PutEnd
-    auto get_result2 = service_->GetReplicaList(key);
+    auto get_result2 = service_->GetReplicaList(key, "default");
     EXPECT_TRUE(get_result2.has_value());
     replica_list = get_result2.value().replicas;
     EXPECT_EQ(random_number, replica_list.size());
@@ -609,7 +1953,7 @@ TEST_F(MasterServiceTest, GetReplicaListByRegex) {
     std::unique_ptr<MasterService> service_(new MasterService(service_config));
     const UUID client_id = generate_uuid();
     // Test getting non-existent key
-    auto get_result = service_->GetReplicaList(".*non_existent.*");
+    auto get_result = service_->GetReplicaList(".*non_existent.*", "default");
     EXPECT_FALSE(get_result.has_value());
     EXPECT_EQ(ErrorCode::OBJECT_NOT_FOUND, get_result.error());
 
@@ -622,19 +1966,19 @@ TEST_F(MasterServiceTest, GetReplicaListByRegex) {
         ReplicateConfig config;
         config.replica_num = 1;
         auto put_start_result =
-            service_->PutStart(client_id, key, value_length, config);
+            service_->PutStart(client_id, key, "default", value_length, config);
         ASSERT_TRUE(put_start_result.has_value());
         auto put_end_result =
-            service_->PutEnd(client_id, key, ReplicaType::MEMORY);
+            service_->PutEnd(client_id, key, "default", ReplicaType::MEMORY);
         ASSERT_TRUE(put_end_result.has_value());
-        auto exist_result = service_->ExistKey(key);
+        auto exist_result = service_->ExistKey(key, "default");
         ASSERT_TRUE(exist_result.has_value());
     }
     // wait for all the lease to expire
     std::this_thread::sleep_for(std::chrono::milliseconds(kv_lease_ttl));
 
     // Test getting existing key
-    auto get_result2 = service_->GetReplicaListByRegex("^test_key");
+    auto get_result2 = service_->GetReplicaListByRegex("^test_key", "default");
     EXPECT_TRUE(get_result2.has_value());
     auto replica_list_local = get_result2.value();
     EXPECT_EQ(10, replica_list_local.size());
@@ -647,13 +1991,14 @@ void put_object(MasterService& service, const UUID& client_id,
     ReplicateConfig config;
     config.replica_num = 1;
     auto put_start_result =
-        service.PutStart(client_id, key, value_length, config);
+        service.PutStart(client_id, key, "default", value_length, config);
     ASSERT_TRUE(put_start_result.has_value())
         << "Failed to PutStart for key: " << key;
-    auto put_end_result = service.PutEnd(client_id, key, ReplicaType::MEMORY);
+    auto put_end_result =
+        service.PutEnd(client_id, key, "default", ReplicaType::MEMORY);
     ASSERT_TRUE(put_end_result.has_value())
         << "Failed to PutEnd for key: " << key;
-    auto exist_result = service.ExistKey(key);
+    auto exist_result = service.ExistKey(key, "default");
     ASSERT_TRUE(exist_result.has_value())
         << "Key does not exist after put: " << key;
 }
@@ -697,7 +2042,7 @@ TEST_F(MasterServiceTest, GetReplicaListByRegexComplex) {
 
     // Test 3.1: Simple prefix matching
     {
-        auto result = service_->GetReplicaListByRegex("^test_key_");
+        auto result = service_->GetReplicaListByRegex("^test_key_", "default");
         ASSERT_TRUE(result.has_value());
         EXPECT_EQ(result.value().size(),
                   3);  // Matches test_key_01, test_key_02, test_key_10
@@ -705,7 +2050,8 @@ TEST_F(MasterServiceTest, GetReplicaListByRegexComplex) {
 
     // Test 3.2: Matching with a wildcard for any number
     {
-        auto result = service_->GetReplicaListByRegex("^test_key_\\d+$");
+        auto result =
+            service_->GetReplicaListByRegex("^test_key_\\d+$", "default");
         ASSERT_TRUE(result.has_value());
         EXPECT_EQ(result.value().size(), 3);
     }
@@ -713,8 +2059,8 @@ TEST_F(MasterServiceTest, GetReplicaListByRegexComplex) {
     // Test 3.3: Matching a specific pattern with wildcards
     {
         // Matches "data_part_1_chunk_a" and "data_part_2_chunk_b"
-        auto result =
-            service_->GetReplicaListByRegex("^data_part_\\d_chunk_.$");
+        auto result = service_->GetReplicaListByRegex("^data_part_\\d_chunk_.$",
+                                                      "default");
         ASSERT_TRUE(result.has_value());
         EXPECT_EQ(result.value().size(), 2);
     }
@@ -722,7 +2068,7 @@ TEST_F(MasterServiceTest, GetReplicaListByRegexComplex) {
     // Test 3.4: Matching keys containing a specific substring
     {
         // Matches all keys with "key" in them
-        auto result = service_->GetReplicaListByRegex("key");
+        auto result = service_->GetReplicaListByRegex("key", "default");
         ASSERT_TRUE(result.has_value());
         // Expected: test_key_01, test_key_02, test_key_10,
         //           prod_key_alpha, prod_key_beta,
@@ -734,7 +2080,7 @@ TEST_F(MasterServiceTest, GetReplicaListByRegexComplex) {
     // Test 3.5: Matching based on file-like paths
     {
         // Match all .log files
-        auto result = service_->GetReplicaListByRegex("\\.log$");
+        auto result = service_->GetReplicaListByRegex("\\.log$", "default");
         ASSERT_TRUE(result.has_value());
         EXPECT_EQ(result.value().size(), 1);
         EXPECT_EQ(result.value().begin()->first, "logs/app-2025-08-13.log");
@@ -743,7 +2089,8 @@ TEST_F(MasterServiceTest, GetReplicaListByRegexComplex) {
     // Test 3.6: OR condition using |
     {
         // Match keys starting with "prod" OR ending with "json"
-        auto result = service_->GetReplicaListByRegex("^prod|\\.json$");
+        auto result =
+            service_->GetReplicaListByRegex("^prod|\\.json$", "default");
         ASSERT_TRUE(result.has_value());
         // Expected: prod_key_alpha, prod_key_beta, config/user/settings.json
         EXPECT_EQ(result.value().size(), 3);
@@ -751,7 +2098,8 @@ TEST_F(MasterServiceTest, GetReplicaListByRegexComplex) {
 
     // Test 3.7: Regex that should not match anything
     {
-        auto result = service_->GetReplicaListByRegex("^non_existent_prefix_");
+        auto result =
+            service_->GetReplicaListByRegex("^non_existent_prefix_", "default");
         // This should succeed but return an empty map.
         ASSERT_TRUE(result.has_value());
         EXPECT_TRUE(result.value().empty());
@@ -759,7 +2107,7 @@ TEST_F(MasterServiceTest, GetReplicaListByRegexComplex) {
 
     // Test 3.8: Exact match regex
     {
-        auto result = service_->GetReplicaListByRegex("^short$");
+        auto result = service_->GetReplicaListByRegex("^short$", "default");
         ASSERT_TRUE(result.has_value());
         EXPECT_EQ(result.value().size(), 1);
         EXPECT_EQ(result.value().begin()->first, "short");
@@ -767,8 +2115,8 @@ TEST_F(MasterServiceTest, GetReplicaListByRegexComplex) {
 
     // Test 3.9: Initial test for non-existent key (as a sanity check)
     {
-        auto get_result =
-            service_->GetReplicaListByRegex(".*absolutely_non_existent.*");
+        auto get_result = service_->GetReplicaListByRegex(
+            ".*absolutely_non_existent.*", "default");
         // Depending on implementation, this could return an empty map or an
         // error. Let's assume it returns an empty map for a valid regex with no
         // matches.
@@ -781,7 +2129,7 @@ TEST_F(MasterServiceTest, GetReplicaList) {
     std::unique_ptr<MasterService> service_(new MasterService());
     const UUID client_id = generate_uuid();
     // Test getting non-existent key
-    auto get_result = service_->GetReplicaList("non_existent");
+    auto get_result = service_->GetReplicaList("non_existent", "default");
     EXPECT_FALSE(get_result.has_value());
     EXPECT_EQ(ErrorCode::OBJECT_NOT_FOUND, get_result.error());
 
@@ -792,13 +2140,14 @@ TEST_F(MasterServiceTest, GetReplicaList) {
     ReplicateConfig config;
     config.replica_num = 1;
     auto put_start_result =
-        service_->PutStart(client_id, key, value_length, config);
+        service_->PutStart(client_id, key, "default", value_length, config);
     ASSERT_TRUE(put_start_result.has_value());
-    auto put_end_result = service_->PutEnd(client_id, key, ReplicaType::MEMORY);
+    auto put_end_result =
+        service_->PutEnd(client_id, key, "default", ReplicaType::MEMORY);
     ASSERT_TRUE(put_end_result.has_value());
 
     // Test getting existing key
-    auto get_result2 = service_->GetReplicaList(key);
+    auto get_result2 = service_->GetReplicaList(key, "default");
     EXPECT_TRUE(get_result2.has_value());
     auto replica_list_local = get_result2.value().replicas;
     EXPECT_FALSE(replica_list_local.empty());
@@ -814,22 +2163,23 @@ TEST_F(MasterServiceTest, RemoveObject) {
     ReplicateConfig config;
     config.replica_num = 1;
     auto put_start_result =
-        service_->PutStart(client_id, key, value_length, config);
+        service_->PutStart(client_id, key, "default", value_length, config);
     ASSERT_TRUE(put_start_result.has_value());
-    auto put_end_result = service_->PutEnd(client_id, key, ReplicaType::MEMORY);
+    auto put_end_result =
+        service_->PutEnd(client_id, key, "default", ReplicaType::MEMORY);
     ASSERT_TRUE(put_end_result.has_value());
 
     // Test removing the object
-    auto remove_result = service_->Remove(key);
+    auto remove_result = service_->Remove(key, "default");
     EXPECT_TRUE(remove_result.has_value());
 
     // Verify object is removed
-    auto get_result = service_->GetReplicaList(key);
+    auto get_result = service_->GetReplicaList(key, "default");
     EXPECT_FALSE(get_result.has_value());
     EXPECT_EQ(ErrorCode::OBJECT_NOT_FOUND, get_result.error());
 
     // Test removing non-existent object
-    auto remove_result2 = service_->Remove("non_existent");
+    auto remove_result2 = service_->Remove("non_existent", "default");
     EXPECT_FALSE(remove_result2.has_value());
     EXPECT_EQ(ErrorCode::OBJECT_NOT_FOUND, remove_result2.error());
 }
@@ -848,18 +2198,18 @@ TEST_F(MasterServiceTest, RandomRemoveObject) {
         ReplicateConfig config;
         config.replica_num = 1;
         auto put_start_result =
-            service_->PutStart(client_id, key, value_length, config);
+            service_->PutStart(client_id, key, "default", value_length, config);
         ASSERT_TRUE(put_start_result.has_value());
         auto put_end_result =
-            service_->PutEnd(client_id, key, ReplicaType::MEMORY);
+            service_->PutEnd(client_id, key, "default", ReplicaType::MEMORY);
         ASSERT_TRUE(put_end_result.has_value());
 
         // Test removing the object
-        auto remove_result = service_->Remove(key);
+        auto remove_result = service_->Remove(key, "default");
         EXPECT_TRUE(remove_result.has_value());
 
         // Verify object is removed
-        auto get_result = service_->GetReplicaList(key);
+        auto get_result = service_->GetReplicaList(key, "default");
         EXPECT_FALSE(get_result.has_value());
         EXPECT_EQ(ErrorCode::OBJECT_NOT_FOUND, get_result.error());
     }
@@ -880,23 +2230,23 @@ TEST_F(MasterServiceTest, RemoveByRegex) {
         ReplicateConfig config;
         config.replica_num = 1;
         auto put_start_result =
-            service_->PutStart(client_id, key, value_length, config);
+            service_->PutStart(client_id, key, "default", value_length, config);
         ASSERT_TRUE(put_start_result.has_value());
         auto put_end_result =
-            service_->PutEnd(client_id, key, ReplicaType::MEMORY);
+            service_->PutEnd(client_id, key, "default", ReplicaType::MEMORY);
         ASSERT_TRUE(put_end_result.has_value());
-        auto exist_result = service_->ExistKey(key);
+        auto exist_result = service_->ExistKey(key, "default");
         ASSERT_TRUE(exist_result.has_value());
     }
     // wait for all the lease to expire
     std::this_thread::sleep_for(std::chrono::milliseconds(kv_lease_ttl));
-    auto res = service_->RemoveByRegex("^test_key");
+    auto res = service_->RemoveByRegex("^test_key", "default");
     ASSERT_TRUE(res.has_value());
     ASSERT_EQ(10, res.value());
     times = 10;
     while (times--) {
         std::string key = "test_key" + std::to_string(times);
-        auto exist_result = service_->ExistKey(key);
+        auto exist_result = service_->ExistKey(key, "default");
         ASSERT_TRUE(exist_result.has_value());
         ASSERT_FALSE(exist_result.value());
     }
@@ -923,8 +2273,8 @@ TEST_F(MasterServiceTest, CopyStart) {
     UUID client_id = generate_uuid();
 
     // Test Case 1: CopyStart a non-existent key, should fail.
-    auto copy_result = service_->CopyStart(client_id, "non_existent_key",
-                                           "segment_1", {"segment_2"});
+    auto copy_result = service_->CopyStart(
+        client_id, "non_existent_key", "default", "segment_1", {"segment_2"});
     EXPECT_FALSE(copy_result.has_value());
     EXPECT_EQ(ErrorCode::OBJECT_NOT_FOUND, copy_result.error());
 
@@ -937,22 +2287,23 @@ TEST_F(MasterServiceTest, CopyStart) {
     config.preferred_segment = "segment_1";
 
     auto put_start_result =
-        service_->PutStart(client_id, key, slice_length, config);
+        service_->PutStart(client_id, key, "default", slice_length, config);
     ASSERT_TRUE(put_start_result.has_value());
 
     // Test Case 2: CopyStart to segment_2 and segment_3, should fail because
     // the only replica is not completed.
-    copy_result = service_->CopyStart(client_id, key, "segment_1",
+    copy_result = service_->CopyStart(client_id, key, "default", "segment_1",
                                       {"segment_2", "segment_3"});
     EXPECT_FALSE(copy_result.has_value());
     EXPECT_EQ(ErrorCode::REPLICA_NOT_FOUND, copy_result.error());
 
     // PutEnd the object.
-    auto put_end_result = service_->PutEnd(client_id, key, ReplicaType::MEMORY);
+    auto put_end_result =
+        service_->PutEnd(client_id, key, "default", ReplicaType::MEMORY);
     ASSERT_TRUE(put_end_result.has_value());
 
     // Test Case 3: CopyStart to segment_2 and segment_3, should success.
-    copy_result = service_->CopyStart(client_id, key, "segment_1",
+    copy_result = service_->CopyStart(client_id, key, "default", "segment_1",
                                       {"segment_2", "segment_3"});
     EXPECT_TRUE(copy_result.has_value());
     auto copy_response = copy_result.value();
@@ -961,40 +2312,41 @@ TEST_F(MasterServiceTest, CopyStart) {
     EXPECT_EQ(2, copy_response.targets.size());
 
     // Test Case 4: Try remove the object, should fail because it is copying.
-    auto remove_result = service_->Remove(key);
+    auto remove_result = service_->Remove(key, "default");
     EXPECT_FALSE(remove_result.has_value());
     EXPECT_EQ(ErrorCode::REPLICA_IS_NOT_READY, remove_result.error());
 
     // Test Case 5: CopyStart to segment_4, should fail because there is an
     // ongoing copy task.
-    copy_result =
-        service_->CopyStart(client_id, key, "segment_1", {"segment_4"});
+    copy_result = service_->CopyStart(client_id, key, "default", "segment_1",
+                                      {"segment_4"});
     EXPECT_FALSE(copy_result.has_value());
     EXPECT_EQ(ErrorCode::OBJECT_HAS_REPLICATION_TASK, copy_result.error());
 
     // Test Case 6: CopyEnd, should success and the object now has 3 replicas.
-    auto copy_end_result = service_->CopyEnd(client_id, key);
+    auto copy_end_result = service_->CopyEnd(client_id, key, "default");
     EXPECT_TRUE(copy_end_result.has_value());
-    auto get_result = service_->GetReplicaList(key);
+    auto get_result = service_->GetReplicaList(key, "default");
     EXPECT_TRUE(get_result.has_value());
     EXPECT_EQ(3, get_result.value().replicas.size());
 
     // Test Case 7: Copy from a non-existent replica to segment_3 and
     // segment_4, should fail.
-    copy_result = service_->CopyStart(client_id, key, "non_existent_segment",
-                                      {"segment_3", "segment_4"});
+    copy_result =
+        service_->CopyStart(client_id, key, "default", "non_existent_segment",
+                            {"segment_3", "segment_4"});
     EXPECT_FALSE(copy_result.has_value());
     EXPECT_EQ(ErrorCode::REPLICA_NOT_FOUND, copy_result.error());
 
     // Test Case 8: Copy to segment_4 and a non-existent segment, should fail.
-    copy_result = service_->CopyStart(client_id, key, "segment_1",
+    copy_result = service_->CopyStart(client_id, key, "default", "segment_1",
                                       {"segment_4", "non_existent_segment"});
     EXPECT_FALSE(copy_result.has_value());
     EXPECT_EQ(ErrorCode::SEGMENT_NOT_FOUND, copy_result.error());
 
     // Test Case 9: Copy to segment_3 and segment_4, should skip segment_3 and
     // successfully copy to segment_4.
-    copy_result = service_->CopyStart(client_id, key, "segment_1",
+    copy_result = service_->CopyStart(client_id, key, "default", "segment_1",
                                       {"segment_3", "segment_4"});
     EXPECT_TRUE(copy_result.has_value());
     copy_response = copy_result.value();
@@ -1007,16 +2359,16 @@ TEST_F(MasterServiceTest, CopyStart) {
                                .buffer_descriptor.transport_endpoint_);
 
     // End the copy operation to clean up state
-    copy_end_result = service_->CopyEnd(client_id, key);
+    copy_end_result = service_->CopyEnd(client_id, key, "default");
     EXPECT_TRUE(copy_end_result.has_value());
-    get_result = service_->GetReplicaList(key);
+    get_result = service_->GetReplicaList(key, "default");
     EXPECT_TRUE(get_result.has_value());
     EXPECT_EQ(4, get_result.value().replicas.size());
 
     // Test Case 10: Copy to segment_4 again, should skip because it's already
     // used.
-    copy_result =
-        service_->CopyStart(client_id, key, "segment_1", {"segment_4"});
+    copy_result = service_->CopyStart(client_id, key, "default", "segment_1",
+                                      {"segment_4"});
     EXPECT_TRUE(copy_result.has_value());
     copy_response = copy_result.value();
     EXPECT_EQ("segment_1", copy_response.source.get_memory_descriptor()
@@ -1029,19 +2381,19 @@ TEST_F(MasterServiceTest, CopyStart) {
     std::this_thread::sleep_for(std::chrono::milliseconds(kv_lease_ttl * 2));
 
     // Test Case 11: Try remove the object, should fail because it is copying.
-    remove_result = service_->Remove(key);
+    remove_result = service_->Remove(key, "default");
     EXPECT_FALSE(remove_result.has_value());
     EXPECT_EQ(ErrorCode::OBJECT_HAS_REPLICATION_TASK, remove_result.error());
 
     // Clean up the copy operation
-    copy_end_result = service_->CopyEnd(client_id, key);
+    copy_end_result = service_->CopyEnd(client_id, key, "default");
     EXPECT_TRUE(copy_end_result.has_value());
 
     // Wait for the lease to expire
     std::this_thread::sleep_for(std::chrono::milliseconds(kv_lease_ttl * 2));
 
     // Test Case 12: Try remove the object, should success.
-    remove_result = service_->Remove(key);
+    remove_result = service_->Remove(key, "default");
     EXPECT_TRUE(remove_result.has_value());
 }
 
@@ -1061,7 +2413,8 @@ TEST_F(MasterServiceTest, CopyEnd) {
     UUID invalid_client_id = generate_uuid();
 
     // Test Case 1: CopyEnd a non-existent key, should fail.
-    auto copy_end_result = service_->CopyEnd(client_id, "non_existent_key");
+    auto copy_end_result =
+        service_->CopyEnd(client_id, "non_existent_key", "default");
     EXPECT_FALSE(copy_end_result.has_value());
     EXPECT_EQ(ErrorCode::OBJECT_NOT_FOUND, copy_end_result.error());
 
@@ -1073,45 +2426,46 @@ TEST_F(MasterServiceTest, CopyEnd) {
     config.preferred_segment = "segment_1";
 
     auto put_start_result =
-        service_->PutStart(client_id, key, slice_length, config);
+        service_->PutStart(client_id, key, "default", slice_length, config);
     ASSERT_TRUE(put_start_result.has_value());
-    auto put_end_result = service_->PutEnd(client_id, key, ReplicaType::MEMORY);
+    auto put_end_result =
+        service_->PutEnd(client_id, key, "default", ReplicaType::MEMORY);
     ASSERT_TRUE(put_end_result.has_value());
 
     // Test Case 2: CopyEnd the object, should fail because there is no ongoing
     // copy task.
-    copy_end_result = service_->CopyEnd(client_id, key);
+    copy_end_result = service_->CopyEnd(client_id, key, "default");
     EXPECT_FALSE(copy_end_result.has_value());
     EXPECT_EQ(ErrorCode::OBJECT_NO_REPLICATION_TASK, copy_end_result.error());
 
     // CopyStart the object to segment_2
-    auto copy_start_result =
-        service_->CopyStart(client_id, key, "segment_1", {"segment_2"});
+    auto copy_start_result = service_->CopyStart(client_id, key, "default",
+                                                 "segment_1", {"segment_2"});
     ASSERT_TRUE(copy_start_result.has_value());
 
     // Test Case 3: CopyEnd with an invalid client id, should fail.
-    copy_end_result = service_->CopyEnd(invalid_client_id, key);
+    copy_end_result = service_->CopyEnd(invalid_client_id, key, "default");
     EXPECT_FALSE(copy_end_result.has_value());
     EXPECT_EQ(ErrorCode::ILLEGAL_CLIENT, copy_end_result.error());
 
     // Test Case 4: MoveEnd the object, should fail because the ongoing task is
     // Copy.
-    auto move_end_result = service_->MoveEnd(client_id, key);
+    auto move_end_result = service_->MoveEnd(client_id, key, "default");
     EXPECT_FALSE(move_end_result.has_value());
     EXPECT_EQ(ErrorCode::INVALID_PARAMS, move_end_result.error());
 
     // Test Case 5: CopyEnd, should success.
-    copy_end_result = service_->CopyEnd(client_id, key);
+    copy_end_result = service_->CopyEnd(client_id, key, "default");
     EXPECT_TRUE(copy_end_result.has_value());
 
     // Verify we now have 2 replicas
-    auto get_result = service_->GetReplicaList(key);
+    auto get_result = service_->GetReplicaList(key, "default");
     EXPECT_TRUE(get_result.has_value());
     EXPECT_EQ(2, get_result.value().replicas.size());
 
     // CopyStart the object from segment_1 to segment_3, then unmount segment_1
-    copy_start_result =
-        service_->CopyStart(client_id, key, "segment_1", {"segment_3"});
+    copy_start_result = service_->CopyStart(client_id, key, "default",
+                                            "segment_1", {"segment_3"});
     ASSERT_TRUE(copy_start_result.has_value());
 
     // Unmount segment_1 to simulate source gone
@@ -1121,10 +2475,10 @@ TEST_F(MasterServiceTest, CopyEnd) {
 
     // Test Case 6: CopyEnd, should fail because the source is gone, the object
     // should have only 1 replica from segment_2.
-    copy_end_result = service_->CopyEnd(client_id, key);
+    copy_end_result = service_->CopyEnd(client_id, key, "default");
     EXPECT_FALSE(copy_end_result.has_value());
     EXPECT_EQ(ErrorCode::REPLICA_IS_GONE, copy_end_result.error());
-    get_result = service_->GetReplicaList(key);
+    get_result = service_->GetReplicaList(key, "default");
     EXPECT_TRUE(get_result.has_value());
     auto& replicas = get_result.value().replicas;
     EXPECT_EQ(1, replicas.size());
@@ -1133,8 +2487,8 @@ TEST_F(MasterServiceTest, CopyEnd) {
                                .buffer_descriptor.transport_endpoint_);
 
     // CopyStart the object from segment_2 to segment_3, then unmount segment_3
-    copy_start_result =
-        service_->CopyStart(client_id, key, "segment_2", {"segment_3"});
+    copy_start_result = service_->CopyStart(client_id, key, "default",
+                                            "segment_2", {"segment_3"});
     ASSERT_TRUE(copy_start_result.has_value());
 
     // Unmount segment_3 to simulate target gone
@@ -1144,10 +2498,10 @@ TEST_F(MasterServiceTest, CopyEnd) {
 
     // Test Case 7: CopyEnd, should fail because the target is gone, the object
     // should have only 1 replica from segment_2.
-    copy_end_result = service_->CopyEnd(client_id, key);
+    copy_end_result = service_->CopyEnd(client_id, key, "default");
     EXPECT_FALSE(copy_end_result.has_value());
     EXPECT_EQ(ErrorCode::REPLICA_IS_GONE, copy_end_result.error());
-    get_result = service_->GetReplicaList(key);
+    get_result = service_->GetReplicaList(key, "default");
     EXPECT_TRUE(get_result.has_value());
     replicas = get_result.value().replicas;
     EXPECT_EQ(1, replicas.size());
@@ -1171,7 +2525,7 @@ TEST_F(MasterServiceTest, CopyRevoke) {
 
     // Test Case 1: CopyRevoke a non-existent key, should fail.
     auto copy_revoke_result =
-        service_->CopyRevoke(client_id, "non_existent_key");
+        service_->CopyRevoke(client_id, "non_existent_key", "default");
     EXPECT_FALSE(copy_revoke_result.has_value());
     EXPECT_EQ(ErrorCode::OBJECT_NOT_FOUND, copy_revoke_result.error());
 
@@ -1183,47 +2537,49 @@ TEST_F(MasterServiceTest, CopyRevoke) {
     config.preferred_segment = "segment_1";
 
     auto put_start_result =
-        service_->PutStart(client_id, key, slice_length, config);
+        service_->PutStart(client_id, key, "default", slice_length, config);
     ASSERT_TRUE(put_start_result.has_value());
-    auto put_end_result = service_->PutEnd(client_id, key, ReplicaType::MEMORY);
+    auto put_end_result =
+        service_->PutEnd(client_id, key, "default", ReplicaType::MEMORY);
     ASSERT_TRUE(put_end_result.has_value());
 
     // Test Case 2: CopyRevoke the object, should fail because there is no
     // ongoing copy task.
-    copy_revoke_result = service_->CopyRevoke(client_id, key);
+    copy_revoke_result = service_->CopyRevoke(client_id, key, "default");
     EXPECT_FALSE(copy_revoke_result.has_value());
     EXPECT_EQ(ErrorCode::OBJECT_NO_REPLICATION_TASK,
               copy_revoke_result.error());
 
     // CopyStart the object to segment_2
-    auto copy_start_result =
-        service_->CopyStart(client_id, key, "segment_1", {"segment_2"});
+    auto copy_start_result = service_->CopyStart(client_id, key, "default",
+                                                 "segment_1", {"segment_2"});
     ASSERT_TRUE(copy_start_result.has_value());
 
     // Test Case 3: CopyRevoke with an invalid client id, should fail.
-    copy_revoke_result = service_->CopyRevoke(invalid_client_id, key);
+    copy_revoke_result =
+        service_->CopyRevoke(invalid_client_id, key, "default");
     EXPECT_FALSE(copy_revoke_result.has_value());
     EXPECT_EQ(ErrorCode::ILLEGAL_CLIENT, copy_revoke_result.error());
 
     // Test Case 4: MoveRevoke the object, should fail because the ongoing task
     // is Copy.
-    auto move_revoke_result = service_->MoveRevoke(client_id, key);
+    auto move_revoke_result = service_->MoveRevoke(client_id, key, "default");
     EXPECT_FALSE(move_revoke_result.has_value());
     EXPECT_EQ(ErrorCode::INVALID_PARAMS, move_revoke_result.error());
 
     // Test Case 5: CopyRevoke, should success.
-    copy_revoke_result = service_->CopyRevoke(client_id, key);
+    copy_revoke_result = service_->CopyRevoke(client_id, key, "default");
     EXPECT_TRUE(copy_revoke_result.has_value());
 
     // Verify we still have 1 replica (the copy was revoked)
-    auto get_result = service_->GetReplicaList(key);
+    auto get_result = service_->GetReplicaList(key, "default");
     EXPECT_TRUE(get_result.has_value());
     EXPECT_EQ(1, get_result.value().replicas.size());
 
     // CopyStart the object from segment_1 to segment_2 again, then unmount
     // segment_1
-    copy_start_result =
-        service_->CopyStart(client_id, key, "segment_1", {"segment_2"});
+    copy_start_result = service_->CopyStart(client_id, key, "default",
+                                            "segment_1", {"segment_2"});
     ASSERT_TRUE(copy_start_result.has_value());
 
     // Unmount segment_1 to simulate source gone
@@ -1233,11 +2589,11 @@ TEST_F(MasterServiceTest, CopyRevoke) {
 
     // Test Case 6: CopyRevoke, should success even though the source is gone,
     // the object should be erased too.
-    copy_revoke_result = service_->CopyRevoke(client_id, key);
+    copy_revoke_result = service_->CopyRevoke(client_id, key, "default");
     EXPECT_TRUE(copy_revoke_result.has_value());
 
     // Verify the object has been removed.
-    get_result = service_->GetReplicaList(key);
+    get_result = service_->GetReplicaList(key, "default");
     EXPECT_FALSE(get_result.has_value());
 }
 
@@ -1260,8 +2616,8 @@ TEST_F(MasterServiceTest, MoveStart) {
     UUID client_id = generate_uuid();
 
     // Test Case 1: MoveStart a non-existent key, should fail.
-    auto move_start_result = service_->MoveStart(client_id, "non_existent_key",
-                                                 "segment_1", "segment_2");
+    auto move_start_result = service_->MoveStart(
+        client_id, "non_existent_key", "default", "segment_1", "segment_2");
     EXPECT_FALSE(move_start_result.has_value());
     EXPECT_EQ(ErrorCode::OBJECT_NOT_FOUND, move_start_result.error());
 
@@ -1273,36 +2629,37 @@ TEST_F(MasterServiceTest, MoveStart) {
     config.preferred_segment = "segment_1";
 
     auto put_start_result =
-        service_->PutStart(client_id, key, slice_length, config);
+        service_->PutStart(client_id, key, "default", slice_length, config);
     ASSERT_TRUE(put_start_result.has_value());
 
     // Test Case 2: MoveStart the object, should fail because the only replica
     // is not completed.
-    move_start_result =
-        service_->MoveStart(client_id, key, "segment_1", "segment_2");
+    move_start_result = service_->MoveStart(client_id, key, "default",
+                                            "segment_1", "segment_2");
     EXPECT_FALSE(move_start_result.has_value());
     EXPECT_EQ(ErrorCode::REPLICA_NOT_FOUND, move_start_result.error());
 
     // PutEnd the object.
-    auto put_end_result = service_->PutEnd(client_id, key, ReplicaType::MEMORY);
+    auto put_end_result =
+        service_->PutEnd(client_id, key, "default", ReplicaType::MEMORY);
     ASSERT_TRUE(put_end_result.has_value());
 
     // Copy the object to segment_3.
-    auto copy_start_result =
-        service_->CopyStart(client_id, key, "segment_1", {"segment_3"});
+    auto copy_start_result = service_->CopyStart(client_id, key, "default",
+                                                 "segment_1", {"segment_3"});
     ASSERT_TRUE(copy_start_result.has_value());
-    auto copy_end_result = service_->CopyEnd(client_id, key);
+    auto copy_end_result = service_->CopyEnd(client_id, key, "default");
     ASSERT_TRUE(copy_end_result.has_value());
 
     // Test Case 3: MoveStart with source and target be the same, should fail.
-    move_start_result =
-        service_->MoveStart(client_id, key, "segment_1", "segment_1");
+    move_start_result = service_->MoveStart(client_id, key, "default",
+                                            "segment_1", "segment_1");
     EXPECT_FALSE(move_start_result.has_value());
     EXPECT_EQ(move_start_result.error(), ErrorCode::INVALID_PARAMS);
 
     // Test Case 4: MoveStart to segment_2, should succeed.
-    move_start_result =
-        service_->MoveStart(client_id, key, "segment_1", "segment_2");
+    move_start_result = service_->MoveStart(client_id, key, "default",
+                                            "segment_1", "segment_2");
     EXPECT_TRUE(move_start_result.has_value());
     auto move_response = move_start_result.value();
     EXPECT_EQ("segment_1", move_response.source.get_memory_descriptor()
@@ -1313,44 +2670,44 @@ TEST_F(MasterServiceTest, MoveStart) {
                                .buffer_descriptor.transport_endpoint_);
 
     // Test Case 5: Try remove the object, should fail because it is moving.
-    auto remove_result = service_->Remove(key);
+    auto remove_result = service_->Remove(key, "default");
     EXPECT_FALSE(remove_result.has_value());
     EXPECT_EQ(ErrorCode::REPLICA_IS_NOT_READY, remove_result.error());
 
     // Test Case 6: MoveStart again, should fail because there is an ongoing
     // move task.
-    move_start_result =
-        service_->MoveStart(client_id, key, "segment_1", "segment_3");
+    move_start_result = service_->MoveStart(client_id, key, "default",
+                                            "segment_1", "segment_3");
     EXPECT_FALSE(move_start_result.has_value());
     EXPECT_EQ(ErrorCode::OBJECT_HAS_REPLICATION_TASK,
               move_start_result.error());
 
     // Test Case 7: MoveEnd, should succeed and the object now has 2 replicas
     // from segment_2 and segment_3
-    auto move_end_result = service_->MoveEnd(client_id, key);
+    auto move_end_result = service_->MoveEnd(client_id, key, "default");
     EXPECT_TRUE(move_end_result.has_value());
 
-    auto get_result = service_->GetReplicaList(key);
+    auto get_result = service_->GetReplicaList(key, "default");
     EXPECT_TRUE(get_result.has_value());
     auto& replicas = get_result.value().replicas;
     EXPECT_EQ(2, replicas.size());
 
     // Test Case 8: Move from a non-existent replica to segment_1, should fail.
     move_start_result = service_->MoveStart(
-        client_id, key, "non_existent_segment", "segment_1");
+        client_id, key, "default", "non_existent_segment", "segment_1");
     EXPECT_FALSE(move_start_result.has_value());
     EXPECT_EQ(ErrorCode::REPLICA_NOT_FOUND, move_start_result.error());
 
     // Test Case 8.5: Move to a non-existent target segment, should fail.
-    move_start_result = service_->MoveStart(client_id, key, "segment_2",
-                                            "non_existent_segment");
+    move_start_result = service_->MoveStart(
+        client_id, key, "default", "segment_2", "non_existent_segment");
     EXPECT_FALSE(move_start_result.has_value());
     EXPECT_EQ(ErrorCode::SEGMENT_NOT_FOUND, move_start_result.error());
 
     // Test Case 9: Move to an already existing segment, should succeed but
     // return nullopt.
-    move_start_result =
-        service_->MoveStart(client_id, key, "segment_2", "segment_3");
+    move_start_result = service_->MoveStart(client_id, key, "default",
+                                            "segment_2", "segment_3");
     EXPECT_TRUE(move_start_result.has_value());
     move_response = move_start_result.value();
     EXPECT_EQ("segment_2", move_response.source.get_memory_descriptor()
@@ -1359,16 +2716,16 @@ TEST_F(MasterServiceTest, MoveStart) {
 
     // Test Case 10: Try remove the object, should fail because it is moving.
     std::this_thread::sleep_for(std::chrono::milliseconds(kv_lease_ttl * 2));
-    remove_result = service_->Remove(key);
+    remove_result = service_->Remove(key, "default");
     EXPECT_FALSE(remove_result.has_value());
     EXPECT_EQ(ErrorCode::OBJECT_HAS_REPLICATION_TASK, remove_result.error());
 
     // End the move.
-    move_end_result = service_->MoveEnd(client_id, key);
+    move_end_result = service_->MoveEnd(client_id, key, "default");
     EXPECT_TRUE(move_end_result.has_value());
 
     // Now the object should have only 1 replica on segment_3.
-    get_result = service_->GetReplicaList(key);
+    get_result = service_->GetReplicaList(key, "default");
     EXPECT_TRUE(get_result.has_value());
     replicas = get_result.value().replicas;
     EXPECT_EQ(1, replicas.size());
@@ -1378,7 +2735,7 @@ TEST_F(MasterServiceTest, MoveStart) {
 
     // Test Case 11: Try remove the object, should succeed after lease expires.
     std::this_thread::sleep_for(std::chrono::milliseconds(kv_lease_ttl * 2));
-    remove_result = service_->Remove(key);
+    remove_result = service_->Remove(key, "default");
     EXPECT_TRUE(remove_result.has_value());
 }
 
@@ -1396,7 +2753,8 @@ TEST_F(MasterServiceTest, MoveEnd) {
     UUID invalid_client_id = generate_uuid();
 
     // Test Case 1: MoveEnd a non-existent key, should fail.
-    auto move_end_result = service_->MoveEnd(client_id, "non_existent_key");
+    auto move_end_result =
+        service_->MoveEnd(client_id, "non_existent_key", "default");
     EXPECT_FALSE(move_end_result.has_value());
     EXPECT_EQ(ErrorCode::OBJECT_NOT_FOUND, move_end_result.error());
 
@@ -1408,46 +2766,47 @@ TEST_F(MasterServiceTest, MoveEnd) {
     config.preferred_segment = "segment_1";
 
     auto put_start_result =
-        service_->PutStart(client_id, key, slice_length, config);
+        service_->PutStart(client_id, key, "default", slice_length, config);
     ASSERT_TRUE(put_start_result.has_value());
-    auto put_end_result = service_->PutEnd(client_id, key, ReplicaType::MEMORY);
+    auto put_end_result =
+        service_->PutEnd(client_id, key, "default", ReplicaType::MEMORY);
     ASSERT_TRUE(put_end_result.has_value());
 
     // Test Case 2: MoveEnd the object, should fail because there is no ongoing
     // move task.
-    move_end_result = service_->MoveEnd(client_id, key);
+    move_end_result = service_->MoveEnd(client_id, key, "default");
     EXPECT_FALSE(move_end_result.has_value());
     EXPECT_EQ(ErrorCode::OBJECT_NO_REPLICATION_TASK, move_end_result.error());
 
     // MoveStart the object to segment_2
-    auto move_start_result =
-        service_->MoveStart(client_id, key, "segment_1", "segment_2");
+    auto move_start_result = service_->MoveStart(client_id, key, "default",
+                                                 "segment_1", "segment_2");
     ASSERT_TRUE(move_start_result.has_value());
 
     // Test Case 3: MoveEnd with an invalid client id, should fail.
-    move_end_result = service_->MoveEnd(invalid_client_id, key);
+    move_end_result = service_->MoveEnd(invalid_client_id, key, "default");
     EXPECT_FALSE(move_end_result.has_value());
     EXPECT_EQ(ErrorCode::ILLEGAL_CLIENT, move_end_result.error());
 
     // Test Case 4: CopyEnd the object, should fail because the ongoing task is
     // Move.
-    auto copy_end_result = service_->CopyEnd(client_id, key);
+    auto copy_end_result = service_->CopyEnd(client_id, key, "default");
     EXPECT_FALSE(copy_end_result.has_value());
     EXPECT_EQ(ErrorCode::INVALID_PARAMS, copy_end_result.error());
 
     // Test Case 5: MoveEnd, should success.
-    move_end_result = service_->MoveEnd(client_id, key);
+    move_end_result = service_->MoveEnd(client_id, key, "default");
     EXPECT_TRUE(move_end_result.has_value());
 
     // Verify we still have 1 replica (the move was successful)
-    auto get_result = service_->GetReplicaList(key);
+    auto get_result = service_->GetReplicaList(key, "default");
     EXPECT_TRUE(get_result.has_value());
     EXPECT_EQ(1, get_result.value().replicas.size());
 
     // MoveStart the object from segment_2 to segment_1 again, then unmount
     // segment_2
-    move_start_result =
-        service_->MoveStart(client_id, key, "segment_2", "segment_1");
+    move_start_result = service_->MoveStart(client_id, key, "default",
+                                            "segment_2", "segment_1");
     ASSERT_TRUE(move_start_result.has_value());
 
     // Unmount segment_2 to simulate source gone
@@ -1456,7 +2815,7 @@ TEST_F(MasterServiceTest, MoveEnd) {
     ASSERT_TRUE(unmount_result.has_value());
 
     // Test Case 6: MoveEnd, should fail because the source is gone.
-    move_end_result = service_->MoveEnd(client_id, key);
+    move_end_result = service_->MoveEnd(client_id, key, "default");
     EXPECT_FALSE(move_end_result.has_value());
     EXPECT_EQ(ErrorCode::REPLICA_IS_GONE, move_end_result.error());
 }
@@ -1475,7 +2834,7 @@ TEST_F(MasterServiceTest, MoveRevoke) {
 
     // Test Case 1: MoveRevoke a non-existent key, should fail.
     auto move_revoke_result =
-        service_->MoveRevoke(client_id, "non_existent_key");
+        service_->MoveRevoke(client_id, "non_existent_key", "default");
     EXPECT_FALSE(move_revoke_result.has_value());
     EXPECT_EQ(ErrorCode::OBJECT_NOT_FOUND, move_revoke_result.error());
 
@@ -1487,40 +2846,42 @@ TEST_F(MasterServiceTest, MoveRevoke) {
     config.preferred_segment = "segment_1";
 
     auto put_start_result =
-        service_->PutStart(client_id, key, slice_length, config);
+        service_->PutStart(client_id, key, "default", slice_length, config);
     ASSERT_TRUE(put_start_result.has_value());
-    auto put_end_result = service_->PutEnd(client_id, key, ReplicaType::MEMORY);
+    auto put_end_result =
+        service_->PutEnd(client_id, key, "default", ReplicaType::MEMORY);
     ASSERT_TRUE(put_end_result.has_value());
 
     // Test Case 2: MoveRevoke the object, should fail because there is no
     // ongoing move task.
-    move_revoke_result = service_->MoveRevoke(client_id, key);
+    move_revoke_result = service_->MoveRevoke(client_id, key, "default");
     EXPECT_FALSE(move_revoke_result.has_value());
     EXPECT_EQ(ErrorCode::OBJECT_NO_REPLICATION_TASK,
               move_revoke_result.error());
 
     // MoveStart the object from segment_1 to segment_2
-    auto move_start_result =
-        service_->MoveStart(client_id, key, "segment_1", "segment_2");
+    auto move_start_result = service_->MoveStart(client_id, key, "default",
+                                                 "segment_1", "segment_2");
     ASSERT_TRUE(move_start_result.has_value());
 
     // Test Case 3: MoveRevoke with an invalid client id, should fail.
-    move_revoke_result = service_->MoveRevoke(invalid_client_id, key);
+    move_revoke_result =
+        service_->MoveRevoke(invalid_client_id, key, "default");
     EXPECT_FALSE(move_revoke_result.has_value());
     EXPECT_EQ(ErrorCode::ILLEGAL_CLIENT, move_revoke_result.error());
 
     // Test Case 4: CopyRevoke the object, should fail because the ongoing task
     // is Move.
-    auto copy_revoke_result = service_->CopyRevoke(client_id, key);
+    auto copy_revoke_result = service_->CopyRevoke(client_id, key, "default");
     EXPECT_FALSE(copy_revoke_result.has_value());
     EXPECT_EQ(ErrorCode::INVALID_PARAMS, copy_revoke_result.error());
 
     // Test Case 5: MoveRevoke, should succeed.
-    move_revoke_result = service_->MoveRevoke(client_id, key);
+    move_revoke_result = service_->MoveRevoke(client_id, key, "default");
     EXPECT_TRUE(move_revoke_result.has_value());
 
     // Verify we still have 1 replica (the move was revoked)
-    auto get_result = service_->GetReplicaList(key);
+    auto get_result = service_->GetReplicaList(key, "default");
     EXPECT_TRUE(get_result.has_value());
     auto& replicas = get_result.value().replicas;
     EXPECT_EQ(1, replicas.size());
@@ -1530,8 +2891,8 @@ TEST_F(MasterServiceTest, MoveRevoke) {
 
     // MoveStart the object from segment_1 to segment_2 again, then unmount
     // segment_1
-    move_start_result =
-        service_->MoveStart(client_id, key, "segment_1", "segment_2");
+    move_start_result = service_->MoveStart(client_id, key, "default",
+                                            "segment_1", "segment_2");
     ASSERT_TRUE(move_start_result.has_value());
 
     // Unmount segment_1 to simulate source gone
@@ -1540,11 +2901,11 @@ TEST_F(MasterServiceTest, MoveRevoke) {
     ASSERT_TRUE(unmount_result.has_value());
 
     // Test Case 6: MoveRevoke, should succeed even though the source is gone.
-    move_revoke_result = service_->MoveRevoke(client_id, key);
+    move_revoke_result = service_->MoveRevoke(client_id, key, "default");
     EXPECT_TRUE(move_revoke_result.has_value());
 
     // The object should be erased as there is no replica left.
-    get_result = service_->GetReplicaList(key);
+    get_result = service_->GetReplicaList(key, "default");
     EXPECT_FALSE(get_result.has_value());
 }
 
@@ -1576,26 +2937,27 @@ TEST_F(MasterServiceTest, ProtectCopyMoveSourceFromEviction) {
     config.preferred_segment = "segment_1";
 
     // Put two objects for move and copy tests.
-    auto put_start_result =
-        service_->PutStart(client_id, copy_key, slice_length, config);
+    auto put_start_result = service_->PutStart(client_id, copy_key, "default",
+                                               slice_length, config);
     ASSERT_TRUE(put_start_result.has_value());
     auto put_end_result =
-        service_->PutEnd(client_id, copy_key, ReplicaType::MEMORY);
+        service_->PutEnd(client_id, copy_key, "default", ReplicaType::MEMORY);
     ASSERT_TRUE(put_end_result.has_value());
 
-    put_start_result =
-        service_->PutStart(client_id, move_key, slice_length, config);
+    put_start_result = service_->PutStart(client_id, move_key, "default",
+                                          slice_length, config);
     ASSERT_TRUE(put_start_result.has_value());
-    put_end_result = service_->PutEnd(client_id, move_key, ReplicaType::MEMORY);
+    put_end_result =
+        service_->PutEnd(client_id, move_key, "default", ReplicaType::MEMORY);
     ASSERT_TRUE(put_end_result.has_value());
 
     // Start copy and move operations.
-    auto copy_start_result =
-        service_->CopyStart(client_id, copy_key, "segment_1", {"segment_2"});
+    auto copy_start_result = service_->CopyStart(client_id, copy_key, "default",
+                                                 "segment_1", {"segment_2"});
     ASSERT_TRUE(copy_start_result.has_value());
 
-    auto move_start_result =
-        service_->MoveStart(client_id, move_key, "segment_1", "segment_2");
+    auto move_start_result = service_->MoveStart(client_id, move_key, "default",
+                                                 "segment_1", "segment_2");
     ASSERT_TRUE(move_start_result.has_value());
 
     // Put more objects to trigger eviction. Do not prefer any segments.
@@ -1603,10 +2965,10 @@ TEST_F(MasterServiceTest, ProtectCopyMoveSourceFromEviction) {
     for (size_t i = 0; i < 128 * (kSegmentSize * 2 / slice_length); ++i) {
         std::string key = "test_key_" + std::to_string(i);
         auto put_start_result =
-            service_->PutStart(client_id, key, slice_length, config);
+            service_->PutStart(client_id, key, "default", slice_length, config);
         if (put_start_result.has_value()) {
-            auto put_end_result =
-                service_->PutEnd(client_id, key, ReplicaType::MEMORY);
+            auto put_end_result = service_->PutEnd(client_id, key, "default",
+                                                   ReplicaType::MEMORY);
             ASSERT_TRUE(put_end_result.has_value());
         } else {
             // wait for eviction to work
@@ -1620,10 +2982,10 @@ TEST_F(MasterServiceTest, ProtectCopyMoveSourceFromEviction) {
     ASSERT_TRUE(remove_all_result > 0);
 
     // Try end copy and move operations, should success.
-    auto copy_end_result = service_->CopyEnd(client_id, copy_key);
+    auto copy_end_result = service_->CopyEnd(client_id, copy_key, "default");
     EXPECT_TRUE(copy_end_result.has_value());
 
-    auto move_end_result = service_->MoveEnd(client_id, move_key);
+    auto move_end_result = service_->MoveEnd(client_id, move_key, "default");
     EXPECT_TRUE(move_end_result.has_value());
 }
 
@@ -1660,26 +3022,27 @@ TEST_F(MasterServiceTest, DiscardTimeoutCopyMove) {
     config.preferred_segment = "segment_1";
 
     // Put two objects for move and copy tests.
-    auto put_start_result =
-        service_->PutStart(client_id, copy_key, slice_length, config);
+    auto put_start_result = service_->PutStart(client_id, copy_key, "default",
+                                               slice_length, config);
     ASSERT_TRUE(put_start_result.has_value());
     auto put_end_result =
-        service_->PutEnd(client_id, copy_key, ReplicaType::MEMORY);
+        service_->PutEnd(client_id, copy_key, "default", ReplicaType::MEMORY);
     ASSERT_TRUE(put_end_result.has_value());
 
-    put_start_result =
-        service_->PutStart(client_id, move_key, slice_length, config);
+    put_start_result = service_->PutStart(client_id, move_key, "default",
+                                          slice_length, config);
     ASSERT_TRUE(put_start_result.has_value());
-    put_end_result = service_->PutEnd(client_id, move_key, ReplicaType::MEMORY);
+    put_end_result =
+        service_->PutEnd(client_id, move_key, "default", ReplicaType::MEMORY);
     ASSERT_TRUE(put_end_result.has_value());
 
     // Start copy and move operations.
-    auto copy_start_result =
-        service_->CopyStart(client_id, copy_key, "segment_1", {"segment_2"});
+    auto copy_start_result = service_->CopyStart(client_id, copy_key, "default",
+                                                 "segment_1", {"segment_2"});
     ASSERT_TRUE(copy_start_result.has_value());
 
-    auto move_start_result =
-        service_->MoveStart(client_id, move_key, "segment_1", "segment_2");
+    auto move_start_result = service_->MoveStart(client_id, move_key, "default",
+                                                 "segment_1", "segment_2");
     ASSERT_TRUE(move_start_result.has_value());
 
     // Wait for the operations timeout.
@@ -1690,10 +3053,10 @@ TEST_F(MasterServiceTest, DiscardTimeoutCopyMove) {
     for (size_t i = 0; i < 128 * (kSegmentSize * 2 / slice_length); ++i) {
         std::string key = "test_key_" + std::to_string(i);
         auto put_start_result =
-            service_->PutStart(client_id, key, slice_length, config);
+            service_->PutStart(client_id, key, "default", slice_length, config);
         if (put_start_result.has_value()) {
-            auto put_end_result =
-                service_->PutEnd(client_id, key, ReplicaType::MEMORY);
+            auto put_end_result = service_->PutEnd(client_id, key, "default",
+                                                   ReplicaType::MEMORY);
             ASSERT_TRUE(put_end_result.has_value());
         } else {
             // wait for eviction to work
@@ -1703,11 +3066,11 @@ TEST_F(MasterServiceTest, DiscardTimeoutCopyMove) {
 
     // Try end copy and move operations, should fail because the objects are
     // evicted.
-    auto copy_end_result = service_->CopyEnd(client_id, copy_key);
+    auto copy_end_result = service_->CopyEnd(client_id, copy_key, "default");
     EXPECT_FALSE(copy_end_result.has_value());
     EXPECT_EQ(copy_end_result.error(), ErrorCode::OBJECT_NOT_FOUND);
 
-    auto move_end_result = service_->MoveEnd(client_id, move_key);
+    auto move_end_result = service_->MoveEnd(client_id, move_key, "default");
     EXPECT_FALSE(move_end_result.has_value());
     EXPECT_EQ(move_end_result.error(), ErrorCode::OBJECT_NOT_FOUND);
 }
@@ -1753,7 +3116,7 @@ TEST_F(MasterServiceTest, RemoveByRegexComplex) {
         populate_store();
 
         // Action: Remove keys starting with "test_key_"
-        auto remove_result = service_->RemoveByRegex("^test_key_");
+        auto remove_result = service_->RemoveByRegex("^test_key_", "default");
         ASSERT_TRUE(remove_result.has_value());
         EXPECT_EQ(remove_result.value(), 3);  // Should remove 3 keys
 
@@ -1761,7 +3124,7 @@ TEST_F(MasterServiceTest, RemoveByRegexComplex) {
         std::vector<std::string> deleted_keys = {"test_key_01", "test_key_02",
                                                  "test_key_10"};
         for (const auto& key : deleted_keys) {
-            auto exist_result = service_->ExistKey(key);
+            auto exist_result = service_->ExistKey(key, "default");
             ASSERT_TRUE(exist_result.has_value());
             EXPECT_FALSE(exist_result.value())
                 << "Key " << key << " should have been deleted.";
@@ -1770,7 +3133,7 @@ TEST_F(MasterServiceTest, RemoveByRegexComplex) {
         std::vector<std::string> remaining_keys = {
             "prod_key_alpha", "short", "test-key-extra"};  // Sample a few
         for (const auto& key : remaining_keys) {
-            auto exist_result = service_->ExistKey(key);
+            auto exist_result = service_->ExistKey(key, "default");
             ASSERT_TRUE(exist_result.has_value());
             EXPECT_TRUE(exist_result.value())
                 << "Key " << key << " should NOT have been deleted.";
@@ -1793,12 +3156,12 @@ TEST_F(MasterServiceTest, RemoveByRegexComplex) {
         size_t total_keys = 13;  // Count from the keys_to_put vector
 
         // Action: Remove all keys
-        auto remove_result = service_->RemoveByRegex(".*");
+        auto remove_result = service_->RemoveByRegex(".*", "default");
         ASSERT_TRUE(remove_result.has_value());
         EXPECT_EQ(remove_result.value(), total_keys);
 
         // Verification: Check that no keys remain
-        auto get_all_result = service_->GetReplicaListByRegex(".*");
+        auto get_all_result = service_->GetReplicaListByRegex(".*", "default");
         ASSERT_TRUE(get_all_result.has_value());
         EXPECT_TRUE(get_all_result.value().empty());
     }
@@ -1818,12 +3181,13 @@ TEST_F(MasterServiceTest, RemoveByRegexComplex) {
         size_t total_keys_before_remove = 13;
 
         // Action: Attempt to remove using a pattern that matches nothing
-        auto remove_result = service_->RemoveByRegex("^nonexistent-pattern-");
+        auto remove_result =
+            service_->RemoveByRegex("^nonexistent-pattern-", "default");
         ASSERT_TRUE(remove_result.has_value());
         EXPECT_EQ(remove_result.value(), 0);  // Should remove 0 keys
 
         // Verification: Check that all keys still exist
-        auto get_all_result = service_->GetReplicaListByRegex(".*");
+        auto get_all_result = service_->GetReplicaListByRegex(".*", "default");
         ASSERT_TRUE(get_all_result.has_value());
         EXPECT_EQ(get_all_result.value().size(), total_keys_before_remove);
     }
@@ -1841,7 +3205,7 @@ TEST_F(MasterServiceTest, RemoveByRegexComplex) {
         populate_store();
 
         // Action: Remove all keys that contain a slash '/' OR end with a number
-        auto remove_result = service_->RemoveByRegex("/|\\d$");
+        auto remove_result = service_->RemoveByRegex("/|\\d$", "default");
         ASSERT_TRUE(remove_result.has_value());
         // Matches: "config/user/settings.json", "logs/app-2025-08-13.log",
         //          "test_key_01", "test_key_02", "test_key_10"
@@ -1866,23 +3230,25 @@ TEST_F(MasterServiceTest, RemoveByRegexComplex) {
         populate_store();
 
         // Action: Remove all keys that contain "chunk" OR "config"
-        auto remove_result = service_->RemoveByRegex("chunk|config");
+        auto remove_result = service_->RemoveByRegex("chunk|config", "default");
         ASSERT_TRUE(remove_result.has_value());
         // Matches: "data_part_1_chunk_a", "data_part_2_chunk_b",
         // "config/user/settings.json"
         EXPECT_EQ(remove_result.value(), 3);
 
         // Verification
-        auto exist_result_chunk = service_->ExistKey("data_part_1_chunk_a");
+        auto exist_result_chunk =
+            service_->ExistKey("data_part_1_chunk_a", "default");
         ASSERT_TRUE(exist_result_chunk.has_value());
         EXPECT_FALSE(exist_result_chunk.value());
 
         auto exist_result_config =
-            service_->ExistKey("config/user/settings.json");
+            service_->ExistKey("config/user/settings.json", "default");
         ASSERT_TRUE(exist_result_config.has_value());
         EXPECT_FALSE(exist_result_config.value());
 
-        auto exist_result_untouched = service_->ExistKey("prod_key_alpha");
+        auto exist_result_untouched =
+            service_->ExistKey("prod_key_alpha", "default");
         ASSERT_TRUE(exist_result_untouched.has_value());
         EXPECT_TRUE(exist_result_untouched.value());
     }
@@ -1903,12 +3269,12 @@ TEST_F(MasterServiceTest, RemoveAll) {
         ReplicateConfig config;
         config.replica_num = 1;
         auto put_start_result =
-            service_->PutStart(client_id, key, value_length, config);
+            service_->PutStart(client_id, key, "default", value_length, config);
         ASSERT_TRUE(put_start_result.has_value());
         auto put_end_result =
-            service_->PutEnd(client_id, key, ReplicaType::MEMORY);
+            service_->PutEnd(client_id, key, "default", ReplicaType::MEMORY);
         ASSERT_TRUE(put_end_result.has_value());
-        auto exist_result = service_->ExistKey(key);
+        auto exist_result = service_->ExistKey(key, "default");
         ASSERT_TRUE(exist_result.has_value());
     }
     // wait for all the lease to expire
@@ -1917,7 +3283,7 @@ TEST_F(MasterServiceTest, RemoveAll) {
     times = 10;
     while (times--) {
         std::string key = "test_key" + std::to_string(times);
-        auto exist_result = service_->ExistKey(key);
+        auto exist_result = service_->ExistKey(key, "default");
         ASSERT_TRUE(exist_result.has_value());
         ASSERT_FALSE(exist_result.value());
     }
@@ -1952,7 +3318,7 @@ TEST_F(MasterServiceTest, SingleSliceMultiReplicaFlow) {
 
     // Test PutStart with multiple slices and replicas
     auto put_start_result =
-        service_->PutStart(client_id, key, slice_length, config);
+        service_->PutStart(client_id, key, "default", slice_length, config);
     ASSERT_TRUE(put_start_result.has_value());
     replica_list = put_start_result.value();
 
@@ -1968,16 +3334,17 @@ TEST_F(MasterServiceTest, SingleSliceMultiReplicaFlow) {
     }
 
     // Test GetReplicaList during processing (should fail)
-    auto get_result = service_->GetReplicaList(key);
+    auto get_result = service_->GetReplicaList(key, "default");
     EXPECT_FALSE(get_result.has_value());
     EXPECT_EQ(ErrorCode::REPLICA_IS_NOT_READY, get_result.error());
 
     // Complete the put operation
-    auto put_end_result = service_->PutEnd(client_id, key, ReplicaType::MEMORY);
+    auto put_end_result =
+        service_->PutEnd(client_id, key, "default", ReplicaType::MEMORY);
     ASSERT_TRUE(put_end_result.has_value());
 
     // Test GetReplicaList after completion
-    auto get_result2 = service_->GetReplicaList(key);
+    auto get_result2 = service_->GetReplicaList(key, "default");
     ASSERT_TRUE(get_result2.has_value());
     auto retrieved_replicas = get_result2.value().replicas;
     ASSERT_EQ(num_replicas, retrieved_replicas.size());
@@ -2011,13 +3378,14 @@ TEST_F(MasterServiceTest, CleanupStaleHandlesTest) {
 
     // Create the object
     auto put_start_result =
-        service_->PutStart(client_id, key, slice_length, config);
+        service_->PutStart(client_id, key, "default", slice_length, config);
     ASSERT_TRUE(put_start_result.has_value());
-    auto put_end_result = service_->PutEnd(client_id, key, ReplicaType::MEMORY);
+    auto put_end_result =
+        service_->PutEnd(client_id, key, "default", ReplicaType::MEMORY);
     ASSERT_TRUE(put_end_result.has_value());
 
     // Verify object exists
-    auto get_result = service_->GetReplicaList(key);
+    auto get_result = service_->GetReplicaList(key, "default");
     ASSERT_TRUE(get_result.has_value());
     auto retrieved_replicas = get_result.value().replicas;
     ASSERT_EQ(1, retrieved_replicas.size());
@@ -2028,7 +3396,7 @@ TEST_F(MasterServiceTest, CleanupStaleHandlesTest) {
 
     // Try to get the object - it should be automatically removed since the
     // replica is invalid
-    auto get_result2 = service_->GetReplicaList(key);
+    auto get_result2 = service_->GetReplicaList(key, "default");
     EXPECT_FALSE(get_result2.has_value());
     EXPECT_EQ(ErrorCode::OBJECT_NOT_FOUND, get_result2.error());
 
@@ -2039,14 +3407,14 @@ TEST_F(MasterServiceTest, CleanupStaleHandlesTest) {
     // Create another object
     std::string key2 = "another_segment_object";
     auto put_start_result2 =
-        service_->PutStart(client_id, key2, slice_length, config);
+        service_->PutStart(client_id, key2, "default", slice_length, config);
     ASSERT_TRUE(put_start_result2.has_value());
     auto put_end_result2 =
-        service_->PutEnd(client_id, key2, ReplicaType::MEMORY);
+        service_->PutEnd(client_id, key2, "default", ReplicaType::MEMORY);
     ASSERT_TRUE(put_end_result2.has_value());
 
     // Verify we can get it
-    auto get_result3 = service_->GetReplicaList(key2);
+    auto get_result3 = service_->GetReplicaList(key2, "default");
     ASSERT_TRUE(get_result3.has_value());
 
     // Unmount the segment
@@ -2054,7 +3422,7 @@ TEST_F(MasterServiceTest, CleanupStaleHandlesTest) {
     ASSERT_TRUE(unmount_result2.has_value());
 
     // Try to remove the object that should already be cleaned up
-    auto remove_result = service_->Remove(key2);
+    auto remove_result = service_->Remove(key2, "default");
     EXPECT_FALSE(remove_result.has_value());
     EXPECT_EQ(ErrorCode::OBJECT_NOT_FOUND, remove_result.error());
 }
@@ -2086,11 +3454,11 @@ TEST_F(MasterServiceTest, ConcurrentWriteAndRemoveAll) {
                 config.replica_num = 1;
                 std::vector<Replica::Descriptor> replica_list;
 
-                auto put_start_result =
-                    service_->PutStart(client_id, key, slice_length, config);
+                auto put_start_result = service_->PutStart(
+                    client_id, key, "default", slice_length, config);
                 if (put_start_result.has_value()) {
-                    auto put_end_result =
-                        service_->PutEnd(client_id, key, ReplicaType::MEMORY);
+                    auto put_end_result = service_->PutEnd(
+                        client_id, key, "default", ReplicaType::MEMORY);
                     if (put_end_result.has_value()) {
                         success_writes++;
                     }
@@ -2156,10 +3524,10 @@ TEST_F(MasterServiceTest, ConcurrentReadAndRemoveAll) {
         config.replica_num = 1;
 
         auto put_start_result =
-            service_->PutStart(client_id, key, slice_length, config);
+            service_->PutStart(client_id, key, "default", slice_length, config);
         ASSERT_TRUE(put_start_result.has_value());
         auto put_end_result =
-            service_->PutEnd(client_id, key, ReplicaType::MEMORY);
+            service_->PutEnd(client_id, key, "default", ReplicaType::MEMORY);
         ASSERT_TRUE(put_end_result.has_value());
     }
 
@@ -2172,7 +3540,7 @@ TEST_F(MasterServiceTest, ConcurrentReadAndRemoveAll) {
         readers.emplace_back([&]() {
             for (int j = 0; j < num_objects; ++j) {
                 std::string key = "pre_key_" + std::to_string(j);
-                auto get_result = service_->GetReplicaList(key);
+                auto get_result = service_->GetReplicaList(key, "default");
                 if (get_result.has_value()) {
                     success_reads++;
                 }
@@ -2213,7 +3581,7 @@ TEST_F(MasterServiceTest, ConcurrentReadAndRemoveAll) {
     // Verify all objects were removed
     for (int i = 0; i < num_objects; ++i) {
         std::string key = "pre_key_" + std::to_string(i);
-        auto get_result = service_->GetReplicaList(key);
+        auto get_result = service_->GetReplicaList(key, "default");
         EXPECT_FALSE(get_result.has_value());
         EXPECT_EQ(ErrorCode::OBJECT_NOT_FOUND, get_result.error());
     }
@@ -2237,10 +3605,10 @@ TEST_F(MasterServiceTest, ConcurrentRemoveAllOperations) {
         config.replica_num = 1;
 
         auto put_start_result =
-            service_->PutStart(client_id, key, slice_length, config);
+            service_->PutStart(client_id, key, "default", slice_length, config);
         ASSERT_TRUE(put_start_result.has_value());
         auto put_end_result =
-            service_->PutEnd(client_id, key, ReplicaType::MEMORY);
+            service_->PutEnd(client_id, key, "default", ReplicaType::MEMORY);
         ASSERT_TRUE(put_end_result.has_value());
     }
 
@@ -2267,7 +3635,7 @@ TEST_F(MasterServiceTest, ConcurrentRemoveAllOperations) {
     // Verify all objects were removed
     for (int i = 0; i < num_objects; ++i) {
         std::string key = "pre_key_" + std::to_string(i);
-        auto get_result = service_->GetReplicaList(key);
+        auto get_result = service_->GetReplicaList(key, "default");
         EXPECT_FALSE(get_result.has_value());
         EXPECT_EQ(ErrorCode::OBJECT_NOT_FOUND, get_result.error());
     }
@@ -2304,23 +3672,23 @@ TEST_F(MasterServiceTest, UnmountSegmentImmediateCleanup) {
     // Umount will remove all objects in the segment, include the key1
     ASSERT_EQ(1, service_->GetKeyCount());
     // Verify objects in segment1 is gone
-    auto get_result1 = service_->GetReplicaList(key1);
+    auto get_result1 = service_->GetReplicaList(key1, "default");
     ASSERT_FALSE(get_result1.has_value());
     ASSERT_EQ(ErrorCode::OBJECT_NOT_FOUND, get_result1.error());
 
     // Verify objects in segment2 is still there
-    auto get_result2 = service_->GetReplicaList(key2);
+    auto get_result2 = service_->GetReplicaList(key2, "default");
     ASSERT_TRUE(get_result2.has_value());
 
     // Verify put key1 will put into segment2 rather than segment1
     auto put_start_result =
-        service_->PutStart(client_id, key1, slice_length, config);
+        service_->PutStart(client_id, key1, "default", slice_length, config);
     ASSERT_TRUE(put_start_result.has_value());
     replica_list = put_start_result.value();
     auto put_end_result =
-        service_->PutEnd(client_id, key1, ReplicaType::MEMORY);
+        service_->PutEnd(client_id, key1, "default", ReplicaType::MEMORY);
     ASSERT_TRUE(put_end_result.has_value());
-    auto get_result3 = service_->GetReplicaList(key1);
+    auto get_result3 = service_->GetReplicaList(key1, "default");
     ASSERT_TRUE(get_result3.has_value());
     auto retrieved = get_result3.value();
     ASSERT_EQ(replica_list[0]
@@ -2353,14 +3721,14 @@ TEST_F(MasterServiceTest, ReadableAfterPartialUnmountWithReplication) {
     config.replica_num = 2;
 
     auto put_start_result =
-        service_->PutStart(client_id, key, slice_length, config);
+        service_->PutStart(client_id, key, "default", slice_length, config);
     ASSERT_TRUE(put_start_result.has_value());
     ASSERT_EQ(2u, put_start_result->size());
-    ASSERT_TRUE(
-        service_->PutEnd(client_id, key, ReplicaType::MEMORY).has_value());
+    ASSERT_TRUE(service_->PutEnd(client_id, key, "default", ReplicaType::MEMORY)
+                    .has_value());
 
     // Verify two replicas exist and they are on distinct segments
-    auto get_result = service_->GetReplicaList(key);
+    auto get_result = service_->GetReplicaList(key, "default");
     ASSERT_TRUE(get_result.has_value());
     auto replicas = get_result.value().replicas;
     ASSERT_EQ(2u, replicas.size());
@@ -2378,7 +3746,7 @@ TEST_F(MasterServiceTest, ReadableAfterPartialUnmountWithReplication) {
     ASSERT_TRUE(service_->UnmountSegment(segment1.id, client_id).has_value());
 
     // Key should still be readable via the remaining replica
-    auto get_after_unmount = service_->GetReplicaList(key);
+    auto get_after_unmount = service_->GetReplicaList(key, "default");
     ASSERT_TRUE(get_after_unmount.has_value())
         << "Object should remain accessible with surviving replica";
 }
@@ -2428,7 +3796,7 @@ TEST_F(MasterServiceTest, UnmountSegmentPerformance) {
 
     // Verify all keys are gone
     for (const auto& key : keys) {
-        auto get_result = service_->GetReplicaList(key);
+        auto get_result = service_->GetReplicaList(key, "default");
         EXPECT_FALSE(get_result.has_value());
         EXPECT_EQ(ErrorCode::OBJECT_NOT_FOUND, get_result.error());
     }
@@ -2459,75 +3827,76 @@ TEST_F(MasterServiceTest, RemoveLeasedObject) {
 
     // Verify lease is granted on ExistsKey
     auto put_start_result =
-        service_->PutStart(client_id, key, slice_length, config);
+        service_->PutStart(client_id, key, "default", slice_length, config);
     ASSERT_TRUE(put_start_result.has_value());
-    auto put_end_result = service_->PutEnd(client_id, key, ReplicaType::MEMORY);
+    auto put_end_result =
+        service_->PutEnd(client_id, key, "default", ReplicaType::MEMORY);
     ASSERT_TRUE(put_end_result.has_value());
-    auto exist_result = service_->ExistKey(key);
+    auto exist_result = service_->ExistKey(key, "default");
     ASSERT_TRUE(exist_result.has_value());
-    auto remove_result = service_->Remove(key);
+    auto remove_result = service_->Remove(key, "default");
     EXPECT_FALSE(remove_result.has_value());
     EXPECT_EQ(ErrorCode::OBJECT_HAS_LEASE, remove_result.error());
     std::this_thread::sleep_for(std::chrono::milliseconds(kv_lease_ttl));
-    auto remove_result2 = service_->Remove(key);
+    auto remove_result2 = service_->Remove(key, "default");
     EXPECT_TRUE(remove_result2.has_value());
 
     // Verify lease is extended on successive ExistsKey
     auto put_start_result2 =
-        service_->PutStart(client_id, key, slice_length, config);
+        service_->PutStart(client_id, key, "default", slice_length, config);
     ASSERT_TRUE(put_start_result2.has_value());
     auto put_end_result2 =
-        service_->PutEnd(client_id, key, ReplicaType::MEMORY);
+        service_->PutEnd(client_id, key, "default", ReplicaType::MEMORY);
     ASSERT_TRUE(put_end_result2.has_value());
-    auto exist_result2 = service_->ExistKey(key);
+    auto exist_result2 = service_->ExistKey(key, "default");
     ASSERT_TRUE(exist_result2.has_value());
     std::this_thread::sleep_for(std::chrono::milliseconds(kv_lease_ttl));
-    auto exist_result3 = service_->ExistKey(key);
+    auto exist_result3 = service_->ExistKey(key, "default");
     ASSERT_TRUE(exist_result3.has_value());
-    auto remove_result3 = service_->Remove(key);
+    auto remove_result3 = service_->Remove(key, "default");
     EXPECT_FALSE(remove_result3.has_value());
     EXPECT_EQ(ErrorCode::OBJECT_HAS_LEASE, remove_result3.error());
     std::this_thread::sleep_for(std::chrono::milliseconds(kv_lease_ttl));
-    auto remove_result4 = service_->Remove(key);
+    auto remove_result4 = service_->Remove(key, "default");
     EXPECT_TRUE(remove_result4.has_value());
 
     // Verify lease is granted on GetReplicaList
     auto put_start_result3 =
-        service_->PutStart(client_id, key, slice_length, config);
+        service_->PutStart(client_id, key, "default", slice_length, config);
     ASSERT_TRUE(put_start_result3.has_value());
     auto put_end_result3 =
-        service_->PutEnd(client_id, key, ReplicaType::MEMORY);
+        service_->PutEnd(client_id, key, "default", ReplicaType::MEMORY);
     ASSERT_TRUE(put_end_result3.has_value());
-    auto get_result = service_->GetReplicaList(key);
+    auto get_result = service_->GetReplicaList(key, "default");
     ASSERT_TRUE(get_result.has_value());
-    auto remove_result5 = service_->Remove(key);
+    auto remove_result5 = service_->Remove(key, "default");
     EXPECT_FALSE(remove_result5.has_value());
     EXPECT_EQ(ErrorCode::OBJECT_HAS_LEASE, remove_result5.error());
     std::this_thread::sleep_for(std::chrono::milliseconds(kv_lease_ttl));
-    auto remove_result6 = service_->Remove(key);
+    auto remove_result6 = service_->Remove(key, "default");
     EXPECT_TRUE(remove_result6.has_value());
 
     // Verify lease is extended on successive GetReplicaList
     auto put_start_result4 =
-        service_->PutStart(client_id, key, slice_length, config);
+        service_->PutStart(client_id, key, "default", slice_length, config);
     ASSERT_TRUE(put_start_result4.has_value());
     auto put_end_result4 =
-        service_->PutEnd(client_id, key, ReplicaType::MEMORY);
+        service_->PutEnd(client_id, key, "default", ReplicaType::MEMORY);
     ASSERT_TRUE(put_end_result4.has_value());
-    auto get_result2 = service_->GetReplicaList(key);
+    auto get_result2 = service_->GetReplicaList(key, "default");
     ASSERT_TRUE(get_result2.has_value());
     std::this_thread::sleep_for(std::chrono::milliseconds(kv_lease_ttl));
-    auto get_result3 = service_->GetReplicaList(key);
+    auto get_result3 = service_->GetReplicaList(key, "default");
     ASSERT_TRUE(get_result3.has_value());
-    auto remove_result7 = service_->Remove(key);
+    auto remove_result7 = service_->Remove(key, "default");
     EXPECT_FALSE(remove_result7.has_value());
     EXPECT_EQ(ErrorCode::OBJECT_HAS_LEASE, remove_result7.error());
     std::this_thread::sleep_for(std::chrono::milliseconds(kv_lease_ttl));
-    auto remove_result8 = service_->Remove(key);
+    auto remove_result8 = service_->Remove(key, "default");
     EXPECT_TRUE(remove_result8.has_value());
 
     // Verify object is removed
-    auto get_result4 = service_->GetReplicaList(key);
+    auto get_result4 = service_->GetReplicaList(key, "default");
     EXPECT_FALSE(get_result4.has_value());
     EXPECT_EQ(ErrorCode::OBJECT_NOT_FOUND, get_result4.error());
 }
@@ -2546,20 +3915,20 @@ TEST_F(MasterServiceTest, RemoveAllLeasedObject) {
         ReplicateConfig config;
         config.replica_num = 1;
         auto put_start_result =
-            service_->PutStart(client_id, key, slice_length, config);
+            service_->PutStart(client_id, key, "default", slice_length, config);
         ASSERT_TRUE(put_start_result.has_value());
         auto put_end_result =
-            service_->PutEnd(client_id, key, ReplicaType::MEMORY);
+            service_->PutEnd(client_id, key, "default", ReplicaType::MEMORY);
         ASSERT_TRUE(put_end_result.has_value());
         if (i >= 5) {
-            auto exist_result = service_->ExistKey(key);
+            auto exist_result = service_->ExistKey(key, "default");
             ASSERT_TRUE(exist_result.has_value());
         }
     }
     ASSERT_EQ(5, service_->RemoveAll());
     for (int i = 0; i < 5; ++i) {
         std::string key = "test_key" + std::to_string(i);
-        auto exist_result = service_->ExistKey(key);
+        auto exist_result = service_->ExistKey(key, "default");
         ASSERT_FALSE(exist_result.value());
     }
     // wait for all the lease to expire
@@ -2567,7 +3936,7 @@ TEST_F(MasterServiceTest, RemoveAllLeasedObject) {
     ASSERT_EQ(5, service_->RemoveAll());
     for (int i = 5; i < 10; ++i) {
         std::string key = "test_key" + std::to_string(i);
-        auto exist_result = service_->ExistKey(key);
+        auto exist_result = service_->ExistKey(key, "default");
         ASSERT_FALSE(exist_result.value());
     }
 }
@@ -2598,10 +3967,10 @@ TEST_F(MasterServiceTest, EvictObject) {
         ReplicateConfig config;
         config.replica_num = 1;
         auto put_start_result =
-            service_->PutStart(client_id, key, slice_length, config);
+            service_->PutStart(client_id, key, "default", slice_length, config);
         if (put_start_result.has_value()) {
-            auto put_end_result =
-                service_->PutEnd(client_id, key, ReplicaType::MEMORY);
+            auto put_end_result = service_->PutEnd(client_id, key, "default",
+                                                   ReplicaType::MEMORY);
             ASSERT_TRUE(put_end_result.has_value());
             success_puts++;
         } else {
@@ -2638,13 +4007,13 @@ TEST_F(MasterServiceTest, TryEvictLeasedObject) {
         ReplicateConfig config;
         config.replica_num = 1;
         auto put_start_result =
-            service_->PutStart(client_id, key, slice_length, config);
+            service_->PutStart(client_id, key, "default", slice_length, config);
         if (put_start_result.has_value()) {
-            auto put_end_result =
-                service_->PutEnd(client_id, key, ReplicaType::MEMORY);
+            auto put_end_result = service_->PutEnd(client_id, key, "default",
+                                                   ReplicaType::MEMORY);
             ASSERT_TRUE(put_end_result.has_value());
             // the object is leased
-            auto get_result = service_->GetReplicaList(key);
+            auto get_result = service_->GetReplicaList(key, "default");
             ASSERT_TRUE(get_result.has_value());
             leased_keys.push_back(key);
             success_puts++;
@@ -2658,7 +4027,7 @@ TEST_F(MasterServiceTest, TryEvictLeasedObject) {
     std::this_thread::sleep_for(std::chrono::milliseconds(50));
     // All leased objects should be accessible
     for (const auto& key : leased_keys) {
-        auto get_result = service_->GetReplicaList(key);
+        auto get_result = service_->GetReplicaList(key, "default");
         ASSERT_TRUE(get_result.has_value());
     }
     std::this_thread::sleep_for(std::chrono::milliseconds(kv_lease_ttl));
@@ -2692,16 +4061,18 @@ TEST_F(MasterServiceTest, RemoveSoftPinObject) {
 
     // Verify soft pin does not block remove
     ASSERT_TRUE(
-        service_->PutStart(client_id, key, slice_length, config).has_value());
-    ASSERT_TRUE(
-        service_->PutEnd(client_id, key, ReplicaType::MEMORY).has_value());
-    EXPECT_TRUE(service_->Remove(key).has_value());
+        service_->PutStart(client_id, key, "default", slice_length, config)
+            .has_value());
+    ASSERT_TRUE(service_->PutEnd(client_id, key, "default", ReplicaType::MEMORY)
+                    .has_value());
+    EXPECT_TRUE(service_->Remove(key, "default").has_value());
 
     // Verify soft pin does not block RemoveAll
     ASSERT_TRUE(
-        service_->PutStart(client_id, key, slice_length, config).has_value());
-    ASSERT_TRUE(
-        service_->PutEnd(client_id, key, ReplicaType::MEMORY).has_value());
+        service_->PutStart(client_id, key, "default", slice_length, config)
+            .has_value());
+    ASSERT_TRUE(service_->PutEnd(client_id, key, "default", ReplicaType::MEMORY)
+                    .has_value());
     EXPECT_EQ(1, service_->RemoveAll());
 }
 
@@ -2739,11 +4110,12 @@ TEST_F(MasterServiceTest, SoftPinObjectsNotEvictedBeforeOtherObjects) {
             soft_pin_config.with_soft_pin = true;
 
             ASSERT_TRUE(service_
-                            ->PutStart(client_id, pin_key, slice_length,
-                                       soft_pin_config)
+                            ->PutStart(client_id, pin_key, "default",
+                                       slice_length, soft_pin_config)
                             .has_value());
             ASSERT_TRUE(
-                service_->PutEnd(client_id, pin_key, ReplicaType::MEMORY)
+                service_
+                    ->PutEnd(client_id, pin_key, "default", ReplicaType::MEMORY)
                     .has_value());
         }
 
@@ -2754,10 +4126,12 @@ TEST_F(MasterServiceTest, SoftPinObjectsNotEvictedBeforeOtherObjects) {
             uint64_t slice_length = value_size;
             ReplicateConfig config;
             config.replica_num = 1;
-            if (service_->PutStart(client_id, key, slice_length, config)
+            if (service_
+                    ->PutStart(client_id, key, "default", slice_length, config)
                     .has_value()) {
                 ASSERT_TRUE(
-                    service_->PutEnd(client_id, key, ReplicaType::MEMORY)
+                    service_
+                        ->PutEnd(client_id, key, "default", ReplicaType::MEMORY)
                         .has_value());
             } else {
                 failed_puts++;
@@ -2770,7 +4144,8 @@ TEST_F(MasterServiceTest, SoftPinObjectsNotEvictedBeforeOtherObjects) {
         // pin_key should still be accessible
         for (int i = 0; i < 2; i++) {
             std::string pin_key = "pin_key" + std::to_string(i);
-            ASSERT_TRUE(service_->GetReplicaList(pin_key).has_value());
+            ASSERT_TRUE(
+                service_->GetReplicaList(pin_key, "default").has_value());
         }
 
         // wait for the lease to expire
@@ -2809,10 +4184,11 @@ TEST_F(MasterServiceTest, SoftPinObjectsCanBeEvicted) {
         ReplicateConfig config;
         config.replica_num = 1;
         config.with_soft_pin = true;
-        if (service_->PutStart(client_id, key, slice_length, config)
+        if (service_->PutStart(client_id, key, "default", slice_length, config)
                 .has_value()) {
-            ASSERT_TRUE(service_->PutEnd(client_id, key, ReplicaType::MEMORY)
-                            .has_value());
+            ASSERT_TRUE(
+                service_->PutEnd(client_id, key, "default", ReplicaType::MEMORY)
+                    .has_value());
             success_puts++;
         } else {
             // wait for eviction to work
@@ -2861,10 +4237,11 @@ TEST_F(MasterServiceTest, SoftPinExtendedOnGet) {
             soft_pin_config.replica_num = 1;
             soft_pin_config.with_soft_pin = true;
 
-            ASSERT_TRUE(service_->PutStart(client_id, pin_key, slice_length,
-                                           soft_pin_config));
+            ASSERT_TRUE(service_->PutStart(client_id, pin_key, "default",
+                                           slice_length, soft_pin_config));
             ASSERT_TRUE(
-                service_->PutEnd(client_id, pin_key, ReplicaType::MEMORY)
+                service_
+                    ->PutEnd(client_id, pin_key, "default", ReplicaType::MEMORY)
                     .has_value());
         }
 
@@ -2874,7 +4251,8 @@ TEST_F(MasterServiceTest, SoftPinExtendedOnGet) {
         // Get the pin_key to extend the soft pin
         for (int i = 0; i < 2; i++) {
             std::string pin_key = "pin_key" + std::to_string(i);
-            ASSERT_TRUE(service_->GetReplicaList(pin_key).has_value());
+            ASSERT_TRUE(
+                service_->GetReplicaList(pin_key, "default").has_value());
         }
 
         // Fill the segment to trigger eviction
@@ -2884,10 +4262,12 @@ TEST_F(MasterServiceTest, SoftPinExtendedOnGet) {
             uint64_t slice_length = value_size;
             ReplicateConfig config;
             config.replica_num = 1;
-            if (service_->PutStart(client_id, key, slice_length, config)
+            if (service_
+                    ->PutStart(client_id, key, "default", slice_length, config)
                     .has_value()) {
                 ASSERT_TRUE(
-                    service_->PutEnd(client_id, key, ReplicaType::MEMORY)
+                    service_
+                        ->PutEnd(client_id, key, "default", ReplicaType::MEMORY)
                         .has_value());
             } else {
                 failed_puts++;
@@ -2901,7 +4281,8 @@ TEST_F(MasterServiceTest, SoftPinExtendedOnGet) {
         // pin_key should still be accessible
         for (int i = 0; i < 2; i++) {
             std::string pin_key = "pin_key" + std::to_string(i);
-            ASSERT_TRUE(service_->GetReplicaList(pin_key).has_value());
+            ASSERT_TRUE(
+                service_->GetReplicaList(pin_key, "default").has_value());
         }
 
         // wait for the lease to expire
@@ -2942,10 +4323,11 @@ TEST_F(MasterServiceTest, SoftPinObjectsNotAllowEvict) {
         ReplicateConfig config;
         config.replica_num = 1;
         config.with_soft_pin = true;
-        if (service_->PutStart(client_id, key, slice_length, config)
+        if (service_->PutStart(client_id, key, "default", slice_length, config)
                 .has_value()) {
-            ASSERT_TRUE(service_->PutEnd(client_id, key, ReplicaType::MEMORY)
-                            .has_value());
+            ASSERT_TRUE(
+                service_->PutEnd(client_id, key, "default", ReplicaType::MEMORY)
+                    .has_value());
             success_keys.push_back(key);
         } else {
             // wait for eviction to work
@@ -2955,7 +4337,7 @@ TEST_F(MasterServiceTest, SoftPinObjectsNotAllowEvict) {
     ASSERT_LE(success_keys.size(), 17);
     // All soft pinned objects should be accessible
     for (const auto& key : success_keys) {
-        ASSERT_TRUE(service_->GetReplicaList(key).has_value());
+        ASSERT_TRUE(service_->GetReplicaList(key, "default").has_value());
     }
     std::this_thread::sleep_for(std::chrono::milliseconds(kv_lease_ttl));
     service_->RemoveAll();
@@ -2981,7 +4363,7 @@ TEST_F(MasterServiceTest, ReplicaSegmentsAreUnique) {
     config.replica_num = 10;
 
     auto put_start_result =
-        service_->PutStart(client_id, key, slice_length, config);
+        service_->PutStart(client_id, key, "default", slice_length, config);
     ASSERT_TRUE(put_start_result.has_value());
     auto replica_list_local = put_start_result.value();
     ASSERT_EQ(config.replica_num, replica_list_local.size());
@@ -2997,8 +4379,8 @@ TEST_F(MasterServiceTest, ReplicaSegmentsAreUnique) {
     EXPECT_EQ(segment_names.size(), config.replica_num)
         << "Duplicate segment found";
 
-    ASSERT_TRUE(
-        service_->PutEnd(client_id, key, ReplicaType::MEMORY).has_value());
+    ASSERT_TRUE(service_->PutEnd(client_id, key, "default", ReplicaType::MEMORY)
+                    .has_value());
 }
 
 TEST_F(MasterServiceTest, ReplicationFactorTwoWithSingleSegment) {
@@ -3011,23 +4393,23 @@ TEST_F(MasterServiceTest, ReplicationFactorTwoWithSingleSegment) {
     [[maybe_unused]] const auto context = PrepareSimpleSegment(
         *service_, "single_segment", kBaseAddr, kSegmentSize);
 
-    // Request replication factor 2 with a single 1KB slice
-    // With best-effort semantics, should succeed with 1 replica
+    // Request replication factor 2 with a single 1KB slice.
+    // With best-effort semantics, should succeed with 1 replica.
     const std::string key = "replication_factor_two_single_segment";
     uint64_t slice_length = 1024;
     ReplicateConfig config;
     config.replica_num = 2;
 
     auto put_start_result =
-        service_->PutStart(client_id, key, slice_length, config);
+        service_->PutStart(client_id, key, "default", slice_length, config);
     ASSERT_TRUE(put_start_result.has_value());
     auto replicas = put_start_result.value();
 
-    // Should get 1 replica instead of the requested 2 (best-effort)
+    // Should get 1 replica instead of the requested 2 (best-effort).
     EXPECT_EQ(1u, replicas.size());
     EXPECT_TRUE(replicas[0].is_memory_replica());
 
-    // Verify the replica is properly allocated on the single segment
+    // Verify the replica is properly allocated on the single segment.
     auto mem_desc = replicas[0].get_memory_descriptor();
     EXPECT_EQ("single_segment", mem_desc.buffer_descriptor.transport_endpoint_);
     EXPECT_EQ(1024u, mem_desc.buffer_descriptor.size_);
@@ -3051,29 +4433,182 @@ TEST_F(MasterServiceTest, BatchExistKeyTest) {
         ReplicateConfig config;
         config.replica_num = 1;
         uint64_t slice_length = value_size;
-        auto put_start_result =
-            service_->PutStart(client_id, test_keys[i], slice_length, config);
+        auto put_start_result = service_->PutStart(
+            client_id, test_keys[i], "default", slice_length, config);
         ASSERT_TRUE(put_start_result.has_value());
-        auto put_end_result =
-            service_->PutEnd(client_id, test_keys[i], ReplicaType::MEMORY);
+        auto put_end_result = service_->PutEnd(client_id, test_keys[i],
+                                               "default", ReplicaType::MEMORY);
         ASSERT_TRUE(put_end_result.has_value());
     }
 
     // Test individual ExistKey calls to verify the underlying functionality
     for (int i = 0; i < test_object_num; ++i) {
-        auto exist_result = service_->ExistKey(test_keys[i]);
+        auto exist_result = service_->ExistKey(test_keys[i], "default");
         EXPECT_TRUE(exist_result.value());
     }
 
     // Tets batch
     test_keys.push_back("non_existent_key");
-    auto exist_resp = service_->BatchExistKey(test_keys);
+    auto exist_resp = service_->BatchExistKey(test_keys, "default");
     for (int i = 0; i < test_object_num; ++i) {
         ASSERT_TRUE(exist_resp[i].value());
     }
     ASSERT_FALSE(exist_resp[test_object_num].value());
 }
 
+TEST_F(MasterServiceTest, BatchExistKeyGroupedAndIncompletePreservesOrder) {
+    std::unique_ptr<MasterService> service_(new MasterService());
+    const UUID client_id = generate_uuid();
+
+    constexpr size_t buffer = 0x300000000;
+    constexpr size_t size = 1024 * 1024 * 128;
+    [[maybe_unused]] const auto context =
+        PrepareSimpleSegment(*service_, "test_segment", buffer, size);
+
+    const std::string grouped_key_a = "batch_grouped_key_a";
+    const std::string grouped_key_b = "batch_grouped_key_b";
+    const std::string group_id = FindGroupIdOnDifferentShard(grouped_key_a);
+    ReplicateConfig grouped_config;
+    grouped_config.replica_num = 1;
+    grouped_config.group_ids = std::vector<std::string>{group_id};
+    PutCompletedObject(*service_, client_id, grouped_key_a, grouped_config);
+    PutCompletedObject(*service_, client_id, grouped_key_b, grouped_config);
+
+    const std::string completed_key = "batch_completed_key";
+    ReplicateConfig config;
+    config.replica_num = 1;
+    PutCompletedObject(*service_, client_id, completed_key, config);
+
+    const std::string incomplete_key = "batch_incomplete_key";
+    ASSERT_TRUE(
+        service_->PutStart(client_id, incomplete_key, "default", 1024, config)
+            .has_value());
+
+    const std::string missing_key = "batch_missing_key";
+    std::vector<std::string> keys = {grouped_key_a, completed_key,
+                                     incomplete_key, missing_key,
+                                     grouped_key_b};
+
+    auto resp = service_->BatchExistKey(keys, "default");
+    ASSERT_EQ(resp.size(), keys.size());
+    ASSERT_TRUE(resp[0].has_value());
+    ASSERT_TRUE(resp[1].has_value());
+    ASSERT_TRUE(resp[2].has_value());
+    ASSERT_TRUE(resp[3].has_value());
+    ASSERT_TRUE(resp[4].has_value());
+    EXPECT_TRUE(resp[0].value());
+    EXPECT_TRUE(resp[1].value());
+    EXPECT_FALSE(resp[2].value());
+    EXPECT_FALSE(resp[3].value());
+    EXPECT_TRUE(resp[4].value());
+}
+
+TEST_F(MasterServiceTest, BatchExistKeyTenantAwarePreservesOrder) {
+    const std::string tenant_id = "tenant_batch_exist";
+    auto service_ = std::make_unique<MasterService>(
+        MakeStrictTenantConfig({"default", tenant_id}));
+    const UUID client_id = generate_uuid();
+
+    constexpr size_t buffer = 0x300000000;
+    constexpr size_t size = 1024 * 1024 * 128;
+    [[maybe_unused]] const auto context =
+        PrepareSimpleSegment(*service_, "test_segment", buffer, size);
+
+    ReplicateConfig config;
+    config.replica_num = 1;
+
+    const std::string tenant_only_key = "batch_tenant_only";
+    const std::string default_only_key = "batch_default_only";
+    const std::string incomplete_key = "batch_tenant_incomplete";
+    const std::string missing_key = "batch_tenant_missing";
+
+    PutCompletedObject(*service_, client_id, tenant_only_key, tenant_id,
+                       config);
+    PutCompletedObject(*service_, client_id, default_only_key, config);
+    ASSERT_TRUE(
+        service_->PutStart(client_id, incomplete_key, tenant_id, 1024, config)
+            .has_value());
+
+    std::vector<std::string> tenant_keys = {tenant_only_key, default_only_key,
+                                            missing_key, incomplete_key,
+                                            tenant_only_key};
+    auto tenant_resp = service_->BatchExistKey(tenant_keys, tenant_id);
+    ASSERT_EQ(tenant_resp.size(), tenant_keys.size());
+    EXPECT_TRUE(tenant_resp[0].value());
+    EXPECT_FALSE(tenant_resp[1].value());
+    EXPECT_FALSE(tenant_resp[2].value());
+    EXPECT_FALSE(tenant_resp[3].value());
+    EXPECT_TRUE(tenant_resp[4].value());
+
+    std::vector<std::string> default_keys = {tenant_only_key, default_only_key};
+    auto default_resp = service_->BatchExistKey(default_keys, "default");
+    ASSERT_EQ(default_resp.size(), default_keys.size());
+    EXPECT_FALSE(default_resp[0].value());
+    EXPECT_TRUE(default_resp[1].value());
+}
+
+TEST_F(MasterServiceTest, WrappedBatchExistKeyUsesTenantAwareBatchPath) {
+    const std::string tenant_id = "wrapped_batch_exist_tenant";
+    auto service_config = MakeStrictWrappedConfig({"default", tenant_id});
+    WrappedMasterService service_(service_config);
+
+    Segment segment = MakeSegment("wrapped_batch_exist_segment");
+    const UUID client_id = generate_uuid();
+    ASSERT_TRUE(service_.MountSegment(segment, client_id).has_value());
+
+    ReplicateConfig config;
+    config.replica_num = 1;
+    const std::string tenant_key_a = "wrapped_batch_tenant_a";
+    const std::string tenant_key_b = "wrapped_batch_tenant_b";
+    const std::string default_only_key = "wrapped_batch_default_only";
+    const std::string missing_key = "wrapped_batch_missing";
+
+    std::vector<std::string> tenant_keys = {tenant_key_a, tenant_key_b};
+    std::vector<uint64_t> tenant_sizes = {1024, 2048};
+    auto tenant_put_start = service_.BatchPutStart(
+        client_id, tenant_keys, tenant_sizes, config, tenant_id);
+    ASSERT_EQ(tenant_put_start.size(), tenant_keys.size());
+    for (const auto& result : tenant_put_start) {
+        ASSERT_TRUE(result.has_value()) << toString(result.error());
+    }
+    auto tenant_put_end = service_.BatchPutEnd(client_id, tenant_keys,
+                                               ReplicaType::MEMORY, tenant_id);
+    ASSERT_EQ(tenant_put_end.size(), tenant_keys.size());
+    for (const auto& result : tenant_put_end) {
+        ASSERT_TRUE(result.has_value()) << toString(result.error());
+    }
+
+    auto default_put_start =
+        service_.PutStart(client_id, default_only_key, 1024, config);
+    ASSERT_TRUE(default_put_start.has_value());
+    ASSERT_TRUE(
+        service_.PutEnd(client_id, default_only_key, ReplicaType::MEMORY)
+            .has_value());
+
+    auto& metrics = MasterMetricManager::instance();
+    const auto base_requests = metrics.get_batch_exist_key_requests();
+    const auto base_items = metrics.get_batch_exist_key_items();
+    const auto base_failures = metrics.get_batch_exist_key_failures();
+    const auto base_partial = metrics.get_batch_exist_key_partial_successes();
+    const auto base_failed_items = metrics.get_batch_exist_key_failed_items();
+
+    std::vector<std::string> lookup_keys = {tenant_key_a, default_only_key,
+                                            missing_key, tenant_key_b};
+    auto resp = service_.BatchExistKey(lookup_keys, tenant_id);
+    ASSERT_EQ(resp.size(), lookup_keys.size());
+    EXPECT_TRUE(resp[0].value());
+    EXPECT_FALSE(resp[1].value());
+    EXPECT_FALSE(resp[2].value());
+    EXPECT_TRUE(resp[3].value());
+
+    EXPECT_EQ(base_requests + 1, metrics.get_batch_exist_key_requests());
+    EXPECT_EQ(base_items + lookup_keys.size(),
+              metrics.get_batch_exist_key_items());
+    EXPECT_EQ(base_failures, metrics.get_batch_exist_key_failures());
+    EXPECT_EQ(base_partial, metrics.get_batch_exist_key_partial_successes());
+    EXPECT_EQ(base_failed_items, metrics.get_batch_exist_key_failed_items());
+}
+
 TEST_F(MasterServiceTest, BatchQueryIpTest) {
     std::unique_ptr<MasterService> service_(new MasterService());
     const UUID client_id = generate_uuid();
@@ -3216,6 +4751,117 @@ TEST_F(MasterServiceTest, BatchQueryIpMultipleSegmentsEmptyTeEndpointTest) {
         << "Client with all empty te_endpoints should have empty IP vector";
 }
 
+TEST_F(MasterServiceTest, BatchQueryIpBracketedIpv6Test) {
+    std::unique_ptr<MasterService> service_(new MasterService());
+    const UUID client_id = generate_uuid();
+
+    // Mount a segment with a bracketed IPv6 endpoint
+    constexpr size_t buffer = 0x300000000;
+    constexpr size_t size = 1024 * 1024 * 16;
+    Segment segment = MakeSegment("test_segment", buffer, size);
+    segment.te_endpoint = "[::1]:17813";
+    auto mount_result = service_->MountSegment(segment, client_id);
+    ASSERT_TRUE(mount_result.has_value());
+
+    std::vector<UUID> client_ids = {client_id};
+    auto query_result = service_->BatchQueryIp(client_ids);
+    ASSERT_TRUE(query_result.has_value());
+
+    const auto& results = query_result.value();
+    auto it = results.find(client_id);
+    ASSERT_NE(it, results.end());
+
+    const auto& ip_addresses = it->second;
+    ASSERT_EQ(1u, ip_addresses.size());
+    EXPECT_EQ("::1", ip_addresses[0]);
+}
+
+TEST_F(MasterServiceTest, BatchQueryIpLinkLocalIpv6WithScopeTest) {
+    std::unique_ptr<MasterService> service_(new MasterService());
+    const UUID client_id = generate_uuid();
+
+    // Mount a segment with a link-local IPv6 address with scope ID
+    constexpr size_t buffer = 0x300000000;
+    constexpr size_t size = 1024 * 1024 * 16;
+    Segment segment = MakeSegment("test_segment", buffer, size);
+    segment.te_endpoint = "fe80::a236:bcff:fecb:a1be%eno2:15773";
+    auto mount_result = service_->MountSegment(segment, client_id);
+    ASSERT_TRUE(mount_result.has_value());
+
+    std::vector<UUID> client_ids = {client_id};
+    auto query_result = service_->BatchQueryIp(client_ids);
+    ASSERT_TRUE(query_result.has_value());
+
+    const auto& results = query_result.value();
+    auto it = results.find(client_id);
+    ASSERT_NE(it, results.end());
+
+    const auto& ip_addresses = it->second;
+    ASSERT_EQ(1u, ip_addresses.size());
+    EXPECT_EQ("fe80::a236:bcff:fecb:a1be%eno2", ip_addresses[0]);
+}
+
+TEST_F(MasterServiceTest, BatchQueryIpIpv6NoPortTest) {
+    std::unique_ptr<MasterService> service_(new MasterService());
+    const UUID client_id = generate_uuid();
+
+    // Mount a segment with an IPv6 address without port
+    constexpr size_t buffer = 0x300000000;
+    constexpr size_t size = 1024 * 1024 * 16;
+    Segment segment = MakeSegment("test_segment", buffer, size);
+    segment.te_endpoint = "::1";
+    auto mount_result = service_->MountSegment(segment, client_id);
+    ASSERT_TRUE(mount_result.has_value());
+
+    std::vector<UUID> client_ids = {client_id};
+    auto query_result = service_->BatchQueryIp(client_ids);
+    ASSERT_TRUE(query_result.has_value());
+
+    const auto& results = query_result.value();
+    auto it = results.find(client_id);
+    ASSERT_NE(it, results.end());
+
+    const auto& ip_addresses = it->second;
+    ASSERT_EQ(1u, ip_addresses.size());
+    EXPECT_EQ("::1", ip_addresses[0]);
+}
+
+TEST_F(MasterServiceTest, BatchQueryIpMixedIpv4AndIpv6Test) {
+    std::unique_ptr<MasterService> service_(new MasterService());
+    const UUID client_id = generate_uuid();
+
+    // Mount segments with IPv4 and IPv6 endpoints for the same client
+    constexpr size_t buffer1 = 0x300000000;
+    constexpr size_t buffer2 = 0x400000000;
+    constexpr size_t size = 1024 * 1024 * 16;
+
+    Segment segment1 = MakeSegment("segment1", buffer1, size);
+    segment1.te_endpoint = "192.168.1.1:12345";
+    auto mount_result1 = service_->MountSegment(segment1, client_id);
+    ASSERT_TRUE(mount_result1.has_value());
+
+    Segment segment2 = MakeSegment("segment2", buffer2, size);
+    segment2.te_endpoint = "[::1]:17813";
+    auto mount_result2 = service_->MountSegment(segment2, client_id);
+    ASSERT_TRUE(mount_result2.has_value());
+
+    std::vector<UUID> client_ids = {client_id};
+    auto query_result = service_->BatchQueryIp(client_ids);
+    ASSERT_TRUE(query_result.has_value());
+
+    const auto& results = query_result.value();
+    auto it = results.find(client_id);
+    ASSERT_NE(it, results.end());
+
+    const auto& ip_addresses = it->second;
+    ASSERT_EQ(2u, ip_addresses.size());
+
+    std::unordered_set<std::string> ip_set(ip_addresses.begin(),
+                                           ip_addresses.end());
+    EXPECT_NE(ip_set.find("192.168.1.1"), ip_set.end());
+    EXPECT_NE(ip_set.find("::1"), ip_set.end());
+}
+
 TEST_F(MasterServiceTest, PutStartExpiringTest) {
     // Reset storage space metrics.
     MasterMetricManager::instance().reset_allocated_mem_size();
@@ -3250,7 +4896,7 @@ TEST_F(MasterServiceTest, PutStartExpiringTest) {
 
     // Put key_1, should success.
     auto put_start_result =
-        service_->PutStart(client_id, key_1, slice_length, config);
+        service_->PutStart(client_id, key_1, "default", slice_length, config);
     EXPECT_TRUE(put_start_result.has_value());
     replica_list = put_start_result.value();
     EXPECT_EQ(replica_list.size(), kReplicaCnt);
@@ -3260,7 +4906,7 @@ TEST_F(MasterServiceTest, PutStartExpiringTest) {
 
     // Put key_1 again, should fail because the key exists.
     put_start_result =
-        service_->PutStart(client_id, key_1, slice_length, config);
+        service_->PutStart(client_id, key_1, "default", slice_length, config);
     EXPECT_FALSE(put_start_result.has_value());
     EXPECT_EQ(put_start_result.error(), ErrorCode::OBJECT_ALREADY_EXISTS);
 
@@ -3276,7 +4922,7 @@ TEST_F(MasterServiceTest, PutStartExpiringTest) {
     // Put key_1 again, should success because the old one has expired and will
     // be discarded by this put.
     put_start_result =
-        service_->PutStart(client_id, key_1, slice_length, config);
+        service_->PutStart(client_id, key_1, "default", slice_length, config);
     EXPECT_TRUE(put_start_result.has_value());
     replica_list = put_start_result.value();
     EXPECT_EQ(replica_list.size(), kReplicaCnt);
@@ -3286,17 +4932,17 @@ TEST_F(MasterServiceTest, PutStartExpiringTest) {
 
     // Complete key_1.
     auto put_end_result =
-        service_->PutEnd(client_id, key_1, ReplicaType::MEMORY);
+        service_->PutEnd(client_id, key_1, "default", ReplicaType::MEMORY);
     EXPECT_TRUE(put_end_result.has_value());
 
     // Protect key_1 from eviction.
-    auto get_result = service_->GetReplicaList(key_1);
+    auto get_result = service_->GetReplicaList(key_1, "default");
     EXPECT_TRUE(get_result.has_value());
 
     // Put key_2, should fail because the key_1 occupied 12MB (6MB processing,
     // 6MB discarded but not yet released) on each segment.
     put_start_result =
-        service_->PutStart(client_id, key_2, slice_length, config);
+        service_->PutStart(client_id, key_2, "default", slice_length, config);
     EXPECT_FALSE(put_start_result.has_value());
     EXPECT_EQ(put_start_result.error(), ErrorCode::NO_AVAILABLE_HANDLE);
 
@@ -3309,7 +4955,7 @@ TEST_F(MasterServiceTest, PutStartExpiringTest) {
             EXPECT_TRUE(result.has_value());
         }
         // Protect key_1 from eviction.
-        auto get_result = service_->GetReplicaList(key_1);
+        auto get_result = service_->GetReplicaList(key_1, "default");
         EXPECT_TRUE(get_result.has_value());
         std::this_thread::sleep_for(std::chrono::seconds(1));
     }
@@ -3317,7 +4963,7 @@ TEST_F(MasterServiceTest, PutStartExpiringTest) {
     // Put key_2 again, should success because the discarded replica has been
     // released.
     put_start_result =
-        service_->PutStart(client_id, key_2, slice_length, config);
+        service_->PutStart(client_id, key_2, "default", slice_length, config);
     EXPECT_TRUE(put_start_result.has_value());
     replica_list = put_start_result.value();
     EXPECT_EQ(replica_list.size(), kReplicaCnt);
@@ -3332,7 +4978,7 @@ TEST_F(MasterServiceTest, PutStartExpiringTest) {
             EXPECT_TRUE(result.has_value());
         }
         // Protect key_1 from eviction.
-        auto get_result = service_->GetReplicaList(key_1);
+        auto get_result = service_->GetReplicaList(key_1, "default");
         EXPECT_TRUE(get_result.has_value());
         std::this_thread::sleep_for(std::chrono::seconds(1));
     }
@@ -3340,7 +4986,7 @@ TEST_F(MasterServiceTest, PutStartExpiringTest) {
     // Put key_2 again, should fail because eviction has not been triggered. And
     // this PutStart should trigger the eviction.
     put_start_result =
-        service_->PutStart(client_id, key_2, slice_length, config);
+        service_->PutStart(client_id, key_2, "default", slice_length, config);
     EXPECT_FALSE(put_start_result.has_value());
     EXPECT_EQ(put_start_result.error(), ErrorCode::NO_AVAILABLE_HANDLE);
 
@@ -3350,7 +4996,7 @@ TEST_F(MasterServiceTest, PutStartExpiringTest) {
     // Put key_2 again, should success because the previous one has been
     // discarded and released.
     put_start_result =
-        service_->PutStart(client_id, key_2, slice_length, config);
+        service_->PutStart(client_id, key_2, "default", slice_length, config);
     EXPECT_TRUE(put_start_result.has_value());
     replica_list = put_start_result.value();
     EXPECT_EQ(replica_list.size(), kReplicaCnt);
@@ -3359,7 +5005,8 @@ TEST_F(MasterServiceTest, PutStartExpiringTest) {
     }
 
     // Complete key_2.
-    put_end_result = service_->PutEnd(client_id, key_2, ReplicaType::MEMORY);
+    put_end_result =
+        service_->PutEnd(client_id, key_2, "default", ReplicaType::MEMORY);
     EXPECT_TRUE(put_end_result.has_value());
 }
 
@@ -3429,8 +5076,11 @@ TEST_F(MasterServiceTest, OffloadObjectHeartbeat) {
     }
     ASSERT_EQ(res->size(), keys.size());
     for (auto& key : keys) {
-        ASSERT_TRUE(res.value().find(key) != res.value().end());
-        ASSERT_EQ(res.value().find(key)->second, 1024);
+        auto it = std::find_if(
+            res->begin(), res->end(),
+            [&key](const OffloadTaskItem& task) { return task.key == key; });
+        ASSERT_TRUE(it != res->end());
+        ASSERT_EQ(it->size, 1024);
     }
 
     keys.clear();
@@ -3446,8 +5096,11 @@ TEST_F(MasterServiceTest, OffloadObjectHeartbeat) {
     }
     ASSERT_EQ(res->size(), keys.size());
     for (auto& key : keys) {
-        ASSERT_TRUE(res.value().find(key) != res.value().end());
-        ASSERT_EQ(res.value().find(key)->second, 1024);
+        auto it = std::find_if(
+            res->begin(), res->end(),
+            [&key](const OffloadTaskItem& task) { return task.key == key; });
+        ASSERT_TRUE(it != res->end());
+        ASSERT_EQ(it->size, 1024);
     }
 }
 
@@ -3470,16 +5123,16 @@ TEST_F(MasterServiceTest, BatchReplicaClearAllSegments) {
         ReplicateConfig config;
         config.replica_num = 1;
         auto put_start_result =
-            service_->PutStart(client_id, key, value_length, config);
+            service_->PutStart(client_id, key, "default", value_length, config);
         ASSERT_TRUE(put_start_result.has_value());
         auto put_end_result =
-            service_->PutEnd(client_id, key, ReplicaType::MEMORY);
+            service_->PutEnd(client_id, key, "default", ReplicaType::MEMORY);
         ASSERT_TRUE(put_end_result.has_value());
     }
 
     // Verify objects exist
     for (const auto& key : keys) {
-        auto exist_result = service_->ExistKey(key);
+        auto exist_result = service_->ExistKey(key, "default");
         ASSERT_TRUE(exist_result.has_value());
         ASSERT_TRUE(exist_result.value());
     }
@@ -3496,7 +5149,7 @@ TEST_F(MasterServiceTest, BatchReplicaClearAllSegments) {
 
     // Verify objects are removed
     for (const auto& key : keys) {
-        auto exist_result = service_->ExistKey(key);
+        auto exist_result = service_->ExistKey(key, "default");
         ASSERT_TRUE(exist_result.has_value());
         ASSERT_FALSE(exist_result.value())
             << "Key " << key << " should be removed";
@@ -3527,9 +5180,10 @@ TEST_F(MasterServiceTest, BatchReplicaClearSpecificSegment) {
     config.preferred_segment =
         segment_name;  // Ensure object is placed on segment1
     auto put_start_result =
-        service_->PutStart(client_id, key, value_length, config);
+        service_->PutStart(client_id, key, "default", value_length, config);
     ASSERT_TRUE(put_start_result.has_value());
-    auto put_end_result = service_->PutEnd(client_id, key, ReplicaType::MEMORY);
+    auto put_end_result =
+        service_->PutEnd(client_id, key, "default", ReplicaType::MEMORY);
     ASSERT_TRUE(put_end_result.has_value());
 
     // 4. Wait for lease to expire and verify it's actually expired
@@ -3567,7 +5221,7 @@ TEST_F(MasterServiceTest, BatchReplicaClearSpecificSegment) {
     const auto& cleared_keys = clear_result.value();
     ASSERT_EQ(1u, cleared_keys.size()) << "Key should be cleared";
 
-    auto exist_result = service_->ExistKey(key);
+    auto exist_result = service_->ExistKey(key, "default");
     ASSERT_TRUE(exist_result.has_value());
     ASSERT_FALSE(exist_result.value())
         << "Key should be removed after being cleared.";
@@ -3588,13 +5242,14 @@ TEST_F(MasterServiceTest, BatchReplicaClearWithLeaseActive) {
     ReplicateConfig config;
     config.replica_num = 1;
     auto put_start_result =
-        service_->PutStart(client_id, key, value_length, config);
+        service_->PutStart(client_id, key, "default", value_length, config);
     ASSERT_TRUE(put_start_result.has_value());
-    auto put_end_result = service_->PutEnd(client_id, key, ReplicaType::MEMORY);
+    auto put_end_result =
+        service_->PutEnd(client_id, key, "default", ReplicaType::MEMORY);
     ASSERT_TRUE(put_end_result.has_value());
 
     // Grant a lease by calling GetReplicaList (similar to normal usage)
-    auto get_result = service_->GetReplicaList(key);
+    auto get_result = service_->GetReplicaList(key, "default");
     ASSERT_TRUE(get_result.has_value());
 
     // Try to clear immediately (lease should still be active)
@@ -3608,7 +5263,7 @@ TEST_F(MasterServiceTest, BatchReplicaClearWithLeaseActive) {
         << "No keys should be cleared when lease is active";
 
     // Verify object still exists
-    auto exist_result = service_->ExistKey(key);
+    auto exist_result = service_->ExistKey(key, "default");
     ASSERT_TRUE(exist_result.has_value());
     ASSERT_TRUE(exist_result.value()) << "Key should still exist";
 }
@@ -3629,10 +5284,10 @@ TEST_F(MasterServiceTest, BatchReplicaClearWithDifferentClientId) {
     ReplicateConfig config;
     config.replica_num = 1;
     auto put_start_result =
-        service_->PutStart(client_id1, key, value_length, config);
+        service_->PutStart(client_id1, key, "default", value_length, config);
     ASSERT_TRUE(put_start_result.has_value());
     auto put_end_result =
-        service_->PutEnd(client_id1, key, ReplicaType::MEMORY);
+        service_->PutEnd(client_id1, key, "default", ReplicaType::MEMORY);
     ASSERT_TRUE(put_end_result.has_value());
 
     // Wait for lease to expire
@@ -3649,7 +5304,7 @@ TEST_F(MasterServiceTest, BatchReplicaClearWithDifferentClientId) {
         << "No keys should be cleared for different client_id";
 
     // Verify object still exists
-    auto exist_result = service_->ExistKey(key);
+    auto exist_result = service_->ExistKey(key, "default");
     ASSERT_TRUE(exist_result.has_value());
     ASSERT_TRUE(exist_result.value()) << "Key should still exist";
 }
@@ -3704,11 +5359,11 @@ TEST_F(MasterServiceTest, BatchReplicaClearWithEmptyStringKeys) {
     uint64_t value_length = 1024;
     ReplicateConfig config;
     config.replica_num = 1;
-    auto put_start_result =
-        service_->PutStart(client_id, valid_key, value_length, config);
+    auto put_start_result = service_->PutStart(client_id, valid_key, "default",
+                                               value_length, config);
     ASSERT_TRUE(put_start_result.has_value());
     auto put_end_result =
-        service_->PutEnd(client_id, valid_key, ReplicaType::MEMORY);
+        service_->PutEnd(client_id, valid_key, "default", ReplicaType::MEMORY);
     ASSERT_TRUE(put_end_result.has_value());
 
     // Wait for lease to expire
@@ -3746,22 +5401,25 @@ TEST_F(MasterServiceTest, BatchReplicaClearMixedScenario) {
 
     // Create key1 and key2 with client_id1
     auto put_start1 =
-        service_->PutStart(client_id1, key1, value_length, config);
+        service_->PutStart(client_id1, key1, "default", value_length, config);
     ASSERT_TRUE(put_start1.has_value());
-    auto put_end1 = service_->PutEnd(client_id1, key1, ReplicaType::MEMORY);
+    auto put_end1 =
+        service_->PutEnd(client_id1, key1, "default", ReplicaType::MEMORY);
     ASSERT_TRUE(put_end1.has_value());
 
     auto put_start2 =
-        service_->PutStart(client_id1, key2, value_length, config);
+        service_->PutStart(client_id1, key2, "default", value_length, config);
     ASSERT_TRUE(put_start2.has_value());
-    auto put_end2 = service_->PutEnd(client_id1, key2, ReplicaType::MEMORY);
+    auto put_end2 =
+        service_->PutEnd(client_id1, key2, "default", ReplicaType::MEMORY);
     ASSERT_TRUE(put_end2.has_value());
 
     // Create key3 with client_id2
     auto put_start3 =
-        service_->PutStart(client_id2, key3, value_length, config);
+        service_->PutStart(client_id2, key3, "default", value_length, config);
     ASSERT_TRUE(put_start3.has_value());
-    auto put_end3 = service_->PutEnd(client_id2, key3, ReplicaType::MEMORY);
+    auto put_end3 =
+        service_->PutEnd(client_id2, key3, "default", ReplicaType::MEMORY);
     ASSERT_TRUE(put_end3.has_value());
 
     // Wait for lease to expire
@@ -3779,16 +5437,16 @@ TEST_F(MasterServiceTest, BatchReplicaClearMixedScenario) {
         << "Only keys belonging to client_id1 should be cleared";
 
     // Verify key1 and key2 are cleared
-    auto exist1 = service_->ExistKey(key1);
+    auto exist1 = service_->ExistKey(key1, "default");
     ASSERT_TRUE(exist1.has_value());
     ASSERT_FALSE(exist1.value()) << "key1 should be cleared";
 
-    auto exist2 = service_->ExistKey(key2);
+    auto exist2 = service_->ExistKey(key2, "default");
     ASSERT_TRUE(exist2.has_value());
     ASSERT_FALSE(exist2.value()) << "key2 should be cleared";
 
     // Verify key3 still exists (different client_id)
-    auto exist3 = service_->ExistKey(key3);
+    auto exist3 = service_->ExistKey(key3, "default");
     ASSERT_TRUE(exist3.has_value());
     ASSERT_TRUE(exist3.value())
         << "key3 should still exist (different client_id)";
@@ -3824,15 +5482,15 @@ TEST_F(MasterServiceTest, CreateCopyTaskTest) {
     config.replica_num = 1;
     config.preferred_segment = "segment_0";
     auto put_start_result =
-        service_->PutStart(client_id, key1, slice_length, config);
+        service_->PutStart(client_id, key1, "default", slice_length, config);
     EXPECT_TRUE(put_start_result.has_value());
     auto put_end_result =
-        service_->PutEnd(client_id, key1, ReplicaType::MEMORY);
+        service_->PutEnd(client_id, key1, "default", ReplicaType::MEMORY);
     EXPECT_TRUE(put_end_result.has_value());
 
     // Copy key1 to "segment_1" and "segment_2"
     auto copy_result =
-        service_->CreateCopyTask(key1, {"segment_1", "segment_2"});
+        service_->CreateCopyTask(key1, "default", {"segment_1", "segment_2"});
     EXPECT_TRUE(copy_result.has_value());
 
     // verify the copy task is created and assigned to the client who executed
@@ -3843,18 +5501,19 @@ TEST_F(MasterServiceTest, CreateCopyTaskTest) {
     EXPECT_EQ(contexts[0].client_id, task.value().assigned_client);
 
     // Copy with empty targets should fail
-    auto copy_result1 = service_->CreateCopyTask(key1, {});
+    auto copy_result1 = service_->CreateCopyTask(key1, "default", {});
     EXPECT_FALSE(copy_result1.has_value());
     EXPECT_EQ(ErrorCode::INVALID_PARAMS, copy_result1.error());
 
     // Copy not exist key should fail
     auto copy_result2 =
-        service_->CreateCopyTask("not_exist_key", {"segment_1"});
+        service_->CreateCopyTask("not_exist_key", "default", {"segment_1"});
     EXPECT_FALSE(copy_result2.has_value());
     EXPECT_EQ(ErrorCode::OBJECT_NOT_FOUND, copy_result2.error());
 
     // Copy to segment that not mounted should fail
-    auto copy_result3 = service_->CreateCopyTask(key1, {"not_mounted_segment"});
+    auto copy_result3 =
+        service_->CreateCopyTask(key1, "default", {"not_mounted_segment"});
     EXPECT_FALSE(copy_result3.has_value());
     EXPECT_EQ(ErrorCode::INVALID_PARAMS, copy_result3.error());
 }
@@ -3889,14 +5548,15 @@ TEST_F(MasterServiceTest, CreateMoveTaskTest) {
     config.replica_num = 1;
     config.preferred_segment = "segment_0";
     auto put_start_result =
-        service_->PutStart(client_id, key1, slice_length, config);
+        service_->PutStart(client_id, key1, "default", slice_length, config);
     EXPECT_TRUE(put_start_result.has_value());
     auto put_end_result =
-        service_->PutEnd(client_id, key1, ReplicaType::MEMORY);
+        service_->PutEnd(client_id, key1, "default", ReplicaType::MEMORY);
     EXPECT_TRUE(put_end_result.has_value());
 
     // Move key1 from "segment_0" to "segment_1"
-    auto move_result = service_->CreateMoveTask(key1, "segment_0", "segment_1");
+    auto move_result =
+        service_->CreateMoveTask(key1, "default", "segment_0", "segment_1");
     EXPECT_TRUE(move_result.has_value());
 
     // Verify the move task is created and assigned to the client owning the
@@ -3907,32 +5567,32 @@ TEST_F(MasterServiceTest, CreateMoveTaskTest) {
     EXPECT_EQ(contexts[0].client_id, task.value().assigned_client);
 
     // Move non-existent key should fail
-    auto move_result1 =
-        service_->CreateMoveTask("not_exist_key", "segment_0", "segment_1");
+    auto move_result1 = service_->CreateMoveTask("not_exist_key", "default",
+                                                 "segment_0", "segment_1");
     EXPECT_FALSE(move_result1.has_value());
     EXPECT_EQ(ErrorCode::OBJECT_NOT_FOUND, move_result1.error());
 
     // Move to segment that is same as source should fail
     auto move_result_same =
-        service_->CreateMoveTask(key1, "segment_1", "segment_1");
+        service_->CreateMoveTask(key1, "default", "segment_1", "segment_1");
     EXPECT_FALSE(move_result_same.has_value());
     EXPECT_EQ(ErrorCode::INVALID_PARAMS, move_result_same.error());
 
     // Move to segment that is not mounted should fail
-    auto move_result2 =
-        service_->CreateMoveTask(key1, "segment_0", "not_mounted_segment");
+    auto move_result2 = service_->CreateMoveTask(key1, "default", "segment_0",
+                                                 "not_mounted_segment");
     EXPECT_FALSE(move_result2.has_value());
     EXPECT_EQ(ErrorCode::INVALID_PARAMS, move_result2.error());
 
     // Move from segment that does not have the replica should fail
     auto move_result3 =
-        service_->CreateMoveTask(key1, "segment_2", "segment_1");
+        service_->CreateMoveTask(key1, "default", "segment_2", "segment_1");
     EXPECT_FALSE(move_result3.has_value());
     EXPECT_EQ(ErrorCode::INVALID_PARAMS, move_result3.error());
 
     // Move from segment that is not mounted should fail
-    auto move_result4 =
-        service_->CreateMoveTask(key1, "not_mounted_segment", "segment_1");
+    auto move_result4 = service_->CreateMoveTask(
+        key1, "default", "not_mounted_segment", "segment_1");
     EXPECT_FALSE(move_result4.has_value());
     EXPECT_EQ(ErrorCode::INVALID_PARAMS, move_result4.error());
 }
@@ -3967,14 +5627,15 @@ TEST_F(MasterServiceTest, QueryTaskTest) {
     config.replica_num = 1;
     config.preferred_segment = "segment_0";
     auto put_start_result =
-        service_->PutStart(client_id, key1, slice_length, config);
+        service_->PutStart(client_id, key1, "default", slice_length, config);
     EXPECT_TRUE(put_start_result.has_value());
     auto put_end_result =
-        service_->PutEnd(client_id, key1, ReplicaType::MEMORY);
+        service_->PutEnd(client_id, key1, "default", ReplicaType::MEMORY);
     EXPECT_TRUE(put_end_result.has_value());
 
     // Move key1 from "segment_0" to "segment_1"
-    auto move_result = service_->CreateMoveTask(key1, "segment_0", "segment_1");
+    auto move_result =
+        service_->CreateMoveTask(key1, "default", "segment_0", "segment_1");
     EXPECT_TRUE(move_result.has_value());
 
     // Query non-existent task should fail
@@ -4014,18 +5675,21 @@ TEST_F(MasterServiceTest, FetchTasksReturnsAssignedTasksOnlyAndDrainsQueue) {
     config.replica_num = 1;
     config.preferred_segment = "segment_0";
 
+    ASSERT_TRUE(service_
+                    ->PutStart(put_client_id, key, "default",
+                               /*slice_length=*/1024, config)
+                    .has_value());
     ASSERT_TRUE(
-        service_->PutStart(put_client_id, key, /*slice_length=*/1024, config)
+        service_->PutEnd(put_client_id, key, "default", ReplicaType::MEMORY)
             .has_value());
-    ASSERT_TRUE(
-        service_->PutEnd(put_client_id, key, ReplicaType::MEMORY).has_value());
 
     // Create two tasks; both should be assigned to the client owning source
     // segment_0.
-    auto copy_task_id = service_->CreateCopyTask(key, {"segment_1"});
+    auto copy_task_id = service_->CreateCopyTask(key, "default", {"segment_1"});
     ASSERT_TRUE(copy_task_id.has_value());
 
-    auto move_task_id = service_->CreateMoveTask(key, "segment_0", "segment_1");
+    auto move_task_id =
+        service_->CreateMoveTask(key, "default", "segment_0", "segment_1");
     ASSERT_TRUE(move_task_id.has_value());
 
     // Fetch from client_0 should get both tasks (order not guaranteed).
@@ -4059,6 +5723,82 @@ TEST_F(MasterServiceTest, FetchTasksReturnsAssignedTasksOnlyAndDrainsQueue) {
     EXPECT_TRUE(fetch0_again->empty());
 }
 
+TEST_F(MasterServiceTest, TenantTasksCarryTenantInPayload) {
+    const std::string tenant_id = "tenant_for_async_task";
+    auto service =
+        std::make_unique<MasterService>(MakeStrictTenantConfig({tenant_id}));
+    const auto ctx0 = PrepareSimpleSegment(*service, "segment_0", 0x300000000,
+                                           kDefaultSegmentSize);
+    [[maybe_unused]] const auto ctx1 = PrepareSimpleSegment(
+        *service, "segment_1", 0x400000000, kDefaultSegmentSize);
+
+    const UUID put_client_id = generate_uuid();
+    const std::string key = "tenant_task_key";
+
+    ReplicateConfig config;
+    config.replica_num = 1;
+    config.preferred_segment = "segment_0";
+
+    ASSERT_TRUE(service
+                    ->PutStart(put_client_id, key, tenant_id,
+                               /*slice_length=*/1024, config)
+                    .has_value());
+    ASSERT_TRUE(
+        service->PutEnd(put_client_id, key, tenant_id, ReplicaType::MEMORY)
+            .has_value());
+
+    auto copy_task_id = service->CreateCopyTask(key, tenant_id, {"segment_1"});
+    ASSERT_TRUE(copy_task_id.has_value());
+    auto move_task_id =
+        service->CreateMoveTask(key, tenant_id, "segment_0", "segment_1");
+    ASSERT_TRUE(move_task_id.has_value());
+
+    auto fetched = service->FetchTasks(ctx0.client_id, /*batch_size=*/16);
+    ASSERT_TRUE(fetched.has_value());
+    ASSERT_EQ(fetched->size(), 2u);
+
+    bool saw_copy = false;
+    bool saw_move = false;
+    for (const auto& assignment : *fetched) {
+        if (assignment.id == copy_task_id.value()) {
+            ReplicaCopyPayload payload;
+            struct_json::from_json(payload, assignment.payload);
+            EXPECT_EQ(payload.tenant_id, tenant_id);
+            EXPECT_EQ(payload.key, key);
+            saw_copy = true;
+        } else if (assignment.id == move_task_id.value()) {
+            ReplicaMovePayload payload;
+            struct_json::from_json(payload, assignment.payload);
+            EXPECT_EQ(payload.tenant_id, tenant_id);
+            EXPECT_EQ(payload.key, key);
+            saw_move = true;
+        }
+    }
+    EXPECT_TRUE(saw_copy);
+    EXPECT_TRUE(saw_move);
+}
+
+TEST_F(MasterServiceTest, LegacyTaskPayloadDefaultsTenant) {
+    ReplicaCopyPayload copy_payload;
+    struct_json::from_json(
+        copy_payload,
+        R"({"key":"legacy_copy_key","source":"segment_0","targets":["segment_1"]})");
+    EXPECT_EQ(copy_payload.tenant_id, "default");
+    EXPECT_EQ(copy_payload.key, "legacy_copy_key");
+    EXPECT_EQ(copy_payload.source, "segment_0");
+    ASSERT_EQ(copy_payload.targets.size(), 1u);
+    EXPECT_EQ(copy_payload.targets[0], "segment_1");
+
+    ReplicaMovePayload move_payload;
+    struct_json::from_json(
+        move_payload,
+        R"({"key":"legacy_move_key","source":"segment_0","target":"segment_1"})");
+    EXPECT_EQ(move_payload.tenant_id, "default");
+    EXPECT_EQ(move_payload.key, "legacy_move_key");
+    EXPECT_EQ(move_payload.source, "segment_0");
+    EXPECT_EQ(move_payload.target, "segment_1");
+}
+
 TEST_F(MasterServiceTest, FetchTasksRespectsBatchSize) {
     std::unique_ptr<MasterService> service_(new MasterService());
 
@@ -4074,15 +5814,18 @@ TEST_F(MasterServiceTest, FetchTasksRespectsBatchSize) {
     config.replica_num = 1;
     config.preferred_segment = "segment_0";
 
+    ASSERT_TRUE(service_
+                    ->PutStart(put_client_id, key, "default",
+                               /*slice_length=*/1024, config)
+                    .has_value());
     ASSERT_TRUE(
-        service_->PutStart(put_client_id, key, /*slice_length=*/1024, config)
+        service_->PutEnd(put_client_id, key, "default", ReplicaType::MEMORY)
             .has_value());
-    ASSERT_TRUE(
-        service_->PutEnd(put_client_id, key, ReplicaType::MEMORY).has_value());
 
-    auto t1 = service_->CreateCopyTask(key, {"segment_1"});
+    auto t1 = service_->CreateCopyTask(key, "default", {"segment_1"});
     ASSERT_TRUE(t1.has_value());
-    auto t2 = service_->CreateMoveTask(key, "segment_0", "segment_1");
+    auto t2 =
+        service_->CreateMoveTask(key, "default", "segment_0", "segment_1");
     ASSERT_TRUE(t2.has_value());
 
     auto fetch_first = service_->FetchTasks(ctx0.client_id, /*batch_size=*/1);
@@ -4123,14 +5866,16 @@ TEST_F(MasterServiceTest, UpdateTaskSuccessFlow) {
     config.replica_num = 1;
     config.preferred_segment = "segment_0";
 
+    ASSERT_TRUE(service_
+                    ->PutStart(put_client_id, key, "default",
+                               /*slice_length=*/1024, config)
+                    .has_value());
     ASSERT_TRUE(
-        service_->PutStart(put_client_id, key, /*slice_length=*/1024, config)
+        service_->PutEnd(put_client_id, key, "default", ReplicaType::MEMORY)
             .has_value());
-    ASSERT_TRUE(
-        service_->PutEnd(put_client_id, key, ReplicaType::MEMORY).has_value());
 
     // Create a task assigned to client owning segment_0.
-    auto task_id_res = service_->CreateCopyTask(key, {"segment_1"});
+    auto task_id_res = service_->CreateCopyTask(key, "default", {"segment_1"});
     ASSERT_TRUE(task_id_res.has_value());
     const UUID task_id = task_id_res.value();
 
@@ -4179,13 +5924,16 @@ TEST_F(MasterServiceTest, UpdateTaskRejectsWrongClient) {
     config.replica_num = 1;
     config.preferred_segment = "segment_0";
 
+    ASSERT_TRUE(service_
+                    ->PutStart(put_client_id, key, "default",
+                               /*slice_length=*/1024, config)
+                    .has_value());
     ASSERT_TRUE(
-        service_->PutStart(put_client_id, key, /*slice_length=*/1024, config)
+        service_->PutEnd(put_client_id, key, "default", ReplicaType::MEMORY)
             .has_value());
-    ASSERT_TRUE(
-        service_->PutEnd(put_client_id, key, ReplicaType::MEMORY).has_value());
 
-    auto task_id_res = service_->CreateMoveTask(key, "segment_0", "segment_1");
+    auto task_id_res =
+        service_->CreateMoveTask(key, "default", "segment_0", "segment_1");
     ASSERT_TRUE(task_id_res.has_value());
     const UUID task_id = task_id_res.value();
 
@@ -4249,7 +5997,7 @@ TEST_F(MasterServiceTest,
     config.preferred_segment = "segment_0";
 
     auto put_result = service_->PutStart(
-        ctx0.client_id, "drain_skip_allocation_key", 1024, config);
+        ctx0.client_id, "drain_skip_allocation_key", "default", 1024, config);
     ASSERT_TRUE(put_result.has_value());
     ASSERT_EQ(put_result->size(), 1u);
     EXPECT_EQ(put_result->front()
@@ -4258,7 +6006,7 @@ TEST_F(MasterServiceTest,
               "segment_1");
     ASSERT_TRUE(service_
                     ->PutEnd(ctx0.client_id, "drain_skip_allocation_key",
-                             ReplicaType::MEMORY)
+                             "default", ReplicaType::MEMORY)
                     .has_value());
 }
 
@@ -4302,7 +6050,7 @@ TEST_F(MasterServiceTest, DrainJobSchedulesMoveTaskAndConvergesToDrained) {
     ASSERT_TRUE(segment_status.has_value());
     EXPECT_EQ(segment_status.value(), SegmentStatus::DRAINED);
 
-    auto replicas = service_->GetReplicaList(key);
+    auto replicas = service_->GetReplicaList(key, "default");
     ASSERT_TRUE(replicas.has_value());
     std::unordered_set<std::string> segment_names;
     for (const auto& replica : replicas->replicas) {
@@ -4316,8 +6064,8 @@ TEST_F(MasterServiceTest, DrainJobSchedulesMoveTaskAndConvergesToDrained) {
 TEST_F(MasterServiceTest, CancelDrainJobRestoresSegmentStatus) {
     auto service_ = std::make_unique<MasterService>();
 
-    const auto ctx0 = PrepareSimpleSegment(*service_, "segment_0", 0x300000000,
-                                           kDefaultSegmentSize);
+    [[maybe_unused]] const auto ctx0 = PrepareSimpleSegment(
+        *service_, "segment_0", 0x300000000, kDefaultSegmentSize);
     [[maybe_unused]] const auto ctx1 = PrepareSimpleSegment(
         *service_, "segment_1", 0x400000000, kDefaultSegmentSize);
 
@@ -4434,27 +6182,28 @@ TEST_F(MasterServiceTest, ForceRemoveLeasedObject) {
     ReplicateConfig config;
     config.replica_num = 1;
     auto put_start_result =
-        service_->PutStart(client_id, key, slice_length, config);
+        service_->PutStart(client_id, key, "default", slice_length, config);
     ASSERT_TRUE(put_start_result.has_value());
-    auto put_end_result = service_->PutEnd(client_id, key, ReplicaType::MEMORY);
+    auto put_end_result =
+        service_->PutEnd(client_id, key, "default", ReplicaType::MEMORY);
     ASSERT_TRUE(put_end_result.has_value());
 
     // Verify object exists
-    auto exist_result = service_->ExistKey(key);
+    auto exist_result = service_->ExistKey(key, "default");
     ASSERT_TRUE(exist_result.has_value());
     ASSERT_TRUE(exist_result.value());
 
     // Normal remove should fail because object has active lease
-    auto remove_result_no_force = service_->Remove(key, false);
+    auto remove_result_no_force = service_->Remove(key, "default", false);
     EXPECT_FALSE(remove_result_no_force.has_value());
     EXPECT_EQ(ErrorCode::OBJECT_HAS_LEASE, remove_result_no_force.error());
 
     // Force remove should succeed even with active lease
-    auto remove_result_force = service_->Remove(key, true);
+    auto remove_result_force = service_->Remove(key, "default", true);
     EXPECT_TRUE(remove_result_force.has_value());
 
     // Verify object is removed
-    auto get_result = service_->GetReplicaList(key);
+    auto get_result = service_->GetReplicaList(key, "default");
     EXPECT_FALSE(get_result.has_value());
     EXPECT_EQ(ErrorCode::OBJECT_NOT_FOUND, get_result.error());
 }
@@ -4477,13 +6226,13 @@ TEST_F(MasterServiceTest, ForceRemoveByRegexLeasedObjects) {
         ReplicateConfig config;
         config.replica_num = 1;
         auto put_start_result =
-            service_->PutStart(client_id, key, slice_length, config);
+            service_->PutStart(client_id, key, "default", slice_length, config);
         ASSERT_TRUE(put_start_result.has_value());
         auto put_end_result =
-            service_->PutEnd(client_id, key, ReplicaType::MEMORY);
+            service_->PutEnd(client_id, key, "default", ReplicaType::MEMORY);
         ASSERT_TRUE(put_end_result.has_value());
         // Grant lease by reading the object
-        auto exist_result = service_->ExistKey(key);
+        auto exist_result = service_->ExistKey(key, "default");
         ASSERT_TRUE(exist_result.has_value());
         ASSERT_TRUE(exist_result.value());
     }
@@ -4491,28 +6240,28 @@ TEST_F(MasterServiceTest, ForceRemoveByRegexLeasedObjects) {
     // Normal RemoveByRegex should remove 0 because all objects have active
     // leases
     auto remove_result_no_force =
-        service_->RemoveByRegex("^force_regex_key_", false);
+        service_->RemoveByRegex("^force_regex_key_", "default", false);
     ASSERT_TRUE(remove_result_no_force.has_value());
     EXPECT_EQ(0, remove_result_no_force.value());
 
     // All objects should still exist
     for (int i = 0; i < 5; ++i) {
         std::string key = "force_regex_key_" + std::to_string(i);
-        auto exist_result = service_->ExistKey(key);
+        auto exist_result = service_->ExistKey(key, "default");
         ASSERT_TRUE(exist_result.has_value());
         ASSERT_TRUE(exist_result.value());
     }
 
     // Force RemoveByRegex should remove all 5 objects
     auto remove_result_force =
-        service_->RemoveByRegex("^force_regex_key_", true);
+        service_->RemoveByRegex("^force_regex_key_", "default", true);
     ASSERT_TRUE(remove_result_force.has_value());
     EXPECT_EQ(5, remove_result_force.value());
 
     // All objects should be removed
     for (int i = 0; i < 5; ++i) {
         std::string key = "force_regex_key_" + std::to_string(i);
-        auto exist_result = service_->ExistKey(key);
+        auto exist_result = service_->ExistKey(key, "default");
         ASSERT_TRUE(exist_result.has_value());
         ASSERT_FALSE(exist_result.value());
     }
@@ -4536,13 +6285,13 @@ TEST_F(MasterServiceTest, ForceRemoveAllLeasedObjects) {
         ReplicateConfig config;
         config.replica_num = 1;
         auto put_start_result =
-            service_->PutStart(client_id, key, slice_length, config);
+            service_->PutStart(client_id, key, "default", slice_length, config);
         ASSERT_TRUE(put_start_result.has_value());
         auto put_end_result =
-            service_->PutEnd(client_id, key, ReplicaType::MEMORY);
+            service_->PutEnd(client_id, key, "default", ReplicaType::MEMORY);
         ASSERT_TRUE(put_end_result.has_value());
         // Grant lease by reading the object
-        auto exist_result = service_->ExistKey(key);
+        auto exist_result = service_->ExistKey(key, "default");
         ASSERT_TRUE(exist_result.has_value());
         ASSERT_TRUE(exist_result.value());
     }
@@ -4553,7 +6302,7 @@ TEST_F(MasterServiceTest, ForceRemoveAllLeasedObjects) {
     // All objects should still exist
     for (int i = 0; i < 10; ++i) {
         std::string key = "force_all_key_" + std::to_string(i);
-        auto exist_result = service_->ExistKey(key);
+        auto exist_result = service_->ExistKey(key, "default");
         ASSERT_TRUE(exist_result.has_value());
         ASSERT_TRUE(exist_result.value());
     }
@@ -4564,7 +6313,7 @@ TEST_F(MasterServiceTest, ForceRemoveAllLeasedObjects) {
     // All objects should be removed
     for (int i = 0; i < 10; ++i) {
         std::string key = "force_all_key_" + std::to_string(i);
-        auto exist_result = service_->ExistKey(key);
+        auto exist_result = service_->ExistKey(key, "default");
         ASSERT_TRUE(exist_result.has_value());
         ASSERT_FALSE(exist_result.value());
     }
@@ -4583,23 +6332,24 @@ TEST_F(MasterServiceTest, UpsertNewKey) {
     config.replica_num = 1;
 
     auto upsert_result =
-        service_->UpsertStart(client_id, key, slice_length, config);
+        service_->UpsertStart(client_id, key, "default", slice_length, config);
     ASSERT_TRUE(upsert_result.has_value());
     auto replicas = upsert_result.value();
     EXPECT_EQ(1, replicas.size());
     EXPECT_EQ(ReplicaStatus::PROCESSING, replicas[0].status);
 
     // During upsert, GetReplicaList should return not ready
-    auto get_result = service_->GetReplicaList(key);
+    auto get_result = service_->GetReplicaList(key, "default");
     EXPECT_FALSE(get_result.has_value());
     EXPECT_EQ(ErrorCode::REPLICA_IS_NOT_READY, get_result.error());
 
     // UpsertEnd completes the operation
-    auto end_result = service_->UpsertEnd(client_id, key, ReplicaType::MEMORY);
+    auto end_result =
+        service_->UpsertEnd(client_id, key, "default", ReplicaType::MEMORY);
     ASSERT_TRUE(end_result.has_value());
 
     // Verify replica is COMPLETE
-    auto final_result = service_->GetReplicaList(key);
+    auto final_result = service_->GetReplicaList(key, "default");
     ASSERT_TRUE(final_result.has_value());
     EXPECT_EQ(1, final_result.value().replicas.size());
     EXPECT_EQ(ReplicaStatus::COMPLETE, final_result.value().replicas[0].status);
@@ -4617,16 +6367,18 @@ TEST_F(MasterServiceTest, UpsertSameSize) {
     config.replica_num = 1;
 
     // First: PutStart + PutEnd to create the object
-    auto put_result = service_->PutStart(client_id, key, slice_length, config);
+    auto put_result =
+        service_->PutStart(client_id, key, "default", slice_length, config);
     ASSERT_TRUE(put_result.has_value());
     auto original_replicas = put_result.value();
-    auto put_end = service_->PutEnd(client_id, key, ReplicaType::MEMORY);
+    auto put_end =
+        service_->PutEnd(client_id, key, "default", ReplicaType::MEMORY);
     ASSERT_TRUE(put_end.has_value());
 
     // UpsertStart with same size — should reuse buffers
     const UUID new_client_id = generate_uuid();
-    auto upsert_result =
-        service_->UpsertStart(new_client_id, key, slice_length, config);
+    auto upsert_result = service_->UpsertStart(new_client_id, key, "default",
+                                               slice_length, config);
     ASSERT_TRUE(upsert_result.has_value());
     auto upsert_replicas = upsert_result.value();
     EXPECT_EQ(1, upsert_replicas.size());
@@ -4642,11 +6394,11 @@ TEST_F(MasterServiceTest, UpsertSameSize) {
 
     // UpsertEnd with the new client_id
     auto end_result =
-        service_->UpsertEnd(new_client_id, key, ReplicaType::MEMORY);
+        service_->UpsertEnd(new_client_id, key, "default", ReplicaType::MEMORY);
     ASSERT_TRUE(end_result.has_value());
 
     // Verify replica is COMPLETE again
-    auto final_result = service_->GetReplicaList(key);
+    auto final_result = service_->GetReplicaList(key, "default");
     ASSERT_TRUE(final_result.has_value());
     EXPECT_EQ(ReplicaStatus::COMPLETE, final_result.value().replicas[0].status);
 }
@@ -4665,23 +6417,26 @@ TEST_F(MasterServiceTest, UpsertSameSizeRefreshesMetadata) {
 
     // Create object with client_a
     auto put_result =
-        service_->PutStart(client_id_a, key, slice_length, config);
+        service_->PutStart(client_id_a, key, "default", slice_length, config);
     ASSERT_TRUE(put_result.has_value());
-    auto put_end = service_->PutEnd(client_id_a, key, ReplicaType::MEMORY);
+    auto put_end =
+        service_->PutEnd(client_id_a, key, "default", ReplicaType::MEMORY);
     ASSERT_TRUE(put_end.has_value());
 
     // UpsertStart with client_b
-    auto upsert_result =
-        service_->UpsertStart(client_id_b, key, slice_length, config);
+    auto upsert_result = service_->UpsertStart(client_id_b, key, "default",
+                                               slice_length, config);
     ASSERT_TRUE(upsert_result.has_value());
 
     // UpsertEnd with client_a should fail (client_id was refreshed to client_b)
-    auto end_fail = service_->UpsertEnd(client_id_a, key, ReplicaType::MEMORY);
+    auto end_fail =
+        service_->UpsertEnd(client_id_a, key, "default", ReplicaType::MEMORY);
     EXPECT_FALSE(end_fail.has_value());
     EXPECT_EQ(ErrorCode::ILLEGAL_CLIENT, end_fail.error());
 
     // UpsertEnd with client_b should succeed
-    auto end_ok = service_->UpsertEnd(client_id_b, key, ReplicaType::MEMORY);
+    auto end_ok =
+        service_->UpsertEnd(client_id_b, key, "default", ReplicaType::MEMORY);
     ASSERT_TRUE(end_ok.has_value());
 }
 
@@ -4698,15 +6453,17 @@ TEST_F(MasterServiceTest, UpsertDifferentSize) {
     config.replica_num = 1;
 
     // Create object with original_size
-    auto put_result = service_->PutStart(client_id, key, original_size, config);
+    auto put_result =
+        service_->PutStart(client_id, key, "default", original_size, config);
     ASSERT_TRUE(put_result.has_value());
     auto original_replicas = put_result.value();
-    auto put_end = service_->PutEnd(client_id, key, ReplicaType::MEMORY);
+    auto put_end =
+        service_->PutEnd(client_id, key, "default", ReplicaType::MEMORY);
     ASSERT_TRUE(put_end.has_value());
 
     // UpsertStart with different size
     auto upsert_result =
-        service_->UpsertStart(client_id, key, new_size, config);
+        service_->UpsertStart(client_id, key, "default", new_size, config);
     ASSERT_TRUE(upsert_result.has_value());
     auto new_replicas = upsert_result.value();
     EXPECT_EQ(1, new_replicas.size());
@@ -4721,11 +6478,12 @@ TEST_F(MasterServiceTest, UpsertDifferentSize) {
                   .buffer_descriptor.buffer_address_);
 
     // UpsertEnd
-    auto end_result = service_->UpsertEnd(client_id, key, ReplicaType::MEMORY);
+    auto end_result =
+        service_->UpsertEnd(client_id, key, "default", ReplicaType::MEMORY);
     ASSERT_TRUE(end_result.has_value());
 
     // Verify the object is complete
-    auto final_result = service_->GetReplicaList(key);
+    auto final_result = service_->GetReplicaList(key, "default");
     ASSERT_TRUE(final_result.has_value());
     EXPECT_EQ(ReplicaStatus::COMPLETE, final_result.value().replicas[0].status);
 }
@@ -4751,19 +6509,21 @@ TEST_F(MasterServiceTest, UpsertConflictReplicationTask) {
     config.preferred_segment = "segment_1";
 
     // Create object
-    auto put_result = service_->PutStart(client_id, key, slice_length, config);
+    auto put_result =
+        service_->PutStart(client_id, key, "default", slice_length, config);
     ASSERT_TRUE(put_result.has_value());
-    auto put_end = service_->PutEnd(client_id, key, ReplicaType::MEMORY);
+    auto put_end =
+        service_->PutEnd(client_id, key, "default", ReplicaType::MEMORY);
     ASSERT_TRUE(put_end.has_value());
 
     // Start a Copy
-    auto copy_result =
-        service_->CopyStart(client_id, key, "segment_1", {"segment_2"});
+    auto copy_result = service_->CopyStart(client_id, key, "default",
+                                           "segment_1", {"segment_2"});
     ASSERT_TRUE(copy_result.has_value());
 
     // UpsertStart should fail with OBJECT_HAS_REPLICATION_TASK
     auto upsert_result =
-        service_->UpsertStart(client_id, key, slice_length, config);
+        service_->UpsertStart(client_id, key, "default", slice_length, config);
     EXPECT_FALSE(upsert_result.has_value());
     EXPECT_EQ(ErrorCode::OBJECT_HAS_REPLICATION_TASK, upsert_result.error());
 }
@@ -4782,27 +6542,30 @@ TEST_F(MasterServiceTest, UpsertPreemptsInProgressPut) {
     config.replica_num = 1;
 
     // Client A starts a Put but doesn't finish
-    auto put_result = service_->PutStart(client_a, key, slice_length, config);
+    auto put_result =
+        service_->PutStart(client_a, key, "default", slice_length, config);
     ASSERT_TRUE(put_result.has_value());
 
     // Client B upserts the same key — should preempt client A
     auto upsert_result =
-        service_->UpsertStart(client_b, key, slice_length, config);
+        service_->UpsertStart(client_b, key, "default", slice_length, config);
     ASSERT_TRUE(upsert_result.has_value());
     auto upsert_replicas = upsert_result.value();
     EXPECT_EQ(1, upsert_replicas.size());
     EXPECT_EQ(ReplicaStatus::PROCESSING, upsert_replicas[0].status);
 
     // Client A's PutEnd should fail
-    auto put_end_a = service_->PutEnd(client_a, key, ReplicaType::MEMORY);
+    auto put_end_a =
+        service_->PutEnd(client_a, key, "default", ReplicaType::MEMORY);
     EXPECT_FALSE(put_end_a.has_value());
 
     // Client B's UpsertEnd should succeed
-    auto upsert_end = service_->UpsertEnd(client_b, key, ReplicaType::MEMORY);
+    auto upsert_end =
+        service_->UpsertEnd(client_b, key, "default", ReplicaType::MEMORY);
     ASSERT_TRUE(upsert_end.has_value());
 
     // Verify final state
-    auto final_result = service_->GetReplicaList(key);
+    auto final_result = service_->GetReplicaList(key, "default");
     ASSERT_TRUE(final_result.has_value());
     EXPECT_EQ(ReplicaStatus::COMPLETE, final_result.value().replicas[0].status);
 }
@@ -4820,16 +6583,16 @@ TEST_F(MasterServiceTest, UpsertRevoke) {
 
     // UpsertStart (Case A — new key)
     auto upsert_result =
-        service_->UpsertStart(client_id, key, slice_length, config);
+        service_->UpsertStart(client_id, key, "default", slice_length, config);
     ASSERT_TRUE(upsert_result.has_value());
 
     // UpsertRevoke
     auto revoke_result =
-        service_->UpsertRevoke(client_id, key, ReplicaType::MEMORY);
+        service_->UpsertRevoke(client_id, key, "default", ReplicaType::MEMORY);
     ASSERT_TRUE(revoke_result.has_value());
 
     // Key should be gone
-    auto exist_result = service_->ExistKey(key);
+    auto exist_result = service_->ExistKey(key, "default");
     ASSERT_TRUE(exist_result.has_value());
     EXPECT_FALSE(exist_result.value());
 }
@@ -4846,24 +6609,26 @@ TEST_F(MasterServiceTest, UpsertInPlaceThenRevoke) {
     config.replica_num = 1;
 
     // Create object first
-    auto put_result = service_->PutStart(client_id, key, slice_length, config);
+    auto put_result =
+        service_->PutStart(client_id, key, "default", slice_length, config);
     ASSERT_TRUE(put_result.has_value());
-    auto put_end = service_->PutEnd(client_id, key, ReplicaType::MEMORY);
+    auto put_end =
+        service_->PutEnd(client_id, key, "default", ReplicaType::MEMORY);
     ASSERT_TRUE(put_end.has_value());
 
     // UpsertStart in-place (same size)
     const UUID new_client = generate_uuid();
     auto upsert_result =
-        service_->UpsertStart(new_client, key, slice_length, config);
+        service_->UpsertStart(new_client, key, "default", slice_length, config);
     ASSERT_TRUE(upsert_result.has_value());
 
     // UpsertRevoke — replicas are PROCESSING, should be erased
     auto revoke_result =
-        service_->UpsertRevoke(new_client, key, ReplicaType::MEMORY);
+        service_->UpsertRevoke(new_client, key, "default", ReplicaType::MEMORY);
     ASSERT_TRUE(revoke_result.has_value());
 
     // Key should be gone (no valid replicas left)
-    auto exist_result = service_->ExistKey(key);
+    auto exist_result = service_->ExistKey(key, "default");
     ASSERT_TRUE(exist_result.has_value());
     EXPECT_FALSE(exist_result.value());
 }
@@ -4878,23 +6643,25 @@ TEST_F(MasterServiceTest, BatchUpsertStart) {
     config.replica_num = 1;
 
     // Create key_1 with size 1024
-    auto put_result = service_->PutStart(client_id, "key_1", 1024, config);
+    auto put_result =
+        service_->PutStart(client_id, "key_1", "default", 1024, config);
     ASSERT_TRUE(put_result.has_value());
-    auto put_end = service_->PutEnd(client_id, "key_1", ReplicaType::MEMORY);
+    auto put_end =
+        service_->PutEnd(client_id, "key_1", "default", ReplicaType::MEMORY);
     ASSERT_TRUE(put_end.has_value());
 
     // BatchUpsertStart: key_1 (same size), key_2 (new)
     std::vector<std::string> keys = {"key_1", "key_2"};
     std::vector<uint64_t> slice_lengths = {1024, 2048};
 
-    auto results =
-        service_->BatchUpsertStart(client_id, keys, slice_lengths, config);
+    auto results = service_->BatchUpsertStart(client_id, keys, "default",
+                                              slice_lengths, config);
     ASSERT_EQ(2, results.size());
     EXPECT_TRUE(results[0].has_value());  // key_1: Case B (in-place)
     EXPECT_TRUE(results[1].has_value());  // key_2: Case A (new)
 
     // Complete both
-    auto end_results = service_->BatchUpsertEnd(client_id, keys);
+    auto end_results = service_->BatchUpsertEnd(client_id, keys, "default");
     ASSERT_EQ(2, end_results.size());
     EXPECT_TRUE(end_results[0].has_value());
     EXPECT_TRUE(end_results[1].has_value());
@@ -4916,36 +6683,42 @@ TEST_F(MasterServiceTest, UpsertPreemptsInProgressUpsert) {
     config.replica_num = 1;
 
     // Step 1: Create the object via Put
-    auto put_result = service_->PutStart(client_a, key, slice_length, config);
+    auto put_result =
+        service_->PutStart(client_a, key, "default", slice_length, config);
     ASSERT_TRUE(put_result.has_value());
-    auto put_end = service_->PutEnd(client_a, key, ReplicaType::MEMORY);
+    auto put_end =
+        service_->PutEnd(client_a, key, "default", ReplicaType::MEMORY);
     ASSERT_TRUE(put_end.has_value());
 
     // Step 2: Client B starts in-place upsert (Case B) — marks COMPLETE →
     // PROCESSING
-    auto upsert_b = service_->UpsertStart(client_b, key, slice_length, config);
+    auto upsert_b =
+        service_->UpsertStart(client_b, key, "default", slice_length, config);
     ASSERT_TRUE(upsert_b.has_value());
 
     // Key should be unreadable now (all replicas are PROCESSING)
-    auto get_mid = service_->GetReplicaList(key);
+    auto get_mid = service_->GetReplicaList(key, "default");
     EXPECT_FALSE(get_mid.has_value());
     EXPECT_EQ(ErrorCode::REPLICA_IS_NOT_READY, get_mid.error());
 
     // Step 3: Client C upserts the same key — preempts Client B
-    auto upsert_c = service_->UpsertStart(client_c, key, slice_length, config);
+    auto upsert_c =
+        service_->UpsertStart(client_c, key, "default", slice_length, config);
     ASSERT_TRUE(upsert_c.has_value());
     EXPECT_EQ(1, upsert_c.value().size());
 
     // Step 4: Client B's UpsertEnd should fail (preempted)
-    auto end_b = service_->UpsertEnd(client_b, key, ReplicaType::MEMORY);
+    auto end_b =
+        service_->UpsertEnd(client_b, key, "default", ReplicaType::MEMORY);
     EXPECT_FALSE(end_b.has_value());
 
     // Step 5: Client C's UpsertEnd should succeed
-    auto end_c = service_->UpsertEnd(client_c, key, ReplicaType::MEMORY);
+    auto end_c =
+        service_->UpsertEnd(client_c, key, "default", ReplicaType::MEMORY);
     ASSERT_TRUE(end_c.has_value());
 
     // Final verification
-    auto final_result = service_->GetReplicaList(key);
+    auto final_result = service_->GetReplicaList(key, "default");
     ASSERT_TRUE(final_result.has_value());
     EXPECT_EQ(1, final_result.value().replicas.size());
     EXPECT_EQ(ReplicaStatus::COMPLETE, final_result.value().replicas[0].status);
@@ -4966,29 +6739,31 @@ TEST_F(MasterServiceTest, UpsertDifferentSizeThenRevoke) {
     config.replica_num = 1;
 
     // Create object with original size
-    auto put_result = service_->PutStart(client_id, key, original_size, config);
+    auto put_result =
+        service_->PutStart(client_id, key, "default", original_size, config);
     ASSERT_TRUE(put_result.has_value());
-    auto put_end = service_->PutEnd(client_id, key, ReplicaType::MEMORY);
+    auto put_end =
+        service_->PutEnd(client_id, key, "default", ReplicaType::MEMORY);
     ASSERT_TRUE(put_end.has_value());
 
     // Verify the key exists
-    auto exist_before = service_->ExistKey(key);
+    auto exist_before = service_->ExistKey(key, "default");
     ASSERT_TRUE(exist_before.has_value());
     EXPECT_TRUE(exist_before.value());
 
     // UpsertStart with different size (Case C) — old replicas discarded,
     // new replicas allocated
     auto upsert_result =
-        service_->UpsertStart(client_id, key, new_size, config);
+        service_->UpsertStart(client_id, key, "default", new_size, config);
     ASSERT_TRUE(upsert_result.has_value());
 
     // Revoke — erase the newly allocated PROCESSING replicas
     auto revoke_result =
-        service_->UpsertRevoke(client_id, key, ReplicaType::MEMORY);
+        service_->UpsertRevoke(client_id, key, "default", ReplicaType::MEMORY);
     ASSERT_TRUE(revoke_result.has_value());
 
     // Key should be gone (old replicas in discarded, new replicas erased)
-    auto exist_after = service_->ExistKey(key);
+    auto exist_after = service_->ExistKey(key, "default");
     ASSERT_TRUE(exist_after.has_value());
     EXPECT_FALSE(exist_after.value());
 }
@@ -5016,12 +6791,13 @@ TEST_F(MasterServiceTest, HardPinObjectNotEvicted) {
         ReplicateConfig config;
         config.replica_num = 1;
         config.with_hard_pin = true;
-        auto result =
-            service_->PutStart(client_id, "pinned_model", value_size, config);
+        auto result = service_->PutStart(client_id, "pinned_model", "default",
+                                         value_size, config);
         ASSERT_TRUE(result.has_value());
-        ASSERT_TRUE(
-            service_->PutEnd(client_id, "pinned_model", ReplicaType::MEMORY)
-                .has_value());
+        ASSERT_TRUE(service_
+                        ->PutEnd(client_id, "pinned_model", "default",
+                                 ReplicaType::MEMORY)
+                        .has_value());
     }
 
     // Fill remaining space with normal objects to trigger eviction
@@ -5029,9 +6805,10 @@ TEST_F(MasterServiceTest, HardPinObjectNotEvicted) {
         std::string key = "filler_" + std::to_string(i);
         ReplicateConfig config;
         config.replica_num = 1;
-        auto result = service_->PutStart(client_id, key, value_size, config);
+        auto result =
+            service_->PutStart(client_id, key, "default", value_size, config);
         if (result.has_value()) {
-            service_->PutEnd(client_id, key, ReplicaType::MEMORY);
+            service_->PutEnd(client_id, key, "default", ReplicaType::MEMORY);
         }
     }
 
@@ -5039,14 +6816,15 @@ TEST_F(MasterServiceTest, HardPinObjectNotEvicted) {
     std::this_thread::sleep_for(std::chrono::milliseconds(kv_lease_ttl + 500));
 
     // Hard-pinned object must still be there
-    auto get_result = service_->GetReplicaList("pinned_model");
+    auto get_result = service_->GetReplicaList("pinned_model", "default");
     ASSERT_TRUE(get_result.has_value())
         << "Hard-pinned object was evicted, but it should never be";
 
     // Explicit Remove should still work on hard-pinned objects
-    auto remove_result = service_->Remove("pinned_model", /*force=*/true);
+    auto remove_result =
+        service_->Remove("pinned_model", "default", /*force=*/true);
     ASSERT_TRUE(remove_result.has_value());
-    auto exist_result = service_->ExistKey("pinned_model");
+    auto exist_result = service_->ExistKey("pinned_model", "default");
     ASSERT_TRUE(exist_result.has_value());
     ASSERT_FALSE(exist_result.value());
 
@@ -5080,12 +6858,14 @@ TEST_F(MasterServiceTest, HardPinWithSoftPinEvictionOrder) {
         ReplicateConfig config;
         config.replica_num = 1;
         config.with_hard_pin = true;
-        ASSERT_TRUE(
-            service_->PutStart(client_id, "hard_pinned", value_size, config)
-                .has_value());
-        ASSERT_TRUE(
-            service_->PutEnd(client_id, "hard_pinned", ReplicaType::MEMORY)
-                .has_value());
+        ASSERT_TRUE(service_
+                        ->PutStart(client_id, "hard_pinned", "default",
+                                   value_size, config)
+                        .has_value());
+        ASSERT_TRUE(service_
+                        ->PutEnd(client_id, "hard_pinned", "default",
+                                 ReplicaType::MEMORY)
+                        .has_value());
     }
 
     // Put a soft-pinned object
@@ -5093,12 +6873,14 @@ TEST_F(MasterServiceTest, HardPinWithSoftPinEvictionOrder) {
         ReplicateConfig config;
         config.replica_num = 1;
         config.with_soft_pin = true;
-        ASSERT_TRUE(
-            service_->PutStart(client_id, "soft_pinned", value_size, config)
-                .has_value());
-        ASSERT_TRUE(
-            service_->PutEnd(client_id, "soft_pinned", ReplicaType::MEMORY)
-                .has_value());
+        ASSERT_TRUE(service_
+                        ->PutStart(client_id, "soft_pinned", "default",
+                                   value_size, config)
+                        .has_value());
+        ASSERT_TRUE(service_
+                        ->PutEnd(client_id, "soft_pinned", "default",
+                                 ReplicaType::MEMORY)
+                        .has_value());
     }
 
     // Fill the rest
@@ -5106,9 +6888,10 @@ TEST_F(MasterServiceTest, HardPinWithSoftPinEvictionOrder) {
         std::string key = "normal_" + std::to_string(i);
         ReplicateConfig config;
         config.replica_num = 1;
-        auto result = service_->PutStart(client_id, key, value_size, config);
+        auto result =
+            service_->PutStart(client_id, key, "default", value_size, config);
         if (result.has_value()) {
-            service_->PutEnd(client_id, key, ReplicaType::MEMORY);
+            service_->PutEnd(client_id, key, "default", ReplicaType::MEMORY);
         }
     }
 
@@ -5116,13 +6899,94 @@ TEST_F(MasterServiceTest, HardPinWithSoftPinEvictionOrder) {
     std::this_thread::sleep_for(std::chrono::milliseconds(kv_lease_ttl + 500));
 
     // Hard-pinned always survives
-    ASSERT_TRUE(service_->GetReplicaList("hard_pinned").has_value())
+    ASSERT_TRUE(service_->GetReplicaList("hard_pinned", "default").has_value())
         << "Hard-pinned object was evicted";
 
     std::this_thread::sleep_for(std::chrono::milliseconds(kv_lease_ttl));
     service_->RemoveAll();
 }
 
+TEST_F(MasterServiceTest,
+       ObjectTypeEvictionScorePolicyRejectsInvalidBuilderInput) {
+    auto builder = MasterServiceConfig::builder();
+    EXPECT_THROW(builder.set_object_type_eviction_score_policy(
+                     ObjectDataType::KVCACHE,
+                     ObjectTypeEvictionScorePolicy{1.0, -1.0, 0}),
+                 std::invalid_argument);
+
+    EXPECT_THROW(builder.set_object_type_eviction_score_policy(
+                     ObjectDataType::KVCACHE,
+                     ObjectTypeEvictionScorePolicy{
+                         std::numeric_limits<double>::quiet_NaN(), 1.0, 0}),
+                 std::invalid_argument);
+
+    EXPECT_THROW(builder.set_object_type_eviction_score_policy(
+                     ObjectDataType::KVCACHE,
+                     ObjectTypeEvictionScorePolicy{
+                         1.0, std::numeric_limits<double>::infinity(), 0}),
+                 std::invalid_argument);
+}
+
+TEST_F(MasterServiceTest,
+       ObjectTypeEvictionScorePolicyRejectsInvalidDirectConfigInput) {
+    auto config = MasterServiceConfig::builder().build();
+    config.object_type_eviction_score_policies[ObjectDataType::KVCACHE] =
+        ObjectTypeEvictionScorePolicy{0.0, 1.0, 0};
+    EXPECT_THROW({ MasterService service(config); }, std::invalid_argument);
+
+    config.object_type_eviction_score_policies[ObjectDataType::KVCACHE] =
+        ObjectTypeEvictionScorePolicy{1.0, -1.0, 0};
+    EXPECT_THROW({ MasterService service(config); }, std::invalid_argument);
+
+    config.object_type_eviction_score_policies[ObjectDataType::KVCACHE] =
+        ObjectTypeEvictionScorePolicy{std::numeric_limits<double>::infinity(),
+                                      1.0, 0};
+    EXPECT_THROW({ MasterService service(config); }, std::invalid_argument);
+}
+
+TEST_F(MasterServiceTest, ObjectTypeEvictionPolicyRejectsInvalidBuilderInput) {
+    auto builder = MasterServiceConfig::builder();
+    EXPECT_THROW(builder.set_object_type_eviction_policy(
+                     ObjectDataType::KVCACHE, ObjectTypeEvictionPolicy{-1.0}),
+                 std::invalid_argument);
+
+    EXPECT_THROW(
+        builder.set_object_type_eviction_policy(
+            ObjectDataType::KVCACHE,
+            ObjectTypeEvictionPolicy{std::numeric_limits<double>::quiet_NaN()}),
+        std::invalid_argument);
+
+    EXPECT_THROW(
+        builder.set_object_type_eviction_policy(
+            ObjectDataType::KVCACHE,
+            ObjectTypeEvictionPolicy{std::numeric_limits<double>::infinity()}),
+        std::invalid_argument);
+
+    EXPECT_THROW(builder.set_object_type_eviction_policy(
+                     ObjectDataType::KVCACHE, ObjectTypeEvictionPolicy{1.1}),
+                 std::invalid_argument);
+}
+
+TEST_F(MasterServiceTest,
+       ObjectTypeEvictionPolicyRejectsInvalidDirectConfigInput) {
+    auto config = MasterServiceConfig::builder().build();
+    config.object_type_eviction_policies[ObjectDataType::KVCACHE] =
+        ObjectTypeEvictionPolicy{-1.0};
+    EXPECT_THROW({ MasterService service(config); }, std::invalid_argument);
+
+    config.object_type_eviction_policies[ObjectDataType::KVCACHE] =
+        ObjectTypeEvictionPolicy{std::numeric_limits<double>::quiet_NaN()};
+    EXPECT_THROW({ MasterService service(config); }, std::invalid_argument);
+
+    config.object_type_eviction_policies[ObjectDataType::KVCACHE] =
+        ObjectTypeEvictionPolicy{std::numeric_limits<double>::infinity()};
+    EXPECT_THROW({ MasterService service(config); }, std::invalid_argument);
+
+    config.object_type_eviction_policies[ObjectDataType::KVCACHE] =
+        ObjectTypeEvictionPolicy{1.1};
+    EXPECT_THROW({ MasterService service(config); }, std::invalid_argument);
+}
+
 TEST_F(MasterServiceTest, HardPinDefaultIsFalse) {
     // Objects created without with_hard_pin should not be hard-pinned
     auto service_config =
@@ -5139,26 +7003,263 @@ TEST_F(MasterServiceTest, HardPinDefaultIsFalse) {
     ReplicateConfig config;
     config.replica_num = 1;
     ASSERT_TRUE(
-        service_->PutStart(client_id, "normal_key", 1024, config).has_value());
-    ASSERT_TRUE(service_->PutEnd(client_id, "normal_key", ReplicaType::MEMORY)
-                    .has_value());
+        service_->PutStart(client_id, "normal_key", "default", 1024, config)
+            .has_value());
+    ASSERT_TRUE(
+        service_
+            ->PutEnd(client_id, "normal_key", "default", ReplicaType::MEMORY)
+            .has_value());
 
     // Put with hard_pin
     ReplicateConfig hp_config;
     hp_config.replica_num = 1;
     hp_config.with_hard_pin = true;
     ASSERT_TRUE(
-        service_->PutStart(client_id, "hp_key", 1024, hp_config).has_value());
+        service_->PutStart(client_id, "hp_key", "default", 1024, hp_config)
+            .has_value());
     ASSERT_TRUE(
-        service_->PutEnd(client_id, "hp_key", ReplicaType::MEMORY).has_value());
+        service_->PutEnd(client_id, "hp_key", "default", ReplicaType::MEMORY)
+            .has_value());
 
     // Both should exist
-    ASSERT_TRUE(service_->GetReplicaList("normal_key").has_value());
-    ASSERT_TRUE(service_->GetReplicaList("hp_key").has_value());
+    ASSERT_TRUE(service_->GetReplicaList("normal_key", "default").has_value());
+    ASSERT_TRUE(service_->GetReplicaList("hp_key", "default").has_value());
 
     service_->RemoveAll();
 }
 
+// ===================== Graceful Unmount Tests =====================
+
+TEST_F(MasterServiceTest, GracefulUnmountSegment_SetsCorrectStatus) {
+    std::unique_ptr<MasterService> service_(new MasterService());
+    auto segment = MakeSegment("graceful_test_segment");
+    UUID client_id = generate_uuid();
+
+    // Mount segment
+    ASSERT_TRUE(service_->MountSegment(segment, client_id).has_value());
+
+    // Verify initial status
+    auto status_before = service_->QuerySegmentStatus(segment.name);
+    ASSERT_TRUE(status_before.has_value());
+    EXPECT_EQ(status_before.value(), SegmentStatus::OK);
+
+    // Graceful unmount with 1 second grace period
+    auto graceful_result = service_->GracefulUnmountSegment(
+        segment.id, client_id, /*grace_period_ms=*/1000);
+    ASSERT_TRUE(graceful_result.has_value())
+        << "Graceful unmount should succeed: "
+        << toString(graceful_result.error());
+
+    // Verify status is GRACEFULLY_UNMOUNTING
+    auto status_after = service_->QuerySegmentStatus(segment.name);
+    ASSERT_TRUE(status_after.has_value());
+    EXPECT_EQ(status_after.value(), SegmentStatus::GRACEFULLY_UNMOUNTING);
+
+    // Wait for timer to expire and clean up
+    std::this_thread::sleep_for(std::chrono::milliseconds(1500));
+}
+
+TEST_F(MasterServiceTest, GracefulUnmountSegment_RejectWrongClient) {
+    std::unique_ptr<MasterService> service_(new MasterService());
+    auto segment = MakeSegment("graceful_owner_segment");
+    UUID owner_client = generate_uuid();
+    UUID wrong_client = generate_uuid();
+
+    ASSERT_TRUE(service_->MountSegment(segment, owner_client).has_value());
+
+    // Wrong client trying to graceful unmount should fail
+    auto graceful_result = service_->GracefulUnmountSegment(
+        segment.id, wrong_client, /*grace_period_ms=*/1000);
+    ASSERT_FALSE(graceful_result.has_value());
+    EXPECT_EQ(graceful_result.error(), ErrorCode::SEGMENT_NOT_FOUND);
+
+    // Owner should still be able to unmount
+    auto owner_result = service_->GracefulUnmountSegment(
+        segment.id, owner_client, /*grace_period_ms=*/1000);
+    EXPECT_TRUE(owner_result.has_value());
+
+    std::this_thread::sleep_for(std::chrono::milliseconds(1500));
+}
+
+TEST_F(MasterServiceTest, GracefulUnmountSegment_Idempotent) {
+    std::unique_ptr<MasterService> service_(new MasterService());
+    auto segment = MakeSegment("graceful_idempotent_segment");
+    UUID client_id = generate_uuid();
+
+    ASSERT_TRUE(service_->MountSegment(segment, client_id).has_value());
+
+    // First graceful unmount should succeed
+    auto result1 = service_->GracefulUnmountSegment(segment.id, client_id,
+                                                    /*grace_period_ms=*/1000);
+    ASSERT_TRUE(result1.has_value());
+
+    // Second graceful unmount on the same segment should also succeed
+    // (idempotent)
+    auto result2 = service_->GracefulUnmountSegment(segment.id, client_id,
+                                                    /*grace_period_ms=*/1000);
+    EXPECT_TRUE(result2.has_value()) << "Graceful unmount should be idempotent";
+
+    std::this_thread::sleep_for(std::chrono::milliseconds(1500));
+}
+
+TEST_F(MasterServiceTest, GracefulUnmountSegment_TimerExpiresAndUnmounts) {
+    std::unique_ptr<MasterService> service_(new MasterService());
+    auto segment = MakeSegment("graceful_timer_segment");
+    UUID client_id = generate_uuid();
+
+    ASSERT_TRUE(service_->MountSegment(segment, client_id).has_value());
+
+    // Graceful unmount with a short grace period (50ms)
+    auto graceful_result = service_->GracefulUnmountSegment(
+        segment.id, client_id, /*grace_period_ms=*/50);
+    ASSERT_TRUE(graceful_result.has_value());
+
+    // Immediately after graceful unmount, segment should still exist
+    auto status_immediate = service_->QuerySegmentStatus(segment.name);
+    ASSERT_TRUE(status_immediate.has_value());
+    EXPECT_EQ(status_immediate.value(), SegmentStatus::GRACEFULLY_UNMOUNTING);
+
+    // Wait for timer to expire and unmount (give some margin)
+    std::this_thread::sleep_for(std::chrono::milliseconds(300));
+
+    // After timer expires, segment should be fully unmounted (UNDEFINED or
+    // error)
+    auto status_after = service_->QuerySegmentStatus(segment.name);
+    // Segment may be UNDEFINED (not found) or return an error
+    EXPECT_TRUE(!status_after.has_value() ||
+                status_after.value() == SegmentStatus::UNDEFINED)
+        << "Segment should be unmounted after timer expires, got status="
+        << (status_after.has_value() ? static_cast<int>(status_after.value())
+                                     : -1);
+}
+
+TEST_F(MasterServiceTest,
+       GracefulUnmountSegment_QueryStatusByIdWithReusedName) {
+    std::unique_ptr<MasterService> service_(new MasterService());
+    auto old_segment = MakeSegment("graceful_reused_name_segment");
+    auto new_segment = MakeSegment(old_segment.name, /*base=*/0x400000000);
+    UUID client_id = generate_uuid();
+
+    ASSERT_TRUE(service_->MountSegment(old_segment, client_id).has_value());
+    ASSERT_TRUE(service_
+                    ->GracefulUnmountSegment(old_segment.id, client_id,
+                                             /*grace_period_ms=*/50)
+                    .has_value());
+    ASSERT_TRUE(service_->MountSegment(new_segment, client_id).has_value());
+
+    auto old_status = service_->QuerySegmentStatusById(old_segment.id);
+    ASSERT_TRUE(old_status.has_value());
+    EXPECT_EQ(old_status.value(), SegmentStatus::GRACEFULLY_UNMOUNTING);
+
+    auto new_status = service_->QuerySegmentStatusById(new_segment.id);
+    ASSERT_TRUE(new_status.has_value());
+    EXPECT_EQ(new_status.value(), SegmentStatus::OK);
+
+    std::this_thread::sleep_for(std::chrono::milliseconds(300));
+
+    EXPECT_FALSE(service_->QuerySegmentStatusById(old_segment.id).has_value());
+    ASSERT_TRUE(service_->QuerySegmentStatusById(new_segment.id).has_value());
+
+    auto status_by_name = service_->QuerySegmentStatus(new_segment.name);
+    ASSERT_TRUE(status_by_name.has_value());
+    EXPECT_EQ(status_by_name.value(), SegmentStatus::OK);
+}
+
+TEST_F(MasterServiceTest, GracefulUnmountSegment_EarlierTimerPreemptsWait) {
+    std::unique_ptr<MasterService> service_(new MasterService());
+    auto long_segment = MakeSegment("graceful_long_timer_segment");
+    auto short_segment =
+        MakeSegment("graceful_short_timer_segment", /*base=*/0x400000000);
+    UUID client_id = generate_uuid();
+
+    ASSERT_TRUE(service_->MountSegment(long_segment, client_id).has_value());
+    ASSERT_TRUE(service_->MountSegment(short_segment, client_id).has_value());
+
+    ASSERT_TRUE(service_
+                    ->GracefulUnmountSegment(long_segment.id, client_id,
+                                             /*grace_period_ms=*/1000)
+                    .has_value());
+    std::this_thread::sleep_for(std::chrono::milliseconds(20));
+    ASSERT_TRUE(service_
+                    ->GracefulUnmountSegment(short_segment.id, client_id,
+                                             /*grace_period_ms=*/50)
+                    .has_value());
+
+    std::this_thread::sleep_for(std::chrono::milliseconds(300));
+
+    auto short_status = service_->QuerySegmentStatus(short_segment.name);
+    EXPECT_TRUE(!short_status.has_value() ||
+                short_status.value() == SegmentStatus::UNDEFINED);
+
+    auto long_status = service_->QuerySegmentStatus(long_segment.name);
+    ASSERT_TRUE(long_status.has_value());
+    EXPECT_EQ(long_status.value(), SegmentStatus::GRACEFULLY_UNMOUNTING);
+}
+
+TEST_F(MasterServiceTest, GracefulUnmountSegment_PreventAllocation) {
+    std::unique_ptr<MasterService> service_(new MasterService());
+    auto segment1 = MakeSegment("graceful_seg1");
+    auto segment2 = MakeSegment("graceful_seg2", /*base=*/0x400000000);
+    UUID client_id = generate_uuid();
+
+    ASSERT_TRUE(service_->MountSegment(segment1, client_id).has_value());
+    ASSERT_TRUE(service_->MountSegment(segment2, client_id).has_value());
+
+    // Put an object on segment1
+    std::string key = "test_key_prevent_alloc";
+    ReplicateConfig config;
+    config.replica_num = 1;
+    config.preferred_segment = segment1.name;
+
+    auto put_start =
+        service_->PutStart(client_id, key, "default", 1024, config);
+    ASSERT_TRUE(put_start.has_value());
+    ASSERT_TRUE(service_->PutEnd(client_id, key, "default", ReplicaType::MEMORY)
+                    .has_value());
+
+    // Graceful unmount segment1
+    ASSERT_TRUE(service_->GracefulUnmountSegment(segment1.id, client_id, 1000)
+                    .has_value());
+
+    // Segment1 status should be GRACEFULLY_UNMOUNTING
+    auto status1 = service_->QuerySegmentStatus(segment1.name);
+    ASSERT_TRUE(status1.has_value());
+    EXPECT_EQ(status1.value(), SegmentStatus::GRACEFULLY_UNMOUNTING);
+
+    // Existing replicas on the graceful segment should remain readable during
+    // the grace window.
+    auto existing_replicas = service_->GetReplicaList(key, "default");
+    ASSERT_TRUE(existing_replicas.has_value());
+    ASSERT_EQ(existing_replicas->replicas.size(), 1u);
+    EXPECT_EQ(existing_replicas->replicas[0]
+                  .get_memory_descriptor()
+                  .buffer_descriptor.transport_endpoint_,
+              segment1.name);
+
+    // Segment2 status should still be OK
+    auto status2 = service_->QuerySegmentStatus(segment2.name);
+    ASSERT_TRUE(status2.has_value());
+    EXPECT_EQ(status2.value(), SegmentStatus::OK);
+
+    // New put without preferred_segment should succeed on segment2
+    std::string key2 = "test_key_after_graceful";
+    ReplicateConfig config2;
+    config2.replica_num = 1;
+
+    auto put_start2 =
+        service_->PutStart(client_id, key2, "default", 1024, config2);
+    ASSERT_TRUE(put_start2.has_value());
+    auto replicas = put_start2.value();
+    ASSERT_EQ(replicas.size(), 1u);
+    // Should be placed on segment2, not segment1
+    EXPECT_EQ(replicas[0]
+                  .get_memory_descriptor()
+                  .buffer_descriptor.transport_endpoint_,
+              segment2.name);
+
+    std::this_thread::sleep_for(std::chrono::milliseconds(1500));
+}
+
 }  // namespace mooncake::test
 
 int main(int argc, char** argv) {
diff --git a/mooncake-store/tests/mmap_arena_fallback_test.cpp b/mooncake-store/tests/mmap_arena_fallback_test.cpp
new file mode 100644
index 00000000..8e6289b1
--- /dev/null
+++ b/mooncake-store/tests/mmap_arena_fallback_test.cpp
@@ -0,0 +1,153 @@
+// Copyright 2026 KVCache.AI
+// Fallback-path tests for the global mmap allocator wrapper.
+
+#include <glog/logging.h>
+#include <gtest/gtest.h>
+
+#include <cstdint>
+#include <cstring>
+#include <cstdlib>
+#include <fstream>
+#include <string>
+
+#include "utils.h"
+
+namespace mooncake {
+
+namespace {
+
+bool HostHasReservedHugepages() {
+    std::ifstream meminfo("/proc/meminfo");
+    std::string key;
+    size_t value = 0;
+    std::string unit;
+    while (meminfo >> key >> value >> unit) {
+        if (key == "HugePages_Total:") {
+            return value > 0;
+        }
+    }
+    return false;
+}
+
+}  // namespace
+
+class MmapArenaFallbackTest : public ::testing::Test {
+   protected:
+    void SetUp() override {
+        FLAGS_logtostderr = 1;
+        FLAGS_minloglevel = google::WARNING;
+        setenv("MC_DISABLE_MMAP_ARENA", "1", 1);
+    }
+
+    void TearDown() override {
+        unsetenv("MC_DISABLE_MMAP_ARENA");
+        unsetenv("MC_MMAP_ARENA_POOL_SIZE");
+        unsetenv("MC_STORE_USE_HUGEPAGE");
+    }
+};
+
+TEST_F(MmapArenaFallbackTest, ArenaInitFailureIsStickyForProcessLifetime) {
+    unsetenv("MC_DISABLE_MMAP_ARENA");
+    unsetenv("MC_STORE_USE_HUGEPAGE");
+
+    // "infinite" parses to UINT64_MAX and deterministically trips the arena's
+    // overflow guard before any mmap attempt.
+    setenv("MC_MMAP_ARENA_POOL_SIZE", "infinite", 1);
+    FLAGS_minloglevel = google::INFO;
+    testing::internal::CaptureStderr();
+    void* first_ptr = allocate_buffer_mmap_memory(64 * 1024, 64);
+    ASSERT_NE(first_ptr, nullptr);
+    free_buffer_mmap_memory(first_ptr, 64 * 1024);
+    const std::string first_logs = testing::internal::GetCapturedStderr();
+    EXPECT_NE(first_logs.find("ARENA INITIALIZATION FAILED"),
+              std::string::npos);
+    EXPECT_NE(first_logs.find("only attempted once per process"),
+              std::string::npos);
+
+    // Fix the env and try again in the same process. std::call_once should
+    // keep us on the fallback path without re-running initialization.
+    setenv("MC_MMAP_ARENA_POOL_SIZE", "2gb", 1);
+    testing::internal::CaptureStderr();
+    void* second_ptr = allocate_buffer_mmap_memory(64 * 1024, 64);
+    ASSERT_NE(second_ptr, nullptr);
+    free_buffer_mmap_memory(second_ptr, 64 * 1024);
+    const std::string second_logs = testing::internal::GetCapturedStderr();
+    EXPECT_EQ(second_logs.find("ARENA INITIALIZATION FAILED"),
+              std::string::npos);
+    EXPECT_EQ(second_logs.find("ARENA ALLOCATOR ENABLED"), std::string::npos);
+}
+
+TEST_F(MmapArenaFallbackTest,
+       ExplicitHugepageRequestDoesNotSilentlyFallbackToRegularPages) {
+    if (HostHasReservedHugepages()) {
+        GTEST_SKIP() << "Host has reserved hugepages; this strict-fallback "
+                        "test only proves itself on hugepage-free hosts";
+    }
+
+    unsetenv("MC_DISABLE_MMAP_ARENA");
+    setenv("MC_STORE_USE_HUGEPAGE", "1", 1);
+    setenv("MC_MMAP_ARENA_POOL_SIZE", "2mb", 1);
+
+    FLAGS_minloglevel = google::INFO;
+    testing::internal::CaptureStderr();
+    void* ptr = allocate_buffer_mmap_memory(64 * 1024, 64);
+    const std::string logs = testing::internal::GetCapturedStderr();
+
+    EXPECT_EQ(ptr, nullptr) << logs;
+    EXPECT_EQ(logs.find("retrying without huge pages"), std::string::npos)
+        << logs;
+}
+
+TEST_F(MmapArenaFallbackTest, HonorsPageAlignment) {
+    const size_t alloc_size = 64 * 1024;
+    constexpr size_t alignment = 64;
+
+    void* ptr = allocate_buffer_mmap_memory(alloc_size, alignment);
+    ASSERT_NE(ptr, nullptr);
+    EXPECT_EQ(reinterpret_cast<uintptr_t>(ptr) % alignment, 0u);
+
+    memset(ptr, 0xAB, alloc_size);
+    EXPECT_EQ(static_cast<uint8_t*>(ptr)[0], 0xAB);
+    EXPECT_EQ(static_cast<uint8_t*>(ptr)[alloc_size - 1], 0xAB);
+
+    free_buffer_mmap_memory(ptr, alloc_size);
+}
+
+TEST_F(MmapArenaFallbackTest, NoHugepagesAllocFree) {
+    unsetenv("MC_STORE_USE_HUGEPAGE");
+
+    const size_t alloc_size = 65000;
+    constexpr size_t alignment = 64;
+
+    void* ptr = allocate_buffer_mmap_memory(alloc_size, alignment);
+    ASSERT_NE(ptr, nullptr);
+    EXPECT_EQ(reinterpret_cast<uintptr_t>(ptr) % 4096, 0u);
+
+    memset(ptr, 0xCD, alloc_size);
+    EXPECT_EQ(static_cast<uint8_t*>(ptr)[0], 0xCD);
+
+    free_buffer_mmap_memory(ptr, alloc_size);
+}
+
+TEST_F(MmapArenaFallbackTest, AllocateFreeCycle) {
+    constexpr int kCycles = 8;
+    constexpr size_t alloc_size = 128 * 1024;
+    constexpr size_t alignment = 64;
+
+    for (int i = 0; i < kCycles; ++i) {
+        void* ptr = allocate_buffer_mmap_memory(alloc_size, alignment);
+        ASSERT_NE(ptr, nullptr) << "Allocation failed on cycle " << i;
+        EXPECT_EQ(reinterpret_cast<uintptr_t>(ptr) % alignment, 0u);
+        memset(ptr, static_cast<uint8_t>(i), alloc_size);
+        free_buffer_mmap_memory(ptr, alloc_size);
+    }
+}
+
+}  // namespace mooncake
+
+int main(int argc, char** argv) {
+    ::testing::InitGoogleTest(&argc, argv);
+    google::InitGoogleLogging(argv[0]);
+    FLAGS_logtostderr = 1;
+    return RUN_ALL_TESTS();
+}
diff --git a/mooncake-store/tests/mmap_arena_test.cpp b/mooncake-store/tests/mmap_arena_test.cpp
new file mode 100644
index 00000000..05425f0b
--- /dev/null
+++ b/mooncake-store/tests/mmap_arena_test.cpp
@@ -0,0 +1,751 @@
+// Copyright 2026 KVCache.AI
+// Unit tests for MmapArena allocator - Production-grade security and
+// correctness tests
+
+#include <gtest/gtest.h>
+#include <glog/logging.h>
+#include "mmap_arena.h"
+#include "utils.h"
+#include <thread>
+#include <vector>
+#include <atomic>
+#include <set>
+#include <cstring>
+#include <sys/mman.h>
+#include <unistd.h>
+
+namespace mooncake {
+
+class MmapArenaTest : public ::testing::Test {
+   protected:
+    void SetUp() override {
+        FLAGS_logtostderr = 1;
+        FLAGS_minloglevel = google::WARNING;  // Reduce log noise in tests
+    }
+};
+
+// ===== BASIC FUNCTIONALITY TESTS =====
+
+TEST_F(MmapArenaTest, BasicInitialization) {
+    MmapArena arena;
+    ASSERT_TRUE(arena.initialize(1024 * 1024));  // 1MB pool
+    ASSERT_TRUE(arena.isInitialized());
+
+    auto stats = arena.getStats();
+    EXPECT_EQ(stats.pool_size, 2 * 1024 * 1024);  // Aligned to 2MB (huge page)
+    EXPECT_EQ(stats.reserved_bytes, 0);
+    EXPECT_EQ(stats.num_allocations, 0);
+    EXPECT_EQ(stats.num_failed_allocs, 0);
+}
+
+TEST_F(MmapArenaTest, BasicAllocation) {
+    MmapArena arena;
+    ASSERT_TRUE(arena.initialize(1024 * 1024));
+
+    void* ptr = arena.allocate(1024);
+    ASSERT_NE(ptr, nullptr);
+
+    auto stats = arena.getStats();
+    EXPECT_EQ(stats.num_allocations, 1);
+    EXPECT_GE(stats.reserved_bytes, 1024);
+    EXPECT_LE(stats.reserved_bytes, 1024 + 64);  // Accounting for alignment
+}
+
+TEST_F(MmapArenaTest, AllocationAlignment) {
+    MmapArena arena;
+    ASSERT_TRUE(arena.initialize(1024 * 1024, 64));
+
+    for (size_t size : {1, 63, 64, 65, 100, 1000}) {
+        void* ptr = arena.allocate(size);
+        ASSERT_NE(ptr, nullptr) << "Failed to allocate size=" << size;
+
+        // Verify 64-byte alignment
+        uintptr_t addr = reinterpret_cast<uintptr_t>(ptr);
+        EXPECT_EQ(addr % 64, 0) << "Pointer not aligned: " << ptr;
+
+        // Verify memory is writable
+        std::memset(ptr, 0xAA, size);
+    }
+}
+
+TEST_F(MmapArenaTest, ZeroSizeAllocation) {
+    MmapArena arena;
+    ASSERT_TRUE(arena.initialize(1024 * 1024));
+
+    void* ptr = arena.allocate(0);
+    EXPECT_EQ(ptr, nullptr);
+
+    auto stats = arena.getStats();
+    EXPECT_EQ(stats.num_allocations, 0);
+}
+
+TEST_F(MmapArenaTest, UninitializedAllocation) {
+    MmapArena arena;
+    EXPECT_FALSE(arena.isInitialized());
+
+    void* ptr = arena.allocate(1024);
+    EXPECT_EQ(ptr, nullptr);
+}
+
+TEST_F(MmapArenaTest, DoubleInitialization) {
+    MmapArena arena;
+    ASSERT_TRUE(arena.initialize(1024 * 1024));
+    EXPECT_FALSE(arena.initialize(1024 * 1024));  // Second init should fail
+
+    // Arena should still work
+    void* ptr = arena.allocate(1024);
+    EXPECT_NE(ptr, nullptr);
+}
+
+// ===== BUG #1: OOM CHECK RACE CONDITION (CRITICAL) =====
+
+TEST_F(MmapArenaTest, OOMDoesNotCorruptCursor) {
+    MmapArena arena;
+    ASSERT_TRUE(arena.initialize(
+        1024));  // Requested 1KB, but aligned to 2MB (huge page)
+
+    auto stats_init = arena.getStats();
+    size_t pool_size =
+        stats_init.pool_size;  // Actual pool size after alignment
+
+    // Fill the arena completely
+    std::vector<void*> ptrs;
+    while (true) {
+        void* ptr = arena.allocate(64);
+        if (ptr == nullptr) break;
+        ptrs.push_back(ptr);
+    }
+
+    // CRITICAL: Cursor should be at or below pool_size, not corrupted
+    auto stats = arena.getStats();
+    ASSERT_LE(stats.reserved_bytes, pool_size);
+    ASSERT_GT(stats.num_failed_allocs, 0);
+
+    // Subsequent allocations should still fail gracefully, not crash
+    void* p_extra = arena.allocate(1);
+    EXPECT_EQ(p_extra, nullptr);
+
+    // Verify cursor didn't go past pool_size
+    stats = arena.getStats();
+    ASSERT_LE(stats.reserved_bytes, pool_size);
+}
+
+TEST_F(MmapArenaTest, ConcurrentOOMStressTest) {
+    MmapArena arena;
+    const size_t requested_pool_size = 1024 * 1024;  // 1MB requested
+    ASSERT_TRUE(arena.initialize(requested_pool_size));
+
+    auto stats_init = arena.getStats();
+    size_t actual_pool_size =
+        stats_init.pool_size;  // Actual pool after alignment
+
+    std::atomic<int> succeeded{0};
+    std::atomic<int> failed{0};
+
+    const int num_threads = 16;
+    // Allocate enough to guarantee OOM
+    const int allocs_per_thread = (actual_pool_size / (64 * num_threads)) + 100;
+    const size_t alloc_size = 64;
+
+    std::vector<std::thread> threads;
+    for (int i = 0; i < num_threads; ++i) {
+        threads.emplace_back([&]() {
+            for (int j = 0; j < allocs_per_thread; ++j) {
+                void* ptr = arena.allocate(alloc_size);
+                if (ptr != nullptr) {
+                    succeeded.fetch_add(1, std::memory_order_relaxed);
+                    // Write pattern to verify no corruption
+                    std::memset(ptr, 0xBB, alloc_size);
+                } else {
+                    failed.fetch_add(1, std::memory_order_relaxed);
+                }
+            }
+        });
+    }
+
+    for (auto& t : threads) t.join();
+
+    auto stats = arena.getStats();
+
+    // Verify cursor didn't go beyond pool
+    ASSERT_LE(stats.reserved_bytes, actual_pool_size);
+
+    // Some allocations should have succeeded
+    ASSERT_GT(succeeded.load(), 0);
+
+    // Some should have failed (pool exhausted)
+    ASSERT_GT(failed.load(), 0);
+
+    // Total attempts should match
+    EXPECT_EQ(succeeded.load() + failed.load(),
+              num_threads * allocs_per_thread);
+
+    LOG(INFO) << "OOM stress test: " << succeeded.load() << " succeeded, "
+              << failed.load() << " failed, pool utilization: "
+              << (100.0 * stats.reserved_bytes / stats.pool_size) << "%";
+}
+
+// ===== BUG #2 & #3: INTEGER OVERFLOW TESTS =====
+
+TEST_F(MmapArenaTest, IntegerOverflowInBoundsCheck) {
+    MmapArena arena;
+    ASSERT_TRUE(arena.initialize(1024 * 1024));
+
+    // Try to allocate SIZE_MAX (should fail, not wrap around)
+    void* ptr = arena.allocate(SIZE_MAX);
+    EXPECT_EQ(ptr, nullptr);
+
+    auto stats = arena.getStats();
+    EXPECT_EQ(stats.num_failed_allocs, 1);
+    EXPECT_EQ(stats.reserved_bytes, 0);
+}
+
+TEST_F(MmapArenaTest, AlignmentOverflow) {
+    MmapArena arena;
+    ASSERT_TRUE(arena.initialize(1024 * 1024));
+
+    // Request size that would overflow during alignment
+    // SIZE_MAX - 10 + 64 - 1 = overflow
+    void* ptr = arena.allocate(SIZE_MAX - 10);
+    EXPECT_EQ(ptr, nullptr);
+
+    auto stats = arena.getStats();
+    EXPECT_EQ(stats.num_failed_allocs, 1);
+    EXPECT_EQ(stats.reserved_bytes, 0);
+}
+
+TEST_F(MmapArenaTest, NearMaxSizeAllocation) {
+    MmapArena arena;
+    ASSERT_TRUE(arena.initialize(1024 * 1024));
+
+    // Try allocating near SIZE_MAX / 2 (should fail gracefully)
+    void* ptr = arena.allocate(SIZE_MAX / 2);
+    EXPECT_EQ(ptr, nullptr);
+
+    auto stats = arena.getStats();
+    EXPECT_EQ(stats.num_failed_allocs, 1);
+}
+
+// ===== BUG #4: ARENA MEMORY CANNOT BE FREED =====
+
+TEST_F(MmapArenaTest, OwnershipDetection) {
+    MmapArena arena;
+    ASSERT_TRUE(arena.initialize(1024 * 1024));
+
+    void* arena_ptr = arena.allocate(1024);
+    ASSERT_NE(arena_ptr, nullptr);
+
+    // Arena should own its allocations
+    EXPECT_TRUE(arena.owns(arena_ptr));
+
+    // Null pointer is not owned
+    EXPECT_FALSE(arena.owns(nullptr));
+
+    // Pointer outside arena range should not be owned
+    char stack_var;
+    EXPECT_FALSE(arena.owns(&stack_var));
+}
+
+TEST_F(MmapArenaTest, OwnershipBoundaryTest) {
+    MmapArena arena;
+    ASSERT_TRUE(arena.initialize(1024 * 1024));
+
+    void* first_ptr = arena.allocate(64);
+    ASSERT_NE(first_ptr, nullptr);
+
+    void* pool_base = arena.getPoolBase();
+    size_t pool_size = arena.getPoolSize();
+
+    // Test boundaries
+    EXPECT_TRUE(arena.owns(pool_base));  // Start of pool
+    EXPECT_TRUE(arena.owns(static_cast<char*>(pool_base) + 100));  // Middle
+    EXPECT_FALSE(arena.owns(static_cast<char*>(pool_base) +
+                            pool_size));  // Just past end
+    EXPECT_FALSE(
+        arena.owns(static_cast<char*>(pool_base) - 1));  // Just before start
+}
+
+// ===== BUG #5: RACE CONDITION IN INITIALIZE =====
+
+TEST_F(MmapArenaTest, ConcurrentInitialization) {
+    MmapArena arena;
+
+    std::atomic<int> init_success{0};
+    std::atomic<int> init_failure{0};
+
+    const int num_threads = 16;
+    std::vector<std::thread> threads;
+
+    // Multiple threads racing to initialize
+    for (int i = 0; i < num_threads; ++i) {
+        threads.emplace_back([&]() {
+            if (arena.initialize(1024 * 1024)) {
+                init_success.fetch_add(1, std::memory_order_relaxed);
+            } else {
+                init_failure.fetch_add(1, std::memory_order_relaxed);
+            }
+        });
+    }
+
+    for (auto& t : threads) t.join();
+
+    // Exactly one initialization should succeed
+    EXPECT_EQ(init_success.load(), 1);
+    EXPECT_EQ(init_failure.load(), num_threads - 1);
+
+    // Arena should be usable
+    EXPECT_TRUE(arena.isInitialized());
+    void* ptr = arena.allocate(1024);
+    EXPECT_NE(ptr, nullptr);
+}
+
+TEST_F(MmapArenaTest, AllocateAfterRacyInit) {
+    MmapArena arena;
+
+    const int num_threads = 10;
+    std::vector<std::thread> threads;
+    std::vector<void*> ptrs(num_threads);
+
+    // Concurrent init + allocate
+    for (int i = 0; i < num_threads; ++i) {
+        threads.emplace_back([&, i]() {
+            arena.initialize(1024 * 1024);
+            ptrs[i] = arena.allocate(1024);
+        });
+    }
+
+    for (auto& t : threads) t.join();
+
+    // All allocations should either succeed or fail consistently
+    int non_null = std::count_if(ptrs.begin(), ptrs.end(),
+                                 [](void* p) { return p != nullptr; });
+    EXPECT_GT(non_null, 0);  // At least some should succeed
+
+    // Verify no duplicate pointers
+    std::set<void*> unique_ptrs(ptrs.begin(), ptrs.end());
+    unique_ptrs.erase(nullptr);
+    EXPECT_EQ(unique_ptrs.size(), static_cast<size_t>(non_null));
+}
+
+// ===== THREAD SAFETY & CONCURRENCY TESTS =====
+
+TEST_F(MmapArenaTest, ConcurrentAllocations) {
+    MmapArena arena;
+    ASSERT_TRUE(arena.initialize(64 * 1024 * 1024));  // 64MB
+
+    const int num_threads = 8;
+    const int allocs_per_thread = 1000;
+    std::vector<std::thread> threads;
+    std::vector<std::vector<void*>> all_pointers(num_threads);
+
+    for (int i = 0; i < num_threads; ++i) {
+        threads.emplace_back([&, i]() {
+            for (int j = 0; j < allocs_per_thread; ++j) {
+                void* ptr = arena.allocate(1024);
+                if (ptr != nullptr) {
+                    all_pointers[i].push_back(ptr);
+                    // Write unique pattern
+                    std::memset(ptr, 0xCC + i, 1024);
+                }
+            }
+        });
+    }
+
+    for (auto& t : threads) t.join();
+
+    // Verify no duplicate pointers (uniqueness)
+    std::set<void*> unique_ptrs;
+    for (const auto& vec : all_pointers) {
+        for (void* ptr : vec) {
+            auto [iter, inserted] = unique_ptrs.insert(ptr);
+            EXPECT_TRUE(inserted) << "Duplicate pointer detected: " << ptr;
+        }
+    }
+
+    LOG(INFO) << "Concurrent allocations: " << unique_ptrs.size()
+              << " unique allocations from " << num_threads << " threads";
+}
+
+TEST_F(MmapArenaTest, StatsConsistencyUnderLoad) {
+    MmapArena arena;
+    ASSERT_TRUE(arena.initialize(64 * 1024 * 1024));
+
+    std::atomic<bool> stop{false};
+    std::atomic<int> invariant_violations{0};
+
+    // Allocator threads
+    std::vector<std::thread> threads;
+    for (int i = 0; i < 8; ++i) {
+        threads.emplace_back([&]() {
+            while (!stop.load(std::memory_order_relaxed)) {
+                void* ptr = arena.allocate(128);
+                (void)ptr;
+            }
+        });
+    }
+
+    // Stats checker thread
+    threads.emplace_back([&]() {
+        for (int i = 0; i < 100; ++i) {
+            auto stats = arena.getStats();
+            // Invariants that must always hold
+            if (stats.reserved_bytes > stats.pool_size) {
+                invariant_violations.fetch_add(1, std::memory_order_relaxed);
+            }
+            // Note: peak_reserved_bytes may temporarily lag behind
+            // reserved_bytes
+            // due to concurrent updates, so we don't check that invariant here
+            std::this_thread::sleep_for(std::chrono::milliseconds(10));
+        }
+        stop.store(true, std::memory_order_relaxed);
+    });
+
+    for (auto& t : threads) t.join();
+
+    auto stats = arena.getStats();
+
+    // Critical invariant: cursor never exceeded pool size
+    EXPECT_EQ(invariant_violations.load(), 0);
+    EXPECT_LE(stats.reserved_bytes, stats.pool_size);
+
+    // After all threads finish, peak should be >= final allocated
+    EXPECT_GE(stats.peak_reserved_bytes, stats.reserved_bytes);
+
+    LOG(INFO) << "Stats consistency test: " << stats.num_allocations
+              << " allocations, " << stats.num_failed_allocs << " failures";
+}
+
+// ===== EDGE CASES & STRESS TESTS =====
+
+TEST_F(MmapArenaTest, NearOOMAllocation) {
+    MmapArena arena;
+    const size_t pool_size = 4096;  // Small pool
+    ASSERT_TRUE(arena.initialize(pool_size));
+
+    // Fill arena almost completely
+    std::vector<void*> ptrs;
+    while (true) {
+        void* ptr = arena.allocate(64);
+        if (ptr == nullptr) break;
+        ptrs.push_back(ptr);
+    }
+
+    auto stats = arena.getStats();
+    EXPECT_LE(stats.reserved_bytes, stats.pool_size);
+    EXPECT_GT(stats.num_failed_allocs, 0);
+
+    LOG(INFO) << "Near-OOM test: " << ptrs.size() << " allocations, "
+              << stats.reserved_bytes << " / " << stats.pool_size
+              << " bytes used";
+}
+
+TEST_F(MmapArenaTest, MixedSizeAllocations) {
+    MmapArena arena;
+    ASSERT_TRUE(arena.initialize(64 * 1024 * 1024));
+
+    std::vector<size_t> sizes = {1, 16, 64, 256, 1024, 4096, 16384, 65536};
+    std::vector<void*> ptrs;
+
+    // Allocate various sizes
+    for (size_t size : sizes) {
+        for (int i = 0; i < 10; ++i) {
+            void* ptr = arena.allocate(size);
+            if (ptr != nullptr) {
+                ptrs.push_back(ptr);
+                std::memset(ptr, 0xDD, size);
+            }
+        }
+    }
+
+    // Verify all pointers are valid and aligned
+    for (void* ptr : ptrs) {
+        EXPECT_NE(ptr, nullptr);
+        EXPECT_EQ(reinterpret_cast<uintptr_t>(ptr) % 64, 0);
+        EXPECT_TRUE(arena.owns(ptr));
+    }
+
+    LOG(INFO) << "Mixed-size test: " << ptrs.size() << " allocations";
+}
+
+TEST_F(MmapArenaTest, PeakAllocationTracking) {
+    MmapArena arena;
+    ASSERT_TRUE(arena.initialize(1024 * 1024));
+
+    void* p1 = arena.allocate(512);
+    (void)p1;
+    auto stats1 = arena.getStats();
+    EXPECT_GE(stats1.peak_reserved_bytes, 512);
+
+    void* p2 = arena.allocate(1024);
+    (void)p2;
+    auto stats2 = arena.getStats();
+    EXPECT_GE(stats2.peak_reserved_bytes, stats1.peak_reserved_bytes);
+    EXPECT_GE(stats2.peak_reserved_bytes, 512 + 1024);
+
+    LOG(INFO) << "Peak tracking: " << stats2.peak_reserved_bytes << " bytes";
+}
+
+TEST_F(MmapArenaTest, ReservedBytesRemainMonotonic) {
+    MmapArena arena;
+    ASSERT_TRUE(arena.initialize(1024 * 1024));
+
+    size_t last_reserved_bytes = 0;
+    for (size_t size : {1UL, 63UL, 64UL, 65UL, 4096UL, 1024UL}) {
+        void* ptr = arena.allocate(size);
+        ASSERT_NE(ptr, nullptr);
+
+        auto stats = arena.getStats();
+        EXPECT_GE(stats.reserved_bytes, last_reserved_bytes);
+        EXPECT_GE(stats.peak_reserved_bytes, stats.reserved_bytes);
+        last_reserved_bytes = stats.reserved_bytes;
+    }
+}
+
+// ===== MIXED ALIGNMENT TESTS =====
+
+TEST_F(MmapArenaTest, MixedAlignmentSequence) {
+    MmapArena arena;
+    ASSERT_TRUE(arena.initialize(64 * 1024 * 1024));  // 64MB
+
+    // First allocation moves cursor to 64 (non-trivially aligned).
+    // This ensures the subsequent 2MB-aligned allocation actually
+    // exercises the offset-alignment logic rather than passing
+    // vacuously because cursor 0 is trivially aligned to everything.
+    void* p1 = arena.allocate(1, 64);
+    ASSERT_NE(p1, nullptr);
+
+    const size_t TWO_MB = 2 * 1024 * 1024;
+    void* p2 = arena.allocate(4 * 1024 * 1024, TWO_MB);
+    ASSERT_NE(p2, nullptr);
+
+    // The returned pointer MUST be 2MB-aligned
+    EXPECT_EQ(reinterpret_cast<uintptr_t>(p2) % TWO_MB, 0)
+        << "Pointer not 2MB-aligned: " << p2;
+
+    // p2 must not overlap p1
+    EXPECT_GT(reinterpret_cast<uintptr_t>(p2), reinterpret_cast<uintptr_t>(p1));
+}
+
+// ===== INPUT VALIDATION TESTS =====
+
+TEST_F(MmapArenaTest, InitializeWithZeroSize) {
+    MmapArena arena;
+    EXPECT_FALSE(arena.initialize(0));
+    EXPECT_FALSE(arena.isInitialized());
+
+    // Allocations should fail gracefully on uninitialized arena
+    void* ptr = arena.allocate(1024);
+    EXPECT_EQ(ptr, nullptr);
+}
+
+TEST_F(MmapArenaTest, NonPowerOfTwoAlignment) {
+    MmapArena arena;
+    // alignment=100 is not a power of 2; should be rejected
+    EXPECT_FALSE(arena.initialize(1024 * 1024, 100));
+    EXPECT_FALSE(arena.isInitialized());
+}
+
+TEST_F(MmapArenaTest, PowerOfTwoAlignmentsAccepted) {
+    // Test several valid power-of-2 alignments
+    for (size_t align : {64, 128, 256, 512, 4096}) {
+        MmapArena arena;
+        ASSERT_TRUE(arena.initialize(4 * 1024 * 1024, align))
+            << "Failed to init with alignment=" << align;
+        void* ptr = arena.allocate(1024);
+        ASSERT_NE(ptr, nullptr);
+        EXPECT_EQ(reinterpret_cast<uintptr_t>(ptr) % align, 0)
+            << "Pointer not aligned to " << align;
+    }
+}
+
+// ===== SIZING REGRESSION TEST =====
+
+TEST_F(MmapArenaTest, ArenaSizingRegression) {
+    // Reproduces the benchmark failure: pool=16MB, allocate all of it,
+    // then any further allocation should OOM.
+    const size_t POOL = 16 * 1024 * 1024;  // 16MB (scaled down from 16GB)
+    MmapArena arena;
+    ASSERT_TRUE(arena.initialize(POOL, 64));
+
+    // First: allocation consumes entire pool
+    void* p1 = arena.allocate(POOL, 64);
+    ASSERT_NE(p1, nullptr);
+
+    // Second: any further allocation should OOM
+    void* p2 = arena.allocate(1024, 64);
+    EXPECT_EQ(p2, nullptr);
+    EXPECT_GE(arena.getStats().num_failed_allocs, 1);
+}
+
+// ===== CONCURRENT INIT METADATA CONSISTENCY =====
+
+TEST_F(MmapArenaTest, ConcurrentInitMetadataConsistency) {
+    // Verify that after racing inits, pool_size and alignment are consistent
+    // with the winning initialization parameters.
+    MmapArena arena;
+
+    const int num_threads = 16;
+    std::vector<std::thread> threads;
+
+    for (int i = 0; i < num_threads; ++i) {
+        threads.emplace_back([&]() { arena.initialize(4 * 1024 * 1024, 128); });
+    }
+    for (auto& t : threads) t.join();
+
+    ASSERT_TRUE(arena.isInitialized());
+    // Pool size should be 4MB (already 2MB-aligned, no rounding needed)
+    EXPECT_EQ(arena.getStats().pool_size, 4 * 1024 * 1024);
+
+    // Verify allocation works and honors alignment
+    void* ptr = arena.allocate(256);
+    ASSERT_NE(ptr, nullptr);
+    EXPECT_EQ(reinterpret_cast<uintptr_t>(ptr) % 128, 0);
+}
+
+// ===== MAP_POPULATE REGRESSION TESTS =====
+// These tests verify the fix for cudaErrorIllegalAddress caused by lazy
+// hugepage faults during GPU DMA.  The arena must pre-fault all pages
+// at initialization time (MAP_POPULATE) so that every byte in the pool
+// is backed by physical memory before any allocation is returned.
+
+TEST_F(MmapArenaTest, PagesArePhysicallyBackedAfterInit) {
+    // Verify that arena pages are resident in physical memory immediately
+    // after initialize() — i.e. MAP_POPULATE is working.
+    // Uses mincore() which reports per-page residency status.
+    MmapArena arena;
+    const size_t POOL = 4 * 1024 * 1024;  // 4MB
+    ASSERT_TRUE(arena.initialize(POOL));
+
+    void* base = arena.getPoolBase();
+    size_t pool_size = arena.getPoolSize();
+    ASSERT_NE(base, nullptr);
+    ASSERT_GT(pool_size, 0);
+
+    // mincore() works on the system page size (typically 4KB), not hugepages.
+    // Query the number of system pages covering the pool.
+    const size_t sys_page_size = sysconf(_SC_PAGESIZE);
+    size_t num_pages = (pool_size + sys_page_size - 1) / sys_page_size;
+
+    std::vector<unsigned char> vec(num_pages);
+    int ret = mincore(base, pool_size, vec.data());
+
+    if (ret == 0) {
+        // mincore succeeded — check that all pages are resident
+        size_t resident = 0;
+        for (size_t i = 0; i < num_pages; ++i) {
+            if (vec[i] & 1) ++resident;
+        }
+        // With MAP_POPULATE, all pages should be resident.
+        // Allow small tolerance for kernel behavior differences.
+        double pct = 100.0 * resident / num_pages;
+        EXPECT_GT(pct, 95.0)
+            << "Only " << pct
+            << "% of pages resident; MAP_POPULATE may not be working. "
+            << resident << "/" << num_pages << " pages.";
+        LOG(INFO) << "mincore: " << resident << "/" << num_pages
+                  << " pages resident (" << pct << "%)";
+    } else {
+        // mincore may fail on some kernels for MAP_HUGETLB regions.
+        // Fall back to verifying that we can read every byte without SIGSEGV.
+        LOG(WARNING) << "mincore() returned " << ret << " (errno=" << errno
+                     << "), falling back to read-verification";
+        // Read every page — if MAP_POPULATE didn't work, this would trigger
+        // page faults (which is fine for CPU but would crash GPU DMA).
+        volatile char sink = 0;
+        for (size_t off = 0; off < pool_size; off += sys_page_size) {
+            sink += static_cast<char*>(base)[off];
+        }
+        (void)sink;
+        // If we get here without SIGSEGV, at least CPU access works.
+        // The real MAP_POPULATE guarantee is that DMA works too, which
+        // can only be tested with actual GPU hardware.
+    }
+}
+
+TEST_F(MmapArenaTest, AllocatedMemoryIsImmediatelyReadableWritable) {
+    // Simulates the GPU DMA scenario: allocate a buffer and immediately
+    // read/write every byte.  Without MAP_POPULATE, a lazy hugepage fault
+    // during DMA would crash.  With MAP_POPULATE, all pages are pre-faulted.
+    MmapArena arena;
+    const size_t POOL = 8 * 1024 * 1024;  // 8MB
+    ASSERT_TRUE(arena.initialize(POOL));
+
+    // Allocate a large buffer (simulates segment allocation)
+    const size_t BUF_SIZE = 4 * 1024 * 1024;  // 4MB
+    void* ptr = arena.allocate(BUF_SIZE);
+    ASSERT_NE(ptr, nullptr);
+
+    // Write a pattern to every byte — would trigger page faults if lazy
+    std::memset(ptr, 0xAB, BUF_SIZE);
+
+    // Read it back — verify no corruption
+    auto* bytes = static_cast<unsigned char*>(ptr);
+    for (size_t i = 0; i < BUF_SIZE; i += 4096) {
+        EXPECT_EQ(bytes[i], 0xAB) << "Memory corruption at offset " << i;
+    }
+
+    // Allocate a second buffer from remaining space
+    void* ptr2 = arena.allocate(BUF_SIZE);
+    if (ptr2 != nullptr) {
+        // Write different pattern
+        std::memset(ptr2, 0xCD, BUF_SIZE);
+        auto* bytes2 = static_cast<unsigned char*>(ptr2);
+        for (size_t i = 0; i < BUF_SIZE; i += 4096) {
+            EXPECT_EQ(bytes2[i], 0xCD)
+                << "Memory corruption in second buffer at offset " << i;
+        }
+        // Verify first buffer wasn't corrupted by second allocation
+        EXPECT_EQ(bytes[0], 0xAB)
+            << "First buffer corrupted after second allocation";
+    }
+}
+
+TEST_F(MmapArenaTest, FallbackMmapRetainsPopulate) {
+    // When huge pages are unavailable, the arena falls back to regular mmap.
+    // Verify that MAP_POPULATE is retained in the fallback path by confirming
+    // the allocated memory is immediately usable (same as above but may
+    // exercise the non-hugepage code path on machines without huge pages).
+    MmapArena arena;
+    const size_t POOL = 2 * 1024 * 1024;  // 2MB — minimum hugepage unit
+    ASSERT_TRUE(arena.initialize(POOL));
+
+    void* ptr = arena.allocate(1024 * 1024);  // 1MB
+    ASSERT_NE(ptr, nullptr);
+
+    // Full read/write cycle
+    std::memset(ptr, 0xEF, 1024 * 1024);
+    auto* bytes = static_cast<unsigned char*>(ptr);
+    EXPECT_EQ(bytes[0], 0xEF);
+    EXPECT_EQ(bytes[1024 * 1024 - 1], 0xEF);
+    EXPECT_EQ(bytes[512 * 1024], 0xEF);  // Middle
+}
+
+// ===== FORK SAFETY TEST =====
+
+TEST_F(MmapArenaTest, MadviseDontForkApplied) {
+    // Verify that the arena applies MADV_DONTFORK to prevent 64GB CoW
+    // page table duplication on fork().  madvise(MADV_DONTFORK) is
+    // idempotent — calling it again on an already-marked region returns 0.
+    MmapArena arena;
+    const size_t POOL = 4 * 1024 * 1024;  // 4MB
+    ASSERT_TRUE(arena.initialize(POOL));
+
+    void* base = arena.getPoolBase();
+    size_t pool_size = arena.getPoolSize();
+    ASSERT_NE(base, nullptr);
+
+    // If MADV_DONTFORK was already applied by initialize(), this is a no-op
+    // and returns 0.  If it wasn't applied, this also returns 0 (first apply).
+    // Either way, verify madvise succeeds on the pool region.
+    int ret = madvise(base, pool_size, MADV_DONTFORK);
+    EXPECT_EQ(ret, 0) << "madvise(MADV_DONTFORK) failed: " << strerror(errno);
+}
+
+}  // namespace mooncake
+
+int main(int argc, char** argv) {
+    ::testing::InitGoogleTest(&argc, argv);
+    google::InitGoogleLogging(argv[0]);
+    FLAGS_logtostderr = 1;
+    return RUN_ALL_TESTS();
+}
diff --git a/mooncake-store/tests/nof_heartbeat_test.cpp b/mooncake-store/tests/nof_heartbeat_test.cpp
new file mode 100644
index 00000000..765e8cdc
--- /dev/null
+++ b/mooncake-store/tests/nof_heartbeat_test.cpp
@@ -0,0 +1,249 @@
+#include "master_service.h"
+
+#include <glog/logging.h>
+#include <gtest/gtest.h>
+
+#include <atomic>
+#include <chrono>
+#include <memory>
+#include <thread>
+
+namespace mooncake::test {
+
+namespace {
+
+constexpr size_t kDefaultNoFSegmentBase = 0x500000000;
+constexpr size_t kDefaultNoFSegmentSize = 1024 * 1024 * 16;
+
+NoFSegment MakeNoFSegment(std::string name, std::string endpoint,
+                          size_t base = kDefaultNoFSegmentBase,
+                          size_t size = kDefaultNoFSegmentSize) {
+    NoFSegment segment;
+    segment.id = generate_uuid();
+    segment.name = std::move(name);
+    segment.base = base;
+    segment.size = size;
+    segment.te_endpoint = std::move(endpoint);
+    return segment;
+}
+
+bool WaitForCondition(std::chrono::milliseconds timeout,
+                      std::chrono::milliseconds interval,
+                      const std::function<bool()>& condition) {
+    auto deadline = std::chrono::steady_clock::now() + timeout;
+    while (std::chrono::steady_clock::now() < deadline) {
+        if (condition()) {
+            return true;
+        }
+        std::this_thread::sleep_for(interval);
+    }
+    return condition();
+}
+
+}  // namespace
+
+class NoFHeartbeatTest : public ::testing::Test {
+   protected:
+    void SetUp() override {
+        google::InitGoogleLogging("NoFHeartbeatTest");
+        FLAGS_logtostderr = true;
+    }
+
+    void TearDown() override { google::ShutdownGoogleLogging(); }
+
+    std::unique_ptr<MasterService> CreateService(int64_t heartbeat_interval_sec,
+                                                 uint32_t probe_timeout_ms,
+                                                 uint32_t failure_threshold,
+                                                 int64_t client_ttl_sec = 10) {
+        auto config =
+            MasterServiceConfig::builder()
+                .set_memory_allocator(BufferAllocatorType::OFFSET)
+                .set_client_live_ttl_sec(client_ttl_sec)
+                .set_nof_heartbeat_interval_sec(heartbeat_interval_sec)
+                .set_nof_heartbeat_probe_timeout_ms(probe_timeout_ms)
+                .set_nof_heartbeat_failures_threshold(failure_threshold)
+                .build();
+        return std::make_unique<MasterService>(config);
+    }
+};
+
+TEST_F(NoFHeartbeatTest, HealthyNoFSegmentDoesNotUnmount) {
+    auto service = CreateService(/*heartbeat_interval_sec=*/1,
+                                 /*probe_timeout_ms=*/50,
+                                 /*failure_threshold=*/3);
+    std::atomic<int> probe_calls{0};
+    service->SetNoFProbeFnForTesting(
+        [&probe_calls](const std::string&, uint32_t, std::string*) {
+            probe_calls.fetch_add(1, std::memory_order_relaxed);
+            return true;
+        });
+
+    UUID client_id = generate_uuid();
+    NoFSegment segment = MakeNoFSegment("nof_seg_ok", "nof_ok");
+    ASSERT_TRUE(service->MountNoFSegment(segment, client_id).has_value());
+
+    ASSERT_TRUE(WaitForCondition(std::chrono::milliseconds(2500),
+                                 std::chrono::milliseconds(50),
+                                 [&]() { return probe_calls.load() >= 1; }));
+    EXPECT_TRUE(service->IsNoFSegmentMountedForTesting(segment.id));
+    EXPECT_EQ(service->GetMountedNoFSegmentCountForTesting(), 1u);
+    auto failure_count =
+        service->GetNoFHeartbeatFailureCountForTesting(segment.id);
+    ASSERT_TRUE(failure_count.has_value());
+    EXPECT_EQ(*failure_count, 0u);
+}
+
+TEST_F(NoFHeartbeatTest, NewlyMountedNoFSegmentHasInitialGracePeriod) {
+    auto service = CreateService(/*heartbeat_interval_sec=*/2,
+                                 /*probe_timeout_ms=*/50,
+                                 /*failure_threshold=*/1);
+    std::atomic<int> probe_calls{0};
+    service->SetNoFProbeFnForTesting(
+        [&probe_calls](const std::string&, uint32_t, std::string* reason) {
+            probe_calls.fetch_add(1, std::memory_order_relaxed);
+            if (reason) {
+                *reason = "submit_fail";
+            }
+            return false;
+        });
+
+    UUID client_id = generate_uuid();
+    NoFSegment segment = MakeNoFSegment("nof_seg_grace", "nof_grace");
+    ASSERT_TRUE(service->MountNoFSegment(segment, client_id).has_value());
+
+    std::this_thread::sleep_for(std::chrono::milliseconds(800));
+
+    EXPECT_EQ(probe_calls.load(), 0);
+    EXPECT_TRUE(service->IsNoFSegmentMountedForTesting(segment.id));
+    auto failure_count =
+        service->GetNoFHeartbeatFailureCountForTesting(segment.id);
+    ASSERT_TRUE(failure_count.has_value());
+    EXPECT_EQ(*failure_count, 0u);
+}
+
+TEST_F(NoFHeartbeatTest, FailedNoFSegmentUnmountsAfterThreshold) {
+    auto service = CreateService(/*heartbeat_interval_sec=*/1,
+                                 /*probe_timeout_ms=*/50,
+                                 /*failure_threshold=*/3);
+    std::atomic<int> probe_calls{0};
+    service->SetNoFProbeFnForTesting(
+        [&probe_calls](const std::string&, uint32_t, std::string* reason) {
+            probe_calls.fetch_add(1, std::memory_order_relaxed);
+            if (reason) {
+                *reason = "submit_fail";
+            }
+            return false;
+        });
+
+    UUID client_id = generate_uuid();
+    NoFSegment segment = MakeNoFSegment("nof_seg_fail", "nof_fail");
+    ASSERT_TRUE(service->MountNoFSegment(segment, client_id).has_value());
+
+    ASSERT_TRUE(WaitForCondition(
+        std::chrono::milliseconds(5000), std::chrono::milliseconds(50),
+        [&]() { return !service->IsNoFSegmentMountedForTesting(segment.id); }));
+    EXPECT_GE(probe_calls.load(), 3);
+    EXPECT_EQ(service->GetMountedNoFSegmentCountForTesting(), 0u);
+    EXPECT_FALSE(
+        service->GetNoFHeartbeatFailureCountForTesting(segment.id).has_value());
+}
+
+TEST_F(NoFHeartbeatTest, FailureCountResetsAfterRecovery) {
+    auto service = CreateService(/*heartbeat_interval_sec=*/1,
+                                 /*probe_timeout_ms=*/50,
+                                 /*failure_threshold=*/3);
+    std::atomic<int> probe_calls{0};
+    service->SetNoFProbeFnForTesting(
+        [&probe_calls](const std::string&, uint32_t, std::string* reason) {
+            int current = probe_calls.fetch_add(1, std::memory_order_relaxed);
+            if (current < 2) {
+                if (reason) {
+                    *reason = "submit_fail";
+                }
+                return false;
+            }
+            return true;
+        });
+
+    UUID client_id = generate_uuid();
+    NoFSegment segment = MakeNoFSegment("nof_seg_recover", "nof_recover");
+    ASSERT_TRUE(service->MountNoFSegment(segment, client_id).has_value());
+
+    ASSERT_TRUE(WaitForCondition(
+        std::chrono::milliseconds(5000), std::chrono::milliseconds(50), [&]() {
+            auto failure_count =
+                service->GetNoFHeartbeatFailureCountForTesting(segment.id);
+            return probe_calls.load() >= 4 && failure_count.has_value() &&
+                   *failure_count == 0;
+        }));
+    EXPECT_TRUE(service->IsNoFSegmentMountedForTesting(segment.id));
+}
+
+TEST_F(NoFHeartbeatTest, OnlyFailedSegmentIsUnmounted) {
+    auto service = CreateService(/*heartbeat_interval_sec=*/1,
+                                 /*probe_timeout_ms=*/50,
+                                 /*failure_threshold=*/2);
+    std::atomic<int> good_probe_calls{0};
+    std::atomic<int> bad_probe_calls{0};
+    service->SetNoFProbeFnForTesting(
+        [&good_probe_calls, &bad_probe_calls](const std::string& endpoint,
+                                              uint32_t, std::string* reason) {
+            if (endpoint == "nof_good") {
+                good_probe_calls.fetch_add(1, std::memory_order_relaxed);
+                return true;
+            }
+            bad_probe_calls.fetch_add(1, std::memory_order_relaxed);
+            if (reason) {
+                *reason = "submit_fail";
+            }
+            return false;
+        });
+
+    UUID client_id = generate_uuid();
+    NoFSegment good_segment = MakeNoFSegment("nof_seg_good", "nof_good");
+    NoFSegment bad_segment =
+        MakeNoFSegment("nof_seg_bad", "nof_bad",
+                       kDefaultNoFSegmentBase + kDefaultNoFSegmentSize,
+                       kDefaultNoFSegmentSize);
+    ASSERT_TRUE(service->MountNoFSegment(good_segment, client_id).has_value());
+    ASSERT_TRUE(service->MountNoFSegment(bad_segment, client_id).has_value());
+
+    ASSERT_TRUE(WaitForCondition(
+        std::chrono::milliseconds(5000), std::chrono::milliseconds(50), [&]() {
+            return service->GetMountedNoFSegmentCountForTesting() == 1u;
+        }));
+    EXPECT_TRUE(service->IsNoFSegmentMountedForTesting(good_segment.id));
+    EXPECT_FALSE(service->IsNoFSegmentMountedForTesting(bad_segment.id));
+    EXPECT_GE(good_probe_calls.load(), 1);
+    EXPECT_GE(bad_probe_calls.load(), 2);
+}
+
+TEST_F(NoFHeartbeatTest, ClientExpiryDoesNotUnmountNoFSegment) {
+    auto service = CreateService(/*heartbeat_interval_sec=*/5,
+                                 /*probe_timeout_ms=*/50,
+                                 /*failure_threshold=*/1,
+                                 /*client_ttl_sec=*/1);
+    service->SetNoFProbeFnForTesting(
+        [](const std::string&, uint32_t, std::string* reason) {
+            if (reason) {
+                *reason = "submit_fail";
+            }
+            return false;
+        });
+
+    UUID client_id = generate_uuid();
+    NoFSegment segment =
+        MakeNoFSegment("nof_seg_ignore_client_ttl", "nof_ignore_client_ttl");
+    ASSERT_TRUE(service->MountNoFSegment(segment, client_id).has_value());
+
+    std::this_thread::sleep_for(std::chrono::milliseconds(1800));
+
+    EXPECT_TRUE(service->IsNoFSegmentMountedForTesting(segment.id));
+    EXPECT_EQ(service->GetMountedNoFSegmentCountForTesting(), 1u);
+    auto failure_count =
+        service->GetNoFHeartbeatFailureCountForTesting(segment.id);
+    ASSERT_TRUE(failure_count.has_value());
+    EXPECT_EQ(*failure_count, 0u);
+}
+
+}  // namespace mooncake::test
diff --git a/mooncake-store/tests/object_data_type_test.cpp b/mooncake-store/tests/object_data_type_test.cpp
new file mode 100644
index 00000000..0c3da07b
--- /dev/null
+++ b/mooncake-store/tests/object_data_type_test.cpp
@@ -0,0 +1,165 @@
+#include "types.h"
+#include "replica.h"
+#include "master_service.h"
+
+#include <glog/logging.h>
+#include <gtest/gtest.h>
+
+#include <sstream>
+#include <vector>
+
+namespace mooncake::test {
+
+class ObjectDataTypeTest : public ::testing::Test {
+   protected:
+    void SetUp() override {
+        google::InitGoogleLogging("ObjectDataTypeTest");
+        FLAGS_logtostderr = true;
+    }
+
+    void TearDown() override { google::ShutdownGoogleLogging(); }
+
+    static constexpr size_t kDefaultSegmentBase = 0x300000000;
+    static constexpr size_t kDefaultSegmentSize = 1024 * 1024 * 16;
+
+    Segment MakeSegment(std::string name = "test_segment",
+                        size_t base = kDefaultSegmentBase,
+                        size_t size = kDefaultSegmentSize) const {
+        Segment segment;
+        segment.id = generate_uuid();
+        segment.name = std::move(name);
+        segment.base = base;
+        segment.size = size;
+        segment.te_endpoint = segment.name;
+        return segment;
+    }
+};
+
+// Verify enum values match the RFC spec
+TEST_F(ObjectDataTypeTest, EnumValues) {
+    EXPECT_EQ(static_cast<uint8_t>(ObjectDataType::UNKNOWN), 0);
+    EXPECT_EQ(static_cast<uint8_t>(ObjectDataType::KVCACHE), 1);
+    EXPECT_EQ(static_cast<uint8_t>(ObjectDataType::TENSOR), 2);
+    EXPECT_EQ(static_cast<uint8_t>(ObjectDataType::WEIGHT), 3);
+    EXPECT_EQ(static_cast<uint8_t>(ObjectDataType::SAMPLE), 4);
+    EXPECT_EQ(static_cast<uint8_t>(ObjectDataType::ACTIVATION), 5);
+    EXPECT_EQ(static_cast<uint8_t>(ObjectDataType::GRADIENT), 6);
+    EXPECT_EQ(static_cast<uint8_t>(ObjectDataType::OPTIMIZER_STATE), 7);
+    EXPECT_EQ(static_cast<uint8_t>(ObjectDataType::METADATA), 8);
+    EXPECT_EQ(static_cast<uint8_t>(ObjectDataType::GENERAL), 9);
+    EXPECT_EQ(static_cast<uint8_t>(ObjectDataType::HIDDEN_STATE), 10);
+}
+
+// Verify stream operator produces readable output
+TEST_F(ObjectDataTypeTest, StreamOperator) {
+    std::ostringstream oss;
+    oss << ObjectDataType::KVCACHE;
+    EXPECT_EQ(oss.str(), "KVCACHE");
+
+    oss.str("");
+    oss << ObjectDataType::UNKNOWN;
+    EXPECT_EQ(oss.str(), "UNKNOWN");
+
+    oss.str("");
+    oss << ObjectDataType::OPTIMIZER_STATE;
+    EXPECT_EQ(oss.str(), "OPTIMIZER_STATE");
+
+    oss.str("");
+    oss << ObjectDataType::GENERAL;
+    EXPECT_EQ(oss.str(), "GENERAL");
+
+    oss.str("");
+    oss << ObjectDataType::HIDDEN_STATE;
+    EXPECT_EQ(oss.str(), "HIDDEN_STATE");
+
+    // Out-of-range value should print "UNKNOWN"
+    oss.str("");
+    oss << static_cast<ObjectDataType>(200);
+    EXPECT_EQ(oss.str(), "UNKNOWN");
+}
+
+// ReplicateConfig defaults to UNKNOWN
+TEST_F(ObjectDataTypeTest, ReplicateConfigDefaultDataType) {
+    ReplicateConfig config;
+    EXPECT_EQ(config.data_type, ObjectDataType::UNKNOWN);
+}
+
+// ReplicateConfig can be set to other types
+TEST_F(ObjectDataTypeTest, ReplicateConfigSetDataType) {
+    ReplicateConfig config;
+    config.data_type = ObjectDataType::WEIGHT;
+    EXPECT_EQ(config.data_type, ObjectDataType::WEIGHT);
+}
+
+// ReplicateConfig stream output includes data_type
+TEST_F(ObjectDataTypeTest, ReplicateConfigStreamIncludesDataType) {
+    ReplicateConfig config;
+    config.data_type = ObjectDataType::TENSOR;
+    std::ostringstream oss;
+    oss << config;
+    EXPECT_NE(oss.str().find("data_type: TENSOR"), std::string::npos);
+}
+
+// PutStart with data_type propagates to ObjectMetadata
+TEST_F(ObjectDataTypeTest, PutStartWithDataType) {
+    std::unique_ptr<MasterService> service(new MasterService());
+    Segment segment = MakeSegment();
+    UUID client_id = generate_uuid();
+    auto mount_result = service->MountSegment(segment, client_id);
+    ASSERT_TRUE(mount_result.has_value());
+
+    UUID put_client = generate_uuid();
+
+    // Put with WEIGHT type
+    ReplicateConfig config;
+    config.replica_num = 1;
+    config.data_type = ObjectDataType::WEIGHT;
+
+    auto result =
+        service->PutStart(put_client, "key_weight", "default", 1024, config);
+    ASSERT_TRUE(result.has_value());
+    EXPECT_FALSE(result.value().empty());
+
+    auto end_result = service->PutEnd(put_client, "key_weight", "default",
+                                      ReplicaType::MEMORY);
+    EXPECT_TRUE(end_result.has_value());
+}
+
+// PutStart with default UNKNOWN data_type still works (backward compat)
+TEST_F(ObjectDataTypeTest, PutStartDefaultDataType) {
+    std::unique_ptr<MasterService> service(new MasterService());
+    Segment segment = MakeSegment();
+    UUID client_id = generate_uuid();
+    auto mount_result = service->MountSegment(segment, client_id);
+    ASSERT_TRUE(mount_result.has_value());
+
+    UUID put_client = generate_uuid();
+    ReplicateConfig config;
+    config.replica_num = 1;
+    // data_type left as default (UNKNOWN)
+
+    auto result =
+        service->PutStart(put_client, "key_default", "default", 1024, config);
+    ASSERT_TRUE(result.has_value());
+    EXPECT_FALSE(result.value().empty());
+}
+
+// Verify all enum values can roundtrip through uint8_t cast
+TEST_F(ObjectDataTypeTest, EnumRoundtrip) {
+    std::vector<ObjectDataType> all_types = {
+        ObjectDataType::UNKNOWN,      ObjectDataType::KVCACHE,
+        ObjectDataType::TENSOR,       ObjectDataType::WEIGHT,
+        ObjectDataType::SAMPLE,       ObjectDataType::ACTIVATION,
+        ObjectDataType::GRADIENT,     ObjectDataType::OPTIMIZER_STATE,
+        ObjectDataType::METADATA,     ObjectDataType::GENERAL,
+        ObjectDataType::HIDDEN_STATE,
+    };
+
+    for (auto type : all_types) {
+        uint8_t raw = static_cast<uint8_t>(type);
+        auto recovered = static_cast<ObjectDataType>(raw);
+        EXPECT_EQ(type, recovered);
+    }
+}
+
+}  // namespace mooncake::test
diff --git a/mooncake-store/tests/offload_on_evict_test.cpp b/mooncake-store/tests/offload_on_evict_test.cpp
index c002d285..fdb687a7 100644
--- a/mooncake-store/tests/offload_on_evict_test.cpp
+++ b/mooncake-store/tests/offload_on_evict_test.cpp
@@ -55,9 +55,11 @@ class OffloadOnEvictTest : public ::testing::Test {
                    const std::string& key, size_t size = 1024) {
         ReplicateConfig config;
         config.replica_num = 1;
-        auto put_start = service.PutStart(client_id, key, size, config);
+        auto put_start =
+            service.PutStart(client_id, key, "default", size, config);
         ASSERT_TRUE(put_start.has_value()) << "PutStart failed for key=" << key;
-        auto put_end = service.PutEnd(client_id, key, ReplicaType::MEMORY);
+        auto put_end =
+            service.PutEnd(client_id, key, "default", ReplicaType::MEMORY);
         ASSERT_TRUE(put_end.has_value()) << "PutEnd failed for key=" << key;
     }
 
@@ -68,7 +70,11 @@ class OffloadOnEvictTest : public ::testing::Test {
         if (!res) {
             return {};
         }
-        return std::move(res.value());
+        std::unordered_map<std::string, int64_t> queued;
+        for (const auto& task : res.value()) {
+            queued[task.key] = task.size;
+        }
+        return queued;
     }
 
     template <typename Predicate>
@@ -97,9 +103,11 @@ class OffloadOnEvictTest : public ::testing::Test {
             std::string key = key_prefix + std::to_string(i);
             ReplicateConfig config;
             config.replica_num = 1;
-            auto result = service.PutStart(client_id, key, object_size, config);
+            auto result = service.PutStart(client_id, key, "default",
+                                           object_size, config);
             if (result.has_value()) {
-                auto end = service.PutEnd(client_id, key, ReplicaType::MEMORY);
+                auto end = service.PutEnd(client_id, key, "default",
+                                          ReplicaType::MEMORY);
                 EXPECT_TRUE(end.has_value());
                 success_puts++;
             } else {
@@ -220,10 +228,11 @@ TEST_F(OffloadOnEvictTest, ComboB_EvictionTriggersOffload) {
         std::string key = "evict_b_" + std::to_string(i);
         ReplicateConfig config;
         config.replica_num = 1;
-        auto result =
-            service->PutStart(ctx.client_id, key, object_size, config);
+        auto result = service->PutStart(ctx.client_id, key, "default",
+                                        object_size, config);
         if (result.has_value()) {
-            auto end = service->PutEnd(ctx.client_id, key, ReplicaType::MEMORY);
+            auto end = service->PutEnd(ctx.client_id, key, "default",
+                                       ReplicaType::MEMORY);
             ASSERT_TRUE(end.has_value());
             success_puts++;
         } else {
diff --git a/mooncake-store/tests/promotion_on_hit_test.cpp b/mooncake-store/tests/promotion_on_hit_test.cpp
new file mode 100644
index 00000000..2743c4e2
--- /dev/null
+++ b/mooncake-store/tests/promotion_on_hit_test.cpp
@@ -0,0 +1,2235 @@
+// Unit tests for the L2->L1 promotion-on-hit master-side path. Exercises
+// the master-service entry points directly without going through the RPC
+// layer.
+
+#include "master_service.h"
+
+#include <glog/logging.h>
+#include <gtest/gtest.h>
+
+#include <algorithm>
+#include <chrono>
+#include <filesystem>
+#include <fstream>
+#include <map>
+#include <memory>
+#include <string>
+#include <thread>
+#include <vector>
+
+#include <unistd.h>
+
+#include "tenant_quota_policy_store.h"
+#include "types.h"
+
+namespace mooncake::test {
+
+size_t CountPromotionTask(const std::vector<PromotionTaskItem>& tasks,
+                          const std::string& key) {
+    return std::count_if(
+        tasks.begin(), tasks.end(),
+        [&key](const PromotionTaskItem& task) { return task.key == key; });
+}
+
+class PromotionOnHitTest : public ::testing::Test {
+   protected:
+    void SetUp() override {
+        google::InitGoogleLogging("PromotionOnHitTest");
+        FLAGS_logtostderr = true;
+    }
+
+    void TearDown() override {
+        for (const auto& path : policy_files_) {
+            std::error_code ec;
+            std::filesystem::remove(path, ec);
+        }
+        google::ShutdownGoogleLogging();
+    }
+
+    // Friend access to MasterService::promotion_admission_threshold_, which
+    // is otherwise private. PromotionOnHitTest is friended; TEST_F-generated
+    // subclasses are not, hence this static funnel.
+    static uint32_t GetPromotionAdmissionThresholdForTesting(
+        MasterService* service) {
+        return service->promotion_admission_threshold_;
+    }
+
+    static constexpr size_t kDefaultSegmentBase = 0x300000000;
+
+    std::string WriteTenantQuotaPolicyFile(
+        const std::map<std::string, uint64_t>& tenant_quotas) {
+        TenantQuotaPolicySnapshot snapshot;
+        snapshot.tenant_quotas = tenant_quotas;
+        auto path =
+            std::filesystem::temp_directory_path() /
+            ("mooncake_promotion_tenant_policy_" + std::to_string(::getpid()) +
+             "_" + std::to_string(next_policy_file_++) + ".yaml");
+        std::ofstream out(path);
+        out << FormatTenantQuotaPolicyYaml(snapshot);
+        out.close();
+        policy_files_.push_back(path.string());
+        return path.string();
+    }
+
+    Segment MakeSegment(std::string name, size_t base, size_t size) const {
+        Segment segment;
+        segment.id = generate_uuid();
+        segment.name = std::move(name);
+        segment.base = base;
+        segment.size = size;
+        segment.te_endpoint = segment.name;
+        return segment;
+    }
+
+    struct MountedSegmentContext {
+        UUID segment_id;
+        UUID client_id;
+        std::string segment_name;
+    };
+
+    MountedSegmentContext PrepareSegment(MasterService& service,
+                                         std::string name, size_t base,
+                                         size_t size) const {
+        Segment segment = MakeSegment(std::move(name), base, size);
+        UUID client_id = generate_uuid();
+        auto mount_result = service.MountSegment(segment, client_id);
+        EXPECT_TRUE(mount_result.has_value());
+        auto mount_ld = service.MountLocalDiskSegment(client_id, true);
+        EXPECT_TRUE(mount_ld.has_value());
+        return {.segment_id = segment.id,
+                .client_id = client_id,
+                .segment_name = segment.name};
+    }
+
+    // Put an object and complete it (creates a MEMORY replica).
+    void PutObject(MasterService& service, const UUID& client_id,
+                   const std::string& key, size_t size = 1024) {
+        ReplicateConfig config;
+        config.replica_num = 1;
+        auto put_start =
+            service.PutStart(client_id, key, "default", size, config);
+        ASSERT_TRUE(put_start.has_value()) << "PutStart failed for key=" << key;
+        auto put_end =
+            service.PutEnd(client_id, key, "default", ReplicaType::MEMORY);
+        ASSERT_TRUE(put_end.has_value()) << "PutEnd failed for key=" << key;
+    }
+
+    // Inject a synthetic LOCAL_DISK replica for `key` on `client_id`'s
+    // segment via NotifyOffloadSuccess. Lets tests put a key into
+    // LOCAL_DISK-only state without running the full offload pipeline.
+    bool InjectLocalDiskReplica(MasterService& service, const UUID& client_id,
+                                const std::string& key, int64_t size,
+                                const std::string& transport_endpoint,
+                                const std::string& tenant_id = "default") {
+        std::vector<OffloadTaskItem> tasks{
+            OffloadTaskItem{.tenant_id = tenant_id, .key = key, .size = size}};
+        StorageObjectMetadata sm;
+        sm.bucket_id = 0;
+        sm.offset = 0;
+        sm.key_size = static_cast<int64_t>(key.size());
+        sm.data_size = size;
+        sm.transport_endpoint = transport_endpoint;
+        std::vector<StorageObjectMetadata> metas{sm};
+        auto res = service.NotifyOffloadSuccess(client_id, tasks, metas);
+        return res.has_value();
+    }
+
+    // Register a client as a LOCAL_DISK holder only (no DRAM segment).
+    // This simulates the cross-host case where the LOCAL_DISK source lives
+    // on a different node than the DRAM target chosen for promotion.
+    UUID PrepareLocalDiskOnlyClient(MasterService& service) const {
+        UUID client_id = generate_uuid();
+        auto mount_ld = service.MountLocalDiskSegment(client_id, true);
+        EXPECT_TRUE(mount_ld.has_value());
+        return client_id;
+    }
+
+    std::vector<std::string> policy_files_;
+    size_t next_policy_file_ = 0;
+};
+
+// Sanity: with promotion disabled, no path mutates promotion_objects.
+TEST_F(PromotionOnHitTest, DefaultOffNoPromotion) {
+    MasterServiceConfig config;
+    config.enable_offload = true;
+    config.promotion_on_hit = false;  // explicitly off
+    config.default_kv_lease_ttl = 2000;
+    auto service = std::make_unique<MasterService>(config);
+
+    constexpr size_t seg_size = 1024 * 1024 * 16;
+    auto ctx =
+        PrepareSegment(*service, "test_segment", kDefaultSegmentBase, seg_size);
+
+    PutObject(*service, ctx.client_id, "k1");
+    // GetReplicaList many times. With promotion_on_hit=false, nothing should
+    // appear in promotion_objects regardless of access count.
+    for (int i = 0; i < 5; ++i) {
+        auto resp = service->GetReplicaList("k1", "default");
+        ASSERT_TRUE(resp.has_value());
+    }
+
+    auto pending = service->PromotionObjectHeartbeat(ctx.client_id);
+    ASSERT_TRUE(pending.has_value());
+    EXPECT_EQ(pending->size(), 0u);
+
+    service->RemoveAll();
+}
+
+// With promotion enabled but no LOCAL_DISK replica present, no promotion
+// should fire (the trigger gate any_local_disk is false).
+TEST_F(PromotionOnHitTest, NoLocalDiskNoPromotion) {
+    MasterServiceConfig config;
+    config.enable_offload = true;
+    config.promotion_on_hit = true;
+    config.promotion_admission_threshold = 1;  // promote on first touch
+    config.default_kv_lease_ttl = 2000;
+    auto service = std::make_unique<MasterService>(config);
+
+    constexpr size_t seg_size = 1024 * 1024 * 16;
+    auto ctx =
+        PrepareSegment(*service, "test_segment", kDefaultSegmentBase, seg_size);
+
+    PutObject(*service, ctx.client_id, "k_mem_only");
+    for (int i = 0; i < 5; ++i) {
+        auto resp = service->GetReplicaList("k_mem_only", "default");
+        ASSERT_TRUE(resp.has_value());
+    }
+
+    auto pending = service->PromotionObjectHeartbeat(ctx.client_id);
+    ASSERT_TRUE(pending.has_value());
+    EXPECT_EQ(pending->size(), 0u)
+        << "Memory-only key should not trigger promotion";
+
+    service->RemoveAll();
+}
+
+// Single-shot Get on a key with both MEMORY and LOCAL_DISK should not
+// promote (any_memory=true → trigger gate fails).
+TEST_F(PromotionOnHitTest, MemoryReplicaPresentNoPromotion) {
+    MasterServiceConfig config;
+    config.enable_offload = true;
+    config.promotion_on_hit = true;
+    config.promotion_admission_threshold = 1;  // first-touch
+    config.default_kv_lease_ttl = 2000;
+    auto service = std::make_unique<MasterService>(config);
+
+    constexpr size_t seg_size = 1024 * 1024 * 16;
+    auto ctx =
+        PrepareSegment(*service, "test_segment", kDefaultSegmentBase, seg_size);
+
+    PutObject(*service, ctx.client_id, "k_dual", 1024);
+    // Add a LOCAL_DISK replica next to the MEMORY one.
+    ASSERT_TRUE(InjectLocalDiskReplica(*service, ctx.client_id, "k_dual", 1024,
+                                       ctx.segment_name));
+
+    for (int i = 0; i < 5; ++i) {
+        auto resp = service->GetReplicaList("k_dual", "default");
+        ASSERT_TRUE(resp.has_value());
+    }
+
+    auto pending = service->PromotionObjectHeartbeat(ctx.client_id);
+    ASSERT_TRUE(pending.has_value());
+    EXPECT_EQ(pending->size(), 0u)
+        << "MEMORY replica still present: should not promote";
+
+    service->RemoveAll();
+}
+
+TEST_F(PromotionOnHitTest, BatchGetReplicaListPromotesLocalDiskOnlyObject) {
+    MasterServiceConfig config;
+    config.enable_offload = true;
+    config.promotion_on_hit = true;
+    config.promotion_admission_threshold = 1;
+    config.promotion_max_per_heartbeat = 2;
+    config.default_kv_lease_ttl = 2000;
+    auto service = std::make_unique<MasterService>(config);
+
+    constexpr size_t seg_size = 1024 * 1024 * 16;
+    auto ctx =
+        PrepareSegment(*service, "test_segment", kDefaultSegmentBase, seg_size);
+
+    const std::string single_key = "k_single_get_promote";
+    const std::string batch_key = "k_batch_get_promote";
+    ASSERT_TRUE(InjectLocalDiskReplica(*service, ctx.client_id, single_key,
+                                       1024, ctx.segment_name));
+    ASSERT_TRUE(InjectLocalDiskReplica(*service, ctx.client_id, batch_key, 1024,
+                                       ctx.segment_name));
+
+    auto& mm = MasterMetricManager::instance();
+    const int64_t admitted_pre = mm.get_promotion_admitted();
+    const int64_t in_flight_pre = mm.get_promotion_in_flight();
+
+    auto single_result = service->GetReplicaList(single_key, "default");
+    ASSERT_TRUE(single_result.has_value());
+    ASSERT_EQ(single_result->replicas.size(), 1u);
+    EXPECT_TRUE(single_result->replicas[0].is_local_disk_replica());
+
+    auto batch_result = service->BatchGetReplicaList(
+        std::vector<std::string>{batch_key}, "default");
+    ASSERT_EQ(batch_result.size(), 1u);
+    ASSERT_TRUE(batch_result[0].has_value());
+    ASSERT_EQ(batch_result[0]->replicas.size(), 1u);
+    EXPECT_TRUE(batch_result[0]->replicas[0].is_local_disk_replica());
+
+    EXPECT_EQ(mm.get_promotion_admitted() - admitted_pre, 2);
+    EXPECT_EQ(mm.get_promotion_in_flight() - in_flight_pre, 2);
+
+    auto pending = service->PromotionObjectHeartbeat(ctx.client_id);
+    ASSERT_TRUE(pending.has_value());
+    EXPECT_EQ(pending->size(), 2u);
+    EXPECT_EQ(CountPromotionTask(*pending, single_key), 1u);
+    EXPECT_EQ(CountPromotionTask(*pending, batch_key), 1u);
+
+    service->RemoveAll();
+}
+
+// PromotionObjectHeartbeat returns an empty task list when called against a
+// client that has no LocalDiskSegment registered.
+TEST_F(PromotionOnHitTest, HeartbeatReturnsErrorForUnknownClient) {
+    MasterServiceConfig config;
+    config.enable_offload = true;
+    config.promotion_on_hit = true;
+    auto service = std::make_unique<MasterService>(config);
+
+    UUID unknown_client = generate_uuid();
+    auto pending = service->PromotionObjectHeartbeat(unknown_client);
+    ASSERT_FALSE(pending.has_value());
+    EXPECT_EQ(pending.error(), ErrorCode::SEGMENT_NOT_FOUND);
+}
+
+// PromotionAllocStart on a non-existent key returns OBJECT_NOT_FOUND.
+TEST_F(PromotionOnHitTest, AllocStartUnknownKey) {
+    MasterServiceConfig config;
+    config.enable_offload = true;
+    config.promotion_on_hit = true;
+    auto service = std::make_unique<MasterService>(config);
+
+    auto resp = service->PromotionAllocStart(generate_uuid(), "nonexistent",
+                                             "default", 1024, {});
+    ASSERT_FALSE(resp.has_value());
+    EXPECT_EQ(resp.error(), ErrorCode::OBJECT_NOT_FOUND);
+}
+
+// NotifyPromotionSuccess on a non-existent key returns OBJECT_NOT_FOUND.
+TEST_F(PromotionOnHitTest, NotifyUnknownKey) {
+    MasterServiceConfig config;
+    config.enable_offload = true;
+    config.promotion_on_hit = true;
+    auto service = std::make_unique<MasterService>(config);
+
+    UUID client_id = generate_uuid();
+    auto resp =
+        service->NotifyPromotionSuccess(client_id, "nonexistent", "default");
+    ASSERT_FALSE(resp.has_value());
+    EXPECT_EQ(resp.error(), ErrorCode::OBJECT_NOT_FOUND);
+}
+
+// Concurrent readers racing into TryPushPromotionQueue must dedupe to a
+// single PromotionTask. Without dedup, the source LOCAL_DISK replica's
+// refcnt would be incremented N times and the per-segment promotion_objects
+// map would either error on the second insert (current behavior) or
+// double-queue.
+TEST_F(PromotionOnHitTest, RacingReadersDedup) {
+    MasterServiceConfig config;
+    config.enable_offload = true;
+    config.promotion_on_hit = true;
+    config.promotion_admission_threshold = 1;  // first-touch fires the gate
+    config.default_kv_lease_ttl = 2000;
+    auto service = std::make_unique<MasterService>(config);
+
+    constexpr size_t seg_size = 1024 * 1024 * 16;
+    auto ctx =
+        PrepareSegment(*service, "test_segment", kDefaultSegmentBase, seg_size);
+
+    // LOCAL_DISK-only key: NotifyOffloadSuccess on a never-PUT key creates
+    // metadata with only the LOCAL_DISK replica (AddReplica's Create-on-
+    // missing path).
+    ASSERT_TRUE(InjectLocalDiskReplica(*service, ctx.client_id, "k_cold", 1024,
+                                       ctx.segment_name));
+
+    constexpr int kThreads = 32;
+    constexpr int kReadsPerThread = 10;
+    std::vector<std::thread> threads;
+    threads.reserve(kThreads);
+    for (int t = 0; t < kThreads; ++t) {
+        threads.emplace_back([&service]() {
+            for (int j = 0; j < kReadsPerThread; ++j) {
+                auto r = service->GetReplicaList("k_cold", "default");
+                EXPECT_TRUE(r.has_value());
+            }
+        });
+    }
+    for (auto& t : threads) t.join();
+
+    auto pending = service->PromotionObjectHeartbeat(ctx.client_id);
+    ASSERT_TRUE(pending.has_value());
+    EXPECT_EQ(pending->size(), 1u)
+        << "Concurrent readers (" << kThreads << " x " << kReadsPerThread
+        << ") must dedupe to a single promotion task";
+    ASSERT_TRUE(CountPromotionTask(*pending, "k_cold"));
+
+    service->RemoveAll();
+}
+
+// A PromotionTask that never receives NotifyPromotionSuccess must be
+// reaped after `put_start_release_timeout_sec` so that:
+//   (a) the source LOCAL_DISK replica's refcnt is decremented (no
+//       permanent pin → eviction can still free it later);
+//   (b) the dedup gate is unblocked, so a subsequent read can re-enqueue
+//       the same key for retry.
+TEST_F(PromotionOnHitTest, StalePromotionReaper) {
+    MasterServiceConfig config;
+    config.enable_offload = true;
+    config.promotion_on_hit = true;
+    config.promotion_admission_threshold = 1;
+    config.default_kv_lease_ttl = 2000;
+    // Short staleness window: the reaper uses put_start_release_timeout_sec
+    // for promotion tasks. The master enforces release > discard, so set
+    // discard to a still-smaller value.
+    config.put_start_discard_timeout_sec = 0;
+    config.put_start_release_timeout_sec = 1;
+    auto service = std::make_unique<MasterService>(config);
+
+    constexpr size_t seg_size = 1024 * 1024 * 16;
+    auto ctx =
+        PrepareSegment(*service, "test_segment", kDefaultSegmentBase, seg_size);
+
+    ASSERT_TRUE(InjectLocalDiskReplica(*service, ctx.client_id, "k_cold", 1024,
+                                       ctx.segment_name));
+
+    auto& mm = MasterMetricManager::instance();
+    const int64_t expired_pre = mm.get_promotion_expired();
+
+    // Trigger #1: enqueue, then drain the per-segment queue. Drain leaves
+    // the per-shard PromotionTask intact (the heartbeat is best-effort GC,
+    // not the authoritative state).
+    {
+        auto r = service->GetReplicaList("k_cold", "default");
+        ASSERT_TRUE(r.has_value());
+        auto pending = service->PromotionObjectHeartbeat(ctx.client_id);
+        ASSERT_TRUE(pending.has_value());
+        EXPECT_EQ(pending->size(), 1u);
+    }
+
+    // Without reap, dedup blocks re-enqueue. Confirm: GetReplicaList again,
+    // heartbeat must be empty because PromotionTask still pins the slot.
+    {
+        auto r = service->GetReplicaList("k_cold", "default");
+        ASSERT_TRUE(r.has_value());
+        auto pending = service->PromotionObjectHeartbeat(ctx.client_id);
+        ASSERT_TRUE(pending.has_value());
+        EXPECT_EQ(pending->size(), 0u)
+            << "Dedup gate should block re-enqueue while task is in flight";
+    }
+
+    // Wait past the staleness window. The eviction thread reaps every
+    // kEvictionThreadSleepMs (10 ms) but gated by
+    // `now - last_discard_time > put_start_release_timeout_sec_` (strict).
+    // With release=1s a 2s sleep leaves only ~1s margin between reaper
+    // firing and the wake-up, which has flaked under CI load. 3s gives
+    // ~2s of margin and matches the schedule's strict-greater comparison.
+    std::this_thread::sleep_for(std::chrono::seconds(3));
+
+    // Trigger #3: with the task reaped, dedup is unblocked and a fresh
+    // GetReplicaList must enqueue again.
+    {
+        auto r = service->GetReplicaList("k_cold", "default");
+        ASSERT_TRUE(r.has_value());
+        auto pending = service->PromotionObjectHeartbeat(ctx.client_id);
+        ASSERT_TRUE(pending.has_value());
+        EXPECT_EQ(pending->size(), 1u)
+            << "After reap, a fresh read must re-enqueue the same key";
+    }
+
+    EXPECT_EQ(mm.get_promotion_expired() - expired_pre, 1)
+        << "Reaper expiry must bump promotion_expired";
+
+    service->RemoveAll();
+}
+
+// Force-Remove on a key with a queued PromotionTask must not corrupt
+// state. The PromotionTask references a Replica*; once the metadata is
+// gone, NotifyPromotionSuccess and the reaper must both tolerate the
+// missing entry.
+TEST_F(PromotionOnHitTest, RemoveDuringPromotion) {
+    MasterServiceConfig config;
+    config.enable_offload = true;
+    config.promotion_on_hit = true;
+    config.promotion_admission_threshold = 1;
+    config.default_kv_lease_ttl = 100;  // short lease: Remove won't block
+    config.put_start_discard_timeout_sec = 0;
+    config.put_start_release_timeout_sec = 1;
+    auto service = std::make_unique<MasterService>(config);
+
+    constexpr size_t seg_size = 1024 * 1024 * 16;
+    auto ctx =
+        PrepareSegment(*service, "test_segment", kDefaultSegmentBase, seg_size);
+
+    ASSERT_TRUE(InjectLocalDiskReplica(*service, ctx.client_id, "k_cold", 1024,
+                                       ctx.segment_name));
+
+    auto& mm = MasterMetricManager::instance();
+    const int64_t cancelled_pre = mm.get_promotion_cancelled();
+
+    // Queue a promotion task.
+    {
+        auto r = service->GetReplicaList("k_cold", "default");
+        ASSERT_TRUE(r.has_value());
+    }
+
+    // Lease must expire before we can call non-force Remove (or use force).
+    auto rm = service->Remove("k_cold", "default", /*force=*/true);
+    // Remove returns REPLICA_IS_NOT_READY if any replica is non-COMPLETE.
+    // The injected LOCAL_DISK replica is COMPLETE, so this should succeed.
+    ASSERT_TRUE(rm.has_value())
+        << "Remove on a LOCAL_DISK-only key with a queued promotion should "
+        << "succeed (all replicas COMPLETE); error=" << rm.error();
+    EXPECT_EQ(mm.get_promotion_cancelled() - cancelled_pre, 1)
+        << "Remove of a key with an in-flight promotion task must bump "
+        << "promotion_cancelled";
+
+    // NotifyPromotionSuccess on the now-removed key must surface the missing
+    // metadata cleanly, not crash.
+    auto notify =
+        service->NotifyPromotionSuccess(ctx.client_id, "k_cold", "default");
+    ASSERT_FALSE(notify.has_value());
+    EXPECT_EQ(notify.error(), ErrorCode::OBJECT_NOT_FOUND);
+
+    // Wait for the reaper; it must tolerate the missing metadata entry
+    // (the source replica it would dec_refcnt is already gone). 3s for
+    // CI-safe margin over the 1s release timeout.
+    std::this_thread::sleep_for(std::chrono::seconds(3));
+
+    // Re-injecting the key and re-triggering must work end-to-end, proving
+    // the per-shard PromotionTask was reaped (not stuck).
+    ASSERT_TRUE(InjectLocalDiskReplica(*service, ctx.client_id, "k_cold", 1024,
+                                       ctx.segment_name));
+    {
+        auto r = service->GetReplicaList("k_cold", "default");
+        ASSERT_TRUE(r.has_value());
+        auto pending = service->PromotionObjectHeartbeat(ctx.client_id);
+        ASSERT_TRUE(pending.has_value());
+        EXPECT_EQ(pending->size(), 1u)
+            << "After Remove + reap, the same key must re-enqueue cleanly";
+    }
+
+    service->RemoveAll();
+}
+
+// Cross-host promotion: the LOCAL_DISK source's holder client has no DRAM
+// segment of its own. The reader's DRAM segment (segment_b) is the only
+// allocator candidate, so PromotionAllocStart must place the new MEMORY
+// replica there — proving the master's allocation strategy crosses the
+// host boundary cleanly.
+TEST_F(PromotionOnHitTest, MultiSegmentAllocPicksAvailableSegment) {
+    MasterServiceConfig config;
+    config.enable_offload = true;
+    config.promotion_on_hit = true;
+    config.promotion_admission_threshold = 1;
+    config.default_kv_lease_ttl = 2000;
+    auto service = std::make_unique<MasterService>(config);
+
+    // Holder: LOCAL_DISK only, no DRAM. Stands in for "host A" that has the
+    // SSD copy but no free RAM.
+    UUID holder_client_id = PrepareLocalDiskOnlyClient(*service);
+
+    // Reader-side DRAM target: "host B" with a fresh DRAM segment.
+    constexpr size_t seg_size = 1024 * 1024 * 16;
+    auto reader_ctx =
+        PrepareSegment(*service, "segment_b", kDefaultSegmentBase, seg_size);
+
+    // The cold key: only replica is LOCAL_DISK on the holder client.
+    ASSERT_TRUE(InjectLocalDiskReplica(*service, holder_client_id, "k_cold",
+                                       1024, "segment_a_endpoint"));
+
+    // Seed the PromotionTask through the gate — PromotionAllocStart now
+    // requires an in-flight task to exist (rejects orphaned-stage path).
+    {
+        auto r = service->GetReplicaList("k_cold", "default");
+        ASSERT_TRUE(r.has_value());
+    }
+
+    // PromotionAllocStart — test the segment-selection logic on top of
+    // the gate-seeded task.
+    auto resp = service->PromotionAllocStart(holder_client_id, "k_cold",
+                                             "default", 1024, {});
+    ASSERT_TRUE(resp.has_value())
+        << "PromotionAllocStart should succeed when any DRAM segment has "
+        << "capacity; error=" << resp.error();
+
+    const auto& mem_desc = resp.value().memory_descriptor;
+    ASSERT_TRUE(mem_desc.is_memory_replica());
+    EXPECT_EQ(
+        mem_desc.get_memory_descriptor().buffer_descriptor.transport_endpoint_,
+        reader_ctx.segment_name)
+        << "New MEMORY replica must be allocated on the only DRAM segment "
+        << "(segment_b), not on the LOCAL_DISK holder which has no DRAM";
+
+    // Sanity: the staged MEMORY replica is PROCESSING (visible only after
+    // NotifyPromotionSuccess flips it COMPLETE).
+    EXPECT_EQ(mem_desc.status, ReplicaStatus::PROCESSING);
+
+    service->RemoveAll();
+}
+
+// preferred_segments is honored: promotion can be steered to a specific
+// DRAM segment (e.g., the reader's local one). TryPushPromotionQueue
+// currently doesn't pass preferred_segments through, but
+// PromotionAllocStart respects them for clients biasing toward locality.
+TEST_F(PromotionOnHitTest, MultiSegmentAllocRespectsPreferred) {
+    MasterServiceConfig config;
+    config.enable_offload = true;
+    config.promotion_on_hit = true;
+    config.promotion_admission_threshold = 1;
+    auto service = std::make_unique<MasterService>(config);
+
+    constexpr size_t seg_size = 1024 * 1024 * 16;
+    auto seg_a =
+        PrepareSegment(*service, "segment_a", kDefaultSegmentBase, seg_size);
+    auto seg_b = PrepareSegment(*service, "segment_b",
+                                kDefaultSegmentBase + seg_size, seg_size);
+
+    ASSERT_TRUE(InjectLocalDiskReplica(*service, seg_a.client_id, "k_cold",
+                                       1024, seg_a.segment_name));
+
+    // Seed the PromotionTask through the gate so AllocStart's
+    // task-existence check passes.
+    {
+        auto r = service->GetReplicaList("k_cold", "default");
+        ASSERT_TRUE(r.has_value());
+    }
+
+    auto resp = service->PromotionAllocStart(
+        seg_a.client_id, "k_cold", "default", 1024, {seg_b.segment_name});
+    ASSERT_TRUE(resp.has_value());
+    const auto& mem_desc = resp.value().memory_descriptor;
+    EXPECT_EQ(
+        mem_desc.get_memory_descriptor().buffer_descriptor.transport_endpoint_,
+        seg_b.segment_name)
+        << "preferred_segments={segment_b} should pin the new MEMORY "
+        << "replica to segment_b";
+
+    service->RemoveAll();
+}
+
+// promotion_queue_limit caps total in-flight tasks cluster-wide
+// (gate: promotion_in_flight_ >= limit). With limit=1 the very first
+// queued task saturates the cap, so a second LOCAL_DISK-only read —
+// even on a key in the same shard — must be silently dropped by the
+// cap gate (reads still succeed; just no new task is enqueued).
+// QueueLimitRejectsCrossShard covers the same-cap-across-different-
+// shards case.
+TEST_F(PromotionOnHitTest, QueueLimitRejectsBeyondCap) {
+    MasterServiceConfig config;
+    config.enable_offload = true;
+    config.promotion_on_hit = true;
+    config.promotion_admission_threshold = 1;
+    config.promotion_queue_limit = 1;  // any 1 task saturates a shard
+    config.default_kv_lease_ttl = 2000;
+    auto service = std::make_unique<MasterService>(config);
+
+    constexpr size_t seg_size = 1024 * 1024 * 16;
+    auto seg = PrepareSegment(*service, "seg_a", kDefaultSegmentBase, seg_size);
+
+    // Find two keys that hash to the same shard. MasterService::
+    // getShardIndex is private but the formula is deterministic
+    // (std::hash<std::string>{}(key) % kNumShards), so we can mirror
+    // it here. kNumShards=1024 (master_service.h:889).
+    constexpr size_t kNumShardsLocal = 1024;
+    auto shard_of = [](const std::string& k) {
+        return std::hash<std::string>{}(k) % kNumShardsLocal;
+    };
+    const std::string k1 = "qlim_first";
+    std::string k2;
+    for (int i = 0; i < 100000 && k2.empty(); ++i) {
+        std::string candidate = "qlim_collide_" + std::to_string(i);
+        if (shard_of(candidate) == shard_of(k1)) {
+            k2 = candidate;
+        }
+    }
+    ASSERT_FALSE(k2.empty())
+        << "could not find a same-shard collision for " << k1;
+
+    ASSERT_TRUE(InjectLocalDiskReplica(*service, seg.client_id, k1, 1024,
+                                       seg.segment_name));
+    ASSERT_TRUE(InjectLocalDiskReplica(*service, seg.client_id, k2, 1024,
+                                       seg.segment_name));
+
+    auto& mm = MasterMetricManager::instance();
+    const int64_t cap_rej_pre = mm.get_promotion_rejected_cap();
+
+    // First read on k1 enqueues a task in shard S.
+    auto r1 = service->GetReplicaList(k1, "default");
+    ASSERT_TRUE(r1.has_value());
+
+    // Second read on k2 (same shard S, different key, so no dedup) must
+    // be dropped by the cap gate: the cluster-wide in-flight counter is
+    // already 1, which meets promotion_queue_limit_ = 1.
+    auto r2 = service->GetReplicaList(k2, "default");
+    ASSERT_TRUE(r2.has_value()) << "read itself must still succeed; "
+                                << "queue gate is silent";
+
+    // Drain the holder's queue: only k1 should appear.
+    auto heartbeat = service->PromotionObjectHeartbeat(seg.client_id);
+    ASSERT_TRUE(heartbeat.has_value());
+    EXPECT_EQ(heartbeat->size(), 1u)
+        << "promotion_queue_limit=1 should admit only the first task "
+        << "globally; k2's enqueue must be dropped";
+    EXPECT_EQ(CountPromotionTask(*heartbeat, k1), 1u)
+        << "k1 was read first and should be the surviving task";
+    EXPECT_EQ(CountPromotionTask(*heartbeat, k2), 0u)
+        << "k2 was rejected by the cap gate; should not appear";
+    EXPECT_EQ(mm.get_promotion_rejected_cap() - cap_rej_pre, 1)
+        << "k2's rejection must increment promotion_rejected_cap";
+
+    service->RemoveAll();
+}
+
+// PromotionObjectHeartbeat caps the per-call response at kMaxPerHeartbeat
+// (1) and leaves the remainder queued for subsequent heartbeats. This is
+// the master-side bound that keeps the client's heartbeat thread inside
+// the liveness window when many keys are queued, while guaranteeing no
+// task is dropped — a regression from the prior client-side cap that
+// silently discarded leftovers.
+TEST_F(PromotionOnHitTest, HeartbeatBoundedBatchPreservesLeftovers) {
+    MasterServiceConfig config;
+    config.enable_offload = true;
+    config.promotion_on_hit = true;
+    config.promotion_admission_threshold = 1;
+    config.default_kv_lease_ttl = 2000;
+    auto service = std::make_unique<MasterService>(config);
+
+    constexpr size_t seg_size = 1024 * 1024 * 16;
+    auto seg = PrepareSegment(*service, "seg_a", kDefaultSegmentBase, seg_size);
+
+    // Push three keys into the holder's promotion_objects via the trigger
+    // path. Each is a LOCAL_DISK-only key (no MEMORY replica), so the
+    // first GetReplicaList per key crosses the admission threshold (=1)
+    // and TryPushPromotionQueue enqueues a task.
+    const std::vector<std::string> keys{"hb_k1", "hb_k2", "hb_k3"};
+    for (const auto& k : keys) {
+        ASSERT_TRUE(InjectLocalDiskReplica(*service, seg.client_id, k, 1024,
+                                           seg.segment_name));
+        auto r = service->GetReplicaList(k, "default");
+        ASSERT_TRUE(r.has_value());
+    }
+
+    // First heartbeat returns at most 1 key.
+    auto tick1 = service->PromotionObjectHeartbeat(seg.client_id);
+    ASSERT_TRUE(tick1.has_value());
+    EXPECT_EQ(tick1->size(), 1u)
+        << "heartbeat should return at most kMaxPerHeartbeat=1 entry";
+    std::string first_key = tick1->begin()->key;
+    EXPECT_TRUE(std::find(keys.begin(), keys.end(), first_key) != keys.end());
+
+    // Second heartbeat returns another key (a different one).
+    auto tick2 = service->PromotionObjectHeartbeat(seg.client_id);
+    ASSERT_TRUE(tick2.has_value());
+    EXPECT_EQ(tick2->size(), 1u);
+    std::string second_key = tick2->begin()->key;
+    EXPECT_NE(second_key, first_key)
+        << "second heartbeat must drain a different leftover key, not "
+        << "re-return the one already extracted";
+
+    // Third heartbeat returns the third key.
+    auto tick3 = service->PromotionObjectHeartbeat(seg.client_id);
+    ASSERT_TRUE(tick3.has_value());
+    EXPECT_EQ(tick3->size(), 1u);
+    std::string third_key = tick3->begin()->key;
+    EXPECT_NE(third_key, first_key);
+    EXPECT_NE(third_key, second_key);
+
+    // Fourth heartbeat: queue is now empty.
+    auto tick4 = service->PromotionObjectHeartbeat(seg.client_id);
+    ASSERT_TRUE(tick4.has_value());
+    EXPECT_TRUE(tick4->empty())
+        << "after draining all queued keys, heartbeat must return empty";
+
+    // Sanity: master-side promotion_tasks records are intact for all keys
+    // (they're cleared by NotifyPromotionSuccess, not by Heartbeat), so the
+    // source refcnts remain pinned until processed.
+    for (const auto& k : keys) {
+        auto rl = service->GetReplicaList(k, "default");
+        ASSERT_TRUE(rl.has_value()) << "key " << k << " should still exist";
+    }
+
+    service->RemoveAll();
+}
+
+// The promotion task reaper must pop the staged PROCESSING MEMORY
+// replica added by PromotionAllocStart. The staged replica is not in
+// shard->processing_keys, so DiscardExpiredProcessingReplicas's main
+// sweep can't see it, and the reaper for promotion tasks is the only
+// place that knows the replica exists. Without this path the orphan
+// holds its allocator buffer indefinitely (until the object is removed
+// or evicted).
+//
+// We can't observe the staged replica via GetReplicaList because the
+// master filters out PROCESSING entries (clients can only read COMPLETE
+// replicas), so we use QuerySegments to watch the DRAM allocator's used
+// bytes: AllocStart bumps it, and the reaper must return it to baseline.
+// NotifyPromotionSuccess on a reaped task must also fail cleanly.
+TEST_F(PromotionOnHitTest, ReaperPopsStagedMemoryReplicaOnExpiry) {
+    MasterServiceConfig config;
+    config.enable_offload = true;
+    config.promotion_on_hit = true;
+    config.promotion_admission_threshold = 1;
+    config.default_kv_lease_ttl = 2000;
+    config.put_start_discard_timeout_sec = 0;
+    config.put_start_release_timeout_sec = 1;
+    auto service = std::make_unique<MasterService>(config);
+
+    constexpr size_t seg_size = 1024 * 1024 * 16;
+    auto ctx = PrepareSegment(*service, "seg_a", kDefaultSegmentBase, seg_size);
+    ASSERT_TRUE(InjectLocalDiskReplica(*service, ctx.client_id, "k_cold", 1024,
+                                       ctx.segment_name));
+
+    // Baseline allocator usage on the DRAM segment.
+    auto seg_baseline = service->QuerySegments(ctx.segment_name);
+    ASSERT_TRUE(seg_baseline.has_value());
+    const size_t used_baseline = seg_baseline->first;
+
+    // Trigger the gate to enqueue a PromotionTask.
+    {
+        auto r = service->GetReplicaList("k_cold", "default");
+        ASSERT_TRUE(r.has_value());
+    }
+    // Drive the AllocStart side so alloc_id != 0 — this is the exact
+    // setup that produces an orphaned PROCESSING MEMORY replica if the
+    // reaper does not pop it.
+    auto alloc = service->PromotionAllocStart(ctx.client_id, "k_cold",
+                                              "default", 1024, {});
+    ASSERT_TRUE(alloc.has_value());
+
+    // After AllocStart, the DRAM allocator must have committed bytes for
+    // the staged PROCESSING MEMORY replica.
+    auto seg_after_alloc = service->QuerySegments(ctx.segment_name);
+    ASSERT_TRUE(seg_after_alloc.has_value());
+    EXPECT_GT(seg_after_alloc->first, used_baseline)
+        << "PromotionAllocStart should bump segment used bytes "
+        << "(allocator-tracked PROCESSING MEMORY replica)";
+
+    // Wait past the staleness window; the eviction thread reaps the
+    // task and (with the fix) pops the staged replica via
+    // EraseReplicaByID, which releases the buffer back to the allocator.
+    // Poll instead of a fixed sleep — the eviction thread cadence and
+    // the test runner's scheduling jitter (especially under suite load)
+    // both vary, so a hard 2s sleep is flaky here.
+    constexpr auto kPollDeadline = std::chrono::seconds(5);
+    constexpr auto kPollInterval = std::chrono::milliseconds(50);
+    const auto poll_start = std::chrono::steady_clock::now();
+    size_t used_after_reap = 0;
+    while (std::chrono::steady_clock::now() - poll_start < kPollDeadline) {
+        auto q = service->QuerySegments(ctx.segment_name);
+        ASSERT_TRUE(q.has_value());
+        used_after_reap = q->first;
+        if (used_after_reap == used_baseline) break;
+        std::this_thread::sleep_for(kPollInterval);
+    }
+    EXPECT_EQ(used_after_reap, used_baseline)
+        << "after reap: staged PROCESSING MEMORY replica's buffer must "
+        << "be freed back to the DRAM allocator. If this fires, the "
+        << "reaper is not popping the staged replica and the buffer "
+        << "leaks until the object itself is removed or evicted.";
+
+    // NotifyPromotionSuccess for a reaped task must not commit anything
+    // and must return REPLICA_IS_NOT_READY (the task entry is gone, so
+    // the alloc_id lookup at the top of NotifyPromotionSuccess fails
+    // fast).
+    auto notify =
+        service->NotifyPromotionSuccess(ctx.client_id, "k_cold", "default");
+    ASSERT_FALSE(notify.has_value());
+    EXPECT_EQ(notify.error(), ErrorCode::REPLICA_IS_NOT_READY);
+
+    service->RemoveAll();
+}
+
+// The cap gate must be cluster-wide, not per-shard. Promotion targets
+// skewed hot keys, which by definition cluster into a small number of
+// shards; a `shard->size() * kNumShards >= limit` heuristic fires
+// roughly kNumShards-times too eagerly on that workload. With a global
+// atomic counter, a task in shard A counts toward the cap that gates a
+// task in shard B.
+TEST_F(PromotionOnHitTest, QueueLimitRejectsCrossShard) {
+    MasterServiceConfig config;
+    config.enable_offload = true;
+    config.promotion_on_hit = true;
+    config.promotion_admission_threshold = 1;
+    config.promotion_queue_limit = 1;  // 1 in-flight task globally
+    config.default_kv_lease_ttl = 2000;
+    auto service = std::make_unique<MasterService>(config);
+
+    constexpr size_t seg_size = 1024 * 1024 * 16;
+    auto seg = PrepareSegment(*service, "seg_a", kDefaultSegmentBase, seg_size);
+
+    // Find two keys hashing to *different* shards. With the old per-shard
+    // heuristic this would let both through (each shard's count is 0
+    // independently). With the global counter, only the first goes in.
+    constexpr size_t kNumShardsLocal = 1024;
+    auto shard_of = [](const std::string& k) {
+        return std::hash<std::string>{}(k) % kNumShardsLocal;
+    };
+    const std::string k1 = "xshard_first";
+    std::string k2;
+    for (int i = 0; i < 100000 && k2.empty(); ++i) {
+        std::string candidate = "xshard_other_" + std::to_string(i);
+        if (shard_of(candidate) != shard_of(k1)) {
+            k2 = candidate;
+        }
+    }
+    ASSERT_FALSE(k2.empty()) << "couldn't find a different-shard key";
+    ASSERT_NE(shard_of(k1), shard_of(k2));
+
+    ASSERT_TRUE(InjectLocalDiskReplica(*service, seg.client_id, k1, 1024,
+                                       seg.segment_name));
+    ASSERT_TRUE(InjectLocalDiskReplica(*service, seg.client_id, k2, 1024,
+                                       seg.segment_name));
+
+    auto r1 = service->GetReplicaList(k1, "default");
+    ASSERT_TRUE(r1.has_value());
+
+    // k2 lives in a different shard, but the global cap is already met
+    // by k1's task — k2 must be rejected.
+    auto r2 = service->GetReplicaList(k2, "default");
+    ASSERT_TRUE(r2.has_value()) << "read itself still succeeds";
+
+    auto heartbeat = service->PromotionObjectHeartbeat(seg.client_id);
+    ASSERT_TRUE(heartbeat.has_value());
+    EXPECT_EQ(heartbeat->size(), 1u)
+        << "with global cap=1 and one task already in shard " << shard_of(k1)
+        << ", a key hashing to shard " << shard_of(k2)
+        << " must be rejected by the global gate. A per-shard heuristic "
+        << "would admit it here since the destination shard's local "
+        << "count is 0.";
+    EXPECT_EQ(CountPromotionTask(*heartbeat, k1), 1u);
+    EXPECT_EQ(CountPromotionTask(*heartbeat, k2), 0u);
+
+    service->RemoveAll();
+}
+
+// PromotionAllocStart must reset PromotionTask.start_time so the reaper
+// TTL covers the active-transfer phase on its own and is not consumed by
+// the queue-wait phase. Without the reset, a task that waited in the
+// holder's promotion_objects queue for most of the original TTL could
+// enter active transfer with little budget left; if the SSD read + RDMA
+// write of a large object then ran past expiry, the reaper would
+// EraseReplicaByID on the staged MEMORY replica mid-flight and the
+// allocator could hand the freed buffer to a concurrent Put while the
+// client's RDMA write is still landing into it (use-after-free in the
+// allocator + silent data corruption for the concurrent Put).
+//
+// We can't directly observe start_time from a black-box test, so we
+// arrange the timing so the reset is the only thing that distinguishes
+// "task still alive" from "task reaped" at the assertion points:
+//
+//   T=0    : admit task   (original start_time = T=0)
+//   T=Wq   : AllocStart   (resets start_time = T=Wq)
+//   T=Wq+Wa: assertion 1  -- alive iff the reset happened
+//                            (Wq + Wa > TTL,  so without reset the
+//                             original start_time has aged past TTL)
+//                            (Wa < TTL,        so with the reset the
+//                             new start_time has NOT aged past TTL)
+//   T=Wq+Wb: assertion 2  -- reaped
+//                            (Wb > TTL, so even with the reset the
+//                             active-transfer phase has now exceeded its
+//                             own full window)
+//
+// Concretely with TTL = 2s, Wq = 1.5s, Wa = 1.5s (-> 3.0s elapsed,
+// 1.5s since AllocStart), Wb = 3.0s (-> 4.5s elapsed, 3.0s since
+// AllocStart). The "alive" assertion at Wq+Wa is the one that proves
+// the reset is wired up correctly.
+//
+// QuerySegments(seg).first (used bytes) is the observable: AllocStart
+// bumps it, the reaper's EraseReplicaByID returns it to baseline. We
+// can't use GetReplicaList because the master filters out PROCESSING
+// replicas from the response.
+TEST_F(PromotionOnHitTest, AllocStartResetsTaskDeadline) {
+    MasterServiceConfig config;
+    config.enable_offload = true;
+    config.promotion_on_hit = true;
+    config.promotion_admission_threshold = 1;
+    config.default_kv_lease_ttl = 8000;
+    config.put_start_discard_timeout_sec = 0;
+    config.put_start_release_timeout_sec = 2;
+    auto service = std::make_unique<MasterService>(config);
+
+    constexpr size_t seg_size = 1024 * 1024 * 16;
+    auto ctx = PrepareSegment(*service, "seg_a", kDefaultSegmentBase, seg_size);
+    ASSERT_TRUE(InjectLocalDiskReplica(*service, ctx.client_id, "k_late", 1024,
+                                       ctx.segment_name));
+
+    auto seg_baseline = service->QuerySegments(ctx.segment_name);
+    ASSERT_TRUE(seg_baseline.has_value());
+    const size_t used_baseline = seg_baseline->first;
+
+    // T=0 : admit. start_time = T=0.
+    {
+        auto r = service->GetReplicaList("k_late", "default");
+        ASSERT_TRUE(r.has_value());
+    }
+
+    // T=Wq : simulate the holder's queue having been backlogged. With
+    // TTL = 2s and Wq = 1.5s the task is still alive at AllocStart
+    // (the queue-wait phase's own window hasn't expired yet — 1.5 < 2).
+    std::this_thread::sleep_for(std::chrono::milliseconds(1500));
+
+    auto alloc = service->PromotionAllocStart(ctx.client_id, "k_late",
+                                              "default", 1024, {});
+    ASSERT_TRUE(alloc.has_value())
+        << "AllocStart must succeed before the queue-wait phase's TTL "
+        << "expires (1.5s elapsed, TTL is 2s). If this fires the test "
+        << "timing is wrong, not the feature.";
+
+    // Buffer has been staged.
+    auto seg_after_alloc = service->QuerySegments(ctx.segment_name);
+    ASSERT_TRUE(seg_after_alloc.has_value());
+    EXPECT_GT(seg_after_alloc->first, used_baseline)
+        << "PromotionAllocStart should bump segment used bytes "
+        << "(allocator-tracked PROCESSING MEMORY replica)";
+
+    // T=Wq+Wa : total elapsed since admission is 3.0s > TTL=2s. If the
+    // reset is missing, the reaper sees the original start_time = 0
+    // and expires the task here, calling EraseReplicaByID which would
+    // free the staged buffer mid-RDMA-write in production. With the
+    // reset, start_time was bumped at AllocStart so age-since-reset is
+    // only 1.5s < TTL=2s and the task stays alive. This is the
+    // assertion that proves the reset is in place.
+    std::this_thread::sleep_for(std::chrono::milliseconds(1500));
+    auto seg_after_wait = service->QuerySegments(ctx.segment_name);
+    ASSERT_TRUE(seg_after_wait.has_value());
+    EXPECT_GT(seg_after_wait->first, used_baseline)
+        << "post-AllocStart staged MEMORY buffer must still be alive: "
+        << "queue-wait + active-transfer wall time (3.0s) has exceeded "
+        << "the bare TTL (2s), but the start_time reset at AllocStart "
+        << "gives the active-transfer phase its own fresh TTL window. "
+        << "If this fires, the reset in PromotionAllocStart is missing "
+        << "or broken — without it the reaper frees the staged buffer "
+        << "here while a client's RDMA write could still be in flight.";
+
+    // T=Wq+Wb : sleep another 3s for a total of 4.5s since AllocStart.
+    // The active-transfer phase's own TTL has now expired regardless
+    // of the reset; the reaper must fire and free the staged buffer.
+    std::this_thread::sleep_for(std::chrono::seconds(3));
+    auto seg_after_reap = service->QuerySegments(ctx.segment_name);
+    ASSERT_TRUE(seg_after_reap.has_value());
+    EXPECT_EQ(seg_after_reap->first, used_baseline)
+        << "after the active-transfer phase's own TTL expires "
+        << "(4.5s > 2s since AllocStart), the reaper must fire and "
+        << "free the staged buffer via EraseReplicaByID";
+
+    service->RemoveAll();
+}
+
+// NotifyPromotionSuccess must decrement the cluster-wide in-flight
+// counter so future admissions can use the slot; otherwise queue_limit
+// remains saturated forever after the first successful promotion. Also
+// exercises the only success-path end-to-end coverage of
+// NotifyPromotionSuccess in the suite.
+TEST_F(PromotionOnHitTest, NotifySuccessDecrementsCounter) {
+    MasterServiceConfig config;
+    config.enable_offload = true;
+    config.promotion_on_hit = true;
+    config.promotion_admission_threshold = 1;
+    config.promotion_queue_limit = 1;  // 1 in-flight task globally
+    config.default_kv_lease_ttl = 2000;
+    auto service = std::make_unique<MasterService>(config);
+
+    constexpr size_t seg_size = 1024 * 1024 * 16;
+    auto seg = PrepareSegment(*service, "seg_a", kDefaultSegmentBase, seg_size);
+
+    ASSERT_TRUE(InjectLocalDiskReplica(*service, seg.client_id, "k_first", 1024,
+                                       seg.segment_name));
+    ASSERT_TRUE(InjectLocalDiskReplica(*service, seg.client_id, "k_second",
+                                       1024, seg.segment_name));
+
+    // Admit task 1. promotion_in_flight_ goes from 0 -> 1.
+    {
+        auto r = service->GetReplicaList("k_first", "default");
+        ASSERT_TRUE(r.has_value());
+    }
+
+    // Drive the full success path: AllocStart stages the PROCESSING MEMORY
+    // replica and records alloc_id; NotifyPromotionSuccess flips it
+    // COMPLETE, drops the source LOCAL_DISK refcnt, erases the task, and
+    // decrements the counter.
+    auto alloc = service->PromotionAllocStart(seg.client_id, "k_first",
+                                              "default", 1024, {});
+    ASSERT_TRUE(alloc.has_value())
+        << "AllocStart should succeed; error=" << alloc.error();
+    auto notify =
+        service->NotifyPromotionSuccess(seg.client_id, "k_first", "default");
+    ASSERT_TRUE(notify.has_value())
+        << "NotifyPromotionSuccess happy path should succeed; if this "
+        << "fires, AllocStart did not record alloc_id, or the staged "
+        << "replica is missing/non-PROCESSING; error=" << notify.error();
+
+    // Counter must now be 0. A second admission on a *different* key must
+    // be allowed. If fetch_sub is missing on the success path, the cap is
+    // still saturated at 1 and TryPushPromotionQueue silently drops this
+    // attempt.
+    {
+        auto r = service->GetReplicaList("k_second", "default");
+        ASSERT_TRUE(r.has_value());
+    }
+    auto pending = service->PromotionObjectHeartbeat(seg.client_id);
+    ASSERT_TRUE(pending.has_value());
+    EXPECT_EQ(pending->size(), 1u)
+        << "k_second must be admitted after k_first succeeds — the global "
+        << "in-flight counter must decrement on the NotifyPromotionSuccess "
+        << "success path. Without fetch_sub, the cap stays saturated and "
+        << "k_second is silently dropped.";
+    EXPECT_EQ(CountPromotionTask(*pending, "k_second"), 1u);
+
+    service->RemoveAll();
+}
+
+// PromotionAllocStart must reject when the in-flight task has been
+// reaped between the holder's heartbeat and the AllocStart RPC arriving
+// (e.g. client stall past put_start_release_timeout_sec_). Without the
+// check, AllocStart would allocate + AddReplicas a PROCESSING MEMORY
+// replica with nothing tracking it: the generic PROCESSING reaper only
+// iterates shard->processing_keys (never populated by promotion) and
+// the promotion-task reaper has nothing left to iterate, so the buffer
+// leaks until the object is removed or evicted.
+TEST_F(PromotionOnHitTest, AllocStartRejectsReapedTask) {
+    MasterServiceConfig config;
+    config.enable_offload = true;
+    config.promotion_on_hit = true;
+    config.promotion_admission_threshold = 1;
+    config.default_kv_lease_ttl = 2000;
+    config.put_start_discard_timeout_sec = 0;
+    config.put_start_release_timeout_sec = 1;
+    auto service = std::make_unique<MasterService>(config);
+
+    constexpr size_t seg_size = 1024 * 1024 * 16;
+    auto ctx = PrepareSegment(*service, "seg_a", kDefaultSegmentBase, seg_size);
+    ASSERT_TRUE(InjectLocalDiskReplica(*service, ctx.client_id, "k_cold", 1024,
+                                       ctx.segment_name));
+
+    // Baseline allocator usage.
+    auto seg_baseline = service->QuerySegments(ctx.segment_name);
+    ASSERT_TRUE(seg_baseline.has_value());
+    const size_t used_baseline = seg_baseline->first;
+
+    // Admit the task.
+    {
+        auto r = service->GetReplicaList("k_cold", "default");
+        ASSERT_TRUE(r.has_value());
+    }
+
+    // Wait past the TTL so the reaper sweeps the task. AllocStart has
+    // not yet been called, so alloc_id is 0; the reaper's
+    // EraseReplicaByID branch is a no-op and only the task entry is
+    // removed. We can't easily poll for reap externally (the only
+    // user-facing observable would be re-admitting through the gate,
+    // which would create a fresh task and defeat the test), so use a
+    // fixed sleep with margin. 3s sleep over 1s release timeout matches
+    // StalePromotionReaper's CI-safe margin.
+    std::this_thread::sleep_for(std::chrono::seconds(3));
+
+    // AllocStart on a reaped task must reject without allocating.
+    // Allocating would leave an orphaned PROCESSING MEMORY replica
+    // attached to the object.
+    auto alloc = service->PromotionAllocStart(ctx.client_id, "k_cold",
+                                              "default", 1024, {});
+    ASSERT_FALSE(alloc.has_value())
+        << "AllocStart must reject when the task has been reaped — "
+        << "otherwise the staged PROCESSING MEMORY replica is orphaned";
+    EXPECT_EQ(alloc.error(), ErrorCode::REPLICA_IS_NOT_READY);
+
+    // The DRAM allocator must be at baseline: no buffer was staged.
+    auto seg_after = service->QuerySegments(ctx.segment_name);
+    ASSERT_TRUE(seg_after.has_value());
+    EXPECT_EQ(seg_after->first, used_baseline)
+        << "AllocStart on a reaped task must not allocate. If used bytes "
+        << "grew here, AllocStart got to AddReplicas before the task-"
+        << "existence check and the staged buffer is now orphaned (no "
+        << "reaper iterates it).";
+
+    service->RemoveAll();
+}
+
+// NotifyPromotionSuccess must reject calls from a client that is not
+// the holder of the source LOCAL_DISK replica. Without the check, any
+// client knowing the key could flip the staged PROCESSING MEMORY replica
+// to COMPLETE before the holder's RDMA write has landed, exposing torn
+// data to readers.
+TEST_F(PromotionOnHitTest, NotifyRejectsNonHolder) {
+    MasterServiceConfig config;
+    config.enable_offload = true;
+    config.promotion_on_hit = true;
+    config.promotion_admission_threshold = 1;
+    config.default_kv_lease_ttl = 2000;
+    auto service = std::make_unique<MasterService>(config);
+
+    constexpr size_t seg_size = 1024 * 1024 * 16;
+    auto holder =
+        PrepareSegment(*service, "seg_a", kDefaultSegmentBase, seg_size);
+    ASSERT_TRUE(InjectLocalDiskReplica(*service, holder.client_id, "k_cold",
+                                       1024, holder.segment_name));
+
+    // Admit + stage.
+    {
+        auto r = service->GetReplicaList("k_cold", "default");
+        ASSERT_TRUE(r.has_value());
+    }
+    auto alloc = service->PromotionAllocStart(holder.client_id, "k_cold",
+                                              "default", 1024, {});
+    ASSERT_TRUE(alloc.has_value());
+
+    // An unrelated client tries to Notify. Must be rejected as
+    // INVALID_PARAMS so the staged replica stays PROCESSING.
+    UUID intruder_id = generate_uuid();
+    ASSERT_NE(intruder_id, holder.client_id);
+    auto bad_notify =
+        service->NotifyPromotionSuccess(intruder_id, "k_cold", "default");
+    ASSERT_FALSE(bad_notify.has_value())
+        << "Notify from a non-holder client must be rejected — otherwise "
+        << "any client knowing the key can commit someone else's "
+        << "still-being-written replica";
+    EXPECT_EQ(bad_notify.error(), ErrorCode::INVALID_PARAMS);
+
+    // The staged replica must still be PROCESSING (not committed by the
+    // rejected call). Readers must not see it via GetReplicaList yet —
+    // GetReplicaList filters PROCESSING replicas, and pre-AllocStart
+    // there are no COMPLETE replicas to read for a LOCAL_DISK-only key
+    // (only the LOCAL_DISK descriptor itself).
+    auto r_check = service->GetReplicaList("k_cold", "default");
+    ASSERT_TRUE(r_check.has_value());
+    bool saw_complete_memory = false;
+    for (const auto& d : r_check.value().replicas) {
+        if (d.is_memory_replica() && d.status == ReplicaStatus::COMPLETE) {
+            saw_complete_memory = true;
+            break;
+        }
+    }
+    EXPECT_FALSE(saw_complete_memory)
+        << "Rejected Notify must not have committed the staged replica";
+
+    // The legitimate holder must still be able to Notify successfully.
+    auto good_notify =
+        service->NotifyPromotionSuccess(holder.client_id, "k_cold", "default");
+    ASSERT_TRUE(good_notify.has_value())
+        << "Holder Notify on the same task must succeed after a rejected "
+        << "intruder Notify — the task entry should be untouched by the "
+        << "rejection path; error=" << good_notify.error();
+
+    service->RemoveAll();
+}
+
+// NotifyPromotionFailure must immediately release the task slot and the
+// staged buffer so that transient client-side errors (SSD throttling,
+// RDMA flakes) do not pin promotion_queue_limit_ for the full reaper
+// TTL. With queue_limit=1 and a holder client that fails after a
+// successful AllocStart, a second admission on a different key would
+// be silently dropped until reaper TTL if Notify-Failure didn't
+// fast-release. The TTL is set high enough here that any test pass
+// can be attributed to the explicit release, not to reaper expiry.
+TEST_F(PromotionOnHitTest, NotifyFailureReleasesStateImmediately) {
+    MasterServiceConfig config;
+    config.enable_offload = true;
+    config.promotion_on_hit = true;
+    config.promotion_admission_threshold = 1;
+    config.promotion_queue_limit = 1;  // cap=1 makes the slot observable
+    config.default_kv_lease_ttl = 5000;
+    // Long TTL so the test's pass-fail signal cannot be attributed to
+    // reaper sweep; only NotifyPromotionFailure could plausibly release.
+    config.put_start_release_timeout_sec = 300;
+    auto service = std::make_unique<MasterService>(config);
+
+    constexpr size_t seg_size = 1024 * 1024 * 16;
+    auto seg = PrepareSegment(*service, "seg_a", kDefaultSegmentBase, seg_size);
+
+    ASSERT_TRUE(InjectLocalDiskReplica(*service, seg.client_id, "k_a", 1024,
+                                       seg.segment_name));
+    ASSERT_TRUE(InjectLocalDiskReplica(*service, seg.client_id, "k_b", 1024,
+                                       seg.segment_name));
+
+    auto seg_baseline = service->QuerySegments(seg.segment_name);
+    ASSERT_TRUE(seg_baseline.has_value());
+    const size_t used_baseline = seg_baseline->first;
+
+    auto& mm = MasterMetricManager::instance();
+    const int64_t failed_pre = mm.get_promotion_failed();
+
+    // Admit + stage k_a. promotion_in_flight_ goes 0 -> 1.
+    {
+        auto r = service->GetReplicaList("k_a", "default");
+        ASSERT_TRUE(r.has_value());
+    }
+    auto alloc =
+        service->PromotionAllocStart(seg.client_id, "k_a", "default", 1024, {});
+    ASSERT_TRUE(alloc.has_value());
+
+    // The staged PROCESSING MEMORY buffer is allocated.
+    auto seg_after_alloc = service->QuerySegments(seg.segment_name);
+    ASSERT_TRUE(seg_after_alloc.has_value());
+    EXPECT_GT(seg_after_alloc->first, used_baseline)
+        << "AllocStart should commit a buffer in the DRAM allocator";
+
+    // Holder reports failure (simulating SSD read error after AllocStart
+    // succeeded). Master must immediately reap the staged replica and
+    // decrement the slot counter.
+    auto failure =
+        service->NotifyPromotionFailure(seg.client_id, "k_a", "default");
+    ASSERT_TRUE(failure.has_value())
+        << "NotifyPromotionFailure on a valid in-flight task from the "
+        << "legitimate holder must succeed; error=" << failure.error();
+    EXPECT_EQ(mm.get_promotion_failed() - failed_pre, 1)
+        << "holder-reported failure must bump promotion_failed";
+
+    // The staged buffer must be freed back to the DRAM allocator. If
+    // this fires, NotifyPromotionFailure did not pop the staged replica
+    // via EraseReplicaByID — same orphan-replica shape that originally
+    // motivated the reaper fix.
+    auto seg_after_release = service->QuerySegments(seg.segment_name);
+    ASSERT_TRUE(seg_after_release.has_value());
+    EXPECT_EQ(seg_after_release->first, used_baseline)
+        << "NotifyPromotionFailure must release the staged buffer back "
+        << "to the allocator; otherwise it leaks until the object is "
+        << "removed or evicted.";
+
+    // The slot must be freed: a second admission on a different key must
+    // succeed even though queue_limit=1. Without the failure-side
+    // decrement the cap would stay saturated until reaper TTL.
+    {
+        auto r = service->GetReplicaList("k_b", "default");
+        ASSERT_TRUE(r.has_value());
+    }
+    auto heartbeat = service->PromotionObjectHeartbeat(seg.client_id);
+    ASSERT_TRUE(heartbeat.has_value());
+    EXPECT_EQ(CountPromotionTask(*heartbeat, "k_b"), 1u)
+        << "k_b admission must succeed after k_a's failure released the "
+        << "slot. If this fires, NotifyPromotionFailure did not decrement "
+        << "promotion_in_flight_, and transient client-side errors "
+        << "would saturate the queue limit for the full reaper TTL.";
+
+    // Idempotency: repeated failure notification on the same key must be
+    // safe (return OK without underflowing the counter).
+    auto failure_again =
+        service->NotifyPromotionFailure(seg.client_id, "k_a", "default");
+    EXPECT_TRUE(failure_again.has_value())
+        << "Repeated NotifyPromotionFailure should be idempotent (return "
+        << "OK on already-released task), not error.";
+
+    service->RemoveAll();
+}
+
+// NotifyPromotionFailure must reject calls from a client that is not the
+// holder. Without this gate, any client knowing the key could prematurely
+// release a legitimate in-flight promotion's master state and free the
+// staged buffer mid-RDMA-write. Mirror of NotifyRejectsNonHolder for the
+// Notify-Success path.
+TEST_F(PromotionOnHitTest, NotifyFailureRejectsNonHolder) {
+    MasterServiceConfig config;
+    config.enable_offload = true;
+    config.promotion_on_hit = true;
+    config.promotion_admission_threshold = 1;
+    config.default_kv_lease_ttl = 5000;
+    config.put_start_release_timeout_sec = 300;
+    auto service = std::make_unique<MasterService>(config);
+
+    constexpr size_t seg_size = 1024 * 1024 * 16;
+    auto holder =
+        PrepareSegment(*service, "seg_a", kDefaultSegmentBase, seg_size);
+    ASSERT_TRUE(InjectLocalDiskReplica(*service, holder.client_id, "k_cold",
+                                       1024, holder.segment_name));
+
+    {
+        auto r = service->GetReplicaList("k_cold", "default");
+        ASSERT_TRUE(r.has_value());
+    }
+    auto alloc = service->PromotionAllocStart(holder.client_id, "k_cold",
+                                              "default", 1024, {});
+    ASSERT_TRUE(alloc.has_value());
+
+    // Intruder calls Failure with the wrong client_id.
+    UUID intruder_id = generate_uuid();
+    ASSERT_NE(intruder_id, holder.client_id);
+    auto bad_failure =
+        service->NotifyPromotionFailure(intruder_id, "k_cold", "default");
+    ASSERT_FALSE(bad_failure.has_value())
+        << "Failure from a non-holder client must be rejected.";
+    EXPECT_EQ(bad_failure.error(), ErrorCode::INVALID_PARAMS);
+
+    // The legitimate holder must still be able to either commit via
+    // Notify-Success or release via Notify-Failure on the same task. Use
+    // Notify-Failure here to exercise the surviving-task path.
+    auto good_failure =
+        service->NotifyPromotionFailure(holder.client_id, "k_cold", "default");
+    ASSERT_TRUE(good_failure.has_value())
+        << "Holder Failure must succeed after a rejected intruder "
+        << "Failure — the task entry should be untouched by the "
+        << "rejection path; error=" << good_failure.error();
+
+    service->RemoveAll();
+}
+
+// PromotionAllocStart must reject callers that aren't the holder. Mirror
+// of the Notify gate. Without this gate a client that drained another's
+// promotion_objects queue via PromotionObjectHeartbeat could call
+// AllocStart to stage arbitrary DRAM allocations on the destination
+// segment, with no path to commit (Notify rejects on holder_id mismatch)
+// — they'd just sit pinned until reaper TTL.
+TEST_F(PromotionOnHitTest, AllocStartRejectsNonHolder) {
+    MasterServiceConfig config;
+    config.enable_offload = true;
+    config.promotion_on_hit = true;
+    config.promotion_admission_threshold = 1;
+    config.default_kv_lease_ttl = 5000;
+    auto service = std::make_unique<MasterService>(config);
+
+    constexpr size_t seg_size = 1024 * 1024 * 16;
+    auto holder =
+        PrepareSegment(*service, "seg_a", kDefaultSegmentBase, seg_size);
+    ASSERT_TRUE(InjectLocalDiskReplica(*service, holder.client_id, "k_cold",
+                                       1024, holder.segment_name));
+
+    auto seg_baseline = service->QuerySegments(holder.segment_name);
+    ASSERT_TRUE(seg_baseline.has_value());
+    const size_t used_baseline = seg_baseline->first;
+
+    {
+        auto r = service->GetReplicaList("k_cold", "default");
+        ASSERT_TRUE(r.has_value());
+    }
+
+    UUID intruder_id = generate_uuid();
+    ASSERT_NE(intruder_id, holder.client_id);
+    auto bad_alloc = service->PromotionAllocStart(intruder_id, "k_cold",
+                                                  "default", 1024, {});
+    ASSERT_FALSE(bad_alloc.has_value())
+        << "AllocStart from a non-holder client must be rejected — "
+        << "otherwise an attacker that drained another's queue could "
+        << "stage arbitrary DRAM allocations.";
+    EXPECT_EQ(bad_alloc.error(), ErrorCode::INVALID_PARAMS);
+
+    // No buffer was allocated.
+    auto seg_after_bad = service->QuerySegments(holder.segment_name);
+    ASSERT_TRUE(seg_after_bad.has_value());
+    EXPECT_EQ(seg_after_bad->first, used_baseline)
+        << "Rejected AllocStart must not have allocated any DRAM.";
+
+    // The legitimate holder must still be able to AllocStart (task
+    // untouched by the rejection).
+    auto good_alloc = service->PromotionAllocStart(holder.client_id, "k_cold",
+                                                   "default", 1024, {});
+    ASSERT_TRUE(good_alloc.has_value())
+        << "Holder AllocStart on the same task must succeed after a "
+        << "rejected intruder AllocStart; error=" << good_alloc.error();
+
+    service->RemoveAll();
+}
+
+// PromotionAllocStart must reject size that doesn't match the task's
+// recorded object_size. The size is captured from the source LOCAL_DISK
+// descriptor at task admission; mismatch indicates a buggy caller or
+// malicious request and would let the caller request an arbitrary-size
+// DRAM buffer (smaller → eventual RDMA-write overflow risk; larger →
+// wasted DRAM pinned until reaper TTL).
+TEST_F(PromotionOnHitTest, AllocStartRejectsSizeMismatch) {
+    MasterServiceConfig config;
+    config.enable_offload = true;
+    config.promotion_on_hit = true;
+    config.promotion_admission_threshold = 1;
+    config.default_kv_lease_ttl = 5000;
+    auto service = std::make_unique<MasterService>(config);
+
+    constexpr size_t seg_size = 1024 * 1024 * 16;
+    auto holder =
+        PrepareSegment(*service, "seg_a", kDefaultSegmentBase, seg_size);
+    constexpr int64_t kRealSize = 1024;
+    ASSERT_TRUE(InjectLocalDiskReplica(*service, holder.client_id, "k_cold",
+                                       kRealSize, holder.segment_name));
+
+    auto seg_baseline = service->QuerySegments(holder.segment_name);
+    ASSERT_TRUE(seg_baseline.has_value());
+    const size_t used_baseline = seg_baseline->first;
+
+    {
+        auto r = service->GetReplicaList("k_cold", "default");
+        ASSERT_TRUE(r.has_value());
+    }
+
+    // Mismatched-size requests must be rejected, regardless of direction.
+    for (uint64_t bad_size : {static_cast<uint64_t>(kRealSize) / 2,
+                              static_cast<uint64_t>(kRealSize) * 4}) {
+        auto bad_alloc = service->PromotionAllocStart(
+            holder.client_id, "k_cold", "default", bad_size, {});
+        ASSERT_FALSE(bad_alloc.has_value())
+            << "AllocStart with size=" << bad_size
+            << " (task.object_size=" << kRealSize << ") must be rejected.";
+        EXPECT_EQ(bad_alloc.error(), ErrorCode::INVALID_PARAMS);
+
+        // No buffer must have been staged.
+        auto seg_after_bad = service->QuerySegments(holder.segment_name);
+        ASSERT_TRUE(seg_after_bad.has_value());
+        EXPECT_EQ(seg_after_bad->first, used_baseline)
+            << "Rejected AllocStart (size=" << bad_size
+            << ") must not have allocated any DRAM.";
+    }
+
+    // Correct size must still work — the task was not consumed by the
+    // rejections.
+    auto good_alloc =
+        service->PromotionAllocStart(holder.client_id, "k_cold", "default",
+                                     static_cast<uint64_t>(kRealSize), {});
+    ASSERT_TRUE(good_alloc.has_value())
+        << "AllocStart with the correct size must succeed after rejected "
+        << "size-mismatch attempts; error=" << good_alloc.error();
+
+    service->RemoveAll();
+}
+
+// When a holder client expires, ClientMonitorFunc must clean up its
+// dangling promotion_tasks entries and decrement the global in-flight
+// counter. Without this cleanup, the entries stay pinned until reaper
+// TTL — and on a rolling restart of many holders the cluster-wide cap
+// promotion_queue_limit_ saturates and blocks all new admissions for
+// the full TTL.
+//
+// Test mechanism: short client_live_ttl_sec, admit a promotion, stop
+// pinging, wait for ClientMonitorFunc to expire the client and call
+// ClearInvalidHandles. Then assert promotion_in_flight_ is back to 0
+// by attempting a second admission with queue_limit=1 on a fresh
+// client.
+TEST_F(PromotionOnHitTest, ClientExpiryClearsPromotionTask) {
+    MasterServiceConfig config;
+    config.enable_offload = true;
+    config.promotion_on_hit = true;
+    config.promotion_admission_threshold = 1;
+    config.promotion_queue_limit = 1;  // cap=1 makes the slot observable
+    config.default_kv_lease_ttl = 5000;
+    // Long task TTL so that any clearing we see must come from
+    // ClearInvalidHandles, not from the promotion-task reaper.
+    config.put_start_release_timeout_sec = 300;
+    // Short client TTL so expiration is fast.
+    config.client_live_ttl_sec = 1;
+    auto service = std::make_unique<MasterService>(config);
+
+    constexpr size_t seg_size = 1024 * 1024 * 16;
+    auto holder =
+        PrepareSegment(*service, "seg_a", kDefaultSegmentBase, seg_size);
+    ASSERT_TRUE(InjectLocalDiskReplica(*service, holder.client_id, "k_cold",
+                                       1024, holder.segment_name));
+
+    auto& mm = MasterMetricManager::instance();
+    const int64_t cancelled_pre = mm.get_promotion_cancelled();
+
+    // Admit the promotion. promotion_in_flight_ goes 0 -> 1.
+    {
+        auto r = service->GetReplicaList("k_cold", "default");
+        ASSERT_TRUE(r.has_value());
+    }
+
+    // Sanity: with queue_limit=1 the cap is saturated. A second
+    // different-shard admission must be rejected right now.
+    auto second_holder = PrepareSegment(
+        *service, "seg_b", kDefaultSegmentBase + seg_size, seg_size);
+    // Promote second_holder into ok_client_ via ReMountSegment so its
+    // LOCAL_DISK replicas survive any ClearInvalidHandles run triggered
+    // by the first holder's expiry. MountSegment alone does not register
+    // the client as alive (only ReMountSegment does), and
+    // CleanupStaleHandles uses ok_client_ to decide which LOCAL_DISK
+    // replicas to erase — without this, second_holder's k_other replica
+    // would be wiped alongside the first holder's k_cold replica when
+    // ClearInvalidHandles runs.
+    {
+        Segment seg_b =
+            MakeSegment("seg_b", kDefaultSegmentBase + seg_size, seg_size);
+        seg_b.id = second_holder.segment_id;
+        std::vector<Segment> segs{seg_b};
+        auto remount = service->ReMountSegment(segs, second_holder.client_id);
+        ASSERT_TRUE(remount.has_value()) << "ReMount failed";
+    }
+    ASSERT_TRUE(InjectLocalDiskReplica(*service, second_holder.client_id,
+                                       "k_other", 1024,
+                                       second_holder.segment_name));
+    {
+        auto r = service->GetReplicaList("k_other", "default");
+        ASSERT_TRUE(r.has_value());
+    }
+    auto pending_pre =
+        service->PromotionObjectHeartbeat(second_holder.client_id);
+    ASSERT_TRUE(pending_pre.has_value());
+    EXPECT_EQ(CountPromotionTask(*pending_pre, "k_other"), 0u)
+        << "Sanity: queue_limit=1 should block the second admission "
+        << "while the first task is in flight.";
+
+    // Wait for the first holder to expire while keeping the second
+    // holder alive via periodic Pings. ClientMonitorFunc runs every
+    // kClientMonitorSleepMs (1s) and expires clients whose last ping is
+    // older than client_live_ttl_sec (1s); we ping second_holder every
+    // 200ms so it stays alive across the 4s wait. Without these pings
+    // both holders would expire together and the second holder's
+    // LOCAL_DISK segment would be unmounted — making "k_other" un-
+    // admittable for reasons unrelated to the promotion-task slot.
+    {
+        const auto deadline =
+            std::chrono::steady_clock::now() + std::chrono::seconds(4);
+        while (std::chrono::steady_clock::now() < deadline) {
+            (void)service->Ping(second_holder.client_id);
+            std::this_thread::sleep_for(std::chrono::milliseconds(200));
+        }
+    }
+
+    // ClearInvalidHandles should have erased the holder's LOCAL_DISK
+    // source replica AND (with the fix) the promotion_tasks entry,
+    // decrementing the global in-flight counter. Re-admit a promotion
+    // on the second holder; with queue_limit=1 this can only succeed if
+    // the slot was freed.
+    {
+        auto r = service->GetReplicaList("k_other", "default");
+        ASSERT_TRUE(r.has_value())
+            << "GetReplicaList(k_other) failed with error=" << r.error();
+    }
+    auto pending_post =
+        service->PromotionObjectHeartbeat(second_holder.client_id);
+    ASSERT_TRUE(pending_post.has_value());
+    EXPECT_EQ(CountPromotionTask(*pending_post, "k_other"), 1u)
+        << "After the holder expired, ClearInvalidHandles must have "
+        << "erased its promotion_tasks entry and decremented "
+        << "promotion_in_flight_. Otherwise the global cap remains "
+        << "saturated by the dead holder's task for "
+        << "put_start_release_timeout_sec_ seconds, and this admission "
+        << "is dropped.";
+    EXPECT_EQ(mm.get_promotion_cancelled() - cancelled_pre, 1)
+        << "Holder-client expiry mid-promotion must bump "
+        << "promotion_cancelled";
+
+    service->RemoveAll();
+}
+
+// promotion_admission_threshold=0 would silently bypass the frequency
+// gate (`freq < 0` is never true for uint8_t freq), letting every Get
+// on a LOCAL_DISK-only key admit a promotion. master.cpp clamps the
+// flag at parse time, and MasterService's constructor adds a
+// defense-in-depth clamp for direct-construction paths (tests,
+// embedded users). Verify the clamp lifts 0 → 1.
+TEST_F(PromotionOnHitTest, AdmissionThresholdZeroClampsToOne) {
+    MasterServiceConfig config;
+    config.enable_offload = true;
+    config.promotion_on_hit = true;
+    config.promotion_admission_threshold = 0;  // would bypass the gate
+    auto service = std::make_unique<MasterService>(config);
+    EXPECT_EQ(GetPromotionAdmissionThresholdForTesting(service.get()), 1u)
+        << "threshold=0 must be clamped to 1; otherwise every Get-on-"
+        << "LOCAL_DISK admits a promotion and the frequency gate is "
+        << "silently disabled.";
+}
+
+// promotion_admission_threshold above the CountMinSketch saturating
+// max (255) would make the gate unreachable (freq saturates at 255,
+// `255 < 256` is true forever → no admissions). Verify the constructor
+// clamps high too.
+TEST_F(PromotionOnHitTest, AdmissionThresholdAboveMaxClampsToMax) {
+    MasterServiceConfig config;
+    config.enable_offload = true;
+    config.promotion_on_hit = true;
+    config.promotion_admission_threshold = 1000;  // > 255
+    auto service = std::make_unique<MasterService>(config);
+    EXPECT_EQ(GetPromotionAdmissionThresholdForTesting(service.get()), 255u)
+        << "threshold above the CountMinSketch counter max (255) must "
+        << "be clamped to 255; otherwise the gate is unreachable.";
+}
+
+// Remove(force=true) on a key with an in-flight PromotionTask must
+// drop the task entry alongside the metadata, so promotion_in_flight_
+// is decremented immediately rather than pinned for ~10 min until the
+// reaper sweeps. With queue_limit=1, the test admits one task, removes
+// the key, then admits a second task on a different key — only
+// succeeds if the slot was freed.
+TEST_F(PromotionOnHitTest, RemoveErasesPromotionTask) {
+    MasterServiceConfig config;
+    config.enable_offload = true;
+    config.promotion_on_hit = true;
+    config.promotion_admission_threshold = 1;
+    config.promotion_queue_limit = 1;  // makes the slot observable
+    config.default_kv_lease_ttl = 5000;
+    // Long task TTL so any cap-slot reclaim must come from the
+    // metadata-erase path, not the reaper.
+    config.put_start_release_timeout_sec = 300;
+    auto service = std::make_unique<MasterService>(config);
+
+    constexpr size_t seg_size = 1024 * 1024 * 16;
+    auto holder =
+        PrepareSegment(*service, "seg_a", kDefaultSegmentBase, seg_size);
+    ASSERT_TRUE(InjectLocalDiskReplica(*service, holder.client_id, "k_first",
+                                       1024, holder.segment_name));
+    ASSERT_TRUE(InjectLocalDiskReplica(*service, holder.client_id, "k_second",
+                                       1024, holder.segment_name));
+
+    // Admit task 1.
+    {
+        auto r = service->GetReplicaList("k_first", "default");
+        ASSERT_TRUE(r.has_value());
+    }
+    {
+        auto pending = service->PromotionObjectHeartbeat(holder.client_id);
+        ASSERT_TRUE(pending.has_value());
+        EXPECT_EQ(CountPromotionTask(*pending, "k_first"), 1u);
+    }
+
+    // Remove k_first with force=true. With the fix, this also wipes
+    // k_first's promotion_tasks entry and decrements
+    // promotion_in_flight_ back to 0.
+    auto rm = service->Remove("k_first", "default", /*force=*/true);
+    ASSERT_TRUE(rm.has_value())
+        << "Remove should succeed; error=" << rm.error();
+
+    // Now admit a different key. With queue_limit=1, this only succeeds
+    // if the slot was freed by Remove. Without the in-flight cleanup,
+    // it would stay pinned for the full 300s TTL.
+    {
+        auto r = service->GetReplicaList("k_second", "default");
+        ASSERT_TRUE(r.has_value());
+    }
+    auto pending_post = service->PromotionObjectHeartbeat(holder.client_id);
+    ASSERT_TRUE(pending_post.has_value());
+    EXPECT_EQ(CountPromotionTask(*pending_post, "k_second"), 1u)
+        << "k_second must be admittable after Remove of k_first — Remove "
+        << "must erase the in-flight promotion_tasks entry and decrement "
+        << "promotion_in_flight_, otherwise queue_limit=1 stays saturated.";
+
+    service->RemoveAll();
+}
+
+// RemoveByRegex on a key with an in-flight PromotionTask must drop the
+// task entry, same as Remove. Mirror of RemoveErasesPromotionTask, but
+// exercises the regex path which iterates shards directly without an
+// accessor object.
+TEST_F(PromotionOnHitTest, RemoveByRegexErasesPromotionTask) {
+    MasterServiceConfig config;
+    config.enable_offload = true;
+    config.promotion_on_hit = true;
+    config.promotion_admission_threshold = 1;
+    config.promotion_queue_limit = 1;
+    config.default_kv_lease_ttl = 5000;
+    config.put_start_release_timeout_sec = 300;
+    auto service = std::make_unique<MasterService>(config);
+
+    constexpr size_t seg_size = 1024 * 1024 * 16;
+    auto holder =
+        PrepareSegment(*service, "seg_a", kDefaultSegmentBase, seg_size);
+    ASSERT_TRUE(InjectLocalDiskReplica(*service, holder.client_id, "regex_k1",
+                                       1024, holder.segment_name));
+    ASSERT_TRUE(InjectLocalDiskReplica(*service, holder.client_id, "other_k2",
+                                       1024, holder.segment_name));
+
+    // Admit task on regex_k1.
+    {
+        auto r = service->GetReplicaList("regex_k1", "default");
+        ASSERT_TRUE(r.has_value());
+    }
+
+    // RemoveByRegex matches regex_k1 only.
+    auto removed = service->RemoveByRegex("^regex_", "default", /*force=*/true);
+    ASSERT_TRUE(removed.has_value())
+        << "RemoveByRegex should succeed; error=" << removed.error();
+    EXPECT_EQ(removed.value(), 1) << "exactly one key (regex_k1) should match";
+
+    // Slot must be free — admit on other_k2 (different shard or same,
+    // doesn't matter because counter is global).
+    {
+        auto r = service->GetReplicaList("other_k2", "default");
+        ASSERT_TRUE(r.has_value());
+    }
+    auto pending_post = service->PromotionObjectHeartbeat(holder.client_id);
+    ASSERT_TRUE(pending_post.has_value());
+    EXPECT_EQ(CountPromotionTask(*pending_post, "other_k2"), 1u)
+        << "other_k2 must be admittable after RemoveByRegex of regex_k1 "
+        << "— RemoveByRegex must erase the in-flight promotion_tasks "
+        << "entry. Otherwise queue_limit=1 stays saturated.";
+
+    service->RemoveAll();
+}
+
+// RemoveAll on a key with an in-flight PromotionTask must drop the
+// task entry alongside the metadata. Same shape as
+// RemoveErasesPromotionTask but exercises the bulk-erase loop in
+// MasterService::RemoveAll, which iterates every shard and erases
+// metadata entries directly.
+TEST_F(PromotionOnHitTest, RemoveAllErasesPromotionTask) {
+    MasterServiceConfig config;
+    config.enable_offload = true;
+    config.promotion_on_hit = true;
+    config.promotion_admission_threshold = 1;
+    config.promotion_queue_limit = 1;
+    config.default_kv_lease_ttl = 5000;
+    config.put_start_release_timeout_sec = 300;
+    auto service = std::make_unique<MasterService>(config);
+
+    constexpr size_t seg_size = 1024 * 1024 * 16;
+    auto holder =
+        PrepareSegment(*service, "seg_a", kDefaultSegmentBase, seg_size);
+    ASSERT_TRUE(InjectLocalDiskReplica(*service, holder.client_id, "k_first",
+                                       1024, holder.segment_name));
+
+    auto& mm = MasterMetricManager::instance();
+    const int64_t cancelled_pre = mm.get_promotion_cancelled();
+
+    {
+        auto r = service->GetReplicaList("k_first", "default");
+        ASSERT_TRUE(r.has_value());
+    }
+    {
+        auto pending = service->PromotionObjectHeartbeat(holder.client_id);
+        ASSERT_TRUE(pending.has_value());
+        EXPECT_EQ(CountPromotionTask(*pending, "k_first"), 1u);
+    }
+
+    auto removed = service->RemoveAll(/*force=*/true);
+    EXPECT_GE(removed, 1) << "RemoveAll should erase k_first";
+
+    EXPECT_EQ(mm.get_promotion_cancelled() - cancelled_pre, 1)
+        << "RemoveAll on a key with an in-flight promotion must route "
+        << "through EraseMetadataEntry and bump promotion_cancelled_total.";
+
+    ASSERT_TRUE(InjectLocalDiskReplica(*service, holder.client_id, "k_second",
+                                       1024, holder.segment_name));
+    {
+        auto r = service->GetReplicaList("k_second", "default");
+        ASSERT_TRUE(r.has_value());
+    }
+    auto pending_post = service->PromotionObjectHeartbeat(holder.client_id);
+    ASSERT_TRUE(pending_post.has_value());
+    EXPECT_EQ(CountPromotionTask(*pending_post, "k_second"), 1u)
+        << "k_second must be admittable after RemoveAll of k_first — "
+        << "otherwise queue_limit=1 stays saturated until reaper TTL.";
+
+    service->RemoveAll(/*force=*/true);
+}
+
+// BatchRemove normal-completion path on a key with an in-flight
+// PromotionTask must drop the task entry. ReMountSegment registers the
+// holder in ok_client_ so CleanupStaleHandles returns false and
+// BatchRemove takes the non-stale branch.
+TEST_F(PromotionOnHitTest, BatchRemoveErasesPromotionTask) {
+    MasterServiceConfig config;
+    config.enable_offload = true;
+    config.promotion_on_hit = true;
+    config.promotion_admission_threshold = 1;
+    config.promotion_queue_limit = 1;
+    config.default_kv_lease_ttl = 5000;
+    config.put_start_release_timeout_sec = 300;
+    auto service = std::make_unique<MasterService>(config);
+
+    constexpr size_t seg_size = 1024 * 1024 * 16;
+    auto holder =
+        PrepareSegment(*service, "seg_a", kDefaultSegmentBase, seg_size);
+    {
+        Segment seg_a = MakeSegment("seg_a", kDefaultSegmentBase, seg_size);
+        seg_a.id = holder.segment_id;
+        std::vector<Segment> segs{seg_a};
+        auto remount = service->ReMountSegment(segs, holder.client_id);
+        ASSERT_TRUE(remount.has_value()) << "ReMount failed";
+    }
+    ASSERT_TRUE(InjectLocalDiskReplica(*service, holder.client_id, "k_first",
+                                       1024, holder.segment_name));
+    ASSERT_TRUE(InjectLocalDiskReplica(*service, holder.client_id, "k_second",
+                                       1024, holder.segment_name));
+
+    auto& mm = MasterMetricManager::instance();
+    const int64_t cancelled_pre = mm.get_promotion_cancelled();
+
+    {
+        auto r = service->GetReplicaList("k_first", "default");
+        ASSERT_TRUE(r.has_value());
+    }
+    {
+        auto pending = service->PromotionObjectHeartbeat(holder.client_id);
+        ASSERT_TRUE(pending.has_value());
+        EXPECT_EQ(CountPromotionTask(*pending, "k_first"), 1u);
+    }
+
+    auto results = service->BatchRemove({"k_first"}, "default", /*force=*/true);
+    ASSERT_EQ(results.size(), 1u);
+    EXPECT_TRUE(results[0].has_value())
+        << "BatchRemove should succeed; error=" << results[0].error();
+
+    EXPECT_EQ(mm.get_promotion_cancelled() - cancelled_pre, 1)
+        << "BatchRemove normal path on a key with an in-flight promotion "
+        << "must bump promotion_cancelled_total.";
+
+    {
+        auto r = service->GetReplicaList("k_second", "default");
+        ASSERT_TRUE(r.has_value());
+    }
+    auto pending_post = service->PromotionObjectHeartbeat(holder.client_id);
+    ASSERT_TRUE(pending_post.has_value());
+    EXPECT_EQ(CountPromotionTask(*pending_post, "k_second"), 1u)
+        << "k_second must be admittable after BatchRemove of k_first — "
+        << "otherwise queue_limit=1 stays saturated until reaper TTL.";
+
+    service->RemoveAll(/*force=*/true);
+}
+
+// BatchRemove stale-handle path on a key with an in-flight
+// PromotionTask must drop the task entry. The holder is mounted via
+// PrepareSegment only (no ReMount), so its client is absent from
+// ok_client_; BatchRemove's CleanupStaleHandles then erases the
+// LOCAL_DISK replica and the stale-handle branch fires.
+TEST_F(PromotionOnHitTest, BatchRemoveStaleHandleErasesPromotionTask) {
+    MasterServiceConfig config;
+    config.enable_offload = true;
+    config.promotion_on_hit = true;
+    config.promotion_admission_threshold = 1;
+    config.promotion_queue_limit = 1;
+    config.default_kv_lease_ttl = 5000;
+    config.put_start_release_timeout_sec = 300;
+    auto service = std::make_unique<MasterService>(config);
+
+    constexpr size_t seg_size = 1024 * 1024 * 16;
+    auto holder =
+        PrepareSegment(*service, "seg_a", kDefaultSegmentBase, seg_size);
+    ASSERT_TRUE(InjectLocalDiskReplica(*service, holder.client_id, "k_first",
+                                       1024, holder.segment_name));
+
+    auto& mm = MasterMetricManager::instance();
+    const int64_t cancelled_pre = mm.get_promotion_cancelled();
+
+    {
+        auto r = service->GetReplicaList("k_first", "default");
+        ASSERT_TRUE(r.has_value());
+    }
+    {
+        auto pending = service->PromotionObjectHeartbeat(holder.client_id);
+        ASSERT_TRUE(pending.has_value());
+        EXPECT_EQ(CountPromotionTask(*pending, "k_first"), 1u);
+    }
+
+    auto results = service->BatchRemove({"k_first"}, "default", /*force=*/true);
+    ASSERT_EQ(results.size(), 1u);
+    EXPECT_FALSE(results[0].has_value())
+        << "stale-handle path should report OBJECT_NOT_FOUND once the "
+        << "LOCAL_DISK replica is wiped.";
+
+    EXPECT_EQ(mm.get_promotion_cancelled() - cancelled_pre, 1)
+        << "BatchRemove stale-handle path must also bump "
+        << "promotion_cancelled_total.";
+
+    auto second_holder = PrepareSegment(
+        *service, "seg_b", kDefaultSegmentBase + seg_size, seg_size);
+    {
+        Segment seg_b =
+            MakeSegment("seg_b", kDefaultSegmentBase + seg_size, seg_size);
+        seg_b.id = second_holder.segment_id;
+        std::vector<Segment> segs{seg_b};
+        auto remount = service->ReMountSegment(segs, second_holder.client_id);
+        ASSERT_TRUE(remount.has_value()) << "ReMount failed";
+    }
+    ASSERT_TRUE(InjectLocalDiskReplica(*service, second_holder.client_id,
+                                       "k_second", 1024,
+                                       second_holder.segment_name));
+    {
+        auto r = service->GetReplicaList("k_second", "default");
+        ASSERT_TRUE(r.has_value());
+    }
+    auto pending_post =
+        service->PromotionObjectHeartbeat(second_holder.client_id);
+    ASSERT_TRUE(pending_post.has_value());
+    EXPECT_EQ(CountPromotionTask(*pending_post, "k_second"), 1u)
+        << "k_second must be admittable after the stale-handle "
+        << "BatchRemove — otherwise queue_limit=1 stays saturated until "
+        << "reaper TTL.";
+
+    service->RemoveAll(/*force=*/true);
+}
+
+// The funnel metrics (admitted, completed, completed_bytes, in_flight)
+// track a successful promotion lifecycle end-to-end: a single
+// admission bumps admitted+1 and in_flight+1; NotifyPromotionSuccess
+// bumps completed+1 and completed_bytes+object_size and brings
+// in_flight back to 0.
+TEST_F(PromotionOnHitTest, MetricsFunnelTracksSuccessfulPromotion) {
+    MasterServiceConfig config;
+    config.enable_offload = true;
+    config.promotion_on_hit = true;
+    config.promotion_admission_threshold = 1;
+    config.default_kv_lease_ttl = 2000;
+    auto service = std::make_unique<MasterService>(config);
+
+    constexpr size_t seg_size = 1024 * 1024 * 16;
+    constexpr int64_t kObjBytes = 4096;
+    auto seg = PrepareSegment(*service, "seg_a", kDefaultSegmentBase, seg_size);
+    ASSERT_TRUE(InjectLocalDiskReplica(*service, seg.client_id, "k_hot",
+                                       kObjBytes, seg.segment_name));
+
+    auto& mm = MasterMetricManager::instance();
+    const int64_t admitted_pre = mm.get_promotion_admitted();
+    const int64_t completed_pre = mm.get_promotion_completed();
+    const int64_t bytes_pre = mm.get_promotion_completed_bytes();
+    const int64_t in_flight_pre = mm.get_promotion_in_flight();
+
+    // Admit.
+    {
+        auto r = service->GetReplicaList("k_hot", "default");
+        ASSERT_TRUE(r.has_value());
+    }
+    EXPECT_EQ(mm.get_promotion_admitted() - admitted_pre, 1);
+    EXPECT_EQ(mm.get_promotion_in_flight() - in_flight_pre, 1);
+
+    // Drive AllocStart + NotifyPromotionSuccess.
+    auto alloc = service->PromotionAllocStart(seg.client_id, "k_hot", "default",
+                                              kObjBytes, {});
+    ASSERT_TRUE(alloc.has_value());
+    auto notify =
+        service->NotifyPromotionSuccess(seg.client_id, "k_hot", "default");
+    ASSERT_TRUE(notify.has_value());
+
+    EXPECT_EQ(mm.get_promotion_completed() - completed_pre, 1);
+    EXPECT_EQ(mm.get_promotion_completed_bytes() - bytes_pre, kObjBytes);
+    EXPECT_EQ(mm.get_promotion_in_flight() - in_flight_pre, 0);
+
+    service->RemoveAll();
+}
+
+// Each rejection gate (frequency / watermark / cap) increments its
+// own counter when its branch fires.
+TEST_F(PromotionOnHitTest, MetricsRejectionCountersIncrementOnGateMiss) {
+    MasterServiceConfig config;
+    config.enable_offload = true;
+    config.promotion_on_hit = true;
+    // Threshold > 1 so the first Get is rejected on frequency.
+    config.promotion_admission_threshold = 2;
+    config.promotion_queue_limit = 1;
+    config.default_kv_lease_ttl = 2000;
+    auto service = std::make_unique<MasterService>(config);
+
+    constexpr size_t seg_size = 1024 * 1024 * 16;
+    auto seg = PrepareSegment(*service, "seg_a", kDefaultSegmentBase, seg_size);
+    ASSERT_TRUE(InjectLocalDiskReplica(*service, seg.client_id, "k_a", 1024,
+                                       seg.segment_name));
+    ASSERT_TRUE(InjectLocalDiskReplica(*service, seg.client_id, "k_b", 1024,
+                                       seg.segment_name));
+
+    auto& mm = MasterMetricManager::instance();
+    const int64_t freq_pre = mm.get_promotion_rejected_frequency();
+    const int64_t cap_pre = mm.get_promotion_rejected_cap();
+
+    // 1st Get on k_a: freq=1, threshold=2 → rejected on frequency.
+    {
+        auto r = service->GetReplicaList("k_a", "default");
+        ASSERT_TRUE(r.has_value());
+    }
+    EXPECT_EQ(mm.get_promotion_rejected_frequency() - freq_pre, 1);
+
+    // 2nd Get on k_a: freq=2, admits. Now in-flight = 1 == limit.
+    {
+        auto r = service->GetReplicaList("k_a", "default");
+        ASSERT_TRUE(r.has_value());
+    }
+    // Get on k_b: freq=1, threshold=2 → rejected on frequency.
+    {
+        auto r = service->GetReplicaList("k_b", "default");
+        ASSERT_TRUE(r.has_value());
+    }
+    // Get on k_b again: freq=2, gets past frequency, but cap=1
+    // saturated → rejected on cap.
+    {
+        auto r = service->GetReplicaList("k_b", "default");
+        ASSERT_TRUE(r.has_value());
+    }
+    EXPECT_EQ(mm.get_promotion_rejected_cap() - cap_pre, 1);
+
+    service->RemoveAll();
+
+    // Watermark gate uses a fresh service configured with
+    // eviction_high_watermark_ratio = 0.0 so any non-negative DRAM
+    // usage trips it. threshold = 1 makes the first Get clear the
+    // frequency gate and reach the watermark check.
+    MasterServiceConfig wm_config;
+    wm_config.enable_offload = true;
+    wm_config.promotion_on_hit = true;
+    wm_config.promotion_admission_threshold = 1;
+    wm_config.promotion_queue_limit = 50;
+    wm_config.eviction_high_watermark_ratio = 0.0;
+    wm_config.default_kv_lease_ttl = 2000;
+    auto wm_service = std::make_unique<MasterService>(wm_config);
+
+    auto wm_seg =
+        PrepareSegment(*wm_service, "wm_seg", kDefaultSegmentBase, seg_size);
+    ASSERT_TRUE(InjectLocalDiskReplica(*wm_service, wm_seg.client_id, "k_w",
+                                       1024, wm_seg.segment_name));
+
+    const int64_t wm_pre = mm.get_promotion_rejected_watermark();
+    {
+        auto r = wm_service->GetReplicaList("k_w", "default");
+        ASSERT_TRUE(r.has_value());
+    }
+    EXPECT_EQ(mm.get_promotion_rejected_watermark() - wm_pre, 1);
+
+    wm_service->RemoveAll();
+}
+
+TEST_F(PromotionOnHitTest, AdmissionFrequencyIsTenantScoped) {
+    MasterServiceConfig config;
+    config.enable_offload = true;
+    config.enable_multi_tenants = true;
+    config.tenant_quota_connector_type = "file";
+    config.promotion_on_hit = true;
+    config.promotion_admission_threshold = 2;
+    config.default_kv_lease_ttl = 2000;
+    const std::string key = "shared_hot_key";
+    const std::string tenant_a = "tenant_promotion_a";
+    const std::string tenant_b = "tenant_promotion_b";
+    config.tenant_quota_connector_uri = WriteTenantQuotaPolicyFile(
+        {{tenant_a, 64 * 1024 * 1024}, {tenant_b, 64 * 1024 * 1024}});
+    auto service = std::make_unique<MasterService>(config);
+
+    constexpr size_t seg_size = 1024 * 1024 * 16;
+    auto seg =
+        PrepareSegment(*service, "seg_tenant", kDefaultSegmentBase, seg_size);
+    ASSERT_TRUE(InjectLocalDiskReplica(*service, seg.client_id, key, 1024,
+                                       seg.segment_name, tenant_a));
+    ASSERT_TRUE(InjectLocalDiskReplica(*service, seg.client_id, key, 1024,
+                                       seg.segment_name, tenant_b));
+
+    {
+        auto r = service->GetReplicaList(key, tenant_a);
+        ASSERT_TRUE(r.has_value());
+    }
+    {
+        auto r = service->GetReplicaList(key, tenant_b);
+        ASSERT_TRUE(r.has_value());
+    }
+    auto pending_after_one_each =
+        service->PromotionObjectHeartbeat(seg.client_id);
+    ASSERT_TRUE(pending_after_one_each.has_value());
+    EXPECT_TRUE(pending_after_one_each->empty())
+        << "same user key in two tenants must not share promotion heat";
+
+    {
+        auto r = service->GetReplicaList(key, tenant_a);
+        ASSERT_TRUE(r.has_value());
+    }
+    auto pending_after_tenant_a_second =
+        service->PromotionObjectHeartbeat(seg.client_id);
+    ASSERT_TRUE(pending_after_tenant_a_second.has_value());
+    ASSERT_EQ(pending_after_tenant_a_second->size(), 1u);
+    EXPECT_EQ((*pending_after_tenant_a_second)[0].tenant_id, tenant_a);
+    EXPECT_EQ((*pending_after_tenant_a_second)[0].key, key);
+
+    service->RemoveAll();
+}
+
+// promotion_max_per_heartbeat controls how many tasks
+// PromotionObjectHeartbeat returns per call. Set the knob to 3 and
+// verify the master returns up to 3 tasks per heartbeat.
+TEST_F(PromotionOnHitTest, MaxPerHeartbeatKnobControlsBatchSize) {
+    MasterServiceConfig config;
+    config.enable_offload = true;
+    config.promotion_on_hit = true;
+    config.promotion_admission_threshold = 1;
+    config.promotion_max_per_heartbeat = 3;  // raise from default 1
+    config.default_kv_lease_ttl = 2000;
+    auto service = std::make_unique<MasterService>(config);
+
+    constexpr size_t seg_size = 1024 * 1024 * 16;
+    auto seg = PrepareSegment(*service, "seg_a", kDefaultSegmentBase, seg_size);
+
+    // Admit 5 tasks on the same holder.
+    for (int i = 0; i < 5; ++i) {
+        const auto key = "k_" + std::to_string(i);
+        ASSERT_TRUE(InjectLocalDiskReplica(*service, seg.client_id, key, 1024,
+                                           seg.segment_name));
+        auto r = service->GetReplicaList(key, "default");
+        ASSERT_TRUE(r.has_value());
+    }
+
+    // First heartbeat: cap=3.
+    auto first = service->PromotionObjectHeartbeat(seg.client_id);
+    ASSERT_TRUE(first.has_value());
+    EXPECT_EQ(first->size(), 3u)
+        << "expected promotion_max_per_heartbeat=3 keys, got " << first->size();
+
+    // Second heartbeat: 2 leftover.
+    auto second = service->PromotionObjectHeartbeat(seg.client_id);
+    ASSERT_TRUE(second.has_value());
+    EXPECT_EQ(second->size(), 2u);
+
+    // Third heartbeat: empty.
+    auto third = service->PromotionObjectHeartbeat(seg.client_id);
+    ASSERT_TRUE(third.has_value());
+    EXPECT_EQ(third->size(), 0u);
+
+    service->RemoveAll();
+}
+
+// promotion_max_per_heartbeat = 0 must be clamped to 1 by the
+// MasterService constructor; otherwise PromotionObjectHeartbeat would
+// return an empty batch every call and silently disable promotion
+// delivery.
+TEST_F(PromotionOnHitTest, MaxPerHeartbeatZeroClampsToOne) {
+    MasterServiceConfig config;
+    config.enable_offload = true;
+    config.promotion_on_hit = true;
+    config.promotion_admission_threshold = 1;
+    config.promotion_max_per_heartbeat = 0;  // pathological
+    config.default_kv_lease_ttl = 2000;
+    auto service = std::make_unique<MasterService>(config);
+
+    constexpr size_t seg_size = 1024 * 1024 * 16;
+    auto seg = PrepareSegment(*service, "seg_a", kDefaultSegmentBase, seg_size);
+    ASSERT_TRUE(InjectLocalDiskReplica(*service, seg.client_id, "k1", 1024,
+                                       seg.segment_name));
+    {
+        auto r = service->GetReplicaList("k1", "default");
+        ASSERT_TRUE(r.has_value());
+    }
+
+    // If max_per_heartbeat=0 weren't clamped, this would return empty.
+    auto hb = service->PromotionObjectHeartbeat(seg.client_id);
+    ASSERT_TRUE(hb.has_value());
+    EXPECT_EQ(hb->size(), 1u)
+        << "max_per_heartbeat=0 must clamp to 1 so promotion delivery "
+        << "isn't silently disabled";
+
+    service->RemoveAll();
+}
+
+// Remove on a key mid-promotion must bump promotion_cancelled so the
+// funnel invariant
+// admitted = completed + failed + expired + cancelled + in_flight
+// holds. Exercises the EraseMetadataEntry path.
+TEST_F(PromotionOnHitTest, MetricsRemoveMidPromotionCountsAsCancelled) {
+    MasterServiceConfig config;
+    config.enable_offload = true;
+    config.promotion_on_hit = true;
+    config.promotion_admission_threshold = 1;
+    config.default_kv_lease_ttl = 2000;
+    auto service = std::make_unique<MasterService>(config);
+
+    constexpr size_t seg_size = 1024 * 1024 * 16;
+    auto seg = PrepareSegment(*service, "seg_a", kDefaultSegmentBase, seg_size);
+    ASSERT_TRUE(InjectLocalDiskReplica(*service, seg.client_id, "k_drop", 1024,
+                                       seg.segment_name));
+
+    auto& mm = MasterMetricManager::instance();
+    const int64_t admitted_pre = mm.get_promotion_admitted();
+    const int64_t cancelled_pre = mm.get_promotion_cancelled();
+    const int64_t in_flight_pre = mm.get_promotion_in_flight();
+
+    // Admit a promotion. in_flight goes from 0 to 1.
+    {
+        auto r = service->GetReplicaList("k_drop", "default");
+        ASSERT_TRUE(r.has_value());
+    }
+    ASSERT_EQ(mm.get_promotion_admitted() - admitted_pre, 1);
+    ASSERT_EQ(mm.get_promotion_in_flight() - in_flight_pre, 1);
+
+    auto rm = service->Remove("k_drop", "default", /*force=*/true);
+    ASSERT_TRUE(rm.has_value()) << "error=" << rm.error();
+
+    EXPECT_EQ(mm.get_promotion_in_flight() - in_flight_pre, 0);
+    EXPECT_EQ(mm.get_promotion_cancelled() - cancelled_pre, 1);
+
+    service->RemoveAll();
+}
+
+}  // namespace mooncake::test
+
+int main(int argc, char** argv) {
+    ::testing::InitGoogleTest(&argc, argv);
+    return RUN_ALL_TESTS();
+}
diff --git a/mooncake-store/tests/pybind_client_test.cpp b/mooncake-store/tests/pybind_client_test.cpp
index b9bc26b6..53492a5d 100644
--- a/mooncake-store/tests/pybind_client_test.cpp
+++ b/mooncake-store/tests/pybind_client_test.cpp
@@ -2,10 +2,17 @@
 #include <glog/logging.h>
 #include <gtest/gtest.h>
 
-#include <memory>
-#include <string>
-#include <random>
 #include <barrier>
+#include <chrono>
+#include <cstdio>
+#include <cstring>
+#include <fcntl.h>
+#include <limits>
+#include <memory>
+#include <random>
+#include <string>
+#include <thread>
+#include <unistd.h>
 
 #include "real_client.h"
 #include "test_server_helpers.h"
@@ -63,8 +70,165 @@ class RealClientTest : public ::testing::Test {
     // In-proc master for tests
     mooncake::testing::InProcMaster master_;
     std::string master_address_;
+
+    void StartMasterAndSetupClient() {
+        ASSERT_TRUE(master_.Start(InProcMasterConfigBuilder().build()))
+            << "Failed to start in-proc master";
+        master_address_ = master_.master_address();
+
+        const std::string rdma_devices = (FLAGS_protocol == std::string("rdma"))
+                                             ? FLAGS_device_name
+                                             : std::string("");
+        ASSERT_EQ(py_client_->setup_real("localhost:17813", "P2PHANDSHAKE",
+                                         16 * 1024 * 1024, 16 * 1024 * 1024,
+                                         FLAGS_protocol, rdma_devices,
+                                         master_address_),
+                  0);
+    }
+
+    ConfigDict MakeConfigDict(const std::string& local_hostname,
+                              const std::string& global_segment_size,
+                              const std::string& local_buffer_size) const {
+        const std::string rdma_devices = (FLAGS_protocol == std::string("rdma"))
+                                             ? FLAGS_device_name
+                                             : std::string("");
+        ConfigDict config;
+        config[CONFIG_KEY_LOCAL_HOSTNAME] = local_hostname;
+        config[CONFIG_KEY_METADATA_SERVER] = "P2PHANDSHAKE";
+        config[CONFIG_KEY_GLOBAL_SEGMENT_SIZE] = global_segment_size;
+        config[CONFIG_KEY_LOCAL_BUFFER_SIZE] = local_buffer_size;
+        config[CONFIG_KEY_PROTOCOL] = FLAGS_protocol;
+        config[CONFIG_KEY_RDMA_DEVICES] = rdma_devices;
+        config[CONFIG_KEY_MASTER_SERVER_ADDR] = master_address_;
+        return config;
+    }
+
+    std::string CreateTempSegmentFile(size_t size) {
+        std::string path = "/tmp/mooncake_real_client_segment_XXXXXX";
+        int fd = mkstemp(path.data());
+        EXPECT_GE(fd, 0) << "Failed to create temp segment file";
+        if (fd < 0) {
+            return "";
+        }
+        EXPECT_EQ(ftruncate(fd, size), 0)
+            << "Failed to resize temp segment file";
+        close(fd);
+        return path;
+    }
 };
 
+TEST_F(RealClientTest, AllocateAndMountSegmentAlignsAndUnmounts) {
+    StartMasterAndSetupClient();
+
+    const size_t slab_size = facebook::cachelib::Slab::kSize;
+    std::vector<std::string> segment_ids;
+    size_t allocated_size = 0;
+
+    ASSERT_EQ(py_client_->allocateAndMountSegment(1, FLAGS_protocol, "",
+                                                  segment_ids, &allocated_size),
+              0);
+    EXPECT_EQ(allocated_size, slab_size);
+    ASSERT_FALSE(segment_ids.empty());
+    EXPECT_EQ(py_client_->unmountAndFreeSegment(segment_ids), 0);
+
+    segment_ids.clear();
+    allocated_size = 0;
+    ASSERT_EQ(
+        py_client_->allocateAndMountSegment(slab_size + 1, FLAGS_protocol, "",
+                                            segment_ids, &allocated_size),
+        0);
+    EXPECT_EQ(allocated_size, slab_size * 2);
+    ASSERT_FALSE(segment_ids.empty());
+    EXPECT_EQ(py_client_->unmountAndFreeSegment(segment_ids), 0);
+}
+
+TEST_F(RealClientTest, AllocateAndMountSegmentRejectsOverflowSize) {
+    StartMasterAndSetupClient();
+
+    std::vector<std::string> segment_ids;
+    size_t allocated_size = 0;
+
+    EXPECT_NE(py_client_->allocateAndMountSegment(
+                  std::numeric_limits<size_t>::max(), FLAGS_protocol, "",
+                  segment_ids, &allocated_size),
+              0);
+    EXPECT_TRUE(segment_ids.empty());
+    EXPECT_EQ(allocated_size, 0);
+}
+
+TEST_F(RealClientTest, AllocateAndMountSegmentFreesOnTearDown) {
+    StartMasterAndSetupClient();
+
+    std::vector<std::string> segment_ids;
+    size_t allocated_size = 0;
+    ASSERT_EQ(py_client_->allocateAndMountSegment(1, FLAGS_protocol, "",
+                                                  segment_ids, &allocated_size),
+              0);
+    ASSERT_FALSE(segment_ids.empty());
+    EXPECT_GT(allocated_size, 0);
+
+    EXPECT_EQ(py_client_->tearDownAll(), 0);
+
+    GLogMuter muter;
+    EXPECT_NE(py_client_->unmountAndFreeSegment(segment_ids), 0);
+}
+
+TEST_F(RealClientTest, MountAndAllocateUnmountApisRejectForeignSegments) {
+    StartMasterAndSetupClient();
+
+    const size_t slab_size = facebook::cachelib::Slab::kSize;
+    std::vector<std::string> allocated_segment_ids;
+    size_t allocated_size = 0;
+
+    ASSERT_EQ(
+        py_client_->allocateAndMountSegment(
+            1, FLAGS_protocol, "", allocated_segment_ids, &allocated_size),
+        0);
+    ASSERT_FALSE(allocated_segment_ids.empty());
+    EXPECT_NE(py_client_->unmountSegment(allocated_segment_ids), 0);
+    EXPECT_EQ(py_client_->unmountAndFreeSegment(allocated_segment_ids), 0);
+
+    std::string path = CreateTempSegmentFile(slab_size);
+    ASSERT_FALSE(path.empty());
+
+    std::vector<std::string> mounted_segment_ids;
+    ASSERT_EQ(py_client_->mountSegment(path, 0, slab_size, FLAGS_protocol, "",
+                                       mounted_segment_ids),
+              0);
+    ASSERT_FALSE(mounted_segment_ids.empty());
+    EXPECT_NE(py_client_->unmountAndFreeSegment(mounted_segment_ids), 0);
+    EXPECT_EQ(py_client_->unmountSegment(mounted_segment_ids), 0);
+
+    EXPECT_EQ(std::remove(path.c_str()), 0);
+}
+
+TEST_F(RealClientTest, MountAndAllocateUnmountApisAcceptGracePeriod) {
+    StartMasterAndSetupClient();
+
+    const size_t slab_size = facebook::cachelib::Slab::kSize;
+    std::vector<std::string> allocated_segment_ids;
+    size_t allocated_size = 0;
+    ASSERT_EQ(
+        py_client_->allocateAndMountSegment(
+            1, FLAGS_protocol, "", allocated_segment_ids, &allocated_size),
+        0);
+    ASSERT_FALSE(allocated_segment_ids.empty());
+    EXPECT_EQ(py_client_->unmountAndFreeSegment(allocated_segment_ids, 1), 0);
+
+    std::string path = CreateTempSegmentFile(slab_size);
+    ASSERT_FALSE(path.empty());
+
+    std::vector<std::string> mounted_segment_ids;
+    ASSERT_EQ(py_client_->mountSegment(path, 0, slab_size, FLAGS_protocol, "",
+                                       mounted_segment_ids),
+              0);
+    ASSERT_FALSE(mounted_segment_ids.empty());
+    EXPECT_EQ(py_client_->unmountSegment(mounted_segment_ids, 1), 0);
+
+    EXPECT_EQ(py_client_->tearDownAll(), 0);
+    EXPECT_EQ(std::remove(path.c_str()), 0);
+}
+
 // Test basic Put and Get operations
 TEST_F(RealClientTest, BasicPutGetOperations) {
     // Start in-proc master
@@ -111,6 +275,39 @@ TEST_F(RealClientTest, BasicPutGetOperations) {
     EXPECT_EQ(exist_result, 1) << "Key should exist";
 }
 
+TEST_F(RealClientTest, GetIntoAcceptsSubrangeOfLocalRegisteredBuffer) {
+    ASSERT_TRUE(master_.Start(InProcMasterConfigBuilder().build()))
+        << "Failed to start in-proc master";
+    master_address_ = master_.master_address();
+
+    const std::string rdma_devices = (FLAGS_protocol == std::string("rdma"))
+                                         ? FLAGS_device_name
+                                         : std::string("");
+    ASSERT_EQ(
+        py_client_->setup_real("localhost:17813", "P2PHANDSHAKE",
+                               16 * 1024 * 1024, 16 * 1024 * 1024,
+                               FLAGS_protocol, rdma_devices, master_address_),
+        0);
+
+    const std::string key = "local_buffer_subrange_key";
+    const std::string test_data = "shared-local-buffer-read";
+    std::span<const char> data_span(test_data.data(), test_data.size());
+    ReplicateConfig config;
+    config.replica_num = 1;
+    ASSERT_EQ(py_client_->put(key, data_span, config), 0);
+
+    auto local_alloc =
+        py_client_->client_buffer_allocator_->allocate(test_data.size() + 32);
+    ASSERT_TRUE(local_alloc.has_value());
+    BufferHandle handle = std::move(*local_alloc);
+    auto* dst = static_cast<char*>(handle.ptr()) + 16;
+    std::memset(handle.ptr(), 0, handle.size());
+
+    auto bytes_read = py_client_->get_into(key, dst, test_data.size());
+    ASSERT_EQ(bytes_read, static_cast<int64_t>(test_data.size()));
+    EXPECT_EQ(std::string(dst, test_data.size()), test_data);
+}
+
 // Test Get Operation will fail if the lease has expired.
 // Set the lease time to 1ms and use large data size to ensure the lease will
 // expire.
@@ -694,16 +891,31 @@ TEST_F(RealClientTest, TestCopyMoveQueryTask) {
     config.preferred_segment = client1_addr;
     ASSERT_EQ(py_client_->put(key, data_span, config), 0);
 
+    auto wait_for_task = [this](const UUID& task_id) {
+        const auto deadline =
+            std::chrono::steady_clock::now() + std::chrono::seconds(10);
+        while (std::chrono::steady_clock::now() < deadline) {
+            auto query_res = py_client_->query_task(task_id);
+            if (query_res.has_value() &&
+                is_finished_status(query_res->status)) {
+                return query_res;
+            }
+            std::this_thread::sleep_for(std::chrono::milliseconds(100));
+        }
+        return py_client_->query_task(task_id);
+    };
+
     // Test create copy task from client 1 to client 2
     auto copy_res = py_client_->create_copy_task(key, {client2_addr});
     ASSERT_TRUE(copy_res.has_value()) << "Copy should return a task ID";
     UUID copy_task_id = copy_res.value();
 
-    // Query Copy Task
-    auto query_copy_res = py_client_->query_task(copy_task_id);
+    auto query_copy_res = wait_for_task(copy_task_id);
     ASSERT_TRUE(query_copy_res.has_value()) << "QueryTask should succeed";
     EXPECT_EQ(query_copy_res->id, copy_task_id);
     EXPECT_EQ(query_copy_res->type, TaskType::REPLICA_COPY);
+    EXPECT_EQ(query_copy_res->status, TaskStatus::SUCCESS)
+        << query_copy_res->message;
 
     // Test create move task from client 1 to client 2
     auto move_res =
@@ -711,11 +923,12 @@ TEST_F(RealClientTest, TestCopyMoveQueryTask) {
     ASSERT_TRUE(move_res.has_value()) << "Move should return a task ID";
     UUID move_task_id = move_res.value();
 
-    // Query Move Task
-    auto query_move_res = py_client_->query_task(move_task_id);
+    auto query_move_res = wait_for_task(move_task_id);
     ASSERT_TRUE(query_move_res.has_value()) << "QueryTask should succeed";
     EXPECT_EQ(query_move_res->id, move_task_id);
     EXPECT_EQ(query_move_res->type, TaskType::REPLICA_MOVE);
+    EXPECT_EQ(query_move_res->status, TaskStatus::SUCCESS)
+        << query_move_res->message;
 
     py_client2->tearDownAll();
 }
@@ -727,22 +940,12 @@ TEST_F(RealClientTest, SetupWithConfigDict) {
     master_address_ = master_.master_address();
     LOG(INFO) << "Started in-proc master at " << master_address_;
 
-    // Setup the client using ConfigDict
-    const std::string rdma_devices = (FLAGS_protocol == std::string("rdma"))
-                                         ? FLAGS_device_name
-                                         : std::string("");
-
     ConfigDict config;
     auto result = py_client_->setup_internal(config);
     ASSERT_FALSE(result.has_value()) << "Setup with empty config should fail";
 
-    config[CONFIG_KEY_LOCAL_HOSTNAME] = "localhost:17813";
-    config[CONFIG_KEY_METADATA_SERVER] = "P2PHANDSHAKE";
-    config[CONFIG_KEY_GLOBAL_SEGMENT_SIZE] = std::to_string(16 * 1024 * 1024);
-    config[CONFIG_KEY_LOCAL_BUFFER_SIZE] = std::to_string(16 * 1024 * 1024);
-    config[CONFIG_KEY_PROTOCOL] = FLAGS_protocol;
-    config[CONFIG_KEY_RDMA_DEVICES] = rdma_devices;
-    config[CONFIG_KEY_MASTER_SERVER_ADDR] = master_address_;
+    config = MakeConfigDict("localhost:17813", std::to_string(16 * 1024 * 1024),
+                            std::to_string(16 * 1024 * 1024));
 
     result = py_client_->setup_internal(config);
     ASSERT_TRUE(result.has_value()) << "Setup with ConfigDict should succeed";
@@ -767,6 +970,57 @@ TEST_F(RealClientTest, SetupWithConfigDict) {
     EXPECT_EQ(retrieved_data, test_data) << "Retrieved data should match";
 }
 
+TEST_F(RealClientTest, SetupWithConfigDictHumanReadableSizes) {
+    ASSERT_TRUE(master_.Start(InProcMasterConfigBuilder().build()))
+        << "Failed to start in-proc master";
+    master_address_ = master_.master_address();
+
+    ConfigDict config = MakeConfigDict("localhost:17814", "16MB", "16 MB");
+    auto result = py_client_->setup_internal(config);
+    ASSERT_TRUE(result.has_value())
+        << "Setup should accept human-readable size strings";
+}
+
+TEST_F(RealClientTest, SetupWithConfigDictAllowsZeroSizes) {
+    ASSERT_TRUE(master_.Start(InProcMasterConfigBuilder().build()))
+        << "Failed to start in-proc master";
+    master_address_ = master_.master_address();
+
+    ConfigDict config = MakeConfigDict("localhost:17815", "0", "0");
+    auto result = py_client_->setup_internal(config);
+    ASSERT_TRUE(result.has_value())
+        << "Setup should preserve zero-size pure client/server semantics";
+}
+
+TEST_F(RealClientTest, ErrSetupWithInvalidConfigDictSize) {
+    GLogMuter muter;
+    ASSERT_TRUE(master_.Start(InProcMasterConfigBuilder().build()))
+        << "Failed to start in-proc master";
+    master_address_ = master_.master_address();
+
+    struct InvalidSizeCase {
+        const char* local_hostname;
+        const char* global_segment_size;
+        const char* local_buffer_size;
+    };
+
+    const InvalidSizeCase invalid_size_cases[] = {
+        {"localhost:17816", "50%", "16MB"},
+        {"localhost:17817", "16MB", "16XB"},
+        {"localhost:17818", "-5", "16MB"},
+    };
+
+    for (const auto& test_case : invalid_size_cases) {
+        ConfigDict config = MakeConfigDict(test_case.local_hostname,
+                                           test_case.global_segment_size,
+                                           test_case.local_buffer_size);
+        auto result = py_client_->setup_internal(config);
+        EXPECT_FALSE(result.has_value())
+            << "Invalid explicit size values should fail instead of being "
+               "partially parsed or silently defaulted";
+    }
+}
+
 TEST_F(RealClientTest, ErrSetupWithInvalidArgument) {
     GLogMuter muter;
     // Case 1: Setup with unreachable master address
@@ -1221,6 +1475,325 @@ TEST_F(RealClientTest, UpsertBatch) {
     }
 }
 
+// ===================== Batch Existence Tests =====================
+
+TEST_F(RealClientTest, BatchIsExistMixed) {
+    StartMasterAndSetupClient();
+
+    ReplicateConfig config;
+    config.replica_num = 1;
+
+    const std::string data = "batch_exist_data";
+    std::span<const char> data_span(data.data(), data.size());
+    ASSERT_EQ(py_client_->put("exist_key_1", data_span, config), 0);
+    ASSERT_EQ(py_client_->put("exist_key_2", data_span, config), 0);
+
+    std::vector<std::string> keys = {"exist_key_1", "missing_key",
+                                     "exist_key_2", "also_missing"};
+    auto results = py_client_->batchIsExist(keys);
+    ASSERT_EQ(results.size(), 4u);
+    EXPECT_EQ(results[0], 1) << "exist_key_1 should exist";
+    EXPECT_EQ(results[1], 0) << "missing_key should not exist";
+    EXPECT_EQ(results[2], 1) << "exist_key_2 should exist";
+    EXPECT_EQ(results[3], 0) << "also_missing should not exist";
+}
+
+TEST_F(RealClientTest, ErrBatchIsExistBeforeSetup) {
+    GLogMuter muter;
+    std::vector<std::string> keys = {"k1", "k2"};
+    auto results = py_client_->batchIsExist(keys);
+    ASSERT_EQ(results.size(), 2u);
+    for (size_t i = 0; i < results.size(); ++i) {
+        EXPECT_LT(results[i], 0)
+            << "batchIsExist[" << i << "] before setup should return negative";
+    }
+}
+
+// ===================== GetSize Tests =====================
+
+TEST_F(RealClientTest, GetSizeBasic) {
+    StartMasterAndSetupClient();
+
+    const std::string data = "getsize_payload_123";
+    const std::string key = "getsize_key";
+    std::span<const char> data_span(data.data(), data.size());
+    ReplicateConfig config;
+    config.replica_num = 1;
+
+    ASSERT_EQ(py_client_->put(key, data_span, config), 0);
+    int64_t size = py_client_->getSize(key);
+    EXPECT_EQ(size, static_cast<int64_t>(data.size()))
+        << "getSize should return exact data length";
+}
+
+TEST_F(RealClientTest, ErrGetSizeBeforeSetup) {
+    GLogMuter muter;
+    EXPECT_LT(py_client_->getSize("any_key"), 0)
+        << "getSize before setup should return negative";
+}
+
+// ===================== RemoveByRegex Tests =====================
+
+TEST_F(RealClientTest, RemoveByRegexBasic) {
+    StartMasterAndSetupClient();
+
+    ReplicateConfig config;
+    config.replica_num = 1;
+    const std::string data = "regex_data";
+    std::span<const char> data_span(data.data(), data.size());
+
+    ASSERT_EQ(py_client_->put("prefix_alpha", data_span, config), 0);
+    ASSERT_EQ(py_client_->put("prefix_beta", data_span, config), 0);
+    ASSERT_EQ(py_client_->put("other_key", data_span, config), 0);
+
+    long removed = py_client_->removeByRegex("^prefix_.*");
+    EXPECT_EQ(removed, 2) << "Should remove exactly the two prefix_ keys";
+
+    EXPECT_EQ(py_client_->isExist("prefix_alpha"), 0);
+    EXPECT_EQ(py_client_->isExist("prefix_beta"), 0);
+    EXPECT_EQ(py_client_->isExist("other_key"), 1)
+        << "Non-matching key should survive";
+}
+
+TEST_F(RealClientTest, RemoveByRegexNoMatch) {
+    StartMasterAndSetupClient();
+
+    ReplicateConfig config;
+    config.replica_num = 1;
+    const std::string data = "no_match_data";
+    std::span<const char> data_span(data.data(), data.size());
+    ASSERT_EQ(py_client_->put("some_key", data_span, config), 0);
+
+    long removed = py_client_->removeByRegex("^nonexistent_pattern_.*");
+    EXPECT_EQ(removed, 0) << "No keys should match";
+    EXPECT_EQ(py_client_->isExist("some_key"), 1)
+        << "Existing key should remain";
+}
+
+TEST_F(RealClientTest, ErrRemoveByRegexBeforeSetup) {
+    GLogMuter muter;
+    long result = py_client_->removeByRegex(".*");
+    EXPECT_LT(result, 0) << "removeByRegex before setup should return negative";
+}
+
+// ===================== RemoveAll Edge Cases =====================
+
+TEST_F(RealClientTest, RemoveAllOnEmptyStore) {
+    StartMasterAndSetupClient();
+    long removed = py_client_->removeAll();
+    EXPECT_EQ(removed, 0) << "removeAll on empty store should return 0";
+}
+
+// ===================== BatchRemove Tests =====================
+
+TEST_F(RealClientTest, BatchRemoveBasic) {
+    StartMasterAndSetupClient();
+
+    ReplicateConfig config;
+    config.replica_num = 1;
+    const std::string data = "batch_rm_data";
+    std::span<const char> data_span(data.data(), data.size());
+
+    ASSERT_EQ(py_client_->put("brm_1", data_span, config), 0);
+    ASSERT_EQ(py_client_->put("brm_2", data_span, config), 0);
+    ASSERT_EQ(py_client_->put("brm_3", data_span, config), 0);
+
+    std::vector<std::string> keys_to_remove = {"brm_1", "brm_3"};
+    auto results = py_client_->batchRemove(keys_to_remove);
+    ASSERT_EQ(results.size(), 2u);
+    for (size_t i = 0; i < results.size(); ++i) {
+        EXPECT_EQ(results[i], 0) << "batchRemove[" << i << "] should succeed";
+    }
+
+    EXPECT_EQ(py_client_->isExist("brm_1"), 0);
+    EXPECT_EQ(py_client_->isExist("brm_2"), 1) << "brm_2 should survive";
+    EXPECT_EQ(py_client_->isExist("brm_3"), 0);
+}
+
+TEST_F(RealClientTest, BatchRemoveNonExistentKeys) {
+    StartMasterAndSetupClient();
+
+    GLogMuter muter;
+    std::vector<std::string> keys = {"never_existed_1", "never_existed_2"};
+    auto results = py_client_->batchRemove(keys);
+    ASSERT_EQ(results.size(), 2u);
+}
+
+TEST_F(RealClientTest, ErrBatchRemoveBeforeSetup) {
+    GLogMuter muter;
+    std::vector<std::string> keys = {"k1", "k2"};
+    auto results = py_client_->batchRemove(keys);
+    ASSERT_EQ(results.size(), 2u);
+    for (size_t i = 0; i < results.size(); ++i) {
+        EXPECT_NE(results[i], 0)
+            << "batchRemove[" << i << "] before setup should fail";
+    }
+}
+
+// ===================== PutParts Tests =====================
+
+TEST_F(RealClientTest, PutPartsBasic) {
+    StartMasterAndSetupClient();
+
+    ReplicateConfig config;
+    config.replica_num = 1;
+    const std::string key = "put_parts_key";
+
+    const std::string part1 = "Hello, ";
+    const std::string part2 = "Parts!";
+    const std::string expected = part1 + part2;
+    std::vector<std::span<const char>> parts;
+    parts.emplace_back(part1.data(), part1.size());
+    parts.emplace_back(part2.data(), part2.size());
+
+    EXPECT_EQ(py_client_->put_parts(key, parts, config), 0);
+
+    auto buf = py_client_->get_buffer(key);
+    ASSERT_NE(buf, nullptr);
+    EXPECT_EQ(buf->size(), expected.size());
+    EXPECT_EQ(std::string(static_cast<const char*>(buf->ptr()), buf->size()),
+              expected);
+}
+
+TEST_F(RealClientTest, ErrPutPartsBeforeSetup) {
+    GLogMuter muter;
+    const std::string part = "data";
+    std::vector<std::span<const char>> parts;
+    parts.emplace_back(part.data(), part.size());
+    ReplicateConfig config;
+    config.replica_num = 1;
+    EXPECT_NE(py_client_->put_parts("key", parts, config), 0)
+        << "put_parts before setup should fail";
+}
+
+// ===================== PutBatch and GetBatch Tests =====================
+
+TEST_F(RealClientTest, PutBatchThenBatchGetBuffer) {
+    StartMasterAndSetupClient();
+
+    ReplicateConfig config;
+    config.replica_num = 1;
+
+    std::vector<std::string> keys = {"batch_kv_0", "batch_kv_1", "batch_kv_2"};
+    std::string val0 = "value_zero";
+    std::string val1 = "value_one!";
+    std::string val2 = "value_two!";
+    std::vector<std::span<const char>> values;
+    values.emplace_back(val0.data(), val0.size());
+    values.emplace_back(val1.data(), val1.size());
+    values.emplace_back(val2.data(), val2.size());
+
+    ASSERT_EQ(py_client_->put_batch(keys, values, config), 0);
+
+    auto handles = py_client_->batch_get_buffer(keys);
+    ASSERT_EQ(handles.size(), 3u);
+
+    std::vector<std::string*> expected = {&val0, &val1, &val2};
+    for (size_t i = 0; i < handles.size(); ++i) {
+        ASSERT_NE(handles[i], nullptr)
+            << "Handle " << i << " should not be null";
+        EXPECT_EQ(std::string(static_cast<const char*>(handles[i]->ptr()),
+                              handles[i]->size()),
+                  *expected[i])
+            << "Data mismatch for key " << keys[i];
+    }
+}
+
+// ===================== HealthCheck Tests =====================
+
+TEST_F(RealClientTest, HealthCheckAfterSetup) {
+    StartMasterAndSetupClient();
+    EXPECT_EQ(py_client_->health_check(), HC_HEALTHY)
+        << "health_check should return HEALTHY after setup";
+}
+
+TEST_F(RealClientTest, HealthCheckBeforeSetup) {
+    int result = py_client_->health_check();
+    EXPECT_EQ(result, HC_NOT_INITIALIZED)
+        << "health_check before setup should return NOT_INITIALIZED";
+}
+
+// ===================== GetHostname Tests =====================
+
+TEST_F(RealClientTest, GetHostnameAfterSetup) {
+    StartMasterAndSetupClient();
+    std::string hostname = py_client_->get_hostname();
+    EXPECT_FALSE(hostname.empty()) << "get_hostname should return non-empty";
+    EXPECT_EQ(hostname, "localhost:17813")
+        << "get_hostname should match the configured hostname";
+}
+
+// ===================== Double TearDown Tests =====================
+
+TEST_F(RealClientTest, DoubleTearDownIsIdempotent) {
+    StartMasterAndSetupClient();
+
+    EXPECT_EQ(py_client_->tearDownAll(), 0) << "First teardown should succeed";
+    EXPECT_EQ(py_client_->tearDownAll(), 0)
+        << "Second teardown should also succeed (idempotent)";
+}
+
+// ===================== Empty Batch Operations =====================
+
+TEST_F(RealClientTest, EmptyBatchOperations) {
+    StartMasterAndSetupClient();
+
+    std::vector<std::string> empty_keys;
+    std::vector<std::span<const char>> empty_values;
+    ReplicateConfig config;
+    config.replica_num = 1;
+
+    EXPECT_EQ(py_client_->put_batch(empty_keys, empty_values, config), 0)
+        << "put_batch with empty input should succeed";
+
+    auto handles = py_client_->batch_get_buffer(empty_keys);
+    EXPECT_TRUE(handles.empty())
+        << "batch_get_buffer with empty input should return empty";
+
+    auto exist_results = py_client_->batchIsExist(empty_keys);
+    EXPECT_TRUE(exist_results.empty())
+        << "batchIsExist with empty input should return empty";
+
+    auto remove_results = py_client_->batchRemove(empty_keys);
+    EXPECT_TRUE(remove_results.empty())
+        << "batchRemove with empty input should return empty";
+}
+
+// ===================== Mount Segment Edge Cases =====================
+
+TEST_F(RealClientTest, ErrMountNonExistentFile) {
+    StartMasterAndSetupClient();
+
+    GLogMuter muter;
+    std::vector<std::string> segment_ids;
+    int ret =
+        py_client_->mountSegment("/tmp/mooncake_nonexistent_file_12345", 0,
+                                 4096, FLAGS_protocol, "", segment_ids);
+    EXPECT_NE(ret, 0) << "Mounting non-existent file should fail";
+    EXPECT_TRUE(segment_ids.empty());
+}
+
+TEST_F(RealClientTest, ErrUnmountInvalidSegmentIds) {
+    StartMasterAndSetupClient();
+
+    GLogMuter muter;
+    std::vector<std::string> bogus_ids = {
+        "00000000-0000-0000-0000-000000000000"};
+    int ret = py_client_->unmountSegment(bogus_ids);
+    EXPECT_NE(ret, 0) << "Unmounting non-existent segment ids should fail";
+}
+
+TEST_F(RealClientTest, ErrUnmountAndFreeInvalidSegmentIds) {
+    StartMasterAndSetupClient();
+
+    GLogMuter muter;
+    std::vector<std::string> bogus_ids = {
+        "00000000-0000-0000-0000-000000000000"};
+    int ret = py_client_->unmountAndFreeSegment(bogus_ids);
+    EXPECT_NE(ret, 0)
+        << "Unmount-and-free of non-existent segment ids should fail";
+}
+
 }  // namespace testing
 
 }  // namespace mooncake
diff --git a/mooncake-store/tests/rpc_timeout_test.cpp b/mooncake-store/tests/rpc_timeout_test.cpp
new file mode 100644
index 00000000..310832b4
--- /dev/null
+++ b/mooncake-store/tests/rpc_timeout_test.cpp
@@ -0,0 +1,120 @@
+// Unit tests for the client->master RPC timeout feature.
+//
+// Covers:
+//   1. The ErrorCode::RPC_TIMEOUT enum value and its toString() mapping.
+//   2. End-to-end: MC_RPC_TIMEOUT_MS shortens the per-request deadline so that
+//      an unresponsive master surfaces ErrorCode::RPC_TIMEOUT (not RPC_FAIL),
+//      and it does so within the configured budget rather than the 30s default.
+//
+// The end-to-end test points a MasterClient at a "black hole" TCP listener: a
+// socket that accepts the connection (so connect() succeeds) but never sends a
+// response. The first RPC issued by Connect() (ServiceReady) therefore blocks
+// until the request timeout fires.
+
+#include <gtest/gtest.h>
+#include <glog/logging.h>
+
+#include <arpa/inet.h>
+#include <netinet/in.h>
+#include <sys/socket.h>
+#include <unistd.h>
+
+#include <chrono>
+#include <cstdlib>
+#include <string>
+
+#include "master_client.h"
+#include "types.h"
+
+namespace mooncake {
+namespace {
+
+// A TCP listener on 127.0.0.1 that accepts connections but never replies.
+// The kernel completes the TCP handshake for backlogged connections, so a
+// client's connect() succeeds; any subsequent request simply never gets a
+// response, which is exactly what we need to exercise the request timeout.
+class BlackHoleServer {
+   public:
+    BlackHoleServer() {
+        fd_ = ::socket(AF_INET, SOCK_STREAM, 0);
+        EXPECT_GE(fd_, 0) << "failed to create socket";
+        int opt = 1;
+        ::setsockopt(fd_, SOL_SOCKET, SO_REUSEADDR, &opt, sizeof(opt));
+
+        sockaddr_in addr{};
+        addr.sin_family = AF_INET;
+        addr.sin_addr.s_addr = ::inet_addr("127.0.0.1");
+        addr.sin_port = 0;  // ephemeral port
+        EXPECT_EQ(::bind(fd_, reinterpret_cast<sockaddr*>(&addr), sizeof(addr)),
+                  0)
+            << "bind failed";
+
+        socklen_t len = sizeof(addr);
+        EXPECT_EQ(::getsockname(fd_, reinterpret_cast<sockaddr*>(&addr), &len),
+                  0);
+        port_ = ntohs(addr.sin_port);
+
+        // Backlog large enough that the kernel handshakes our single client.
+        EXPECT_EQ(::listen(fd_, 16), 0) << "listen failed";
+    }
+
+    ~BlackHoleServer() {
+        if (fd_ >= 0) {
+            ::close(fd_);
+        }
+    }
+
+    std::string address() const { return "127.0.0.1:" + std::to_string(port_); }
+
+   private:
+    int fd_ = -1;
+    uint16_t port_ = 0;
+};
+
+}  // namespace
+
+// Change #1 and #2: the new error code exists with the documented value and a
+// matching human-readable string.
+TEST(RpcTimeoutTest, ErrorCodeValueAndString) {
+    EXPECT_EQ(static_cast<int32_t>(ErrorCode::RPC_TIMEOUT), -901);
+    EXPECT_EQ(toString(ErrorCode::RPC_TIMEOUT), "RPC_TIMEOUT");
+    // Distinct from the generic RPC failure code.
+    EXPECT_NE(ErrorCode::RPC_TIMEOUT, ErrorCode::RPC_FAIL);
+}
+
+// End-to-end: a small MC_RPC_TIMEOUT_MS must be honored by every RPC and turn
+// an unanswered call into ErrorCode::RPC_TIMEOUT well before the 30s default.
+TEST(RpcTimeoutTest, RpcTimesOutAgainstUnresponsiveMaster) {
+    constexpr int kTimeoutMs = 500;
+
+    // The constructor reads MC_RPC_TIMEOUT_MS, so it must be set beforehand.
+    ASSERT_EQ(::setenv("MC_RPC_TIMEOUT_MS", std::to_string(kTimeoutMs).c_str(),
+                       /*overwrite=*/1),
+              0);
+
+    BlackHoleServer server;
+
+    MasterClient client(generate_uuid());
+
+    const auto start = std::chrono::steady_clock::now();
+    ErrorCode rc = client.Connect(server.address());
+    const auto elapsed = std::chrono::duration_cast<std::chrono::milliseconds>(
+                             std::chrono::steady_clock::now() - start)
+                             .count();
+
+    ::unsetenv("MC_RPC_TIMEOUT_MS");
+
+    // The unanswered ServiceReady RPC must be reported as a timeout, not as a
+    // generic RPC failure.
+    EXPECT_EQ(rc, ErrorCode::RPC_TIMEOUT)
+        << "expected RPC_TIMEOUT, got " << toString(rc);
+
+    // It must give up around the configured budget, proving the env override
+    // took effect (the built-in default would be ~30s).
+    EXPECT_GE(elapsed, kTimeoutMs - 100)
+        << "returned too early to be the configured timeout";
+    EXPECT_LT(elapsed, 5000) << "did not honor MC_RPC_TIMEOUT_MS (default 30s "
+                                "timeout still active?)";
+}
+
+}  // namespace mooncake
diff --git a/mooncake-store/tests/storage_backend_test.cpp b/mooncake-store/tests/storage_backend_test.cpp
index db2d7d13..1c66f376 100644
--- a/mooncake-store/tests/storage_backend_test.cpp
+++ b/mooncake-store/tests/storage_backend_test.cpp
@@ -2343,7 +2343,7 @@ TEST_F(StorageBackendTest,
     // Deleter thread: calls DeleteBucket
     std::thread deleter_thread([&]() {
         delete_started.store(true);
-        auto delete_result = storage_backend.DeleteBucket(bucket_id);
+        std::ignore = storage_backend.DeleteBucket(bucket_id);
         delete_completed.store(true);
     });
 
diff --git a/mooncake-store/tests/task_manager_test.cpp b/mooncake-store/tests/task_manager_test.cpp
index ad094519..663a4144 100644
--- a/mooncake-store/tests/task_manager_test.cpp
+++ b/mooncake-store/tests/task_manager_test.cpp
@@ -269,7 +269,7 @@ TEST_F(ClientTaskManagerTest, PruneExpiredTasksProcessingTimeoutFreesSlot) {
 }
 
 TEST_F(ClientTaskManagerTest, SerializerRoundTrip) {
-    ClientTaskManager manager({10000, 10000, 10000, 0, 0});
+    ClientTaskManager manager({10000, 10000, 10000, 0, 0, 0});
     UUID client_id1 = generate_uuid();
     UUID client_id2 = generate_uuid();
 
@@ -277,11 +277,13 @@ TEST_F(ClientTaskManagerTest, SerializerRoundTrip) {
     auto t1 =
         manager.get_write_access().submit_task_typed<TaskType::REPLICA_COPY>(
             client_id1,
-            ReplicaCopyPayload{.key = "pending_key", .targets = {"seg1"}});
+            ReplicaCopyPayload{
+                .key = "pending_key", .source = "", .targets = {"seg1"}});
     auto t2 =
         manager.get_write_access().submit_task_typed<TaskType::REPLICA_COPY>(
             client_id2,
-            ReplicaCopyPayload{.key = "processing_key", .targets = {"seg2"}});
+            ReplicaCopyPayload{
+                .key = "processing_key", .source = "", .targets = {"seg2"}});
     auto t3 =
         manager.get_write_access().submit_task_typed<TaskType::REPLICA_MOVE>(
             client_id1, ReplicaMovePayload{
@@ -303,9 +305,9 @@ TEST_F(ClientTaskManagerTest, SerializerRoundTrip) {
     manager.get_write_access().pop_tasks(client_id2, 2);
 
     // Complete t1 as SUCCESS, t2 as FAILED
-    auto ec1 = manager.get_write_access().complete_task(
+    std::ignore = manager.get_write_access().complete_task(
         client_id1, t1.value(), TaskStatus::SUCCESS, "Done");
-    auto ec2 = manager.get_write_access().complete_task(
+    std::ignore = manager.get_write_access().complete_task(
         client_id2, t2.value(), TaskStatus::FAILED, "Failed");
 
     // Get original tasks
@@ -320,7 +322,7 @@ TEST_F(ClientTaskManagerTest, SerializerRoundTrip) {
     ASSERT_TRUE(serialized.has_value());
 
     // Deserialize into new manager
-    ClientTaskManager manager2({10000, 10000, 10000, 0, 0});
+    ClientTaskManager manager2({10000, 10000, 10000, 0, 0, 0});
     TaskManagerSerializer serializer2(&manager2);
     auto result = serializer2.Deserialize(serialized.value());
     ASSERT_TRUE(result.has_value());
@@ -377,13 +379,13 @@ TEST_F(ClientTaskManagerTest, SerializerRoundTrip) {
 }
 
 TEST_F(ClientTaskManagerTest, SerializerEmptyManager) {
-    ClientTaskManager manager({10000, 10000, 10000, 0, 0});
+    ClientTaskManager manager({10000, 10000, 10000, 0, 0, 0});
 
     TaskManagerSerializer serializer(&manager);
     auto serialized = serializer.Serialize();
     ASSERT_TRUE(serialized.has_value());
 
-    ClientTaskManager manager2({10000, 10000, 10000, 0, 0});
+    ClientTaskManager manager2({10000, 10000, 10000, 0, 0, 0});
     TaskManagerSerializer serializer2(&manager2);
     auto result = serializer2.Deserialize(serialized.value());
     ASSERT_TRUE(result.has_value());
@@ -394,12 +396,13 @@ TEST_F(ClientTaskManagerTest, SerializerEmptyManager) {
 }
 
 TEST_F(ClientTaskManagerTest, SerializerReset) {
-    ClientTaskManager manager({10000, 10000, 10000, 0, 0});
+    ClientTaskManager manager({10000, 10000, 10000, 0, 0, 0});
     UUID client_id = generate_uuid();
 
     auto t1 =
         manager.get_write_access().submit_task_typed<TaskType::REPLICA_COPY>(
-            client_id, ReplicaCopyPayload{.key = "key1", .targets = {"seg1"}});
+            client_id, ReplicaCopyPayload{
+                           .key = "key1", .source = "", .targets = {"seg1"}});
     ASSERT_TRUE(t1.has_value());
 
     TaskManagerSerializer serializer(&manager);
diff --git a/mooncake-store/tests/tenant_quota_test.cpp b/mooncake-store/tests/tenant_quota_test.cpp
new file mode 100644
index 00000000..7faba55a
--- /dev/null
+++ b/mooncake-store/tests/tenant_quota_test.cpp
@@ -0,0 +1,260 @@
+#include "tenant_quota.h"
+#include "tenant_quota_policy_store.h"
+#include "types.h"
+
+#include <filesystem>
+#include <fstream>
+#include <limits>
+#include <map>
+#include <string>
+#include <vector>
+
+#include <gtest/gtest.h>
+#include <unistd.h>
+
+namespace mooncake {
+namespace {
+
+TenantQuotaSnapshot Snapshot(const TenantQuotaTable& table,
+                             const std::string& tenant_id) {
+    auto snapshot = table.GetTenantSnapshot(tenant_id);
+    EXPECT_TRUE(snapshot.has_value());
+    return *snapshot;
+}
+
+uint64_t SumEffectiveQuotas(const TenantQuotaTable& table) {
+    uint64_t sum = 0;
+    for (const auto& snapshot : table.ListTenantSnapshots()) {
+        sum += snapshot.effective_quota_bytes;
+    }
+    return sum;
+}
+
+std::filesystem::path MakeTempPolicyPath(const std::string& suffix) {
+    return std::filesystem::temp_directory_path() /
+           ("mooncake_tenant_quota_policy_store_test_" +
+            std::to_string(::getpid()) + "_" + suffix + ".yaml");
+}
+
+void MakeOrphanTenant(TenantQuotaTable* table, const std::string& tenant_id,
+                      uint64_t bytes) {
+    ASSERT_TRUE(table->UpsertTenantPolicy(tenant_id, bytes).has_value());
+    table->RecomputeEffectiveQuotas(bytes);
+    ASSERT_TRUE(table->Reserve(tenant_id, bytes).has_value());
+    ASSERT_TRUE(table->Commit(tenant_id, bytes).has_value());
+    table->EraseTenantPolicy(tenant_id);
+}
+
+TEST(TenantQuotaTableTest, NormalizesEmptyExplicitTenantIdToDefault) {
+    TenantQuotaTable table;
+
+    ASSERT_TRUE(table.UpsertTenantPolicy("", 1024).has_value());
+    table.RecomputeEffectiveQuotas(4096);
+
+    auto snapshot = Snapshot(table, "");
+    EXPECT_EQ(snapshot.tenant_id, "default");
+    EXPECT_TRUE(snapshot.has_explicit_policy);
+    EXPECT_EQ(snapshot.requested_quota_bytes, 1024);
+    EXPECT_EQ(snapshot.effective_quota_bytes, 1024);
+}
+
+TEST(TenantQuotaTableTest, RejectsZeroExplicitQuotaWithoutChangingState) {
+    TenantQuotaTable table;
+
+    ASSERT_TRUE(table.UpsertTenantPolicy("tenant-a", 100).has_value());
+    auto result = table.UpsertTenantPolicy("tenant-a", 0);
+
+    ASSERT_FALSE(result.has_value());
+    EXPECT_EQ(result.error(), TenantQuotaError::kInvalidArgument);
+    auto snapshot = Snapshot(table, "tenant-a");
+    EXPECT_TRUE(snapshot.has_explicit_policy);
+    EXPECT_EQ(snapshot.requested_quota_bytes, 100);
+}
+
+TEST(TenantQuotaTableTest, EraseExplicitPolicyCreatesOrphanState) {
+    TenantQuotaTable table;
+    MakeOrphanTenant(&table, "tenant-a", 40);
+    table.RecomputeEffectiveQuotas(1000);
+
+    auto snapshot = Snapshot(table, "tenant-a");
+    EXPECT_FALSE(snapshot.has_explicit_policy);
+    EXPECT_EQ(snapshot.requested_quota_bytes, 0);
+    EXPECT_EQ(snapshot.effective_quota_bytes, 0);
+    EXPECT_EQ(snapshot.used_bytes, 40);
+    EXPECT_EQ(snapshot.committed_count, 1);
+    EXPECT_TRUE(snapshot.over_quota);
+}
+
+TEST(TenantQuotaTableTest, EraseMissingPolicyDoesNotCreateLazyState) {
+    TenantQuotaTable table;
+
+    table.EraseTenantPolicy("missing");
+
+    EXPECT_FALSE(table.GetTenantSnapshot("missing").has_value());
+    EXPECT_TRUE(table.ListTenantSnapshots().empty());
+}
+
+TEST(TenantQuotaTableTest, PolicyMutationDoesNotRecomputeEffectiveQuota) {
+    TenantQuotaTable table;
+    ASSERT_TRUE(table.UpsertTenantPolicy("tenant-a", 100).has_value());
+    table.RecomputeEffectiveQuotas(1000);
+    EXPECT_EQ(Snapshot(table, "tenant-a").effective_quota_bytes, 100);
+
+    ASSERT_TRUE(table.UpsertTenantPolicy("tenant-a", 200).has_value());
+    EXPECT_EQ(Snapshot(table, "tenant-a").effective_quota_bytes, 100);
+
+    table.RecomputeEffectiveQuotas(1000);
+    EXPECT_EQ(Snapshot(table, "tenant-a").effective_quota_bytes, 200);
+}
+
+TEST(TenantQuotaTableTest, ListSnapshotsSortedAndSkipsLazyEmptyTenants) {
+    TenantQuotaTable table;
+    ASSERT_TRUE(table.Reserve("z-empty", 0).has_value());
+    ASSERT_TRUE(table.UpsertTenantPolicy("b", 10).has_value());
+    ASSERT_TRUE(table.UpsertTenantPolicy("a", 10).has_value());
+    table.RecomputeEffectiveQuotas(100);
+
+    auto snapshots = table.ListTenantSnapshots();
+    ASSERT_EQ(snapshots.size(), 2);
+    EXPECT_EQ(snapshots[0].tenant_id, "a");
+    EXPECT_EQ(snapshots[1].tenant_id, "b");
+}
+
+TEST(TenantQuotaTableTest, ExplicitTenantsReceiveRequestedWhenCapacityFits) {
+    TenantQuotaTable table;
+    ASSERT_TRUE(table.UpsertTenantPolicy("tenant-a", 100).has_value());
+    ASSERT_TRUE(table.UpsertTenantPolicy("tenant-b", 200).has_value());
+
+    table.RecomputeEffectiveQuotas(1000);
+
+    EXPECT_EQ(Snapshot(table, "tenant-a").effective_quota_bytes, 100);
+    EXPECT_EQ(Snapshot(table, "tenant-b").effective_quota_bytes, 200);
+    EXPECT_EQ(SumEffectiveQuotas(table), 300);
+}
+
+TEST(TenantQuotaTableTest, OverCapacityScalesOnlyExplicitTenants) {
+    TenantQuotaTable table;
+    MakeOrphanTenant(&table, "orphan", 20);
+    ASSERT_TRUE(table.UpsertTenantPolicy("b", 200).has_value());
+    ASSERT_TRUE(table.UpsertTenantPolicy("a", 100).has_value());
+
+    table.RecomputeEffectiveQuotas(150);
+
+    EXPECT_EQ(Snapshot(table, "a").effective_quota_bytes, 50);
+    EXPECT_EQ(Snapshot(table, "b").effective_quota_bytes, 100);
+    EXPECT_EQ(Snapshot(table, "orphan").effective_quota_bytes, 0);
+    EXPECT_TRUE(Snapshot(table, "orphan").over_quota);
+}
+
+TEST(TenantQuotaTableTest, LazyEmptyOrphansDoNotAppearInList) {
+    TenantQuotaTable table;
+    ASSERT_TRUE(table.UpsertTenantPolicy("team-a", 30).has_value());
+    ASSERT_TRUE(table.UpsertTenantPolicy("ghost", 10).has_value());
+    table.EraseTenantPolicy("ghost");
+
+    table.RecomputeEffectiveQuotas(100);
+
+    EXPECT_EQ(Snapshot(table, "team-a").effective_quota_bytes, 30);
+    EXPECT_EQ(Snapshot(table, "ghost").effective_quota_bytes, 0);
+
+    auto snapshots = table.ListTenantSnapshots();
+    ASSERT_EQ(snapshots.size(), 1);
+    EXPECT_EQ(snapshots[0].tenant_id, "team-a");
+}
+
+TEST(TenantQuotaPolicyStoreTest, ParsesValidYamlUnits) {
+    const char* yaml = R"yaml(
+version: 1
+
+tenants:
+  - name: tenant-a
+    quota: 200GB
+  - name: tenant-b
+    quota: 500MB
+  - name: experiment
+    quota: 12345
+)yaml";
+
+    auto snapshot = ParseTenantQuotaPolicyYaml(yaml);
+
+    ASSERT_TRUE(snapshot.has_value()) << snapshot.error();
+    EXPECT_EQ(snapshot->tenant_quotas.at("tenant-a"),
+              200ULL * 1024 * 1024 * 1024);
+    EXPECT_EQ(snapshot->tenant_quotas.at("tenant-b"), 500ULL * 1024 * 1024);
+    EXPECT_EQ(snapshot->tenant_quotas.at("experiment"), 12345);
+}
+
+TEST(TenantQuotaPolicyStoreTest, RejectsInvalidYamlPolicies) {
+    std::vector<std::string> invalid_policies = {
+        "version: 2\n\ntenants: []\n",
+        "version: 1\n\ntenants:\n  - name: tenant-a\n    quota: 1XB\n",
+        "version: 1\n\ntenants:\n  - name: tenant-a\n    quota: 0\n",
+        "version: 1\n\ntenants:\n  - name: \"\"\n    quota: 1KB\n",
+        "version: 1\n\ntenants:\n  - name: _system\n    quota: 1KB\n",
+        "version: 1\n\ntenants:\n  - name: \"tenant\\0bad\"\n    quota: "
+        "1KB\n",
+        "version: 1\n\ntenants:\n  - name: \"tenant\\nline\"\n    quota: "
+        "1KB\n",
+        "version: 1\n\ntenants:\n  - name: \"tenant\\x7f\"\n    quota: 1KB\n",
+        "version: 1\n\ntenants:\n  - name: tenant-a\n    quota: 1KB\n  - name: "
+        "tenant-a\n    quota: 2KB\n",
+        "version: 1\n\ntenants:\n  - name: tenant-a\n    quota: "
+        "18446744073709551616\n",
+        "version: 1\n\ntenants:\n  - name: tenant-a\n    quota: "
+        "18446744073709551615TB\n",
+    };
+
+    for (const auto& policy : invalid_policies) {
+        auto snapshot = ParseTenantQuotaPolicyYaml(policy);
+        EXPECT_FALSE(snapshot.has_value()) << policy;
+    }
+}
+
+TEST(TenantQuotaPolicyStoreTest, RoundTripsYamlFile) {
+    const auto path = MakeTempPolicyPath("roundtrip");
+    std::filesystem::remove(path);
+
+    YamlTenantQuotaPolicyStore store(path.string());
+    TenantQuotaPolicySnapshot snapshot;
+    snapshot.tenant_quotas = {{"tenant-a", 1024}, {"tenant-b", 2048}};
+
+    auto save = store.Save(snapshot);
+    ASSERT_TRUE(save.has_value()) << save.error();
+
+    auto loaded = store.Load();
+    ASSERT_TRUE(loaded.has_value()) << loaded.error();
+    EXPECT_EQ(loaded->tenant_quotas, snapshot.tenant_quotas);
+
+    std::filesystem::remove(path);
+}
+
+TEST(TenantQuotaPolicyStoreTest, RoundTripsYamlSpecialScalarNames) {
+    TenantQuotaPolicySnapshot snapshot;
+    snapshot.tenant_quotas = {{"foo#bar", 1},
+                              {"true", 2},
+                              {"[a, b]", 3},
+                              {"key: val", 4},
+                              {"quote\"slash\\", 5}};
+
+    auto parsed =
+        ParseTenantQuotaPolicyYaml(FormatTenantQuotaPolicyYaml(snapshot));
+
+    ASSERT_TRUE(parsed.has_value()) << parsed.error();
+    EXPECT_EQ(parsed->tenant_quotas, snapshot.tenant_quotas);
+}
+
+TEST(TenantQuotaPolicyStoreTest, SaveFailureReturnsError) {
+    const auto path = MakeTempPolicyPath("missing-dir").parent_path() /
+                      ("missing_dir_" + std::to_string(::getpid())) /
+                      "policy.yaml";
+    YamlTenantQuotaPolicyStore store(path.string());
+
+    TenantQuotaPolicySnapshot snapshot;
+    snapshot.tenant_quotas = {{"tenant-a", 1024}};
+
+    auto save = store.Save(snapshot);
+    EXPECT_FALSE(save.has_value());
+}
+
+}  // namespace
+}  // namespace mooncake
diff --git a/mooncake-store/tests/test_server_helpers.h b/mooncake-store/tests/test_server_helpers.h
index 40b3fd10..765eb442 100644
--- a/mooncake-store/tests/test_server_helpers.h
+++ b/mooncake-store/tests/test_server_helpers.h
@@ -1,8 +1,10 @@
 #pragma once
 
+#include <algorithm>
 #include <cstdlib>
 #include <memory>
 #include <optional>
+#include <vector>
 #include <string>
 #include <thread>
 #include <csignal>
@@ -10,6 +12,7 @@
 #include <ylt/coro_rpc/coro_rpc_server.hpp>
 
 #include "http_metadata_server.h"
+#include "master_admin_service.h"
 #include "master_config.h"
 #include "rpc_service.h"
 #include "types.h"
@@ -35,17 +38,44 @@ class InProcMaster {
             int needed = (!config.rpc_port.has_value()) +
                          (!config.http_metrics_port.has_value()) +
                          (!config.http_metadata_port.has_value());
-            auto free_ports = getFreeTcpPorts(needed);
-            if (static_cast<int>(free_ports.size()) < needed) return false;
+            std::vector<int> available_ports;
+            if (needed > 0) {
+                std::vector<int> reserved_ports;
+                if (config.rpc_port.has_value()) {
+                    reserved_ports.push_back(config.rpc_port.value());
+                }
+                if (config.http_metrics_port.has_value()) {
+                    reserved_ports.push_back(config.http_metrics_port.value());
+                }
+                if (config.http_metadata_port.has_value()) {
+                    reserved_ports.push_back(config.http_metadata_port.value());
+                }
+                auto free_ports = getFreeTcpPorts(
+                    needed + static_cast<int>(reserved_ports.size()));
+                available_ports.reserve(needed);
+                for (int port : free_ports) {
+                    if (std::find(reserved_ports.begin(), reserved_ports.end(),
+                                  port) == reserved_ports.end()) {
+                        available_ports.push_back(port);
+                        if (static_cast<int>(available_ports.size()) ==
+                            needed) {
+                            break;
+                        }
+                    }
+                }
+                if (static_cast<int>(available_ports.size()) < needed) {
+                    return false;
+                }
+            }
             int idx = 0;
             rpc_port_ = config.rpc_port.has_value() ? config.rpc_port.value()
-                                                    : free_ports[idx++];
+                                                    : available_ports[idx++];
             http_metrics_port_ = config.http_metrics_port.has_value()
                                      ? config.http_metrics_port.value()
-                                     : free_ports[idx++];
+                                     : available_ports[idx++];
             http_metadata_port_ = config.http_metadata_port.has_value()
                                       ? config.http_metadata_port.value()
-                                      : free_ports[idx++];
+                                      : available_ports[idx++];
 
             // Optional HTTP metadata server
             if (http_metadata_port_ > 0) {
diff --git a/mooncake-store/tests/transfer_task_test.cpp b/mooncake-store/tests/transfer_task_test.cpp
index caa7c7ad..f841831c 100644
--- a/mooncake-store/tests/transfer_task_test.cpp
+++ b/mooncake-store/tests/transfer_task_test.cpp
@@ -144,11 +144,45 @@ TEST_F(TransferTaskTest, MemcpyWorkerPoolMultipleOperations) {
     }
 }
 
+// Test the locality decision used by TransferSubmitter::isLocalTransfer.
+// Same-host different-process pairs share an IP but have distinct ports;
+// they must NOT be treated as locally addressable, otherwise memcpy in the
+// caller process would dereference a virtual address belonging to a peer
+// process and segfault.
+TEST_F(TransferTaskTest, IsSameProcessEndpoint) {
+    // Empty inputs -> not same-process (cannot prove locality).
+    EXPECT_FALSE(TransferSubmitter::isSameProcessEndpoint("", ""));
+    EXPECT_FALSE(
+        TransferSubmitter::isSameProcessEndpoint("", "192.168.1.10:12345"));
+    EXPECT_FALSE(
+        TransferSubmitter::isSameProcessEndpoint("192.168.1.10:12345", ""));
+
+    // Identical ip:port -> same process.
+    EXPECT_TRUE(TransferSubmitter::isSameProcessEndpoint("192.168.1.10:12345",
+                                                         "192.168.1.10:12345"));
+
+    // Same host, different port -> different process, NOT local.
+    // This is the regression case fixed by this change.
+    EXPECT_FALSE(TransferSubmitter::isSameProcessEndpoint(
+        "192.168.1.10:12345", "192.168.1.10:12346"));
+
+    // Different hosts -> not local.
+    EXPECT_FALSE(TransferSubmitter::isSameProcessEndpoint(
+        "192.168.1.10:12345", "192.168.1.11:12345"));
+
+    // Hostname endpoints (non-P2P metadata mode) compare as full strings.
+    EXPECT_TRUE(TransferSubmitter::isSameProcessEndpoint("host-a", "host-a"));
+    EXPECT_FALSE(TransferSubmitter::isSameProcessEndpoint("host-a", "host-b"));
+}
+
 // Test TransferStrategy enum and stream operator
 TEST_F(TransferTaskTest, TransferStrategyEnum) {
     // Test enum values
     EXPECT_EQ(static_cast<int>(TransferStrategy::LOCAL_MEMCPY), 0);
     EXPECT_EQ(static_cast<int>(TransferStrategy::TRANSFER_ENGINE), 1);
+    EXPECT_EQ(static_cast<int>(TransferStrategy::FILE_READ), 2);
+    EXPECT_EQ(static_cast<int>(TransferStrategy::EMPTY), 3);
+    EXPECT_EQ(static_cast<int>(TransferStrategy::SPDK_NVMF), 4);
 
     // Test stream operator
     std::ostringstream oss;
@@ -158,6 +192,18 @@ TEST_F(TransferTaskTest, TransferStrategyEnum) {
     oss.str("");
     oss << TransferStrategy::TRANSFER_ENGINE;
     EXPECT_EQ(oss.str(), "TRANSFER_ENGINE");
+
+    oss.str("");
+    oss << TransferStrategy::SPDK_NVMF;
+    EXPECT_EQ(oss.str(), "SPDK_NVMF");
+
+    oss.str("");
+    oss << TransferStrategy::FILE_READ;
+    EXPECT_EQ(oss.str(), "FILE_READ");
+
+    oss.str("");
+    oss << TransferStrategy::EMPTY;
+    EXPECT_EQ(oss.str(), "EMPTY");
 }
 
 }  // namespace mooncake
diff --git a/mooncake-store/tests/utils_test.cpp b/mooncake-store/tests/utils_test.cpp
index 55a02b98..b3102fd9 100644
--- a/mooncake-store/tests/utils_test.cpp
+++ b/mooncake-store/tests/utils_test.cpp
@@ -23,6 +23,33 @@ TEST(UtilsTest, ByteSizeToString) {
     EXPECT_EQ(byte_size_to_string(15 * 1024 * 1024 + 44048), "15.04 MB");
 }
 
+TEST(UtilsTest, StringToByteSize) {
+    auto parsed = try_string_to_byte_size("16 MB");
+    ASSERT_TRUE(parsed.has_value());
+    EXPECT_EQ(parsed.value(), 16ULL * 1024 * 1024);
+
+    parsed = try_string_to_byte_size("0");
+    ASSERT_TRUE(parsed.has_value());
+    EXPECT_EQ(parsed.value(), 0);
+
+    EXPECT_FALSE(try_string_to_byte_size("-5").has_value());
+    EXPECT_FALSE(try_string_to_byte_size("16XB").has_value());
+    EXPECT_EQ(string_to_byte_size("-5"), 0);
+}
+
+TEST(UtilsTest, StringToBool) {
+    EXPECT_EQ(string_to_bool("1"), true);
+    EXPECT_EQ(string_to_bool("true"), true);
+    EXPECT_EQ(string_to_bool("YES"), true);
+    EXPECT_EQ(string_to_bool(" on "), true);
+    EXPECT_EQ(string_to_bool("0"), false);
+    EXPECT_EQ(string_to_bool("false"), false);
+    EXPECT_EQ(string_to_bool("No"), false);
+    EXPECT_EQ(string_to_bool(" off "), false);
+    EXPECT_EQ(string_to_bool("maybe"), std::nullopt);
+    EXPECT_EQ(string_to_bool(""), std::nullopt);
+}
+
 TEST(UtilsTest, IsPortAvailable) {
     // Find an available port
     int test_port = -1;
diff --git a/mooncake-transfer-engine/CMakeLists.txt b/mooncake-transfer-engine/CMakeLists.txt
index 87193c0e..121c3972 100644
--- a/mooncake-transfer-engine/CMakeLists.txt
+++ b/mooncake-transfer-engine/CMakeLists.txt
@@ -61,8 +61,9 @@ if (WITH_RUST_EXAMPLE)
   add_subdirectory(rust)
 endif()
 
-# Non-peermem requires cuMem allocated memory pointers
-if (USE_MNNVL OR NOT WITH_NVIDIA_PEERMEM)
+# Non-peermem requires cuMem allocated memory pointers; since WITH_NVIDIA_PEERMEM
+# is now a runtime environment variable, always build when GPU support is present.
+if (USE_MNNVL OR USE_CUDA OR USE_HIP OR USE_MUSA OR USE_MACA)
   add_subdirectory(nvlink-allocator)
 endif()
 
@@ -73,5 +74,12 @@ endif()
 
 if (USE_TENT)
   add_subdirectory(tent)
+endif()
+
+# tebench only needs the classic Transfer Engine for its default "classic"
+# backend; the TENT backend is compiled in conditionally (see benchmark/
+# CMakeLists.txt and main.cpp). Keep it gated solely on BUILD_BENCHMARK so it
+# can be built without enabling USE_TENT.
+if (BUILD_BENCHMARK)
   add_subdirectory(benchmark)
 endif()
diff --git a/mooncake-transfer-engine/benchmark/CMakeLists.txt b/mooncake-transfer-engine/benchmark/CMakeLists.txt
index 4a05f3d7..b17a8a80 100644
--- a/mooncake-transfer-engine/benchmark/CMakeLists.txt
+++ b/mooncake-transfer-engine/benchmark/CMakeLists.txt
@@ -7,14 +7,59 @@ else()
   message(STATUS "CUDA: Disabled")
 endif()
 
+if(USE_HIP)
+  list(APPEND CMAKE_PREFIX_PATH "/opt/rocm/lib/cmake")
+  find_package(HIP REQUIRED)
+  add_definitions(-DUSE_HIP -D__HIP_PLATFORM_AMD__)
+  message(STATUS "ROCm/HIP: Enabled")
+else()
+  message(STATUS "ROCm/HIP: Disabled")
+endif()
+
+if(USE_SUNRISE)
+  add_definitions(-DUSE_SUNRISE)
+  include_directories(${MC_TANGRT_ROOT}/include)
+  link_directories(${MC_TANGRT_ROOT}/lib)
+  message(STATUS "Sunrise/Tang: Enabled")
+else()
+  message(STATUS "Sunrise/Tang: Disabled")
+endif()
+
 file(GLOB TEBENCH_SOURCES "*.cpp")
+# The TENT backend is only available when USE_TENT is enabled; drop its
+# translation unit (which pulls in tent/ headers) from non-TENT builds.
+if(NOT USE_TENT)
+  list(REMOVE_ITEM TEBENCH_SOURCES "${CMAKE_CURRENT_SOURCE_DIR}/tent_backend.cpp")
+endif()
 add_executable(tebench ${TEBENCH_SOURCES})
-target_link_libraries(tebench PUBLIC transfer_engine tent_link_group)
+target_link_libraries(tebench PUBLIC transfer_engine)
+if(USE_TENT)
+  target_compile_definitions(tebench PRIVATE USE_TENT)
+  # tent_link_group's interface brings in the tent/ header search path.
+  target_link_libraries(tebench PUBLIC tent_link_group)
+else()
+  # The classic backend still uses a couple of header-only helpers that live
+  # under tent/include (SimpleRandom in utils.h, bindToSocket in te_backend.cpp).
+  # Expose just the header path so tebench builds without the TENT library.
+  target_include_directories(tebench PRIVATE
+    "${CMAKE_CURRENT_SOURCE_DIR}/../tent/include")
+endif()
 if(USE_CUDA)
   target_link_libraries(tebench PUBLIC CUDA::cudart)
 endif()
+if(USE_HIP)
+  target_link_libraries(tebench PUBLIC hip::host)
+endif()
+if(USE_SUNRISE)
+  target_link_libraries(tebench PUBLIC tangrt_shared)
+endif()
 
 # Set RPATH for finding libasio.so at runtime
+if(USE_SUNRISE AND MC_TANGRT_ROOT)
+  set(TANGRT_RPATH ":${MC_TANGRT_ROOT}/lib")
+else()
+  set(TANGRT_RPATH "")
+endif()
 set_target_properties(
   tebench PROPERTIES BUILD_WITH_INSTALL_RPATH TRUE
-                     INSTALL_RPATH "$ORIGIN/../lib:$ORIGIN/../../mooncake-common")
+                     INSTALL_RPATH "$ORIGIN/../lib:$ORIGIN/../../mooncake-common${TANGRT_RPATH}")
diff --git a/mooncake-transfer-engine/benchmark/main.cpp b/mooncake-transfer-engine/benchmark/main.cpp
index bb3b792b..6e9d76cb 100644
--- a/mooncake-transfer-engine/benchmark/main.cpp
+++ b/mooncake-transfer-engine/benchmark/main.cpp
@@ -16,7 +16,9 @@
 
 #include "bench_runner.h"
 #include "te_backend.h"
+#ifdef USE_TENT
 #include "tent_backend.h"
+#endif
 
 using namespace mooncake::tent;
 
@@ -101,10 +103,18 @@ int main(int argc, char* argv[]) {
     gflags::ParseCommandLineFlags(&argc, &argv, true);
     XferBenchConfig::loadFromFlags();
     std::unique_ptr<BenchRunner> runner;
-    if (XferBenchConfig::backend == "classic")
+    if (XferBenchConfig::backend == "classic") {
         runner = std::make_unique<TEBenchRunner>();
-    else
+    } else {
+#ifdef USE_TENT
         runner = std::make_unique<TENTBenchRunner>();
+#else
+        LOG(ERROR) << "Backend '" << XferBenchConfig::backend
+                   << "' requires building with -DUSE_TENT=ON; only the "
+                      "'classic' backend is available in this build";
+        return EXIT_FAILURE;
+#endif
+    }
     if (XferBenchConfig::target_seg_name.empty()) {
         std::cout << "\033[33mTo start initiators, run " << std::endl
                   << "  ./tebench --target_seg_name="
@@ -119,7 +129,10 @@ int main(int argc, char* argv[]) {
     for (int num_threads = XferBenchConfig::start_num_threads;
          !interrupted && num_threads <= XferBenchConfig::max_num_threads;
          num_threads *= 2) {
-        runner->startInitiator(num_threads);
+        if (runner->startInitiator(num_threads) != 0) {
+            interrupted = true;
+            break;
+        }
         for (size_t block_size = XferBenchConfig::start_block_size;
              !interrupted && block_size <= XferBenchConfig::max_block_size;
              block_size *= 2) {
diff --git a/mooncake-transfer-engine/benchmark/te_backend.cpp b/mooncake-transfer-engine/benchmark/te_backend.cpp
index af380690..7b742e94 100644
--- a/mooncake-transfer-engine/benchmark/te_backend.cpp
+++ b/mooncake-transfer-engine/benchmark/te_backend.cpp
@@ -17,14 +17,18 @@
 #include "te_backend.h"
 #include "utils.h"
 #include "common.h"
+#include "char_util.h"
 
-#ifdef USE_CUDA
+#if defined(USE_CUDA)
 #include <bits/stdint-uintn.h>
 #include <cuda.h>
 #include <cuda_runtime.h>
+#include "cuda_alike.h"
 #ifdef USE_MNNVL
 #include "transport/nvlink_transport/nvlink_transport.h"
 #endif
+#elif defined(USE_SUNRISE)
+#include "cuda_alike.h"
 #endif
 
 namespace mooncake {
@@ -47,9 +51,29 @@ static inline int getCudaDeviceNumaID(int cuda_id) {
         LOG(WARNING) << "cudaDeviceGetPCIBusId: " << cudaGetErrorString(err);
         return 0;
     }
-    for (char* ch = pci_bus_id; (*ch = tolower(*ch)); ch++);
+    for (char* ch = pci_bus_id; (*ch = to_lower(*ch)); ch++);
     return getNumaNodeFromPciDevice(pci_bus_id);
 }
+#elif defined(USE_SUNRISE)
+static inline int getSunriseDeviceNumaID(int dev_id) {
+    char pci_bus_id[20];
+    tangError_t err =
+        tangDeviceGetPCIBusId(pci_bus_id, sizeof(pci_bus_id), dev_id);
+    if (err != tangSuccess) {
+        LOG(WARNING) << "tangDeviceGetPCIBusId failed: "
+                     << tangGetErrorString(err);
+        return 0;
+    }
+    for (char* ch = pci_bus_id; (*ch = to_lower(*ch)); ch++);
+    std::string sysfs_path =
+        "/sys/bus/pci/devices/" + std::string(pci_bus_id) + "/numa_node";
+    std::ifstream numa_file(sysfs_path);
+    if (!numa_file.is_open()) return -1;
+    int numa_node = -1;
+    numa_file >> numa_node;
+    if (numa_file.fail()) return -1;
+    return numa_node;
+}
 #else
 static inline int getCudaDeviceNumaID(int cuda_id) { return 0; }
 #endif
@@ -70,7 +94,7 @@ void signalHandlerV0(int signum) {
 
 static void* allocateMemoryPool(size_t size, int buffer_id,
                                 bool from_vram = false) {
-#ifdef USE_CUDA
+#if defined(USE_CUDA) || defined(USE_SUNRISE)
     if (from_vram) {
         int gpu_id = buffer_id;
         void* d_buf;
@@ -79,8 +103,22 @@ static void* allocateMemoryPool(size_t size, int buffer_id,
 #ifdef USE_MNNVL
         d_buf = mooncake::NvlinkTransport::allocatePinnedLocalMemory(size);
 #else
-        cudaMalloc(&d_buf, size);
+        cudaError_t alloc_err = cudaMalloc(&d_buf, size);
+        if (alloc_err != cudaSuccess) {
+            LOG(ERROR) << "cudaMalloc failed on GPU " << gpu_id
+                       << ": err=" << alloc_err << " ("
+                       << cudaGetErrorString(alloc_err) << ")"
+                       << " size=" << size;
+#ifdef USE_SUNRISE
+            LOG(ERROR) << "Sunrise classic benchmark VRAM allocation failed. "
+                       << "Try a smaller -total_buffer_size, for example "
+                       << "-total_buffer_size=4096 for a minimal smoke test.";
 #endif
+            return nullptr;
+        }
+#endif
+        LOG(INFO) << "Allocated GPU memory at " << d_buf << " on GPU "
+                  << gpu_id;
         return d_buf;
     }
 #endif
@@ -88,8 +126,7 @@ static void* allocateMemoryPool(size_t size, int buffer_id,
 }
 
 static void freeMemoryPool(void* addr, size_t size) {
-#ifdef USE_CUDA
-#ifdef USE_MNNVL
+#if defined(USE_CUDA) && defined(USE_MNNVL)
     CUmemGenericAllocationHandle handle;
     auto result = cuMemRetainAllocationHandle(&handle, addr);
     if (result == CUDA_SUCCESS) {
@@ -97,7 +134,7 @@ static void freeMemoryPool(void* addr, size_t size) {
         return;
     }
 #endif
-    // check pointer on GPU
+#if defined(USE_CUDA) || defined(USE_SUNRISE)
     cudaPointerAttributes attributes;
     cudaPointerGetAttributes(&attributes, addr);
 
@@ -123,10 +160,11 @@ int TEBenchRunner::allocateBuffers() {
             auto location = "cpu:" + std::to_string(i);
             pinned_buffer_list_[i] =
                 allocateMemoryPool(total_buffer_size, i, false);
+            if (!pinned_buffer_list_[i]) return -1;
             engine_->registerLocalMemory(pinned_buffer_list_[i],
                                          total_buffer_size, location);
         }
-#ifdef USE_CUDA
+#if defined(USE_CUDA) || defined(USE_SUNRISE)
     } else if (XferBenchConfig::seg_type == "VRAM") {
         int gpu_count = 0;
         cudaGetDeviceCount(&gpu_count);
@@ -141,9 +179,10 @@ int TEBenchRunner::allocateBuffers() {
         pinned_buffer_list_.resize(num_buffers, nullptr);
         for (int i = 0; i < num_buffers; ++i) {
             int gpu_id = start_gpu + i;
-            auto location = "cuda:" + std::to_string(gpu_id);
+            auto location = std::string(GPU_PREFIX) + std::to_string(gpu_id);
             pinned_buffer_list_[i] =
                 allocateMemoryPool(total_buffer_size, gpu_id, true);
+            if (!pinned_buffer_list_[i]) return -1;
             engine_->registerLocalMemory(pinned_buffer_list_[i],
                                          total_buffer_size, location);
         }
@@ -157,6 +196,7 @@ int TEBenchRunner::allocateBuffers() {
 int TEBenchRunner::freeBuffers() {
     auto total_buffer_size = XferBenchConfig::total_buffer_size;
     for (size_t i = 0; i < pinned_buffer_list_.size(); ++i) {
+        if (!pinned_buffer_list_[i]) continue;
         engine_->unregisterLocalMemory(pinned_buffer_list_[i]);
         freeMemoryPool(pinned_buffer_list_[i], total_buffer_size);
     }
@@ -175,17 +215,38 @@ TEBenchRunner::TEBenchRunner() {
                         ? mooncake::getHostname()
                         : XferBenchConfig::seg_name;
     engine_->init(conn_str, seg_name);
-    allocateBuffers();
+#ifdef USE_SUNRISE
+    // Sunrise IPC handles must be produced by the installed transport before
+    // the benchmark starts issuing transfers, so classic tebench installs the
+    // transport explicitly instead of relying on TransferEngineImpl auto-setup.
+    auto* sunrise_transport =
+        engine_->installTransport("sunrise_link", nullptr);
+    LOG_ASSERT(sunrise_transport != nullptr)
+        << "Failed to install SunriseLink transport";
+#endif
+    init_ok_ = (allocateBuffers() == 0);
+    if (!init_ok_) {
+        LOG(ERROR) << "TEBenchRunner initialization failed";
+        return;
+    }
 }
 
 TEBenchRunner::~TEBenchRunner() { freeBuffers(); }
 
 int TEBenchRunner::runTarget() {
+    if (!init_ok_) {
+        LOG(ERROR) << "Target cannot start: initialization failed";
+        return -1;
+    }
     while (g_te_running) sleep(1);
     return 0;
 }
 
 int TEBenchRunner::startInitiator(int num_threads) {
+    if (!init_ok_) {
+        LOG(ERROR) << "Initiator cannot start: initialization failed";
+        return -1;
+    }
     handle_ = engine_->openSegment(XferBenchConfig::target_seg_name);
     info_ = engine_->getMetadata()->getSegmentDescByID(handle_);
     std::sort(
@@ -222,16 +283,36 @@ static int parseIndex(const std::string& loc) {
 }
 
 void TEBenchRunner::pinThread(int thread_id) {
+#ifdef USE_SUNRISE
+    if (XferBenchConfig::seg_type == "VRAM" && !pinned_buffer_list_.empty()) {
+        // Tang pointer queries are device-context sensitive on worker threads.
+        int base_gpu = std::max(0, XferBenchConfig::local_gpu_id);
+        int device_id =
+            base_gpu +
+            (thread_id % static_cast<int>(pinned_buffer_list_.size()));
+        auto err = cudaSetDevice(device_id);
+        LOG_ASSERT(err == cudaSuccess)
+            << "cudaSetDevice failed before getMemoryLocation: "
+            << cudaGetErrorString(err) << " device_id=" << device_id;
+    }
+#endif
     uint64_t addr =
         (uint64_t)pinned_buffer_list_[thread_id % pinned_buffer_list_.size()];
     auto result = getMemoryLocation((void*)addr, 1, true);
     if (result[0].location.starts_with("cpu")) {
         auto socket_id = parseIndex(result[0].location);
         bindToSocket(socket_id);
+#ifdef USE_CUDA
     } else if (result[0].location.starts_with("cuda")) {
         auto device_id = parseIndex(result[0].location);
         auto socket_id = getCudaDeviceNumaID(device_id);
         bindToSocket(socket_id);
+#elif defined(USE_SUNRISE)
+    } else if (result[0].location.starts_with("cuda")) {
+        auto device_id = parseIndex(result[0].location);
+        auto socket_id = getSunriseDeviceNumaID(device_id);
+        bindToSocket(socket_id);
+#endif
     }
 }
 
diff --git a/mooncake-transfer-engine/benchmark/te_backend.h b/mooncake-transfer-engine/benchmark/te_backend.h
index 92cbbf34..7d48a1d9 100644
--- a/mooncake-transfer-engine/benchmark/te_backend.h
+++ b/mooncake-transfer-engine/benchmark/te_backend.h
@@ -90,9 +90,10 @@ class TEBenchRunner : public BenchRunner {
     std::condition_variable cv_task_;
     std::condition_variable cv_done_;
     int pending_ = 0;
+    bool init_ok_ = true;
 };
 
 }  // namespace tent
 }  // namespace mooncake
 
-#endif  // TEV0_BACKEND_H
\ No newline at end of file
+#endif  // TEV0_BACKEND_H
diff --git a/mooncake-transfer-engine/benchmark/tent_backend.cpp b/mooncake-transfer-engine/benchmark/tent_backend.cpp
index 435d553d..3aa24db0 100644
--- a/mooncake-transfer-engine/benchmark/tent_backend.cpp
+++ b/mooncake-transfer-engine/benchmark/tent_backend.cpp
@@ -14,12 +14,18 @@
 
 #include "tent_backend.h"
 #include "utils.h"
+#include "char_util.h"
 #include "tent/common/types.h"
 #include "tent/runtime/platform.h"
 #include "tent/runtime/topology.h"
+#include "tent/runtime/transport_selector.h"
 
-#ifdef USE_CUDA
-#include <cuda_runtime.h>
+#if defined(USE_CUDA) || defined(USE_SUNRISE)
+#include "cuda_alike.h"
+#endif
+
+#ifdef USE_HIP
+#include <hip/hip_runtime.h>
 #endif
 
 namespace mooncake {
@@ -50,9 +56,14 @@ std::shared_ptr<Config> loadConfig() {
     if (!XferBenchConfig::xport_type.empty()) {
         // Map of transport names to their config keys (handle name mismatches)
         std::unordered_map<std::string, std::string> transport_map = {
-            {"rdma", "rdma"},        {"tcp", "tcp"},    {"shm", "shm"},
+            {"rdma", "rdma"},
+            {"tcp", "tcp"},
+            {"shm", "shm"},
             {"iouring", "io_uring"},  // Note: iouring -> io_uring
-            {"gds", "gds"},          {"mnnvl", "mnnvl"}};
+            {"gds", "gds"},
+            {"mnnvl", "mnnvl"},
+            {"nvlink", "nvlink"},
+            {"sunrise_link", "sunrise_link"}};
 
         // Disable all transports by default
         for (const auto& entry : transport_map) {
@@ -74,8 +85,10 @@ static TransportType getTransportType(const std::string& xport_type) {
     if (xport_type == "shm") return SHM;
     if (xport_type == "gds") return GDS;
     if (xport_type == "mnnvl") return MNNVL;
+    if (xport_type == "nvlink") return NVLINK;
     if (xport_type == "tcp") return TCP;
     if (xport_type == "iouring") return IOURING;
+    if (xport_type == "sunrise_link") return SUNRISE_LINK;
     return UNSPEC;
 }
 
@@ -91,11 +104,27 @@ int TENTBenchRunner::allocateBuffers() {
     if (seg_type == "DRAM") {
         device_prefix = "cpu";
         num_buffers = numa_num_configured_nodes();
-#ifdef USE_CUDA
+#if defined(USE_CUDA) || defined(USE_SUNRISE)
     } else if (seg_type == "VRAM") {
         device_prefix = "cuda";
         int gpu_count = 0;
-        cudaGetDeviceCount(&gpu_count);
+        auto err = cudaGetDeviceCount(&gpu_count);
+        LOG_ASSERT(err == cudaSuccess && gpu_count > 0)
+            << "cudaGetDeviceCount failed: " << cudaGetErrorString(err);
+        start_idx = 0;
+        num_buffers = gpu_count;
+        if (XferBenchConfig::local_gpu_id != -1) {
+            start_idx = XferBenchConfig::local_gpu_id;
+            num_buffers = 1;
+            LOG_ASSERT(start_idx >= 0 && start_idx < gpu_count)
+                << "local_gpu_id " << start_idx << " out of range [0, "
+                << gpu_count << ")";
+        }
+#elif defined(USE_HIP)
+    } else if (seg_type == "VRAM") {
+        device_prefix = "rocm";
+        int gpu_count = 0;
+        hipGetDeviceCount(&gpu_count);
         start_idx = 0;
         num_buffers = gpu_count;
         if (XferBenchConfig::local_gpu_id != -1) {
@@ -111,38 +140,47 @@ int TENTBenchRunner::allocateBuffers() {
         return -1;
     }
 
-    // Allocate
     pinned_buffer_list_.resize(num_buffers, nullptr);
-    auto start_ts = getCurrentTimeInNano();
+    uint64_t alloc_ns = 0, reg_ns = 0;
     for (int i = 0; i < num_buffers; ++i) {
         auto location = device_prefix + ":" + std::to_string(start_idx + i);
+        MemoryOptions options;
         if (!xport_type.empty()) {
-            MemoryOptions options;
             options.type = getTransportType(xport_type);
-            options.location = std::move(location);
+            options.location = location;
+        }
+
+        auto t0 = getCurrentTimeInNano();
+        if (!xport_type.empty()) {
+            options.location = location;
             CHECK_FAIL(engine_->allocateLocalMemory(
                 &pinned_buffer_list_[i], total_buffer_size, options));
         } else {
             CHECK_FAIL(engine_->allocateLocalMemory(
                 &pinned_buffer_list_[i], total_buffer_size, location));
         }
-    }
+        auto t1 = getCurrentTimeInNano();
 
-    // Register
-    auto allocated_ts = getCurrentTimeInNano();
-    std::vector<size_t> buffers_size(num_buffers, total_buffer_size);
-    MemoryOptions options;
-    if (!xport_type.empty()) {
-        options.type = getTransportType(xport_type);
+#ifdef USE_SUNRISE
+        if (seg_type == "VRAM") {
+            auto err = cudaSetDevice(start_idx + i);
+            CHECK_FAIL(err == cudaSuccess ? Status::OK()
+                                          : Status::InternalError(
+                                                "Failed to set Sunrise device "
+                                                "before registerLocalMemory"));
+        }
+#endif
+        CHECK_FAIL(engine_->registerLocalMemory(pinned_buffer_list_[i],
+                                                total_buffer_size, options));
+        auto t2 = getCurrentTimeInNano();
+
+        alloc_ns += (t1 - t0);
+        reg_ns += (t2 - t1);
     }
-    CHECK_FAIL(engine_->registerLocalMemory(pinned_buffer_list_, buffers_size,
-                                            options));
-    auto registered_ts = getCurrentTimeInNano();
 
     LOG(INFO) << "Allocated " << total_buffer_size * num_buffers << " bytes "
-              << seg_type << " buffers in " << (allocated_ts - start_ts) / 1e6
-              << " ms, registered in " << (registered_ts - allocated_ts) / 1e6
-              << " ms";
+              << seg_type << " buffers in " << alloc_ns / 1e6
+              << " ms, registered in " << reg_ns / 1e6 << " ms";
     return 0;
 }
 
@@ -161,6 +199,8 @@ TENTBenchRunner::TENTBenchRunner() {
     signal(SIGINT, signalHandlerV1);
     signal(SIGTERM, signalHandlerV1);
     engine_ = std::make_unique<TransferEngine>(loadConfig());
+    transport_hint_ = TransportSelector::parseTransportType(
+        XferBenchConfig::tent_transport_hint);
     allocateBuffers();
 }
 
@@ -200,7 +240,6 @@ int TENTBenchRunner::stopInitiator() {
     return 0;
 }
 
-#ifdef USE_CUDA
 static inline int getNumaNodeFromPciDevice(const std::string& pci_bdf) {
     std::string sysfs_path = "/sys/bus/pci/devices/" + pci_bdf + "/numa_node";
     std::ifstream numa_file(sysfs_path);
@@ -210,21 +249,46 @@ static inline int getNumaNodeFromPciDevice(const std::string& pci_bdf) {
     if (numa_file.fail()) return -1;
     return numa_node;
 }
-static inline int getCudaDeviceNumaID(int cuda_id) {
+
+#if defined(USE_CUDA) || defined(USE_SUNRISE)
+static inline int getGpuDeviceNumaID(int gpu_id) {
     char pci_bus_id[20];
-    auto err = cudaDeviceGetPCIBusId(pci_bus_id, sizeof(pci_bus_id), cuda_id);
+    auto err = cudaDeviceGetPCIBusId(pci_bus_id, sizeof(pci_bus_id), gpu_id);
     if (err != cudaSuccess) {
         LOG(WARNING) << "cudaDeviceGetPCIBusId: " << cudaGetErrorString(err);
         return 0;
     }
-    for (char* ch = pci_bus_id; (*ch = tolower(*ch)); ch++);
+    for (char* ch = pci_bus_id; (*ch = to_lower(*ch)); ch++);
+    return getNumaNodeFromPciDevice(pci_bus_id);
+}
+#elif defined(USE_HIP)
+static inline int getGpuDeviceNumaID(int gpu_id) {
+    hipDeviceProp_t prop;
+    if (hipGetDeviceProperties(&prop, gpu_id) != hipSuccess) return 0;
+    char pci_bus_id[20];
+    snprintf(pci_bus_id, sizeof(pci_bus_id), "%04x:%02x:%02x.0",
+             prop.pciDomainID, prop.pciBusID, prop.pciDeviceID);
     return getNumaNodeFromPciDevice(pci_bus_id);
 }
 #else
-static inline int getCudaDeviceNumaID(int cuda_id) { return 0; }
+static inline int getGpuDeviceNumaID(int gpu_id) { return 0; }
 #endif
 
 void TENTBenchRunner::pinThread(int thread_id) {
+#ifdef USE_SUNRISE
+    if (XferBenchConfig::seg_type == "VRAM" && !pinned_buffer_list_.empty()) {
+        int base_gpu = std::max(0, XferBenchConfig::local_gpu_id);
+        int device_id =
+            base_gpu +
+            (thread_id % static_cast<int>(pinned_buffer_list_.size()));
+        auto err = cudaSetDevice(device_id);
+        LOG_ASSERT(err == cudaSuccess)
+            << "cudaSetDevice failed before getLocation: "
+            << cudaGetErrorString(err) << " device_id=" << device_id;
+        bindToSocket(getGpuDeviceNumaID(device_id));
+        return;
+    }
+#endif
     uint64_t addr =
         (uint64_t)pinned_buffer_list_[thread_id % pinned_buffer_list_.size()];
     auto result = Platform::getLoader().getLocation((void*)addr, 1);
@@ -232,9 +296,9 @@ void TENTBenchRunner::pinThread(int thread_id) {
     if (location.type() == "cpu") {
         auto socket_id = location.index();
         bindToSocket(socket_id);
-    } else if (location.type() == "cuda") {
+    } else if (location.type() == "cuda" || location.type() == "rocm") {
         auto device_id = location.index();
-        auto socket_id = getCudaDeviceNumaID(device_id);
+        auto socket_id = getGpuDeviceNumaID(device_id);
         bindToSocket(socket_id);
     }
 }
@@ -283,6 +347,7 @@ double TENTBenchRunner::runSingleTransfer(uint64_t local_addr,
         entry.source = (void*)(local_addr + block_size * i);
         entry.target_id = handle_;
         entry.target_offset = target_addr + block_size * i;
+        entry.transport_hint = transport_hint_;
         requests.emplace_back(entry);
     }
     XferBenchTimer timer;
diff --git a/mooncake-transfer-engine/benchmark/tent_backend.h b/mooncake-transfer-engine/benchmark/tent_backend.h
index 555ac995..1648ad92 100644
--- a/mooncake-transfer-engine/benchmark/tent_backend.h
+++ b/mooncake-transfer-engine/benchmark/tent_backend.h
@@ -86,6 +86,7 @@ class TENTBenchRunner : public BenchRunner {
     std::vector<void*> pinned_buffer_list_;
     SegmentID handle_;
     SegmentInfo info_;
+    TransportType transport_hint_{UNSPEC};
 
     std::vector<std::function<int(int)>> current_task_;
     std::vector<std::thread> threads_;
diff --git a/mooncake-transfer-engine/benchmark/utils.cpp b/mooncake-transfer-engine/benchmark/utils.cpp
index 2d7ca5e5..d907a2bd 100644
--- a/mooncake-transfer-engine/benchmark/utils.cpp
+++ b/mooncake-transfer-engine/benchmark/utils.cpp
@@ -44,10 +44,15 @@ DEFINE_string(metadata_url_list, "",
 DEFINE_int32(
     rpc_server_port, 0,
     "RPC server port used for p2p metadata service (0 = auto-select).");
-DEFINE_string(xport_type, "", "Transport type: rdma|shm|mnnvl|gds|iouring");
+DEFINE_string(xport_type, "",
+              "Transport type: rdma|shm|mnnvl|gds|iouring|sunrise_link");
 DEFINE_string(backend, "tent", "Transport backend: classic|tent");
 DEFINE_bool(notifi, false,
             "Enable RDMA notification for performance measurement.");
+DEFINE_string(
+    tent_transport_hint, "unspec",
+    "tent only: per-request transport_hint. "
+    "unspec|rdma|tcp|shm|nvlink|gds|io_uring|mnnvl|ascend|sunrise_link");
 
 namespace mooncake {
 namespace tent {
@@ -72,6 +77,7 @@ int XferBenchConfig::rpc_server_port = 0;
 std::string XferBenchConfig::xport_type;
 std::string XferBenchConfig::backend;
 bool XferBenchConfig::notifi = false;
+std::string XferBenchConfig::tent_transport_hint;
 
 int XferBenchConfig::local_gpu_id = 0;
 int XferBenchConfig::target_gpu_id = 0;
@@ -99,6 +105,7 @@ void XferBenchConfig::loadFromFlags() {
     xport_type = FLAGS_xport_type;
     backend = FLAGS_backend;
     notifi = FLAGS_notifi;
+    tent_transport_hint = FLAGS_tent_transport_hint;
 
     local_gpu_id = FLAGS_local_gpu_id;
     target_gpu_id = FLAGS_target_gpu_id;
diff --git a/mooncake-transfer-engine/benchmark/utils.h b/mooncake-transfer-engine/benchmark/utils.h
index ae48d43b..6c862cd6 100644
--- a/mooncake-transfer-engine/benchmark/utils.h
+++ b/mooncake-transfer-engine/benchmark/utils.h
@@ -30,8 +30,14 @@
 #include "tent/common/utils/os.h"
 #include "tent/common/utils/random.h"
 
-#ifdef USE_CUDA
+#if defined(USE_CUDA)
 #include <cuda_runtime.h>
+#elif defined(USE_SUNRISE)
+#include "cuda_alike.h"
+#endif
+
+#ifdef USE_HIP
+#include <hip/hip_runtime.h>
 #endif
 
 #define CHECK_FAIL(call)                                        \
@@ -69,6 +75,7 @@ struct XferBenchConfig {
     static std::string xport_type;
     static std::string backend;
     static bool notifi;
+    static std::string tent_transport_hint;
 
     static int local_gpu_id;
     static int target_gpu_id;
@@ -146,7 +153,7 @@ void printStatsHeader();
 void printStats(size_t block_size, size_t batch_size, XferBenchStats& stats,
                 int num_threads);
 
-#ifdef USE_CUDA
+#if defined(USE_CUDA) || defined(USE_SUNRISE)
 static inline bool isCudaMemory(void* ptr) {
     cudaPointerAttributes attr;
     auto ret = cudaPointerGetAttributes(&attr, ptr);
@@ -154,33 +161,86 @@ static inline bool isCudaMemory(void* ptr) {
 }
 #endif
 
+#ifdef USE_HIP
+static inline bool isHipMemory(void* ptr) {
+    hipPointerAttribute_t attr;
+    auto ret = hipPointerGetAttributes(&attr, ptr);
+    return ret == hipSuccess && attr.type == hipMemoryTypeDevice;
+}
+#endif
+
+static inline bool isGpuMemory(void* ptr) {
+#if defined(USE_CUDA) || defined(USE_SUNRISE)
+    if (isCudaMemory(ptr)) return true;
+#endif
+#ifdef USE_HIP
+    if (isHipMemory(ptr)) return true;
+#endif
+    return false;
+}
+
 static inline uint8_t fillData(void* addr, size_t length) {
     uint8_t seed = (uint8_t)SimpleRandom::Get().next(256);
-#ifdef USE_CUDA
+#if defined(USE_CUDA)
     if (isCudaMemory(addr)) {
-        std::vector<uint8_t> ref_data;
-        ref_data.resize(length, seed);
+        std::vector<uint8_t> ref_data(length, seed);
         cudaMemcpy(addr, ref_data.data(), length, cudaMemcpyDefault);
         return seed;
     }
+#elif defined(USE_SUNRISE)
+    if (isCudaMemory(addr)) {
+        std::vector<uint8_t> ref_data(length, seed);
+        auto err =
+            cudaMemcpy(addr, ref_data.data(), length, cudaMemcpyHostToDevice);
+        LOG_ASSERT(err == cudaSuccess)
+            << "cudaMemcpy failed: " << cudaGetErrorString(err);
+        return seed;
+    }
+#endif
+#ifdef USE_HIP
+    if (isHipMemory(addr)) {
+        std::vector<uint8_t> ref_data(length, seed);
+        hipMemcpy(addr, ref_data.data(), length, hipMemcpyDefault);
+        return seed;
+    }
 #endif
     memset(addr, seed, length);
     return seed;
 }
 
 static inline void verifyData(void* addr, size_t length, uint8_t seed) {
-    std::vector<uint8_t> ref_data;
-    ref_data.resize(length, seed);
-#ifdef USE_CUDA
+    std::vector<uint8_t> ref_data(length, seed);
+#if defined(USE_CUDA)
     if (isCudaMemory(addr)) {
-        std::vector<uint8_t> act_data;
-        act_data.resize(length);
+        std::vector<uint8_t> act_data(length);
         cudaMemcpy(act_data.data(), addr, length, cudaMemcpyDefault);
         if (memcmp(act_data.data(), ref_data.data(), length)) {
             LOG(FATAL) << "Inconsistent data detected";
         }
         return;
     }
+#elif defined(USE_SUNRISE)
+    if (isCudaMemory(addr)) {
+        std::vector<uint8_t> act_data(length);
+        auto err =
+            cudaMemcpy(act_data.data(), addr, length, cudaMemcpyDeviceToHost);
+        LOG_ASSERT(err == cudaSuccess)
+            << "cudaMemcpy failed: " << cudaGetErrorString(err);
+        if (memcmp(act_data.data(), ref_data.data(), length)) {
+            LOG(FATAL) << "Inconsistent data detected";
+        }
+        return;
+    }
+#endif
+#ifdef USE_HIP
+    if (isHipMemory(addr)) {
+        std::vector<uint8_t> act_data(length);
+        hipMemcpy(act_data.data(), addr, length, hipMemcpyDefault);
+        if (memcmp(act_data.data(), ref_data.data(), length)) {
+            LOG(FATAL) << "Inconsistent data detected";
+        }
+        return;
+    }
 #endif
     if (memcmp(addr, ref_data.data(), length)) {
         LOG(FATAL) << "Inconsistent data detected";
diff --git a/mooncake-transfer-engine/example/CMakeLists.txt b/mooncake-transfer-engine/example/CMakeLists.txt
index 54a6cdfb..4adfcdba 100644
--- a/mooncake-transfer-engine/example/CMakeLists.txt
+++ b/mooncake-transfer-engine/example/CMakeLists.txt
@@ -24,6 +24,12 @@ add_executable(transfer_engine_bench_with_notify
                ${WORKSPACE}/transfer_engine_bench_with_notify.cpp)
 target_link_libraries(transfer_engine_bench_with_notify PUBLIC transfer_engine)
 
+if(USE_EFA)
+  add_executable(efa_first_submit_probe
+                 ${WORKSPACE}/efa_first_submit_probe.cpp)
+  target_link_libraries(efa_first_submit_probe PUBLIC transfer_engine)
+endif()
+
 add_executable(memory_pool ${WORKSPACE}/memory_pool.cpp)
 target_link_libraries(memory_pool PUBLIC transfer_engine)
 
@@ -55,3 +61,48 @@ endif()
 if(USE_UBSHMEM)
   target_link_libraries(transfer_engine_bench PUBLIC transfer_engine)
 endif()
+
+# Device Transport Example — two-rank P2P write + signal via Device API.
+# Requires 2 GPUs with P2P access.  Built but not registered with CTest
+# (manual execution only, like rdma_transport_test).
+#
+# Off by default: this .cu target pulls in transfer_engine.h, whose common.h
+# uses C++20 std::string APIs (e.g. starts_with), so it needs the CUDA20
+# dialect.  Older CMake (such as the CI Docker image) cannot enable CUDA20 and
+# fails at generate time.  Opt in explicitly with -DBUILD_DEVICE_TRANSPORT_EXAMPLE=ON
+# when building manually on a host with a recent CMake and 2 GPUs.
+option(BUILD_DEVICE_TRANSPORT_EXAMPLE
+       "Build the two-rank Device API P2P example (requires CUDA20-capable CMake and 2 GPUs)"
+       OFF)
+if(USE_CUDA AND BUILD_DEVICE_TRANSPORT_EXAMPLE)
+  enable_language(CUDA)
+  add_executable(device_transport_example
+                 ${WORKSPACE}/device_transport_example.cu)
+  set_source_files_properties(
+    ${WORKSPACE}/device_transport_example.cu
+    PROPERTIES LANGUAGE CUDA)
+  target_include_directories(device_transport_example PRIVATE
+    ${CMAKE_SOURCE_DIR}/mooncake-transfer-engine/include)
+  target_link_libraries(device_transport_example PUBLIC
+    transfer_engine gflags::gflags glog::glog)
+  # common.h (pulled in via transfer_engine.h) uses C++20 std::string APIs
+  # such as starts_with; nvcc does not inherit CMAKE_CXX_STANDARD, so set the
+  # CUDA standard explicitly on this target.  CUDA_EXTENSIONS OFF requests
+  # plain -std=c++20 (nvcc has no gnu++20 dialect mapping).
+  set_target_properties(device_transport_example PROPERTIES
+    CUDA_STANDARD 20 CUDA_STANDARD_REQUIRED ON CUDA_EXTENSIONS OFF)
+  # Convert torch-style arch list (e.g. "8.0;9.0") to CMake CUDA format
+  # (e.g. "80;90") by stripping the dot.
+  if(TORCH_CUDA_ARCH_LIST)
+    set(_cuda_arch_list "")
+    foreach(_arch IN LISTS TORCH_CUDA_ARCH_LIST)
+      string(REPLACE "." "" _arch_clean "${_arch}")
+      list(APPEND _cuda_arch_list "${_arch_clean}")
+    endforeach()
+    set_target_properties(device_transport_example PROPERTIES
+      CUDA_ARCHITECTURES "${_cuda_arch_list}")
+  else()
+    set_target_properties(device_transport_example PROPERTIES
+      CUDA_ARCHITECTURES "80;90")
+  endif()
+endif()
diff --git a/mooncake-transfer-engine/example/device_transport_example.cu b/mooncake-transfer-engine/example/device_transport_example.cu
new file mode 100644
index 00000000..f06ade3b
--- /dev/null
+++ b/mooncake-transfer-engine/example/device_transport_example.cu
@@ -0,0 +1,305 @@
+// Copyright 2024 KVCache.AI
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+// Device Transport Example — two-rank P2P write + signal via Device API.
+//
+// This example demonstrates the full Device API lifecycle:
+//   1. Host side:  P2pTransport for IPC handle exchange and peer mapping.
+//   2. Device side: CommCtx + mc_route_put + mc_signal for GPU-initiated
+//      P2P data transfer and notification.
+//
+// Usage (two terminals on the same node):
+//   $ ./device_transport_example --rank=0
+//   $ ./device_transport_example --rank=1
+//
+// The example uses file-based IPC handle exchange (no external dependencies).
+// Requires 2 GPUs with P2P access (NVLink or PCIe).
+
+#include <gflags/gflags.h>
+#include <glog/logging.h>
+#include <cuda_runtime.h>
+
+#include <atomic>
+#include <chrono>
+#include <cstring>
+#include <fstream>
+#include <thread>
+#include <vector>
+
+#include "transfer_engine.h"
+#include "transport/device/comm_device.cuh"
+#include "transport/device/device_transport.h"
+
+DEFINE_int32(rank, -1, "Rank of this process (0 or 1)");
+DEFINE_int32(world_size, 2, "Total number of ranks");
+DEFINE_int32(gpu_id, -1, "GPU ID (defaults to rank)");
+DEFINE_string(metadata_server, "P2PHANDSHAKE",
+              "Metadata server (P2PHANDSHAKE for no external deps)");
+DEFINE_string(local_server_name, "",
+              "Local server name (default: 127.0.0.1:<port>)");
+DEFINE_string(ipc_dir, "/tmp", "Directory for IPC handle exchange files");
+DEFINE_int32(kDataBytes, 4096, "Bytes to transfer in the P2P write");
+DEFINE_int32(kSignalOffset, 0, "Offset within buffer for the signal word");
+
+namespace {
+
+static void checkCuda(cudaError_t err, const char* msg) {
+    if (err != cudaSuccess) {
+        LOG(FATAL) << msg << ": " << cudaGetErrorString(err);
+    }
+}
+
+// ---------------------------------------------------------------------------
+// File-based IPC handle exchange
+// ---------------------------------------------------------------------------
+
+static std::string ipcFilePath(int rank) {
+    return FLAGS_ipc_dir + "/device_api_ex_rank_" + std::to_string(rank) +
+           ".bin";
+}
+
+static void writeIpcHandle(int rank, const std::vector<int32_t>& handle) {
+    std::string path = ipcFilePath(rank);
+    std::ofstream ofs(path, std::ios::binary | std::ios::trunc);
+    CHECK(ofs) << "Failed to open " << path << " for writing";
+    uint32_t size = static_cast<uint32_t>(handle.size());
+    ofs.write(reinterpret_cast<const char*>(&size), sizeof(size));
+    ofs.write(reinterpret_cast<const char*>(handle.data()),
+              handle.size() * sizeof(int32_t));
+    ofs.close();
+    LOG(INFO) << "Rank " << rank << " wrote IPC handle (" << handle.size()
+              << " int32s) to " << path;
+}
+
+static std::vector<int32_t> readIpcHandle(int rank) {
+    std::string path = ipcFilePath(rank);
+    // Poll until the file appears (peer may not have written it yet).
+    for (int attempt = 0; attempt < 300; ++attempt) {
+        std::ifstream ifs(path, std::ios::binary);
+        if (ifs) {
+            uint32_t size = 0;
+            ifs.read(reinterpret_cast<char*>(&size), sizeof(size));
+            CHECK(size > 0 && size <= 256)
+                << "Invalid IPC handle size: " << size;
+            std::vector<int32_t> handle(size);
+            ifs.read(reinterpret_cast<char*>(handle.data()),
+                     size * sizeof(int32_t));
+            ifs.close();
+            LOG(INFO) << "Rank " << FLAGS_rank << " read IPC handle from rank "
+                      << rank << " (" << size << " int32s)";
+            return handle;
+        }
+        std::this_thread::sleep_for(std::chrono::milliseconds(100));
+    }
+    LOG(FATAL) << "Timeout waiting for IPC handle from rank " << rank;
+    return {};
+}
+
+// ---------------------------------------------------------------------------
+// Device-side kernel
+// ---------------------------------------------------------------------------
+
+// GDR buffer layout:
+//   [0 .. kDataBytes)              — data region
+//   [kDataBytes .. kDataBytes+4)   — signal word (int32_t)
+static constexpr int kSignalWordOffset = 4096;  // must match kDataBytes default
+
+// Rank 0: P2P-write data to rank 1, then signal.
+__global__ void senderKernel(mooncake::device::CommCtx ctx, int dst_rank,
+                             int data_bytes) {
+    using namespace mooncake::device;
+    if (threadIdx.x != 0 || blockIdx.x != 0) return;
+
+    // Write a known pattern to the peer's data region via P2P.
+    char* local_data = reinterpret_cast<char*>(ctx.p2p.local_base);
+    void* peer_data = mc_route_put(ctx, dst_rank, local_data);
+    if (peer_data == nullptr) {
+        printf("[Rank 0] FAIL: P2P not available to rank %d\n", dst_rank);
+        asm("trap;");
+        return;
+    }
+
+    // Fill local buffer with pattern, then copy to peer via P2P store.
+    for (int i = 0; i < data_bytes; ++i) {
+        local_data[i] = static_cast<char>(i & 0xFF);
+    }
+    // Use 16-byte non-temporal stores for the bulk copy.
+    for (int i = 0; i < data_bytes; i += 16) {
+        int4 val = mc_ld_nc(reinterpret_cast<const int4*>(local_data + i));
+        mc_st_na(
+            reinterpret_cast<int4*>(reinterpret_cast<char*>(peer_data) + i),
+            val);
+    }
+
+    // Signal: write 1 to the peer's signal word.
+    int* local_sig = reinterpret_cast<int*>(local_data + kSignalWordOffset);
+    *local_sig = 1;
+    mc_signal(ctx, dst_rank, 0 /*channel*/, 1 /*qps_per_rank*/, local_sig, 1);
+}
+
+// Rank 1: wait for signal from rank 0, then verify data.
+__global__ void receiverKernel(mooncake::device::CommCtx ctx, int src_rank,
+                               int data_bytes) {
+    using namespace mooncake::device;
+    if (threadIdx.x != 0 || blockIdx.x != 0) return;
+
+    char* local_data = reinterpret_cast<char*>(ctx.p2p.local_base);
+    int* local_sig = reinterpret_cast<int*>(local_data + kSignalWordOffset);
+
+    // Spin-wait for the signal.
+    // The sender writes the signal via P2P store (mc_p2p_signal →
+    // mc_st_release). We use ld_acquire to observe it.
+    while (mc_ld_acquire(local_sig) == 0) {
+        __threadfence_block();
+    }
+
+    // Verify the data written by rank 0.
+    int mismatches = 0;
+    for (int i = 0; i < data_bytes; ++i) {
+        if (local_data[i] != static_cast<char>(i & 0xFF)) {
+            ++mismatches;
+        }
+    }
+    if (mismatches > 0) {
+        printf("[Rank 1] FAIL: %d byte mismatches in received data\n",
+               mismatches);
+    } else {
+        printf("[Rank 1] PASS: all %d bytes match expected pattern\n",
+               data_bytes);
+    }
+}
+
+}  // namespace
+
+int main(int argc, char** argv) {
+    gflags::ParseCommandLineFlags(&argc, &argv, true);
+    google::InitGoogleLogging(argv[0]);
+
+    int rank = FLAGS_rank;
+    int world_size = FLAGS_world_size;
+    CHECK(rank == 0 || rank == 1) << "--rank must be 0 or 1";
+    CHECK_EQ(world_size, 2) << "This example requires exactly 2 ranks";
+    CHECK_EQ(FLAGS_kDataBytes % 16, 0) << "kDataBytes must be a multiple of 16";
+    CHECK_LE(FLAGS_kDataBytes, kSignalWordOffset)
+        << "kDataBytes cannot exceed " << kSignalWordOffset;
+
+    int gpu_id = FLAGS_gpu_id >= 0 ? FLAGS_gpu_id : rank;
+    checkCuda(cudaSetDevice(gpu_id), "cudaSetDevice");
+    LOG(INFO) << "Rank " << rank << " using GPU " << gpu_id;
+
+    // -----------------------------------------------------------------------
+    // 1. Create TransferEngine and get P2pTransport.
+    // -----------------------------------------------------------------------
+    auto engine = std::make_unique<mooncake::TransferEngine>(false);
+
+    std::string local_name = FLAGS_local_server_name;
+    if (local_name.empty()) {
+        local_name = "127.0.0.1:" + std::to_string(12345 + rank);
+    }
+    engine->init(FLAGS_metadata_server, local_name);
+
+    auto* p2p = engine->getOrCreateP2pTransport(world_size);
+    CHECK_NOTNULL(p2p);
+
+    // -----------------------------------------------------------------------
+    // 2. Allocate GDR buffer and fill with initial pattern.
+    // -----------------------------------------------------------------------
+    constexpr size_t kBufSize = kSignalWordOffset + sizeof(int32_t);
+    void* gdr_buffer = p2p->allocateBuffer(kBufSize);
+    CHECK_NOTNULL(gdr_buffer);
+    LOG(INFO) << "Rank " << rank << " allocated GDR buffer: " << gdr_buffer;
+
+    // Zero-initialize the signal word.
+    checkCuda(cudaMemset(gdr_buffer, 0, kBufSize), "cudaMemset");
+
+    // -----------------------------------------------------------------------
+    // 3. Export IPC handle and exchange with peer.
+    // -----------------------------------------------------------------------
+    auto my_handle = p2p->exportIpcHandle(gdr_buffer);
+    CHECK(!my_handle.empty()) << "exportIpcHandle returned empty vector";
+    writeIpcHandle(rank, my_handle);
+
+    // Read peer's handle.
+    int peer_rank = 1 - rank;
+    auto peer_handle = readIpcHandle(peer_rank);
+
+    // Build remote_handles vector: index by rank.
+    std::vector<std::vector<int32_t>> remote_handles(world_size);
+    remote_handles[rank] = my_handle;
+    remote_handles[peer_rank] = peer_handle;
+
+    std::vector<int> active_ranks_mask(world_size, 1);
+
+    p2p->importPeerHandles(gdr_buffer, rank, world_size, remote_handles,
+                           active_ranks_mask);
+
+    CHECK(p2p->allPeersAccessible())
+        << "P2P not accessible between rank " << rank << " and peer";
+    LOG(INFO) << "Rank " << rank << " P2P peer mapping complete";
+
+    // -----------------------------------------------------------------------
+    // 4. Build CommCtx and launch kernel.
+    // -----------------------------------------------------------------------
+    mooncake::device::CommCtx ctx{};
+    ctx.rank = rank;
+    ctx.p2p.available = p2p->availableTablePtr();
+    ctx.p2p.peer_ptrs = p2p->peerPtrsTablePtr();
+    ctx.p2p.local_base = gdr_buffer;
+    // IBGDA fields left as nullptr (not used for P2P-only example).
+
+    // Pass CommCtx by value — CUDA copies it to kernel parameter space.
+    if (rank == 0) {
+        LOG(INFO) << "Rank 0 launching sender kernel...";
+        senderKernel<<<1, 1>>>(ctx, peer_rank, FLAGS_kDataBytes);
+    } else {
+        LOG(INFO) << "Rank 1 launching receiver kernel...";
+        receiverKernel<<<1, 1>>>(ctx, peer_rank, FLAGS_kDataBytes);
+    }
+    checkCuda(cudaDeviceSynchronize(), "cudaDeviceSynchronize");
+
+    // -----------------------------------------------------------------------
+    // 5. Barrier: rank 1 signals completion, rank 0 waits before cleanup.
+    //    The IPC handle is only valid while the original allocation exists,
+    //    so rank 0 must not free its buffer until rank 1 has opened it.
+    // -----------------------------------------------------------------------
+    if (rank == 1) {
+        // Signal rank 0 that we're done.
+        std::string done_path = FLAGS_ipc_dir + "/device_api_ex_done.bin";
+        std::ofstream ofs(done_path);
+        ofs << "1";
+        ofs.close();
+    } else {
+        // Wait for rank 1 to finish.
+        std::string done_path = FLAGS_ipc_dir + "/device_api_ex_done.bin";
+        for (int attempt = 0; attempt < 300; ++attempt) {
+            std::ifstream ifs(done_path);
+            if (ifs) break;
+            std::this_thread::sleep_for(std::chrono::milliseconds(100));
+        }
+    }
+
+    // -----------------------------------------------------------------------
+    // 6. Cleanup.
+    // -----------------------------------------------------------------------
+    p2p->freeBuffer(gdr_buffer);
+
+    // Remove IPC handle files.
+    std::remove(ipcFilePath(rank).c_str());
+    if (rank == 0) {
+        std::remove((FLAGS_ipc_dir + "/device_api_ex_done.bin").c_str());
+    }
+
+    LOG(INFO) << "Rank " << rank << " done.";
+    return 0;
+}
diff --git a/mooncake-transfer-engine/example/efa_first_submit_probe.cpp b/mooncake-transfer-engine/example/efa_first_submit_probe.cpp
new file mode 100644
index 00000000..ffef327e
--- /dev/null
+++ b/mooncake-transfer-engine/example/efa_first_submit_probe.cpp
@@ -0,0 +1,199 @@
+// Copyright 2024 KVCache.AI
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// First-submit latency probe for EFA transport.
+//
+// transfer_engine_bench reports a 10-second throughput average, which
+// hides the two costs that matter for first-request latency: the
+// handshake / fi_av_insert that fires on the first send to each
+// (local_NIC, peer_NIC) pair, and the optional warmupSegment() that
+// pre-pays that cost outside the critical path.  This probe times
+// those two things separately so the cost of each is directly visible.
+//
+// Reports:
+//   - warmup:             time for EfaTransport::warmupSegment() (if enabled)
+//   - submit #0..#N-1:    time for a single 1 MB submitTransfer + poll
+//
+// The first submit carries the handshake cost when warmup is off; the
+// rest are steady-state.  It's a two-host tool (not hooked into ctest)
+// used to:
+//   (a) decide whether to call warmupSegment() in your application —
+//       does it actually remove the stall?
+//   (b) compare this PR's shared-endpoint cost against upstream
+//       per-peer-endpoint cost (see the numbers in docs/.../efa_transport.md
+//       under "First-request latency").
+//
+// Usage:
+//   Target:     ./efa_first_submit_probe --mode=target ...
+//   Initiator:  ./efa_first_submit_probe --mode=initiator \
+//                   --segment_id=<target_host>:<port> \
+//                   --warmup=1 --iters=5
+
+#include <gflags/gflags.h>
+#include <glog/logging.h>
+#include <numa.h>
+
+#include <chrono>
+#include <cstring>
+#include <memory>
+#include <string>
+
+#include "transfer_engine.h"
+#include "transport/efa_transport/efa_transport.h"
+#include "transport/transport.h"
+
+using namespace mooncake;
+
+DEFINE_string(mode, "initiator", "initiator or target");
+DEFINE_string(metadata_server, "P2PHANDSHAKE", "metadata backend");
+DEFINE_string(local_server_name, "",
+              "local server name (host:port); default = hostname:12345");
+DEFINE_string(
+    segment_id, "",
+    "target <host>:<port> from target's startup log (initiator only)");
+DEFINE_bool(warmup, true, "call warmupSegment before first submit");
+DEFINE_int32(iters, 5, "number of timed submits after warmup");
+DEFINE_uint64(xfer_size, 1 << 20, "bytes per submit (default 1 MB)");
+DEFINE_uint64(buffer_size, 1ULL << 30, "registered buffer size (default 1 GB)");
+
+static std::string defaultLocalServerName() {
+    char host[256];
+    if (gethostname(host, sizeof(host)) != 0) return "127.0.0.1:12345";
+    return std::string(host) + ":12345";
+}
+
+static double toMs(std::chrono::steady_clock::duration d) {
+    return std::chrono::duration<double, std::milli>(d).count();
+}
+
+static int runTarget() {
+    auto engine = std::make_unique<TransferEngine>(false);
+    engine->getLocalTopology()->discover({});
+
+    std::string name = FLAGS_local_server_name.empty()
+                           ? defaultLocalServerName()
+                           : FLAGS_local_server_name;
+    auto hp = parseHostNameWithPort(name);
+    int rc =
+        engine->init(FLAGS_metadata_server, name, hp.first.c_str(), hp.second);
+    LOG_ASSERT(!rc) << "engine init failed";
+
+    Transport *xport = engine->installTransport("efa", nullptr);
+    LOG_ASSERT(xport) << "installTransport(efa) failed";
+
+    void *buf = numa_alloc_onnode(FLAGS_buffer_size, 0);
+    LOG_ASSERT(buf) << "numa_alloc_onnode failed";
+
+    rc = engine->registerLocalMemory(buf, FLAGS_buffer_size, "cpu:0");
+    LOG_ASSERT(!rc) << "registerLocalMemory failed";
+
+    LOG(INFO) << "[target] ready, addr=" << engine->getLocalIpAndPort();
+    LOG(INFO) << "[target] Ctrl-C to stop";
+
+    pause();  // block until SIGINT
+
+    engine->unregisterLocalMemory(buf);
+    numa_free(buf, FLAGS_buffer_size);
+    return 0;
+}
+
+static int runInitiator() {
+    if (FLAGS_segment_id.empty()) {
+        LOG(ERROR) << "--segment_id is required for initiator";
+        return 1;
+    }
+
+    auto engine = std::make_unique<TransferEngine>(false);
+    engine->getLocalTopology()->discover({});
+
+    std::string name = FLAGS_local_server_name.empty()
+                           ? defaultLocalServerName()
+                           : FLAGS_local_server_name;
+    auto hp = parseHostNameWithPort(name);
+    int rc =
+        engine->init(FLAGS_metadata_server, name, hp.first.c_str(), hp.second);
+    LOG_ASSERT(!rc) << "engine init failed";
+
+    Transport *xport = engine->installTransport("efa", nullptr);
+    LOG_ASSERT(xport) << "installTransport(efa) failed";
+
+    void *buf = numa_alloc_onnode(FLAGS_buffer_size, 0);
+    LOG_ASSERT(buf) << "numa_alloc_onnode failed";
+    memset(buf, 0xAB, FLAGS_xfer_size);
+
+    rc = engine->registerLocalMemory(buf, FLAGS_buffer_size, "cpu:0");
+    LOG_ASSERT(!rc) << "registerLocalMemory failed";
+
+    auto segment_id = engine->openSegment(FLAGS_segment_id);
+    LOG_ASSERT(segment_id != (SegmentID)-1) << "openSegment failed";
+
+    auto seg_desc = engine->getMetadata()->getSegmentDescByID(segment_id);
+    LOG_ASSERT(seg_desc) << "getSegmentDescByID failed";
+    uint64_t remote_base = (uint64_t)seg_desc->buffers[0].addr;
+
+    LOG(INFO) << "[initiator] peer=" << FLAGS_segment_id
+              << " warmup=" << (FLAGS_warmup ? "ON" : "OFF")
+              << " iters=" << FLAGS_iters << " xfer_size=" << FLAGS_xfer_size;
+
+    if (FLAGS_warmup) {
+        auto *efa = dynamic_cast<EfaTransport *>(xport);
+        LOG_ASSERT(efa) << "transport is not EfaTransport";
+        auto t0 = std::chrono::steady_clock::now();
+        int wrc = efa->warmupSegment(FLAGS_segment_id);
+        auto elapsed = std::chrono::steady_clock::now() - t0;
+        LOG(INFO) << "warmup: " << toMs(elapsed) << " ms (rc=" << wrc << ")";
+    }
+
+    for (int i = 0; i < FLAGS_iters; ++i) {
+        auto t0 = std::chrono::steady_clock::now();
+
+        auto batch_id = engine->allocateBatchID(1);
+        TransferRequest req;
+        req.opcode = TransferRequest::WRITE;
+        req.length = FLAGS_xfer_size;
+        req.source = (uint8_t *)buf;
+        req.target_id = segment_id;
+        req.target_offset = remote_base;
+
+        Status s = engine->submitTransfer(batch_id, {req});
+        LOG_ASSERT(s.ok()) << "submitTransfer #" << i
+                           << " failed: " << s.ToString();
+
+        TransferStatus status;
+        for (int poll = 0; poll < 20'000'000; ++poll) {
+            s = engine->getTransferStatus(batch_id, 0, status);
+            LOG_ASSERT(s.ok());
+            if (status.s == TransferStatusEnum::COMPLETED ||
+                status.s == TransferStatusEnum::FAILED)
+                break;
+        }
+        LOG_ASSERT(status.s == TransferStatusEnum::COMPLETED)
+            << "submit #" << i << " did not complete";
+
+        engine->freeBatchID(batch_id);
+
+        auto elapsed = std::chrono::steady_clock::now() - t0;
+        LOG(INFO) << "submit #" << i << ": " << toMs(elapsed) << " ms";
+    }
+
+    engine->unregisterLocalMemory(buf);
+    numa_free(buf, FLAGS_buffer_size);
+    return 0;
+}
+
+int main(int argc, char **argv) {
+    gflags::ParseCommandLineFlags(&argc, &argv, false);
+    google::InitGoogleLogging(argv[0]);
+    FLAGS_logtostderr = 1;
+
+    if (FLAGS_mode == "target") return runTarget();
+    if (FLAGS_mode == "initiator") return runInitiator();
+
+    LOG(ERROR) << "--mode must be 'target' or 'initiator'";
+    return 1;
+}
diff --git a/mooncake-transfer-engine/example/efa_latency_bench.png b/mooncake-transfer-engine/example/efa_latency_bench.png
deleted file mode 100644
index 34693c07..00000000
Binary files a/mooncake-transfer-engine/example/efa_latency_bench.png and /dev/null differ
diff --git a/mooncake-transfer-engine/example/efa_latency_bench.py b/mooncake-transfer-engine/example/efa_latency_bench.py
deleted file mode 100755
index 42f288e2..00000000
--- a/mooncake-transfer-engine/example/efa_latency_bench.py
+++ /dev/null
@@ -1,454 +0,0 @@
-#!/usr/bin/env python3
-# Copyright 2024 KVCache.AI
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-"""
-EFA Latency vs Cache Size Benchmark
-
-Runs transfer_engine_bench on two machines via SSH to measure throughput
-for multiple EFA configurations. For each cache size point, an independent
-benchmark run measures throughput, then latency = cache_size / throughput
-is computed. Running separate measurements per point captures natural
-throughput variation between runs.
-
-Produces a chart similar to image/transfer-engine-performance.png.
-
-Usage:
-    python3 efa_latency_bench.py
-    python3 efa_latency_bench.py --target_host=HOST_A --initiator_host=HOST_B
-    python3 efa_latency_bench.py --output=my_chart.png --duration=10
-"""
-
-import argparse
-import os
-import re
-import subprocess
-import sys
-import time
-
-
-def parse_args():
-    parser = argparse.ArgumentParser(
-        description="EFA Latency vs Cache Size Benchmark"
-    )
-    parser.add_argument(
-        "--target_host",
-        default="ip-172-31-22-204",
-        help="Hostname or IP of the target machine (default: ip-172-31-22-204)",
-    )
-    parser.add_argument(
-        "--initiator_host",
-        default="ip-172-31-26-160",
-        help="Hostname or IP of the initiator machine (default: ip-172-31-26-160)",
-    )
-    parser.add_argument(
-        "--build_dir",
-        default="/home/ubuntu/Mooncake-efa/build-efa",
-        help="Path to the Mooncake build directory on both machines",
-    )
-    parser.add_argument(
-        "--duration",
-        type=int,
-        default=10,
-        help="Benchmark duration in seconds per measurement point (default: 10)",
-    )
-    parser.add_argument(
-        "--output",
-        default="efa_latency_bench.png",
-        help="Output chart filename (default: efa_latency_bench.png)",
-    )
-    parser.add_argument(
-        "--cache_sizes",
-        default="1,5,10,20,30,40,50,60,70,80,90,100",
-        help="Comma-separated cache sizes in GB (default: 1,5,10,20,...,100)",
-    )
-    parser.add_argument(
-        "--operation",
-        default="write",
-        choices=["read", "write"],
-        help="Transfer operation type (default: write)",
-    )
-    parser.add_argument(
-        "--ssh_user",
-        default="ubuntu",
-        help="SSH username for remote connections (default: ubuntu)",
-    )
-    parser.add_argument(
-        "--ssh_opts",
-        default="-o StrictHostKeyChecking=no -o ConnectTimeout=10",
-        help="Additional SSH options",
-    )
-    parser.add_argument(
-        "--annotation_gb",
-        type=int,
-        default=40,
-        help="Cache size (GB) at which to annotate speedup ratio (default: 40)",
-    )
-    parser.add_argument(
-        "--buffer_size",
-        type=int,
-        default=1073741824,
-        help="Buffer size in bytes for bench tool (default: 1073741824 = 1 GiB)",
-    )
-    return parser.parse_args()
-
-
-# Benchmark configurations
-CONFIGS = [
-    {
-        "label": "EFA (tuned)",
-        "env": {"MC_SLICE_SIZE": "262144"},
-        "flags": {
-            "threads": 48,
-            "block_size": 131072,  # 128KB
-            "batch_size": 128,
-        },
-        "color": "#1f77b4",  # blue
-        "marker": "o",
-    },
-    {
-        "label": "EFA (default)",
-        "env": {},
-        "flags": {
-            "threads": 8,
-            "block_size": 65536,  # 64KB
-            "batch_size": 128,
-        },
-        "color": "#ff7f0e",  # orange
-        "marker": "^",
-    },
-]
-
-
-def run_ssh(host, command, user="ubuntu", ssh_opts="", timeout=None):
-    """Run a command on a remote host via SSH."""
-    ssh_args = [
-        "ssh",
-        *ssh_opts.split(),
-        f"{user}@{host}",
-        command,
-    ]
-    try:
-        result = subprocess.run(
-            ssh_args, capture_output=True, text=True, timeout=timeout,
-        )
-        return result.returncode, result.stdout, result.stderr
-    except subprocess.TimeoutExpired:
-        return -1, "", "Command timed out"
-
-
-def run_cmd(cmd, timeout=None):
-    """Run a local shell command."""
-    try:
-        result = subprocess.run(
-            cmd, shell=True, capture_output=True, text=True, timeout=timeout,
-        )
-        return result.returncode, result.stdout, result.stderr
-    except subprocess.TimeoutExpired:
-        return -1, "", "Command timed out"
-
-
-def kill_bench(target_host, user, ssh_opts):
-    """Kill any running transfer_engine_bench processes."""
-    kill_cmd = (
-        "ps aux | grep '[t]ransfer_engine_bench' | awk '{print $2}' "
-        "| xargs -r kill 2>/dev/null; sleep 1; echo done"
-    )
-    run_ssh(target_host, kill_cmd, user, ssh_opts, timeout=15)
-    run_cmd(
-        "ps aux | grep '[t]ransfer_engine_bench' | awk '{print $2}' "
-        "| xargs -r kill 2>/dev/null",
-        timeout=5,
-    )
-    time.sleep(2)
-
-
-def start_target(target_host, build_dir, buffer_size, user, ssh_opts):
-    """Start the target process via SSH. Returns target address or None."""
-    bench_bin = os.path.join(
-        build_dir, "mooncake-transfer-engine/example/transfer_engine_bench"
-    )
-    log_file = "/tmp/efa_bench_target.log"
-
-    target_cmd = (
-        f"cd {build_dir} && "
-        f"env MC_METADATA_SERVER=P2PHANDSHAKE "
-        f"{bench_bin} "
-        f"--mode=target --protocol=efa --metadata_server=P2PHANDSHAKE "
-        f"--buffer_size={buffer_size} "
-        f"> {log_file} 2>&1"
-    )
-    ssh_args = [
-        "ssh", "-n",
-        *ssh_opts.split(),
-        f"{user}@{target_host}",
-        target_cmd,
-    ]
-    subprocess.Popen(
-        ssh_args,
-        stdout=subprocess.DEVNULL,
-        stderr=subprocess.DEVNULL,
-        stdin=subprocess.DEVNULL,
-    )
-
-    for _ in range(20):
-        time.sleep(1)
-        rc, stdout, _ = run_ssh(
-            target_host,
-            f"grep 'listening on' {log_file} 2>/dev/null",
-            user, ssh_opts, timeout=10,
-        )
-        if rc == 0 and "listening on" in stdout:
-            match = re.search(r"listening on (\S+:\d+)", stdout)
-            if match:
-                return match.group(1)
-
-    _, log_out, _ = run_ssh(
-        target_host, f"tail -20 {log_file}", user, ssh_opts, timeout=10
-    )
-    print(f"  Target log:\n{log_out}", file=sys.stderr)
-    return None
-
-
-def run_initiator(initiator_host, build_dir, target_addr, config,
-                  buffer_size, duration, operation, user, ssh_opts):
-    """Run the initiator benchmark. Returns throughput in GB/s or None."""
-    bench_bin = os.path.join(
-        build_dir, "mooncake-transfer-engine/example/transfer_engine_bench"
-    )
-    flags = config["flags"]
-
-    env_parts = ["MC_METADATA_SERVER=P2PHANDSHAKE"]
-    for key, val in config.get("env", {}).items():
-        env_parts.append(f"{key}={val}")
-    env_str = " ".join(env_parts)
-
-    bench_cmd = (
-        f"cd {build_dir} && "
-        f"{env_str} {bench_bin} "
-        f"--mode=initiator --protocol=efa --metadata_server=P2PHANDSHAKE "
-        f"--segment_id={target_addr} "
-        f"--operation={operation} "
-        f"--duration={duration} "
-        f"--threads={flags['threads']} "
-        f"--block_size={flags['block_size']} "
-        f"--batch_size={flags['batch_size']} "
-        f"--buffer_size={buffer_size} "
-        f"--report_unit=GB "
-        f"2>&1"
-    )
-
-    timeout = duration + 60
-    rc, stdout, stderr = run_ssh(
-        initiator_host, bench_cmd, user, ssh_opts, timeout=timeout
-    )
-
-    combined = stdout + "\n" + stderr
-    match = re.search(r"throughput\s+([\d.]+)\s+GB/s", combined)
-    if match:
-        return float(match.group(1))
-
-    print(f"  WARNING: Could not parse throughput", file=sys.stderr)
-    lines = combined.strip().split("\n")
-    for line in lines[-3:]:
-        print(f"    {line}", file=sys.stderr)
-    return None
-
-
-def plot_results(results, cache_sizes_gb, annotation_gb, output_path):
-    """Generate and save the latency vs cache size chart."""
-    import matplotlib
-    matplotlib.use("Agg")
-    import matplotlib.pyplot as plt
-
-    fig, ax = plt.subplots(1, 1, figsize=(7, 5))
-
-    for entry in results:
-        ax.plot(
-            cache_sizes_gb,
-            entry["latencies"],
-            color=entry["color"],
-            marker=entry["marker"],
-            markersize=5,
-            linewidth=1.5,
-            label=entry["label"],
-        )
-
-    ax.set_xlabel("Cache Size (GB)", fontsize=12)
-    ax.set_ylabel("Latency (s)", fontsize=12)
-    ax.set_xlim(0, max(cache_sizes_gb))
-    ax.set_ylim(bottom=0)
-    ax.grid(True, alpha=0.3)
-
-    # Speedup annotation
-    if len(results) >= 2 and annotation_gb in cache_sizes_gb:
-        idx = cache_sizes_gb.index(annotation_gb)
-        tuned_lat = results[0]["latencies"][idx]
-        default_lat = results[1]["latencies"][idx]
-
-        if tuned_lat > 0:
-            speedup = default_lat / tuned_lat
-            ax.axvline(x=annotation_gb, color="gray", linestyle="--", alpha=0.5)
-            mid_y = (tuned_lat + default_lat) / 2
-            ax.annotate(
-                "",
-                xy=(annotation_gb + 1, tuned_lat),
-                xytext=(annotation_gb + 1, default_lat),
-                arrowprops=dict(arrowstyle="<->", color="black", lw=1.5),
-            )
-            ax.text(
-                annotation_gb + 3, mid_y,
-                f"{speedup:.1f}x",
-                fontsize=11, fontweight="bold", va="center",
-            )
-
-    ax.text(
-        0.02, 0.98, "8 x 400 Gbps EFA NICs",
-        transform=ax.transAxes, fontsize=10,
-        verticalalignment="top", fontweight="bold",
-    )
-    ax.legend(fontsize=10, loc="upper left", bbox_to_anchor=(0.0, 0.90))
-
-    plt.tight_layout()
-    plt.savefig(output_path, dpi=150, bbox_inches="tight")
-    print(f"\nChart saved to: {output_path}")
-
-
-def main():
-    args = parse_args()
-    cache_sizes_gb = [float(x) for x in args.cache_sizes.split(",")]
-
-    print("=" * 60)
-    print("EFA Latency vs Cache Size Benchmark")
-    print("=" * 60)
-    print(f"  Target host:    {args.target_host}")
-    print(f"  Initiator host: {args.initiator_host}")
-    print(f"  Build dir:      {args.build_dir}")
-    print(f"  Duration:       {args.duration}s per measurement")
-    print(f"  Operation:      {args.operation}")
-    print(f"  Buffer size:    {args.buffer_size} bytes")
-    print(f"  Cache sizes:    {cache_sizes_gb} GB")
-    print(f"  Output:         {args.output}")
-    print(f"  Configs:        {len(CONFIGS)}")
-    n_runs = len(CONFIGS) * len(cache_sizes_gb)
-    print(f"  Total runs:     {n_runs} "
-          f"(each {args.duration}s + overhead)")
-    print()
-
-    results = []
-
-    for ci, config in enumerate(CONFIGS):
-        label = config["label"]
-        print(f"[{ci+1}/{len(CONFIGS)}] Config: {label}")
-        print(f"  threads={config['flags']['threads']}, "
-              f"block_size={config['flags']['block_size']}, "
-              f"batch_size={config['flags']['batch_size']}")
-        if config.get("env"):
-            print(f"  env: {config['env']}")
-        print()
-
-        throughputs = []
-        latencies = []
-
-        # Start target once for all cache sizes of this config
-        kill_bench(args.target_host, args.ssh_user, args.ssh_opts)
-        print(f"  Starting target...", end="", flush=True)
-        target_addr = start_target(
-            args.target_host, args.build_dir, args.buffer_size,
-            args.ssh_user, args.ssh_opts,
-        )
-        if not target_addr:
-            print(" FAILED")
-            continue
-        print(f" ready ({target_addr})")
-        print()
-
-        for si, cache_gb in enumerate(cache_sizes_gb):
-            tag = f"  [{si+1}/{len(cache_sizes_gb)}] {cache_gb:6.0f} GB"
-            print(f"{tag} ...", end="", flush=True)
-
-            tp = run_initiator(
-                args.initiator_host, args.build_dir,
-                target_addr, config, args.buffer_size,
-                args.duration, args.operation,
-                args.ssh_user, args.ssh_opts,
-            )
-
-            if tp is None:
-                print(" FAILED")
-                throughputs.append(None)
-                latencies.append(None)
-                continue
-
-            lat = cache_gb / tp
-            throughputs.append(tp)
-            latencies.append(lat)
-            print(f" {tp:7.2f} GB/s  lat={lat:.3f}s")
-
-        # Cleanup
-        kill_bench(args.target_host, args.ssh_user, args.ssh_opts)
-
-        valid = [t for t in throughputs if t is not None]
-        if not valid:
-            print(f"  ERROR: No results for {label}\n", file=sys.stderr)
-            continue
-
-        results.append({
-            "label": label,
-            "throughputs": throughputs,
-            "latencies": latencies,
-            "color": config["color"],
-            "marker": config["marker"],
-        })
-        print()
-
-    if not results:
-        print("ERROR: No results collected.", file=sys.stderr)
-        sys.exit(1)
-
-    # Summary table
-    print("=" * 60)
-    print("Results Summary")
-    print("=" * 60)
-    header = f"{'Cache(GB)':>10}"
-    for r in results:
-        header += f"  {r['label']:>24}"
-    print(header)
-    print("-" * len(header))
-    for i, cache_gb in enumerate(cache_sizes_gb):
-        row = f"{cache_gb:>10.0f}"
-        for r in results:
-            tp = r["throughputs"][i]
-            lat = r["latencies"][i]
-            if tp is not None:
-                row += f"  {tp:7.2f} GB/s  lat={lat:.3f}s"
-            else:
-                row += f"  {'N/A':>24}"
-        print(row)
-    print()
-
-    # Plot only valid points
-    plot_cache = []
-    plot_data = [{**r, "latencies": []} for r in results]
-    for i, cache_gb in enumerate(cache_sizes_gb):
-        if all(r["latencies"][i] is not None for r in results):
-            plot_cache.append(cache_gb)
-            for j in range(len(results)):
-                plot_data[j]["latencies"].append(results[j]["latencies"][i])
-
-    plot_results(plot_data, plot_cache, args.annotation_gb, args.output)
-
-
-if __name__ == "__main__":
-    main()
diff --git a/mooncake-transfer-engine/example/transfer_engine_bench.cpp b/mooncake-transfer-engine/example/transfer_engine_bench.cpp
index e47aafd9..5c653227 100644
--- a/mooncake-transfer-engine/example/transfer_engine_bench.cpp
+++ b/mooncake-transfer-engine/example/transfer_engine_bench.cpp
@@ -42,8 +42,9 @@
 #endif
 #endif
 
-#if defined(USE_CUDA) || defined(USE_MUSA) || defined(USE_HIP) || \
-    defined(USE_MACA) || defined(USE_UBSHMEM)
+#if defined(USE_CUDA) || defined(USE_MUSA) || defined(USE_HIP) ||    \
+    defined(USE_MACA) || defined(USE_HYGON) || defined(USE_COREX) || \
+    defined(USE_UBSHMEM) || defined(USE_SUNRISE)
 #include <cassert>
 
 #if defined(USE_MNNVL) || defined(USE_UBSHMEM)
@@ -55,16 +56,16 @@
 #endif
 
 #if defined(USE_UBSHMEM)
-static void checkAclError(aclError result, const char *message) {
+static void checkAclError(aclError result, const char* message) {
     if (result != ACL_ERROR_NONE) {
-        const char *errMsg = aclGetRecentErrMsg();
+        const char* errMsg = aclGetRecentErrMsg();
         LOG(ERROR) << message << " (Error code: " << result << " - " << errMsg
                    << ")";
         exit(EXIT_FAILURE);
     }
 }
 #else
-static void checkCudaError(cudaError_t result, const char *message) {
+static void checkCudaError(cudaError_t result, const char* message) {
     if (result != cudaSuccess) {
         LOG(ERROR) << message << " (Error code: " << result << " - "
                    << cudaGetErrorString(result) << ")" << std::endl;
@@ -88,7 +89,8 @@ DEFINE_string(mode, "initiator",
 DEFINE_string(operation, "read", "Operation type: read or write");
 
 DEFINE_string(protocol, "rdma",
-              "Transfer protocol: rdma|barex|tcp|efa|nvlink|nvlink_intra|hip");
+              "Transfer protocol: "
+              "rdma|barex|tcp|efa|nvlink|nvlink_intra|hip|sunrise_link");
 
 DEFINE_string(device_name, "mlx5_2",
               "Device name to use, valid if protocol=rdma");
@@ -106,8 +108,9 @@ DEFINE_string(report_unit, "GB", "Report unit: GB|GiB|Gb|MB|MiB|Mb|KB|KiB|Kb");
 DEFINE_uint32(report_precision, 2, "Report precision");
 DEFINE_string(backend, "classic", "Backend to use: classic|tent");
 
-#if defined(USE_CUDA) || defined(USE_MUSA) || defined(USE_HIP) || \
-    defined(USE_MACA) || defined(USE_UBSHMEM)
+#if defined(USE_CUDA) || defined(USE_MUSA) || defined(USE_HIP) ||    \
+    defined(USE_MACA) || defined(USE_HYGON) || defined(USE_COREX) || \
+    defined(USE_UBSHMEM) || defined(USE_SUNRISE)
 DEFINE_bool(use_vram, true, "Allocate memory from GPU/NPU VRAM");
 DEFINE_bool(init_mem, true, "Initialize allocated memory");
 DEFINE_int32(gpu_id, 0,
@@ -116,10 +119,11 @@ DEFINE_int32(gpu_id, 0,
 
 using namespace mooncake;
 
-static void *allocateMemoryPool(size_t size, int buffer_id,
+static void* allocateMemoryPool(size_t size, int buffer_id,
                                 bool from_vram = false) {
-#if defined(USE_CUDA) || defined(USE_MUSA) || defined(USE_HIP) || \
-    defined(USE_MACA) || defined(USE_UBSHMEM)
+#if defined(USE_CUDA) || defined(USE_MUSA) || defined(USE_HIP) ||    \
+    defined(USE_MACA) || defined(USE_HYGON) || defined(USE_COREX) || \
+    defined(USE_UBSHMEM) || defined(USE_SUNRISE)
     if (from_vram) {
         int gpu_id;
         if (FLAGS_gpu_id == -1) {
@@ -127,7 +131,7 @@ static void *allocateMemoryPool(size_t size, int buffer_id,
         } else {
             gpu_id = FLAGS_gpu_id;
         }
-        void *d_buf;
+        void* d_buf;
 #if defined(USE_UBSHMEM)
         LOG(INFO) << "Allocating memory on NPU " << gpu_id;
         checkAclError(aclrtSetDevice(gpu_id), "Failed to set device");
@@ -188,9 +192,10 @@ static void *allocateMemoryPool(size_t size, int buffer_id,
     return numa_alloc_onnode(size, buffer_id);
 }
 
-static void freeMemoryPool(void *addr, size_t size) {
-#if defined(USE_CUDA) || defined(USE_MUSA) || defined(USE_HIP) || \
-    defined(USE_MACA) || defined(USE_UBSHMEM)
+static void freeMemoryPool(void* addr, size_t size) {
+#if defined(USE_CUDA) || defined(USE_MUSA) || defined(USE_HIP) ||    \
+    defined(USE_MACA) || defined(USE_HYGON) || defined(USE_COREX) || \
+    defined(USE_UBSHMEM) || defined(USE_SUNRISE)
     if (FLAGS_protocol == "nvlink" || FLAGS_protocol == "hip") {
 #ifdef USE_MNNVL
         if (FLAGS_use_vram) {
@@ -214,19 +219,35 @@ static void freeMemoryPool(void *addr, size_t size) {
 #endif
     } else {
 #ifndef USE_UBSHMEM
-        // check pointer on GPU
-        cudaPointerAttributes attributes;
-        checkCudaError(cudaPointerGetAttributes(&attributes, addr),
-                       "Failed to get pointer attributes");
-
-        if (attributes.type == cudaMemoryTypeDevice) {
-            cudaFree(addr);
-        } else if (attributes.type == cudaMemoryTypeHost ||
-                   attributes.type == cudaMemoryTypeUnregistered) {
+        // Check FLAGS_use_vram first to avoid unnecessary CUDA calls in non-GPU
+        // environments
+        if (!FLAGS_use_vram) {
+            // Memory was allocated via numa_alloc_onnode, free it directly
             numa_free(addr, size);
         } else {
-            LOG(ERROR) << "Unknown memory type, " << addr << " "
-                       << attributes.type;
+            // Memory may be on GPU, check pointer attributes
+            // Use graceful error handling like transfer engine core
+            // (memory_location.cpp)
+            cudaPointerAttributes attributes;
+            cudaError_t result = cudaPointerGetAttributes(&attributes, addr);
+
+            if (result != cudaSuccess) {
+                // CUDA call failed when FLAGS_use_vram is true - this is an
+                // error
+                LOG(ERROR) << "cudaPointerGetAttributes failed (Error code: "
+                           << result << " - " << cudaGetErrorString(result)
+                           << ")";
+                numa_free(addr, size);
+            } else if (attributes.type == cudaMemoryTypeDevice) {
+                cudaFree(addr);
+            } else if (attributes.type == cudaMemoryTypeHost ||
+                       attributes.type == cudaMemoryTypeUnregistered) {
+                numa_free(addr, size);
+            } else {
+                LOG(ERROR) << "Unknown memory type, " << addr << " "
+                           << attributes.type << ", assuming CPU memory";
+                numa_free(addr, size);
+            }
         }
 #endif
     }
@@ -271,10 +292,23 @@ static inline std::string calculateRate(uint64_t data_bytes, double duration) {
 volatile bool running = true;
 std::atomic<size_t> total_batch_count(0);
 
+// Ensure each worker thread has a valid GPU context before issuing transfers.
+static inline void setWorkerDeviceIfNeeded() {
+#if defined(USE_CUDA) || defined(USE_MUSA) || defined(USE_HIP) ||    \
+    defined(USE_MACA) || defined(USE_HYGON) || defined(USE_COREX) || \
+    defined(USE_SUNRISE)
+    if (FLAGS_use_vram && FLAGS_gpu_id >= 0) {
+        checkCudaError(cudaSetDevice(FLAGS_gpu_id),
+                       "Failed to set device in worker");
+    }
+#endif
+}
+
 // Common helper to determine buffer count based on GPU/NUMA configuration
 static int determineBufferCount() {
-#if defined(USE_CUDA) || defined(USE_MUSA) || defined(USE_HIP) || \
-    defined(USE_MACA)
+#if defined(USE_CUDA) || defined(USE_MUSA) || defined(USE_HIP) ||    \
+    defined(USE_MACA) || defined(USE_HYGON) || defined(USE_COREX) || \
+    defined(USE_SUNRISE)
     if (FLAGS_use_vram) {
         int gpu_num;
         LOG(INFO) << "VRAM is used";
@@ -301,11 +335,12 @@ static int determineBufferCount() {
 }
 
 // Common helper to allocate memory buffers
-static std::vector<void *> allocateBuffers() {
+static std::vector<void*> allocateBuffers() {
     buffer_num = determineBufferCount();
-    std::vector<void *> addr(buffer_num);
-#if defined(USE_CUDA) || defined(USE_MUSA) || defined(USE_HIP) || \
-    defined(USE_MACA) || defined(USE_UBSHMEM)
+    std::vector<void*> addr(buffer_num);
+#if defined(USE_CUDA) || defined(USE_MUSA) || defined(USE_HIP) ||    \
+    defined(USE_MACA) || defined(USE_HYGON) || defined(USE_COREX) || \
+    defined(USE_UBSHMEM) || defined(USE_SUNRISE)
     for (int i = 0; i < buffer_num; ++i) {
         addr[i] = allocateMemoryPool(FLAGS_buffer_size, i, FLAGS_use_vram);
     }
@@ -318,7 +353,7 @@ static std::vector<void *> allocateBuffers() {
 }
 
 // Common helper to free memory buffers
-static void freeBuffers(std::vector<void *> &addr) {
+static void freeBuffers(std::vector<void*>& addr) {
     for (int i = 0; i < buffer_num; ++i) {
         freeMemoryPool(addr[i], FLAGS_buffer_size);
     }
@@ -327,8 +362,9 @@ static void freeBuffers(std::vector<void *> &addr) {
 
 // Helper to get location name for classic backend
 static std::string getLocationName(int buffer_id) {
-#if defined(USE_CUDA) || defined(USE_MUSA) || defined(USE_HIP) || \
-    defined(USE_MACA) || defined(USE_UBSHMEM)
+#if defined(USE_CUDA) || defined(USE_MUSA) || defined(USE_HIP) ||    \
+    defined(USE_MACA) || defined(USE_HYGON) || defined(USE_COREX) || \
+    defined(USE_UBSHMEM) || defined(USE_SUNRISE)
     if (FLAGS_use_vram) {
         int name_suffix = (FLAGS_gpu_id == -1) ? buffer_id : FLAGS_gpu_id;
         return std::string(GPU_PREFIX) + std::to_string(name_suffix);
@@ -337,8 +373,8 @@ static std::string getLocationName(int buffer_id) {
     return "cpu:" + std::to_string(buffer_id);
 }
 
-Status initiatorWorker(TransferEngine *engine, SegmentID segment_id,
-                       int thread_id, void *addr) {
+Status initiatorWorker(TransferEngine* engine, SegmentID segment_id,
+                       int thread_id, void* addr) {
     bindToSocket(thread_id % NR_SOCKETS);
     TransferRequest::OpCode opcode;
     if (FLAGS_operation == "read")
@@ -367,7 +403,7 @@ Status initiatorWorker(TransferEngine *engine, SegmentID segment_id,
             TransferRequest entry;
             entry.opcode = opcode;
             entry.length = FLAGS_block_size;
-            entry.source = (uint8_t *)(addr) +
+            entry.source = (uint8_t*)(addr) +
                            FLAGS_block_size * (i * FLAGS_threads + thread_id);
             entry.target_id = segment_id;
             entry.target_offset =
@@ -404,7 +440,7 @@ Status initiatorWorker(TransferEngine *engine, SegmentID segment_id,
     return Status::OK();
 }
 
-std::string formatDeviceNames(const std::string &device_names) {
+std::string formatDeviceNames(const std::string& device_names) {
     std::stringstream ss(device_names);
     std::string item;
     std::vector<std::string> tokens;
@@ -451,19 +487,19 @@ std::string loadNicPriorityMatrix() {
 
 // Common helper to install transport based on protocol flag
 // Returns the installed transport, or nullptr if auto_discovery is enabled
-static Transport *installTransportFromFlags(TransferEngine *engine) {
+static Transport* installTransportFromFlags(TransferEngine* engine) {
     if (FLAGS_auto_discovery) {
         return nullptr;
     }
 
-    Transport *xport = nullptr;
+    Transport* xport = nullptr;
     std::string nic_priority_matrix;
-    std::unique_ptr<void *, decltype(&free)> args(nullptr, free);
+    std::unique_ptr<void*, decltype(&free)> args(nullptr, free);
 
     if (FLAGS_protocol == "rdma" || FLAGS_protocol == "barex") {
         nic_priority_matrix = loadNicPriorityMatrix();
-        args.reset(static_cast<void **>(malloc(2 * sizeof(void *))));
-        args.get()[0] = const_cast<char *>(nic_priority_matrix.c_str());
+        args.reset(static_cast<void**>(malloc(2 * sizeof(void*))));
+        args.get()[0] = const_cast<char*>(nic_priority_matrix.c_str());
         args.get()[1] = nullptr;
         xport = engine->installTransport(FLAGS_protocol.c_str(), args.get());
     } else if (FLAGS_protocol == "ub") {
@@ -476,7 +512,8 @@ static Transport *installTransportFromFlags(TransferEngine *engine) {
         xport = engine->installTransport("efa", nullptr);
     } else if (FLAGS_protocol == "tcp" || FLAGS_protocol == "nvlink" ||
                FLAGS_protocol == "hip" || FLAGS_protocol == "nvlink_intra" ||
-               FLAGS_protocol == "ubshmem") {
+               FLAGS_protocol == "ubshmem" ||
+               FLAGS_protocol == "sunrise_link") {
         xport = engine->installTransport(FLAGS_protocol.c_str(), nullptr);
     } else {
         LOG(ERROR) << "Unsupported protocol: " << FLAGS_protocol;
@@ -494,7 +531,7 @@ int initiator() {
                  hostname_port.first.c_str(), hostname_port.second);
 
     if (!FLAGS_auto_discovery) {
-        Transport *xport = installTransportFromFlags(engine.get());
+        Transport* xport = installTransportFromFlags(engine.get());
         LOG_ASSERT(xport);
     }
 
@@ -590,8 +627,8 @@ int target() {
 namespace tent_backend {
 
 // Helper function to register buffers for TENT backend
-static void registerBuffers(mooncake::tent::TransferEngine *engine,
-                            std::vector<void *> &addr) {
+static void registerBuffers(mooncake::tent::TransferEngine* engine,
+                            std::vector<void*>& addr) {
     for (int i = 0; i < buffer_num; ++i) {
         auto status = engine->registerLocalMemory(addr[i], FLAGS_buffer_size);
         LOG_ASSERT(status.ok())
@@ -600,8 +637,8 @@ static void registerBuffers(mooncake::tent::TransferEngine *engine,
 }
 
 // Helper function to unregister buffers for TENT backend
-static void unregisterBuffers(mooncake::tent::TransferEngine *engine,
-                              std::vector<void *> &addr) {
+static void unregisterBuffers(mooncake::tent::TransferEngine* engine,
+                              std::vector<void*>& addr) {
     for (int i = 0; i < buffer_num; ++i) {
         engine->unregisterLocalMemory(addr[i], FLAGS_buffer_size);
     }
@@ -639,11 +676,12 @@ std::shared_ptr<mooncake::tent::Config> createTentConfig() {
     return config;
 }
 
-void initiatorWorker(mooncake::tent::TransferEngine *engine,
+void initiatorWorker(mooncake::tent::TransferEngine* engine,
                      mooncake::tent::SegmentID segment_id, int thread_id,
-                     void *addr,
-                     const mooncake::tent::SegmentInfo &segment_info) {
+                     void* addr,
+                     const mooncake::tent::SegmentInfo& segment_info) {
     bindToSocket(thread_id % NR_SOCKETS);
+    setWorkerDeviceIfNeeded();
     mooncake::tent::Request::OpCode opcode;
     if (FLAGS_operation == "read")
         opcode = mooncake::tent::Request::READ;
@@ -672,7 +710,7 @@ void initiatorWorker(mooncake::tent::TransferEngine *engine,
             mooncake::tent::Request entry;
             entry.opcode = opcode;
             entry.length = FLAGS_block_size;
-            entry.source = (uint8_t *)(addr) +
+            entry.source = (uint8_t*)(addr) +
                            FLAGS_block_size * (i * FLAGS_threads + thread_id);
             entry.target_id = segment_id;
             entry.target_offset =
@@ -806,7 +844,7 @@ void check_total_buffer_size() {
     }
 }
 
-int main(int argc, char **argv) {
+int main(int argc, char** argv) {
     gflags::ParseCommandLineFlags(&argc, &argv, false);
     check_total_buffer_size();
 
diff --git a/mooncake-transfer-engine/example/transfer_engine_bench_with_notify.cpp b/mooncake-transfer-engine/example/transfer_engine_bench_with_notify.cpp
index 07ead7d8..2d76210f 100644
--- a/mooncake-transfer-engine/example/transfer_engine_bench_with_notify.cpp
+++ b/mooncake-transfer-engine/example/transfer_engine_bench_with_notify.cpp
@@ -38,7 +38,7 @@
 #endif
 
 #if defined(USE_CUDA) || defined(USE_MUSA) || defined(USE_HIP) || \
-    defined(USE_MACA)
+    defined(USE_MACA) || defined(USE_HYGON) || defined(USE_COREX)
 #include <cassert>
 
 #ifdef USE_MNNVL
@@ -59,7 +59,7 @@ static void checkCudaError(cudaError_t result, const char *message) {
 #endif
 
 #if defined(USE_CUDA) || defined(USE_MUSA) || defined(USE_HIP) || \
-    defined(USE_MACA)
+    defined(USE_MACA) || defined(USE_HYGON) || defined(USE_COREX)
 const static int NR_SOCKETS = 4;
 #else
 const static int NR_SOCKETS =
@@ -92,7 +92,7 @@ DEFINE_string(report_unit, "GB", "Report unit: GB|GiB|Gb|MB|MiB|Mb|KB|KiB|Kb");
 DEFINE_uint32(report_precision, 2, "Report precision");
 
 #if defined(USE_CUDA) || defined(USE_MUSA) || defined(USE_HIP) || \
-    defined(USE_MACA)
+    defined(USE_MACA) || defined(USE_HYGON) || defined(USE_COREX)
 DEFINE_bool(use_vram, true, "Allocate memory from GPU VRAM");
 DEFINE_bool(init_mem, true, "Initialize allocated memory");
 DEFINE_int32(gpu_id, 0, "GPU ID to use");
@@ -103,7 +103,7 @@ using namespace mooncake;
 static void *allocateMemoryPool(size_t size, int socket_id,
                                 bool from_vram = false) {
 #if defined(USE_CUDA) || defined(USE_MUSA) || defined(USE_HIP) || \
-    defined(USE_MACA)
+    defined(USE_MACA) || defined(USE_HYGON) || defined(USE_COREX)
     if (from_vram) {
         int gpu_id = FLAGS_gpu_id;
         void *d_buf;
@@ -129,7 +129,7 @@ static void *allocateMemoryPool(size_t size, int socket_id,
 
 static void freeMemoryPool(void *addr, size_t size) {
 #if defined(USE_CUDA) || defined(USE_MUSA) || defined(USE_HIP) || \
-    defined(USE_MACA)
+    defined(USE_MACA) || defined(USE_HYGON) || defined(USE_COREX)
 #ifdef USE_MNNVL
     if (FLAGS_use_vram) {
         freeFabricMemory(addr);
@@ -342,7 +342,7 @@ int initiator() {
     int buffer_num = NR_SOCKETS;
 
 #if defined(USE_CUDA) || defined(USE_MUSA) || defined(USE_HIP) || \
-    defined(USE_MACA)
+    defined(USE_MACA) || defined(USE_HYGON) || defined(USE_COREX)
     if (FLAGS_use_vram) LOG(INFO) << "VRAM is used";
     for (int i = 0; i < buffer_num; ++i) {
         addr[i] = allocateMemoryPool(FLAGS_buffer_size, i, FLAGS_use_vram);
@@ -442,7 +442,7 @@ int target() {
     int buffer_num = NR_SOCKETS;
 
 #if defined(USE_CUDA) || defined(USE_MUSA) || defined(USE_HIP) || \
-    defined(USE_MACA)
+    defined(USE_MACA) || defined(USE_HYGON) || defined(USE_COREX)
     if (FLAGS_use_vram) LOG(INFO) << "VRAM is used";
     for (int i = 0; i < buffer_num; ++i) {
         addr[i] = allocateMemoryPool(FLAGS_buffer_size, i, FLAGS_use_vram);
diff --git a/mooncake-transfer-engine/example/transfer_engine_bench_with_retry.cpp b/mooncake-transfer-engine/example/transfer_engine_bench_with_retry.cpp
index c7011633..c0e2d313 100644
--- a/mooncake-transfer-engine/example/transfer_engine_bench_with_retry.cpp
+++ b/mooncake-transfer-engine/example/transfer_engine_bench_with_retry.cpp
@@ -36,7 +36,7 @@
 #endif
 
 #if defined(USE_CUDA) || defined(USE_MUSA) || defined(USE_HIP) || \
-    defined(USE_MACA)
+    defined(USE_MACA) || defined(USE_HYGON) || defined(USE_COREX)
 #include <cassert>
 
 static void checkCudaError(cudaError_t result, const char *message) {
@@ -77,7 +77,7 @@ DEFINE_string(report_unit, "GB", "Report unit: GB|GiB|Gb|MB|MiB|Mb|KB|KiB|Kb");
 DEFINE_uint32(report_precision, 2, "Report precision");
 
 #if defined(USE_CUDA) || defined(USE_MUSA) || defined(USE_HIP) || \
-    defined(USE_MACA)
+    defined(USE_MACA) || defined(USE_HYGON) || defined(USE_COREX)
 DEFINE_bool(use_vram, true, "Allocate memory from GPU VRAM");
 DEFINE_bool(init_mem, true, "Initialize allocated memory");
 DEFINE_int32(gpu_id, 0, "GPU ID to use");
@@ -88,7 +88,7 @@ using namespace mooncake;
 static void *allocateMemoryPool(size_t size, int socket_id,
                                 bool from_vram = false) {
 #if defined(USE_CUDA) || defined(USE_MUSA) || defined(USE_HIP) || \
-    defined(USE_MACA)
+    defined(USE_MACA) || defined(USE_HYGON) || defined(USE_COREX)
     if (from_vram) {
         int gpu_id = FLAGS_gpu_id;
         void *d_buf;
@@ -111,7 +111,7 @@ static void *allocateMemoryPool(size_t size, int socket_id,
 
 static void freeMemoryPool(void *addr, size_t size) {
 #if defined(USE_CUDA) || defined(USE_MUSA) || defined(USE_HIP) || \
-    defined(USE_MACA)
+    defined(USE_MACA) || defined(USE_HYGON) || defined(USE_COREX)
     // check pointer on GPU
     cudaPointerAttributes attributes;
     checkCudaError(cudaPointerGetAttributes(&attributes, addr),
@@ -304,7 +304,7 @@ int initiator() {
     int buffer_num = NR_SOCKETS;
 
 #if defined(USE_CUDA) || defined(USE_MUSA) || defined(USE_HIP) || \
-    defined(USE_MACA)
+    defined(USE_MACA) || defined(USE_HYGON) || defined(USE_COREX)
     buffer_num = FLAGS_use_vram ? 1 : NR_SOCKETS;
     if (FLAGS_use_vram) LOG(INFO) << "VRAM is used";
     for (int i = 0; i < buffer_num; ++i) {
diff --git a/mooncake-transfer-engine/example/transfer_engine_validator.cpp b/mooncake-transfer-engine/example/transfer_engine_validator.cpp
index 2ac32d7f..0756012a 100644
--- a/mooncake-transfer-engine/example/transfer_engine_validator.cpp
+++ b/mooncake-transfer-engine/example/transfer_engine_validator.cpp
@@ -38,7 +38,7 @@
 #endif
 
 #if defined(USE_CUDA) || defined(USE_MUSA) || defined(USE_HIP) || \
-    defined(USE_MACA)
+    defined(USE_MACA) || defined(USE_HYGON) || defined(USE_COREX)
 #include <cassert>
 
 #ifdef USE_MNNVL
@@ -88,7 +88,7 @@ DEFINE_string(report_unit, "GB", "Report unit: GB|GiB|Gb|MB|MiB|Mb|KB|KiB|Kb");
 DEFINE_uint32(report_precision, 2, "Report precision");
 
 #if defined(USE_CUDA) || defined(USE_MUSA) || defined(USE_HIP) || \
-    defined(USE_MACA)
+    defined(USE_MACA) || defined(USE_HYGON) || defined(USE_COREX)
 DEFINE_bool(use_vram, true, "Allocate memory from GPU VRAM");
 DEFINE_int32(gpu_id, 0, "GPU ID to use, -1 for all GPUs");
 #endif
@@ -98,7 +98,7 @@ using namespace mooncake;
 static void *allocateMemoryPool(size_t size, int buffer_id,
                                 bool from_vram = false) {
 #if defined(USE_CUDA) || defined(USE_MUSA) || defined(USE_HIP) || \
-    defined(USE_MACA)
+    defined(USE_MACA) || defined(USE_HYGON) || defined(USE_COREX)
     if (from_vram) {
         int gpu_id;
         if (FLAGS_gpu_id == -1) {
@@ -123,7 +123,7 @@ static void *allocateMemoryPool(size_t size, int buffer_id,
 
 static void freeMemoryPool(void *addr, size_t size) {
 #if defined(USE_CUDA) || defined(USE_MUSA) || defined(USE_HIP) || \
-    defined(USE_MACA)
+    defined(USE_MACA) || defined(USE_HYGON) || defined(USE_COREX)
 #ifdef USE_MNNVL
     if (FLAGS_use_vram) {
         freeFabricMemory(addr);
@@ -224,7 +224,7 @@ Status submitRequestSync(TransferEngine *engine, SegmentID handle,
 }
 
 #if defined(USE_CUDA) || defined(USE_MUSA) || defined(USE_HIP) || \
-    defined(USE_MACA)
+    defined(USE_MACA) || defined(USE_HYGON) || defined(USE_COREX)
 class PinnedBuffer {
    public:
     explicit PinnedBuffer(size_t size) : size_(size), ptr_(nullptr) {
@@ -260,7 +260,7 @@ thread_local std::vector<uint8_t> user_buf(FLAGS_block_size);
 
 void fillData(int thread_id, void *addr, uint8_t seed) {
 #if defined(USE_CUDA) || defined(USE_MUSA) || defined(USE_HIP) || \
-    defined(USE_MACA)
+    defined(USE_MACA) || defined(USE_HYGON) || defined(USE_COREX)
     memset(ref_buf.data(), seed, FLAGS_block_size);
     cudaStream_t s;
     cudaStreamCreate(&s);
@@ -290,7 +290,7 @@ void checkData(int thread_id, void *addr, uint8_t seed) {
             (uint8_t *)(addr) +
             FLAGS_block_size * (i * FLAGS_threads + thread_id);
 #if defined(USE_CUDA) || defined(USE_MUSA) || defined(USE_HIP) || \
-    defined(USE_MACA)
+    defined(USE_MACA) || defined(USE_HYGON) || defined(USE_COREX)
         cudaStream_t s;
         cudaStreamCreate(&s);
         cudaMemcpyAsync(user_buf.data(), local_addr, FLAGS_block_size,
@@ -413,6 +413,8 @@ int initiator() {
             xport = engine->installTransport("nvlink", nullptr);
         } else if (FLAGS_protocol == "nvlink_intra") {
             xport = engine->installTransport("nvlink_intra", nullptr);
+        } else if (FLAGS_protocol == "maca") {
+            xport = engine->installTransport("maca", nullptr);
         } else if (FLAGS_protocol == "hip") {
             xport = engine->installTransport("hip", nullptr);
         } else {
@@ -423,7 +425,7 @@ int initiator() {
 
     std::vector<void *> addr;
 #if defined(USE_CUDA) || defined(USE_MUSA) || defined(USE_HIP) || \
-    defined(USE_MACA)
+    defined(USE_MACA) || defined(USE_HYGON) || defined(USE_COREX)
     if (FLAGS_use_vram) {
         int gpu_num;
         LOG(INFO) << "VRAM is used";
@@ -536,6 +538,8 @@ int target() {
             engine->installTransport("tcp", nullptr);
         } else if (FLAGS_protocol == "nvlink") {
             engine->installTransport("nvlink", nullptr);
+        } else if (FLAGS_protocol == "maca") {
+            engine->installTransport("maca", nullptr);
         } else if (FLAGS_protocol == "hip") {
             engine->installTransport("hip", nullptr);
         } else {
@@ -545,7 +549,7 @@ int target() {
 
     std::vector<void *> addr;
 #if defined(USE_CUDA) || defined(USE_MUSA) || defined(USE_HIP) || \
-    defined(USE_MACA)
+    defined(USE_MACA) || defined(USE_HYGON) || defined(USE_COREX)
     if (FLAGS_use_vram) {
         int gpu_num;
         LOG(INFO) << "VRAM is used";
diff --git a/mooncake-transfer-engine/fabric_allocator.cmake b/mooncake-transfer-engine/fabric_allocator.cmake
new file mode 100644
index 00000000..62ee1095
--- /dev/null
+++ b/mooncake-transfer-engine/fabric_allocator.cmake
@@ -0,0 +1,40 @@
+function(add_fabric_allocator_build_target)
+  set(options)
+  set(oneValueArgs TARGET_NAME BUILD_SCRIPT COMMENT ENABLE_BUILD)
+  set(multiValueArgs BUILD_ARGS)
+  cmake_parse_arguments(FAB "${options}" "${oneValueArgs}" "${multiValueArgs}"
+                        ${ARGN})
+
+  if(NOT FAB_TARGET_NAME)
+    message(
+      FATAL_ERROR
+        "TARGET_NAME is required for add_fabric_allocator_build_target")
+  endif()
+  if(NOT FAB_BUILD_SCRIPT)
+    message(
+      FATAL_ERROR
+        "BUILD_SCRIPT is required for add_fabric_allocator_build_target")
+  endif()
+
+  add_custom_target(${FAB_TARGET_NAME} DEPENDS transfer_engine)
+
+  get_target_property(_include_dirs ${FAB_TARGET_NAME} INCLUDE_DIRECTORIES)
+  if(NOT _include_dirs)
+    set(_include_dirs "")
+  endif()
+  string(REPLACE ";" " " _include_dirs_str "${_include_dirs}")
+
+  if(FAB_ENABLE_BUILD)
+    add_custom_command(
+      TARGET ${FAB_TARGET_NAME}
+      POST_BUILD
+      COMMAND ${CMAKE_COMMAND} -E make_directory ${CMAKE_CURRENT_BINARY_DIR}
+      COMMAND bash ${FAB_BUILD_SCRIPT} ${FAB_BUILD_ARGS}
+              ${CMAKE_CURRENT_BINARY_DIR} "${_include_dirs_str}"
+      WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}
+      COMMENT "${FAB_COMMENT}"
+      VERBATIM)
+  endif()
+
+  set_property(TARGET ${FAB_TARGET_NAME} PROPERTY EXCLUDE_FROM_ALL FALSE)
+endfunction()
diff --git a/mooncake-transfer-engine/include/CMakeLists.txt b/mooncake-transfer-engine/include/CMakeLists.txt
index 4f52e1e5..c9486b9b 100644
--- a/mooncake-transfer-engine/include/CMakeLists.txt
+++ b/mooncake-transfer-engine/include/CMakeLists.txt
@@ -7,5 +7,19 @@ install(FILES multi_transport.h DESTINATION include)
 install(FILES topology.h DESTINATION include)
 install(FILES transfer_engine.h DESTINATION include)
 install(FILES transfer_metadata.h DESTINATION include)
+install(FILES ub_allocator.h DESTINATION include)
 install(FILES common/base/status.h DESTINATION include/common/base)
 install(FILES transport/transport.h DESTINATION include/transport)
+
+# Device API headers (header-only, consumed by EP kernel)
+install(FILES transport/device/device_transport.h DESTINATION include/transport/device)
+install(FILES transport/device/device_ops.cuh DESTINATION include/transport/device)
+install(FILES transport/device/comm_device.cuh DESTINATION include/transport/device)
+install(FILES transport/device/p2p_device.cuh DESTINATION include/transport/device)
+install(FILES transport/device/ibgda_device.cuh DESTINATION include/transport/device)
+install(FILES transport/device/cuda/cuda_ops.cuh DESTINATION include/transport/device/cuda)
+install(FILES transport/device/musa/musa_ops.cuh DESTINATION include/transport/device/musa)
+
+# IBGDA library headers
+install(DIRECTORY transport/device/ibgda/ DESTINATION include/transport/device/ibgda
+        FILES_MATCHING PATTERN "*.h")
diff --git a/mooncake-transfer-engine/include/char_util.h b/mooncake-transfer-engine/include/char_util.h
new file mode 100644
index 00000000..ab5f0ee5
--- /dev/null
+++ b/mooncake-transfer-engine/include/char_util.h
@@ -0,0 +1,31 @@
+// Copyright 2024 KVCache.AI
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#ifndef CHAR_UTIL_H
+#define CHAR_UTIL_H
+
+#include <cctype>
+
+namespace mooncake {
+
+// Lowercase a single byte safely. Passing a (possibly signed) char straight
+// to std::tolower is UB when the byte is > 0x7F; the argument must be
+// representable as unsigned char or equal EOF.
+static inline char to_lower(char c) {
+    return static_cast<char>(std::tolower(static_cast<unsigned char>(c)));
+}
+
+}  // namespace mooncake
+
+#endif  // CHAR_UTIL_H
diff --git a/mooncake-transfer-engine/include/common.h b/mooncake-transfer-engine/include/common.h
index f98ba147..5ae9a664 100644
--- a/mooncake-transfer-engine/include/common.h
+++ b/mooncake-transfer-engine/include/common.h
@@ -267,6 +267,33 @@ static inline std::pair<std::string, uint16_t> parseHostNameWithPort(
             getPortFromString(server_name.substr(colon_pos + 1), port)};
 }
 
+static inline bool hasExplicitPort(const std::string &server_name) {
+    auto result = extractIPv6HostAndPort(server_name);
+    if (result.matched) {
+        return !result.port.empty();
+    }
+    return server_name.rfind(':') != std::string::npos;
+}
+
+static inline std::string getHostNameWithoutPort(
+    const std::string &server_name) {
+    auto result = extractIPv6HostAndPort(server_name);
+    if (result.matched) {
+        return std::move(result.host);
+    }
+
+    const size_t colon_pos = server_name.rfind(':');
+    if (colon_pos == std::string::npos) {
+        return server_name;
+    }
+    return server_name.substr(0, colon_pos);
+}
+
+static inline std::string buildHostNameWithPort(const std::string &host,
+                                                uint16_t port) {
+    return maybeWrapIpV6(host) + ":" + std::to_string(port);
+}
+
 static inline uint16_t parsePortAndDevice(std::string_view suffix,
                                           uint16_t default_port,
                                           int *device_id) {
@@ -431,6 +458,11 @@ static inline std::pair<HandShakeRequestType, std::string> readString(int fd) {
         return {type, ""};
     }
 
+    if (length == 0) {
+        LOG(ERROR) << "readString: zero length from socket";
+        return {type, ""};
+    }
+
     std::string str;
     std::vector<char> buffer(length);
     n = readFully(fd, buffer.data(), length);
diff --git a/mooncake-transfer-engine/include/config.h b/mooncake-transfer-engine/include/config.h
index bf632ad3..b363c3c3 100644
--- a/mooncake-transfer-engine/include/config.h
+++ b/mooncake-transfer-engine/include/config.h
@@ -22,6 +22,9 @@
 #include <cstdint>
 #include <mutex>
 #include <string>
+#include <unordered_map>
+#include <utility>
+#include <vector>
 
 namespace mooncake {
 
@@ -34,7 +37,8 @@ struct GlobalConfig {
     size_t num_cq_per_ctx = 1;
     size_t num_comp_channels_per_ctx = 1;
     uint8_t port = 1;
-    int gid_index = -1;  // -1 for auto-selection, >=0 for user-specified
+    int gid_index = -1;       // -1 for auto-selection, >=0 for user-specified
+    uint16_t pkey_index = 0;  // QP attr.pkey_index; override via MC_PKEY_INDEX
     uint64_t max_mr_size = 0x10000000000;
     size_t max_cqe = 4096;
     int max_ep_per_ctx = 65536;
@@ -47,7 +51,14 @@ struct GlobalConfig {
     int workers_per_ctx = 2;
     size_t slice_size = 65536;
     int retry_cnt = 9;
+    int auto_gid_max_retries = 2;
     int handshake_listen_backlog = 128;
+    // Connect timeout (seconds) for outbound handshake-port RPCs (QP
+    // handshake, probe, notify, metadata exchange). A plain blocking
+    // connect() has no deadline: to an unroutable address (e.g. a
+    // torn-down pod IP) it stalls for the kernel's full SYN-retry cycle,
+    // which is minutes. Override via MC_HANDSHAKE_CONNECT_TIMEOUT.
+    int handshake_connect_timeout = 5;
     bool metacache = true;
     int log_level = google::INFO;
     bool trace = false;
@@ -57,15 +68,38 @@ struct GlobalConfig {
     bool use_ipv6 = false;
     size_t fragment_limit = 16384;
     bool enable_dest_device_affinity = false;
+    bool enable_hca_peer_affinity = false;
+    std::unordered_map<std::string, std::vector<std::string>> nic_peer_affinity;
+    bool log_rdma_slice_affinity = false;
     int parallel_reg_mr = -1;
     size_t eic_max_block_size = 64UL * 1024 * 1024;
     EndpointStoreType endpoint_store_type = EndpointStoreType::SIEVE;
     int ib_traffic_class = -1;
+    // InfiniBand Service Level (SL), 0-15. -1 = use default (0).
+    // Maps to a Virtual Lane on the switch for QoS isolation, e.g. to
+    // steer KV-cache traffic into a different VL than EP all-to-all.
+    int ib_service_level = -1;
+    // mlx5 QP UDP source ports for ECMP path diversification.
+    // Empty = no modification. QP at index i uses
+    // mlx5_qp_udp_sports[i % size]. Requires mlx5 device + RoCEv2,
+    // and the binary must be built with USE_MLX5DV.
+    std::vector<uint16_t> mlx5_qp_udp_sports;
+    // mlx5 QP LAG port balancing. When enabled, QPs are distributed across
+    // physical LAG ports: QP at index i is pinned to port (i % num_lag_ports)
+    // + 1. num_lag_ports is queried from hardware; if the device is not in LAG
+    // mode the setting is a no-op. Requires USE_MLX5DV.
+    bool mlx5_qp_lag_port_balance = false;
     // ib_pci_relaxed_ordering_mode: 0: off, 1: on if supported, 2: auto
     int ib_pci_relaxed_ordering_mode = 0;
     bool ascend_use_fabric_mem = false;
     bool ascend_agent_mode = false;
-    size_t efa_striping_threshold = 2 * 1024 * 1024;  // 2MB default
+    // Transient flag scoped to a single TE init: set true by the Store entry
+    // (Client::InitTransferEngine) before installing the ascend transport, and
+    // reset to false right after. Lets ascend_direct distinguish a Store-init
+    // TE from a normal/P2P TE so each can resolve its own
+    // ASCEND_GLOBAL_RESOURCE_CONFIG (e.g. Store=RoCE, P2P=HCCS). Assumes TE
+    // inits are serialized within the process.
+    bool ascend_store_te_init = false;
     // ub config parameters
     size_t num_jfc_per_ctx = 2;
     size_t num_jfce_per_ctx = 2;
@@ -92,6 +126,11 @@ GlobalConfig& globalConfig();
 
 uint16_t getDefaultHandshakePort();
 
+// Validates a port range. Returns {default_min, default_max} on invalid input.
+// Rejects: min > max, well-known ports (0-1023), ephemeral ports (32768-60999).
+std::pair<int, int> ValidatePortRange(int min_port, int max_port,
+                                      int default_min, int default_max);
+
 }  // namespace mooncake
 
 #endif  // CONFIG_H
diff --git a/mooncake-transfer-engine/include/cuda_alike.h b/mooncake-transfer-engine/include/cuda_alike.h
index 39f62897..2ec35f41 100644
--- a/mooncake-transfer-engine/include/cuda_alike.h
+++ b/mooncake-transfer-engine/include/cuda_alike.h
@@ -13,10 +13,18 @@
 #include "gpu_vendor/ubshmem.h"
 #elif defined(USE_MACA)
 #include "gpu_vendor/maca.h"
+#elif defined(USE_SUNRISE)
+#include "gpu_vendor/sunrise.h"
+#elif defined(USE_HYGON)
+#include <cuda.h>
+#include <cuda_runtime.h>
+#elif defined(USE_COREX)
+#include <cuda.h>
+#include <cuda_runtime.h>
 #endif
 
 #if !defined(USE_HIP) && !defined(USE_MUSA) && !defined(USE_MLU) && \
-    !defined(USE_UBSHMEM) && !defined(USE_MACA)
+    !defined(USE_UBSHMEM) && !defined(USE_MACA) && !defined(USE_SUNRISE)
 #include <string>
 const static std::string GPU_PREFIX = "cuda:";
 #endif
diff --git a/mooncake-transfer-engine/include/gpu_vendor/hip.h b/mooncake-transfer-engine/include/gpu_vendor/hip.h
index b613c87a..6797fce6 100644
--- a/mooncake-transfer-engine/include/gpu_vendor/hip.h
+++ b/mooncake-transfer-engine/include/gpu_vendor/hip.h
@@ -1,3 +1,5 @@
+#pragma once
+
 #include <hip/hip_runtime.h>
 #include <string>
 
diff --git a/mooncake-transfer-engine/include/gpu_vendor/maca.h b/mooncake-transfer-engine/include/gpu_vendor/maca.h
index df470dbb..56384149 100644
--- a/mooncake-transfer-engine/include/gpu_vendor/maca.h
+++ b/mooncake-transfer-engine/include/gpu_vendor/maca.h
@@ -8,13 +8,16 @@
 // First-stage MACA integration: provide a CUDA-like API surface.
 const static std::string GPU_PREFIX = "maca:";
 
+// torch-maca's cu-bridge already defines CUDA/CU compatibility macros after
+// ATen/cuda headers are included. Keep this fallback mapping for translation
+// units that include Mooncake's cuda_alike.h without cu-bridge first.
+#ifndef __CUDA_TO_MACA_ADAPTOR_H__
 #define CUdevice MCdevice
-#define CUdeviceptr MCdeviceptr
+#define CUdeviceptr mcDeviceptr_t
 #define CUmemorytype MCmemorytype
 #define CUresult mcError_t
 #define cuDeviceGet mcDeviceGet
 #define cuDeviceGetAttribute mcDeviceGetAttribute
-#define cuGetErrorString mcGetErrorString
 #define cuMemAddressFree mcMemAddressFree
 #define cuMemAddressReserve mcMemAddressReserve
 #define cuMemCreate mcMemCreate
@@ -34,18 +37,73 @@ const static std::string GPU_PREFIX = "maca:";
 #define CUDA_ERROR_NOT_PERMITTED mcErrorNotPermitted
 #define CUDA_ERROR_NOT_SUPPORTED mcErrorNotSupported
 
+#define CUmemFabricHandle mcMemFabricHandle_t
+#define CUmemGenericAllocationHandle mcMemGenericAllocationHandle
+#define CUmemAllocationProp mcMemAllocationProp
+#define CUmemAllocationHandleType mcMemAllocationHandleType
+#define CUmemAccessDesc mcMemAccessDesc
+#define CUmemAllocationType mcMemAllocationType
+#define CU_MEM_ALLOCATION_TYPE_PINNED mcMemAllocationTypePinned
+#define CU_MEM_LOCATION_TYPE_DEVICE mcMemLocationTypeDevice
+#define CU_MEM_HANDLE_TYPE_FABRIC mcMemHandleTypeFabric
+#define CU_MEM_HANDLE_TYPE_POSIX_FILE_DESCRIPTOR \
+    mcMemHandleTypePosixFileDescriptor
+#define CU_DEVICE_ATTRIBUTE_HANDLE_TYPE_FABRIC_SUPPORTED \
+    mcDeviceAttributeHandleTypeFabricSupported
+#define CU_DEVICE_ATTRIBUTE_GPU_DIRECT_RDMA_WITH_CUDA_VMM_SUPPORTED \
+    mcDeviceAttributeHandleTypePosixFileDescriptorSupported
+#define CU_MEM_ACCESS_FLAGS_PROT_READWRITE mcMemAccessFlagsProtReadWrite
+#define CU_MEM_ALLOC_GRANULARITY_MINIMUM MC_MEM_ALLOC_GRANULARITY_MINIMUM
+
+#define CU_MEMORYTYPE_HOST mcMemoryTypeHost
+#define CU_MEMORYTYPE_DEVICE mcMemoryTypeDevice
+#define CU_POINTER_ATTRIBUTE_MEMORY_TYPE mcPointerAttributeMemoryType
+#define CU_POINTER_ATTRIBUTE_DEVICE_ORDINAL mcPointerAttributeDevice
+#define CU_POINTER_ATTRIBUTE_RANGE_START_ADDR mcPointerAttributeRangeStartAddr
+#define CU_POINTER_ATTRIBUTE_RANGE_SIZE mcPointerAttributeRangeSize
+#define CU_MEM_RANGE_HANDLE_TYPE_DMA_BUF_FD mcMemHandleTypePosixFileDescriptor
+#define CU_DEVICE_ATTRIBUTE_DMA_BUF_SUPPORTED \
+    mcDeviceAttributeHandleTypePosixFileDescriptorSupported
+
+static inline CUresult cuGetErrorString(CUresult error, const char **err_str) {
+    if (err_str) {
+        *err_str = mcGetErrorString(error);
+    }
+    return CUDA_SUCCESS;
+}
+
 #define cudaDeviceCanAccessPeer mcDeviceCanAccessPeer
 #define cudaDeviceEnablePeerAccess mcDeviceEnablePeerAccess
+#define cudaDeviceGetAttribute mcDeviceGetAttribute
 #define cudaDeviceGetPCIBusId mcDeviceGetPCIBusId
+#define cudaDeviceSynchronize mcDeviceSynchronize
+#define cudaDevAttrMultiProcessorCount mcDeviceAttributeMultiProcessorCount
+#define cudaErrorNotReady mcErrorNotReady
 #define cudaErrorPeerAccessAlreadyEnabled mcErrorPeerAccessAlreadyEnabled
 #define cudaError_t mcError_t
+#define cudaEvent_t mcEvent_t
+#define cudaEventCreate mcEventCreate
+#define cudaEventCreateWithFlags mcEventCreateWithFlags
+#define cudaEventDestroy mcEventDestroy
+#define cudaEventDisableTiming mcEventDisableTiming
+#define cudaEventElapsedTime mcEventElapsedTime
+#define cudaEventQuery mcEventQuery
+#define cudaEventRecord mcEventRecord
+#define cudaEventSynchronize mcEventSynchronize
 #define cudaFree mcFree
 #define cudaFreeHost mcFreeHost
+#define cudaFuncAttributes mcFuncAttributes
 #define cudaGetDevice mcGetDevice
 #define cudaGetDeviceCount mcGetDeviceCount
 #define cudaGetErrorString mcGetErrorString
 #define cudaGetLastError mcGetLastError
+#define cudaHostAlloc mcMallocHost
+#define cudaHostAllocDefault mcMallocHostDefault
+#define cudaHostAllocMapped mcMallocHostMapped
+#define cudaHostAllocPortable mcMallocHostPortable
+#define cudaHostAllocWriteCombined mcMallocHostWriteCombined
 #define cudaHostRegister mcHostRegister
+#define cudaHostRegisterMapped mcHostRegisterMapped
 #define cudaHostRegisterPortable mcHostRegisterPortable
 #define cudaHostUnregister mcHostUnregister
 #define cudaIpcCloseMemHandle mcIpcCloseMemHandle
@@ -58,8 +116,10 @@ const static std::string GPU_PREFIX = "maca:";
 #define cudaMemcpy mcMemcpy
 #define cudaMemcpyAsync mcMemcpyAsync
 #define cudaMemcpyDefault mcMemcpyDefault
+#define cudaMemcpyDeviceToDevice mcMemcpyDeviceToDevice
 #define cudaMemcpyDeviceToHost mcMemcpyDeviceToHost
 #define cudaMemcpyHostToDevice mcMemcpyHostToDevice
+#define cudaMemcpyKind mcMemcpyKind
 #define cudaMemset mcMemset
 #define cudaMemsetAsync mcMemsetAsync
 #define cudaMemoryTypeDevice mcMemoryTypeDevice
@@ -69,7 +129,37 @@ const static std::string GPU_PREFIX = "maca:";
 #define cudaPointerGetAttributes mcPointerGetAttributes
 #define cudaSetDevice mcSetDevice
 #define cudaStreamCreate mcStreamCreate
+#define cudaStreamCreateWithFlags mcStreamCreateWithFlags
 #define cudaStreamDestroy mcStreamDestroy
+#define cudaStreamNonBlocking mcStreamNonBlocking
+#define cudaStreamPerThread mcStreamPerThread
+#define cudaStreamQuery mcStreamQuery
 #define cudaStreamSynchronize mcStreamSynchronize
 #define cudaStream_t mcStream_t
+#define cudaStreamWaitEvent mcStreamWaitEvent
 #define cudaSuccess mcSuccess
+
+static inline mcError_t cudaFuncGetAttributes(mcFuncAttributes *attr,
+                                              const void *func) {
+    return mcFuncGetAttributes(attr, func);
+}
+
+#ifdef __cplusplus
+template <class T>
+static inline mcError_t cudaFuncGetAttributes(mcFuncAttributes *attr, T func) {
+    return mcFuncGetAttributes(attr, reinterpret_cast<const void *>(func));
+}
+
+template <class T>
+static inline mcError_t cudaHostGetDevicePointer(T **dev_ptr, void *host_ptr,
+                                                 unsigned int flags) {
+    return mcHostGetDevicePointer(reinterpret_cast<void **>(dev_ptr), host_ptr,
+                                  flags);
+}
+#else
+static inline mcError_t cudaHostGetDevicePointer(void **dev_ptr, void *host_ptr,
+                                                 unsigned int flags) {
+    return mcHostGetDevicePointer(dev_ptr, host_ptr, flags);
+}
+#endif
+#endif  // __CUDA_TO_MACA_ADAPTOR_H__
diff --git a/mooncake-transfer-engine/include/gpu_vendor/mlu.h b/mooncake-transfer-engine/include/gpu_vendor/mlu.h
index d2daf500..249594bf 100644
--- a/mooncake-transfer-engine/include/gpu_vendor/mlu.h
+++ b/mooncake-transfer-engine/include/gpu_vendor/mlu.h
@@ -23,7 +23,6 @@ using CUresult = CNresult;
 #define CU_MEMORYTYPE_HOST CN_MEMORYTYPE_HOST
 #define CU_MEMORYTYPE_DEVICE CN_MEMORYTYPE_DEVICE
 #define CU_POINTER_ATTRIBUTE_MEMORY_TYPE CN_MEM_ATTRIBUTE_TYPE
-#define CU_POINTER_ATTRIBUTE_RANGE_SIZE CN_MEM_ATTRIBUTE_RANGE_SIZE
 
 #define cudaSuccess cnrtSuccess
 #define cudaMemoryTypeUnregistered cnrtMemTypeUnregistered
@@ -158,18 +157,21 @@ static inline CUresult cuPointerGetAttribute(void *data, int attribute,
             }
             return CUDA_SUCCESS;
         }
-        case CU_POINTER_ATTRIBUTE_RANGE_SIZE: {
-            if (attributes.type != cudaMemoryTypeDevice) {
-                return CN_ERROR_INVALID_VALUE;
-            }
-            *static_cast<size_t *>(data) = attributes.size;
-            return CUDA_SUCCESS;
-        }
         default:
             return CN_ERROR_INVALID_VALUE;
     }
 }
 
+static inline CUresult cuMemGetAddressRange(CUdeviceptr *base, size_t *size,
+                                            CUdeviceptr ptr) {
+    cn_uint64_t bytes = 0;
+    auto ret = cnMemGetAddressRange(base, size ? &bytes : nullptr, ptr);
+    if (ret == CN_SUCCESS && size) {
+        *size = static_cast<size_t>(bytes);
+    }
+    return ret;
+}
+
 static inline CUresult cuMemGetHandleForAddressRange(void *handle,
                                                      CUdeviceptr ptr,
                                                      size_t size,
diff --git a/mooncake-transfer-engine/include/gpu_vendor/musa.h b/mooncake-transfer-engine/include/gpu_vendor/musa.h
index 45d4428d..58f2d23f 100644
--- a/mooncake-transfer-engine/include/gpu_vendor/musa.h
+++ b/mooncake-transfer-engine/include/gpu_vendor/musa.h
@@ -1,3 +1,5 @@
+#pragma once
+
 #include <string>
 #include <musa.h>
 #include <musa_runtime.h>
@@ -20,6 +22,11 @@ const static std::string GPU_PREFIX = "musa:";
 #define CU_POINTER_ATTRIBUTE_RANGE_SIZE MU_POINTER_ATTRIBUTE_RANGE_SIZE
 #define CU_DEVICE_ATTRIBUTE_MEMORY_POOLS_SUPPORTED \
     MU_DEVICE_ATTRIBUTE_MEMORY_POOLS_SUPPORTED
+#define CU_DEVICE_ATTRIBUTE_DMA_BUF_SUPPORTED \
+    MU_DEVICE_ATTRIBUTE_DMA_BUF_SUPPORTED
+#define CU_POINTER_ATTRIBUTE_DEVICE_ORDINAL MU_POINTER_ATTRIBUTE_DEVICE_ORDINAL
+
+#define CUcontext MUcontext
 
 #define CUdevice MUdevice
 #define CUdeviceptr MUdeviceptr
@@ -86,7 +93,37 @@ const static std::string GPU_PREFIX = "musa:";
 #define cudaPointerGetAttributes musaPointerGetAttributes
 #define cudaSetDevice musaSetDevice
 #define cudaStreamCreate musaStreamCreate
+#define cudaStreamCreateWithFlags musaStreamCreateWithFlags
+#define cudaStreamNonBlocking musaStreamNonBlocking
 #define cudaStreamDestroy musaStreamDestroy
+#define cudaStreamPerThread musaStreamPerThread
+#define cudaStreamQuery musaStreamQuery
+#define cudaStreamWaitEvent musaStreamWaitEvent
+#define cudaDeviceSynchronize musaDeviceSynchronize
 #define cudaStreamSynchronize musaStreamSynchronize
 #define cudaStream_t musaStream_t
 #define cudaSuccess musaSuccess
+#define cudaErrorNotReady musaErrorNotReady
+#define cudaDeviceGetAttribute musaDeviceGetAttribute
+#define cudaEvent_t musaEvent_t
+#define cudaEventCreateWithFlags musaEventCreateWithFlags
+#define cudaEventDisableTiming musaEventDisableTiming
+#define cudaEventDestroy musaEventDestroy
+#define cudaEventRecord musaEventRecord
+#define cudaMemcpyDeviceToDevice musaMemcpyDeviceToDevice
+#define cudaDevAttrClockRate musaDevAttrClockRate
+#define cudaLaunchConfig_t musaLaunchConfig_t
+#define cudaLaunchAttribute musaLaunchAttribute
+#define cudaLaunchAttributeCooperative musaLaunchAttributeCooperative
+#define cudaLaunchKernelEx musaLaunchKernelEx
+#define CUDA_R_16BF MUSA_R_16BF
+#define CUDA_R_32F MUSA_R_32F
+
+// IBGDA-specific mappings
+#define cuInit muInit
+#define cuDevicePrimaryCtxRetain muDevicePrimaryCtxRetain
+#define cuDevicePrimaryCtxRelease muDevicePrimaryCtxRelease
+#define cuCtxSetCurrent muCtxSetCurrent
+#define cudaHostRegisterMapped musaHostRegisterMapped
+#define cudaHostRegisterIoMemory musaHostRegisterIoMemory
+#define cudaHostGetDevicePointer musaHostGetDevicePointer
diff --git a/mooncake-transfer-engine/include/gpu_vendor/sunrise.h b/mooncake-transfer-engine/include/gpu_vendor/sunrise.h
new file mode 100644
index 00000000..5ff525b9
--- /dev/null
+++ b/mooncake-transfer-engine/include/gpu_vendor/sunrise.h
@@ -0,0 +1,32 @@
+#pragma once
+
+#include <string>
+
+#include <tang_runtime_api.h>
+
+// SunriseLink uses CUDA-style location naming in metadata.
+const static std::string GPU_PREFIX = "cuda:";
+
+// Minimal CUDA-like API shim for benchmark paths.
+#define cudaError_t tangError_t
+#define cudaSuccess tangSuccess
+#define cudaGetErrorString tangGetErrorString
+
+#define cudaSetDevice tangSetDevice
+#define cudaGetDeviceCount tangGetDeviceCount
+#define cudaDeviceGetPCIBusId tangDeviceGetPCIBusId
+
+#define cudaMalloc tangMalloc
+#define cudaFree tangFree
+#define cudaMemset tangMemset
+#define cudaMemcpy tangMemcpy
+#define cudaMemcpyHostToDevice tangMemcpyHostToDevice
+#define cudaMemcpyDeviceToHost tangMemcpyDeviceToHost
+
+#define cudaPointerAttributes tangPointerAttributes
+#define cudaPointerGetAttributes tangPointerGetAttributes
+#define cudaMemoryTypeDevice tangMemoryTypeDevice
+#define cudaMemoryTypeHost tangMemoryTypeHost
+#define cudaMemoryTypeUnregistered tangMemoryTypeUnregistered
+
+#define cudaStreamSynchronize tangStreamSynchronize
diff --git a/mooncake-transfer-engine/include/gpu_vendor/ubshmem.h b/mooncake-transfer-engine/include/gpu_vendor/ubshmem.h
index 7be526a7..584680e6 100644
--- a/mooncake-transfer-engine/include/gpu_vendor/ubshmem.h
+++ b/mooncake-transfer-engine/include/gpu_vendor/ubshmem.h
@@ -1,3 +1,5 @@
+#pragma once
+
 #include <string>
 #include <acl/acl.h>
 
diff --git a/mooncake-transfer-engine/include/hip_device_guard.h b/mooncake-transfer-engine/include/hip_device_guard.h
new file mode 100644
index 00000000..49b6b506
--- /dev/null
+++ b/mooncake-transfer-engine/include/hip_device_guard.h
@@ -0,0 +1,52 @@
+// Copyright 2025 Mooncake Authors
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#pragma once
+
+#if defined(USE_HIP) || defined(USE_HIP_DMABUF)
+
+#include <hip/hip_runtime.h>
+
+namespace mooncake {
+
+// RAII: saves the active HIP device and restores it on scope exit, so a
+// function that calls hipSetDevice() doesn't leave the caller on the wrong GPU
+// (which would fail its next kernel launch with hipErrorInvalidDevice).
+class HipDeviceGuard {
+   public:
+    HipDeviceGuard() { saved_ = (hipGetDevice(&prev_device_) == hipSuccess); }
+
+    // Also switch to target_device; set_ok() reports whether that succeeded.
+    explicit HipDeviceGuard(int target_device) : HipDeviceGuard() {
+        set_ok_ = (hipSetDevice(target_device) == hipSuccess);
+    }
+
+    ~HipDeviceGuard() {
+        if (saved_) (void)hipSetDevice(prev_device_);
+    }
+
+    bool set_ok() const { return set_ok_; }
+
+    HipDeviceGuard(const HipDeviceGuard&) = delete;
+    HipDeviceGuard& operator=(const HipDeviceGuard&) = delete;
+
+   private:
+    int prev_device_ = 0;
+    bool saved_ = false;
+    bool set_ok_ = true;
+};
+
+}  // namespace mooncake
+
+#endif  // USE_HIP || USE_HIP_DMABUF
diff --git a/mooncake-transfer-engine/include/topology.h b/mooncake-transfer-engine/include/topology.h
index 123c1749..e9d16605 100644
--- a/mooncake-transfer-engine/include/topology.h
+++ b/mooncake-transfer-engine/include/topology.h
@@ -30,6 +30,7 @@
 #include <string>
 #include <thread>
 #include <unordered_map>
+#include <vector>
 
 #include "common.h"
 
@@ -86,6 +87,8 @@ class Topology {
     int selectDevice(const std::string storage_type, int retry_count = 0);
     int selectDevice(const std::string storage_type, std::string_view hint,
                      int retry_count = 0);
+    int selectDeviceByLocalHca(const std::string storage_type,
+                               std::string_view local_hca, int retry_count = 0);
 
     TopologyMatrix getMatrix() const { return matrix_; }
 
@@ -121,8 +124,12 @@ class Topology {
     };
     std::unordered_map<std::string /* storage type */, ResolvedTopologyEntry>
         resolved_matrix_;
+    std::unordered_map<std::string /* local HCA */,
+                       std::unordered_map<std::string /* storage type */,
+                                          std::vector<int> /* peer HCA ids */>>
+        resolved_hca_peer_affinity_by_local_;
 };
 
 }  // namespace mooncake
 
-#endif  // TOPOLOGY_H
\ No newline at end of file
+#endif  // TOPOLOGY_H
diff --git a/mooncake-transfer-engine/include/transfer_engine.h b/mooncake-transfer-engine/include/transfer_engine.h
index 65bc8741..4f3e1d73 100644
--- a/mooncake-transfer-engine/include/transfer_engine.h
+++ b/mooncake-transfer-engine/include/transfer_engine.h
@@ -25,6 +25,12 @@ class TransferEngineImpl;
 namespace tent {
 class TransferEngine;
 };
+#if defined(USE_CUDA) || defined(USE_MUSA)
+namespace device {
+class P2pTransport;
+class RdmaTransport;
+}  // namespace device
+#endif
 using TransferRequest = Transport::TransferRequest;
 using TransferStatus = Transport::TransferStatus;
 using TransferStatusEnum = Transport::TransferStatusEnum;
@@ -65,6 +71,10 @@ class TransferEngine {
 
     TransferEngine(bool auto_discover, const std::vector<std::string>& filter);
 
+    TransferEngine(TransferEngine&&) = default;
+
+    TransferEngine& operator=(TransferEngine&&) = default;
+
     ~TransferEngine();
 
     int init(const std::string& metadata_conn_string,
@@ -82,6 +92,8 @@ class TransferEngine {
 
     int getRpcPort();
 
+    bool isUsingTent() const { return use_tent_; }
+
     SegmentHandle openSegment(const std::string& segment_name);
 
     Status CheckSegmentStatus(SegmentID sid);
@@ -150,6 +162,16 @@ class TransferEngine {
 
     Transport* getTransport(const std::string& proto);
 
+#if defined(USE_CUDA) || defined(USE_MUSA)
+    // Device transport accessors (P2P + IBGDA).  Lazily created on first
+    // call and owned by the TransferEngine.  These allow EP (and future
+    // CPU-proxy paths) to obtain device transports from an engine instance
+    // instead of calling the global factory functions directly.
+    device::P2pTransport* getOrCreateP2pTransport(int num_ranks);
+    device::RdmaTransport* getOrCreateRdmaTransport(
+        const std::vector<std::string>& device_filter = {});
+#endif
+
     /**
      * @brief Check if TCP is the only installed transport.
      *
diff --git a/mooncake-transfer-engine/include/transfer_engine_impl.h b/mooncake-transfer-engine/include/transfer_engine_impl.h
index 4016d822..b1e4fff7 100644
--- a/mooncake-transfer-engine/include/transfer_engine_impl.h
+++ b/mooncake-transfer-engine/include/transfer_engine_impl.h
@@ -21,6 +21,7 @@
 #include <atomic>
 #include <cstddef>
 #include <cstdint>
+#include <map>
 #include <memory>
 #include <shared_mutex>
 #include <string>
@@ -32,6 +33,9 @@
 #include "transfer_metadata.h"
 #include "transfer_engine.h"
 #include "transport/transport.h"
+#if defined(USE_CUDA) || defined(USE_MUSA)
+#include "transport/device/device_transport.h"
+#endif
 #ifdef WITH_METRICS
 #include "ylt/metric/counter.hpp"
 #include "ylt/metric/histogram.hpp"
@@ -132,6 +136,9 @@ class TransferEngineImpl {
     Status submitTransferWithNotify(BatchID batch_id,
                                     const std::vector<TransferRequest>& entries,
                                     TransferMetadata::NotifyDesc notify_msg) {
+        if (entries.empty()) {
+            return Status::InvalidArgument("entries must not be empty");
+        }
         auto target_id = entries[0].target_id;
         Status s = multi_transports_->submitTransfer(batch_id, entries);
         if (!s.ok()) {
@@ -190,6 +197,9 @@ class TransferEngineImpl {
     Status mp_submitTransferWithNotify(
         BatchID batch_id, const std::vector<TransferRequest>& entries,
         TransferMetadata::NotifyDesc notify_msg, std::string& proto) {
+        if (entries.empty()) {
+            return Status::InvalidArgument("entries must not be empty");
+        }
         auto target_id = entries[0].target_id;
         Status s =
             multi_transports_->mp_submitTransfer(batch_id, entries, proto);
@@ -334,6 +344,13 @@ class TransferEngineImpl {
         return multi_transports_->getTransport(proto);
     }
 
+#if defined(USE_CUDA) || defined(USE_MUSA)
+    // Device transport accessors — lazily created, owned by this impl.
+    device::P2pTransport* getOrCreateP2pTransport(int num_ranks);
+    device::RdmaTransport* getOrCreateRdmaTransport(
+        const std::vector<std::string>& device_filter = {});
+#endif
+
     bool isTcpOnly() const { return multi_transports_->isTcpOnly(); }
 
     int syncSegmentCache(const std::string& segment_name = "") {
@@ -379,11 +396,24 @@ class TransferEngineImpl {
         bool remote_accessible;
     };
 
+    using MemoryRegionMap = std::map<uintptr_t, MemoryRegion>;
+
+    MemoryRegionMap::iterator findMemoryRegionContaining(uintptr_t addr);
+
+    MemoryRegionMap::const_iterator findMemoryRegionContaining(
+        uintptr_t addr) const;
+
+    bool hasOverlapLocked(uintptr_t addr, uint64_t length) const;
+
+    void insertMemoryRegionLocked(const MemoryRegion& region);
+
+    void eraseMemoryRegionLocked(void* addr);
+
     std::shared_ptr<TransferMetadata> metadata_;
     std::string local_server_name_;
     std::shared_ptr<MultiTransport> multi_transports_;
     std::shared_mutex mutex_;
-    std::vector<MemoryRegion> local_memory_regions_;
+    MemoryRegionMap local_memory_regions_;
     std::shared_ptr<Topology> local_topology_;
 
     RWSpinlock send_notifies_lock_;
@@ -397,6 +427,13 @@ class TransferEngineImpl {
     std::vector<std::string> filter_;
     bool use_barex_ = false;
 
+#if defined(USE_CUDA) || defined(USE_MUSA)
+    // Device transports (P2P + IBGDA) — lazily created, owned by this impl.
+    // Referenced by EP and future CPU-proxy paths.
+    std::unique_ptr<device::P2pTransport> p2p_transport_;
+    std::unique_ptr<device::RdmaTransport> rdma_transport_;
+#endif
+
 #ifdef WITH_METRICS
     // Latency bucket in microseconds
     inline static const std::vector<double> kTaskLatencyBuckets = {
diff --git a/mooncake-transfer-engine/include/transfer_metadata.h b/mooncake-transfer-engine/include/transfer_metadata.h
index a4522a22..715cd7ca 100644
--- a/mooncake-transfer-engine/include/transfer_metadata.h
+++ b/mooncake-transfer-engine/include/transfer_metadata.h
@@ -56,8 +56,16 @@ class TransferMetadata {
 #ifdef ENABLE_MULTI_PROTOCOL
         std::string protocol;  // for multi-protocol mode (cxl/tcp/rdma)
 #endif
-        std::vector<uint32_t> lkey;         // for rdma
-        std::vector<uint32_t> rkey;         // for rdma
+        // EFA/CXI's libfabric provider returns 64-bit MR keys (fi_mr_key()), so
+        // these must be 64-bit wide to avoid truncation. RDMA verbs keys are
+        // 32-bit and the non-EFA/CXI path keeps them as such.
+#if defined(USE_EFA) || defined(USE_CXI)
+        using mr_key_t = uint64_t;
+#else
+        using mr_key_t = uint32_t;
+#endif
+        std::vector<mr_key_t> lkey;         // for rdma/efa
+        std::vector<mr_key_t> rkey;         // for rdma/efa
         std::string shm_name;               // for nvlink and hip
         uint64_t offset;                    // for cxl
         std::vector<std::string> tseg;      // for ub/urma
@@ -104,6 +112,19 @@ class TransferMetadata {
 
         int tcp_data_port;
 
+        // In dual-NIC setups (MC_RDMA_BIND_ADDRESS), the RDMA-reachable
+        // address may differ from the TCP-routable segment name.  When
+        // non-empty, NIC paths are constructed using this value instead
+        // of `name`.
+        std::string rdma_server_name;
+
+        // Returns the server name to use for NIC path construction.
+        // Uses rdma_server_name when available, otherwise falls back
+        // to name.
+        const std::string &nicPathServerName() const {
+            return rdma_server_name.empty() ? name : rdma_server_name;
+        }
+
         void dump() const;
     };
 
@@ -118,6 +139,8 @@ class TransferMetadata {
 
     struct HandShakeDesc {
         std::string local_nic_path;
+        uint16_t local_lid = 0;
+        std::string local_gid;
         std::string peer_nic_path;
 #ifdef USE_UB
         std::vector<uint32_t> jetty_num;  // for ub/urma
@@ -129,6 +152,9 @@ class TransferMetadata {
         std::string reply_msg;  // on error
 #ifdef USE_EFA
         std::string efa_addr;  // EFA endpoint address (hex encoded)
+#endif
+#ifdef USE_CXI
+        std::string cxi_addr;
 #endif
     };
 
@@ -176,6 +202,9 @@ class TransferMetadata {
 
     int removeRpcMetaEntry(const std::string &server_name);
 
+    // Re-publish the local RPC meta entry to the HTTP metadata server.
+    int rePublishRpcMetaEntry(const std::string &server_name);
+
     int getRpcMetaEntry(const std::string &server_name, RpcMetaDesc &desc);
     int getNotifies(std::vector<NotifyDesc> &notifies);
 
diff --git a/mooncake-transfer-engine/include/transport/ascend_transport/ascend_direct_transport/adxl_compat.h b/mooncake-transfer-engine/include/transport/ascend_transport/ascend_direct_transport/adxl_compat.h
index 9bd397d2..1c40901b 100644
--- a/mooncake-transfer-engine/include/transport/ascend_transport/ascend_direct_transport/adxl_compat.h
+++ b/mooncake-transfer-engine/include/transport/ascend_transport/ascend_direct_transport/adxl_compat.h
@@ -16,10 +16,12 @@
 #ifndef ASCEND_DIRECT_TRANSPORT_ADXL_COMPAT_H
 #define ASCEND_DIRECT_TRANSPORT_ADXL_COMPAT_H
 
-#include <adxl/adxl_types.h>
+#include <graph/ascend_string.h>
 
-#include <memory>
+#include <cstddef>
+#include <cstdint>
 #include <map>
+#include <memory>
 #include <vector>
 
 #ifdef FUNC_VISIBILITY
@@ -29,6 +31,50 @@
 #endif
 
 namespace adxl {
+using Status = uint32_t;
+using AscendString = ge::AscendString;
+using TransferReq = void*;
+using MemHandle = void*;
+
+constexpr Status SUCCESS = 0U;
+constexpr Status PARAM_INVALID = 103900U;
+constexpr Status TIMEOUT = 103901U;
+constexpr Status NOT_CONNECTED = 103902U;
+constexpr Status ALREADY_CONNECTED = 103903U;
+constexpr Status NOTIFY_FAILED = 103904U;
+constexpr Status UNSUPPORTED = 103905U;
+constexpr Status FAILED = 503900U;
+constexpr Status RESOURCE_EXHAUSTED = 203900U;
+
+enum MemType { MEM_DEVICE, MEM_HOST };
+
+enum TransferOp { READ, WRITE };
+
+struct MemDesc {
+    uintptr_t addr;
+    size_t len;
+    uint8_t reserved[128] = {};
+};
+
+struct TransferOpDesc {
+    uintptr_t local_addr;
+    uintptr_t remote_addr;
+    size_t len;
+};
+
+enum class TransferStatus { WAITING, COMPLETED, TIMEOUT, FAILED };
+
+struct TransferArgs {
+    uint8_t reserved[128] = {};
+};
+
+// FeatureType values must be explicitly assigned. Append new capabilities at
+// the end only.
+enum FeatureType : int32_t {
+    AUTO_CONNECT = 0,
+    CLIENT_SERVER_COMM = 1,
+};
+
 class ASCEND_FUNC_VISIBILITY AdxlEngine {
    public:
     /**
@@ -145,10 +191,28 @@ class ASCEND_FUNC_VISIBILITY AdxlEngine {
      */
     __attribute__((weak)) static Status FreeMem(void* ptr);
 
+    /**
+     * @brief 查询库能力特性
+     * @param [in] feature_type 特性类型
+     * @param [out] value 1=支持，0=不支持
+     * @return 成功:SUCCESS；未知特性:UNSUPPORTED；参数非法:PARAM_INVALID
+     */
+    __attribute__((weak)) static Status GetCapability(FeatureType feature_type,
+                                                      int32_t& value);
+
    private:
     class AdxlEngineImpl;
     std::unique_ptr<AdxlEngineImpl> impl_;
 };
+
+inline bool IsAdxlFeatureSupported(FeatureType feature) {
+    if (&AdxlEngine::GetCapability == nullptr) {
+        return false;
+    }
+    int32_t val = 0;
+    const Status st = AdxlEngine::GetCapability(feature, val);
+    return st == SUCCESS && val == 1;
+}
 }  // namespace adxl
 
 #endif  // ASCEND_DIRECT_TRANSPORT_ADXL_COMPAT_H
diff --git a/mooncake-transfer-engine/include/transport/ascend_transport/ascend_direct_transport/ascend_direct_transport.h b/mooncake-transfer-engine/include/transport/ascend_transport/ascend_direct_transport/ascend_direct_transport.h
index 383d090a..b6970814 100644
--- a/mooncake-transfer-engine/include/transport/ascend_transport/ascend_direct_transport/ascend_direct_transport.h
+++ b/mooncake-transfer-engine/include/transport/ascend_transport/ascend_direct_transport/ascend_direct_transport.h
@@ -77,8 +77,22 @@ class AscendDirectTransport : public Transport {
                                const std::string &host_ip, SegmentDesc *desc);
 
     int32_t base_port_ = 20000;
-    bool dummy_real_mode_{false};
+    // Mirrors globalConfig().ascend_agent_mode. ascend_agent_mode is enabled
+    // both by the dummy client (inference/GPU process) and by the standalone
+    // real client (real_client_main). The dummy client never installs this
+    // transport, so within AscendDirectTransport this flag effectively means
+    // "I am the independently-launched real client". In this mode the local
+    // segment enumerates every local NPU device via ContextManager (the real
+    // client owns all devices) instead of using the single current-device
+    // context.
+    bool agent_mode_{false};
     bool roce_mode_{false};
+    // Whether this TE uses Ascend fabric memory (host mem for Store, e.g. RoCE
+    // D2H/H2D). Captured once at install time from globalConfig() and gated on
+    // ascend_store_te_init, so a non-Store (e.g. P2P/HCCS) TE never inherits a
+    // Store TE's fabric setting that leaked into the process-global flag. All
+    // transfer-time decisions read this member, not the global.
+    bool use_fabric_mem_{false};
     std::vector<aclrtContext> local_engine_contexts_;
 
     std::unique_ptr<TransferExecutorBase> transfer_executor_;
diff --git a/mooncake-transfer-engine/include/transport/ascend_transport/ascend_direct_transport/async_transfer_executor.h b/mooncake-transfer-engine/include/transport/ascend_transport/ascend_direct_transport/async_transfer_executor.h
index f33d6879..4b96d9c2 100644
--- a/mooncake-transfer-engine/include/transport/ascend_transport/ascend_direct_transport/async_transfer_executor.h
+++ b/mooncake-transfer-engine/include/transport/ascend_transport/ascend_direct_transport/async_transfer_executor.h
@@ -18,10 +18,11 @@
 
 #include <atomic>
 #include <condition_variable>
+#include <cstdint>
 #include <mutex>
 #include <queue>
+#include <string>
 #include <thread>
-#include <unordered_map>
 
 #include "transfer_executor_base.h"
 #include "transport/transport.h"
@@ -45,15 +46,38 @@ class AsyncTransferExecutor : public TransferExecutorBase {
         const std::vector<Transport::Slice*>& slice_list) override;
 
    private:
+    struct QueryBatch {
+        std::vector<Transport::Slice*> slices;
+        size_t engine_idx;
+        std::string target_adxl_engine_name;
+    };
+
+    struct BatchPollResult {
+        bool done = false;
+        // When true, disconnect once and fail every batch already in
+        // pending_batches for this route. Batches enqueued after disconnect
+        // are not affected.
+        bool fail_entire_route = false;
+        // When true with auto_connect, ADXL DisconnectOnError already ran;
+        // only forgetConnectedSegment is needed (if the route was recorded).
+        bool adxl_already_disconnected = false;
+        std::string fail_reason;
+    };
+
     void queryThreadLoop();
-    void fetchPendingBatches(
-        std::vector<std::vector<Transport::Slice*>>& out_batches);
-    bool processOneBatch(std::vector<Transport::Slice*>& slice_list);
-    void handleTaskFinished(adxl::TransferReq handle);
+    void fetchPendingBatches(std::vector<QueryBatch>& out_batches);
+    void processOneBatch(QueryBatch& batch, BatchPollResult& result);
+    void failAllPendingOnRoute(size_t engine_idx, const std::string& target,
+                               std::vector<QueryBatch>& pending,
+                               const std::string& reason,
+                               bool adxl_already_disconnected);
+    void markBatchFailed(QueryBatch& batch, const std::string& reason,
+                         bool log_error = true);
+    void handleTaskFinished();
 
     std::atomic<bool> running_{false};
     std::thread query_thread_;
-    std::queue<std::vector<Transport::Slice*>> query_slice_queue_;
+    std::queue<QueryBatch> query_slice_queue_;
     std::mutex query_mutex_;
     std::condition_variable query_cv_;
 
@@ -61,9 +85,9 @@ class AsyncTransferExecutor : public TransferExecutorBase {
     size_t active_async_tasks_{0};
     std::mutex async_task_mutex_;
     std::condition_variable async_task_cv_;
-    std::mutex async_handle_map_mutex_;
-    std::unordered_map<uintptr_t, size_t> async_handle_to_engine_idx_;
     bool finalized_{false};
+
+    size_t last_context_engine_idx_{SIZE_MAX};
 };
 
 }  // namespace mooncake
diff --git a/mooncake-transfer-engine/include/transport/ascend_transport/ascend_direct_transport/transfer_executor_base.h b/mooncake-transfer-engine/include/transport/ascend_transport/ascend_direct_transport/transfer_executor_base.h
index 687bfece..e07e9716 100644
--- a/mooncake-transfer-engine/include/transport/ascend_transport/ascend_direct_transport/transfer_executor_base.h
+++ b/mooncake-transfer-engine/include/transport/ascend_transport/ascend_direct_transport/transfer_executor_base.h
@@ -54,11 +54,12 @@ class TransferExecutorBase {
         int32_t connect_timeout = 10000;
         int32_t transfer_timeout = 10000;
         bool use_async_transfer = false;
-        bool auto_connect = false;
+        bool auto_connect = true;
         bool use_short_connection = false;
         bool use_buffer_pool = false;
-        bool dummy_real_mode = false;
+        bool agent_mode = false;
         bool roce_mode = false;
+        bool use_fabric_mem = false;
     };
 
     explicit TransferExecutorBase(const InitParams& params);
@@ -77,7 +78,7 @@ class TransferExecutorBase {
     void processSliceList(const std::vector<Transport::Slice*>& slice_list);
 
     int registerMem(void* addr, size_t length, adxl::MemType mem_type,
-                    bool use_buffer_pool, bool roce_mode, bool dummy_real_mode);
+                    bool use_buffer_pool, bool roce_mode, bool agent_mode);
     int deregisterMem(void* addr);
 
     const size_t getNumEngines() const { return adxl_engines_.size(); }
@@ -93,6 +94,8 @@ class TransferExecutorBase {
     void cleanupConnections();
     void finalizeEngines();
     void disconnectAllForEngine(size_t engine_idx);
+    void recordConnectedSegment(size_t engine_idx, const std::string& remote);
+    void forgetConnectedSegment(size_t engine_idx, const std::string& remote);
 
     int checkAndConnect(size_t engine_idx,
                         const std::string& target_adxl_engine_name);
diff --git a/mooncake-transfer-engine/include/transport/ascend_transport/ascend_direct_transport/utils.h b/mooncake-transfer-engine/include/transport/ascend_transport/ascend_direct_transport/utils.h
index eb455909..8b5470d5 100644
--- a/mooncake-transfer-engine/include/transport/ascend_transport/ascend_direct_transport/utils.h
+++ b/mooncake-transfer-engine/include/transport/ascend_transport/ascend_direct_transport/utils.h
@@ -79,6 +79,23 @@ std::string GenAdxlEngineName(const std::string &ip, const uint64_t port);
 
 uint16_t FindAdxlListenPort(int32_t base_port, int32_t device_id);
 
+// True when HCCL_INTRA_ROCE_ENABLE=1 or ASCEND_GLOBAL_RESOURCE_CONFIG sets
+// comm_resource_config.protocol_desc to a roce:* entry (string or array).
+bool IsRoceModeEnabled();
+
+// Parse ASCEND_GLOBAL_RESOURCE_CONFIG JSON for roce protocol_desc only.
+bool HasRoceProtocolDescInGlobalResourceConfig(const char *config_str);
+
+// Resolve the effective ASCEND_GLOBAL_RESOURCE_CONFIG for the current TE role.
+// Schema: the top-level object is the default config (passed verbatim to adxl
+// for normal/P2P TEs); an optional "store" sub-object overrides it for a
+// Store-init TE (globalConfig().ascend_store_te_init == true).
+//   - null/empty or no "store" key  -> returned verbatim (backward compatible)
+//   - has "store" and store-init TE  -> the "store" sub-object, serialized
+//   - has "store" and normal TE      -> root with "store" removed, serialized
+// A Store TE with no "store" key falls back to the default config.
+std::string ResolveAscendGlobalResourceConfig(const char *config_str);
+
 int SetDeviceAndGetContext(int32_t device_id, aclrtContext *out_context);
 
 /**
diff --git a/mooncake-transfer-engine/include/transport/cxi_transport/cxi_context.h b/mooncake-transfer-engine/include/transport/cxi_transport/cxi_context.h
new file mode 100644
index 00000000..e606e45b
--- /dev/null
+++ b/mooncake-transfer-engine/include/transport/cxi_transport/cxi_context.h
@@ -0,0 +1,216 @@
+// Copyright 2024 KVCache.AI
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#ifndef CXI_CONTEXT_H
+#define CXI_CONTEXT_H
+
+#include <gflags/gflags.h>
+#include <glog/logging.h>
+#include <rdma/fabric.h>
+#include <rdma/fi_domain.h>
+#include <rdma/fi_endpoint.h>
+#include <rdma/fi_cm.h>
+#include <rdma/fi_rma.h>
+#include <rdma/fi_errno.h>
+
+#include <atomic>
+#include <condition_variable>
+#include <cstdint>
+#include <list>
+#include <memory>
+#include <string>
+#include <thread>
+#include <unordered_map>
+
+#include "common.h"
+#include "cxi_transport.h"
+#include "transport/transport.h"
+
+namespace mooncake {
+
+class CxiEndpoint;
+class CxiTransport;
+
+struct CxiCq {
+    CxiCq() : cq(nullptr), outstanding(0) {}
+    struct fid_cq* cq;
+    std::atomic<int> outstanding;
+};
+
+struct CxiMemoryRegionMeta {
+    void* addr;
+    size_t length;
+    struct fid_mr* mr;
+    uint64_t key;
+    bool associated;
+};
+
+// CxiContext represents the set of resources controlled by each local Cxi
+// device, including Memory Region, CQ, EndPoint, etc. using libfabric
+class CxiContext {
+   public:
+    CxiContext(CxiTransport& engine, const std::string& device_name);
+
+    ~CxiContext();
+
+    int construct(size_t num_cq_list = 1, size_t max_cqe = 4096,
+                  int max_endpoints = 65536);
+
+   private:
+    int deconstruct();
+    int buildSharedEndpoint(size_t max_wr, size_t max_inline);
+
+   public:
+    // Memory Region Management
+    int registerMemoryRegion(void* addr, size_t length, int access);
+    int unregisterMemoryRegion(void* addr);
+    int preTouchMemory(void* addr, size_t length);
+    uint64_t rkey(void* addr);
+    uint64_t lkey(void* addr);
+    void* mrDesc(void* addr);  // Get MR descriptor for fi_write local_desc
+
+   private:
+    int registerMemoryRegionInternal(void* addr, size_t length, int access,
+                                     CxiMemoryRegionMeta& mrMeta);
+
+   public:
+    bool active() const { return active_; }
+    void set_active(bool flag) { active_ = flag; }
+
+   public:
+    // Get or create a per-peer handle.  Does NOT open an fid_ep or call
+    // fi_enable — the shared endpoint was created once at construct() time.
+    // The returned CxiEndpoint only carries {peer_fi_addr_t, status, mutex}
+    // and, when connected, routes sends through this context's shared_ep_.
+    std::shared_ptr<CxiEndpoint> endpoint(const std::string& peer_nic_path);
+
+    // Non-creating lookup under the normalized key.  Returns nullptr if the
+    // peer handle does not yet exist.  Safe for idempotency checks.
+    std::shared_ptr<CxiEndpoint> peekEndpoint(const std::string& peer_nic_path);
+
+    int deleteEndpoint(const std::string& peer_nic_path);
+    int disconnectAllEndpoints();
+
+    // Number of live peer handles.  Historically named "QP number"; with the
+    // shared endpoint model the actual QP count is always 1 per context.
+    size_t getTotalQPNumber() const;
+
+   public:
+    // Access to engine for endpoint handshake
+    CxiTransport& engine() { return engine_; }
+    const CxiTransport& engine() const { return engine_; }
+
+    // Submit slices for transfer
+    int submitPostSend(const std::vector<Transport::Slice*>& slice_list);
+
+    // Hot-path submit: post a batch of slices to `peer_fi_addr` via the
+    // shared endpoint.  Handles WR / CQ reservation, MR descriptor prep,
+    // and the fi_write / fi_read burst under post_lock_.  Called by
+    // CxiEndpoint::submitPostSend once the peer is connected.
+    int submitSlicesOnPeer(fi_addr_t peer_fi_addr,
+                           std::vector<Transport::Slice*>& slice_list,
+                           std::vector<Transport::Slice*>& failed_slice_list);
+
+    // Poll completion queue for completed operations
+    int pollCq(int max_entries, int cq_index = 0);
+
+    // Get CQ count
+    size_t cqCount() const { return cq_list_.size(); }
+
+   public:
+    // Device name, such as `rdmap0s2`
+    std::string deviceName() const { return device_name_; }
+
+    // NIC Path, such as `192.168.3.76@rdmap0s2`
+    std::string nicPath() const;
+
+   public:
+    // Libfabric accessors
+    struct fid_fabric* fabric() const { return fabric_; }
+    struct fid_domain* domain() const { return domain_; }
+    struct fid_av* av() const { return av_; }
+    struct fi_info* info() const { return fi_info_; }
+    std::string localAddr() const;
+
+    // Local (shared-endpoint) address in hex, for inclusion in handshake.
+    // Populated by construct() -> buildSharedEndpoint().
+    std::string localEpAddr() const;
+
+    // Raw bytes of the local endpoint address.  Use this for loopback
+    // (skip the hex encode/decode round-trip) or for any caller that
+    // already has the bytes.
+    const std::vector<uint8_t>& localEpAddrBytes() const {
+        return local_ep_addr_;
+    }
+
+    // Insert a peer's hex-encoded CXI address into this context's AV and
+    // return the resulting fi_addr_t.  Thread-safe (fi_av_insert is safe
+    // under libfabric's domain-level threading).
+    int insertPeerAddr(const std::string& peer_hex_addr, fi_addr_t& out);
+
+    // Binary variant — avoids the hex-decode when the caller already
+    // has the raw address bytes (e.g. loopback).
+    int insertPeerAddrBytes(const uint8_t* addr, size_t len, fi_addr_t& out);
+
+    // Remove a peer from the AV.  No-op if fi_addr is FI_ADDR_UNSPEC.
+    void removePeerAddr(fi_addr_t fi_addr);
+
+    // Compatibility methods (libfabric doesn't use lid/gid like ibverbs)
+    uint16_t lid() const { return 0; }
+    std::string gid() const { return localAddr(); }
+
+   private:
+    CxiTransport& engine_;
+    std::string device_name_;
+
+    // Libfabric objects
+    struct fi_info* fi_info_;
+    struct fi_info* hints_;
+    struct fid_fabric* fabric_;
+    struct fid_domain* domain_;
+    struct fid_av* av_;  // Address vector for peer addressing
+
+    bool active_;
+
+    // ---- Shared endpoint (one per local NIC, serves ALL peers) ----
+    struct fid_ep* shared_ep_;
+    std::vector<uint8_t> local_ep_addr_;  // bytes returned by fi_getname()
+    // Pacing for outstanding work requests on the shared endpoint.  Shared
+    // across all peers routed through this context.  std::atomic<int> so the
+    // submit-path fetch_add and the CQ-poller fetch_sub obey the C++ memory
+    // model; plain `volatile int` + __sync_* was UB under the current
+    // standard.
+    std::atomic<int> wr_depth_;
+    int max_wr_depth_;
+    // CQ that shared_ep_ is bound to (FI_TRANSMIT|FI_RECV).  Points into
+    // cq_list_[0]; kept here to avoid re-indexing on the hot path.
+    std::shared_ptr<CxiCq> shared_cq_;
+    // Serializes fi_write / fi_read calls on shared_ep_.  libfabric's RDM
+    // endpoints are not thread-safe for concurrent post, even with
+    // FI_THREAD_SAFE at the domain level.
+    std::atomic_flag post_lock_;
+
+    std::vector<std::shared_ptr<CxiCq>> cq_list_;
+
+    // ---- Peer handles (one entry per peer, each ~constant size) ----
+    mutable RWSpinlock peer_map_lock_;
+    std::unordered_map<std::string, std::shared_ptr<CxiEndpoint>> peer_map_;
+
+    RWSpinlock mr_lock_;
+    std::map<uint64_t, CxiMemoryRegionMeta> mr_map_;
+};
+
+}  // namespace mooncake
+
+#endif  // CXI_CONTEXT_H
diff --git a/mooncake-transfer-engine/include/transport/cxi_transport/cxi_endpoint.h b/mooncake-transfer-engine/include/transport/cxi_transport/cxi_endpoint.h
new file mode 100644
index 00000000..9a9febaf
--- /dev/null
+++ b/mooncake-transfer-engine/include/transport/cxi_transport/cxi_endpoint.h
@@ -0,0 +1,132 @@
+// Copyright 2024 KVCache.AI
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#ifndef CXI_ENDPOINT_H
+#define CXI_ENDPOINT_H
+
+#include <glog/logging.h>
+#include <rdma/fabric.h>
+#include <rdma/fi_domain.h>
+#include <rdma/fi_endpoint.h>
+#include <rdma/fi_cm.h>
+#include <rdma/fi_rma.h>
+
+#include <atomic>
+#include <cstdint>
+#include <queue>
+#include <string>
+#include <vector>
+
+#include "common.h"
+#include "cxi_context.h"
+#include "transfer_metadata.h"
+#include "transport/transport.h"
+
+namespace mooncake {
+
+class CxiContext;
+
+// Custom context for libfabric operations - stores slice pointer for
+// completion handling.  This struct MUST have fi_context as its first member.
+struct CxiOpContext {
+    struct fi_context fi_ctx;  // Must be first member
+    Transport::Slice* slice;   // Slice pointer for completion handling
+    // Pointer to the context's wr_depth_ for CQ completion decrement.
+    // std::atomic<int> (not volatile int) so the CQ-poller decrement
+    // and the submit-path fetch_add play by the C++ memory model.
+    std::atomic<int>* wr_depth;
+};
+
+// Per-peer handle in the shared-endpoint SRD model.
+//
+// This class does NOT own an fid_ep.  A single `shared_ep_` on the owning
+// CxiTransport services every peer; each CxiEndpoint just carries one
+// fi_addr_t (AV index) plus the handshake state needed to populate it.
+//
+// Lifecycle:
+//   1. construct()            : trivial; no libfabric resources allocated.
+//   2. setupConnectionsByActive(peer): handshake -> fi_av_insert -> CONNECTED.
+//   3. submitPostSend()       : delegates to CxiTransport::submitSlicesOnPeer.
+//   4. disconnect() / dtor    : fi_av_remove(peer_fi_addr_).
+class CxiEndpoint {
+   public:
+    using HandShakeDesc = TransferMetadata::HandShakeDesc;
+
+    enum Status { INITIALIZING, UNCONNECTED, CONNECTED };
+
+    explicit CxiEndpoint(CxiContext& context);
+    ~CxiEndpoint();
+
+   public:
+    void setPeerNicPath(const std::string& peer_nic_path);
+
+    int setupConnectionsByActive();
+
+    int setupConnectionsByActive(const std::string& peer_nic_path) {
+        setPeerNicPath(peer_nic_path);
+        return setupConnectionsByActive();
+    }
+
+    int setupConnectionsByPassive(const HandShakeDesc& peer_desc,
+                                  HandShakeDesc& local_desc);
+
+    // Always false under the shared-endpoint model: outstanding work is
+    // tracked at the context (shared) level, not per-peer.  Retained for
+    // API compatibility with callers that still ask.
+    bool hasOutstandingSlice() const { return false; }
+
+    bool connected() const {
+        return status_.load(std::memory_order_relaxed) == CONNECTED;
+    }
+
+    void disconnect();
+
+    // Called during CxiTransport teardown: forget the AV slot WITHOUT calling
+    // fi_av_remove().  The AV itself is about to be closed, which invalidates
+    // every slot in one shot; calling fi_av_remove after the shared endpoint
+    // has been closed trips an assertion inside the CXI provider.
+    void markDetachedForTeardown();
+
+   private:
+    void disconnectUnlocked();
+
+   public:
+    const std::string toString() const;
+
+    // Submit a batch of slices bound for this peer.  Internally establishes
+    // the connection if needed, then delegates to
+    // CxiTransport::submitSlicesOnPeer using this peer's fi_addr_t.
+    int submitPostSend(std::vector<Transport::Slice*>& slice_list,
+                       std::vector<Transport::Slice*>& failed_slice_list);
+
+    // Always 1: the shared endpoint backs every peer.  Kept so callers that
+    // sum "QP count" across endpoints don't break.
+    size_t getQPNumber() const { return 1; }
+
+    fi_addr_t getPeerFiAddr() const { return peer_fi_addr_; }
+
+    CxiContext& context() { return context_; }
+
+   private:
+    CxiContext& context_;
+    std::atomic<Status> status_;
+
+    RWSpinlock lock_;  // protects peer_nic_path_ and status_
+    std::string peer_nic_path_;
+    fi_addr_t peer_fi_addr_;  // slot in context_.av()
+};
+
+}  // namespace mooncake
+
+#endif  // CXI_ENDPOINT_H
diff --git a/mooncake-transfer-engine/include/transport/cxi_transport/cxi_transport.h b/mooncake-transfer-engine/include/transport/cxi_transport/cxi_transport.h
new file mode 100644
index 00000000..0243b038
--- /dev/null
+++ b/mooncake-transfer-engine/include/transport/cxi_transport/cxi_transport.h
@@ -0,0 +1,174 @@
+// Copyright 2024 KVCache.AI
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#ifndef CXI_TRANSPORT_H_
+#define CXI_TRANSPORT_H_
+
+#include <infiniband/verbs.h>
+
+#include <atomic>
+#include <cstddef>
+#include <map>
+#include <memory>
+#include <mutex>
+#include <string>
+#include <thread>
+#include <unordered_map>
+#include <unordered_set>
+#include <vector>
+
+#include "topology.h"
+#include "transfer_metadata.h"
+#include "transport/transport.h"
+
+namespace mooncake {
+
+class CxiContext;
+class CxiEndpoint;
+class TransferMetadata;
+
+enum NicReplicaPolicy { REPLICATE_ALL, NUMA_AWARE, DEBUG };
+
+class CxiTransport : public Transport {
+    friend class CxiContext;
+    friend class CxiEndpoint;
+
+   public:
+    using BufferDesc = TransferMetadata::BufferDesc;
+    using SegmentDesc = TransferMetadata::SegmentDesc;
+    using HandShakeDesc = TransferMetadata::HandShakeDesc;
+
+   public:
+    CxiTransport();
+
+    ~CxiTransport();
+
+    int install(std::string& local_server_name,
+                std::shared_ptr<TransferMetadata> meta,
+                std::shared_ptr<Topology> topo) override;
+
+    const char* getName() const override { return "cxi"; }
+
+    int registerLocalMemory(void* addr, size_t length,
+                            const std::string& location, bool remote_accessible,
+                            bool update_metadata) override;
+
+    int unregisterLocalMemory(void* addr, bool update_metadata = true) override;
+
+    int registerLocalMemoryBatch(const std::vector<BufferEntry>& buffer_list,
+                                 const std::string& location) override;
+
+    int unregisterLocalMemoryBatch(
+        const std::vector<void*>& addr_list) override;
+
+    // Eagerly populate the address vector with every (local_ctx, peer_nic)
+    // handshake for `segment_name`.
+    //
+    // Under the shared-endpoint model each warmup is handshake RPC +
+    // fi_av_insert — no fi_endpoint/fi_enable.  This removes the first-
+    // submit stall (handshakes on the data path) without consuming QPs.
+    //
+    // Safe to call multiple times (idempotent: endpoint() + setup are both
+    // idempotent).  Re-run after any openSegment() on a new peer.
+    int warmupSegment(const std::string& segment_name);
+
+   private:
+    // Internal version with force_sequential option to avoid nested parallelism
+    int registerLocalMemoryInternal(void* addr, size_t length,
+                                    const std::string& location,
+                                    bool remote_accessible,
+                                    bool update_metadata,
+                                    bool force_sequential);
+
+    int unregisterLocalMemoryInternal(void* addr, bool update_metadata,
+                                      bool force_sequential);
+
+    // TRANSFER
+
+    Status submitTransfer(BatchID batch_id,
+                          const std::vector<TransferRequest>& entries) override;
+
+    Status submitTransferTask(
+        const std::vector<TransferTask*>& task_list) override;
+
+    Status getTransferStatus(BatchID batch_id,
+                             std::vector<TransferStatus>& status);
+
+    Status getTransferStatus(BatchID batch_id, size_t task_id,
+                             TransferStatus& status) override;
+
+    SegmentID getSegmentID(const std::string& segment_name);
+
+    NicReplicaPolicy getReplicaPolicy();
+
+   private:
+    int allocateLocalSegmentID();
+
+    int preTouchMemory(void* addr, size_t length);
+
+   public:
+    int onSetupCxiConnections(const HandShakeDesc& peer_desc,
+                              HandShakeDesc& local_desc);
+
+    int sendHandshake(const std::string& peer_server_name,
+                      const HandShakeDesc& local_desc,
+                      HandShakeDesc& peer_desc) {
+        return metadata_->sendHandshake(peer_server_name, local_desc,
+                                        peer_desc);
+    }
+
+    const std::string& local_server_name() const { return local_server_name_; }
+
+    std::shared_ptr<TransferMetadata> meta() { return metadata_; }
+
+   private:
+    int initializeCxiResources();
+
+    int startHandshakeDaemon(std::string& local_server_name);
+
+   public:
+    static int selectDevice(SegmentDesc* desc, uint64_t offset, size_t length,
+                            int& buffer_id, int& device_id, int retry_cnt = 0);
+    static int selectDevice(SegmentDesc* desc, uint64_t offset, size_t length,
+                            std::string_view hint, int& buffer_id,
+                            int& device_id, int retry_cnt = 0);
+
+   private:
+    // Start/stop CQ polling worker threads
+    void startWorkerThreads();
+    void stopWorkerThreads();
+    void workerThreadFunc(int thread_id);
+
+   private:
+    std::vector<std::shared_ptr<CxiContext>> context_list_;
+    std::shared_ptr<Topology> local_topology_;
+
+    // Track chunked MR registrations for per-NIC partitioned buffers.
+    // When a buffer exceeds max_mr_size, it is split into chunks, each
+    // registered on a disjoint subset of NICs (per-NIC partition).
+    struct ChunkRegistration {
+        uint64_t addr;
+        std::vector<size_t> nic_indices;
+    };
+    std::mutex chunk_map_mutex_;
+    std::unordered_map<uint64_t, std::vector<ChunkRegistration>> chunk_map_;
+
+    // CQ polling worker threads
+    std::atomic<bool> worker_running_{false};
+    std::vector<std::thread> worker_threads_;
+};
+
+}  // namespace mooncake
+
+#endif  // CXI_TRANSPORT_H_
diff --git a/mooncake-transfer-engine/include/transport/device/comm_device.cuh b/mooncake-transfer-engine/include/transport/device/comm_device.cuh
new file mode 100644
index 00000000..29e6c6c6
--- /dev/null
+++ b/mooncake-transfer-engine/include/transport/device/comm_device.cuh
@@ -0,0 +1,136 @@
+// Communication device API — top-level context and routing.
+//
+// CommCtx bundles all transport state the kernel needs.  The kernel
+// constructs one from the raw pointers passed in and calls mc_route_put /
+// mc_signal / mc_red_add instead of touching transport internals directly.
+#pragma once
+
+#include "transport/device/device_ops.cuh"
+#include "transport/device/p2p_device.cuh"
+#include "transport/device/ibgda_device.cuh"
+
+namespace mooncake {
+namespace device {
+
+// ---------------------------------------------------------------------------
+// CommCtx
+// ---------------------------------------------------------------------------
+
+struct CommCtx {
+    P2PContext p2p;
+    IbgdaContext ibgda;
+    int rank;
+};
+
+// Construct CommCtx from the raw kernel arguments.
+// raddrs/rkeys/qp_devctxs may be nullptr on MUSA (ignored).
+__device__ __forceinline__ CommCtx make_comm_ctx(
+    void* gdr_buffer, const int32_t* nvlink_available,
+    void* const* ipc_peer_ptrs, void* raddrs, void* rkeys, void* qp_devctxs,
+    const void* rdma_send_signal_buffer, const void* rdma_recv_signal_buffer,
+    int rank, int num_ranks, int num_qps) {
+    CommCtx ctx;
+    ctx.rank = rank;
+
+    ctx.p2p.available = nvlink_available;
+    ctx.p2p.peer_ptrs = ipc_peer_ptrs;
+    ctx.p2p.local_base = gdr_buffer;
+
+    ctx.ibgda.qp_devctxs = reinterpret_cast<mlx5gda_qp_devctx*>(qp_devctxs);
+    ctx.ibgda.raddrs = reinterpret_cast<const uint64_t*>(raddrs);
+    ctx.ibgda.rkeys = reinterpret_cast<const uint32_t*>(rkeys);
+    ctx.ibgda.local_atomic_base = rdma_send_signal_buffer;
+    ctx.ibgda.remote_atomic_base = rdma_recv_signal_buffer;
+
+    return ctx;
+}
+
+// ---------------------------------------------------------------------------
+// Routing helpers
+// ---------------------------------------------------------------------------
+
+__device__ __forceinline__ bool mc_comm_p2p_available(const CommCtx& ctx,
+                                                      int dst_rank) {
+    return mc_p2p_available(ctx.p2p, dst_rank);
+}
+
+// Translate a local GDR pointer to the peer's mapped VA.
+__device__ __forceinline__ void* mc_comm_peer_ptr(const CommCtx& ctx,
+                                                  int dst_rank,
+                                                  const void* local_ptr) {
+    return mc_p2p_peer_ptr(ctx.p2p, dst_rank, local_ptr);
+}
+
+// ---------------------------------------------------------------------------
+// mc_route_put
+//
+// Returns the destination pointer for a warp-cooperative copy:
+//   - local rank:  recv_ptr itself (caller does UNROLLED_WARP_COPY)
+//   - P2P rank:    peer-mapped recv_ptr (caller does UNROLLED_WARP_COPY)
+//   - IBGDA rank:  nullptr (caller must stage data then call mc_rdma_put)
+// ---------------------------------------------------------------------------
+__device__ __forceinline__ void* mc_route_put(const CommCtx& ctx, int dst_rank,
+                                              void* recv_ptr) {
+    if (dst_rank == ctx.rank) return recv_ptr;
+    if (mc_comm_p2p_available(ctx, dst_rank))
+        return mc_comm_peer_ptr(ctx, dst_rank, recv_ptr);
+    return nullptr;  // IBGDA path
+}
+
+// Issue an IBGDA RDMA WRITE.  Call only when mc_route_put returned nullptr.
+// lane_id: only lane 0 issues the WQE.
+__device__ __forceinline__ void mc_rdma_put(
+    const CommCtx& ctx, int channel, int dst_rank, int qps_per_rank,
+    const void* send_ptr,
+    void* recv_ptr,  // local VA of the recv slot (for raddr computation)
+    uint32_t nbytes, int lane_id) {
+    if (lane_id == 0) {
+        uint64_t recv_raddr =
+            ctx.ibgda.raddrs[dst_rank] +
+            (reinterpret_cast<const char*>(recv_ptr) -
+             reinterpret_cast<const char*>(ctx.p2p.local_base));
+        mc_ibgda_put(ctx.ibgda, channel, dst_rank, ctx.rank, qps_per_rank,
+                     send_ptr, recv_raddr, nbytes);
+    }
+}
+
+// ---------------------------------------------------------------------------
+// mc_signal / mc_red_add
+//
+// Route a signal (store) or reduction (atomic add) to dst_rank.
+// sig_ptr is a local VA within the GDR buffer.
+// ---------------------------------------------------------------------------
+
+__device__ __forceinline__ void mc_signal(const CommCtx& ctx, int dst_rank,
+                                          int channel, int qps_per_rank,
+                                          int* sig_ptr, int32_t val) {
+    if (dst_rank == ctx.rank) {
+        mc_st_release(sig_ptr, val);
+        return;
+    }
+    if (mc_comm_p2p_available(ctx, dst_rank)) {
+        mc_p2p_signal(ctx.p2p, dst_rank, sig_ptr, val);
+    } else {
+        uint64_t recv_raddr =
+            ctx.ibgda.raddrs[dst_rank] +
+            (reinterpret_cast<const char*>(sig_ptr) -
+             reinterpret_cast<const char*>(ctx.p2p.local_base));
+        uint64_t laddr =
+            ctx.ibgda.raddrs[ctx.rank] +
+            (reinterpret_cast<const char*>(sig_ptr) -
+             reinterpret_cast<const char*>(ctx.ibgda.remote_atomic_base)) +
+            (reinterpret_cast<const char*>(ctx.ibgda.local_atomic_base) -
+             reinterpret_cast<const char*>(ctx.p2p.local_base));
+        mc_ibgda_red_add(ctx.ibgda, channel, dst_rank, ctx.rank, qps_per_rank,
+                         laddr, recv_raddr, val);
+    }
+}
+
+__device__ __forceinline__ void mc_red_add(const CommCtx& ctx, int dst_rank,
+                                           int channel, int qps_per_rank,
+                                           int* sig_ptr, int32_t val) {
+    mc_signal(ctx, dst_rank, channel, qps_per_rank, sig_ptr, val);
+}
+
+}  // namespace device
+}  // namespace mooncake
diff --git a/mooncake-transfer-engine/include/transport/device/cuda/cuda_ops.cuh b/mooncake-transfer-engine/include/transport/device/cuda/cuda_ops.cuh
new file mode 100644
index 00000000..f48d1744
--- /dev/null
+++ b/mooncake-transfer-engine/include/transport/device/cuda/cuda_ops.cuh
@@ -0,0 +1,157 @@
+// CUDA PTX implementations of device-side memory ordering primitives.
+// Included by device_ops.cuh when MOONCAKE_EP_USE_MUSA is not defined.
+#pragma once
+
+#include <cuda_runtime.h>
+#include <cooperative_groups.h>
+
+namespace mooncake {
+namespace device {
+
+// ---------------------------------------------------------------------------
+// Acquire loads — cross-GPU visibility (sys scope)
+// ---------------------------------------------------------------------------
+__device__ __forceinline__ int mc_ld_acquire(const int* ptr) {
+    int ret;
+    asm volatile("ld.acquire.sys.global.s32 %0, [%1];" : "=r"(ret) : "l"(ptr));
+    return ret;
+}
+
+__device__ __forceinline__ uint64_t mc_ld_acquire_u64(const uint64_t* ptr) {
+    uint64_t ret;
+    asm volatile("ld.acquire.sys.global.u64 %0, [%1];" : "=l"(ret) : "l"(ptr));
+    return ret;
+}
+
+// ---------------------------------------------------------------------------
+// Release stores — cross-GPU visibility (sys scope), non-temporal (no alloc)
+// ---------------------------------------------------------------------------
+__device__ __forceinline__ void mc_st_release(const int* ptr, int val) {
+    asm volatile("st.release.sys.global.L1::no_allocate.s32 [%0], %1;"
+                 :
+                 : "l"(ptr), "r"(val));
+}
+
+__device__ __forceinline__ void mc_st_release_u32(const uint32_t* ptr,
+                                                  uint32_t val) {
+    asm volatile("st.release.sys.global.L1::no_allocate.b32 [%0], %1;"
+                 :
+                 : "l"(ptr), "r"(val));
+}
+
+__device__ __forceinline__ void mc_st_release_u64(const uint64_t* ptr,
+                                                  uint64_t val) {
+    asm volatile("st.release.sys.global.L1::no_allocate.b64 [%0], %1;"
+                 :
+                 : "l"(ptr), "l"(val));
+}
+
+// ---------------------------------------------------------------------------
+// Atomic add — release semantics, sys scope
+// ---------------------------------------------------------------------------
+__device__ __forceinline__ int mc_atomic_add_release(const int* ptr, int val) {
+    int ret;
+    asm volatile("atom.add.release.sys.global.s32 %0, [%1], %2;"
+                 : "=r"(ret)
+                 : "l"(ptr), "r"(val));
+    return ret;
+}
+
+// ---------------------------------------------------------------------------
+// Non-coherent loads (read-only cache, no L1 alloc) — for bulk data reads
+// ---------------------------------------------------------------------------
+__device__ __forceinline__ int4 mc_ld_nc(const int4* ptr) {
+    int4 ret;
+    asm volatile(
+        "ld.global.nc.L1::no_allocate.L2::256B.v4.s32 {%0,%1,%2,%3}, [%4];"
+        : "=r"(ret.x), "=r"(ret.y), "=r"(ret.z), "=r"(ret.w)
+        : "l"(ptr));
+    return ret;
+}
+
+__device__ __forceinline__ int mc_ld_nc_s32(const int* ptr) {
+    int ret;
+    asm volatile("ld.global.nc.L1::no_allocate.s32 %0, [%1];"
+                 : "=r"(ret)
+                 : "l"(ptr));
+    return ret;
+}
+
+__device__ __forceinline__ float mc_ld_nc_f32(const float* ptr) {
+    float ret;
+    asm volatile("ld.global.nc.L1::no_allocate.f32 %0, [%1];"
+                 : "=f"(ret)
+                 : "l"(ptr));
+    return ret;
+}
+
+__device__ __forceinline__ int64_t mc_ld_nc_s64(const int64_t* ptr) {
+    int64_t ret;
+    asm volatile("ld.global.nc.L1::no_allocate.s64 %0, [%1];"
+                 : "=l"(ret)
+                 : "l"(ptr));
+    return ret;
+}
+
+// ---------------------------------------------------------------------------
+// Non-temporal stores (no L1 alloc) — for bulk data writes
+// ---------------------------------------------------------------------------
+__device__ __forceinline__ void mc_st_na(const int4* ptr, const int4& val) {
+    asm volatile("st.global.L1::no_allocate.v4.s32 [%0], {%1,%2,%3,%4};"
+                 :
+                 : "l"(ptr), "r"(val.x), "r"(val.y), "r"(val.z), "r"(val.w));
+}
+
+// ---------------------------------------------------------------------------
+// Named barrier init — no-op on CUDA (hardware named barriers need no setup).
+// ---------------------------------------------------------------------------
+__device__ __forceinline__ void mc_bar_init() {}
+
+// ---------------------------------------------------------------------------
+// Named barrier (warp-group scope) — CUDA PTX bar.sync
+// On MUSA this maps to __syncthreads(); see musa_ops.cuh for details.
+// ---------------------------------------------------------------------------
+__device__ __forceinline__ void mc_bar_sync(int bar_id, int num_threads) {
+    asm volatile("bar.sync %0, %1;" : : "r"(bar_id), "r"(num_threads));
+}
+
+// ---------------------------------------------------------------------------
+// Grid-level sync — cooperative_groups::this_grid().sync().
+// On MUSA this is a no-op because the host always uses separate kernel
+// launches (return_recv_hook=true), so SEND and RECV never run in the same
+// kernel invocation.
+// ---------------------------------------------------------------------------
+__device__ __forceinline__ void mc_grid_sync() {
+    cooperative_groups::this_grid().sync();
+}
+
+// ---------------------------------------------------------------------------
+// System-level memory fence: no-op on CUDA (hardware guarantees P2P
+// visibility through NVLink without explicit fences).
+// ---------------------------------------------------------------------------
+__device__ __forceinline__ void mc_fence() {}
+
+// ---------------------------------------------------------------------------
+// Fence/barrier/fence: ensures all threads' writes are globally visible
+// before any thread proceeds.  On CUDA, __syncthreads() implies a memory
+// fence, so a single __syncthreads() is sufficient.
+// ---------------------------------------------------------------------------
+__device__ __forceinline__ void mc_fence_barrier_fence() { __syncthreads(); }
+
+// ---------------------------------------------------------------------------
+// Byte-swap helpers (for mlx5 big-endian WQE fields)
+// ---------------------------------------------------------------------------
+__device__ __forceinline__ uint16_t mc_bswap16(uint16_t x) {
+    return __byte_perm(x, x, 0x2301);
+}
+__device__ __forceinline__ uint32_t mc_bswap32(uint32_t x) {
+    return __byte_perm(x, x, 0x0123);
+}
+__device__ __forceinline__ uint64_t mc_bswap64(uint64_t x) {
+    uint32_t hi = __byte_perm((uint32_t)(x >> 32), 0, 0x0123);
+    uint32_t lo = __byte_perm((uint32_t)(x), 0, 0x0123);
+    return ((uint64_t)lo << 32) | hi;
+}
+
+}  // namespace device
+}  // namespace mooncake
diff --git a/mooncake-transfer-engine/include/transport/device/device_ops.cuh b/mooncake-transfer-engine/include/transport/device/device_ops.cuh
new file mode 100644
index 00000000..2835f815
--- /dev/null
+++ b/mooncake-transfer-engine/include/transport/device/device_ops.cuh
@@ -0,0 +1,12 @@
+// Device-side memory ordering primitives — platform-portable selector.
+//
+// This is the ONLY file in the device API that contains a platform #ifdef.
+// All other device API headers include this file and use the mc_* functions.
+// The kernel itself has zero platform #ifdef branches.
+#pragma once
+
+#ifdef MOONCAKE_EP_USE_MUSA
+#include "transport/device/musa/musa_ops.cuh"
+#else
+#include "transport/device/cuda/cuda_ops.cuh"
+#endif
diff --git a/mooncake-transfer-engine/include/transport/device/device_transport.h b/mooncake-transfer-engine/include/transport/device/device_transport.h
new file mode 100644
index 00000000..82556e7d
--- /dev/null
+++ b/mooncake-transfer-engine/include/transport/device/device_transport.h
@@ -0,0 +1,168 @@
+// Copyright 2024 KVCache.AI
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#pragma once
+
+// EP DeviceTransport — platform-agnostic host-side abstraction for the two
+// communication paths used by the EP kernel:
+//
+//   P2pTransport   — intra-node GPU-initiated P2P (NVLink on CUDA, MTLink on
+//                    MUSA).  Manages IPC handle exchange and peer pointer
+//                    table.
+//   RdmaTransport  — inter-node GPU-initiated RDMA (IBGDA / mlx5gda).
+//                    Manages QP lifecycle, MR, and device context table.
+//
+// EP code includes only this header and calls the abstract interface.
+// Platform-specific implementations live in device/ and are
+// selected at build time by the factory functions at the bottom of this file.
+//
+// The header intentionally avoids including cuda_alike.h so it can be included
+// from pure C++ translation units.  Implementations include cuda_alike.h.
+
+#include <cstdint>
+#include <memory>
+#include <string>
+#include <vector>
+
+namespace mooncake {
+namespace device {
+
+// ---------------------------------------------------------------------------
+// P2pTransport
+//
+// Manages intra-node GPU P2P IPC handles and the device-visible peer pointer
+// table used by the EP kernel.
+// ---------------------------------------------------------------------------
+class P2pTransport {
+   public:
+    virtual ~P2pTransport() = default;
+
+    // Allocate the GDR buffer that will be shared via IPC.
+    // Returns a device pointer; size must be > 0.
+    virtual void* allocateBuffer(size_t bytes) = 0;
+
+    // Free a buffer previously returned by allocateBuffer.
+    virtual void freeBuffer(void* ptr) = 0;
+
+    // Export an IPC handle for the buffer allocated by this rank.
+    // Returns a byte blob (serialised as int32_t array for Python compat).
+    // Returns empty vector if IPC is not needed (e.g. fabric memory).
+    virtual std::vector<int32_t> exportIpcHandle(void* ptr) = 0;
+
+    // Import peer IPC handles and populate the device-visible tables.
+    // remote_handles[i] is the handle exported by rank i (may be empty for
+    // ranks that use fabric memory or are on a different node).
+    // active_ranks_mask[i] == 1 means rank i is participating.
+    // After this call, availableTablePtr() and peerPtrsTablePtr() are valid.
+    virtual void importPeerHandles(
+        void* local_ptr, int rank, int num_ranks,
+        const std::vector<std::vector<int32_t>>& remote_handles,
+        const std::vector<int>& active_ranks_mask) = 0;
+
+    // Device pointer to int32_t[num_ranks]: 1 if P2P is available to rank i.
+    virtual int32_t* availableTablePtr() = 0;
+
+    // Device pointer to void*[num_ranks]: peer GDR buffer base pointers.
+    virtual void** peerPtrsTablePtr() = 0;
+
+    // True if all active ranks have P2P access (fast-path condition).
+    virtual bool allPeersAccessible() const = 0;
+
+    // Verify that peer-mapped memory is writable (via memcpy).
+    // Returns true if a small test write to each peer's mapped buffer
+    // succeeds.  On failure, sets all_peers_accessible to false.
+    virtual bool verifyPeerAccess() = 0;
+};
+
+// ---------------------------------------------------------------------------
+// RdmaLocalMetadata — exchanged between ranks during IBGDA bootstrap.
+// ---------------------------------------------------------------------------
+struct RdmaLocalMetadata {
+    int64_t raddr;
+    int32_t rkey;
+    int64_t subnet_prefix;
+    int64_t interface_id;
+    std::vector<int32_t> qpns;
+    std::vector<int32_t> lids;
+};
+
+// ---------------------------------------------------------------------------
+// RdmaTransport
+//
+// Manages IBGDA QP lifecycle, MR registration, and the device-visible RDMA
+// context tables (raddrs, rkeys, qp_devctxs) used by the EP kernel.
+// ---------------------------------------------------------------------------
+class RdmaTransport {
+   public:
+    virtual ~RdmaTransport() = default;
+
+    // Initialise the RDMA transport for the given NIC.
+    // device_name: e.g. "mlx5_1".  Pass empty string for auto-detect.
+    // Returns 0 on success, non-zero on failure (IBGDA disabled).
+    virtual int initialize(const std::string& device_name, int num_ranks,
+                           int num_qps) = 0;
+
+    // Register the GDR buffer for RDMA access.
+    virtual int registerMemory(void* ptr, size_t bytes) = 0;
+
+    // Allocate the GPU-side control buffer (QP/CQ structures).
+    virtual int allocateControlBuffer() = 0;
+
+    // Create QPs in RST→INIT state.  Call after allocateControlBuffer.
+    // stream is a cudaStream_t / musaStream_t cast to void*.
+    virtual int createQueuePairs(void* stream) = 0;
+
+    // Destroy and recreate QPs (used when active_ranks changes).
+    virtual int recreateQueuePairs(void* stream) = 0;
+
+    // Connect QPs to peers using exchanged metadata.
+    // is_roce: true for RoCE, false for IB.
+    virtual int connectPeers(int local_rank, bool is_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<int32_t>& remote_lids,
+                             const std::vector<int64_t>& subnet_prefixes,
+                             const std::vector<int64_t>& interface_ids,
+                             const std::vector<int>& active_ranks_mask) = 0;
+
+    // Metadata for this rank, to be exchanged with peers.
+    virtual RdmaLocalMetadata localMetadata() const = 0;
+
+    // Device pointers to the tables consumed by the EP kernel.
+    virtual void* raddrsPtr() = 0;     // uint64_t[num_ranks]
+    virtual void* rkeysPtr() = 0;      // uint32_t[num_ranks]
+    virtual void* qpDevCtxsPtr() = 0;  // mlx5gda_qp_devctx[num_qps]
+
+    virtual bool isRoce() const = 0;
+    virtual int gidIndex() const = 0;
+};
+
+// ---------------------------------------------------------------------------
+// Factory functions — implemented in device_transport.cpp.
+// Returns nullptr if the transport is not available on this platform.
+// ---------------------------------------------------------------------------
+
+// Create the platform-native P2P transport (NVLink on CUDA, MTLink on MUSA).
+std::unique_ptr<P2pTransport> createP2pDeviceTransport(int num_ranks);
+
+// Create the IBGDA RDMA transport backed by TE's RdmaContext.
+// device_filter: optional whitelist of NIC names (e.g. {"mlx5_1", "mlx5_2"}).
+//   Empty vector = auto-detect via TE's Topology::discover() with no filter.
+//   Non-empty = restrict discovery to these NICs, then pick the closest one.
+std::unique_ptr<RdmaTransport> createIbgdaDeviceTransport(
+    const std::vector<std::string>& device_filter = {});
+
+}  // namespace device
+}  // namespace mooncake
diff --git a/mooncake-ep/include/mooncake_ibgda/compiler.h b/mooncake-transfer-engine/include/transport/device/ibgda/compiler.h
similarity index 100%
rename from mooncake-ep/include/mooncake_ibgda/compiler.h
rename to mooncake-transfer-engine/include/transport/device/ibgda/compiler.h
diff --git a/mooncake-ep/include/mooncake_ibgda/memheap.h b/mooncake-transfer-engine/include/transport/device/ibgda/memheap.h
similarity index 100%
rename from mooncake-ep/include/mooncake_ibgda/memheap.h
rename to mooncake-transfer-engine/include/transport/device/ibgda/memheap.h
diff --git a/mooncake-ep/include/mooncake_ibgda/mlx5_ifc.h b/mooncake-transfer-engine/include/transport/device/ibgda/mlx5_ifc.h
similarity index 100%
rename from mooncake-ep/include/mooncake_ibgda/mlx5_ifc.h
rename to mooncake-transfer-engine/include/transport/device/ibgda/mlx5_ifc.h
diff --git a/mooncake-ep/include/mooncake_ibgda/mlx5_prm.h b/mooncake-transfer-engine/include/transport/device/ibgda/mlx5_prm.h
similarity index 100%
rename from mooncake-ep/include/mooncake_ibgda/mlx5_prm.h
rename to mooncake-transfer-engine/include/transport/device/ibgda/mlx5_prm.h
diff --git a/mooncake-ep/include/mooncake_ibgda/mlx5gda.h b/mooncake-transfer-engine/include/transport/device/ibgda/mlx5gda.h
similarity index 97%
rename from mooncake-ep/include/mooncake_ibgda/mlx5gda.h
rename to mooncake-transfer-engine/include/transport/device/ibgda/mlx5gda.h
index bca275da..82030954 100644
--- a/mooncake-ep/include/mooncake_ibgda/mlx5gda.h
+++ b/mooncake-transfer-engine/include/transport/device/ibgda/mlx5gda.h
@@ -3,7 +3,12 @@
 
 #include <stddef.h>
 #include <stdint.h>
+#ifdef USE_MUSA
+#include <musa_runtime.h>
+#define cudaStream_t musaStream_t
+#else
 #include <cuda_runtime.h>
+#endif
 #include <infiniband/verbs.h>
 #include <infiniband/mlx5dv.h>
 
diff --git a/mooncake-ep/include/mooncake_ibgda/os.h b/mooncake-transfer-engine/include/transport/device/ibgda/os.h
similarity index 100%
rename from mooncake-ep/include/mooncake_ibgda/os.h
rename to mooncake-transfer-engine/include/transport/device/ibgda/os.h
diff --git a/mooncake-transfer-engine/include/transport/device/ibgda_device.cuh b/mooncake-transfer-engine/include/transport/device/ibgda_device.cuh
new file mode 100644
index 00000000..15fcb990
--- /dev/null
+++ b/mooncake-transfer-engine/include/transport/device/ibgda_device.cuh
@@ -0,0 +1,206 @@
+// IBGDA device context and operations.
+//
+// Wraps mlx5gda_qp_devctx and issues RDMA writes/atomics via
+// device-side WQE construction.
+//
+// On MUSA: BF (Blue Flame) doorbell is not available (musaHostRegisterIoMemory
+// doesn't support MMIO), so qp->bf is NULL and the kernel uses DBR-only mode.
+// All other IBGDA logic (WQE construction, CQ polling, DBR write) is shared.
+#pragma once
+
+#include <cstdint>
+#include "transport/device/device_ops.cuh"
+
+#ifndef MOONCAKE_EP_USE_MUSA
+#include <cuda/atomic>
+#endif
+#include <transport/device/ibgda/mlx5gda.h>
+
+// mlx5 32-bit atomic-add WQE segment (not in mlx5gda.h; defined here for use
+// in mc_ibgda_write_rdma_atomic_add_wqe).
+struct mlx5_wqe_atomic_add_32_seg {
+    __be32 add_data;
+    __be32 field_boundary;
+    __be64 compare;
+};
+
+namespace mooncake {
+namespace device {
+
+// ---------------------------------------------------------------------------
+// IbgdaContext
+// ---------------------------------------------------------------------------
+
+struct IbgdaContext {
+    mlx5gda_qp_devctx* qp_devctxs;   // device ptr: [num_qps]
+    const uint64_t* raddrs;          // device ptr: [num_ranks] remote GDR base
+    const uint32_t* rkeys;           // device ptr: [num_ranks] remote rkey
+    const void* local_atomic_base;   // local scratch base for atomic responses
+    const void* remote_atomic_base;  // symmetric remote signal base
+};
+
+// ---------------------------------------------------------------------------
+// Helpers
+// ---------------------------------------------------------------------------
+
+__device__ __forceinline__ mlx5gda_qp_devctx* mc_ibgda_channel(
+    const IbgdaContext& ctx, int channel, int dst_rank, int qps_per_rank) {
+    int qp_idx = dst_rank * qps_per_rank + (channel % qps_per_rank);
+    return ctx.qp_devctxs + qp_idx;
+}
+
+__device__ __forceinline__ void mc_ibgda_lock(mlx5gda_qp_devctx* qp) {
+#ifdef MOONCAKE_EP_USE_MUSA
+    uint32_t old;
+    do {
+        old = atomicCAS(&qp->mutex, 0u, 1u);
+    } while (old != 0);
+#else
+    cuda::atomic_ref<uint32_t, cuda::thread_scope_system> lock(qp->mutex);
+    while (lock.exchange(1u, cuda::memory_order_acquire) != 0u);
+#endif
+}
+
+__device__ __forceinline__ void mc_ibgda_unlock(mlx5gda_qp_devctx* qp) {
+#ifdef MOONCAKE_EP_USE_MUSA
+    mc_st_release_u32(&qp->mutex, 0u);
+#else
+    cuda::atomic_ref<uint32_t, cuda::thread_scope_system> lock(qp->mutex);
+    lock.store(0u, cuda::memory_order_release);
+#endif
+}
+
+__device__ __forceinline__ void mc_ibgda_poll_cq(mlx5gda_qp_devctx* qp,
+                                                 uint16_t expect) {
+    uint16_t wq_tail = qp->wq_tail;
+    while (static_cast<int16_t>(wq_tail - expect) <= 0) {
+        uint16_t cq_be =
+            *reinterpret_cast<volatile uint16_t*>(&qp->cq->wqe_counter);
+        uint8_t opcode = qp->cq->op_own >> 4;
+        if (opcode == 0xD)
+            printf("[EP IBGDA] Requester error: syndrome=0x%lx\n",
+                   qp->cq->timestamp >> 56);
+        if (!(opcode == 0x0 || opcode == 0xF)) {
+            printf("[EP IBGDA] Unexpected CQE opcode=0x%x, trapping\n", opcode);
+            __trap();
+        }
+        wq_tail = mc_bswap16(cq_be) + 1;
+    }
+    if (wq_tail != qp->wq_tail) qp->wq_tail = wq_tail;
+}
+
+__device__ __forceinline__ void mc_ibgda_post_send_db(mlx5gda_qp_devctx* qp) {
+    uint32_t num_posted = static_cast<uint32_t>(qp->wq_head);
+    // DBR write — always done (NIC polls doorbell record in GPU memory)
+    mc_st_release_u32(reinterpret_cast<uint32_t*>(&qp->dbr->send_counter),
+                      mc_bswap32(num_posted));
+    // BF (Blue Flame) doorbell — only if BF register is mapped into GPU VA.
+    // On MUSA, musaHostRegisterIoMemory fails for MMIO addresses, so bf is
+    // NULL and we rely on DBR-only mode (slightly higher latency).
+    if (qp->bf != nullptr) {
+        auto* last_wqe = qp->wq + ((num_posted - 1) & qp->wqeid_mask);
+        mc_st_release_u64(reinterpret_cast<uint64_t*>(qp->bf + qp->bf_offset),
+                          *reinterpret_cast<uint64_t*>(last_wqe));
+        qp->bf_offset ^= MLX5GDA_BF_SIZE;
+    }
+}
+
+// Issue an RDMA WRITE WQE.  laddr/raddr are device VAs; keys are big-endian.
+__device__ __forceinline__ void mc_ibgda_write_rdma_write_wqe(
+    mlx5gda_qp_devctx* qp, uint64_t laddr, __be32 lkey, uint64_t raddr,
+    __be32 rkey, uint32_t bytes) {
+    auto* wqe = reinterpret_cast<mlx5gda_rdma_write_wqe*>(
+        qp->wq + (qp->wq_head & qp->wqeid_mask));
+
+    wqe->ctrl = {};
+    wqe->ctrl.qpn_ds = mc_bswap32((qp->qpn << 8) | 3);
+    wqe->ctrl.fm_ce_se = MLX5_WQE_CTRL_CQ_UPDATE;
+    wqe->ctrl.opmod_idx_opcode = mc_bswap32(
+        (static_cast<uint32_t>(qp->wq_head) << 8) | MLX5_OPCODE_RDMA_WRITE);
+
+    wqe->raddr.raddr = mc_bswap64(raddr);
+    wqe->raddr.rkey = rkey;
+    wqe->raddr.reserved = 0;
+
+    wqe->data.byte_count = mc_bswap32(bytes);
+    wqe->data.lkey = lkey;
+    wqe->data.addr = mc_bswap64(laddr);
+
+    ++qp->wq_head;
+}
+
+// Issue an RDMA ATOMIC MASKED FETCH-AND-ADD WQE (32-bit add_data).
+// This matches the original CUDA IBGDA EP kernel. A regular 64-bit
+// MLX5_OPCODE_ATOMIC_FA does not implement the 32-bit signal-buffer add used
+// by dispatch/combine.
+__device__ __forceinline__ void mc_ibgda_write_rdma_atomic_add_wqe(
+    mlx5gda_qp_devctx* qp, int32_t value, uint64_t laddr, __be32 lkey,
+    uint64_t raddr, __be32 rkey) {
+    auto* wqe = reinterpret_cast<mlx5gda_rdma_atomic_wqe*>(
+        qp->wq + (qp->wq_head & qp->wqeid_mask));
+
+    wqe->ctrl = {};
+    wqe->ctrl.qpn_ds = mc_bswap32((qp->qpn << 8) | 4);
+    wqe->ctrl.fm_ce_se = MLX5_WQE_CTRL_CQ_UPDATE;
+    wqe->ctrl.opmod_idx_opcode =
+        mc_bswap32(MLX5_OPCODE_ATOMIC_MASKED_FA |
+                   (static_cast<uint32_t>(qp->wq_head) << 8) | 0x08000000);
+
+    wqe->raddr.raddr = mc_bswap64(raddr);
+    wqe->raddr.rkey = rkey;
+    wqe->raddr.reserved = 0;
+
+    // atomic_seg: add_data (32-bit, big-endian), field_boundary=0, compare=0
+    auto* atomic_seg =
+        reinterpret_cast<mlx5_wqe_atomic_add_32_seg*>(&wqe->atomic);
+    atomic_seg->add_data = mc_bswap32(static_cast<uint32_t>(value));
+    atomic_seg->field_boundary = 0;
+    atomic_seg->compare = 0;
+
+    wqe->data.byte_count = mc_bswap32(static_cast<uint32_t>(4));
+    wqe->data.lkey = lkey;
+    wqe->data.addr = mc_bswap64(laddr);
+
+    ++qp->wq_head;
+}
+
+// ---------------------------------------------------------------------------
+// High-level IBGDA operations
+// ---------------------------------------------------------------------------
+
+// RDMA WRITE: send `nbytes` from `send_ptr` to `recv_ptr` on `dst_rank`.
+// Must be called by lane 0 only.
+__device__ __forceinline__ void mc_ibgda_put(const IbgdaContext& ctx,
+                                             int channel, int dst_rank,
+                                             int src_rank, int qps_per_rank,
+                                             const void* send_ptr,
+                                             uint64_t recv_raddr,
+                                             uint32_t nbytes) {
+    auto* qp = mc_ibgda_channel(ctx, channel, dst_rank, qps_per_rank);
+    mc_ibgda_lock(qp);
+    mc_ibgda_write_rdma_write_wqe(qp, reinterpret_cast<uint64_t>(send_ptr),
+                                  mc_bswap32(ctx.rkeys[src_rank]), recv_raddr,
+                                  mc_bswap32(ctx.rkeys[dst_rank]), nbytes);
+    mc_ibgda_post_send_db(qp);
+    mc_ibgda_unlock(qp);
+}
+
+// RDMA ATOMIC ADD: add `value` to the 32-bit word at `recv_raddr` on
+// `dst_rank`. Must be called by lane 0 only.
+__device__ __forceinline__ void mc_ibgda_red_add(
+    const IbgdaContext& ctx, int channel, int dst_rank, int src_rank,
+    int qps_per_rank,
+    uint64_t laddr,       // local scratch VA for the atomic result
+    uint64_t recv_raddr,  // remote VA of the signal word
+    int32_t value) {
+    auto* qp = mc_ibgda_channel(ctx, channel, dst_rank, qps_per_rank);
+    mc_ibgda_lock(qp);
+    mc_ibgda_write_rdma_atomic_add_wqe(
+        qp, value, laddr, mc_bswap32(ctx.rkeys[src_rank]), recv_raddr,
+        mc_bswap32(ctx.rkeys[dst_rank]));
+    mc_ibgda_post_send_db(qp);
+    mc_ibgda_unlock(qp);
+}
+
+}  // namespace device
+}  // namespace mooncake
diff --git a/mooncake-transfer-engine/include/transport/device/musa/musa_ops.cuh b/mooncake-transfer-engine/include/transport/device/musa/musa_ops.cuh
new file mode 100644
index 00000000..dec9c271
--- /dev/null
+++ b/mooncake-transfer-engine/include/transport/device/musa/musa_ops.cuh
@@ -0,0 +1,160 @@
+// MUSA implementations of device-side memory ordering primitives.
+// Included by device_ops.cuh when MOONCAKE_EP_USE_MUSA is defined.
+//
+// MUSA has no PTX-style acquire/release instructions.  We emulate them with
+// plain loads/stores + __threadfence_system() for cross-GPU (MTLink)
+// visibility.
+//
+// Known MUSA SDK 4.3.3 compiler bugs to avoid:
+//   - atomicAdd_system / atomicCAS_system → infinite SelectionDAG loop.
+//     Use block-scope atomicAdd + __threadfence_system() instead.
+//   - Named barriers (bar.sync) → not available; use __syncthreads().
+//   - cooperative_groups::this_grid().sync() → not available; host MUSA paths
+//     split SEND/RECV into separate kernel launches, so grid sync is a no-op.
+#pragma once
+
+#include <musa_runtime.h>
+
+namespace mooncake {
+namespace device {
+
+// ---------------------------------------------------------------------------
+// Acquire loads
+// ---------------------------------------------------------------------------
+__device__ __forceinline__ int mc_ld_acquire(const int* ptr) {
+    int ret = *const_cast<volatile const int*>(ptr);
+    __threadfence_system();
+    return ret;
+}
+
+__device__ __forceinline__ uint64_t mc_ld_acquire_u64(const uint64_t* ptr) {
+    uint64_t ret = *const_cast<volatile const uint64_t*>(ptr);
+    __threadfence_system();
+    return ret;
+}
+
+// ---------------------------------------------------------------------------
+// Release stores
+// ---------------------------------------------------------------------------
+__device__ __forceinline__ void mc_st_release(const int* ptr, int val) {
+    __threadfence_system();
+    *const_cast<volatile int*>(ptr) = val;
+    __threadfence_system();
+}
+
+__device__ __forceinline__ void mc_st_release_u32(const uint32_t* ptr,
+                                                  uint32_t val) {
+    __threadfence_system();
+    *const_cast<volatile uint32_t*>(ptr) = val;
+    __threadfence_system();
+}
+
+__device__ __forceinline__ void mc_st_release_u64(const uint64_t* ptr,
+                                                  uint64_t val) {
+    __threadfence_system();
+    *const_cast<volatile uint64_t*>(ptr) = val;
+    __threadfence_system();
+}
+
+// ---------------------------------------------------------------------------
+// Atomic add — block-scope atomicAdd + system fence (avoids SDK bug)
+// ---------------------------------------------------------------------------
+__device__ __forceinline__ int mc_atomic_add_release(const int* ptr, int val) {
+    __threadfence_system();
+    int ret = atomicAdd(const_cast<int*>(ptr), val);
+    __threadfence_system();
+    return ret;
+}
+
+// ---------------------------------------------------------------------------
+// Non-coherent loads — MUSA has no nc/no_allocate cache hints; use volatile.
+// int4 volatile copy not supported by MUSA compiler; copy field-by-field.
+// ---------------------------------------------------------------------------
+__device__ __forceinline__ int4 mc_ld_nc(const int4* ptr) {
+    const volatile int* vp = reinterpret_cast<const volatile int*>(ptr);
+    int4 ret;
+    ret.x = vp[0];
+    ret.y = vp[1];
+    ret.z = vp[2];
+    ret.w = vp[3];
+    return ret;
+}
+
+__device__ __forceinline__ int mc_ld_nc_s32(const int* ptr) {
+    return *const_cast<volatile const int*>(ptr);
+}
+
+__device__ __forceinline__ float mc_ld_nc_f32(const float* ptr) {
+    return *const_cast<volatile const float*>(ptr);
+}
+
+__device__ __forceinline__ int64_t mc_ld_nc_s64(const int64_t* ptr) {
+    return *const_cast<volatile const int64_t*>(ptr);
+}
+
+// ---------------------------------------------------------------------------
+// Non-temporal stores — MUSA has no nt/no_allocate hints; plain store.
+// int4 volatile assignment not supported by MUSA compiler; copy field-by-field.
+// ---------------------------------------------------------------------------
+__device__ __forceinline__ void mc_st_na(const int4* ptr, const int4& val) {
+    volatile int* vp = reinterpret_cast<volatile int*>(const_cast<int4*>(ptr));
+    vp[0] = val.x;
+    vp[1] = val.y;
+    vp[2] = val.z;
+    vp[3] = val.w;
+}
+
+// ---------------------------------------------------------------------------
+// Named barrier — MUSA has no bar.sync.
+// On MUSA, mc_bar_sync is implemented as __syncthreads() (full CTA barrier).
+// Kernels that call mc_bar_sync from a subset of threads must ensure the
+// remaining threads call it the same number of times from another code path.
+// See mooncake_ep_kernel.cu for the warp-31 dispatch-send workaround.
+// ---------------------------------------------------------------------------
+__device__ __forceinline__ void mc_bar_sync(int /*bar_id*/,
+                                            int /*num_threads*/) {
+    __syncthreads();
+}
+
+// ---------------------------------------------------------------------------
+// Grid sync — not available on MUSA.  Host always uses separate kernel
+// launches, so SEND and RECV never share a kernel invocation.  No-op.
+// ---------------------------------------------------------------------------
+__device__ __forceinline__ void mc_grid_sync() {}
+
+// ---------------------------------------------------------------------------
+// System-level memory fence: MUSA requires explicit __threadfence_system()
+// for cross-GPU (MTLink) visibility.  CUDA hardware guarantees this without
+// explicit fences, so mc_fence() is a no-op there.
+// ---------------------------------------------------------------------------
+__device__ __forceinline__ void mc_fence() { __threadfence_system(); }
+
+// ---------------------------------------------------------------------------
+// Fence/barrier/fence: ensures all threads' writes are globally visible
+// before any thread proceeds.  On MUSA, __syncthreads() does NOT imply a
+// memory fence, so explicit fences are needed on both sides.
+// ---------------------------------------------------------------------------
+__device__ __forceinline__ void mc_fence_barrier_fence() {
+    mc_fence();
+    mc_bar_sync(0, 0);
+    mc_fence();
+}
+
+// ---------------------------------------------------------------------------
+// Byte-swap helpers — MUSA has no __byte_perm; implement manually.
+// ---------------------------------------------------------------------------
+__device__ __forceinline__ uint16_t mc_bswap16(uint16_t x) {
+    return (uint16_t)(((x & 0x00FFu) << 8) | ((x & 0xFF00u) >> 8));
+}
+__device__ __forceinline__ uint32_t mc_bswap32(uint32_t x) {
+    return ((x & 0x000000FFu) << 24) | ((x & 0x0000FF00u) << 8) |
+           ((x & 0x00FF0000u) >> 8) | ((x & 0xFF000000u) >> 24);
+}
+__device__ __forceinline__ uint64_t mc_bswap64(uint64_t x) {
+    uint32_t hi = mc_bswap32((uint32_t)(x >> 32));
+    uint32_t lo = mc_bswap32((uint32_t)(x));
+    return ((uint64_t)lo << 32) | hi;
+}
+
+}  // namespace device
+}  // namespace mooncake
diff --git a/mooncake-transfer-engine/include/transport/device/p2p_device.cuh b/mooncake-transfer-engine/include/transport/device/p2p_device.cuh
new file mode 100644
index 00000000..d47ca0a2
--- /dev/null
+++ b/mooncake-transfer-engine/include/transport/device/p2p_device.cuh
@@ -0,0 +1,46 @@
+// P2P device context and operations — unified NVLink (CUDA) / MTLink (MUSA).
+//
+// P2PContext holds the device-visible tables populated by P2pTransport on
+// the host side.  The kernel constructs one from the raw pointers passed in
+// and calls mc_p2p_* helpers instead of accessing the tables directly.
+#pragma once
+
+#include <cstdint>
+#include "transport/device/device_ops.cuh"
+
+namespace mooncake {
+namespace device {
+
+struct P2PContext {
+    const int32_t* available;  // device ptr: [num_ranks], 1 = P2P reachable
+    void* const* peer_ptrs;    // device ptr: [num_ranks], peer GDR base ptrs
+    void* local_base;          // this rank's GDR buffer base (for offset math)
+};
+
+__device__ __forceinline__ bool mc_p2p_available(const P2PContext& ctx,
+                                                 int dst_rank) {
+    return ctx.available[dst_rank] != 0 && ctx.peer_ptrs[dst_rank] != nullptr;
+}
+
+// Translate a local pointer (within the GDR buffer) to the peer's mapped VA.
+__device__ __forceinline__ void* mc_p2p_peer_ptr(const P2PContext& ctx,
+                                                 int dst_rank,
+                                                 const void* local_ptr) {
+    const auto offset = reinterpret_cast<const char*>(local_ptr) -
+                        reinterpret_cast<const char*>(ctx.local_base);
+    return reinterpret_cast<char*>(ctx.peer_ptrs[dst_rank]) + offset;
+}
+
+// Write a 32-bit signal value to the peer's signal slot via P2P store.
+// Single-writer assumption: uses release store, not atomic add.
+__device__ __forceinline__ void mc_p2p_signal(const P2PContext& ctx,
+                                              int dst_rank,
+                                              const int* local_sig_ptr,
+                                              int32_t val) {
+    auto* peer_sig =
+        reinterpret_cast<int*>(mc_p2p_peer_ptr(ctx, dst_rank, local_sig_ptr));
+    mc_st_release(peer_sig, val);
+}
+
+}  // namespace device
+}  // namespace mooncake
diff --git a/mooncake-transfer-engine/include/transport/efa_transport/efa_context.h b/mooncake-transfer-engine/include/transport/efa_transport/efa_context.h
index 77aa09d6..067c46d1 100644
--- a/mooncake-transfer-engine/include/transport/efa_transport/efa_context.h
+++ b/mooncake-transfer-engine/include/transport/efa_transport/efa_context.h
@@ -33,6 +33,7 @@
 #include <thread>
 #include <map>
 #include <unordered_map>
+#include <vector>
 
 #include "common.h"
 #include "efa_transport.h"
@@ -46,7 +47,7 @@ class EfaTransport;
 struct EfaCq {
     EfaCq() : cq(nullptr), outstanding(0) {}
     struct fid_cq* cq;
-    volatile int outstanding;
+    std::atomic<int> outstanding;
 };
 
 struct EfaMemoryRegionMeta {
@@ -56,63 +57,30 @@ struct EfaMemoryRegionMeta {
     uint64_t key;
 };
 
-// Endpoint store for EFA with LRU eviction support.
-// When the store approaches the AV capacity (max_endpoints), stale endpoints
-// that have been inactive beyond `inactive_timeout_sec` are evicted to free
-// AV slots.  This prevents AV exhaustion in long-running services that
-// communicate with many transient peers.
-class EfaEndpointStore {
-   public:
-    static constexpr double kDefaultInactiveTimeoutSec = 5.0;  // 5 seconds
-
-    explicit EfaEndpointStore(
-        size_t max_endpoints = 65536,
-        double inactive_timeout_sec = kDefaultInactiveTimeoutSec);
-
-    std::shared_ptr<EfaEndPoint> get(const std::string& peer_nic_path);
-    // Atomically get-or-insert: returns existing endpoint or inserts new_ep.
-    // Prevents duplicate endpoint creation from concurrent callers.
-    // Triggers eviction when the store is at or above max_endpoints.
-    std::shared_ptr<EfaEndPoint> getOrInsert(
-        const std::string& peer_nic_path, std::shared_ptr<EfaEndPoint> new_ep);
-    void add(const std::string& peer_nic_path,
-             std::shared_ptr<EfaEndPoint> endpoint);
-    void remove(const std::string& peer_nic_path);
-    int disconnectAll();
-    size_t size() const;
-
-    // Evict endpoints that have been inactive longer than the configured
-    // timeout.  Returns the number of evicted endpoints.
-    size_t evictStale();
-
-    // Remove endpoints whose connections are broken (not connected and not
-    // initializing).  Returns the number of removed endpoints.
-    size_t removeDisconnected();
-
-   private:
-    // Must be called with write lock held.
-    size_t evictStaleLocked();
-
-    mutable RWSpinlock lock_;
-    std::unordered_map<std::string, std::shared_ptr<EfaEndPoint>> endpoints_;
-    size_t max_endpoints_;
-    double inactive_timeout_sec_;
-};
-
 // EfaContext represents the set of resources controlled by each local EFA
-// device, including Memory Region, CQ, EndPoint, etc. using libfabric
+// device: one libfabric domain, one address vector (AV), one shared endpoint
+// (fid_ep), one or more CQs, and the MR table.
+//
+// Key design point (SRD shared-endpoint model): there is exactly ONE fid_ep
+// per local NIC for the entire process.  Every peer lives as an fi_addr_t
+// entry inside the AV.  Adding a peer is an O(1) `fi_av_insert` into the
+// existing AV — it does NOT consume a QP slot.  Consequences:
+//   * QP usage is a constant 1 per local NIC, regardless of peer count.
+//   * Cold warmup is ~ms per peer (handshake + fi_av_insert) instead of
+//     ~35 ms per peer (fi_endpoint + fi_enable).
+//   * Scale-out is bounded only by AV capacity (65536) instead of 768/NIC.
 class EfaContext {
    public:
     EfaContext(EfaTransport& engine, const std::string& device_name);
 
     ~EfaContext();
 
-    int construct(size_t num_cq_list = 1, size_t num_comp_channels = 1,
-                  uint8_t port = 1, int gid_index = -1, size_t max_cqe = 4096,
+    int construct(size_t num_cq_list = 1, size_t max_cqe = 4096,
                   int max_endpoints = 65536);
 
    private:
     int deconstruct();
+    int buildSharedEndpoint(size_t max_wr, size_t max_inline);
 
    public:
     // Memory Region Management
@@ -132,10 +100,21 @@ class EfaContext {
     void set_active(bool flag) { active_ = flag; }
 
    public:
-    // EndPoint Management
+    // Get or create a per-peer handle.  Does NOT open an fid_ep or call
+    // fi_enable — the shared endpoint was created once at construct() time.
+    // The returned EfaEndPoint only carries {peer_fi_addr_t, status, mutex}
+    // and, when connected, routes sends through this context's shared_ep_.
     std::shared_ptr<EfaEndPoint> endpoint(const std::string& peer_nic_path);
+
+    // Non-creating lookup under the normalized key.  Returns nullptr if the
+    // peer handle does not yet exist.  Safe for idempotency checks.
+    std::shared_ptr<EfaEndPoint> peekEndpoint(const std::string& peer_nic_path);
+
     int deleteEndpoint(const std::string& peer_nic_path);
     int disconnectAllEndpoints();
+
+    // Number of live peer handles.  Historically named "QP number"; with the
+    // shared endpoint model the actual QP count is always 1 per context.
     size_t getTotalQPNumber() const;
 
    public:
@@ -146,30 +125,33 @@ class EfaContext {
     // Submit slices for transfer
     int submitPostSend(const std::vector<Transport::Slice*>& slice_list);
 
+    // Same-process self-loopback fast path.  When a slice's resolved peer
+    // NIC path equals our own (same server_name — which embeds the
+    // per-process RPC port — AND same device), source and destination are
+    // both valid pointers in THIS process's address space.  We satisfy the
+    // copy with a local memcpy / cudaMemcpy instead of issuing it over EFA.
+    //
+    // This avoids libfabric's EFA SHM intra-node path, which performs a
+    // host memcpy into FI_HMEM_CUDA device buffers and segfaults
+    // (ofiwg/libfabric#12328).  Returns true if the slice was handled here
+    // (and marked success/failed); false if it should fall through to the
+    // normal EFA submit path.
+    bool tryLoopbackCopy(Transport::Slice* slice);
+
+    // Hot-path submit: post a batch of slices to `peer_fi_addr` via the
+    // shared endpoint.  Handles WR / CQ reservation, MR descriptor prep,
+    // and the fi_write / fi_read burst under post_lock_.  Called by
+    // EfaEndPoint::submitPostSend once the peer is connected.
+    int submitSlicesOnPeer(fi_addr_t peer_fi_addr,
+                           std::vector<Transport::Slice*>& slice_list,
+                           std::vector<Transport::Slice*>& failed_slice_list);
+
     // Poll completion queue for completed operations
     int pollCq(int max_entries, int cq_index = 0);
 
-    // Evict stale endpoints to free EFA resources
-    size_t evictStaleEndpoints() {
-        return endpoint_store_ ? endpoint_store_->evictStale() : 0;
-    }
-
     // Get CQ count
     size_t cqCount() const { return cq_list_.size(); }
 
-    // Round-robin CQ assignment for new endpoints (mirrors RDMA transport)
-    std::shared_ptr<EfaCq> nextCq() {
-        if (cq_list_.empty()) return nullptr;
-        int index = (next_cq_index_++) % cq_list_.size();
-        return cq_list_[index];
-    }
-
-    // Get CQ outstanding count pointer
-    volatile int* cqOutstandingCount(int cq_index) {
-        if (cq_index < 0 || (size_t)cq_index >= cq_list_.size()) return nullptr;
-        return &cq_list_[cq_index]->outstanding;
-    }
-
    public:
     // Device name, such as `rdmap0s2`
     std::string deviceName() const { return device_name_; }
@@ -185,6 +167,29 @@ class EfaContext {
     struct fi_info* info() const { return fi_info_; }
     std::string localAddr() const;
 
+    // Local (shared-endpoint) address in hex, for inclusion in handshake.
+    // Populated by construct() -> buildSharedEndpoint().
+    std::string localEpAddr() const;
+
+    // Raw bytes of the local endpoint address.  Use this for loopback
+    // (skip the hex encode/decode round-trip) or for any caller that
+    // already has the bytes.
+    const std::vector<uint8_t>& localEpAddrBytes() const {
+        return local_ep_addr_;
+    }
+
+    // Insert a peer's hex-encoded EFA address into this context's AV and
+    // return the resulting fi_addr_t.  Thread-safe (fi_av_insert is safe
+    // under libfabric's domain-level threading).
+    int insertPeerAddr(const std::string& peer_hex_addr, fi_addr_t& out);
+
+    // Binary variant — avoids the hex-decode when the caller already
+    // has the raw address bytes (e.g. loopback).
+    int insertPeerAddrBytes(const uint8_t* addr, size_t len, fi_addr_t& out);
+
+    // Remove a peer from the AV.  No-op if fi_addr is FI_ADDR_UNSPEC.
+    void removePeerAddr(fi_addr_t fi_addr);
+
     // Compatibility methods (libfabric doesn't use lid/gid like ibverbs)
     uint16_t lid() const { return 0; }
     std::string gid() const { return localAddr(); }
@@ -202,9 +207,29 @@ class EfaContext {
 
     bool active_;
 
-    std::shared_ptr<EfaEndpointStore> endpoint_store_;
+    // ---- Shared endpoint (one per local NIC, serves ALL peers) ----
+    struct fid_ep* shared_ep_;
+    std::vector<uint8_t> local_ep_addr_;  // bytes returned by fi_getname()
+    // Pacing for outstanding work requests on the shared endpoint.  Shared
+    // across all peers routed through this context.  std::atomic<int> so the
+    // submit-path fetch_add and the CQ-poller fetch_sub obey the C++ memory
+    // model; plain `volatile int` + __sync_* was UB under the current
+    // standard.
+    std::atomic<int> wr_depth_;
+    int max_wr_depth_;
+    // CQ that shared_ep_ is bound to (FI_TRANSMIT|FI_RECV).  Points into
+    // cq_list_[0]; kept here to avoid re-indexing on the hot path.
+    std::shared_ptr<EfaCq> shared_cq_;
+    // Serializes fi_write / fi_read calls on shared_ep_.  libfabric's RDM
+    // endpoints are not thread-safe for concurrent post, even with
+    // FI_THREAD_SAFE at the domain level.
+    std::atomic_flag post_lock_;
+
     std::vector<std::shared_ptr<EfaCq>> cq_list_;
-    std::atomic<int> next_cq_index_{0};
+
+    // ---- Peer handles (one entry per peer, each ~constant size) ----
+    mutable RWSpinlock peer_map_lock_;
+    std::unordered_map<std::string, std::shared_ptr<EfaEndPoint>> peer_map_;
 
     RWSpinlock mr_lock_;
     std::map<uint64_t, EfaMemoryRegionMeta> mr_map_;
diff --git a/mooncake-transfer-engine/include/transport/efa_transport/efa_endpoint.h b/mooncake-transfer-engine/include/transport/efa_transport/efa_endpoint.h
index 89f04238..7b886be4 100644
--- a/mooncake-transfer-engine/include/transport/efa_transport/efa_endpoint.h
+++ b/mooncake-transfer-engine/include/transport/efa_transport/efa_endpoint.h
@@ -24,7 +24,6 @@
 
 #include <atomic>
 #include <cstdint>
-#include <queue>
 #include <string>
 #include <vector>
 
@@ -37,35 +36,37 @@ namespace mooncake {
 
 class EfaContext;
 
-// Custom context for libfabric operations - stores slice pointer for completion
-// handling This struct MUST have fi_context as its first member
+// Custom context for libfabric operations - stores slice pointer for
+// completion handling.  This struct MUST have fi_context as its first member.
 struct EfaOpContext {
     struct fi_context fi_ctx;  // Must be first member
     Transport::Slice* slice;   // Slice pointer for completion handling
-    volatile int* wr_depth;    // Pointer to endpoint's wr_depth_ for CQ
-                               // completion decrement
+    // Pointer to the context's wr_depth_ for CQ completion decrement.
+    // std::atomic<int> (not volatile int) so the CQ-poller decrement
+    // and the submit-path fetch_add play by the C++ memory model.
+    std::atomic<int>* wr_depth;
 };
 
-// EfaEndPoint represents a libfabric endpoint for EFA communication.
-// Unlike RDMA QPs, EFA uses RDM (Reliable Datagram) endpoints with
-// an address vector for peer addressing.
+// Per-peer handle in the shared-endpoint SRD model.
+//
+// This class does NOT own an fid_ep.  A single `shared_ep_` on the owning
+// EfaContext services every peer; each EfaEndPoint just carries one
+// fi_addr_t (AV index) plus the handshake state needed to populate it.
+//
+// Lifecycle:
+//   1. construct()            : trivial; no libfabric resources allocated.
+//   2. setupConnectionsByActive(peer): handshake -> fi_av_insert -> CONNECTED.
+//   3. submitPostSend()       : delegates to EfaContext::submitSlicesOnPeer.
+//   4. disconnect() / dtor    : fi_av_remove(peer_fi_addr_).
 class EfaEndPoint {
    public:
     using HandShakeDesc = TransferMetadata::HandShakeDesc;
 
     enum Status { INITIALIZING, UNCONNECTED, CONNECTED };
 
-    EfaEndPoint(EfaContext& context);
+    explicit EfaEndPoint(EfaContext& context);
     ~EfaEndPoint();
 
-    // Construct endpoint with specified completion queue
-    int construct(struct fid_cq* cq, volatile int* cq_outstanding,
-                  size_t num_qp_list = 1, size_t max_sge = 4,
-                  size_t max_wr = 256, size_t max_inline = 64);
-
-   private:
-    int deconstruct();
-
    public:
     void setPeerNicPath(const std::string& peer_nic_path);
 
@@ -79,34 +80,22 @@ class EfaEndPoint {
     int setupConnectionsByPassive(const HandShakeDesc& peer_desc,
                                   HandShakeDesc& local_desc);
 
-    bool hasOutstandingSlice() const;
+    // Always false under the shared-endpoint model: outstanding work is
+    // tracked at the context (shared) level, not per-peer.  Retained for
+    // API compatibility with callers that still ask.
+    bool hasOutstandingSlice() const { return false; }
 
-    bool active() const { return active_; }
-
-    void set_active(bool flag) {
-        RWSpinlock::WriteGuard guard(lock_);
-        active_ = flag;
-        if (!flag) inactive_time_ = getCurrentTimeInNano();
-    }
-
-    double inactiveTime() {
-        if (active_) return 0.0;
-        return (getCurrentTimeInNano() - inactive_time_) / 1000000000.0;
-    }
-
-    void touchLastUsed() { last_used_time_ = getCurrentTimeInNano(); }
-
-    double lastUsedAge() const {
-        return (getCurrentTimeInNano() - last_used_time_) / 1000000000.0;
-    }
-
-   public:
     bool connected() const {
         return status_.load(std::memory_order_relaxed) == CONNECTED;
     }
 
     void disconnect();
-    int destroyQP();
+
+    // Called during EfaContext teardown: forget the AV slot WITHOUT calling
+    // fi_av_remove().  The AV itself is about to be closed, which invalidates
+    // every slot in one shot; calling fi_av_remove after the shared endpoint
+    // has been closed trips an assertion inside the EFA provider.
+    void markDetachedForTeardown();
 
    private:
     void disconnectUnlocked();
@@ -114,57 +103,34 @@ class EfaEndPoint {
    public:
     const std::string toString() const;
 
-    // Submit RDMA write/read operations via libfabric
+    // Submit a batch of slices bound for this peer.  Internally establishes
+    // the connection if needed, then delegates to
+    // EfaContext::submitSlicesOnPeer using this peer's fi_addr_t.
     int submitPostSend(std::vector<Transport::Slice*>& slice_list,
                        std::vector<Transport::Slice*>& failed_slice_list);
 
-    // Get the number of endpoints (always 1 for EFA RDM)
+    // Always 1: the shared endpoint backs every peer.  Kept so callers that
+    // sum "QP count" across endpoints don't break.
     size_t getQPNumber() const { return 1; }
 
-    // Get local endpoint address for handshake
-    std::string getLocalAddr() const;
-
-    // Get peer's fi_addr
     fi_addr_t getPeerFiAddr() const { return peer_fi_addr_; }
 
     EfaContext& context() { return context_; }
 
-   private:
-    // Setup connection using peer's address from handshake
-    int doSetupConnection(const std::string& peer_addr,
-                          std::string* reply_msg = nullptr);
-
-    // Insert peer address into address vector
-    int insertPeerAddr(const std::string& peer_addr);
-
    private:
     EfaContext& context_;
     std::atomic<Status> status_;
 
-    RWSpinlock lock_;
+    RWSpinlock lock_;  // protects peer_nic_path_ and status_
     std::string peer_nic_path_;
-
-    // Libfabric endpoint
-    struct fid_ep* ep_;
-    struct fid_cq* tx_cq_;
-    struct fid_cq* rx_cq_;
-    fi_addr_t peer_fi_addr_;  // Peer's address in the AV
-
-    // Local endpoint address (for handshake)
-    std::vector<uint8_t> local_addr_;
-    size_t local_addr_len_;
-
-    volatile int wr_depth_;
-    int max_wr_depth_;
-    volatile int* cq_outstanding_;
-
-    // Spinlock to serialize fi_write/fi_read calls on this endpoint.
-    // libfabric RDM endpoints are not thread-safe by default.
-    std::atomic_flag post_lock_ = ATOMIC_FLAG_INIT;
-
-    volatile bool active_;
-    volatile uint64_t inactive_time_;
-    volatile uint64_t last_used_time_;  // Updated on connection and I/O
+    fi_addr_t peer_fi_addr_;  // slot in context_.av()
+    // Last peer EFA address successfully inserted into the AV.  Used to
+    // make passive handshakes idempotent: when both peers initiate a
+    // handshake concurrently (the common case under bilateral sglang PD
+    // traffic), the second passive handshake carries the same EFA
+    // address as the cached value, and we can skip the fi_av_remove/
+    // fi_av_insert churn entirely.
+    std::string cached_peer_addr_;
 };
 
 }  // namespace mooncake
diff --git a/mooncake-transfer-engine/include/transport/efa_transport/efa_transport.h b/mooncake-transfer-engine/include/transport/efa_transport/efa_transport.h
index 63c6d4ae..a1711986 100644
--- a/mooncake-transfer-engine/include/transport/efa_transport/efa_transport.h
+++ b/mooncake-transfer-engine/include/transport/efa_transport/efa_transport.h
@@ -70,18 +70,15 @@ class EfaTransport : public Transport {
     int unregisterLocalMemoryBatch(
         const std::vector<void*>& addr_list) override;
 
-    // Eagerly establish EFA endpoints to every NIC of `segment_name`.
+    // Eagerly populate the address vector with every (local_ctx, peer_nic)
+    // handshake for `segment_name`.
     //
-    // Rationale: libfabric FI_EP_RDM endpoints resolve peer addresses lazily
-    // via fi_av_insert() on first send. With 16 local NICs × N peer NICs,
-    // the first submitTransfer() of a batch serializes ~N*16
-    // handshake+fi_av_insert round-trips, producing a single-digit-second
-    // stall (observed ~6 s on B300 for the first 100 × 0.5 MB batch). After
-    // this call returns, every (local_ctx, peer_nic) endpoint is CONNECTED
-    // and the first real submitTransfer() goes straight to fi_write/fi_read.
+    // Under the shared-endpoint model each warmup is handshake RPC +
+    // fi_av_insert — no fi_endpoint/fi_enable.  This removes the first-
+    // submit stall (handshakes on the data path) without consuming QPs.
     //
     // Safe to call multiple times (idempotent: endpoint() + setup are both
-    // idempotent). Re-run after any openSegment() on a new peer.
+    // idempotent).  Re-run after any openSegment() on a new peer.
     int warmupSegment(const std::string& segment_name);
 
    private:
diff --git a/mooncake-transfer-engine/include/transport/kunpeng_transport/ub_context.h b/mooncake-transfer-engine/include/transport/kunpeng_transport/ub_context.h
index 7240ff7b..56767c9d 100644
--- a/mooncake-transfer-engine/include/transport/kunpeng_transport/ub_context.h
+++ b/mooncake-transfer-engine/include/transport/kunpeng_transport/ub_context.h
@@ -22,6 +22,7 @@
 #include <memory>
 #include <string>
 #include <thread>
+#include <unordered_map>
 #include <unordered_set>
 #include <sys/epoll.h>
 
@@ -83,6 +84,7 @@ class UbEndpointStore {
     virtual std::shared_ptr<UbEndPoint> insertEndpoint(
         const std::string& peer_nic_path, UbContext* context) = 0;
     virtual int deleteEndpoint(const std::string& peer_nic_path) = 0;
+    virtual int deleteEndpointByPtr(UbEndPoint* point_ptr) = 0;
     virtual void evictEndpoint() = 0;
     virtual void reclaimEndpoint() = 0;
     virtual size_t getSize() = 0;
@@ -102,6 +104,7 @@ class UbSIEVEEndpointStore : public UbEndpointStore {
     std::shared_ptr<UbEndPoint> insertEndpoint(const std::string& peer_nic_path,
                                                UbContext* context) override;
     int deleteEndpoint(const std::string& peer_nic_path) override;
+    int deleteEndpointByPtr(UbEndPoint* point_ptr) override;
     void evictEndpoint() override;
     void reclaimEndpoint() override;
     size_t getSize() override;
@@ -180,7 +183,18 @@ class UbContext {
     virtual int openDevice(const std::string& device_name, uint8_t port,
                            int& eid_index) = 0;
 
-    virtual int poll(int num_entries, Transport::Slice** cr,
+    // Polls one JFC and processes the slices internally:
+    //   * Aggregates each completion's jetty depth into jetty_depth_set.
+    //   * Successful slices have markSuccess() called in place and are NOT
+    //     returned (they may be recycled by the submitting thread the
+    //     moment markSuccess() runs).
+    //   * Failed slices are returned in failed_slices[0..num_failed-1] for
+    //     the caller to apply retry / markFailed.
+    // Returns the total number of completions polled (>= 0), or a negative
+    // error code.
+    virtual int poll(int num_entries, Transport::Slice** failed_slices,
+                     int& num_failed,
+                     std::unordered_map<volatile int*, int>& jetty_depth_set,
                      int jfc_index = 0) = 0;
 
     virtual volatile int* outstandingCount(int jfc_index) = 0;
@@ -229,6 +243,10 @@ class UbContext {
         return endpoint_store_->deleteEndpoint(peer_nic_path);
     }
 
+    int deleteEndpointByPtr(UbEndPoint* point_ptr) {
+        return endpoint_store_->deleteEndpointByPtr(point_ptr);
+    }
+
     int disconnectAllEndpoints() { return endpoint_store_->disconnect(); }
 
     // Device name, such as `mlx5_3`
diff --git a/mooncake-transfer-engine/include/transport/kunpeng_transport/urma/urma_endpoint.h b/mooncake-transfer-engine/include/transport/kunpeng_transport/urma/urma_endpoint.h
index 096e151d..6b6fa707 100644
--- a/mooncake-transfer-engine/include/transport/kunpeng_transport/urma/urma_endpoint.h
+++ b/mooncake-transfer-engine/include/transport/kunpeng_transport/urma/urma_endpoint.h
@@ -58,7 +58,9 @@ class UrmaContext : public UbContext {
     int unregisterMemoryRegion(uint64_t va) override;
     int doProcessContextEvents() override;
     void* retrieveRemoteSeg(const std::string& value) override;
-    int poll(int num_entries, Transport::Slice** cr, int jfc_index) override;
+    int poll(int num_entries, Transport::Slice** failed_slices, int& num_failed,
+             std::unordered_map<volatile int*, int>& jetty_depth_set,
+             int jfc_index) override;
     volatile int* outstandingCount(int jfc_index) override;
     int submitPostSend(
         const std::vector<Transport::Slice*>& slice_list) override;
diff --git a/mooncake-transfer-engine/include/transport/maca_transport/maca_transport.h b/mooncake-transfer-engine/include/transport/maca_transport/maca_transport.h
new file mode 100644
index 00000000..fed0c240
--- /dev/null
+++ b/mooncake-transfer-engine/include/transport/maca_transport/maca_transport.h
@@ -0,0 +1,97 @@
+// Copyright 2024 KVCache.AI
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#ifndef MACA_TRANSPORT_H_
+#define MACA_TRANSPORT_H_
+
+#include "cuda_alike.h"
+
+#include <functional>
+#include <iostream>
+#include <memory>
+#include <mutex>
+#include <string>
+#include <unordered_set>
+#include <vector>
+#include <utility>
+
+#include "common/hash_utils.h"
+#include "topology.h"
+#include "transfer_metadata.h"
+#include "transport/transport.h"
+
+namespace mooncake {
+
+class TransferMetadata;
+
+class MacaTransport : public Transport {
+   public:
+    MacaTransport();
+
+    ~MacaTransport();
+
+    Status submitTransfer(BatchID batch_id,
+                          const std::vector<TransferRequest>& entries) override;
+
+    Status submitTransferTask(
+        const std::vector<TransferTask*>& task_list) override;
+
+    Status getTransferStatus(BatchID batch_id, size_t task_id,
+                             TransferStatus& status) override;
+
+    static void* allocatePinnedLocalMemory(size_t length);
+
+    static void freePinnedLocalMemory(void* addr);
+
+   protected:
+    int install(std::string& local_server_name,
+                std::shared_ptr<TransferMetadata> meta,
+                std::shared_ptr<Topology> topo) override;
+
+    int registerLocalMemory(void* addr, size_t length,
+                            const std::string& location, bool remote_accessible,
+                            bool update_metadata = true) override;
+
+    int unregisterLocalMemory(void* addr, bool update_metadata = true) override;
+
+    int registerLocalMemoryBatch(const std::vector<BufferEntry>& buffer_list,
+                                 const std::string& location) override;
+
+    int unregisterLocalMemoryBatch(
+        const std::vector<void*>& addr_list) override;
+
+    int relocateSharedMemoryAddress(uint64_t& dest_addr, uint64_t length,
+                                    uint64_t target_id);
+
+    const char* getName() const override { return "maca"; }
+
+   private:
+    std::atomic_bool running_;
+
+    struct OpenedShmEntry {
+        void* shm_addr;
+        uint64_t length;
+    };
+
+    std::unordered_map<std::pair<uint64_t, uint64_t>, OpenedShmEntry, PairHash>
+        remap_entries_;
+    RWSpinlock remap_lock_;
+
+    std::mutex register_mutex_;
+    std::unordered_set<uint64_t> registered_base_addrs_;
+};
+
+}  // namespace mooncake
+
+#endif  // MACA_TRANSPORT_H_
diff --git a/mooncake-transfer-engine/include/transport/rdma_transport/endpoint_store.h b/mooncake-transfer-engine/include/transport/rdma_transport/endpoint_store.h
index 87798000..3d99bb2a 100644
--- a/mooncake-transfer-engine/include/transport/rdma_transport/endpoint_store.h
+++ b/mooncake-transfer-engine/include/transport/rdma_transport/endpoint_store.h
@@ -43,7 +43,10 @@ class EndpointStore {
     virtual std::shared_ptr<RdmaEndPoint> insertEndpoint(
         const std::string &peer_nic_path, RdmaContext *context) = 0;
     virtual int deleteEndpoint(const std::string &peer_nic_path) = 0;
+    virtual int deleteEndpointByPtr(const RdmaEndPoint *endpoint_ptr) = 0;
     virtual void evictEndpoint() = 0;
+    // Takes endpoint_map_lock_; caller must not hold it (RWSpinlock is
+    // non-reentrant, so recursive acquisition deadlocks).
     virtual void reclaimEndpoint() = 0;
     virtual size_t getSize() = 0;
 
@@ -52,12 +55,21 @@ class EndpointStore {
 
     // Get the total number of QPs across all endpoints
     virtual size_t getTotalQPNumber() = 0;
+
+    // Number of endpoints awaiting reclaim (evicted or explicitly deleted but
+    // not yet destructed). Exposed for tests and for operator observability.
+    virtual size_t waitingListSize() const = 0;
+
+    // Test-only: push a pre-constructed endpoint into waiting_list_ so reclaim
+    // logic can be exercised without standing up an RDMA device.
+    virtual void testOnlyInsertWaiting(std::shared_ptr<RdmaEndPoint> ep) = 0;
 };
 
 // FIFO
 class FIFOEndpointStore : public EndpointStore {
    public:
-    FIFOEndpointStore(size_t max_size) : max_size_(max_size) {}
+    FIFOEndpointStore(size_t max_size)
+        : waiting_list_len_(0), max_size_(max_size) {}
     std::shared_ptr<RdmaEndPoint> getEndpoint(
         const std::string &peer_nic_path) override;
     std::shared_ptr<RdmaEndPoint> getEndpointByPtr(
@@ -65,6 +77,7 @@ class FIFOEndpointStore : public EndpointStore {
     std::shared_ptr<RdmaEndPoint> insertEndpoint(
         const std::string &peer_nic_path, RdmaContext *context) override;
     int deleteEndpoint(const std::string &peer_nic_path) override;
+    int deleteEndpointByPtr(const RdmaEndPoint *endpoint_ptr) override;
     void evictEndpoint() override;
     void reclaimEndpoint() override;
     size_t getSize() override;
@@ -73,6 +86,11 @@ class FIFOEndpointStore : public EndpointStore {
     int disconnectQPs() override;
 
     size_t getTotalQPNumber() override;
+    size_t waitingListSize() const override {
+        return waiting_list_len_.load(std::memory_order_relaxed);
+    }
+
+    void testOnlyInsertWaiting(std::shared_ptr<RdmaEndPoint> ep) override;
 
    private:
     RWSpinlock endpoint_map_lock_;
@@ -82,6 +100,7 @@ class FIFOEndpointStore : public EndpointStore {
     std::list<std::string> fifo_list_;
 
     std::unordered_set<std::shared_ptr<RdmaEndPoint>> waiting_list_;
+    std::atomic<size_t> waiting_list_len_;
 
     size_t max_size_;
 };
@@ -98,6 +117,7 @@ class SIEVEEndpointStore : public EndpointStore {
     std::shared_ptr<RdmaEndPoint> insertEndpoint(
         const std::string &peer_nic_path, RdmaContext *context) override;
     int deleteEndpoint(const std::string &peer_nic_path) override;
+    int deleteEndpointByPtr(const RdmaEndPoint *endpoint_ptr) override;
     void evictEndpoint() override;
     void reclaimEndpoint() override;
     size_t getSize() override;
@@ -106,6 +126,11 @@ class SIEVEEndpointStore : public EndpointStore {
     int disconnectQPs() override;
 
     size_t getTotalQPNumber() override;
+    size_t waitingListSize() const override {
+        return waiting_list_len_.load(std::memory_order_relaxed);
+    }
+
+    void testOnlyInsertWaiting(std::shared_ptr<RdmaEndPoint> ep) override;
 
    private:
     RWSpinlock endpoint_map_lock_;
@@ -119,7 +144,7 @@ class SIEVEEndpointStore : public EndpointStore {
     std::optional<std::list<std::string>::iterator> hand_;
 
     std::unordered_set<std::shared_ptr<RdmaEndPoint>> waiting_list_;
-    std::atomic<int> waiting_list_len_;
+    std::atomic<size_t> waiting_list_len_;
 
     size_t max_size_;
 };
diff --git a/mooncake-transfer-engine/include/transport/rdma_transport/rdma_context.h b/mooncake-transfer-engine/include/transport/rdma_transport/rdma_context.h
index 825de6a1..162b7fc9 100644
--- a/mooncake-transfer-engine/include/transport/rdma_transport/rdma_context.h
+++ b/mooncake-transfer-engine/include/transport/rdma_transport/rdma_context.h
@@ -19,16 +19,23 @@
 #include <glog/logging.h>
 #include <infiniband/verbs.h>
 
+#ifdef USE_MLX5DV
+#include <infiniband/mlx5dv.h>
+#endif
+
 #include <atomic>
 #include <condition_variable>
 #include <cstdint>
 #include <list>
+#include <map>
 #include <memory>
+#include <mutex>
 #include <string>
 #include <thread>
 #include <unordered_map>
 
 #include "common.h"
+#include "rdma_gid_probe.h"
 #include "rdma_transport.h"
 #include "transport/transport.h"
 
@@ -36,6 +43,7 @@ namespace mooncake {
 
 class RdmaEndPoint;
 class RdmaTransport;
+class RdmaContextTestPeer;
 class WorkerPool;
 class EndpointStore;
 
@@ -46,6 +54,11 @@ enum class GidNetworkState {
     GID_NOT_FOUND = 2         // No suitable GID found
 };
 
+struct GidSelectionSnapshot {
+    std::string gid;
+    int gid_index = -1;
+};
+
 struct RdmaCq {
     RdmaCq() : native(nullptr), outstanding(0) {}
     ibv_cq *native;
@@ -63,6 +76,8 @@ struct MemoryRegionMeta {
 // including Memory Region, CQ, EndPoint (QPs), etc.
 class RdmaContext {
    public:
+    friend class RdmaContextTestPeer;
+
     RdmaContext(RdmaTransport &engine, const std::string &device_name);
 
     ~RdmaContext();
@@ -90,6 +105,13 @@ class RdmaContext {
     int registerMemoryRegionInternal(void *addr, size_t length, int access,
                                      MemoryRegionMeta &mrMeta);
 
+    using MemoryRegionMap = std::map<uintptr_t, MemoryRegionMeta>;
+
+    MemoryRegionMap::iterator findMemoryRegionContaining(uintptr_t addr);
+
+    MemoryRegionMap::const_iterator findMemoryRegionContaining(
+        uintptr_t addr) const;
+
    public:
     bool active() const { return active_; }
 
@@ -103,6 +125,22 @@ class RdmaContext {
         const RdmaEndPoint *endpoint_ptr);
 
     int deleteEndpoint(const std::string &peer_nic_path);
+    int deleteEndpointByPtr(const RdmaEndPoint *endpoint_ptr);
+
+    // Drain the endpoint store's waiting list. Safe to call on any thread;
+    // intended to be invoked periodically from monitorWorker so reclaim is
+    // not gated on new endpoint insertions (which can stall under failure
+    // load while evictions/deletions continue). See issue #1845.
+    void reclaimEndpoints();
+
+    // Number of endpoints awaiting reclaim. For tests and operator
+    // observability.
+    size_t waitingListSize() const;
+
+    // Test-only: push a pre-constructed endpoint into the store's
+    // waiting_list_ so the reclaim path can be exercised without standing up
+    // a real RDMA QP.
+    void testOnlyInsertWaiting(std::shared_ptr<RdmaEndPoint> ep);
 
     int disconnectAllEndpoints();
 
@@ -121,7 +159,16 @@ class RdmaContext {
 
     std::string gid() const;
 
-    int gidIndex() const { return gid_index_; }
+    GidSelectionSnapshot gidSelection() const;
+
+    int gidIndex() const;
+
+    bool autoGidSelectionEnabled() const { return auto_gid_selection_enabled_; }
+
+    bool reprobeAutoGid(
+        const GidSelectionSnapshot &expected_selection,
+        const std::vector<AutoGidSelectionIdentity> &tried_selections = {},
+        std::string *previous_gid = nullptr, std::string *next_gid = nullptr);
 
     ibv_context *context() const { return context_; }
 
@@ -131,6 +178,8 @@ class RdmaContext {
 
     uint8_t portNum() const { return port_; }
 
+    uint8_t numLagPorts() const { return num_lag_ports_; }
+
     int activeSpeed() const { return active_speed_; }
 
     ibv_mtu activeMTU() const { return active_mtu_; }
@@ -184,10 +233,14 @@ class RdmaContext {
     int gid_index_ = -1;
     int active_speed_ = -1;
     ibv_mtu active_mtu_;
+    uint8_t num_lag_ports_ = 0;  // 0/1 = not in LAG; ≥2 = LAG active
     ibv_gid gid_;
+    mutable std::mutex gid_lock_;
+    mutable std::mutex gid_reprobe_lock_;
+    bool auto_gid_selection_enabled_ = false;
 
     RWSpinlock memory_regions_lock_;
-    std::vector<struct MemoryRegionMeta> memory_region_list_;
+    MemoryRegionMap memory_region_map_;
     std::vector<RdmaCq> cq_list_;
 
     std::shared_ptr<EndpointStore> endpoint_store_;
diff --git a/mooncake-transfer-engine/include/transport/rdma_transport/rdma_endpoint.h b/mooncake-transfer-engine/include/transport/rdma_transport/rdma_endpoint.h
index 042d8661..e98fa8e7 100644
--- a/mooncake-transfer-engine/include/transport/rdma_transport/rdma_endpoint.h
+++ b/mooncake-transfer-engine/include/transport/rdma_transport/rdma_endpoint.h
@@ -43,6 +43,8 @@ class RdmaEndPoint {
         UNCONNECTED,
         CONNECTING,
         CONNECTED,
+        DESTROYING,
+        DESTROYED,
     };
 
    public:
@@ -56,6 +58,8 @@ class RdmaEndPoint {
    private:
     int reconstruct();
     int deconstruct();
+    int deconstructLocked();
+    void beginDestroyLocked();
 
    public:
     void setPeerNicPath(const std::string &peer_nic_path);
@@ -71,8 +75,6 @@ class RdmaEndPoint {
     int setupConnectionsByPassive(const HandShakeDesc &peer_desc,
                                   HandShakeDesc &local_desc);
 
-    bool hasOutstandingSlice() const;
-
     bool active() const { return active_; }
 
     void set_active(bool flag) {
@@ -91,6 +93,11 @@ class RdmaEndPoint {
         return status_.load(std::memory_order_relaxed) == CONNECTED;
     }
 
+    bool retired() const {
+        auto status = status_.load(std::memory_order_relaxed);
+        return status == DESTROYING || status == DESTROYED;
+    }
+
     // Interrupts the connection, which can be triggered by user or by internal
     // error. Use setupConnectionsByActive or setupConnectionsByPassive to
     // reconnect
@@ -99,28 +106,21 @@ class RdmaEndPoint {
     // Destroy QPs before CQs (in RDMA Context)
     int destroyQP();
 
+    // Two-phase QP destruction to avoid use-after-free in concurrent
+    // submitPostSend. Phase 1 (beginDestroy): sets active_=false and
+    // status_=DESTROYING, transitions QPs to ERR state so hardware flushes
+    // inflight WRs to CQ. Does not block. Phase 2 (finishDestroy): called
+    // after all outstanding WRs have been drained (wr_depth_list_ all zero),
+    // actually destroys QPs and frees resources. Returns true if destruction
+    // is complete, false if outstanding WRs remain.
+    void beginDestroy();
+    bool finishDestroy();
+
    private:
     int disconnectUnlocked();
 
-    // Resets the connection.
-    //
-    // The main difference between this function and `disconnectUnlocked`
-    // is that it will reconstruct QPs when `CONFIG_ERDMA` is defined.
-    // Without `CONFIG_ERDMA`, it is essentially the same as
-    // `disconnectUnlocked` but with additional logging.
-    //
-    // This serves as a workaround for Aliyun eRDMA devices (i.e., once a QP is
-    // transitioned to the RTS state, it cannot be reset to RTS again directly).
-    // For more details:
-    // https://github.com/kvcache-ai/Mooncake/pull/1733#discussion_r2992088663
-    //
-    // In practice:
-    // - Call `resetConnection` if the QPs' state may have transitioned to RTS.
-    // - Call `disconnectUnlocked` otherwise.
-    //
-    // This is mainly used in `setupConnectionsByActive` or
-    // `setupConnectionsByPassive`. It is NOT invoked in the normal execution
-    // flow, so a `reason` argument is passed for internal logging purposes.
+    // Resets only pre-connected handshake attempts. Once an endpoint has ever
+    // reached CONNECTED, it is retired instead of being reused.
     int resetConnection(const std::string &reason);
 
    public:
@@ -137,15 +137,31 @@ class RdmaEndPoint {
     size_t getQPNumber() const;
 
    private:
+    enum class SetupConnectionFailureStage {
+        kNone,
+        kPeerValidation,
+        kReset,
+        kInit,
+        kRtr,
+        kRts,
+    };
+
+    struct SetupConnectionFailureInfo {
+        SetupConnectionFailureStage stage = SetupConnectionFailureStage::kNone;
+        int sys_errno = 0;
+    };
+
     std::vector<uint32_t> qpNum() const;
 
     int doSetupConnection(const std::string &peer_gid, uint16_t peer_lid,
                           std::vector<uint32_t> peer_qp_num_list,
-                          std::string *reply_msg = nullptr);
+                          std::string *reply_msg = nullptr,
+                          SetupConnectionFailureInfo *failure_info = nullptr);
 
-    int doSetupConnection(int qp_index, const std::string &peer_gid,
+    int doSetupConnection(int qp_index, const ibv_gid &peer_gid,
                           uint16_t peer_lid, uint32_t peer_qp_num,
-                          std::string *reply_msg = nullptr);
+                          int local_gid_index, std::string *reply_msg = nullptr,
+                          SetupConnectionFailureInfo *failure_info = nullptr);
 
    private:
     static constexpr uint64_t kWaitExistingHandshakeTimeoutNano =
@@ -154,6 +170,16 @@ class RdmaEndPoint {
     static constexpr uint32_t kWaitExistingHandshakeInitialSleepUs = 50;
     static constexpr uint32_t kWaitExistingHandshakeMaxSleepUs = 2000;
 
+    // Maximum time (in seconds) to wait for outstanding WRs to drain in
+    // finishDestroy before forcing QP destruction. This guards against
+    // ibv_modify_qp-to-ERR failures that prevent WR flushing.
+    static constexpr double kFinishDestroyTimeoutSec = 30.0;
+
+    // Maximum number of deconstructLocked retries in finishDestroy before
+    // giving up and marking the endpoint as DESTROYED. Prevents infinite
+    // retry loops and log flooding when ibv_destroy_qp fails permanently.
+    static constexpr int kFinishDestroyMaxRetries = 3;
+
     RdmaContext &context_;
     std::atomic<Status> status_;
 
@@ -162,6 +188,7 @@ class RdmaEndPoint {
 
     std::string peer_nic_path_;
     std::vector<uint32_t> peer_qp_num_list_;
+    bool has_connected_;
 
     volatile int *wr_depth_list_;
     int max_wr_depth_;
@@ -171,6 +198,7 @@ class RdmaEndPoint {
     volatile bool active_;
     volatile int *cq_outstanding_;
     volatile uint64_t inactive_time_;
+    int finish_destroy_retries_ = 0;
 };
 
 }  // namespace mooncake
diff --git a/mooncake-transfer-engine/include/transport/rdma_transport/rdma_gid_probe.h b/mooncake-transfer-engine/include/transport/rdma_transport/rdma_gid_probe.h
new file mode 100644
index 00000000..a4ee60db
--- /dev/null
+++ b/mooncake-transfer-engine/include/transport/rdma_transport/rdma_gid_probe.h
@@ -0,0 +1,247 @@
+// Copyright 2026 KVCache.AI
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#ifndef RDMA_GID_PROBE_H
+#define RDMA_GID_PROBE_H
+
+#include <cerrno>
+#include <infiniband/verbs.h>
+
+#include <algorithm>
+#include <optional>
+#include <string>
+#include <string_view>
+#include <vector>
+
+namespace mooncake {
+
+enum class AutoGidCandidateClass {
+    kNetworkRoutable = 0,
+    kNoNetworkRoutable = 1,
+    kNetworkDegraded = 2,
+    kNoNetworkDegraded = 3,
+    kFallbackNonzero = 4,
+};
+
+enum class AutoGidRetryAction {
+    kDoNotRetry = 0,
+    kRetryWithReprobedGid = 1,
+    kRetryWithObservedChange = 2,
+};
+
+struct AutoGidCandidate {
+    int gid_index = -1;
+    std::string gid;
+    uint32_t gid_type = 0;
+    bool has_network_device = false;
+    bool is_ipv4_mapped = false;
+    bool is_link_local_ipv6 = false;
+    bool is_overlay_network = false;
+    bool is_overlay_ipv4 = false;
+    bool is_null_gid = false;
+    bool query_succeeded = true;
+};
+
+struct AutoGidSelection {
+    int gid_index = -1;
+    std::string gid;
+    AutoGidCandidateClass candidate_class =
+        AutoGidCandidateClass::kFallbackNonzero;
+};
+
+struct AutoGidSelectionIdentity {
+    int gid_index = -1;
+    std::string gid;
+};
+
+inline const char* autoGidCandidateClassToString(
+    AutoGidCandidateClass candidate_class) {
+    switch (candidate_class) {
+        case AutoGidCandidateClass::kNetworkRoutable:
+            return "network-routable";
+        case AutoGidCandidateClass::kNoNetworkRoutable:
+            return "no-network-routable";
+        case AutoGidCandidateClass::kNetworkDegraded:
+            return "network-degraded";
+        case AutoGidCandidateClass::kNoNetworkDegraded:
+            return "no-network-degraded";
+        case AutoGidCandidateClass::kFallbackNonzero:
+            return "fallback-nonzero";
+    }
+    return "unknown";
+}
+
+inline std::optional<AutoGidCandidateClass> classifyAutoGidCandidate(
+    const AutoGidCandidate& candidate) {
+    if (!candidate.query_succeeded || candidate.gid_index < 0 ||
+        candidate.is_null_gid) {
+        return std::nullopt;
+    }
+
+    if (candidate.gid_type != IBV_GID_TYPE_ROCE_V2 &&
+        candidate.gid_type != IBV_GID_TYPE_IB) {
+        return std::nullopt;
+    }
+
+    const bool is_roce_v2 = candidate.gid_type == IBV_GID_TYPE_ROCE_V2;
+    const bool is_overlay =
+        is_roce_v2 && (candidate.is_overlay_network ||
+                       (candidate.is_ipv4_mapped && candidate.is_overlay_ipv4));
+    const bool is_link_local =
+        is_roce_v2 && !candidate.is_ipv4_mapped && candidate.is_link_local_ipv6;
+    const bool is_degraded = is_overlay || is_link_local;
+
+    if (candidate.has_network_device) {
+        return is_degraded ? AutoGidCandidateClass::kNetworkDegraded
+                           : AutoGidCandidateClass::kNetworkRoutable;
+    }
+
+    return is_degraded ? AutoGidCandidateClass::kNoNetworkDegraded
+                       : AutoGidCandidateClass::kNoNetworkRoutable;
+}
+
+inline int autoGidCandidateClassPriority(
+    AutoGidCandidateClass candidate_class) {
+    switch (candidate_class) {
+        case AutoGidCandidateClass::kNetworkRoutable:
+            return 0;
+        case AutoGidCandidateClass::kNoNetworkRoutable:
+            return 1;
+        case AutoGidCandidateClass::kNetworkDegraded:
+            return 2;
+        case AutoGidCandidateClass::kNoNetworkDegraded:
+            return 3;
+        case AutoGidCandidateClass::kFallbackNonzero:
+            return 4;
+    }
+    return 5;
+}
+
+inline std::vector<AutoGidSelection> rankAutoGidCandidates(
+    const std::vector<AutoGidCandidate>& candidates) {
+    std::vector<AutoGidSelection> ranked;
+    int first_query_success_fallback = -1;
+    std::string first_query_success_fallback_gid;
+
+    for (const auto& candidate : candidates) {
+        auto candidate_class = classifyAutoGidCandidate(candidate);
+        if (candidate_class.has_value()) {
+            ranked.push_back(
+                {candidate.gid_index, candidate.gid, *candidate_class});
+            continue;
+        }
+
+        if (candidate.query_succeeded && !candidate.is_null_gid &&
+            first_query_success_fallback < 0) {
+            first_query_success_fallback = candidate.gid_index;
+            first_query_success_fallback_gid = candidate.gid;
+        }
+    }
+
+    std::stable_sort(
+        ranked.begin(), ranked.end(),
+        [](const AutoGidSelection& lhs, const AutoGidSelection& rhs) {
+            int lhs_priority =
+                autoGidCandidateClassPriority(lhs.candidate_class);
+            int rhs_priority =
+                autoGidCandidateClassPriority(rhs.candidate_class);
+            if (lhs_priority != rhs_priority) {
+                return lhs_priority < rhs_priority;
+            }
+            return lhs.gid_index < rhs.gid_index;
+        });
+
+    if (first_query_success_fallback >= 0) {
+        ranked.push_back({first_query_success_fallback,
+                          first_query_success_fallback_gid,
+                          AutoGidCandidateClass::kFallbackNonzero});
+    }
+
+    return ranked;
+}
+
+inline std::optional<AutoGidSelection> selectBestAutoGidCandidate(
+    const std::vector<AutoGidCandidate>& candidates) {
+    auto ranked = rankAutoGidCandidates(candidates);
+    if (ranked.empty()) {
+        return std::nullopt;
+    }
+    return ranked.front();
+}
+
+inline bool shouldAttemptAutoGidHandshakeRetry(bool auto_gid_selection_enabled,
+                                               int retry_count, int max_retries,
+                                               bool failure_happened_at_rtr,
+                                               int sys_errno) {
+    return auto_gid_selection_enabled && retry_count < max_retries &&
+           failure_happened_at_rtr && sys_errno == EINVAL;
+}
+
+inline bool didAutoGidSelectionChange(int previous_gid_index,
+                                      std::string_view previous_gid,
+                                      int current_gid_index,
+                                      std::string_view current_gid) {
+    return previous_gid_index != current_gid_index ||
+           previous_gid != current_gid;
+}
+
+inline bool matchesAutoGidSelection(const AutoGidSelectionIdentity& identity,
+                                    const AutoGidSelection& selection) {
+    return identity.gid_index == selection.gid_index &&
+           identity.gid == selection.gid;
+}
+
+inline bool hasTriedAutoGidSelection(
+    const std::vector<AutoGidSelectionIdentity>& tried_selections,
+    const AutoGidSelection& selection) {
+    return std::any_of(tried_selections.begin(), tried_selections.end(),
+                       [&](const AutoGidSelectionIdentity& identity) {
+                           return matchesAutoGidSelection(identity, selection);
+                       });
+}
+
+inline std::optional<AutoGidSelection> reselectAutoGidCandidate(
+    const std::vector<AutoGidCandidate>& candidates, int current_gid_index,
+    std::string_view current_gid,
+    const std::vector<AutoGidSelectionIdentity>& tried_selections = {}) {
+    auto ranked = rankAutoGidCandidates(candidates);
+    for (const auto& selection : ranked) {
+        if (!didAutoGidSelectionChange(current_gid_index, current_gid,
+                                       selection.gid_index, selection.gid)) {
+            continue;
+        }
+        if (hasTriedAutoGidSelection(tried_selections, selection)) {
+            continue;
+        }
+        return selection;
+    }
+    return std::nullopt;
+}
+
+inline AutoGidRetryAction decideAutoGidRetryAction(
+    bool reprobe_changed, int previous_gid_index, std::string_view previous_gid,
+    int current_gid_index, std::string_view current_gid) {
+    if (reprobe_changed) {
+        return AutoGidRetryAction::kRetryWithReprobedGid;
+    }
+    if (didAutoGidSelectionChange(previous_gid_index, previous_gid,
+                                  current_gid_index, current_gid)) {
+        return AutoGidRetryAction::kRetryWithObservedChange;
+    }
+    return AutoGidRetryAction::kDoNotRetry;
+}
+
+}  // namespace mooncake
+
+#endif  // RDMA_GID_PROBE_H
diff --git a/mooncake-transfer-engine/include/transport/rdma_transport/rdma_transport.h b/mooncake-transfer-engine/include/transport/rdma_transport/rdma_transport.h
index d292a27e..e7a171b8 100644
--- a/mooncake-transfer-engine/include/transport/rdma_transport/rdma_transport.h
+++ b/mooncake-transfer-engine/include/transport/rdma_transport/rdma_transport.h
@@ -36,11 +36,13 @@ namespace mooncake {
 class RdmaContext;
 class RdmaEndPoint;
 class TransferMetadata;
+class RdmaTransportTestPeer;
 class WorkerPool;
 
 class RdmaTransport : public Transport {
     friend class RdmaContext;
     friend class RdmaEndPoint;
+    friend class RdmaTransportTestPeer;
     friend class WorkerPool;
 
    public:
@@ -101,6 +103,9 @@ class RdmaTransport : public Transport {
    private:
     int allocateLocalSegmentID();
 
+    int refreshLocalDeviceDesc(const std::string &device_name, uint16_t lid,
+                               const std::string &gid);
+
     int preTouchMemory(void *addr, size_t length);
 
    public:
@@ -125,10 +130,20 @@ class RdmaTransport : public Transport {
     static int selectDevice(SegmentDesc *desc, uint64_t offset, size_t length,
                             std::string_view hint, int &buffer_id,
                             int &device_id, int retry_cnt = 0);
+    static int selectDeviceByLocalHca(SegmentDesc *desc, uint64_t offset,
+                                      size_t length, std::string_view local_hca,
+                                      int &buffer_id, int &device_id,
+                                      int retry_cnt = 0);
 
    private:
     std::vector<std::shared_ptr<RdmaContext>> context_list_;
     std::shared_ptr<Topology> local_topology_;
+    // When MC_RDMA_BIND_ADDRESS is set in a dual-NIC environment,
+    // rdma_server_name_ holds the RDMA-reachable address (e.g.
+    // "192.168.0.y:port") for NIC path construction, while
+    // local_server_name_ keeps the TCP-reachable address for P2P routing.
+    std::string rdma_server_name_;
+    std::mutex local_desc_lock_;
 };
 
 using TransferRequest = Transport::TransferRequest;
@@ -139,4 +154,4 @@ using BatchID = Transport::BatchID;
 
 }  // namespace mooncake
 
-#endif  // RDMA_TRANSPORT_H_
\ No newline at end of file
+#endif  // RDMA_TRANSPORT_H_
diff --git a/mooncake-transfer-engine/include/transport/rdma_transport/worker_pool.h b/mooncake-transfer-engine/include/transport/rdma_transport/worker_pool.h
index f6b24248..e8ada2a7 100644
--- a/mooncake-transfer-engine/include/transport/rdma_transport/worker_pool.h
+++ b/mooncake-transfer-engine/include/transport/rdma_transport/worker_pool.h
@@ -18,6 +18,7 @@
 #include <queue>
 #include <unordered_set>
 
+#include "config.h"
 #include "rdma_context.h"
 
 namespace mooncake {
@@ -39,18 +40,56 @@ class WorkerPool {
 
     void transferWorker(int thread_id);
 
+    bool hasOutstandingCq(int thread_id);
+
     void monitorWorker();
 
     int doProcessContextEvents();
 
+    // Simplified rail monitor: pause problematic paths for a cooldown period
+    struct RailState {
+        int error_count = 0;
+        uint64_t pause_until_ns = 0;  // Timestamp (ns) when pause expires
+    };
+
+    void markRailFailed(const std::string &peer_nic_path);
+    bool isRailAvailable(const std::string &peer_nic_path);
+
+    // Retry helper: increment retry count and return whether retry is allowed
+    static bool shouldRetrySlice(Transport::Slice *slice);
+
+    // Unified path failure handler: marks rail failed, notifies other workers,
+    // and optionally deletes the endpoint
+    void handlePathFailure(const std::string &peer_nic_path,
+                           RdmaEndPoint *endpoint = nullptr);
+
+    // Context-level health tracking for catastrophic hardware failure.
+    // When all rails through a local RNIC are unavailable, increment the
+    // failure counter. Reset on any success. Mark context inactive after
+    // consecutive failures exceed threshold.
+    bool contextHealthy() const {
+        return context_failure_count_ < kContextFailureThreshold;
+    }
+    void markContextSuccess() { context_failure_count_ = 0; }
+    void markContextFailure() {
+        context_failure_count_++;
+        if (context_failure_count_ >= kContextFailureThreshold) {
+            LOG(WARNING) << "All rails failed for context "
+                         << context_.deviceName() << " for "
+                         << context_failure_count_
+                         << " consecutive attempts, marking inactive";
+            context_.set_active(false);
+        }
+    }
+
    private:
     RdmaContext &context_;
     const int numa_socket_id_;
 
     std::vector<std::thread> worker_thread_;
     std::atomic<bool> workers_running_;
-    std::atomic<int> suspended_flag_;
 
+    std::atomic<int> parked_worker_count_;
     std::atomic<int> redispatch_counter_;
 
     std::mutex cond_mutex_;
@@ -68,7 +107,18 @@ class WorkerPool {
 
     std::atomic<uint64_t> submitted_slice_count_, processed_slice_count_;
 
-    uint64_t success_nr_polls = 0, failed_nr_polls = 0;
+    // Rail state management: peer_nic_path -> RailState
+    std::unordered_map<std::string, RailState> rail_states_;
+    std::mutex rail_state_lock_;
+
+    // Rail monitor configuration
+    const static int kRailErrorThreshold = 5;            // Errors before pause
+    const static uint64_t kRailPauseNs = 1000000000ull;  // 1 second pause
+
+    // Context-level health tracking
+    int context_failure_count_ = 0;
+    const static int kContextFailureThreshold =
+        32;  // consecutive all-rails-failed
 };
 }  // namespace mooncake
 
diff --git a/mooncake-transfer-engine/include/transport/sunrise_link_transport/sunrise_link_transport.h b/mooncake-transfer-engine/include/transport/sunrise_link_transport/sunrise_link_transport.h
new file mode 100644
index 00000000..321384ec
--- /dev/null
+++ b/mooncake-transfer-engine/include/transport/sunrise_link_transport/sunrise_link_transport.h
@@ -0,0 +1,92 @@
+// Copyright 2024 KVCache.AI
+
+#ifndef SUNRISE_LINK_TRANSPORT_H_
+#define SUNRISE_LINK_TRANSPORT_H_
+
+#include <mutex>
+#include <unordered_map>
+#include <utility>
+#include <vector>
+
+#include "common/hash_utils.h"
+#include "topology.h"
+#include "transfer_metadata.h"
+#include "transport/transport.h"
+
+namespace mooncake {
+
+class TransferMetadata;
+
+class SunriseLinkTransport : public Transport {
+   public:
+    SunriseLinkTransport();
+
+    ~SunriseLinkTransport() override;
+
+    Status submitTransfer(BatchID batch_id,
+                          const std::vector<TransferRequest>& entries) override;
+
+    Status submitTransferTask(
+        const std::vector<TransferTask*>& task_list) override;
+
+    Status getTransferStatus(BatchID batch_id, size_t task_id,
+                             TransferStatus& status) override;
+
+    static int parseDeviceIdForTest(const std::string& location);
+    static bool requestRangeFitsInBufferForTest(uint64_t dest_addr,
+                                                uint64_t length,
+                                                uint64_t buffer_addr,
+                                                uint64_t buffer_length);
+
+   protected:
+    int install(std::string& local_server_name,
+                std::shared_ptr<TransferMetadata> meta,
+                std::shared_ptr<Topology> topo) override;
+
+    int registerLocalMemory(void* addr, size_t length,
+                            const std::string& location, bool remote_accessible,
+                            bool update_metadata = true) override;
+
+    int unregisterLocalMemory(void* addr, bool update_metadata = true) override;
+
+    int registerLocalMemoryBatch(const std::vector<BufferEntry>& buffer_list,
+                                 const std::string& location) override;
+
+    int unregisterLocalMemoryBatch(
+        const std::vector<void*>& addr_list) override;
+
+    const char* getName() const override { return "sunrise_link"; }
+
+   private:
+    struct OpenedShmEntry {
+        void* shm_addr;
+        uint64_t length;
+        int gpu_id;
+        bool is_raw_addr = false;
+    };
+
+    int relocateSharedMemoryAddress(uint64_t& dest_addr, uint64_t length,
+                                    uint64_t target_id, int* target_gpu_id);
+    int submitRequest(TransferTask& task, const TransferRequest& request,
+                      bool keep_slice_ref);
+    int startCopy(void* src, void* dst, size_t length, int remote_dev,
+                  int local_dev);
+    int inferRegisteredDevice(void* ptr) const;
+    int inferPointerDeviceBestEffort(void* ptr, int preferred_dev) const;
+    bool loadRuntime();
+
+    mutable std::mutex register_mutex_;
+    struct RegisteredRegion {
+        size_t length;
+        int gpu_id;
+    };
+    std::unordered_map<void*, RegisteredRegion> registered_regions_;
+
+    std::unordered_map<std::pair<uint64_t, uint64_t>, OpenedShmEntry, PairHash>
+        remap_entries_;
+    RWSpinlock remap_lock_;
+};
+
+}  // namespace mooncake
+
+#endif  // SUNRISE_LINK_TRANSPORT_H_
diff --git a/mooncake-transfer-engine/include/transport/tcp_transport/tcp_transport.h b/mooncake-transfer-engine/include/transport/tcp_transport/tcp_transport.h
index 94dc2d9d..3e537cfe 100644
--- a/mooncake-transfer-engine/include/transport/tcp_transport/tcp_transport.h
+++ b/mooncake-transfer-engine/include/transport/tcp_transport/tcp_transport.h
@@ -104,6 +104,8 @@ class TcpTransport : public Transport {
 
     void startTransfer(Slice *slice);
 
+    bool validateAddress(uint64_t addr, uint64_t size) const;
+
     const char *getName() const override { return "tcp"; }
 
    private:
diff --git a/mooncake-transfer-engine/include/transport/transport.h b/mooncake-transfer-engine/include/transport/transport.h
index e3a22c7e..532eebc6 100644
--- a/mooncake-transfer-engine/include/transport/transport.h
+++ b/mooncake-transfer-engine/include/transport/transport.h
@@ -35,6 +35,8 @@
 #include "transfer_metadata.h"
 
 namespace mooncake {
+
+class RdmaEndPoint;
 class TransferMetadata;
 /// By default, these functions return 0 (or non-null pointer) on success and
 /// return -1 (or null pointer) on failure. The errno is set accordingly on
@@ -64,6 +66,8 @@ class Transport {
         uint64_t target_offset;
         size_t length;
         int advise_retry_cnt = 0;
+        // Per-request transport pin, TENT only.
+        int transport_hint = 0;
     };
 
     enum TransferStatusEnum {
@@ -109,21 +113,30 @@ class Transport {
         TransferRequest::OpCode opcode;
         SegmentID target_id;
         std::string peer_nic_path;
+        std::string source_location;
         SliceStatus status;
         TransferTask *task;
-        std::vector<uint32_t> dest_rkeys;
+        // EFA/CXI's libfabric MR keys are 64-bit (fi_mr_key()); RDMA verbs keys
+        // are 32-bit. Use a scoped alias so the width is defined in one place.
+#if defined(USE_EFA) || defined(USE_CXI)
+        using mr_key_t = uint64_t;
+#else
+        using mr_key_t = uint32_t;
+#endif
+        std::vector<mr_key_t> dest_rkeys;
         bool from_cache;
 
         union {
             struct {
                 uint64_t dest_addr;
-                uint32_t source_lkey;
-                uint32_t dest_rkey;
+                mr_key_t source_lkey;
+                mr_key_t dest_rkey;
                 int lkey_index;
                 int rkey_index;
                 volatile int *qp_depth;
                 uint32_t retry_cnt;
                 uint32_t max_retry_cnt;
+                RdmaEndPoint *endpoint;  // Endpoint used for this transfer
             } rdma;
             struct {
                 uint64_t dest_addr;
@@ -132,9 +145,12 @@ class Transport {
                 uint32_t max_retry_cnt;
                 void *r_seg;
                 void *l_seg;
+                void *endpoint;
             } ub;
             struct {
                 void *dest_addr;
+                void *cuda_stream;  // cudaStream_t, used by async NVLink
+                                    // transport
             } local;
             struct {
                 uint64_t dest_addr;
@@ -287,6 +303,12 @@ class Transport {
         uint64_t total_bytes = 0;
         BatchID batch_id = 0;
 
+        // Pointer to the transport that handles this task, set by
+        // MultiTransport::submitTransfer(). Used to delegate
+        // transport-specific completion polling (e.g., CUDA stream
+        // query for NVLink async transfers) in getTransferStatus().
+        Transport *transport_ = nullptr;
+
 #ifdef WITH_METRICS
         std::chrono::steady_clock::time_point start_time;
 #endif
diff --git a/mooncake-transfer-engine/include/ub_allocator.h b/mooncake-transfer-engine/include/ub_allocator.h
new file mode 100644
index 00000000..a753f421
--- /dev/null
+++ b/mooncake-transfer-engine/include/ub_allocator.h
@@ -0,0 +1,11 @@
+#pragma once
+
+namespace mooncake {
+
+void* ub_allocate_memory(size_t alignment, size_t total_size);
+
+void ub_free_memory(void* ptr);
+
+bool ub_is_store_memory(void* addr, size_t length);
+
+}  // namespace mooncake
\ No newline at end of file
diff --git a/mooncake-transfer-engine/nvlink-allocator/CMakeLists.txt b/mooncake-transfer-engine/nvlink-allocator/CMakeLists.txt
index d2b4e7c8..69e0d714 100644
--- a/mooncake-transfer-engine/nvlink-allocator/CMakeLists.txt
+++ b/mooncake-transfer-engine/nvlink-allocator/CMakeLists.txt
@@ -1,25 +1,30 @@
-# Build nvlink allocator and output to build directory
-
-add_custom_target(build_nvlink_allocator DEPENDS transfer_engine)
-
-get_target_property(INCLUDE_DIRS build_nvlink_allocator INCLUDE_DIRECTORIES)
-string(REPLACE ";" " " INCLUDE_DIRS_STR "${INCLUDE_DIRS}")
+include(${CMAKE_CURRENT_SOURCE_DIR}/../fabric_allocator.cmake)
 
 set(_extra_build_opts "")
-if (USE_HIP)
-    set(_extra_build_opts "--use-hipcc")
-elseif (USE_MUSA)
-    set(_extra_build_opts "--use-mcc")
+if(USE_HIP)
+  list(APPEND _extra_build_opts --use-hipcc)
+elseif(USE_MUSA)
+  list(APPEND _extra_build_opts --use-mcc)
+elseif(USE_MACA)
+  list(APPEND _extra_build_opts --use-maca)
 endif()
 
-if (USE_CUDA OR USE_HIP OR USE_MUSA)
-    add_custom_command(
-        TARGET build_nvlink_allocator
-        COMMAND ${CMAKE_COMMAND} -E make_directory ${CMAKE_CURRENT_BINARY_DIR}
-        COMMAND bash build.sh ${_extra_build_opts} ${CMAKE_CURRENT_BINARY_DIR} "${INCLUDE_DIRS_STR}"
-        WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}
-        COMMENT "Building nvlink allocator to ${CMAKE_CURRENT_BINARY_DIR}"
-    )
+set(_enable_nvlink_allocator_build FALSE)
+if(USE_CUDA
+   OR USE_HIP
+   OR USE_MUSA
+   OR USE_MACA)
+  set(_enable_nvlink_allocator_build TRUE)
 endif()
 
-set_property(TARGET build_nvlink_allocator PROPERTY EXCLUDE_FROM_ALL FALSE)
+add_fabric_allocator_build_target(
+  TARGET_NAME
+  build_nvlink_allocator
+  BUILD_SCRIPT
+  ${CMAKE_CURRENT_SOURCE_DIR}/build.sh
+  BUILD_ARGS
+  ${_extra_build_opts}
+  COMMENT
+  "Building nvlink allocator to ${CMAKE_CURRENT_BINARY_DIR}"
+  ENABLE_BUILD
+  ${_enable_nvlink_allocator_build})
diff --git a/mooncake-transfer-engine/nvlink-allocator/build.sh b/mooncake-transfer-engine/nvlink-allocator/build.sh
index 5f342466..018e1133 100755
--- a/mooncake-transfer-engine/nvlink-allocator/build.sh
+++ b/mooncake-transfer-engine/nvlink-allocator/build.sh
@@ -2,10 +2,13 @@
 
 set -e
 
+source "$(dirname "$(readlink -f "$0")")/../scripts/allocator_build_common.sh"
+
 # Check for flags
 USE_NVCC=false
 USE_HIPCC=false
 USE_MCC=false
+USE_MACA=false
 CI_BUILD=false
 
 if [[ "$1" == "--use-nvcc" ]]; then
@@ -17,6 +20,9 @@ elif [[ "$1" == "--use-hipcc" ]]; then
 elif [[ "$1" == "--use-mcc" ]]; then
     USE_MCC=true
     shift
+elif [[ "$1" == "--use-maca" ]]; then
+    USE_MACA=true
+    shift
 elif [[ "$1" == "--ci-build" ]]; then
     CI_BUILD=true
     shift
@@ -25,27 +31,11 @@ fi
 # Get output directory from command line argument, default to current directory
 OUTPUT_DIR=${1:-.}
 
-# Get include directories from second argument (if provided)
-INCLUDE_LIST=""
-if [ $# -ge 2 ]; then
-    INCLUDE_LIST=${2}
-fi
-
-# Add include directory for cuda_alike.h (relative to build.sh location)
-SCRIPT_DIR=$(cd "$(dirname "${BASH_SOURCE[0]}")" &>/dev/null && pwd)
-INCLUDE_LIST="${INCLUDE_LIST:+${INCLUDE_LIST} }${SCRIPT_DIR}/../include"
-
-# Process include directories into flags
-INCLUDE_FLAGS=""
-if [ -n "$INCLUDE_LIST" ]; then
-    INCLUDE_FLAGS=$(echo "$INCLUDE_LIST" | tr ' ' '\n' | sed 's/^/-I/' | paste -sd' ' -)
-fi
+prepare_allocator_build_env "$OUTPUT_DIR" "${2:-}"
 
 echo "Building nvlink allocator to: $OUTPUT_DIR"
-# Create output directory if it doesn't exist
-mkdir -p "$OUTPUT_DIR"
 
-CPP_FILE=$(dirname $(readlink -f $0))/nvlink_allocator.cpp  # get cpp file path, under same dir with this script
+CPP_FILE="${SCRIPT_DIR}/nvlink_allocator.cpp"
 
 # Choose build command based on flags
 if [ "$CI_BUILD" = true ]; then
@@ -59,6 +49,11 @@ elif [ "$USE_HIPCC" = true ]; then
     hipcc "$OUTPUT_DIR/nvlink_allocator.cpp" -o "$OUTPUT_DIR/nvlink_allocator.so" -shared -fPIC -lamdhip64 -I/opt/rocm/include ${INCLUDE_FLAGS} -DUSE_HIP=1
 elif [ "$USE_MCC" = true ]; then
     mcc "$CPP_FILE" -o "$OUTPUT_DIR/nvlink_allocator.so" --shared -fPIC -lmusa -I/usr/local/musa/include ${INCLUDE_FLAGS} -DUSE_MUSA=1
+elif [ "$USE_MACA" = true ]; then
+    MACA_ROOT=${MACA_HOME:-/opt/maca}
+    g++ "$CPP_FILE" -o "$OUTPUT_DIR/nvlink_allocator.so" --shared -fPIC \
+        -I"${MACA_ROOT}/include" ${INCLUDE_FLAGS} \
+        -L"${MACA_ROOT}/lib64" -L"${MACA_ROOT}/lib" -lmcruntime -DUSE_MACA=1
 else
     # Default g++ build
     g++ "$CPP_FILE" -o "$OUTPUT_DIR/nvlink_allocator.so" --shared -fPIC -lcuda -I/usr/local/cuda/include ${INCLUDE_FLAGS} -DUSE_CUDA=1
diff --git a/mooncake-transfer-engine/nvlink-allocator/nvlink_allocator.cpp b/mooncake-transfer-engine/nvlink-allocator/nvlink_allocator.cpp
index 78c5c290..3c347eea 100644
--- a/mooncake-transfer-engine/nvlink-allocator/nvlink_allocator.cpp
+++ b/mooncake-transfer-engine/nvlink-allocator/nvlink_allocator.cpp
@@ -19,16 +19,15 @@ static CUresult cuMemCreateTryFabric(CUmemGenericAllocationHandle *handle,
 
 enum class MemoryBackendType { use_cudamalloc, use_cumemcreate, unknown };
 
-extern "C" {
+namespace {
 
-MemoryBackendType mc_probe_fabric_support(int device_id) {
+MemoryBackendType ProbeAllocatorBackend(int device_id) {
     CUdevice dev;
     CUresult res = cuDeviceGet(&dev, device_id);
     if (res != CUDA_SUCCESS) {
         return MemoryBackendType::unknown;
     }
 
-    // Check device attribute first
     int fabric_attr = 0;
     res = cuDeviceGetAttribute(
         &fabric_attr, CU_DEVICE_ATTRIBUTE_HANDLE_TYPE_FABRIC_SUPPORTED, dev);
@@ -40,28 +39,26 @@ MemoryBackendType mc_probe_fabric_support(int device_id) {
     prop.type = CU_MEM_ALLOCATION_TYPE_PINNED;
     prop.location.type = CU_MEM_LOCATION_TYPE_DEVICE;
     prop.location.id = dev;
-    prop.requestedHandleTypes = CU_MEM_HANDLE_TYPE_FABRIC;  // require fabric
+    prop.requestedHandleTypes = CU_MEM_HANDLE_TYPE_FABRIC;
 
     CUmemGenericAllocationHandle handle;
     size_t size = 4096;
 
     res = cuMemCreate(&handle, size, &prop, 0);
-
     if (res == CUDA_SUCCESS) {
-        cuMemRelease(handle);  // success → clean up
+        cuMemRelease(handle);
         return MemoryBackendType::use_cumemcreate;
-    } else {
-        return MemoryBackendType::use_cudamalloc;
     }
+    return MemoryBackendType::use_cudamalloc;
 }
 
-void *mc_nvlink_malloc(ssize_t size, int device, cudaStream_t stream) {
+void *AllocateFabricMemory(ssize_t size, int device, cudaStream_t stream) {
+    (void)stream;
     size_t granularity = 0;
     CUdevice currentDev;
     CUmemAllocationProp prop = {};
     CUmemGenericAllocationHandle handle;
     void *ptr = nullptr;
-    int cudaDev;
     int flag = 0;
     CUresult result = cuDeviceGet(&currentDev, device);
     if (result != CUDA_SUCCESS) {
@@ -98,7 +95,6 @@ void *mc_nvlink_malloc(ssize_t size, int device, cudaStream_t stream) {
         std::cerr << "cuMemGetAllocationGranularity failed: " << result;
         return nullptr;
     }
-    // fix size
     size = (size + granularity - 1) & ~(granularity - 1);
     if (size == 0) size = granularity;
     result = cuMemCreateTryFabric(&handle, size, &prop, 0);
@@ -138,7 +134,11 @@ void *mc_nvlink_malloc(ssize_t size, int device, cudaStream_t stream) {
     return ptr;
 }
 
-void mc_nvlink_free(void *ptr, ssize_t ssize, int device, cudaStream_t stream) {
+void FreeFabricMemory(void *ptr, ssize_t ssize, int device,
+                      cudaStream_t stream) {
+    (void)ssize;
+    (void)device;
+    (void)stream;
     CUmemGenericAllocationHandle handle;
     size_t size = 0;
     if (!ptr) return;
@@ -154,4 +154,33 @@ void mc_nvlink_free(void *ptr, ssize_t ssize, int device, cudaStream_t stream) {
     }
     cuMemRelease(handle);
 }
+
+}  // namespace
+
+extern "C" {
+
+MemoryBackendType mc_probe_fabric_support(int device_id) {
+    return ProbeAllocatorBackend(device_id);
+}
+
+int mc_allocator_probe(int device_id) {
+    return static_cast<int>(ProbeAllocatorBackend(device_id));
+}
+
+void *mc_allocator_malloc(ssize_t size, int device, cudaStream_t stream) {
+    return AllocateFabricMemory(size, device, stream);
+}
+
+void *mc_nvlink_malloc(ssize_t size, int device, cudaStream_t stream) {
+    return mc_allocator_malloc(size, device, stream);
+}
+
+void mc_allocator_free(void *ptr, ssize_t ssize, int device,
+                       cudaStream_t stream) {
+    FreeFabricMemory(ptr, ssize, device, stream);
+}
+
+void mc_nvlink_free(void *ptr, ssize_t ssize, int device, cudaStream_t stream) {
+    mc_allocator_free(ptr, ssize, device, stream);
+}
 }
diff --git a/mooncake-transfer-engine/rust/src/transfer_engine.rs b/mooncake-transfer-engine/rust/src/transfer_engine.rs
index 241c18a0..8c6186df 100644
--- a/mooncake-transfer-engine/rust/src/transfer_engine.rs
+++ b/mooncake-transfer-engine/rust/src/transfer_engine.rs
@@ -150,6 +150,9 @@ impl TransferEngine {
         buffer_list: &[BufferEntry],
         location: &str,
     ) -> Result<()> {
+        if buffer_list.is_empty() {
+            return Ok(());
+        }
         let location_c = CString::new(location).map_err(|_| anyhow!("CString::new failed"))?;
         let mut buffer_list_c: Vec<bindings::buffer_entry_t> = vec![];
         let buffer_len_c = buffer_list.len();
@@ -175,6 +178,9 @@ impl TransferEngine {
     }
 
     pub fn unregister_local_memory_batch(&self, buffer_list: &[BufferEntry]) -> Result<()> {
+        if buffer_list.is_empty() {
+            return Ok(());
+        }
         let mut addr_list: Vec<*mut c_void> = buffer_list.iter().map(|entry| entry.addr).collect();
         let addr_len = buffer_list.len();
         let ret = unsafe {
@@ -201,6 +207,9 @@ impl TransferEngine {
         batch_id: BatchID,
         requests: &mut [TransferRequest],
     ) -> Result<()> {
+        if requests.is_empty() {
+            return Ok(());
+        }
         let mut requests_c: Vec<bindings::transfer_request_t> = vec![];
         for i in 0..requests.len() {
             requests_c.push(bindings::transfer_request_t {
diff --git a/mooncake-transfer-engine/scripts/allocator_build_common.sh b/mooncake-transfer-engine/scripts/allocator_build_common.sh
new file mode 100644
index 00000000..410ce0ce
--- /dev/null
+++ b/mooncake-transfer-engine/scripts/allocator_build_common.sh
@@ -0,0 +1,20 @@
+#!/bin/bash
+
+resolve_allocator_script_dir() {
+    cd "$(dirname "$(readlink -f "$0")")" &>/dev/null && pwd
+}
+
+prepare_allocator_build_env() {
+    OUTPUT_DIR=${1:-.}
+    local include_list=${2:-}
+
+    SCRIPT_DIR=$(resolve_allocator_script_dir)
+    include_list="${include_list:+${include_list} }${SCRIPT_DIR}/../include"
+
+    INCLUDE_FLAGS=""
+    if [ -n "$include_list" ]; then
+        INCLUDE_FLAGS=$(printf '%s\n' "$include_list" | tr ' ' '\n' | sed '/^$/d; s/^/-I/' | paste -sd' ' -)
+    fi
+
+    mkdir -p "$OUTPUT_DIR"
+}
diff --git a/mooncake-transfer-engine/src/CMakeLists.txt b/mooncake-transfer-engine/src/CMakeLists.txt
index 95590dc6..b151901f 100644
--- a/mooncake-transfer-engine/src/CMakeLists.txt
+++ b/mooncake-transfer-engine/src/CMakeLists.txt
@@ -48,6 +48,7 @@ target_link_libraries(
   PUBLIC base
          transport
          rdma_transport
+         mooncake_common
          ibverbs
          glog::glog
          gflags::gflags
@@ -63,7 +64,11 @@ endif()
 
 if(USE_CUDA)
   target_include_directories(transfer_engine PRIVATE /usr/local/cuda/include)
-  target_link_libraries(transfer_engine PUBLIC cuda cudart rt)
+  if(USE_CXI)
+    target_link_libraries(transfer_engine PUBLIC cuda cudart rt)
+  else()
+    target_link_libraries(transfer_engine PUBLIC cuda cudart rt mlx5)
+  endif()
   if(USE_NVMEOF)
     target_link_libraries(transfer_engine PUBLIC nvmeof_transport cufile)
   endif()
@@ -82,9 +87,36 @@ if(USE_MUSA)
   target_link_libraries(transfer_engine PUBLIC musa musart rt)
 endif()
 
+if(USE_HYGON)
+  target_include_directories(transfer_engine PRIVATE ${DTK_INCLUDE_DIR})
+  target_link_libraries(transfer_engine PUBLIC cuda cudart rt)
+endif()
+
+if(USE_COREX)
+  target_include_directories(transfer_engine PRIVATE ${COREX_INCLUDE_DIR})
+  target_link_libraries(transfer_engine PUBLIC cuda cudart rt)
+endif()
+
 if(USE_HIP)
   target_include_directories(transfer_engine PRIVATE ${HIP_INCLUDE_DIRS})
   target_link_libraries(transfer_engine PUBLIC hip::host rt)
+
+  # Optional dmabuf MR registration path (requires hsa-runtime64 and a kernel
+  # with CONFIG_PCI_P2PDMA + CONFIG_DMABUF_MOVE_NOTIFY).
+  option(USE_HIP_DMABUF "Enable HIP dmabuf RDMA MR registration" ON)
+  if(USE_HIP_DMABUF)
+    find_package(hsa-runtime64 CONFIG)
+    if(hsa-runtime64_FOUND)
+      target_compile_definitions(transfer_engine PRIVATE USE_HIP_DMABUF)
+      target_link_libraries(transfer_engine PUBLIC hsa-runtime64::hsa-runtime64)
+      message(STATUS "HIP dmabuf MR registration enabled (hsa-runtime64 found)")
+    else()
+      message(
+        STATUS "HIP dmabuf MR registration disabled (hsa-runtime64 not found)")
+    endif()
+  else()
+    message(STATUS "HIP dmabuf MR registration disabled (USE_HIP_DMABUF=OFF)")
+  endif()
 endif()
 
 if(USE_MLU)
@@ -122,10 +154,20 @@ if(USE_INTRA_NVLINK)
   target_compile_definitions(transfer_engine PUBLIC USE_INTRA_NVLINK)
 endif()
 
+if(USE_MLX5DV)
+  message(STATUS "Enabled USE_MLX5DV support (mlx5 direct verbs)")
+  target_link_libraries(transfer_engine PUBLIC mlx5)
+endif()
+
 if(USE_EFA)
   message(STATUS "Enabled USE_EFA (AWS Elastic Fabric Adapter) support")
   target_link_libraries(transfer_engine PUBLIC fabric efa_transport)
 endif()
+if(USE_CXI)
+  message(STATUS "Enabled USE_CXI (HPE Cray Slingshot) support")
+  target_link_libraries(transfer_engine PUBLIC fabric cxi_transport)
+endif()
+
 if(USE_UB)
   message(STATUS "Enabled USE_UB protocol support")
   target_link_libraries(transfer_engine PUBLIC ub_transport)
diff --git a/mooncake-transfer-engine/src/config.cpp b/mooncake-transfer-engine/src/config.cpp
index 93815b3e..6adb9352 100644
--- a/mooncake-transfer-engine/src/config.cpp
+++ b/mooncake-transfer-engine/src/config.cpp
@@ -17,10 +17,88 @@
 #include <cstring>
 #include <cstdio>
 #include <cstdlib>
-#include <dirent.h>
+#include <filesystem>
+#include <sstream>
+#include <strings.h>
 #include <unistd.h>
 
 namespace mooncake {
+namespace {
+
+std::string trimConfigToken(const std::string& value) {
+    const auto begin = value.find_first_not_of(" \t\n\r");
+    if (begin == std::string::npos) return "";
+    const auto end = value.find_last_not_of(" \t\n\r");
+    return value.substr(begin, end - begin + 1);
+}
+
+std::vector<std::string> splitConfigString(const std::string& value,
+                                           char delim) {
+    std::vector<std::string> result;
+    std::stringstream stream(value);
+    std::string item;
+    while (std::getline(stream, item, delim)) {
+        result.push_back(trimConfigToken(item));
+    }
+    return result;
+}
+
+bool parseBoolConfigEnv(const char* value, const char* env_name, bool& output) {
+    if (strcmp(value, "1") == 0 || strcasecmp(value, "true") == 0) {
+        output = true;
+        return true;
+    }
+    if (strcmp(value, "0") == 0 || strcasecmp(value, "false") == 0) {
+        output = false;
+        return true;
+    }
+    LOG(WARNING) << "Ignore value from environment variable " << env_name
+                 << ", it should be 0|1|true|false";
+    return false;
+}
+
+void parseNicPeerAffinity(
+    const char* env,
+    std::unordered_map<std::string, std::vector<std::string>>& affinity) {
+    affinity.clear();
+    if (!env || env[0] == '\0') return;
+
+    for (const auto& raw_rule : splitConfigString(env, ';')) {
+        const auto rule = trimConfigToken(raw_rule);
+        if (rule.empty()) continue;
+
+        auto delim = rule.find('=');
+        if (delim == std::string::npos) {
+            LOG(WARNING) << "Invalid MC_NIC_PEER_AFFINITY rule '" << rule
+                         << "'. Expected local_hca=peer_hca[,peer_hca].";
+            continue;
+        }
+
+        auto local_hca = trimConfigToken(rule.substr(0, delim));
+        if (local_hca.empty()) {
+            LOG(WARNING) << "Invalid MC_NIC_PEER_AFFINITY rule '" << rule
+                         << "': local HCA is empty.";
+            continue;
+        }
+
+        std::vector<std::string> peer_hcas;
+        for (const auto& peer_hca :
+             splitConfigString(rule.substr(delim + 1), ',')) {
+            if (!peer_hca.empty()) peer_hcas.push_back(peer_hca);
+        }
+
+        if (peer_hcas.empty()) {
+            LOG(WARNING) << "Invalid MC_NIC_PEER_AFFINITY rule '" << rule
+                         << "': peer HCA list is empty.";
+            continue;
+        }
+
+        affinity[local_hca] = std::move(peer_hcas);
+    }
+}
+
+}  // namespace
+
 void loadGlobalConfig(GlobalConfig& config) {
     const char* num_cq_per_ctx_env = std::getenv("MC_NUM_CQ_PER_CTX");
     if (num_cq_per_ctx_env) {
@@ -67,6 +145,24 @@ void loadGlobalConfig(GlobalConfig& config) {
                 << "Ignore value from environment variable MC_GID_INDEX";
     }
 
+    const char* pkey_index_env = std::getenv("MC_PKEY_INDEX");
+    if (pkey_index_env) {
+        try {
+            int val = std::stoi(pkey_index_env);
+            if (val >= 0 && val <= UINT16_MAX) {
+                config.pkey_index = static_cast<uint16_t>(val);
+            } else {
+                LOG(WARNING)
+                    << "Ignore value from environment variable MC_PKEY_INDEX, "
+                    << "value " << pkey_index_env
+                    << " out of range (should be 0-65535)";
+            }
+        } catch (const std::exception& e) {
+            LOG(WARNING) << "Invalid MC_PKEY_INDEX environment value: "
+                         << pkey_index_env << ". Error: " << e.what();
+        }
+    }
+
     const char* max_cqe_per_ctx_env = std::getenv("MC_MAX_CQE_PER_CTX");
     if (max_cqe_per_ctx_env) {
         size_t val = atoi(max_cqe_per_ctx_env);
@@ -139,7 +235,6 @@ void loadGlobalConfig(GlobalConfig& config) {
         else {
             LOG(ERROR) << "Ignore value from environment variable MC_MTU, it "
                           "should be 512|1024|2048|4096";
-            exit(EXIT_FAILURE);
         }
     }
 
@@ -195,16 +290,6 @@ void loadGlobalConfig(GlobalConfig& config) {
         }
     }
 
-    const char* efa_striping_env = std::getenv("MC_EFA_STRIPING_THRESHOLD");
-    if (efa_striping_env) {
-        size_t val = strtoull(efa_striping_env, nullptr, 10);
-        if (val > 0)
-            config.efa_striping_threshold = val;
-        else
-            LOG(WARNING) << "Ignore value from environment variable "
-                            "MC_EFA_STRIPING_THRESHOLD";
-    }
-
     const char* retry_cnt_env = std::getenv("MC_RETRY_CNT");
     if (retry_cnt_env) {
         size_t val = atoi(retry_cnt_env);
@@ -215,6 +300,18 @@ void loadGlobalConfig(GlobalConfig& config) {
                 << "Ignore value from environment variable MC_RETRY_CNT";
     }
 
+    const char* auto_gid_max_retries_env =
+        std::getenv("MC_AUTO_GID_MAX_RETRIES");
+    if (auto_gid_max_retries_env) {
+        int val = atoi(auto_gid_max_retries_env);
+        if (val >= 0 && val <= 16) {
+            config.auto_gid_max_retries = val;
+        } else {
+            LOG(WARNING) << "Ignore value from environment variable "
+                            "MC_AUTO_GID_MAX_RETRIES";
+        }
+    }
+
     const char* disable_metacache = std::getenv("MC_DISABLE_METACACHE");
     if (disable_metacache) {
         config.metacache = false;
@@ -223,15 +320,32 @@ void loadGlobalConfig(GlobalConfig& config) {
     const char* handshake_listen_backlog =
         std::getenv("MC_HANDSHAKE_LISTEN_BACKLOG");
     if (handshake_listen_backlog) {
-        int val = std::stoi(handshake_listen_backlog);
-        if (val > 0) {
-            config.handshake_listen_backlog = val;
-        } else {
-            LOG(WARNING) << "Ignore value from environment variable "
-                            "MC_HANDSHAKE_LISTEN_BACKLOG";
+        try {
+            int val = std::stoi(handshake_listen_backlog);
+            if (val > 0) {
+                config.handshake_listen_backlog = val;
+            } else {
+                LOG(WARNING) << "Ignore value from environment variable "
+                                "MC_HANDSHAKE_LISTEN_BACKLOG";
+            }
+        } catch (const std::exception& e) {
+            LOG(WARNING) << "Invalid MC_HANDSHAKE_LISTEN_BACKLOG environment "
+                            "value: "
+                         << handshake_listen_backlog << ". Error: " << e.what();
         }
     }
 
+    const char* handshake_connect_timeout =
+        std::getenv("MC_HANDSHAKE_CONNECT_TIMEOUT");
+    if (handshake_connect_timeout) {
+        int val = atoi(handshake_connect_timeout);
+        if (val > 0 && val < 3600)
+            config.handshake_connect_timeout = val;
+        else
+            LOG(WARNING) << "Ignore value from environment variable "
+                            "MC_HANDSHAKE_CONNECT_TIMEOUT";
+    }
+
     const char* log_level = std::getenv("MC_LOG_LEVEL");
     config.trace = false;
     if (log_level) {
@@ -261,7 +375,8 @@ void loadGlobalConfig(GlobalConfig& config) {
     const char* log_dir_path = std::getenv("MC_LOG_DIR");
     if (log_dir_path) {
         google::InitGoogleLogging("mooncake-transfer-engine");
-        if (opendir(log_dir_path) == NULL) {
+        std::error_code ec;
+        if (!std::filesystem::is_directory(log_dir_path, ec)) {
             LOG(WARNING)
                 << "Path [" << log_dir_path
                 << "] is not a valid directory path. Still logging to stderr.";
@@ -277,24 +392,17 @@ void loadGlobalConfig(GlobalConfig& config) {
         }
     }
 
-    const char* min_port_env = std::getenv("MC_MIN_PRC_PORT");
-    if (min_port_env) {
-        int val = atoi(min_port_env);
-        if (val > 0 && val < 65536)
-            config.rpc_min_port = val;
-        else
-            LOG(WARNING)
-                << "Ignore value from environment variable MC_PRC_MIN_PORT";
-    }
-
-    const char* max_port_env = std::getenv("MC_MAX_PRC_PORT");
-    if (max_port_env) {
-        int val = atoi(max_port_env);
-        if (val > 0 && val < 65536)
-            config.rpc_max_port = val;
-        else
-            LOG(WARNING)
-                << "Ignore value from environment variable MC_PRC_MAX_PORT";
+    const char* min_port_env = std::getenv("MC_MIN_RPC_PORT");
+    if (!min_port_env) min_port_env = std::getenv("MC_MIN_PRC_PORT");
+    const char* max_port_env = std::getenv("MC_MAX_RPC_PORT");
+    if (!max_port_env) max_port_env = std::getenv("MC_MAX_PRC_PORT");
+    {
+        int raw_min = min_port_env ? atoi(min_port_env) : config.rpc_min_port;
+        int raw_max = max_port_env ? atoi(max_port_env) : config.rpc_max_port;
+        auto [validated_min, validated_max] =
+            ValidatePortRange(raw_min, raw_max, 15000, 17000);
+        config.rpc_min_port = validated_min;
+        config.rpc_max_port = validated_max;
     }
 
     if (std::getenv("MC_USE_IPV6")) {
@@ -317,6 +425,34 @@ void loadGlobalConfig(GlobalConfig& config) {
         config.enable_dest_device_affinity = true;
     }
 
+    const char* enable_hca_peer_affinity_env =
+        std::getenv("MC_ENABLE_HCA_PEER_AFFINITY");
+    if (enable_hca_peer_affinity_env) {
+        parseBoolConfigEnv(enable_hca_peer_affinity_env,
+                           "MC_ENABLE_HCA_PEER_AFFINITY",
+                           config.enable_hca_peer_affinity);
+    }
+
+    parseNicPeerAffinity(std::getenv("MC_NIC_PEER_AFFINITY"),
+                         config.nic_peer_affinity);
+
+    if (config.enable_hca_peer_affinity && config.enable_dest_device_affinity) {
+        LOG(ERROR) << "MC_ENABLE_HCA_PEER_AFFINITY and "
+                      "MC_ENABLE_DEST_DEVICE_AFFINITY cannot be enabled at "
+                      "the same time; falling back to default peer device "
+                      "selection.";
+        config.enable_hca_peer_affinity = false;
+        config.enable_dest_device_affinity = false;
+    }
+
+    const char* log_rdma_slice_affinity_env =
+        std::getenv("MC_LOG_RDMA_SLICE_AFFINITY");
+    if (log_rdma_slice_affinity_env) {
+        parseBoolConfigEnv(log_rdma_slice_affinity_env,
+                           "MC_LOG_RDMA_SLICE_AFFINITY",
+                           config.log_rdma_slice_affinity);
+    }
+
     const char* enable_parallel_reg_mr =
         std::getenv("MC_ENABLE_PARALLEL_REG_MR");
     if (enable_parallel_reg_mr) {
@@ -359,6 +495,24 @@ void loadGlobalConfig(GlobalConfig& config) {
         }
     }
 
+    const char* service_level_env = std::getenv("MC_IB_SL");
+    if (service_level_env) {
+        try {
+            int val = std::stoi(service_level_env);
+            if (val >= 0 && val <= 15) {
+                config.ib_service_level = val;
+            } else {
+                LOG(WARNING)
+                    << "Ignore value from environment variable MC_IB_SL, "
+                    << "value " << service_level_env
+                    << " out of range (should be 0-15)";
+            }
+        } catch (const std::exception& e) {
+            LOG(WARNING) << "Invalid MC_IB_SL environment value: "
+                         << service_level_env << ". Error: " << e.what();
+        }
+    }
+
     const char* ib_relaxed_ordering_env =
         std::getenv("MC_IB_PCI_RELAXED_ORDERING");
     if (ib_relaxed_ordering_env) {
@@ -369,6 +523,62 @@ void loadGlobalConfig(GlobalConfig& config) {
             LOG(WARNING) << "Ignore value from environment variable "
                             "MC_IB_PCI_RELAXED_ORDERING, it should be 0|1|2";
     }
+
+    const char* mlx5_qp_udp_sports_env = std::getenv("MC_MLX5_QP_UDP_SPORTS");
+    if (mlx5_qp_udp_sports_env && *mlx5_qp_udp_sports_env) {
+        std::vector<uint16_t> ports;
+        std::stringstream ss(mlx5_qp_udp_sports_env);
+        std::string item;
+        bool ok = true;
+        while (std::getline(ss, item, ',')) {
+            // Trim leading/trailing whitespace.
+            auto l = item.find_first_not_of(" \t");
+            auto r = item.find_last_not_of(" \t");
+            if (l == std::string::npos) continue;
+            item = item.substr(l, r - l + 1);
+            try {
+                int val = std::stoi(item);
+                if (val < 0 || val > 65535) {
+                    LOG(WARNING)
+                        << "MC_MLX5_QP_UDP_SPORTS entry out of range: " << item;
+                    ok = false;
+                    break;
+                }
+                ports.push_back(static_cast<uint16_t>(val));
+            } catch (const std::exception& e) {
+                LOG(WARNING) << "Invalid MC_MLX5_QP_UDP_SPORTS entry: " << item
+                             << ". Error: " << e.what();
+                ok = false;
+                break;
+            }
+        }
+        if (ok && !ports.empty()) {
+            config.mlx5_qp_udp_sports = std::move(ports);
+        } else if (!ok) {
+            LOG(WARNING) << "Ignore MC_MLX5_QP_UDP_SPORTS entirely due to "
+                            "parse errors";
+        }
+    }
+
+    const char* mlx5_qp_lag_port_balance_env =
+        std::getenv("MC_MLX5_QP_LAG_PORT_BALANCE");
+    if (mlx5_qp_lag_port_balance_env && *mlx5_qp_lag_port_balance_env) {
+        std::string val(mlx5_qp_lag_port_balance_env);
+        // Trim leading/trailing whitespace.
+        auto l = val.find_first_not_of(" \t");
+        auto r = val.find_last_not_of(" \t");
+        if (l != std::string::npos) {
+            val = val.substr(l, r - l + 1);
+            if (val == "1" || val == "true")
+                config.mlx5_qp_lag_port_balance = true;
+            else if (val == "0" || val == "false")
+                config.mlx5_qp_lag_port_balance = false;
+            else
+                LOG(WARNING) << "Ignore MC_MLX5_QP_LAG_PORT_BALANCE: expected "
+                                "0/1/true/false, got: "
+                             << val;
+        }
+    }
 }
 
 std::string mtuLengthToString(ibv_mtu mtu) {
@@ -409,6 +619,7 @@ void dumpGlobalConfig() {
               << config.num_comp_channels_per_ctx;
     LOG(INFO) << "port = " << config.port;
     LOG(INFO) << "gid_index = " << config.gid_index;
+    LOG(INFO) << "pkey_index = " << config.pkey_index;
     LOG(INFO) << "max_mr_size = " << config.max_mr_size;
     LOG(INFO) << "max_cqe = " << config.max_cqe;
     LOG(INFO) << "max_ep_per_ctx = " << config.max_ep_per_ctx;
@@ -419,6 +630,21 @@ void dumpGlobalConfig() {
     LOG(INFO) << "mtu_length = " << mtuLengthToString(config.mtu_length);
     LOG(INFO) << "parallel_reg_mr = " << config.parallel_reg_mr;
     LOG(INFO) << "ib_traffic_class = " << config.ib_traffic_class;
+    LOG(INFO) << "ib_service_level = " << config.ib_service_level;
+    {
+        std::ostringstream oss;
+        for (size_t i = 0; i < config.mlx5_qp_udp_sports.size(); ++i) {
+            if (i) oss << ",";
+            oss << config.mlx5_qp_udp_sports[i];
+        }
+        LOG(INFO) << "mlx5_qp_udp_sports = ["
+                  << (config.mlx5_qp_udp_sports.empty() ? "<unset>" : oss.str())
+                  << "]";
+    }
+    LOG(INFO) << "mlx5_qp_lag_port_balance = "
+              << (config.mlx5_qp_lag_port_balance ? "true" : "false");
+    LOG(INFO) << "log_rdma_slice_affinity = "
+              << (config.log_rdma_slice_affinity ? "true" : "false");
 }
 
 GlobalConfig& globalConfig() {
@@ -429,4 +655,26 @@ GlobalConfig& globalConfig() {
 }
 
 uint16_t getDefaultHandshakePort() { return globalConfig().handshake_port; }
+
+std::pair<int, int> ValidatePortRange(int min_port, int max_port,
+                                      int default_min, int default_max) {
+    constexpr int kMinAllowed = 1024;
+    constexpr int kEphemeralStart = 32768;
+    constexpr int kEphemeralEnd = 60999;
+    constexpr int kMaxAllowed = 65535;
+
+    auto is_valid_port = [&](int p) {
+        return p >= kMinAllowed && p <= kMaxAllowed &&
+               !(p >= kEphemeralStart && p <= kEphemeralEnd);
+    };
+
+    if (!is_valid_port(min_port) || !is_valid_port(max_port) ||
+        min_port > max_port) {
+        LOG(WARNING) << "Invalid port range [" << min_port << ", " << max_port
+                     << "], falling back to default [" << default_min << ", "
+                     << default_max << "]";
+        return {default_min, default_max};
+    }
+    return {min_port, max_port};
+}
 }  // namespace mooncake
diff --git a/mooncake-transfer-engine/src/memory_location.cpp b/mooncake-transfer-engine/src/memory_location.cpp
index 321a63c3..a18b2460 100644
--- a/mooncake-transfer-engine/src/memory_location.cpp
+++ b/mooncake-transfer-engine/src/memory_location.cpp
@@ -35,8 +35,9 @@ const std::vector<MemoryLocationEntry> getMemoryLocation(void *start,
                                                          bool only_first_page) {
     std::vector<MemoryLocationEntry> entries;
 
-#if defined(USE_CUDA) || defined(USE_MUSA) || defined(USE_HIP) || \
-    defined(USE_MLU) || defined(USE_MACA)
+#if defined(USE_CUDA) || defined(USE_MUSA) || defined(USE_HIP) ||  \
+    defined(USE_MLU) || defined(USE_MACA) || defined(USE_HYGON) || \
+    defined(USE_COREX) || defined(USE_SUNRISE)
     cudaPointerAttributes attributes;
     cudaError_t result = cudaPointerGetAttributes(&attributes, start);
     if (result != cudaSuccess) {
diff --git a/mooncake-transfer-engine/src/multi_transport.cpp b/mooncake-transfer-engine/src/multi_transport.cpp
index 812ce821..70a6c749 100644
--- a/mooncake-transfer-engine/src/multi_transport.cpp
+++ b/mooncake-transfer-engine/src/multi_transport.cpp
@@ -44,6 +44,9 @@
 #ifdef USE_HIP
 #include "transport/hip_transport/hip_transport.h"
 #endif
+#ifdef USE_MACA
+#include "transport/maca_transport/maca_transport.h"
+#endif
 #ifdef USE_MNNVL
 #include "transport/nvlink_transport/nvlink_transport.h"
 #endif
@@ -56,6 +59,12 @@
 #ifdef USE_EFA
 #include "transport/efa_transport/efa_transport.h"
 #endif
+#ifdef USE_CXI
+#include "transport/cxi_transport/cxi_transport.h"
+#endif
+#ifdef USE_SUNRISE
+#include "transport/sunrise_link_transport/sunrise_link_transport.h"
+#endif
 #ifdef USE_UB
 #include "transport/kunpeng_transport/ub_transport.h"
 #endif
@@ -121,6 +130,7 @@ Status MultiTransport::submitTransfer(
         assert(transport);
         auto& task = batch_desc.task_list[task_id];
         task.batch_id = batch_id;
+        task.transport_ = transport;
 #ifdef USE_ASCEND_HETEROGENEOUS
         task.request = const_cast<Transport::TransferRequest*>(&request);
 #else
@@ -163,6 +173,7 @@ Status MultiTransport::mp_submitTransfer(
         assert(transport);
         auto& task = batch_desc.task_list[task_id];
         task.batch_id = batch_id;
+        task.transport_ = transport;
 #ifdef USE_ASCEND_HETEROGENEOUS
         task.request = const_cast<Transport::TransferRequest*>(&request);
 #else
@@ -192,6 +203,45 @@ Status MultiTransport::getTransferStatus(BatchID batch_id, size_t task_id,
         return Status::InvalidArgument("Task ID out of range");
     }
     auto& task = batch_desc.task_list[task_id];
+
+    // Helper: check if any slice has exceeded the configured timeout.
+    // Returns true if a timeout was detected (and logs it).
+    auto checkSliceTimeout = [&](const Transport::TransferTask& t) -> bool {
+        if (globalConfig().slice_timeout <= 0) return false;
+        auto current_ts = getCurrentTimeInNano();
+        const int64_t kPacketDeliveryTimeout =
+            globalConfig().slice_timeout * 1000000000;
+        for (auto& slice : t.slice_list) {
+            auto ts = slice->ts;
+            if (ts > 0 && current_ts > ts &&
+                current_ts - ts > kPacketDeliveryTimeout) {
+                LOG(INFO) << "Slice timeout detected";
+                return true;
+            }
+        }
+        return false;
+    };
+
+    // If the task has an associated transport, delegate to its
+    // getTransferStatus() to trigger transport-specific completion
+    // polling. For example, the NVLink async transport polls CUDA
+    // streams via cudaStreamQuery() here; without this call the
+    // slice statuses (and therefore success/failed_slice_count)
+    // would never be updated.
+    if (task.transport_) {
+        auto ret =
+            task.transport_->getTransferStatus(batch_id, task_id, status);
+        if (!ret.ok()) return ret;
+
+        // Apply timeout check on top of the transport's result.
+        if (status.s == Transport::TransferStatusEnum::WAITING &&
+            checkSliceTimeout(task)) {
+            status.s = Transport::TransferStatusEnum::TIMEOUT;
+        }
+        return Status::OK();
+    }
+
+    // Fallback for tasks without a transport pointer (legacy path)
     status.transferred_bytes = task.transferred_bytes;
     uint64_t success_slice_count = task.success_slice_count;
     uint64_t failed_slice_count = task.failed_slice_count;
@@ -204,21 +254,11 @@ Status MultiTransport::getTransferStatus(BatchID batch_id, size_t task_id,
         }
         task.is_finished = true;
     } else {
-        if (globalConfig().slice_timeout > 0) {
-            auto current_ts = getCurrentTimeInNano();
-            const int64_t kPacketDeliveryTimeout =
-                globalConfig().slice_timeout * 1000000000;
-            for (auto& slice : task.slice_list) {
-                auto ts = slice->ts;
-                if (ts > 0 && current_ts > ts &&
-                    current_ts - ts > kPacketDeliveryTimeout) {
-                    LOG(INFO) << "Slice timeout detected";
-                    status.s = Transport::TransferStatusEnum::TIMEOUT;
-                    return Status::OK();
-                }
-            }
+        if (checkSliceTimeout(task)) {
+            status.s = Transport::TransferStatusEnum::TIMEOUT;
+        } else {
+            status.s = Transport::TransferStatusEnum::WAITING;
         }
-        status.s = Transport::TransferStatusEnum::WAITING;
     }
     return Status::OK();
 }
@@ -322,6 +362,11 @@ Transport* MultiTransport::installTransport(const std::string& proto,
         transport = new HipTransport();
     }
 #endif
+#ifdef USE_MACA
+    else if (std::string(proto) == "maca") {
+        transport = new MacaTransport();
+    }
+#endif
 #ifdef USE_MNNVL
     else if (std::string(proto) == "nvlink") {
         transport = new NvlinkTransport();
@@ -342,6 +387,16 @@ Transport* MultiTransport::installTransport(const std::string& proto,
         transport = new EfaTransport();
     }
 #endif
+#ifdef USE_CXI
+    else if (std::string(proto) == "cxi") {
+        transport = new CxiTransport();
+    }
+#endif
+#ifdef USE_SUNRISE
+    else if (std::string(proto) == "sunrise_link") {
+        transport = new SunriseLinkTransport();
+    }
+#endif
 
     if (!transport) {
         LOG(ERROR) << "Unsupported transport " << proto
@@ -384,6 +439,7 @@ Transport* MultiTransport::installTransport(const std::string& proto,
     }
 #endif
     if (transport->install(local_server_name_, metadata_, topo)) {
+        delete transport;
         return nullptr;
     }
 
diff --git a/mooncake-transfer-engine/src/topology.cpp b/mooncake-transfer-engine/src/topology.cpp
index 2037b3f9..38435f43 100644
--- a/mooncake-transfer-engine/src/topology.cpp
+++ b/mooncake-transfer-engine/src/topology.cpp
@@ -14,14 +14,15 @@
 
 #include <glog/logging.h>
 
+#include <algorithm>
 #include <fstream>
 #include <iostream>
 #include <map>
 #include <set>
 #include <string>
+#include <unordered_set>
 #include <utility>
 #include <vector>
-#include <ctype.h>
 #include <dirent.h>
 #include <infiniband/verbs.h>
 #include <limits.h>
@@ -31,8 +32,10 @@
 #include <sys/stat.h>
 
 #include "cuda_alike.h"
+#include "config.h"
 #include "memory_location.h"
 #include "topology.h"
+#include "char_util.h"
 #ifdef USE_UB
 #include <libgen.h>
 #include <urma_api.h>
@@ -142,10 +145,56 @@ struct InfinibandDevice {
     int numa_node;
 };
 
+static inline void ltrim(std::string &s) {
+    s.erase(s.begin(), std::find_if(s.begin(), s.end(), [](unsigned char ch) {
+                return !std::isspace(ch);
+            }));
+}
+
+static inline void rtrim(std::string &s) {
+    s.erase(std::find_if(s.rbegin(), s.rend(),
+                         [](unsigned char ch) { return !std::isspace(ch); })
+                .base(),
+            s.end());
+}
+
+static std::set<std::string> getIbvDeviceWhitelist() {
+    std::set<std::string> whitelist;
+    const char *env = std::getenv("MC_TE_FILTERS");
+    if (!env || !*env) {
+        return whitelist;
+    }
+
+    LOG(INFO) << "IB device whitelist: " << env;
+    std::string env_str(env);
+    size_t start = 0;
+    size_t pos = 0;
+    while ((pos = env_str.find(',', start)) != std::string::npos) {
+        std::string token = env_str.substr(start, pos - start);
+        ltrim(token);
+        rtrim(token);
+        if (!token.empty()) {
+            whitelist.insert(std::move(token));
+        }
+        start = pos + 1;
+    }
+    if (start < env_str.length()) {
+        std::string token = env_str.substr(start);
+        ltrim(token);
+        rtrim(token);
+        if (!token.empty()) {
+            whitelist.insert(std::move(token));
+        }
+    }
+    return whitelist;
+}
+
 static std::vector<InfinibandDevice> listInfiniBandDevices(
     const std::vector<std::string> &filter) {
+#ifndef USE_CXI
     int num_devices = 0;
     std::vector<InfinibandDevice> devices;
+    const auto whitelist = getIbvDeviceWhitelist();
 
     struct ibv_device **device_list = ibv_get_device_list(&num_devices);
     if (!device_list) {
@@ -164,6 +213,12 @@ static std::vector<InfinibandDevice> listInfiniBandDevices(
                                          device_name) == filter.end())
             continue;
 
+        if (!whitelist.empty() &&
+            whitelist.find(device_name) == whitelist.end()) {
+            LOG(INFO) << "Skipping device: " << device_name;
+            continue;
+        }
+
         // Check device availability before adding to the list
         if (!isIbDeviceAvailable(device_list[i])) {
             LOG(WARNING) << "Skipping unavailable device: " << device_name;
@@ -194,6 +249,99 @@ static std::vector<InfinibandDevice> listInfiniBandDevices(
     }
     ibv_free_device_list(device_list);
     return devices;
+#else
+    // assume devices available and working, later use libcxi to check if they
+    // actually work
+    std::vector<std::string> device_names;
+    std::vector<InfinibandDevice> devices;
+    const char *nic_path = "/sys/class/cxi";
+
+    DIR *dir = opendir(nic_path);
+    if (!dir) {
+        perror("opendir");
+        LOG(WARNING) << "No CXI devices were found, check you system";
+        return {};
+    }
+
+    struct dirent *entry;
+    while ((entry = readdir(dir)) != nullptr) {
+        if (entry->d_type == DT_LNK && std::string(entry->d_name) != "." &&
+            std::string(entry->d_name) != "..") {
+            device_names.emplace_back(entry->d_name);
+        }
+    }
+    closedir(dir);
+
+    for (auto &dv_name : device_names) {
+        if (!filter.empty() &&
+            std::find(filter.begin(), filter.end(), dv_name) == filter.end())
+            continue;
+
+        char path[PATH_MAX + 32];
+        char resolved_path[PATH_MAX];
+
+        snprintf(path, sizeof(path), "/sys/class/cxi/%s/device",
+                 dv_name.c_str());
+        if (realpath(path, resolved_path) == NULL) {
+            PLOG(ERROR) << "Can't resolve CXI device path for " << path;
+            continue;
+        }
+        std::string pci_bus_id = basename(resolved_path);
+
+        int numa_node = -1;
+
+        snprintf(path, sizeof(path), "%s/numa_node", resolved_path);
+        std::ifstream(path) >> numa_node;
+
+        devices.push_back(InfinibandDevice{.name = std::move(dv_name),
+                                           .pci_bus_id = std::move(pci_bus_id),
+                                           .numa_node = numa_node});
+    }
+    return devices;
+#endif
+}
+
+static void precomputeResolvedHcaPeerAffinity(
+    const TopologyMatrix &matrix, const std::map<std::string, int> &hca_id_map,
+    std::unordered_map<std::string,
+                       std::unordered_map<std::string, std::vector<int>>>
+        &affinity_by_local) {
+    affinity_by_local.clear();
+    const auto &config = globalConfig();
+    if (!config.enable_hca_peer_affinity) return;
+    if (config.nic_peer_affinity.empty()) {
+        LOG(WARNING) << "MC_ENABLE_HCA_PEER_AFFINITY is set, but "
+                        "MC_NIC_PEER_AFFINITY has no valid mappings.";
+        return;
+    }
+
+    for (const auto &entry : matrix) {
+        if (entry.first.rfind("cuda:", 0) != 0) continue;
+
+        std::unordered_set<std::string> preferred_set(
+            entry.second.preferred_hca.begin(),
+            entry.second.preferred_hca.end());
+
+        for (const auto &mapping : config.nic_peer_affinity) {
+            const auto &local_hca = mapping.first;
+            const auto &peer_hcas = mapping.second;
+
+            std::vector<int> peer_preferred_hca;
+            std::unordered_set<std::string> peer_preferred_set;
+            for (const auto &peer_hca : peer_hcas) {
+                if (!preferred_set.count(peer_hca)) continue;
+                auto hca_id_it = hca_id_map.find(peer_hca);
+                if (hca_id_it == hca_id_map.end()) continue;
+                if (peer_preferred_set.insert(peer_hca).second) {
+                    peer_preferred_hca.push_back(hca_id_it->second);
+                }
+            }
+
+            if (peer_preferred_hca.empty()) continue;
+            affinity_by_local[local_hca][entry.first] =
+                std::move(peer_preferred_hca);
+        }
+    }
 }
 
 #ifdef USE_UB
@@ -337,8 +485,9 @@ static std::vector<TopologyEntry> discoverCpuTopology(
     return topology;
 }
 
-#if defined(USE_CUDA) || defined(USE_MUSA) || defined(USE_HIP) || \
-    defined(USE_MLU) || defined(USE_MACA)
+#if defined(USE_CUDA) || defined(USE_MUSA) || defined(USE_HIP) ||  \
+    defined(USE_MLU) || defined(USE_MACA) || defined(USE_HYGON) || \
+    defined(USE_COREX)
 
 static int getPciDistance(const char *bus1, const char *bus2) {
     char buf[PATH_MAX];
@@ -394,7 +543,7 @@ static std::vector<TopologyEntry> discoverCudaTopology(
             cudaSuccess) {
             continue;
         }
-        for (char *ch = pci_bus_id; (*ch = tolower(*ch)); ch++);
+        for (char *ch = pci_bus_id; (*ch = to_lower(*ch)); ch++);
 
         std::vector<std::string> preferred_hca;
         std::vector<std::string> avail_hca;
@@ -468,16 +617,19 @@ void Topology::clear() {
     matrix_.clear();
     hca_list_.clear();
     resolved_matrix_.clear();
+    resolved_hca_peer_affinity_by_local_.clear();
 }
 
 int Topology::discover(const std::vector<std::string> &filter) {
     matrix_.clear();
+    resolved_hca_peer_affinity_by_local_.clear();
     auto all_hca = listInfiniBandDevices(filter);
     for (auto &ent : discoverCpuTopology(all_hca)) {
         matrix_[ent.name] = ent;
     }
-#if defined(USE_CUDA) || defined(USE_MUSA) || defined(USE_HIP) || \
-    defined(USE_MLU) || defined(USE_MACA)
+#if defined(USE_CUDA) || defined(USE_MUSA) || defined(USE_HIP) ||  \
+    defined(USE_MLU) || defined(USE_MACA) || defined(USE_HYGON) || \
+    defined(USE_COREX)
     for (auto &ent : discoverCudaTopology(all_hca)) {
         matrix_[ent.name] = ent;
     }
@@ -512,6 +664,7 @@ int Topology::parse(const std::string &topology_json) {
     }
 
     matrix_.clear();
+    resolved_hca_peer_affinity_by_local_.clear();
     for (const auto &key : root.getMemberNames()) {
         const Json::Value &value = root[key];
         if (value.isArray() && value.size() == 2) {
@@ -569,6 +722,36 @@ int Topology::selectDevice(const std::string storage_type,
     return selectDevice(storage_type, retry_count);
 }
 
+int Topology::selectDeviceByLocalHca(const std::string storage_type,
+                                     std::string_view local_hca,
+                                     int retry_count) {
+    if (!local_hca.empty()) {
+        auto local_it =
+            resolved_hca_peer_affinity_by_local_.find(std::string(local_hca));
+        if (local_it != resolved_hca_peer_affinity_by_local_.end()) {
+            auto hints_it = local_it->second.find(std::string(storage_type));
+            if (hints_it != local_it->second.end() &&
+                !hints_it->second.empty()) {
+                const auto &candidates = hints_it->second;
+                if (retry_count == 0) {
+                    int rand_value;
+                    if (use_round_robin_) {
+                        thread_local int tl_counter = 0;
+                        rand_value = tl_counter;
+                        tl_counter = (tl_counter + 1) % 10000;
+                    } else {
+                        rand_value = SimpleRandom::Get().next();
+                    }
+                    return candidates[rand_value % candidates.size()];
+                }
+                return candidates[(retry_count - 1) % candidates.size()];
+            }
+        }
+    }
+
+    return selectDevice(storage_type, retry_count);
+}
+
 int Topology::selectDevice(const std::string storage_type, int retry_count) {
     if (resolved_matrix_.count(storage_type) == 0) return ERR_DEVICE_NOT_FOUND;
 
@@ -600,6 +783,7 @@ int Topology::selectDevice(const std::string storage_type, int retry_count) {
 
 int Topology::resolve() {
     resolved_matrix_.clear();
+    resolved_hca_peer_affinity_by_local_.clear();
     hca_list_.clear();
     std::map<std::string, int> hca_id_map;
     int next_hca_map_index = 0;
@@ -634,6 +818,8 @@ int Topology::resolve() {
                 hca_id_map[hca];
         }
     }
+    precomputeResolvedHcaPeerAffinity(matrix_, hca_id_map,
+                                      resolved_hca_peer_affinity_by_local_);
     return 0;
 }
 
diff --git a/mooncake-transfer-engine/src/transfer_engine.cpp b/mooncake-transfer-engine/src/transfer_engine.cpp
index 4800c005..16fbc601 100644
--- a/mooncake-transfer-engine/src/transfer_engine.cpp
+++ b/mooncake-transfer-engine/src/transfer_engine.cpp
@@ -179,6 +179,17 @@ Transport* TransferEngine::getTransport(const std::string& proto) {
     return impl_->getTransport(proto);
 }
 
+#if (defined(USE_CUDA) || defined(USE_MUSA)) && !defined(USE_CXI)
+device::P2pTransport* TransferEngine::getOrCreateP2pTransport(int num_ranks) {
+    return impl_->getOrCreateP2pTransport(num_ranks);
+}
+
+device::RdmaTransport* TransferEngine::getOrCreateRdmaTransport(
+    const std::vector<std::string>& device_filter) {
+    return impl_->getOrCreateRdmaTransport(device_filter);
+}
+#endif
+
 bool TransferEngine::isTcpOnly() const { return impl_->isTcpOnly(); }
 
 int TransferEngine::syncSegmentCache(const std::string& segment_name) {
@@ -449,6 +460,8 @@ Status TransferEngine::submitTransfer(
             req.source = item.source;
             req.target_id = item.target_id;
             req.target_offset = item.target_offset;
+            req.transport_hint =
+                mooncake::tent::c_to_transport_hint(item.transport_hint);
             requests.push_back(req);
         }
         auto status = impl_tent_->submitTransfer(batch_id, requests);
@@ -473,6 +486,8 @@ Status TransferEngine::submitTransferWithNotify(
             req.source = item.source;
             req.target_id = item.target_id;
             req.target_offset = item.target_offset;
+            req.transport_hint =
+                mooncake::tent::c_to_transport_hint(item.transport_hint);
             requests.push_back(req);
         }
         mooncake::tent::Notification notifi;
@@ -580,6 +595,19 @@ Transport* TransferEngine::getTransport(const std::string& proto) {
         return impl_->getTransport(proto);
 }
 
+#if defined(USE_CUDA) || defined(USE_MUSA)
+device::P2pTransport* TransferEngine::getOrCreateP2pTransport(int num_ranks) {
+    if (use_tent_) return nullptr;
+    return impl_->getOrCreateP2pTransport(num_ranks);
+}
+
+device::RdmaTransport* TransferEngine::getOrCreateRdmaTransport(
+    const std::vector<std::string>& device_filter) {
+    if (use_tent_) return nullptr;
+    return impl_->getOrCreateRdmaTransport(device_filter);
+}
+#endif
+
 bool TransferEngine::isTcpOnly() const {
     if (use_tent_)
         // TENT already rejects TCP loopback transfers when MC_STORE_MEMCPY
diff --git a/mooncake-transfer-engine/src/transfer_engine_c.cpp b/mooncake-transfer-engine/src/transfer_engine_c.cpp
index 246f1ab3..b643a153 100644
--- a/mooncake-transfer-engine/src/transfer_engine_c.cpp
+++ b/mooncake-transfer-engine/src/transfer_engine_c.cpp
@@ -203,8 +203,10 @@ notify_msg_t *getNotifsFromEngine(transfer_engine_t engine, int *size) {
     std::vector<TransferMetadata::NotifyDesc> notifies_desc;
     native->getNotifies(notifies_desc);
     *size = notifies_desc.size();
+    if (*size == 0) return nullptr;
     notify_msg_t *notifies =
         (notify_msg_t *)malloc(*size * sizeof(notify_msg_t));
+    if (!notifies) return nullptr;
     memset(notifies, 0, *size * sizeof(notify_msg_t));
     for (int i = 0; i < *size; i++) {
         notifies[i].name = (char *)malloc(notifies_desc[i].name.size() + 1);
diff --git a/mooncake-transfer-engine/src/transfer_engine_impl.cpp b/mooncake-transfer-engine/src/transfer_engine_impl.cpp
index 9caeffda..3bab6677 100644
--- a/mooncake-transfer-engine/src/transfer_engine_impl.cpp
+++ b/mooncake-transfer-engine/src/transfer_engine_impl.cpp
@@ -35,6 +35,14 @@
 
 namespace mooncake {
 
+namespace {
+bool overlapWithRegion(uintptr_t addr, uint64_t length, void* region_addr,
+                       uint64_t region_length) {
+    return overlap(reinterpret_cast<void*>(addr), length, region_addr,
+                   region_length);
+}
+}  // namespace
+
 static bool setFilesLimit() {
     struct rlimit filesLimit;
     if (getrlimit(RLIMIT_NOFILE, &filesLimit) != 0) {
@@ -194,6 +202,27 @@ int TransferEngineImpl::init(const std::string& metadata_conn_string,
     int ret = metadata_->addRpcMetaEntry(local_server_name_, desc);
     if (ret) return ret;
 
+    // Universal TCP force mechanism: if MC_FORCE_TCP is set, skip all other
+    // transport installation logic and use TCP transport only. This allows
+    // running metadata-only instances without requiring specialized hardware
+    // (e.g., NPU for Ascend Direct, RDMA HCAs, etc.).
+    if (getenv("MC_FORCE_TCP")) {
+#ifdef USE_TCP
+        Transport* tcp_transport =
+            multi_transports_->installTransport("tcp", nullptr);
+        if (!tcp_transport) {
+            LOG(ERROR)
+                << "MC_FORCE_TCP is set but failed to install TCP transport";
+            return -1;
+        }
+        LOG(INFO) << "MC_FORCE_TCP is set, using TCP transport only";
+        return 0;
+#else
+        LOG(ERROR) << "MC_FORCE_TCP is set but USE_TCP is not compiled in";
+        return -1;
+#endif
+    }
+
 #if defined(USE_ASCEND) || defined(USE_ASCEND_DIRECT)
     Transport* ascend_transport =
         multi_transports_->installTransport("ascend", local_topology_);
@@ -260,6 +289,43 @@ int TransferEngineImpl::init(const std::string& metadata_conn_string,
             LOG(ERROR) << "Failed to install Ascend transport";
             return -1;
         }
+#elif defined(USE_MACA)
+
+        if (getenv("MC_MACA_HOST_TRANSPORT")) {
+            if ((local_topology_->getHcaList().size() > 0 &&
+                 !getenv("MC_FORCE_TCP")) ||
+                getenv("MC_FORCE_HCA")) {
+                Transport* t = multi_transports_->installTransport(
+                    "rdma", local_topology_);
+                if (!t) {
+                    LOG(ERROR) << "Failed to install RDMA transport for MACA";
+                    return -1;
+                }
+                LOG(INFO) << "Using RDMA host transport for MACA";
+            } else {
+#ifdef USE_TCP
+                Transport* t =
+                    multi_transports_->installTransport("tcp", nullptr);
+                if (!t) {
+                    LOG(ERROR) << "Failed to install TCP transport for MACA";
+                    return -1;
+                }
+                LOG(INFO) << "Using TCP host transport for MACA";
+#else
+                LOG(ERROR)
+                    << "MC_MACA_HOST_TRANSPORT requires RDMA HCAs or USE_TCP";
+                return -1;
+#endif
+            }
+        } else {
+            Transport* t = multi_transports_->installTransport("maca", nullptr);
+            if (!t) {
+                LOG(ERROR) << "Failed to install MACA transport";
+                return -1;
+            }
+            LOG(INFO) << "Using MACA transport";
+        }
+
 #elif defined(USE_MNNVL) || defined(USE_INTRA_NVLINK)
 
         const char* force_mnnvl = getenv("MC_FORCE_MNNVL");
@@ -293,9 +359,12 @@ int TransferEngineImpl::init(const std::string& metadata_conn_string,
             LOG(INFO) << "Using RDMA transport (RoCE/iWARP)";
         }
 
-#else
-        if (local_topology_->getHcaList().size() > 0 &&
-                !getenv("MC_FORCE_TCP") ||
+#elif !defined(USE_SUNRISE)
+        // Sunrise classic installs its transport explicitly from tebench after
+        // benchmark-specific setup, so it skips the default auto transport
+        // path.
+        if ((local_topology_->getHcaList().size() > 0 &&
+             !getenv("MC_FORCE_TCP")) ||
             getenv("MC_FORCE_HCA")) {
             // only install RDMA transport when there is at least one HCA
             Transport* rdma_transport = nullptr;
@@ -384,7 +453,7 @@ Transport* TransferEngineImpl::installTransport(const std::string& proto,
     // shared lock here. If future modifications allow installTransport() to be
     // invoked concurrently, a std::shared_lock<std::shared_mutex> should be
     // added to ensure thread safety.
-    for (auto& entry : local_memory_regions_) {
+    for (auto& [_, entry] : local_memory_regions_) {
         int ret = transport->registerLocalMemory(
             entry.addr, entry.length, entry.location, entry.remote_accessible);
         if (ret < 0) return nullptr;
@@ -396,6 +465,24 @@ int TransferEngineImpl::uninstallTransport(const std::string& proto) {
     return 0;
 }
 
+#if (defined(USE_CUDA) || defined(USE_MUSA)) && !defined(USE_CXI)
+device::P2pTransport* TransferEngineImpl::getOrCreateP2pTransport(
+    int num_ranks) {
+    if (!p2p_transport_) {
+        p2p_transport_ = device::createP2pDeviceTransport(num_ranks);
+    }
+    return p2p_transport_.get();
+}
+
+device::RdmaTransport* TransferEngineImpl::getOrCreateRdmaTransport(
+    const std::vector<std::string>& device_filter) {
+    if (!rdma_transport_) {
+        rdma_transport_ = device::createIbgdaDeviceTransport(device_filter);
+    }
+    return rdma_transport_.get();
+}
+#endif
+
 int TransferEngineImpl::getRpcPort() {
     return metadata_->localRpcMeta().rpc_port;
 }
@@ -413,6 +500,10 @@ int TransferEngineImpl::getNotifies(
 int TransferEngineImpl::sendNotifyByID(
     SegmentID target_id, TransferMetadata::NotifyDesc notify_msg) {
     auto desc = metadata_->getSegmentDescByID(target_id);
+    if (!desc) {
+        LOG(ERROR) << "sendNotifyByID: invalid segment ID " << target_id;
+        return ERR_METADATA;
+    }
     Transport::NotifyDesc peer_desc;
     int ret = metadata_->sendNotify(desc->name, notify_msg, peer_desc);
     return ret;
@@ -488,13 +579,7 @@ int TransferEngineImpl::removeLocalSegment(const std::string& segment_name) {
 
 bool TransferEngineImpl::checkOverlap(void* addr, uint64_t length) {
     std::shared_lock<std::shared_mutex> lock(mutex_);
-    for (auto& local_memory_region : local_memory_regions_) {
-        if (overlap(addr, length, local_memory_region.addr,
-                    local_memory_region.length)) {
-            return true;
-        }
-    }
-    return false;
+    return hasOverlapLocked(reinterpret_cast<uintptr_t>(addr), length);
 }
 
 int TransferEngineImpl::registerLocalMemory(void* addr, size_t length,
@@ -518,8 +603,7 @@ int TransferEngineImpl::registerLocalMemory(void* addr, size_t length,
     }
 
     std::unique_lock<std::shared_mutex> lock(mutex_);
-    local_memory_regions_.push_back(
-        {addr, length, location, remote_accessible});
+    insertMemoryRegionLocked({addr, length, location, remote_accessible});
     return 0;
 }
 
@@ -531,13 +615,7 @@ int TransferEngineImpl::unregisterLocalMemory(void* addr,
     }
 
     std::unique_lock<std::shared_mutex> lock(mutex_);
-    for (auto it = local_memory_regions_.begin();
-         it != local_memory_regions_.end(); ++it) {
-        if (it->addr == addr) {
-            local_memory_regions_.erase(it);
-            break;
-        }
-    }
+    eraseMemoryRegionLocked(addr);
     return 0;
 }
 
@@ -611,9 +689,9 @@ int TransferEngineImpl::mp_registerLocalMemory(
     {
         std::unique_lock<std::shared_mutex> lock(mutex_);
         for (const auto& record : success_records) {
-            local_memory_regions_.push_back({record.addr, record.length,
-                                             record.location,
-                                             record.remote_accessible});
+            insertMemoryRegionLocked({record.addr, record.length,
+                                      record.location,
+                                      record.remote_accessible});
         }
     }
 
@@ -654,13 +732,7 @@ int TransferEngineImpl::mp_unregisterLocalMemory(
 
         std::unique_lock<std::shared_mutex> lock(mutex_);
         for (const auto& buffer : buffer_list) {
-            for (auto it = local_memory_regions_.begin();
-                 it != local_memory_regions_.end(); ++it) {
-                if (it->addr == buffer.addr) {
-                    local_memory_regions_.erase(it);
-                    break;
-                }
-            }
+            eraseMemoryRegionLocked(buffer.addr);
         }
     }
     return 0;
@@ -683,8 +755,7 @@ int TransferEngineImpl::registerLocalMemoryBatch(
 
     std::unique_lock<std::shared_mutex> lock(mutex_);
     for (auto& buffer : buffer_list) {
-        local_memory_regions_.push_back(
-            {buffer.addr, buffer.length, location, true});
+        insertMemoryRegionLocked({buffer.addr, buffer.length, location, true});
     }
     return 0;
 }
@@ -698,17 +769,74 @@ int TransferEngineImpl::unregisterLocalMemoryBatch(
 
     std::unique_lock<std::shared_mutex> lock(mutex_);
     for (auto& addr : addr_list) {
-        for (auto it = local_memory_regions_.begin();
-             it != local_memory_regions_.end(); ++it) {
-            if (it->addr == addr) {
-                local_memory_regions_.erase(it);
-                break;
-            }
-        }
+        eraseMemoryRegionLocked(addr);
     }
     return 0;
 }
 
+TransferEngineImpl::MemoryRegionMap::iterator
+TransferEngineImpl::findMemoryRegionContaining(uintptr_t addr) {
+    auto upper = local_memory_regions_.upper_bound(addr);
+    if (upper == local_memory_regions_.begin()) {
+        return local_memory_regions_.end();
+    }
+    auto candidate = std::prev(upper);
+    return overlapWithRegion(addr, 1, candidate->second.addr,
+                             candidate->second.length)
+               ? candidate
+               : local_memory_regions_.end();
+}
+
+TransferEngineImpl::MemoryRegionMap::const_iterator
+TransferEngineImpl::findMemoryRegionContaining(uintptr_t addr) const {
+    auto upper = local_memory_regions_.upper_bound(addr);
+    if (upper == local_memory_regions_.begin()) {
+        return local_memory_regions_.end();
+    }
+    auto candidate = std::prev(upper);
+    return overlapWithRegion(addr, 1, candidate->second.addr,
+                             candidate->second.length)
+               ? candidate
+               : local_memory_regions_.end();
+}
+
+bool TransferEngineImpl::hasOverlapLocked(uintptr_t addr,
+                                          uint64_t length) const {
+    if (length == 0) {
+        return false;
+    }
+
+    auto containing = findMemoryRegionContaining(addr);
+    if (containing != local_memory_regions_.end()) {
+        return true;
+    }
+
+    auto next = local_memory_regions_.lower_bound(addr);
+    if (next != local_memory_regions_.end() &&
+        overlapWithRegion(addr, length, next->second.addr,
+                          next->second.length)) {
+        return true;
+    }
+
+    if (next != local_memory_regions_.begin()) {
+        auto prev = std::prev(next);
+        if (overlapWithRegion(addr, length, prev->second.addr,
+                              prev->second.length)) {
+            return true;
+        }
+    }
+
+    return false;
+}
+
+void TransferEngineImpl::insertMemoryRegionLocked(const MemoryRegion& region) {
+    local_memory_regions_[reinterpret_cast<uintptr_t>(region.addr)] = region;
+}
+
+void TransferEngineImpl::eraseMemoryRegionLocked(void* addr) {
+    local_memory_regions_.erase(reinterpret_cast<uintptr_t>(addr));
+}
+
 #ifdef WITH_METRICS
 // Helper function to convert string to lowercase for case-insensitive
 // comparison
diff --git a/mooncake-transfer-engine/src/transfer_metadata.cpp b/mooncake-transfer-engine/src/transfer_metadata.cpp
index 7946c69f..88aab478 100644
--- a/mooncake-transfer-engine/src/transfer_metadata.cpp
+++ b/mooncake-transfer-engine/src/transfer_metadata.cpp
@@ -67,6 +67,8 @@ struct TransferHandshakeUtil {
     static Json::Value encode(const TransferMetadata::HandShakeDesc &desc) {
         Json::Value root;
         root["local_nic_path"] = desc.local_nic_path;
+        root["local_lid"] = desc.local_lid;
+        root["local_gid"] = desc.local_gid;
         root["peer_nic_path"] = desc.peer_nic_path;
 #ifdef USE_BAREX
         root["barex_port"] = desc.barex_port;
@@ -78,6 +80,9 @@ struct TransferHandshakeUtil {
 #ifdef USE_EFA
         root["efa_addr"] = desc.efa_addr;  // EFA endpoint address
 #endif
+#ifdef USE_CXI
+        root["cxi_addr"] = desc.cxi_addr;
+#endif
 
 #ifdef USE_UB
         Json::Value jettyNums(Json::arrayValue);
@@ -92,6 +97,16 @@ struct TransferHandshakeUtil {
 
     static int decode(Json::Value root, TransferMetadata::HandShakeDesc &desc) {
         desc.local_nic_path = root["local_nic_path"].asString();
+        if (root.isMember("local_lid") && root["local_lid"].isUInt()) {
+            desc.local_lid = root["local_lid"].asUInt();
+        } else {
+            desc.local_lid = 0;
+        }
+        if (root.isMember("local_gid") && root["local_gid"].isString()) {
+            desc.local_gid = root["local_gid"].asString();
+        } else {
+            desc.local_gid.clear();
+        }
         desc.peer_nic_path = root["peer_nic_path"].asString();
 #ifdef USE_BAREX
         desc.barex_port = root["barex_port"].asInt();
@@ -103,6 +118,10 @@ struct TransferHandshakeUtil {
         desc.efa_addr = root["efa_addr"].asString();  // EFA endpoint address
 #endif
 
+#ifdef USE_CXI
+        desc.cxi_addr = root["cxi_addr"].asString();
+#endif
+
 #ifdef USE_UB
         for (const auto &jetty : root["jetty_num"]) {
             desc.jetty_num.push_back(jetty.asUInt());
@@ -207,6 +226,9 @@ static int encodeMultiProtocolSegmentDesc(
     const TransferMetadata::SegmentDesc &desc, Json::Value &segmentJSON) {
     // Multi-protocol encoding for CXL+TCP or CXL+RDMA combination
     segmentJSON["name"] = desc.name;
+    if (!desc.rdma_server_name.empty()) {
+        segmentJSON["rdma_server_name"] = desc.rdma_server_name;
+    }
     Json::Value protocolJSON(Json::arrayValue);
     for (const auto &proto : protocols) {
         if (proto == "rdma") {
@@ -303,10 +325,13 @@ int TransferMetadata::encodeSegmentDesc(const SegmentDesc &desc,
     segmentJSON["protocol"] = desc.protocol;
     segmentJSON["tcp_data_port"] = desc.tcp_data_port;
     segmentJSON["timestamp"] = getCurrentDateTime();
+    if (!desc.rdma_server_name.empty()) {
+        segmentJSON["rdma_server_name"] = desc.rdma_server_name;
+    }
 
     if (segmentJSON["protocol"] == "rdma" ||
         segmentJSON["protocol"] == "barex" ||
-        segmentJSON["protocol"] == "efa") {
+        segmentJSON["protocol"] == "efa" || segmentJSON["protocol"] == "cxi") {
         Json::Value devicesJSON(Json::arrayValue);
         for (const auto &device : desc.devices) {
             Json::Value deviceJSON;
@@ -411,7 +436,9 @@ int TransferMetadata::encodeSegmentDesc(const SegmentDesc &desc,
     } else if (segmentJSON["protocol"] == "nvlink" ||
                segmentJSON["protocol"] == "nvlink_intra" ||
                segmentJSON["protocol"] == "hip" ||
-               segmentJSON["protocol"] == "ubshmem") {
+               segmentJSON["protocol"] == "maca" ||
+               segmentJSON["protocol"] == "ubshmem" ||
+               segmentJSON["protocol"] == "sunrise_link") {
         Json::Value buffersJSON(Json::arrayValue);
         for (const auto &buffer : desc.buffers) {
             Json::Value bufferJSON;
@@ -495,6 +522,8 @@ decodeMultiProtocolSegmentDesc(Json::Value &segmentJSON,
     desc->tcp_data_port = segmentJSON["tcp_data_port"].asInt();
     if (segmentJSON.isMember("timestamp"))
         desc->timestamp = segmentJSON["timestamp"].asString();
+    if (segmentJSON.isMember("rdma_server_name"))
+        desc->rdma_server_name = segmentJSON["rdma_server_name"].asString();
 
     for (const auto &protocolStr : segmentJSON["protocol"]) {
         std::string proto = protocolStr.asString();
@@ -552,9 +581,13 @@ decodeMultiProtocolSegmentDesc(Json::Value &segmentJSON,
             buffer.length = bufferJSON["length"].asUInt64();
             buffer.protocol = buffer_protocol;
             for (const auto &rkeyJSON : bufferJSON["rkey"])
-                buffer.rkey.push_back(rkeyJSON.asUInt());
+                buffer.rkey.push_back(
+                    static_cast<decltype(buffer.rkey)::value_type>(
+                        rkeyJSON.asUInt64()));
             for (const auto &lkeyJSON : bufferJSON["lkey"])
-                buffer.lkey.push_back(lkeyJSON.asUInt());
+                buffer.lkey.push_back(
+                    static_cast<decltype(buffer.lkey)::value_type>(
+                        lkeyJSON.asUInt64()));
             if (buffer.name.empty() || !buffer.addr || !buffer.length ||
                 buffer.rkey.empty() ||
                 buffer.rkey.size() != buffer.lkey.size()) {
@@ -638,9 +671,11 @@ TransferMetadata::decodeSegmentDesc(Json::Value &segmentJSON,
     desc->tcp_data_port = segmentJSON["tcp_data_port"].asInt();
     if (segmentJSON.isMember("timestamp"))
         desc->timestamp = segmentJSON["timestamp"].asString();
+    if (segmentJSON.isMember("rdma_server_name"))
+        desc->rdma_server_name = segmentJSON["rdma_server_name"].asString();
 
     if (desc->protocol == "rdma" || desc->protocol == "barex" ||
-        desc->protocol == "efa") {
+        desc->protocol == "efa" || desc->protocol == "cxi") {
         for (const auto &deviceJSON : segmentJSON["devices"]) {
             DeviceDesc device;
             device.name = deviceJSON["name"].asString();
@@ -660,9 +695,13 @@ TransferMetadata::decodeSegmentDesc(Json::Value &segmentJSON,
             buffer.addr = bufferJSON["addr"].asUInt64();
             buffer.length = bufferJSON["length"].asUInt64();
             for (const auto &rkeyJSON : bufferJSON["rkey"])
-                buffer.rkey.push_back(rkeyJSON.asUInt());
+                buffer.rkey.push_back(
+                    static_cast<decltype(buffer.rkey)::value_type>(
+                        rkeyJSON.asUInt64()));
             for (const auto &lkeyJSON : bufferJSON["lkey"])
-                buffer.lkey.push_back(lkeyJSON.asUInt());
+                buffer.lkey.push_back(
+                    static_cast<decltype(buffer.lkey)::value_type>(
+                        lkeyJSON.asUInt64()));
             if (buffer.name.empty() || !buffer.addr || !buffer.length ||
                 buffer.rkey.empty() ||
                 buffer.rkey.size() != buffer.lkey.size()) {
@@ -732,7 +771,9 @@ TransferMetadata::decodeSegmentDesc(Json::Value &segmentJSON,
             desc->buffers.push_back(buffer);
         }
     } else if (desc->protocol == "nvlink" || desc->protocol == "nvlink_intra" ||
-               desc->protocol == "hip" || desc->protocol == "ubshmem") {
+               desc->protocol == "hip" || desc->protocol == "maca" ||
+               desc->protocol == "ubshmem" ||
+               desc->protocol == "sunrise_link") {
         for (const auto &bufferJSON : segmentJSON["buffers"]) {
             BufferDesc buffer;
             buffer.name = bufferJSON["name"].asString();
@@ -883,7 +924,30 @@ std::shared_ptr<TransferMetadata::SegmentDesc> TransferMetadata::getSegmentDesc(
         }
     }
 
-    return decodeSegmentDesc(peer_json, segment_name);
+    auto result = decodeSegmentDesc(peer_json, segment_name);
+
+    // In P2P mode with dual-NIC setups (MC_RDMA_BIND_ADDRESS), the peer's
+    // segment descriptor may contain an rdma_server_name that differs from
+    // the TCP-routable segment_name. Cache the mapping so subsequent
+    // sendHandshake() calls can resolve the peer's TCP address from the
+    // RDMA server name extracted from NIC paths.
+    if (p2p_handshake_mode_ && result && !result->rdma_server_name.empty() &&
+        result->rdma_server_name != segment_name) {
+        auto [tcp_ip, tcp_port] = parseHostNameWithPort(segment_name);
+        RWSpinlock::WriteGuard guard(rpc_meta_lock_);
+        if (!rpc_meta_map_.count(result->rdma_server_name)) {
+            RpcMetaDesc meta;
+            meta.ip_or_host_name = tcp_ip;
+            meta.rpc_port = tcp_port;
+            meta.sockfd = -1;
+            rpc_meta_map_[result->rdma_server_name] = meta;
+            LOG(INFO) << "P2P: cached RDMA->TCP mapping: "
+                      << result->rdma_server_name << " -> " << tcp_ip << ":"
+                      << tcp_port;
+        }
+    }
+
+    return result;
 }
 
 int TransferMetadata::syncSegmentCache(const std::string &segment_name) {
@@ -1114,7 +1178,7 @@ int TransferMetadata::addRpcMetaEntry(const std::string &server_name,
 
     Json::Value rpcMetaJSON;
     rpcMetaJSON["ip_or_host_name"] = desc.ip_or_host_name;
-    rpcMetaJSON["rpc_port"] = static_cast<Json::UInt64>(desc.rpc_port);
+    rpcMetaJSON["rpc_port"] = static_cast<Json::UInt>(desc.rpc_port);
     if (!storage_plugin_->set(rpc_meta_prefix_ + server_name, rpcMetaJSON)) {
         LOG(ERROR) << "Failed to set location of " << server_name;
         return ERR_METADATA;
@@ -1133,6 +1197,32 @@ int TransferMetadata::removeRpcMetaEntry(const std::string &server_name) {
     return 0;
 }
 
+int TransferMetadata::rePublishRpcMetaEntry(const std::string &server_name) {
+    if (p2p_handshake_mode_) {
+        return 0;
+    }
+    const std::string full_key = rpc_meta_prefix_ + server_name;
+
+    Json::Value existing;
+    if (storage_plugin_->get(full_key, existing)) {
+        if (existing["ip_or_host_name"].asString() ==
+                local_rpc_meta_.ip_or_host_name &&
+            existing["rpc_port"].asUInt() == local_rpc_meta_.rpc_port) {
+            return 0;
+        }
+    }
+
+    LOG(INFO) << "Re-publishing RPC meta entry for " << server_name;
+    Json::Value rpcMetaJSON;
+    rpcMetaJSON["ip_or_host_name"] = local_rpc_meta_.ip_or_host_name;
+    rpcMetaJSON["rpc_port"] = static_cast<Json::UInt>(local_rpc_meta_.rpc_port);
+    if (!storage_plugin_->set(full_key, rpcMetaJSON)) {
+        LOG(ERROR) << "Failed to re-publish RPC meta entry for " << server_name;
+        return ERR_METADATA;
+    }
+    return 0;
+}
+
 int TransferMetadata::getRpcMetaEntry(const std::string &server_name,
                                       RpcMetaDesc &desc) {
     {
diff --git a/mooncake-transfer-engine/src/transfer_metadata_dump.cpp b/mooncake-transfer-engine/src/transfer_metadata_dump.cpp
index 4035cdfb..093cd502 100644
--- a/mooncake-transfer-engine/src/transfer_metadata_dump.cpp
+++ b/mooncake-transfer-engine/src/transfer_metadata_dump.cpp
@@ -18,24 +18,27 @@
 namespace mooncake {
 void TransferMetadata::SegmentDesc::dump() const {
     LOG(INFO) << "  segment name: " << name;
+    if (!rdma_server_name.empty() && rdma_server_name != name) {
+        LOG(INFO) << "  rdma server name: " << rdma_server_name;
+    }
     LOG(INFO) << "  protocol: " << protocol;
     LOG(INFO) << "  topology: " << topology.toString();
     LOG(INFO) << "  devices: ";
-    for (auto &device : devices) {
+    for (auto& device : devices) {
         LOG(INFO) << "    device name " << device.name << ", lid " << device.lid
                   << ", " << device.gid;
     }
     LOG(INFO) << "  buffers: ";
-    for (auto &buffer : buffers) {
+    for (auto& buffer : buffers) {
         LOG(INFO) << "    buffer type " << buffer.name << ", address "
-                  << (void *)buffer.addr << "--"
-                  << (void *)(buffer.addr + buffer.length);
+                  << (void*)buffer.addr << "--"
+                  << (void*)(buffer.addr + buffer.length);
     }
     LOG(INFO) << "  nvmeof buffers: " << nvmeof_buffers.size() << " items";
     LOG(INFO) << "  timestamp: " << timestamp;
 }
 
-void TransferMetadata::dumpMetadataContent(const std::string &segment_name,
+void TransferMetadata::dumpMetadataContent(const std::string& segment_name,
                                            uint64_t offset, uint64_t length) {
     thread_local uint64_t last_ts = 0;
     uint64_t current_ts = getCurrentTimeInNano();
@@ -44,13 +47,15 @@ void TransferMetadata::dumpMetadataContent(const std::string &segment_name,
     auto segment_locked = segment_lock_.tryLockShared();
     auto rpc_meta_locked = rpc_meta_lock_.tryLockShared();
     if (!segment_locked || !rpc_meta_locked) {
+        if (rpc_meta_locked) rpc_meta_lock_.unlockShared();
+        if (segment_locked) segment_lock_.unlockShared();
         return;
     }
 
     if (current_ts - last_ts > kMinDisplayThreshold || globalConfig().trace) {
         LOG(INFO) << "Failed to get segment descriptor for segment "
-                  << segment_name << " address " << (void *)offset << "--"
-                  << (void *)(offset + length);
+                  << segment_name << " address " << (void*)offset << "--"
+                  << (void*)(offset + length);
         dumpMetadataContentUnlocked();
         last_ts = current_ts;
     }
@@ -64,8 +69,8 @@ void TransferMetadata::dumpMetadataContentUnlocked() {
     LOG(INFO) << "TransferMetadata::dumpMetadataContent";
     LOG(INFO) << "-----------------------------------------------------------";
     LOG(INFO) << "=== Cached Segment Descriptors ===";
-    for (auto &entry : segment_id_to_desc_map_) {
-        auto &desc = entry.second;
+    for (auto& entry : segment_id_to_desc_map_) {
+        auto& desc = entry.second;
         if (!desc) {
             LOG(INFO) << "segment id: " << entry.first << ", ref object nil";
         } else {
@@ -78,7 +83,7 @@ void TransferMetadata::dumpMetadataContentUnlocked() {
     LOG(INFO) << "location: " << local_rpc_meta_.ip_or_host_name << ":"
               << local_rpc_meta_.rpc_port;
     LOG(INFO) << "=== Remote RPC Routes ===";
-    for (auto &entry : rpc_meta_map_) {
+    for (auto& entry : rpc_meta_map_) {
         LOG(INFO) << "segment name: " << entry.first
                   << ", location: " << entry.second.ip_or_host_name << ":"
                   << entry.second.rpc_port;
diff --git a/mooncake-transfer-engine/src/transfer_metadata_plugin.cpp b/mooncake-transfer-engine/src/transfer_metadata_plugin.cpp
index ab40c538..4df028d2 100644
--- a/mooncake-transfer-engine/src/transfer_metadata_plugin.cpp
+++ b/mooncake-transfer-engine/src/transfer_metadata_plugin.cpp
@@ -16,10 +16,12 @@
 
 #include <arpa/inet.h>
 #include <bits/stdint-uintn.h>
+#include <fcntl.h>
 #include <ifaddrs.h>
 #include <json/value.h>
 #include <net/if.h>
 #include <netdb.h>
+#include <poll.h>
 #include <sys/socket.h>
 
 #include <random>
@@ -856,6 +858,7 @@ struct SocketHandShakePlugin : public HandShakePlugin {
                 return 0;
             }
             if (ret == ERR_MALFORMED_JSON) {
+                freeaddrinfo(result);
                 return ret;
             }
         }
@@ -891,6 +894,7 @@ struct SocketHandShakePlugin : public HandShakePlugin {
                 return 0;
             }
             if (ret == ERR_MALFORMED_JSON) {
+                freeaddrinfo(result);
                 return ret;
             }
         }
@@ -926,6 +930,7 @@ struct SocketHandShakePlugin : public HandShakePlugin {
                 return 0;
             }
             if (ret == ERR_MALFORMED_JSON) {
+                freeaddrinfo(result);
                 return ret;
             }
         }
@@ -957,9 +962,77 @@ struct SocketHandShakePlugin : public HandShakePlugin {
             return ERR_SOCKET;
         }
 
+        // SO_RCVTIMEO does not apply to connect(). A blocking connect() to
+        // an unroutable address (e.g. a torn-down pod IP) stalls for the
+        // kernel's full SYN-retry cycle -- minutes -- and this runs on RDMA
+        // worker threads, where the stall also blocks CQ polling. Connect in
+        // non-blocking mode and bound the wait with poll().
+        int flags = fcntl(conn_fd, F_GETFL, 0);
+        if (flags == -1 || fcntl(conn_fd, F_SETFL, flags | O_NONBLOCK) == -1) {
+            PLOG(ERROR) << "SocketHandShakePlugin: fcntl(O_NONBLOCK)";
+            close(conn_fd);
+            return ERR_SOCKET;
+        }
+
         if (connect(conn_fd, addr->ai_addr, addr->ai_addrlen)) {
-            PLOG(ERROR) << "SocketHandShakePlugin: connect()"
-                        << getNetworkAddress(addr->ai_addr);
+            if (errno != EINPROGRESS) {
+                PLOG(ERROR) << "SocketHandShakePlugin: connect()"
+                            << getNetworkAddress(addr->ai_addr);
+                close(conn_fd);
+                return ERR_SOCKET;
+            }
+
+            const int64_t deadline_ms =
+                getCurrentTimeInMilli() +
+                globalConfig().handshake_connect_timeout * 1000;
+            struct pollfd pfd;
+            pfd.fd = conn_fd;
+            pfd.events = POLLOUT;
+            while (true) {
+                const int64_t remaining_ms =
+                    deadline_ms - getCurrentTimeInMilli();
+                // poll() returning 0 already means the timeout expired; an
+                // exhausted deadline (only reachable after EINTR) is the
+                // same condition.
+                int ret =
+                    remaining_ms <= 0 ? 0 : poll(&pfd, 1, (int)remaining_ms);
+                if (ret > 0) break;
+                if (ret == 0) {
+                    errno = ETIMEDOUT;
+                    PLOG(ERROR) << "SocketHandShakePlugin: connect() "
+                                << getNetworkAddress(addr->ai_addr);
+                    close(conn_fd);
+                    return ERR_SOCKET;
+                }
+                if (errno != EINTR) {
+                    PLOG(ERROR) << "SocketHandShakePlugin: poll()";
+                    close(conn_fd);
+                    return ERR_SOCKET;
+                }
+                // EINTR: retry with the remaining time.
+            }
+
+            int conn_err = 0;
+            socklen_t err_len = sizeof(conn_err);
+            if (getsockopt(conn_fd, SOL_SOCKET, SO_ERROR, &conn_err,
+                           &err_len)) {
+                PLOG(ERROR) << "SocketHandShakePlugin: getsockopt(SO_ERROR)";
+                close(conn_fd);
+                return ERR_SOCKET;
+            }
+            if (conn_err) {
+                errno = conn_err;
+                PLOG(ERROR) << "SocketHandShakePlugin: connect()"
+                            << getNetworkAddress(addr->ai_addr);
+                close(conn_fd);
+                return ERR_SOCKET;
+            }
+        }
+
+        // Restore blocking mode; the request/response exchange relies on
+        // blocking reads bounded by SO_RCVTIMEO.
+        if (fcntl(conn_fd, F_SETFL, flags) == -1) {
+            PLOG(ERROR) << "SocketHandShakePlugin: fcntl(restore flags)";
             close(conn_fd);
             return ERR_SOCKET;
         }
@@ -1034,6 +1107,7 @@ struct SocketHandShakePlugin : public HandShakePlugin {
                 return 0;
             }
             if (ret == ERR_MALFORMED_JSON) {
+                freeaddrinfo(result);
                 return ret;
             }
         }
diff --git a/mooncake-transfer-engine/src/transport/CMakeLists.txt b/mooncake-transfer-engine/src/transport/CMakeLists.txt
index 96b333e6..ecee46b5 100644
--- a/mooncake-transfer-engine/src/transport/CMakeLists.txt
+++ b/mooncake-transfer-engine/src/transport/CMakeLists.txt
@@ -2,68 +2,96 @@ file(GLOB XPORT_SOURCES "*.cpp")
 
 add_subdirectory(rdma_transport)
 add_subdirectory(rpc_communicator)
-add_library(transport OBJECT ${XPORT_SOURCES} $<TARGET_OBJECTS:rdma_transport> $<TARGET_OBJECTS:rpc_communicator>)
-target_link_libraries(transport PRIVATE JsonCpp::JsonCpp yalantinglibs::yalantinglibs glog::glog pthread)
+add_library(transport OBJECT ${XPORT_SOURCES} $<TARGET_OBJECTS:rdma_transport>
+                             $<TARGET_OBJECTS:rpc_communicator>)
+target_link_libraries(
+  transport PRIVATE JsonCpp::JsonCpp yalantinglibs::yalantinglibs glog::glog
+                    pthread)
 
 if(USE_UB)
   add_subdirectory(kunpeng_transport)
   target_sources(transport PUBLIC $<TARGET_OBJECTS:ub_transport>)
 endif()
 
-if (USE_TCP)
+if(USE_TCP)
   add_subdirectory(tcp_transport)
   target_sources(transport PUBLIC $<TARGET_OBJECTS:tcp_transport>)
 endif()
 
-if (USE_BAREX)
+if(USE_BAREX)
   add_subdirectory(barex_transport)
   target_sources(transport PUBLIC $<TARGET_OBJECTS:barex_transport>)
 endif()
 
-if (USE_NVMEOF)
+if(USE_NVMEOF)
   add_subdirectory(nvmeof_transport)
   target_sources(transport PUBLIC $<TARGET_OBJECTS:nvmeof_transport>)
 endif()
 
-if (USE_CXL)
+if(USE_CXL)
   add_subdirectory(cxl_transport)
   target_sources(transport PUBLIC $<TARGET_OBJECTS:cxl_transport>)
 endif()
 
-if (USE_ASCEND_DIRECT)
+if(USE_ASCEND_DIRECT)
   add_subdirectory(ascend_transport)
 elseif(USE_ASCEND)
   add_subdirectory(ascend_transport)
   target_sources(transport PUBLIC $<TARGET_OBJECTS:ascend_transport>)
 endif()
 
-if (USE_ASCEND_HETEROGENEOUS)
+if(USE_ASCEND_HETEROGENEOUS)
   add_subdirectory(ascend_transport)
   target_sources(transport PUBLIC $<TARGET_OBJECTS:ascend_transport>)
 endif()
 
-if (USE_HIP)
+if(USE_HIP)
   add_subdirectory(hip_transport)
   target_sources(transport PUBLIC $<TARGET_OBJECTS:hip_transport>)
 endif()
 
-if (USE_MNNVL AND NOT USE_HIP)
+if(USE_MACA)
+  add_subdirectory(maca_transport)
+  target_sources(transport PUBLIC $<TARGET_OBJECTS:maca_transport>)
+endif()
+
+if(USE_SUNRISE)
+  add_subdirectory(sunrise_link)
+  target_sources(transport PUBLIC $<TARGET_OBJECTS:sunrise_link_transport>)
+endif()
+
+if(USE_MNNVL AND NOT USE_HIP)
   add_subdirectory(nvlink_transport)
   target_sources(transport PUBLIC $<TARGET_OBJECTS:nvlink_transport>)
 endif()
 
-if (USE_INTRA_NVLINK)
-      add_subdirectory(intranode_nvlink_transport)
-      target_sources(transport PUBLIC $<TARGET_OBJECTS:intranode_nvlink_transport>)
+if(USE_INTRA_NVLINK)
+  add_subdirectory(intranode_nvlink_transport)
+  target_sources(transport PUBLIC $<TARGET_OBJECTS:intranode_nvlink_transport>)
 endif()
 
-if (USE_UBSHMEM)
+if(USE_UBSHMEM)
   add_subdirectory(ascend_transport)
   target_sources(transport PUBLIC $<TARGET_OBJECTS:ascend_transport>)
 endif()
 
-if (USE_EFA)
+if(USE_EFA)
   add_subdirectory(efa_transport)
   target_sources(transport PUBLIC $<TARGET_OBJECTS:efa_transport>)
   target_link_libraries(transport PRIVATE fabric)
 endif()
+
+if(USE_CXI)
+  add_subdirectory(cxi_transport)
+  target_sources(transport PUBLIC $<TARGET_OBJECTS:cxi_transport>)
+  target_link_libraries(transport PRIVATE fabric)
+endif()
+if(USE_CUDA OR USE_MUSA)
+  add_subdirectory(device)
+  target_sources(transport PUBLIC $<TARGET_OBJECTS:device_transport>)
+  # device_transport (ibgda_device_transport.cpp / mlx5gda.cpp) calls libmlx5
+  # DevX symbols (mlx5dv_devx_*, mlx5dv_init_obj) directly.
+  if(NOT USE_CXI)
+    target_link_libraries(transport PUBLIC mlx5)
+  endif()
+endif()
diff --git a/mooncake-transfer-engine/src/transport/ascend_transport/ascend_allocator.cpp b/mooncake-transfer-engine/src/transport/ascend_transport/ascend_allocator.cpp
index 9e3599bb..1071999a 100644
--- a/mooncake-transfer-engine/src/transport/ascend_transport/ascend_allocator.cpp
+++ b/mooncake-transfer-engine/src/transport/ascend_transport/ascend_allocator.cpp
@@ -66,10 +66,6 @@ int allocate_physical_memory(size_t total_size, aclrtDrvMemHandle &handle) {
 // Direct ACL VMM allocation (always bypasses adxl MallocMem).
 // Used by shm_helper when ascend_agent_mode && ascend_use_fabric_mem.
 void *allocate_vmm_memory_direct_impl(size_t total_size) {
-    if (total_size % kFabricMemPageSize != 0) {
-        LOG(ERROR) << "VMM memory size must be a multiple of 1GB";
-        return nullptr;
-    }
     aclrtDrvMemHandle handle = nullptr;
     if (allocate_physical_memory(total_size, handle) != 0) {
         return nullptr;
@@ -137,10 +133,6 @@ void *ascend_allocate_memory(size_t total_size, const std::string &protocol) {
     if (globalConfig().ascend_use_fabric_mem) {
         void *va = nullptr;
 #ifdef ASCEND_SUPPORT_FABRIC_MEM
-        if (total_size % kFabricMemPageSize != 0) {
-            LOG(ERROR) << "Local buffer size must be a multiple of 1GB";
-            return nullptr;
-        }
         if (&adxl::AdxlEngine::MallocMem != nullptr) {
             // Try to use adxl_engine's MallocMem
             auto status = adxl::AdxlEngine::MallocMem(adxl::MemType::MEM_HOST,
@@ -202,18 +194,19 @@ bool ascend_is_store_memory(void *addr, size_t length) {
 }
 
 void ascend_free_memory(const std::string &protocol, void *ptr) {
-    // make sure ascend_use_fabric_mem do not change between malloc and free
-    if (globalConfig().ascend_use_fabric_mem) {
+    // Route by the per-allocation record, not the process-global fabric flag:
+    // ascend_use_fabric_mem can differ between alloc time and free time (it is
+    // a process-global toggled during TE init), so it cannot be trusted here.
+    // Only fabric/VMM allocations are recorded in g_vmm_alloc_records;
+    // everything else is a plain aclrtMallocHost buffer.
 #ifdef ASCEND_SUPPORT_FABRIC_MEM
-        remove_store_memory_range(ptr);
-        std::lock_guard<std::mutex> lock(g_vmm_alloc_mutex);
-        auto it = g_vmm_alloc_records.find(ptr);
-        if (it == g_vmm_alloc_records.end()) {
-            LOG(INFO) << "Can not find record for va:" << ptr;
-            return;
-        }
+    std::unique_lock<std::mutex> lock(g_vmm_alloc_mutex);
+    auto it = g_vmm_alloc_records.find(ptr);
+    if (it != g_vmm_alloc_records.end()) {
         const AllocRecord alloc_record = it->second;
         g_vmm_alloc_records.erase(it);
+        lock.unlock();
+        remove_store_memory_range(ptr);
         if (!alloc_record.is_direct_alloc) {
             if (&adxl::AdxlEngine::FreeMem != nullptr) {
                 auto status = adxl::AdxlEngine::FreeMem(ptr);
@@ -239,9 +232,10 @@ void ascend_free_memory(const std::string &protocol, void *ptr) {
             LOG(ERROR) << "Failed to free physical mem: " << ptr;
             return;
         }
-#endif
         return;
     }
+    lock.unlock();
+#endif
 
     if (protocol == "ascend") {
         remove_store_memory_range(ptr);
diff --git a/mooncake-transfer-engine/src/transport/ascend_transport/ascend_direct_transport/ascend_direct_transport.cpp b/mooncake-transfer-engine/src/transport/ascend_transport/ascend_direct_transport/ascend_direct_transport.cpp
index 6c61ee0e..f5869b34 100644
--- a/mooncake-transfer-engine/src/transport/ascend_transport/ascend_direct_transport/ascend_direct_transport.cpp
+++ b/mooncake-transfer-engine/src/transport/ascend_transport/ascend_direct_transport/ascend_direct_transport.cpp
@@ -35,8 +35,8 @@
 namespace mooncake {
 namespace {
 
-int32_t ResolveCurrentEngineId(bool dummy_real_mode) {
-    if (!dummy_real_mode) {
+int32_t ResolveCurrentEngineId(bool agent_mode) {
+    if (!agent_mode) {
         return 0;
     }
     int32_t current_device_id = 0;
@@ -105,8 +105,9 @@ int AscendDirectTransport::install(std::string &local_server_name,
     TransferExecutorBase::InitParams exec_params;
     exec_params.metadata = metadata_;
     exec_params.local_engine_contexts = local_engine_contexts_;
-    exec_params.dummy_real_mode = dummy_real_mode_;
+    exec_params.agent_mode = agent_mode_;
     exec_params.roce_mode = roce_mode_;
+    exec_params.use_fabric_mem = use_fabric_mem_;
 
     transfer_executor_ = TransferExecutorBase::Create(exec_params);
     ret = transfer_executor_->initialize();
@@ -117,7 +118,7 @@ int AscendDirectTransport::install(std::string &local_server_name,
         return ret;
     }
 
-    if (dummy_real_mode_ && roce_mode_) {
+    if (agent_mode_ && roce_mode_) {
         dispatcher_ = std::make_unique<RoceDummyRealSliceDispatcher>(
             transfer_executor_.get(), local_engine_contexts_);
     } else {
@@ -148,19 +149,21 @@ int AscendDirectTransport::allocateLocalSegmentID() {
     desc->name = local_server_name_;
     desc->protocol = "ascend";
 
-    dummy_real_mode_ = globalConfig().ascend_agent_mode;
-    char *roce_enable_str = std::getenv("HCCL_INTRA_ROCE_ENABLE");
-    if (roce_enable_str) {
-        std::optional<int32_t> roce_enable =
-            parseFromString<int32_t>(roce_enable_str);
-        if (roce_enable.has_value() && roce_enable.value() == 1) {
-            roce_mode_ = true;
-            LOG(INFO) << "Roce mode is enabled.";
-        } else {
-            LOG(WARNING) << "HCCL_INTRA_ROCE_ENABLE is not valid, value:"
-                         << roce_enable_str;
-        }
-    }
+    agent_mode_ = globalConfig().ascend_agent_mode;
+    roce_mode_ = IsRoceModeEnabled();
+    // Only a Store-init TE may use fabric mem; gate on ascend_store_te_init so
+    // a P2P/HCCS TE does not inherit a Store TE's fabric flag left in the
+    // process-global config.
+    use_fabric_mem_ = globalConfig().ascend_use_fabric_mem &&
+                      globalConfig().ascend_store_te_init;
+    LOG(INFO) << "[AscendTE] init local segment, te is created for store="
+              << (globalConfig().ascend_store_te_init ? "true" : "false")
+              << ", roce_mode=" << (roce_mode_ ? "true" : "false")
+              << ", use_fabric_mem=" << (use_fabric_mem_ ? "true" : "false")
+              << (agent_mode_
+                      ? ", launched as standalone real client (manages all "
+                        "local NPU devices)"
+                      : "");
     char *adxl_base_port = std::getenv("ASCEND_BASE_PORT");
     if (adxl_base_port) {
         std::optional<int32_t> base_port =
@@ -175,9 +178,10 @@ int AscendDirectTransport::allocateLocalSegmentID() {
     }
     const auto [host_ip, port] = parseHostNameWithPort(local_server_name_);
     (void)port;
+    desc->rank_info.hostIp = host_ip;
     uint32_t device_count = 0;
     CHECK_ACL(aclrtGetDeviceCount(&device_count));
-    if (dummy_real_mode_) {
+    if (agent_mode_) {
         auto &ctx_mgr = ContextManager::getInstance();
         if (!ctx_mgr.isInitialized()) {
             LOG(ERROR) << "ContextManager is not initialized.";
@@ -224,7 +228,7 @@ Status AscendDirectTransport::submitTransfer(
             std::to_string(batch_id));
     }
 
-    const int32_t current_engine_id = ResolveCurrentEngineId(dummy_real_mode_);
+    const int32_t current_engine_id = ResolveCurrentEngineId(agent_mode_);
     if (current_engine_id < 0) {
         return Status::Context("aclrtGetDevice failed");
     }
@@ -252,7 +256,7 @@ Status AscendDirectTransport::submitTransfer(
 
 Status AscendDirectTransport::submitTransferTask(
     const std::vector<TransferTask *> &task_list) {
-    const int32_t current_engine_id = ResolveCurrentEngineId(dummy_real_mode_);
+    const int32_t current_engine_id = ResolveCurrentEngineId(agent_mode_);
     if (current_engine_id < 0) {
         return Status::Context("aclrtGetDevice failed");
     }
@@ -356,7 +360,7 @@ int AscendDirectTransport::registerLocalMemory(void *addr, size_t length,
 
     const int register_ret = transfer_executor_->registerMem(
         addr, length, mem_type, transfer_executor_->getUseBufferPool(),
-        roce_mode_, dummy_real_mode_);
+        roce_mode_, agent_mode_);
     if (register_ret == 0) {
         return 0;
     }
diff --git a/mooncake-transfer-engine/src/transport/ascend_transport/ascend_direct_transport/async_transfer_executor.cpp b/mooncake-transfer-engine/src/transport/ascend_transport/ascend_direct_transport/async_transfer_executor.cpp
index b0d44574..3455d121 100644
--- a/mooncake-transfer-engine/src/transport/ascend_transport/ascend_direct_transport/async_transfer_executor.cpp
+++ b/mooncake-transfer-engine/src/transport/ascend_transport/ascend_direct_transport/async_transfer_executor.cpp
@@ -14,21 +14,20 @@
 // limitations under the License.
 
 #include "transport/ascend_transport/ascend_direct_transport/async_transfer_executor.h"
-#include "transport/ascend_transport/ascend_direct_transport/utils.h"
 
 #include <glog/logging.h>
 
 #include <chrono>
+#include <utility>
 
 #include "common.h"
-#include "transfer_metadata.h"
 
 namespace mooncake {
 namespace {
 constexpr int64_t kMillisToNano = 1000000;
-constexpr size_t kAsyncTaskLimit = 100U;
+constexpr size_t kAsyncTaskLimit = 100;
 constexpr int32_t kDefaultDisconnectTime = 1000;
-constexpr int64_t kQueryPollIntervalMicros = 10;
+constexpr int64_t kQueryPollIntervalMicros = 1;
 }  // namespace
 
 AsyncTransferExecutor::AsyncTransferExecutor(const InitParams& params)
@@ -115,21 +114,24 @@ TransferExecutorBase::ExecuteResult AsyncTransferExecutor::execute(
         adxl::TransferArgs(), req_handle);
 
     if (status == adxl::SUCCESS) {
-        {
-            std::lock_guard<std::mutex> lock(async_handle_map_mutex_);
-            async_handle_to_engine_idx_[reinterpret_cast<uintptr_t>(
-                req_handle)] = local_engine_idx;
-        }
+        recordConnectedSegment(local_engine_idx, target_adxl_engine_name);
+        QueryBatch batch;
+        batch.slices = slice_list;
+        batch.engine_idx = local_engine_idx;
+        batch.target_adxl_engine_name = target_adxl_engine_name;
+
         for (auto* slice : slice_list) {
             slice->ascend_direct.handle = req_handle;
         }
+
         {
             std::unique_lock<std::mutex> lock(query_mutex_);
-            query_slice_queue_.push(slice_list);
+            query_slice_queue_.push(std::move(batch));
         }
         query_cv_.notify_one();
         return {.ret = 0, .status = status, .retryable = false};
     }
+
     {
         std::lock_guard<std::mutex> lock(async_task_mutex_);
         if (active_async_tasks_ > 0) {
@@ -137,27 +139,51 @@ TransferExecutorBase::ExecuteResult AsyncTransferExecutor::execute(
         }
         async_task_cv_.notify_one();
     }
-    disconnect(local_engine_idx, target_adxl_engine_name,
-               kDefaultDisconnectTime);
+    if (!params_.auto_connect) {
+        disconnect(local_engine_idx, target_adxl_engine_name,
+                   kDefaultDisconnectTime);
+    }
     return {.ret = -1, .status = status, .retryable = true};
 }
 
 void AsyncTransferExecutor::queryThreadLoop() {
-    std::vector<std::vector<Transport::Slice*>> pending_batches;
+    std::vector<QueryBatch> pending_batches;
+    pending_batches.reserve(64);
+
     while (running_) {
         fetchPendingBatches(pending_batches);
         if (pending_batches.empty()) {
             continue;
         }
 
-        auto it = pending_batches.begin();
-        while (it != pending_batches.end()) {
-            auto& slice_list = *it;
-            if (processOneBatch(slice_list)) {
-                it = pending_batches.erase(it);
-            } else {
-                ++it;
+        size_t i = 0;
+        while (i < pending_batches.size()) {
+            auto& batch = pending_batches[i];
+            BatchPollResult poll_result;
+            processOneBatch(batch, poll_result);
+            if (!poll_result.done) {
+                ++i;
+                continue;
             }
+
+            if (poll_result.fail_entire_route) {
+                const std::string& reason =
+                    poll_result.fail_reason.empty()
+                        ? "Route failed for target: " +
+                              batch.target_adxl_engine_name
+                        : poll_result.fail_reason;
+                failAllPendingOnRoute(batch.engine_idx,
+                                      batch.target_adxl_engine_name,
+                                      pending_batches, reason,
+                                      poll_result.adxl_already_disconnected);
+                // failAllPendingOnRoute swap-pops may move unvisited batches to
+                // indices < i; restart the scan so none are skipped this cycle.
+                i = 0;
+                continue;
+            }
+
+            std::swap(pending_batches[i], pending_batches.back());
+            pending_batches.pop_back();
         }
 
         if (!pending_batches.empty()) {
@@ -168,7 +194,7 @@ void AsyncTransferExecutor::queryThreadLoop() {
 }
 
 void AsyncTransferExecutor::fetchPendingBatches(
-    std::vector<std::vector<Transport::Slice*>>& out_batches) {
+    std::vector<QueryBatch>& out_batches) {
     std::unique_lock<std::mutex> lock(query_mutex_);
     if (out_batches.empty()) {
         query_cv_.wait(
@@ -183,110 +209,121 @@ void AsyncTransferExecutor::fetchPendingBatches(
     }
 }
 
-bool AsyncTransferExecutor::processOneBatch(
-    std::vector<Transport::Slice*>& slice_list) {
-    if (slice_list.empty()) {
-        return true;
+void AsyncTransferExecutor::markBatchFailed(QueryBatch& batch,
+                                            const std::string& reason,
+                                            bool log_error) {
+    if (log_error) {
+        LOG(ERROR) << reason;
+    }
+    for (auto* slice : batch.slices) {
+        slice->markFailed();
+    }
+    handleTaskFinished();
+}
+
+void AsyncTransferExecutor::failAllPendingOnRoute(
+    size_t engine_idx, const std::string& target,
+    std::vector<QueryBatch>& pending, const std::string& reason,
+    bool adxl_already_disconnected) {
+    LOG(ERROR) << reason;
+    bool cleaned_up = false;
+    for (size_t j = 0; j < pending.size();) {
+        auto& batch = pending[j];
+        if (batch.engine_idx != engine_idx ||
+            batch.target_adxl_engine_name != target) {
+            ++j;
+            continue;
+        }
+
+        markBatchFailed(batch, reason, false);
+        if (!cleaned_up && engine_idx < adxl_engines_.size() &&
+            !target.empty()) {
+            if (params_.auto_connect && adxl_already_disconnected) {
+                forgetConnectedSegment(engine_idx, target);
+            } else {
+                disconnect(engine_idx, target, params_.connect_timeout);
+            }
+            cleaned_up = true;
+        }
+
+        std::swap(pending[j], pending.back());
+        pending.pop_back();
+    }
+}
+
+void AsyncTransferExecutor::processOneBatch(QueryBatch& batch,
+                                            BatchPollResult& result) {
+    result = {};
+
+    if (batch.slices.empty()) {
+        result.done = true;
+        return;
+    }
+
+    if (batch.engine_idx >= adxl_engines_.size() ||
+        batch.engine_idx >= local_engine_contexts_.size()) {
+        markBatchFailed(
+            batch, "Invalid engine_idx: " + std::to_string(batch.engine_idx));
+        result.done = true;
+        return;
+    }
+
+    if (last_context_engine_idx_ != batch.engine_idx) {
+        auto context_ret =
+            aclrtSetCurrentContext(local_engine_contexts_[batch.engine_idx]);
+        if (context_ret != ACL_ERROR_NONE) {
+            markBatchFailed(
+                batch, "aclrtSetCurrentContext failed, ret: " +
+                           std::to_string(context_ret) +
+                           ", engine_idx: " + std::to_string(batch.engine_idx));
+            result.done = true;
+            return;
+        }
+        last_context_engine_idx_ = batch.engine_idx;
     }
 
     auto handle =
-        static_cast<adxl::TransferReq>(slice_list[0]->ascend_direct.handle);
-    size_t engine_idx = 0;
-    bool found_engine_idx = false;
-    {
-        std::lock_guard<std::mutex> lock(async_handle_map_mutex_);
-        auto it_handle = async_handle_to_engine_idx_.find(
-            reinterpret_cast<uintptr_t>(handle));
-        if (it_handle != async_handle_to_engine_idx_.end()) {
-            engine_idx = it_handle->second;
-            found_engine_idx = true;
-        }
-    }
-
-    auto fail_batch = [&](const std::string& message,
-                          const std::string& target_adxl_engine_name) {
-        LOG(ERROR) << message;
-        for (auto* slice : slice_list) {
-            slice->markFailed();
-        }
-        handleTaskFinished(handle);
-        if (engine_idx < adxl_engines_.size() &&
-            !target_adxl_engine_name.empty()) {
-            disconnect(engine_idx, target_adxl_engine_name,
-                       params_.connect_timeout);
-        }
-        return true;
-    };
-
-    if (!found_engine_idx) {
-        return fail_batch("Cannot resolve async handle to engine index", "");
-    }
-
-    if (engine_idx >= adxl_engines_.size() ||
-        engine_idx >= local_engine_contexts_.size()) {
-        return fail_batch(
-            "Invalid async engine index: " + std::to_string(engine_idx), "");
-    }
-
-    auto target_segment_desc =
-        metadata_->getSegmentDescByID(slice_list[0]->target_id);
-    if (!target_segment_desc) {
-        return fail_batch(
-            "Cannot find target segment descriptor for target_id: " +
-                std::to_string(slice_list[0]->target_id),
-            "");
-    }
-
-    std::string target_adxl_engine_name =
-        resolveTargetAdxlEngineName(target_segment_desc, engine_idx);
-    if (target_adxl_engine_name.empty()) {
-        return fail_batch("Cannot resolve target adxl engine name", "");
-    }
-
-    auto context_ret =
-        aclrtSetCurrentContext(local_engine_contexts_[engine_idx]);
-    if (context_ret != ACL_ERROR_NONE) {
-        return fail_batch("Call aclrtSetCurrentContext failed, ret: " +
-                              std::to_string(context_ret) +
-                              ", engine_idx: " + std::to_string(engine_idx),
-                          target_adxl_engine_name);
-    }
+        static_cast<adxl::TransferReq>(batch.slices[0]->ascend_direct.handle);
 
     adxl::TransferStatus task_status;
     auto ret =
-        adxl_engines_[engine_idx]->GetTransferStatus(handle, task_status);
+        adxl_engines_[batch.engine_idx]->GetTransferStatus(handle, task_status);
 
     if (ret != adxl::SUCCESS || task_status == adxl::TransferStatus::FAILED) {
-        return fail_batch("Get transfer status failed, ret: " +
-                              std::to_string(static_cast<int>(ret)) +
-                              ", errmsg: " + aclGetRecentErrMsg(),
-                          target_adxl_engine_name);
+        result.done = true;
+        result.fail_entire_route = true;
+        result.adxl_already_disconnected = true;
+        result.fail_reason = "Get transfer status failed, ret: " +
+                             std::to_string(static_cast<int>(ret)) +
+                             ", errmsg: " + aclGetRecentErrMsg();
+        return;
     }
 
     if (task_status == adxl::TransferStatus::COMPLETED) {
-        for (auto* slice : slice_list) {
+        for (auto* slice : batch.slices) {
             slice->markSuccess();
         }
-        handleTaskFinished(handle);
+        handleTaskFinished();
         if (params_.use_short_connection) {
-            disconnect(engine_idx, target_adxl_engine_name,
+            disconnect(batch.engine_idx, batch.target_adxl_engine_name,
                        params_.connect_timeout);
         }
-        return true;
+        result.done = true;
+        return;
     }
 
     auto now = getCurrentTimeInNano();
-    if (now - slice_list[0]->ascend_direct.start_time >
+    if (now - batch.slices[0]->ascend_direct.start_time >
         transfer_timeout_in_nano_) {
-        return fail_batch("Transfer timeout", target_adxl_engine_name);
+        result.done = true;
+        result.fail_entire_route = true;
+        result.fail_reason =
+            "Transfer timeout to: " + batch.target_adxl_engine_name;
+        return;
     }
-
-    return false;
 }
 
-void AsyncTransferExecutor::handleTaskFinished(adxl::TransferReq handle) {
-    std::lock_guard<std::mutex> handle_lock(async_handle_map_mutex_);
-    async_handle_to_engine_idx_.erase(reinterpret_cast<uintptr_t>(handle));
+void AsyncTransferExecutor::handleTaskFinished() {
     std::lock_guard<std::mutex> lock(async_task_mutex_);
     if (active_async_tasks_ > 0) {
         active_async_tasks_--;
diff --git a/mooncake-transfer-engine/src/transport/ascend_transport/ascend_direct_transport/sync_transfer_executor.cpp b/mooncake-transfer-engine/src/transport/ascend_transport/ascend_direct_transport/sync_transfer_executor.cpp
index 8e71dfc3..434235a2 100644
--- a/mooncake-transfer-engine/src/transport/ascend_transport/ascend_direct_transport/sync_transfer_executor.cpp
+++ b/mooncake-transfer-engine/src/transport/ascend_transport/ascend_direct_transport/sync_transfer_executor.cpp
@@ -76,11 +76,15 @@ TransferExecutorBase::ExecuteResult SyncTransferExecutor::execute(
                              op_descs, params_.transfer_timeout);
 
     if (status != adxl::SUCCESS) {
-        disconnect(local_engine_idx, target_adxl_engine_name,
-                   kDefaultDisconnectTime);
+        if (!params_.auto_connect) {
+            disconnect(local_engine_idx, target_adxl_engine_name,
+                       kDefaultDisconnectTime);
+        }
         return {.ret = -1, .status = status, .retryable = true};
     }
 
+    recordConnectedSegment(local_engine_idx, target_adxl_engine_name);
+
     if (params_.use_short_connection) {
         disconnect(local_engine_idx, target_adxl_engine_name,
                    params_.connect_timeout);
diff --git a/mooncake-transfer-engine/src/transport/ascend_transport/ascend_direct_transport/transfer_executor_base.cpp b/mooncake-transfer-engine/src/transport/ascend_transport/ascend_direct_transport/transfer_executor_base.cpp
index b3ab63aa..6f08a325 100644
--- a/mooncake-transfer-engine/src/transport/ascend_transport/ascend_direct_transport/transfer_executor_base.cpp
+++ b/mooncake-transfer-engine/src/transport/ascend_transport/ascend_direct_transport/transfer_executor_base.cpp
@@ -98,9 +98,18 @@ void TransferExecutorBase::ParseExecutorEnvIntoInitParams(InitParams& params) {
     char* auto_connect = std::getenv("ASCEND_AUTO_CONNECT");
     if (auto_connect) {
         auto auto_connect_opt = parseFromString<int32_t>(auto_connect);
-        if (auto_connect_opt.has_value() && *auto_connect_opt == 1) {
-            params.auto_connect = true;
+        if (!auto_connect_opt.has_value()) {
+            LOG(WARNING) << "ASCEND_AUTO_CONNECT is not valid, value:"
+                         << auto_connect << ", defaulting to disabled";
+            params.auto_connect = false;
+        } else {
+            params.auto_connect = (*auto_connect_opt == 1);
+            LOG(INFO) << "Set auto_connect from ASCEND_AUTO_CONNECT to: "
+                      << *auto_connect_opt;
         }
+    } else if (adxl::IsAdxlFeatureSupported(adxl::AUTO_CONNECT)) {
+        params.auto_connect = true;
+        LOG(INFO) << "AutoConnect enabled by capability probe";
     }
     char* buffer_pool = std::getenv("ASCEND_BUFFER_POOL");
     if (buffer_pool && std::strcmp(buffer_pool, "0:0") != 0) {
@@ -166,15 +175,8 @@ int TransferExecutorBase::initEngines() {
         LOG(INFO) << "Set LocalCommRes to:" << local_comm_res;
     }
 
-    char* auto_connect = std::getenv("ASCEND_AUTO_CONNECT");
-    if (auto_connect) {
-        auto auto_connect_opt = parseFromString<int32_t>(auto_connect);
-        if (auto_connect_opt.has_value()) {
-            options[kAutoConnect] =
-                (*auto_connect_opt == 1) ? kEnabled : kDisabled;
-            LOG(INFO) << "Set AutoConnect to: " << auto_connect;
-        }
-    }
+    options[kAutoConnect] = params_.auto_connect ? kEnabled : kDisabled;
+    LOG(INFO) << "Set AutoConnect to: " << (params_.auto_connect ? "1" : "0");
 
     options["adxl.BufferPool"] = "0:0";
     char* buffer_pool = std::getenv("ASCEND_BUFFER_POOL");
@@ -189,15 +191,24 @@ int TransferExecutorBase::initEngines() {
         }
     }
 
-    if (globalConfig().ascend_use_fabric_mem) {
+    if (params_.use_fabric_mem) {
         options["EnableUseFabricMem"] = "1";
         LOG(INFO) << "Fabric mem mode is enabled.";
     }
 
+    const char* store_te =
+        globalConfig().ascend_store_te_init ? "true" : "false";
     char* global_resource_config = std::getenv("ASCEND_GLOBAL_RESOURCE_CONFIG");
-    if (global_resource_config) {
-        options["GlobalResourceConfig"] = global_resource_config;
-        LOG(INFO) << "Set GlobalResourceConfig to:" << global_resource_config;
+    std::string resolved_resource_config =
+        ResolveAscendGlobalResourceConfig(global_resource_config);
+    if (!resolved_resource_config.empty()) {
+        options["GlobalResourceConfig"] = resolved_resource_config.c_str();
+        LOG(INFO) << "[AscendTE] init adxl, te is created for store="
+                  << store_te
+                  << ", GlobalResourceConfig: " << resolved_resource_config;
+    } else {
+        LOG(INFO) << "[AscendTE] init adxl, te is created for store="
+                  << store_te << ", GlobalResourceConfig unset.";
     }
 
     if (params_.use_async_transfer) {
@@ -270,18 +281,48 @@ int TransferExecutorBase::checkAndConnect(
     return 0;
 }
 
+void TransferExecutorBase::recordConnectedSegment(size_t engine_idx,
+                                                  const std::string& remote) {
+    if (!params_.auto_connect) {
+        return;
+    }
+    std::lock_guard<std::mutex> lock(connection_mutex_);
+    connected_segments_[engine_idx].insert(remote);
+}
+
+void TransferExecutorBase::forgetConnectedSegment(size_t engine_idx,
+                                                  const std::string& remote) {
+    std::lock_guard<std::mutex> lock(connection_mutex_);
+    auto it = connected_segments_.find(engine_idx);
+    if (it == connected_segments_.end()) {
+        return;
+    }
+    it->second.erase(remote);
+    if (it->second.empty()) {
+        connected_segments_.erase(it);
+    }
+}
+
 int TransferExecutorBase::disconnect(size_t engine_idx,
                                      const std::string& target_adxl_engine_name,
                                      int32_t timeout_in_millis) {
     if (params_.auto_connect) {
         auto status = adxl_engines_[engine_idx]->Disconnect(
             target_adxl_engine_name.c_str(), timeout_in_millis);
-        if (status != adxl::SUCCESS) {
+        if (status != adxl::SUCCESS && status != adxl::NOT_CONNECTED) {
             LOG(ERROR) << "Failed to disconnect to: " << target_adxl_engine_name
                        << ", status: " << status
                        << ", errmsg: " << aclGetRecentErrMsg();
             return -1;
         }
+        std::lock_guard<std::mutex> lock(connection_mutex_);
+        auto it = connected_segments_.find(engine_idx);
+        if (it != connected_segments_.end()) {
+            it->second.erase(target_adxl_engine_name);
+            if (it->second.empty()) {
+                connected_segments_.erase(it);
+            }
+        }
         return 0;
     }
     std::lock_guard<std::mutex> lock(connection_mutex_);
@@ -319,7 +360,7 @@ void TransferExecutorBase::disconnectAllForEngine(size_t engine_idx) {
     for (const auto& connected_segment : it->second) {
         auto status = adxl_engines_[engine_idx]->Disconnect(
             connected_segment.c_str(), params_.connect_timeout);
-        if (status != adxl::SUCCESS) {
+        if (status != adxl::SUCCESS && status != adxl::NOT_CONNECTED) {
             LOG(ERROR) << "Failed to disconnect AdxlEngine: "
                        << connected_segment
                        << ", errmsg: " << aclGetRecentErrMsg();
@@ -381,7 +422,7 @@ void TransferExecutorBase::rollbackRegisteredMem(
 int TransferExecutorBase::registerMem(void* addr, size_t length,
                                       adxl::MemType mem_type,
                                       bool use_buffer_pool, bool roce_mode,
-                                      bool dummy_real_mode) {
+                                      bool agent_mode) {
     if (mem_type == adxl::MEM_HOST && use_buffer_pool) {
         LOG(INFO) << "Ignore register host mem:" << addr
                   << " when buffer pool is enabled.";
@@ -403,7 +444,7 @@ int TransferExecutorBase::registerMem(void* addr, size_t length,
 
     std::vector<size_t> engine_indices;
     bool register_to_all =
-        (roce_mode && dummy_real_mode && ascend_is_store_memory(addr, length));
+        (roce_mode && agent_mode && ascend_is_store_memory(addr, length));
     if (register_to_all || adxl_engines_.size() == 1U) {
         engine_indices.resize(adxl_engines_.size());
         std::iota(engine_indices.begin(), engine_indices.end(), 0);
@@ -475,23 +516,31 @@ int TransferExecutorBase::deregisterMem(void* addr) {
         return 0;
     }
 
-    if (params_.dummy_real_mode) {
-        std::set<size_t> engine_indices;
-        for (const auto& [engine_idx, mem_handle] : it->second) {
-            (void)mem_handle;
-            engine_indices.insert(engine_idx);
-        }
-        for (size_t engine_idx : engine_indices) {
-            disconnectAllForEngine(engine_idx);
-        }
+    std::set<size_t> engine_indices;
+    for (const auto& [engine_idx, mem_handle] : it->second) {
+        (void)mem_handle;
+        engine_indices.insert(engine_idx);
+    }
+    for (size_t engine_idx : engine_indices) {
+        disconnectAllForEngine(engine_idx);
     }
 
+    bool deregister_failed = false;
     for (const auto& [engine_idx, mem_handle] : it->second) {
         if (engine_idx >= local_engine_contexts_.size()) {
             continue;
         }
         CHECK_ACL(aclrtSetCurrentContext(local_engine_contexts_[engine_idx]));
-        (void)adxl_engines_[engine_idx]->DeregisterMem(mem_handle);
+        auto status = adxl_engines_[engine_idx]->DeregisterMem(mem_handle);
+        if (status != adxl::SUCCESS) {
+            LOG(ERROR) << "DeregisterMem failed, engine_idx: " << engine_idx
+                       << ", status: " << status
+                       << ", errmsg: " << aclGetRecentErrMsg();
+            deregister_failed = true;
+        }
+    }
+    if (deregister_failed) {
+        return -1;
     }
     addr_to_mem_handles_.erase(it);
     return 0;
@@ -501,16 +550,74 @@ std::string TransferExecutorBase::resolveTargetAdxlEngineName(
     const std::shared_ptr<TransferMetadata::SegmentDesc>& segment_desc,
     size_t engine_idx) const {
     const auto& endpoints = segment_desc->rank_info.endpoints;
-    if (params_.dummy_real_mode && params_.roce_mode) {
-        if (engine_idx >= endpoints.size()) {
-            return {};
-        }
+    if (endpoints.empty()) return {};
+
+    // Standard dummy-real RoCE: same-index pairing (unchanged)
+    if (params_.agent_mode && params_.roce_mode) {
+        if (engine_idx >= endpoints.size()) return {};
         return endpoints[engine_idx];
     }
-    if (!endpoints.empty()) {
-        return endpoints.front();
+
+    // Standalone mode: non-dummy-real thin client without fabric mem.
+    // RoCE/HCCS all use phy_dev mapping; same-host offset +1 avoids
+    // connecting to the same physical device across processes.
+    if (!params_.agent_mode && !params_.use_fabric_mem) {
+        aclrtContext saved_ctx = nullptr;
+        if (aclrtGetCurrentContext(&saved_ctx) != ACL_ERROR_NONE) {
+            LOG(ERROR) << "aclrtGetCurrentContext failed in standalone resolve";
+            return endpoints.front();
+        }
+        MAKE_GUARD(ctx_restore,
+                   [saved_ctx]() { (void)aclrtSetCurrentContext(saved_ctx); });
+
+        int32_t logic_dev = 0;
+        if (engine_idx < local_engine_contexts_.size() &&
+            local_engine_contexts_[engine_idx] != nullptr) {
+            if (aclrtSetCurrentContext(local_engine_contexts_[engine_idx]) !=
+                ACL_ERROR_NONE) {
+                LOG(ERROR) << "aclrtSetCurrentContext failed in standalone "
+                              "resolve, engine_idx="
+                           << engine_idx;
+                return endpoints.front();
+            }
+        }
+        if (aclrtGetDevice(&logic_dev) != ACL_ERROR_NONE) {
+            LOG(ERROR) << "aclrtGetDevice failed in standalone resolve";
+            return endpoints.front();
+        }
+        if (logic_dev < 0) {
+            LOG(ERROR) << "Invalid logic device ID: " << logic_dev;
+            return endpoints.front();
+        }
+
+        int32_t phy_dev = logic_dev;
+        auto acl_ret = aclrtGetPhyDevIdByLogicDevId(logic_dev, &phy_dev);
+        if (acl_ret != ACL_ERROR_NONE) {
+            LOG(WARNING)
+                << "aclrtGetPhyDevIdByLogicDevId failed, using logic dev id";
+        }
+        if (phy_dev < 0) {
+            LOG(ERROR) << "Invalid physical device ID: " << phy_dev;
+            return endpoints.front();
+        }
+        size_t base_idx = static_cast<size_t>(phy_dev);
+
+        auto local_desc = metadata_->getSegmentDescByID(LOCAL_SEGMENT_ID);
+        const auto& remote_host_ip = segment_desc->rank_info.hostIp;
+        if (local_desc && !local_desc->rank_info.hostIp.empty() &&
+            !remote_host_ip.empty() &&
+            local_desc->rank_info.hostIp == remote_host_ip &&
+            endpoints.size() > 1) {
+            base_idx = (base_idx + 1) % endpoints.size();
+            VLOG(1) << "Standalone same-host offset: phy_dev=" << phy_dev
+                    << " -> target_idx=" << base_idx;
+        }
+        if (base_idx >= endpoints.size()) return endpoints.front();
+        return endpoints[base_idx];
     }
-    return {};
+
+    // Default: dummy-real non-RoCE or fabric-mem standalone
+    return endpoints.front();
 }
 
 void TransferExecutorBase::processSliceList(
@@ -519,7 +626,7 @@ void TransferExecutorBase::processSliceList(
         return;
     }
     size_t local_engine_idx =
-        params_.dummy_real_mode ? slice_list[0]->ascend_direct.engine_id : 0;
+        params_.agent_mode ? slice_list[0]->ascend_direct.engine_id : 0;
     VLOG(1) << "processSliceList for dev:" << local_engine_idx;
     auto local_segment_desc = metadata_->getSegmentDescByID(LOCAL_SEGMENT_ID);
     if (!local_segment_desc ||
@@ -564,7 +671,7 @@ void TransferExecutorBase::processSliceList(
             return;
         }
 
-        auto need_local_copy = !globalConfig().ascend_use_fabric_mem &&
+        auto need_local_copy = !params_.use_fabric_mem &&
                                (target_adxl_engine_name == local_engine_name);
         if (need_local_copy && local_copy_engine_) {
             auto start = std::chrono::steady_clock::now();
diff --git a/mooncake-transfer-engine/src/transport/ascend_transport/ascend_direct_transport/utils.cpp b/mooncake-transfer-engine/src/transport/ascend_transport/ascend_direct_transport/utils.cpp
index a022286f..e1e62372 100644
--- a/mooncake-transfer-engine/src/transport/ascend_transport/ascend_direct_transport/utils.cpp
+++ b/mooncake-transfer-engine/src/transport/ascend_transport/ascend_direct_transport/utils.cpp
@@ -15,8 +15,18 @@
 
 #include "transport/ascend_transport/ascend_direct_transport/utils.h"
 
+#include <cstring>
+#include <memory>
+#include <pthread.h>
 #include <random>
+#include <string>
+
 #include <glog/logging.h>
+#if __has_include(<jsoncpp/json/json.h>)
+#include <jsoncpp/json/json.h>  // Ubuntu
+#else
+#include <json/json.h>  // CentOS
+#endif
 
 #include "common.h"
 #include "config.h"
@@ -25,12 +35,139 @@ namespace mooncake {
 namespace {
 constexpr int32_t kPortRange = 100;
 constexpr int32_t kMaxGenPortAttempts = 500;
+constexpr const char* kProtocolDescFlatKey =
+    "comm_resource_config.protocol_desc";
+constexpr const char* kRocePrefix = "roce:";
+constexpr const char* kStoreConfigKey = "store";
+
+std::string SerializeCompactJson(const Json::Value& value) {
+    Json::StreamWriterBuilder writer;
+    writer["indentation"] = "";
+    return Json::writeString(writer, value);
+}
+
+bool IsRoceProtocolDesc(const std::string& desc) {
+    return desc.size() >= std::strlen(kRocePrefix) &&
+           desc.compare(0, std::strlen(kRocePrefix), kRocePrefix) == 0;
+}
+
+bool ProtocolDescValueContainsRoce(const Json::Value& value) {
+    if (value.isString()) {
+        return IsRoceProtocolDesc(value.asString());
+    }
+    if (value.isArray()) {
+        for (const auto& item : value) {
+            if (item.isString() && IsRoceProtocolDesc(item.asString())) {
+                return true;
+            }
+        }
+    }
+    return false;
+}
+
+const Json::Value* FindProtocolDescValue(const Json::Value& root) {
+    if (root.isMember(kProtocolDescFlatKey)) {
+        return &root[kProtocolDescFlatKey];
+    }
+    if (root.isMember("comm_resource_config") &&
+        root["comm_resource_config"].isObject() &&
+        root["comm_resource_config"].isMember("protocol_desc")) {
+        return &root["comm_resource_config"]["protocol_desc"];
+    }
+    return nullptr;
+}
 }  // namespace
 
+bool HasRoceProtocolDescInGlobalResourceConfig(const char* config_str) {
+    if (config_str == nullptr || config_str[0] == '\0') {
+        return false;
+    }
+    Json::CharReaderBuilder builder;
+    builder["collectComments"] = false;
+    Json::Value root;
+    std::string errs;
+    const std::unique_ptr<Json::CharReader> reader(builder.newCharReader());
+    if (!reader->parse(config_str, config_str + std::strlen(config_str), &root,
+                       &errs)) {
+        LOG(WARNING)
+            << "Failed to parse ASCEND_GLOBAL_RESOURCE_CONFIG as JSON: "
+            << errs;
+        return false;
+    }
+    if (!root.isObject()) {
+        LOG(WARNING) << "ASCEND_GLOBAL_RESOURCE_CONFIG must be a JSON object";
+        return false;
+    }
+    const Json::Value* protocol_desc = FindProtocolDescValue(root);
+    if (protocol_desc == nullptr) {
+        return false;
+    }
+    return ProtocolDescValueContainsRoce(*protocol_desc);
+}
+
+std::string ResolveAscendGlobalResourceConfig(const char* config_str) {
+    if (config_str == nullptr || config_str[0] == '\0') {
+        return std::string();
+    }
+    Json::CharReaderBuilder builder;
+    builder["collectComments"] = false;
+    Json::Value root;
+    std::string errs;
+    const std::unique_ptr<Json::CharReader> reader(builder.newCharReader());
+    if (!reader->parse(config_str, config_str + std::strlen(config_str), &root,
+                       &errs)) {
+        // Not valid JSON: leave it untouched and let adxl reject/handle it.
+        return std::string(config_str);
+    }
+    if (!root.isObject() || !root.isMember(kStoreConfigKey)) {
+        // Plain (legacy) config without a "store" override: pass verbatim.
+        return std::string(config_str);
+    }
+    if (globalConfig().ascend_store_te_init) {
+        return SerializeCompactJson(root[kStoreConfigKey]);
+    }
+    Json::Value normal = root;
+    normal.removeMember(kStoreConfigKey);
+    return SerializeCompactJson(normal);
+}
+
+bool IsRoceModeEnabled() {
+    const char* store_te =
+        globalConfig().ascend_store_te_init ? "true" : "false";
+    char* roce_enable_str = std::getenv("HCCL_INTRA_ROCE_ENABLE");
+    LOG(INFO) << "[AscendTE] resolving link config, te is created for store="
+              << store_te << ", HCCL_INTRA_ROCE_ENABLE="
+              << (roce_enable_str ? roce_enable_str : "(unset)");
+    if (roce_enable_str) {
+        const auto roce_enable = parseFromString<int32_t>(roce_enable_str);
+        if (roce_enable.has_value() && roce_enable.value() == 1) {
+            LOG(INFO) << "[AscendTE] roce mode enabled via "
+                         "HCCL_INTRA_ROCE_ENABLE.";
+            return true;
+        }
+        LOG(WARNING) << "HCCL_INTRA_ROCE_ENABLE is not valid, value:"
+                     << roce_enable_str;
+    }
+    char* global_resource_config = std::getenv("ASCEND_GLOBAL_RESOURCE_CONFIG");
+    std::string resolved =
+        ResolveAscendGlobalResourceConfig(global_resource_config);
+    LOG(INFO) << "[AscendTE] resolved ASCEND_GLOBAL_RESOURCE_CONFIG "
+                 "(protocol_desc): "
+              << (resolved.empty() ? "(none)" : resolved);
+    if (HasRoceProtocolDescInGlobalResourceConfig(resolved.c_str())) {
+        LOG(INFO) << "[AscendTE] roce mode enabled via "
+                     "ASCEND_GLOBAL_RESOURCE_CONFIG protocol_desc.";
+        return true;
+    }
+    return false;
+}
+
 // AscendThreadPool implementation
 AscendThreadPool::AscendThreadPool(size_t num_threads) : running_(true) {
     for (size_t i = 0; i < num_threads; ++i) {
-        workers_.emplace_back([this] {
+        workers_.emplace_back([this, i] {
+            auto thread_name = "ascend-wk-" + std::to_string(i);
+            pthread_setname_np(pthread_self(), thread_name.c_str());
             while (true) {
                 std::function<void()> task;
                 {
diff --git a/mooncake-transfer-engine/src/transport/cxi_transport/CMakeLists.txt b/mooncake-transfer-engine/src/transport/cxi_transport/CMakeLists.txt
new file mode 100644
index 00000000..51d41977
--- /dev/null
+++ b/mooncake-transfer-engine/src/transport/cxi_transport/CMakeLists.txt
@@ -0,0 +1,14 @@
+file(GLOB CXI_SOURCES "*.cpp")
+add_library(cxi_transport OBJECT ${CXI_SOURCES})
+
+# Link against libfabric (fabric) instead of ibverbs for AWS EFA
+target_link_libraries(cxi_transport PRIVATE fabric glog::glog)
+
+target_include_directories(
+  cxi_transport
+  PRIVATE
+    ${CMAKE_SOURCE_DIR}/mooncake-transfer-engine/include
+    ${CMAKE_SOURCE_DIR}/mooncake-transfer-engine/include/transport/cxi_transport
+    ${CMAKE_SOURCE_DIR}/mooncake-common/include ${LIBFABRIC_INCLUDE_DIR})
+
+# libfabric library path is set globally via common.cmake (LIBFABRIC_LIB_DIR)
diff --git a/mooncake-transfer-engine/src/transport/cxi_transport/cxi_context.cpp b/mooncake-transfer-engine/src/transport/cxi_transport/cxi_context.cpp
new file mode 100644
index 00000000..a1b9e791
--- /dev/null
+++ b/mooncake-transfer-engine/src/transport/cxi_transport/cxi_context.cpp
@@ -0,0 +1,967 @@
+// Copyright 2024 KVCache.AI
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "transport/cxi_transport/cxi_transport.h"
+
+#include <fcntl.h>
+#include <sys/epoll.h>
+
+#include <atomic>
+#include <cassert>
+#include <fstream>
+#include <memory>
+#include <thread>
+#include <cstring>
+#include <iomanip>
+#include <sstream>
+
+#include "config.h"
+#include "cuda_alike.h"
+#include "transport/cxi_transport/cxi_endpoint.h"
+#include "transport/cxi_transport/cxi_transport.h"
+#include "transport/transport.h"
+
+namespace mooncake {
+
+// CxiContext implementation
+CxiContext::CxiContext(CxiTransport& engine, const std::string& device_name)
+    : engine_(engine),
+      device_name_(device_name),
+      fi_info_(nullptr),
+      hints_(nullptr),
+      fabric_(nullptr),
+      domain_(nullptr),
+      av_(nullptr),
+      active_(true),
+      shared_ep_(nullptr),
+      wr_depth_(0),
+      max_wr_depth_(0),
+      post_lock_(ATOMIC_FLAG_INIT) {}
+
+CxiContext::~CxiContext() {
+    if (fabric_) {
+        deconstruct();
+    }
+}
+
+int CxiContext::construct(size_t num_cq_list, size_t max_cqe,
+                          int max_endpoints) {
+    // Setup hints for CXI provider
+    hints_ = fi_allocinfo();
+    if (!hints_) {
+        LOG(ERROR) << "Failed to allocate fi_info hints";
+        return ERR_CONTEXT;
+    }
+
+    hints_->caps =
+        FI_MSG | FI_RMA | FI_READ | FI_WRITE | FI_REMOTE_READ | FI_REMOTE_WRITE
+#if defined(USE_CUDA) || defined(USE_HIP)
+        | FI_HMEM  // must be added to the fabric caps to allow cuda/hip on cxi
+#endif
+        ;
+    hints_->mode = FI_CONTEXT;
+    hints_->ep_attr->type = FI_EP_RDM;  // CXI uses RDM endpoints
+    hints_->fabric_attr->prov_name = strdup("cxi");
+
+    // Specify the domain (device) name - append "-rdm" for RDM endpoint
+    std::string domain_name = device_name_;
+    // IMPORTANT: FI_MR_ENDPOINT must be specified when using CXI, otherwise
+    // fi_getinfo returns ENODATA CXI does not support FI_MR_VIRT_ADDR, so
+    // addresses must be offsets compared to MR base address
+    hints_->domain_attr->mr_mode = FI_MR_ENDPOINT | FI_MR_LOCAL |
+                                   FI_MR_ALLOCATED | FI_MR_PROV_KEY
+#if defined(USE_CUDA) || defined(USE_HIP)
+                                   | FI_MR_HMEM
+#endif
+        ;
+    hints_->domain_attr->name = strdup(domain_name.c_str());
+    hints_->domain_attr->threading = FI_THREAD_SAFE;
+
+    // Get fabric info
+    int ret = fi_getinfo(fi_version(), nullptr, nullptr, 0, hints_, &fi_info_);
+    if (ret) {
+        LOG(ERROR) << "fi_getinfo failed for device " << device_name_ << ": "
+                   << fi_strerror(-ret);
+        fi_freeinfo(hints_);
+        hints_ = nullptr;
+        return ERR_CONTEXT;
+    }
+
+    // Open fabric
+    ret = fi_fabric(fi_info_->fabric_attr, &fabric_, nullptr);
+    if (ret) {
+        LOG(ERROR) << "fi_fabric failed: " << fi_strerror(-ret);
+        fi_freeinfo(fi_info_);
+        fi_freeinfo(hints_);
+        fi_info_ = nullptr;
+        hints_ = nullptr;
+        return ERR_CONTEXT;
+    }
+
+    // Open domain
+    ret = fi_domain(fabric_, fi_info_, &domain_, nullptr);
+    if (ret) {
+        LOG(ERROR) << "fi_domain failed: " << fi_strerror(-ret);
+        fi_close(&fabric_->fid);
+        fi_freeinfo(fi_info_);
+        fi_freeinfo(hints_);
+        fabric_ = nullptr;
+        fi_info_ = nullptr;
+        hints_ = nullptr;
+        return ERR_CONTEXT;
+    }
+
+    // Create address vector
+    struct fi_av_attr av_attr = {};
+    av_attr.type = FI_AV_TABLE;
+    av_attr.count = max_endpoints;
+
+    ret = fi_av_open(domain_, &av_attr, &av_, nullptr);
+    if (ret) {
+        LOG(ERROR) << "fi_av_open failed: " << fi_strerror(-ret);
+        fi_close(&domain_->fid);
+        fi_close(&fabric_->fid);
+        fi_freeinfo(fi_info_);
+        fi_freeinfo(hints_);
+        domain_ = nullptr;
+        fabric_ = nullptr;
+        fi_info_ = nullptr;
+        hints_ = nullptr;
+        return ERR_CONTEXT;
+    }
+
+    // Create completion queues
+    cq_list_.resize(num_cq_list);
+    for (size_t i = 0; i < num_cq_list; ++i) {
+        auto cq = std::make_shared<CxiCq>();
+
+        struct fi_cq_attr cq_attr = {};
+        cq_attr.size = max_cqe;
+        cq_attr.format = FI_CQ_FORMAT_DATA;
+        cq_attr.wait_obj = FI_WAIT_NONE;
+
+        ret = fi_cq_open(domain_, &cq_attr, &cq->cq, nullptr);
+        if (ret) {
+            LOG(ERROR) << "fi_cq_open failed: " << fi_strerror(-ret);
+            return ERR_CONTEXT;
+        }
+        cq_list_[i] = cq;
+    }
+
+    ret = buildSharedEndpoint(globalConfig().max_wr, 16);
+    if (ret) {
+        LOG(ERROR) << "CxiContext::construct: buildSharedEndpoint failed for "
+                   << device_name_;
+        return ret;
+    }
+
+    LOG(INFO) << "CXI device (libfabric): " << device_name_
+              << ", domain: " << fi_info_->domain_attr->name
+              << ", provider: " << fi_info_->fabric_attr->prov_name;
+
+    return 0;
+}
+
+int CxiContext::buildSharedEndpoint(size_t max_wr, size_t max_inline) {
+    (void)max_inline;
+    if (shared_ep_) return 0;
+
+    shared_cq_ = cq_list_.empty() ? nullptr : cq_list_[0];
+    if (!shared_cq_) {
+        LOG(ERROR) << "CxiContext::buildSharedEndpoint: no CQ available";
+        return ERR_CONTEXT;
+    }
+    max_wr_depth_ = static_cast<int>(max_wr);
+
+    int ret = fi_endpoint(domain_, fi_info_, &shared_ep_, nullptr);
+    if (ret) {
+        LOG(ERROR) << "fi_endpoint failed: " << fi_strerror(-ret);
+        shared_ep_ = nullptr;
+        return ERR_ENDPOINT;
+    }
+
+    ret = fi_ep_bind(shared_ep_, &av_->fid, 0);
+    if (ret) {
+        LOG(ERROR) << "fi_ep_bind(av) failed: " << fi_strerror(-ret);
+        fi_close(&shared_ep_->fid);
+        shared_ep_ = nullptr;
+        return ERR_ENDPOINT;
+    }
+
+    ret = fi_ep_bind(shared_ep_, &shared_cq_->cq->fid, FI_TRANSMIT);
+    if (ret) {
+        LOG(ERROR) << "fi_ep_bind(tx_cq) failed: " << fi_strerror(-ret);
+        fi_close(&shared_ep_->fid);
+        shared_ep_ = nullptr;
+        return ERR_ENDPOINT;
+    }
+
+    ret = fi_ep_bind(shared_ep_, &shared_cq_->cq->fid, FI_RECV);
+    if (ret) {
+        LOG(ERROR) << "fi_ep_bind(rx_cq) failed: " << fi_strerror(-ret);
+        fi_close(&shared_ep_->fid);
+        shared_ep_ = nullptr;
+        return ERR_ENDPOINT;
+    }
+
+    ret = fi_enable(shared_ep_);
+    if (ret) {
+        LOG(ERROR) << "fi_enable failed: " << fi_strerror(-ret);
+        fi_close(&shared_ep_->fid);
+        shared_ep_ = nullptr;
+        return ERR_ENDPOINT;
+    }
+
+    // Cache our own libfabric address for handshake advertisement.
+    size_t addr_len = 8;
+    local_ep_addr_.assign(addr_len, 0);
+    ret = fi_getname(&shared_ep_->fid, local_ep_addr_.data(), &addr_len);
+    if (ret) {
+        LOG(ERROR) << "fi_getname failed: " << fi_strerror(-ret);
+        fi_close(&shared_ep_->fid);
+        shared_ep_ = nullptr;
+        local_ep_addr_.clear();
+        return ERR_ENDPOINT;
+    }
+    LOG(INFO) << "device " << deviceName()
+              << " ptr address: " << &shared_ep_->fid;
+    local_ep_addr_.resize(addr_len);
+    return 0;
+}
+
+int CxiContext::deconstruct() {
+    // Destroy all endpoints before closing domain/fabric/AV.
+    // Endpoints hold fi_ep handles that reference the domain, so they must
+    // be closed first.
+    if (shared_ep_) {
+        fi_close(&shared_ep_->fid);
+        shared_ep_ = nullptr;
+    }
+
+    {
+        RWSpinlock::WriteGuard guard(peer_map_lock_);
+        for (auto& entry : peer_map_) {
+            if (entry.second) entry.second->markDetachedForTeardown();
+        }
+        peer_map_.clear();
+    }
+
+    {
+        RWSpinlock::WriteGuard guard(mr_lock_);
+        for (auto& entry : mr_map_) {
+            if (entry.second.mr) {
+                fi_close(&entry.second.mr->fid);
+            }
+        }
+        mr_map_.clear();
+    }
+
+    for (auto& cq : cq_list_) {
+        if (cq && cq->cq) {
+            fi_close(&cq->cq->fid);
+            cq->cq = nullptr;
+        }
+    }
+    cq_list_.clear();
+    shared_cq_.reset();
+
+    if (av_) {
+        fi_close(&av_->fid);
+        av_ = nullptr;
+    }
+
+    if (domain_) {
+        fi_close(&domain_->fid);
+        domain_ = nullptr;
+    }
+
+    if (fabric_) {
+        fi_close(&fabric_->fid);
+        fabric_ = nullptr;
+    }
+
+    if (fi_info_) {
+        fi_freeinfo(fi_info_);
+        fi_info_ = nullptr;
+    }
+
+    if (hints_) {
+        fi_freeinfo(hints_);
+        hints_ = nullptr;
+    }
+
+    return 0;
+}
+
+int CxiContext::registerMemoryRegionInternal(void* addr, size_t length,
+                                             int access,
+                                             CxiMemoryRegionMeta& mrMeta) {
+    (void)access;
+    if (length > (size_t)globalConfig().max_mr_size) {
+        LOG(ERROR) << "Buffer length " << length
+                   << " exceeds device max_mr_size "
+                   << globalConfig().max_mr_size
+                   << ". Use CxiTransport::registerLocalMemory() which "
+                      "auto-splits large buffers.";
+        return ERR_CONTEXT;
+    } else if (!shared_ep_) {
+        LOG(ERROR)
+            << "shared ep must be built before registering memory regions!";
+    }
+
+    mrMeta.addr = addr;
+    mrMeta.length = length;
+
+    uint64_t fi_access = FI_READ | FI_WRITE | FI_REMOTE_READ | FI_REMOTE_WRITE;
+
+    enum fi_hmem_iface iface = FI_HMEM_SYSTEM;
+    int device_ordinal = 0;
+    int current_device = 0;
+
+#if defined(USE_CUDA)
+    cudaPointerAttributes attributes;
+    cudaError_t cuda_ret = cudaPointerGetAttributes(&attributes, addr);
+    if (cuda_ret == cudaSuccess && attributes.type == cudaMemoryTypeDevice) {
+        iface = FI_HMEM_CUDA;
+        device_ordinal = attributes.device;
+        cudaGetDevice(&current_device);
+        cudaSetDevice(device_ordinal);  // if by any chance we are setting a
+                                        // device pointer from a different
+                                        // device, fi_mr_regattr will fail
+    }
+#elif defined(USE_HIP)
+    hipPointerAttribute_t attributes;
+    hipError_t hip_ret = hipPointerGetAttributes(&attributes, addr);
+    if (hip_ret == hipSuccess && attributes.type == hipMemoryTypeDevice) {
+        iface = FI_HMEM_ROCR;
+        device_ordinal = attributes.device;
+        hipGetDevice(&current_device);
+        hipSetDevice(device_ordinal);
+    }
+#endif
+
+    int ret;
+    if (iface != FI_HMEM_SYSTEM) {
+        // GPU memory: use fi_mr_regattr with explicit iface and device
+        struct iovec iov = {.iov_base = addr, .iov_len = length};
+        struct fi_mr_attr attr = {};
+        attr.mr_iov = &iov;
+        attr.iov_count = 1;
+        attr.access = fi_access;
+        attr.iface = iface;
+        attr.device.cuda = device_ordinal;
+
+        ret = fi_mr_regattr(domain_, &attr, 0, &mrMeta.mr);
+        if (ret) {
+            LOG(ERROR) << "fi_mr_regattr failed for GPU memory " << addr
+                       << " (device " << device_ordinal
+                       << "): " << fi_strerror(-ret);
+#if defined(USE_CUDA)
+            cudaSetDevice(current_device);
+#elif defined(USE_HIP)
+            hipSetDevice(current_device);
+#endif
+            return ERR_CONTEXT;
+        }
+#if defined(USE_CUDA)
+        cudaSetDevice(current_device);
+#elif defined(USE_HIP)
+        hipSetDevice(current_device);
+#endif
+
+    } else {
+        // CPU memory: fi_mr_reg is sufficient
+        ret = fi_mr_reg(domain_, addr, length, fi_access, 0, 0, 0, &mrMeta.mr,
+                        nullptr);
+        if (ret) {
+            LOG(ERROR) << "fi_mr_reg failed for " << addr << ": "
+                       << fi_strerror(-ret);
+            return ERR_CONTEXT;
+        }
+    }
+
+    // cxi requires FI_MR_ENDPOINT
+    ret = fi_mr_bind(mrMeta.mr, &shared_ep_->fid, 0);
+    if (ret) {
+        LOG(ERROR) << "fi_mr_bind failed for " << addr << ": "
+                   << fi_strerror(-ret);
+        return ERR_CONTEXT;
+    }
+    ret = fi_mr_enable(mrMeta.mr);
+    if (ret) {
+        LOG(ERROR) << "fi_mr_enable failed for " << addr << ": "
+                   << fi_strerror(-ret);
+        return ERR_CONTEXT;
+    }
+    mrMeta.key = fi_mr_key(mrMeta.mr);
+    return 0;
+}
+
+int CxiContext::registerMemoryRegion(void* addr, size_t length, int access) {
+    CxiMemoryRegionMeta mrMeta;
+    int ret = registerMemoryRegionInternal(addr, length, access, mrMeta);
+    if (ret != 0) {
+        return ret;
+    }
+    {
+        RWSpinlock::WriteGuard guard(mr_lock_);
+        mr_map_[(uint64_t)addr] = mrMeta;
+    }
+    return 0;
+}
+
+int CxiContext::unregisterMemoryRegion(void* addr) {
+    RWSpinlock::WriteGuard guard(mr_lock_);
+    auto it = mr_map_.find((uint64_t)addr);
+    if (it == mr_map_.end()) {
+        return 0;
+    }
+    if (it->second.mr) {
+        fi_close(&it->second.mr->fid);
+    }
+    mr_map_.erase(it);
+    return 0;
+}
+
+int CxiContext::preTouchMemory(void* addr, size_t length) {
+    volatile char* ptr = (volatile char*)addr;
+    for (size_t i = 0; i < length; i += 4096) {
+        ptr[i] = ptr[i];
+    }
+    return 0;
+}
+
+uint64_t CxiContext::rkey(void* addr) {
+    RWSpinlock::ReadGuard guard(mr_lock_);
+    auto it = mr_map_.upper_bound((uint64_t)addr);
+    if (it != mr_map_.begin()) {
+        --it;
+        if ((uint64_t)addr < it->first + it->second.length && it->second.mr) {
+            return it->second.key;
+        }
+    }
+    return 0;
+}
+
+uint64_t CxiContext::lkey(void* addr) {
+    RWSpinlock::ReadGuard guard(mr_lock_);
+    auto it = mr_map_.upper_bound((uint64_t)addr);
+    if (it != mr_map_.begin()) {
+        --it;
+        if ((uint64_t)addr < it->first + it->second.length && it->second.mr) {
+            return it->second.key;
+        }
+    }
+    return 0;
+}
+
+void* CxiContext::mrDesc(void* addr) {
+    RWSpinlock::ReadGuard guard(mr_lock_);
+    // Find the MR that contains this address
+    auto it = mr_map_.upper_bound((uint64_t)addr);
+    if (it != mr_map_.begin()) {
+        --it;
+        if ((uint64_t)addr < it->first + it->second.length && it->second.mr) {
+            return fi_mr_desc(it->second.mr);
+        }
+    }
+    return nullptr;
+}
+
+std::shared_ptr<CxiEndpoint> CxiContext::endpoint(
+    const std::string& peer_nic_path) {
+    // Use normalized key (strip port) so the same physical peer reuses its
+    // handle across reconnections.  Each P2PHANDSHAKE run picks a random
+    // port, producing a different peer_nic_path for the same peer host+NIC.
+    std::string key = peer_nic_path;
+
+    {
+        RWSpinlock::ReadGuard guard(peer_map_lock_);
+        auto it = peer_map_.find(key);
+        if (it != peer_map_.end()) {
+            it->second->setPeerNicPath(peer_nic_path);
+            return it->second;
+        }
+    }
+
+    auto new_ep = std::make_shared<CxiEndpoint>(*this);
+    new_ep->setPeerNicPath(peer_nic_path);
+
+    RWSpinlock::WriteGuard guard(peer_map_lock_);
+    auto it = peer_map_.find(key);
+    if (it != peer_map_.end()) {
+        it->second->setPeerNicPath(peer_nic_path);
+        return it->second;
+    }
+    peer_map_[key] = new_ep;
+    return new_ep;
+}
+
+std::shared_ptr<CxiEndpoint> CxiContext::peekEndpoint(
+    const std::string& peer_nic_path) {
+    RWSpinlock::ReadGuard guard(peer_map_lock_);
+    auto it = peer_map_.find(peer_nic_path);
+    if (it == peer_map_.end()) return nullptr;
+    return it->second;
+}
+
+int CxiContext::deleteEndpoint(const std::string& peer_nic_path) {
+    std::shared_ptr<CxiEndpoint> ep;
+    {
+        RWSpinlock::WriteGuard guard(peer_map_lock_);
+        auto it = peer_map_.find(peer_nic_path);
+        if (it == peer_map_.end()) return 0;
+        ep = it->second;
+        peer_map_.erase(it);
+    }
+    if (ep) ep->disconnect();  // runs fi_av_remove
+    return 0;
+}
+
+int CxiContext::disconnectAllEndpoints() {
+    RWSpinlock::WriteGuard guard(peer_map_lock_);
+    for (auto& entry : peer_map_) {
+        if (entry.second) entry.second->disconnect();
+    }
+    return 0;
+}
+
+size_t CxiContext::getTotalQPNumber() const {
+    RWSpinlock::ReadGuard guard(peer_map_lock_);
+    return peer_map_.size();
+}
+
+std::string CxiContext::nicPath() const {
+    return engine_.local_server_name() + "@" + device_name_;
+}
+
+std::string CxiContext::localAddr() const {
+    // Return a hex string representation of the local address info
+    if (!fi_info_ || !fi_info_->src_addr) {
+        return "";
+    }
+
+    std::ostringstream oss;
+    const uint8_t* addr = static_cast<const uint8_t*>(fi_info_->src_addr);
+    for (size_t i = 0; i < fi_info_->src_addrlen; ++i) {
+        oss << std::hex << std::setw(2) << std::setfill('0') << (int)addr[i];
+    }
+    return oss.str();
+}
+
+std::string CxiContext::localEpAddr() const {
+    static constexpr char kHex[] = "0123456789abcdef";
+    std::string out;
+    out.resize(local_ep_addr_.size() * 2);
+    for (size_t i = 0; i < local_ep_addr_.size(); ++i) {
+        out[2 * i] = kHex[(local_ep_addr_[i] >> 4) & 0xF];
+        out[2 * i + 1] = kHex[local_ep_addr_[i] & 0xF];
+    }
+    return out;
+}
+
+// Decode one hex nibble, -1 on invalid input.
+static inline int hexNibble(char c) {
+    if (c >= '0' && c <= '9') return c - '0';
+    if (c >= 'a' && c <= 'f') return c - 'a' + 10;
+    if (c >= 'A' && c <= 'F') return c - 'A' + 10;
+    return -1;
+}
+
+int CxiContext::insertPeerAddr(const std::string& peer_hex_addr,
+                               fi_addr_t& out) {
+    if (peer_hex_addr.empty() || (peer_hex_addr.size() % 2) != 0) {
+        LOG(ERROR) << "insertPeerAddr: invalid hex length "
+                   << peer_hex_addr.size();
+        return ERR_INVALID_ARGUMENT;
+    }
+    const size_t n = peer_hex_addr.size() / 2;
+    std::vector<uint8_t> bin(n);
+    for (size_t i = 0; i < n; ++i) {
+        int hi = hexNibble(peer_hex_addr[2 * i]);
+        int lo = hexNibble(peer_hex_addr[2 * i + 1]);
+        if (hi < 0 || lo < 0) {
+            LOG(ERROR) << "insertPeerAddr: non-hex char at offset " << (2 * i);
+            return ERR_INVALID_ARGUMENT;
+        }
+        bin[i] = static_cast<uint8_t>((hi << 4) | lo);
+    }
+    return insertPeerAddrBytes(bin.data(), bin.size(), out);
+}
+
+int CxiContext::insertPeerAddrBytes(const uint8_t* addr, size_t len,
+                                    fi_addr_t& out) {
+    if (!addr || len == 0) return ERR_INVALID_ARGUMENT;
+    int ret = fi_av_insert(av_, addr, 1, &out, 0, nullptr);
+    if (ret != 1) {
+        LOG(ERROR) << "fi_av_insert failed: " << fi_strerror(-ret);
+        return ERR_ENDPOINT;
+    }
+    return 0;
+}
+
+void CxiContext::removePeerAddr(fi_addr_t fi_addr) {
+    if (fi_addr == FI_ADDR_UNSPEC) return;
+    int ret = fi_av_remove(av_, &fi_addr, 1, 0);
+    if (ret) {
+        LOG(WARNING) << "fi_av_remove failed: " << fi_strerror(-ret);
+    }
+}
+
+int CxiContext::submitPostSend(
+    const std::vector<Transport::Slice*>& slice_list) {
+    // Route slices to appropriate peer handles.  Group by peer NIC path.
+    std::unordered_map<std::string, std::vector<Transport::Slice*>>
+        slices_by_peer;
+
+    for (auto* slice : slice_list) {
+        if (!slice) continue;
+
+        // Fast path: peer info already filled in by the caller
+        // (dest_rkey and peer_nic_path set on the slice before dispatch).
+        if (!slice->peer_nic_path.empty()) {
+            slices_by_peer[slice->peer_nic_path].push_back(slice);
+            continue;
+        }
+
+        // Slow path: resolve peer info per-slice.
+        auto peer_segment_desc =
+            engine_.meta()->getSegmentDescByID(slice->target_id);
+        if (!peer_segment_desc) {
+            LOG(ERROR) << "Cannot get segment descriptor for target "
+                       << slice->target_id;
+            slice->markFailed();
+            continue;
+        }
+
+        int buffer_id = -1, device_id = -1;
+        if (CxiTransport::selectDevice(peer_segment_desc.get(),
+                                       slice->rdma.dest_addr, slice->length,
+                                       buffer_id, device_id)) {
+            // The cached peer descriptor may be stale: the remote can free
+            // and re-register buffers at new virtual addresses between
+            // iterations (e.g. a fresh per-iter H2D buffer). Force-refresh
+            // the segment descriptor once and retry before giving up.
+            // Mirrors the RDMA transport's retry-on-miss in worker_pool.cpp.
+            peer_segment_desc =
+                engine_.meta()->getSegmentDescByID(slice->target_id, true);
+            if (!peer_segment_desc ||
+                CxiTransport::selectDevice(peer_segment_desc.get(),
+                                           slice->rdma.dest_addr, slice->length,
+                                           buffer_id, device_id)) {
+                LOG(ERROR) << "Cannot select device for dest_addr "
+                           << (void*)slice->rdma.dest_addr;
+                slice->markFailed();
+                continue;
+            }
+        }
+
+        // no FI_VIRT_ADDR support on slingshot, must be offset of memory region
+        slice->rdma.dest_addr -= peer_segment_desc->buffers[buffer_id].addr;
+        slice->rdma.dest_rkey =
+            peer_segment_desc->buffers[buffer_id].rkey[device_id];
+
+        std::string peer_nic_path = peer_segment_desc->name + "@" +
+                                    peer_segment_desc->devices[device_id].name;
+        slice->peer_nic_path = peer_nic_path;
+        slices_by_peer[peer_nic_path].push_back(slice);
+    }
+
+    for (auto& entry : slices_by_peer) {
+        const std::string& peer_nic_path = entry.first;
+        auto& peer_slices = entry.second;
+
+        auto ep = endpoint(peer_nic_path);
+        if (!ep) {
+            LOG(ERROR) << "Cannot create peer handle for " << peer_nic_path;
+            for (auto* slice : peer_slices) slice->markFailed();
+            continue;
+        }
+
+        std::vector<Transport::Slice*> failed_slice_list;
+        int rc = ep->submitPostSend(peer_slices, failed_slice_list);
+        for (auto* slice : failed_slice_list) slice->markFailed();
+
+        // Drop peer handle if it is no longer connected after submit,
+        // freeing its AV entry for reuse.  Under the shared-endpoint model
+        // this is cheap (no fid_ep to destroy).
+        if (rc != 0 && !ep->connected()) {
+            deleteEndpoint(peer_nic_path);
+        }
+    }
+
+    return 0;
+}
+
+int CxiContext::submitSlicesOnPeer(
+    fi_addr_t peer_fi_addr, std::vector<Transport::Slice*>& slice_list,
+    std::vector<Transport::Slice*>& failed_slice_list) {
+    // Batched submission against the shared endpoint.  Mirrors the previous
+    // per-endpoint submit path but uses context-level wr_depth / post_lock.
+    //
+    // 1. Reserve N WR+CQ slots in bulk (single CAS each)
+    // 2. Prepare MR descriptors and op contexts outside the lock
+    // 3. Hold post_lock_ once for the entire batch of fi_write calls
+    const int kMaxBackoffYields = 100000;
+    const int cq_limit = static_cast<int>(globalConfig().max_cqe);
+    std::atomic<int>* cq_outstanding =
+        shared_cq_ ? &shared_cq_->outstanding : nullptr;
+
+    struct BatchEntry {
+        Transport::Slice* slice;
+        void* local_desc;
+        CxiOpContext* op_ctx;
+    };
+
+    // Consume slice_list via a moving index instead of erase-from-front,
+    // which was O(N^2) on large batches.  retry_slices accumulate across
+    // passes and are applied by rewinding the cursor.
+    size_t cursor = 0;
+    std::vector<Transport::Slice*> retry_slices;
+    while (cursor < slice_list.size() || !retry_slices.empty()) {
+        if (!retry_slices.empty()) {
+            // Splice retry slices back in at the current cursor so the next
+            // pass picks them up.  O(retry_slices.size()) per retry wave,
+            // which is bounded by valid_count of the last batch.
+            slice_list.insert(slice_list.begin() + cursor, retry_slices.begin(),
+                              retry_slices.end());
+            retry_slices.clear();
+            std::this_thread::yield();
+        }
+
+        const size_t remaining = slice_list.size() - cursor;
+        int batch_count = 0;
+        int backoff = 0;
+        bool timed_out = false;
+        while (batch_count == 0) {
+            int cur_wr = wr_depth_.load(std::memory_order_relaxed);
+            int wr_avail = max_wr_depth_ - cur_wr;
+            if (wr_avail <= 0) {
+                if (++backoff > kMaxBackoffYields) {
+                    timed_out = true;
+                    break;
+                }
+                std::this_thread::yield();
+                continue;
+            }
+            int want = std::min(wr_avail, (int)remaining);
+            if (cq_outstanding) {
+                int cur_cq = cq_outstanding->load(std::memory_order_relaxed);
+                int cq_avail = cq_limit - cur_cq;
+                if (cq_avail <= 0) {
+                    if (++backoff > kMaxBackoffYields) {
+                        timed_out = true;
+                        break;
+                    }
+                    std::this_thread::yield();
+                    continue;
+                }
+                want = std::min(want, cq_avail);
+                if (!wr_depth_.compare_exchange_weak(
+                        cur_wr, cur_wr + want, std::memory_order_acq_rel,
+                        std::memory_order_relaxed)) {
+                    continue;
+                }
+                cur_cq = cq_outstanding->load(std::memory_order_relaxed);
+                cq_avail = cq_limit - cur_cq;
+                if (cq_avail < want) {
+                    wr_depth_.fetch_sub(want, std::memory_order_acq_rel);
+                    continue;
+                }
+                if (!cq_outstanding->compare_exchange_weak(
+                        cur_cq, cur_cq + want, std::memory_order_acq_rel,
+                        std::memory_order_relaxed)) {
+                    wr_depth_.fetch_sub(want, std::memory_order_acq_rel);
+                    continue;
+                }
+            } else {
+                if (!wr_depth_.compare_exchange_weak(
+                        cur_wr, cur_wr + want, std::memory_order_acq_rel,
+                        std::memory_order_relaxed)) {
+                    continue;
+                }
+            }
+            batch_count = want;
+        }
+
+        if (timed_out) {
+            LOG(WARNING) << "CXI submitSlicesOnPeer: timed out waiting for CQ"
+                         << " drain (wr_depth="
+                         << wr_depth_.load(std::memory_order_relaxed)
+                         << ", max=" << max_wr_depth_ << ", cq_outstanding="
+                         << (cq_outstanding ? cq_outstanding->load(
+                                                  std::memory_order_relaxed)
+                                            : -1)
+                         << ", max_cqe=" << cq_limit << ")";
+            for (size_t i = cursor; i < slice_list.size(); ++i) {
+                failed_slice_list.push_back(slice_list[i]);
+            }
+            slice_list.clear();
+            return 0;
+        }
+
+        std::vector<BatchEntry> batch(batch_count);
+        int valid_count = 0;
+
+        for (int i = 0; i < batch_count; i++) {
+            Transport::Slice* slice = slice_list[cursor + i];
+            void* local_desc = mrDesc(slice->source_addr);
+            if (!local_desc) {
+                LOG(ERROR) << "No MR descriptor found for address "
+                           << slice->source_addr;
+                failed_slice_list.push_back(slice);
+                continue;
+            }
+            CxiOpContext* op_ctx = new CxiOpContext();
+            memset(op_ctx, 0, sizeof(CxiOpContext));
+            op_ctx->slice = slice;
+            op_ctx->wr_depth = &wr_depth_;
+            batch[valid_count++] = {slice, local_desc, op_ctx};
+        }
+
+        int mr_failures = batch_count - valid_count;
+        if (mr_failures > 0) {
+            wr_depth_.fetch_sub(mr_failures, std::memory_order_acq_rel);
+            if (cq_outstanding)
+                cq_outstanding->fetch_sub(mr_failures,
+                                          std::memory_order_acq_rel);
+        }
+
+        if (valid_count > 0) {
+            while (post_lock_.test_and_set(std::memory_order_acquire)) {
+            }
+            for (int i = 0; i < valid_count; i++) {
+                auto& entry = batch[i];
+                ssize_t ret;
+                if (entry.slice->opcode == Transport::TransferRequest::READ) {
+                    ret = fi_read(shared_ep_, (void*)entry.slice->source_addr,
+                                  entry.slice->length, entry.local_desc,
+                                  peer_fi_addr, entry.slice->rdma.dest_addr,
+                                  entry.slice->rdma.dest_rkey,
+                                  &entry.op_ctx->fi_ctx);
+                } else {
+                    ret = fi_write(shared_ep_, (void*)entry.slice->source_addr,
+                                   entry.slice->length, entry.local_desc,
+                                   peer_fi_addr, entry.slice->rdma.dest_addr,
+                                   entry.slice->rdma.dest_rkey,
+                                   &entry.op_ctx->fi_ctx);
+                }
+                if (ret == 0) {
+                    entry.slice->status = Transport::Slice::PENDING;
+                } else if (ret == -FI_EAGAIN) {
+                    delete entry.op_ctx;
+                    int not_posted = valid_count - i;
+                    wr_depth_.fetch_sub(not_posted, std::memory_order_acq_rel);
+                    if (cq_outstanding)
+                        cq_outstanding->fetch_sub(not_posted,
+                                                  std::memory_order_acq_rel);
+                    for (int j = i; j < valid_count; j++) {
+                        if (j > i) delete batch[j].op_ctx;
+                        retry_slices.push_back(batch[j].slice);
+                    }
+                    break;
+                } else {
+                    LOG(ERROR)
+                        << "fi_read/fi_write failed: " << fi_strerror(-ret)
+                        << " (source=" << entry.slice->source_addr
+                        << ", len=" << entry.slice->length
+                        << ", dest=" << (void*)entry.slice->rdma.dest_addr
+                        << ", rkey=" << entry.slice->rdma.dest_rkey << ")";
+                    delete entry.op_ctx;
+                    wr_depth_.fetch_sub(1, std::memory_order_acq_rel);
+                    if (cq_outstanding)
+                        cq_outstanding->fetch_sub(1, std::memory_order_acq_rel);
+                    failed_slice_list.push_back(entry.slice);
+                }
+            }
+            post_lock_.clear(std::memory_order_release);
+        }
+
+        cursor += batch_count;
+    }
+
+    slice_list.clear();
+    return 0;
+}
+
+int CxiContext::pollCq(int max_entries, int cq_index) {
+    if (cq_index < 0 || (size_t)cq_index >= cq_list_.size()) {
+        return 0;
+    }
+
+    struct fid_cq* cq = cq_list_[cq_index]->cq;
+    if (!cq) return 0;
+
+    struct fi_cq_data_entry entries[64];
+    int to_poll = std::min(max_entries, 64);
+
+    ssize_t ret = fi_cq_read(cq, entries, to_poll);
+
+    if (ret > 0) {
+        std::unordered_map<std::atomic<int>*, int> wr_depth_set;
+        for (ssize_t i = 0; i < ret; i++) {
+            CxiOpContext* op_ctx =
+                reinterpret_cast<CxiOpContext*>(entries[i].op_context);
+            if (op_ctx && op_ctx->slice) {
+                op_ctx->slice->markSuccess();
+                if (op_ctx->wr_depth) {
+                    wr_depth_set[op_ctx->wr_depth]++;
+                }
+                delete op_ctx;
+            }
+        }
+        for (auto& entry : wr_depth_set) {
+            entry.first->fetch_sub(entry.second, std::memory_order_acq_rel);
+        }
+        cq_list_[cq_index]->outstanding.fetch_sub(static_cast<int>(ret),
+                                                  std::memory_order_acq_rel);
+        return static_cast<int>(ret);
+    } else if (ret == -FI_EAGAIN) {
+        return 0;
+    } else if (ret < 0) {
+        int err_count = 0;
+        struct fi_cq_err_entry err_entry;
+        std::unordered_map<std::atomic<int>*, int> wr_depth_set;
+
+        while ((ret = fi_cq_readerr(cq, &err_entry, 0)) > 0) {
+            CxiOpContext* op_ctx =
+                reinterpret_cast<CxiOpContext*>(err_entry.op_context);
+            if (op_ctx && op_ctx->slice) {
+                LOG(ERROR) << "CXI CQ error: "
+                           << fi_cq_strerror(cq, err_entry.prov_errno,
+                                             err_entry.err_data, nullptr, 0)
+                           << " for slice at " << op_ctx->slice->source_addr;
+                op_ctx->slice->markFailed();
+                if (op_ctx->wr_depth) {
+                    wr_depth_set[op_ctx->wr_depth]++;
+                }
+                delete op_ctx;
+            }
+            err_count++;
+        }
+
+        for (auto& entry : wr_depth_set) {
+            entry.first->fetch_sub(entry.second, std::memory_order_acq_rel);
+        }
+        if (err_count > 0) {
+            cq_list_[cq_index]->outstanding.fetch_sub(
+                err_count, std::memory_order_acq_rel);
+        }
+        return err_count;
+    }
+
+    return 0;
+}
+
+}  // namespace mooncake
diff --git a/mooncake-transfer-engine/src/transport/cxi_transport/cxi_endpoint.cpp b/mooncake-transfer-engine/src/transport/cxi_transport/cxi_endpoint.cpp
new file mode 100644
index 00000000..f17326e6
--- /dev/null
+++ b/mooncake-transfer-engine/src/transport/cxi_transport/cxi_endpoint.cpp
@@ -0,0 +1,186 @@
+// Copyright 2024 KVCache.AI
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "transport/cxi_transport/cxi_endpoint.h"
+
+#include <glog/logging.h>
+
+#include <cassert>
+#include <cstddef>
+#include <cstring>
+#include <iomanip>
+#include <sstream>
+#include <thread>
+
+#include "config.h"
+
+namespace mooncake {
+
+CxiEndpoint::CxiEndpoint(CxiContext& context)
+    : context_(context), status_(INITIALIZING), peer_fi_addr_(FI_ADDR_UNSPEC) {}
+
+CxiEndpoint::~CxiEndpoint() { disconnect(); }
+
+void CxiEndpoint::setPeerNicPath(const std::string& peer_nic_path) {
+    RWSpinlock::WriteGuard guard(lock_);
+    if (peer_nic_path_ == peer_nic_path) return;  // No change
+    if (status_.load(std::memory_order_relaxed) == CONNECTED) {
+        LOG(INFO) << "Peer reconnected with new address, re-establishing: "
+                  << peer_nic_path_ << " -> " << peer_nic_path;
+        disconnectUnlocked();
+    }
+    peer_nic_path_ = peer_nic_path;
+}
+
+int CxiEndpoint::setupConnectionsByActive() {
+    RWSpinlock::WriteGuard guard(lock_);
+    if (status_.load(std::memory_order_relaxed) == CONNECTED) return 0;
+
+    // Loopback: handshake against ourselves.  Use the binary overload
+    // so we avoid hex-encoding the local address just to have
+    // insertPeerAddr decode it right back.
+    if (context_.nicPath() == peer_nic_path_) {
+        const auto& bytes = context_.localEpAddrBytes();
+        int ret = context_.insertPeerAddrBytes(bytes.data(), bytes.size(),
+                                               peer_fi_addr_);
+        if (ret != 0) return ret;
+        status_.store(CONNECTED, std::memory_order_release);
+        LOG(INFO) << "CXI loopback connection established: " << toString();
+        return 0;
+    }
+
+    // Exchange addresses via the transfer-metadata handshake RPC.
+    TransferMetadata::HandShakeDesc local_desc, peer_desc;
+    local_desc.local_nic_path = context_.nicPath();
+    local_desc.peer_nic_path = peer_nic_path_;
+    local_desc.cxi_addr = context_.localEpAddr();
+
+    auto peer_server_name = getServerNameFromNicPath(peer_nic_path_);
+    auto peer_nic_name = getNicNameFromNicPath(peer_nic_path_);
+    if (peer_server_name.empty() || peer_nic_name.empty()) {
+        LOG(ERROR) << "Parse peer CXI nic path failed: " << peer_nic_path_;
+        return ERR_INVALID_ARGUMENT;
+    }
+
+    int rc = context_.engine().sendHandshake(peer_server_name, local_desc,
+                                             peer_desc);
+    if (rc) return rc;
+
+    if (peer_desc.cxi_addr.empty()) {
+        LOG(ERROR) << "Peer did not provide CXI address in handshake";
+        return ERR_REJECT_HANDSHAKE;
+    }
+
+    rc = context_.insertPeerAddr(peer_desc.cxi_addr, peer_fi_addr_);
+    if (rc != 0) return rc;
+
+    status_.store(CONNECTED, std::memory_order_release);
+    VLOG(1) << "CXI connection established: " << toString()
+            << " peer_fi_addr=" << peer_fi_addr_;
+    return 0;
+}
+
+int CxiEndpoint::setupConnectionsByPassive(const HandShakeDesc& peer_desc,
+                                           HandShakeDesc& local_desc) {
+    RWSpinlock::WriteGuard guard(lock_);
+    if (status_.load(std::memory_order_relaxed) == CONNECTED) {
+        LOG(WARNING) << "Re-establish CXI connection: " << toString();
+        disconnectUnlocked();
+    }
+
+    if (peer_desc.peer_nic_path != context_.nicPath() ||
+        peer_desc.local_nic_path != peer_nic_path_) {
+        local_desc.reply_msg = "CXI nic path inconsistency";
+        LOG(ERROR) << "Invalid argument: peer CXI nic path inconsistency"
+                   << " peer_nic_path=" << peer_desc.peer_nic_path
+                   << " context_.nicPath()=" << context_.nicPath()
+                   << " local_nic_path=" << peer_desc.local_nic_path
+                   << " peer_nic_path_=" << peer_nic_path_;
+        return ERR_REJECT_HANDSHAKE;
+    }
+
+    if (peer_desc.cxi_addr.empty()) {
+        local_desc.reply_msg = "No CXI address provided";
+        LOG(ERROR) << "Peer did not provide CXI address";
+        return ERR_REJECT_HANDSHAKE;
+    }
+
+    int ret = context_.insertPeerAddr(peer_desc.cxi_addr, peer_fi_addr_);
+    if (ret != 0) {
+        local_desc.reply_msg = "Failed to insert peer address";
+        return ret;
+    }
+
+    local_desc.local_nic_path = context_.nicPath();
+    local_desc.peer_nic_path = peer_nic_path_;
+    local_desc.cxi_addr = context_.localEpAddr();
+    // reply_msg empty on success
+
+    status_.store(CONNECTED, std::memory_order_release);
+    VLOG(1) << "CXI connection established (passive): " << toString();
+    return 0;
+}
+
+const std::string CxiEndpoint::toString() const {
+    return "CxiEndpoint[" + context_.nicPath() + " <-> " + peer_nic_path_ + "]";
+}
+
+void CxiEndpoint::disconnect() {
+    RWSpinlock::WriteGuard guard(lock_);
+    disconnectUnlocked();
+}
+
+void CxiEndpoint::disconnectUnlocked() {
+    if (peer_fi_addr_ != FI_ADDR_UNSPEC) {
+        context_.removePeerAddr(peer_fi_addr_);
+        peer_fi_addr_ = FI_ADDR_UNSPEC;
+    }
+    status_.store(UNCONNECTED, std::memory_order_release);
+}
+
+void CxiEndpoint::markDetachedForTeardown() {
+    RWSpinlock::WriteGuard guard(lock_);
+    peer_fi_addr_ = FI_ADDR_UNSPEC;
+    status_.store(UNCONNECTED, std::memory_order_release);
+}
+
+int CxiEndpoint::submitPostSend(
+    std::vector<Transport::Slice*>& slice_list,
+    std::vector<Transport::Slice*>& failed_slice_list) {
+    if (status_.load(std::memory_order_relaxed) != CONNECTED) {
+        int ret = setupConnectionsByActive();
+        if (ret != 0) {
+            for (auto* slice : slice_list) {
+                failed_slice_list.push_back(slice);
+            }
+            slice_list.clear();
+            return ret;
+        }
+    }
+
+    fi_addr_t peer;
+    {
+        RWSpinlock::ReadGuard guard(lock_);
+        peer = peer_fi_addr_;
+    }
+    if (peer == FI_ADDR_UNSPEC) {
+        for (auto* slice : slice_list) failed_slice_list.push_back(slice);
+        slice_list.clear();
+        return ERR_ENDPOINT;
+    }
+
+    return context_.submitSlicesOnPeer(peer, slice_list, failed_slice_list);
+}
+
+}  // namespace mooncake
diff --git a/mooncake-transfer-engine/src/transport/cxi_transport/cxi_transport.cpp b/mooncake-transfer-engine/src/transport/cxi_transport/cxi_transport.cpp
new file mode 100644
index 00000000..c106bead
--- /dev/null
+++ b/mooncake-transfer-engine/src/transport/cxi_transport/cxi_transport.cpp
@@ -0,0 +1,1007 @@
+// Copyright 2024 KVCache.AI
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "transport/cxi_transport/cxi_transport.h"
+
+#include <glog/logging.h>
+#include <sys/mman.h>
+#include <sys/time.h>
+#include <unistd.h>
+
+#include <algorithm>
+#include <cassert>
+#include <chrono>
+#include <cstddef>
+#include <fstream>
+#include <future>
+#include <set>
+#include <thread>
+
+#include <dlfcn.h>
+
+#include "common.h"
+#include "config.h"
+#include "memory_location.h"
+#include "topology.h"
+#include "transport/cxi_transport/cxi_context.h"
+#include "transport/cxi_transport/cxi_endpoint.h"
+
+namespace mooncake {
+
+NicReplicaPolicy CxiTransport::getReplicaPolicy() {
+    return NicReplicaPolicy::NUMA_AWARE;
+}
+
+CxiTransport::CxiTransport() {
+    LOG(INFO) << "[CXI] CXI Slingshot transport initialized";
+}
+
+CxiTransport::~CxiTransport() {
+    stopWorkerThreads();
+    metadata_->removeSegmentDesc(local_server_name_);
+    batch_desc_set_.clear();
+    context_list_.clear();
+}
+
+void CxiTransport::startWorkerThreads() {
+    if (worker_running_) return;
+
+    worker_running_ = true;
+    // One poller thread per context for responsive CQ draining under load
+    size_t num_threads = context_list_.size();
+    for (size_t i = 0; i < num_threads; i++) {
+        worker_threads_.emplace_back(&CxiTransport::workerThreadFunc, this, i);
+    }
+    LOG(INFO) << "CxiTransport: Started " << num_threads
+              << " CQ polling worker threads";
+}
+
+void CxiTransport::stopWorkerThreads() {
+    if (!worker_running_) return;
+
+    worker_running_ = false;
+    for (auto& thread : worker_threads_) {
+        if (thread.joinable()) {
+            thread.join();
+        }
+    }
+    worker_threads_.clear();
+    LOG(INFO) << "CxiTransport: Stopped CQ polling worker threads";
+}
+
+void CxiTransport::workerThreadFunc(int thread_id) {
+    const int kPollBatchSize = 64;
+
+    while (worker_running_) {
+        bool did_work = false;
+
+        // Poll CQs from all contexts
+        for (size_t ctx_idx = thread_id; ctx_idx < context_list_.size();
+             ctx_idx += worker_threads_.size()) {
+            auto& context = context_list_[ctx_idx];
+            if (!context || !context->active()) continue;
+
+            for (size_t cq_idx = 0; cq_idx < context->cqCount(); cq_idx++) {
+                int completed = context->pollCq(kPollBatchSize, cq_idx);
+                if (completed > 0) {
+                    did_work = true;
+                }
+            }
+        }
+
+        // Under the shared-endpoint model there is no per-peer QP to evict:
+        // a new peer costs one AV entry (~bytes), not an fi_endpoint slot.
+        // Stale peers are reclaimed when submitPostSend() drops a peer whose
+        // handshake failed.
+
+        // If no work was done, yield CPU briefly
+        if (!did_work) {
+            std::this_thread::yield();
+        }
+    }
+}
+
+int CxiTransport::install(std::string& local_server_name,
+                          std::shared_ptr<TransferMetadata> meta,
+                          std::shared_ptr<Topology> topo) {
+    if (topo == nullptr) {
+        LOG(ERROR) << "CxiTransport: missing topology";
+        return ERR_INVALID_ARGUMENT;
+    }
+
+    metadata_ = meta;
+    local_server_name_ = local_server_name;
+    local_topology_ = topo;
+
+    auto ret = initializeCxiResources();
+    if (ret) {
+        LOG(ERROR) << "CxiTransport: cannot initialize CXI resources";
+        return ret;
+    }
+
+    ret = allocateLocalSegmentID();
+    if (ret) {
+        LOG(ERROR) << "Transfer engine cannot be initialized: cannot "
+                      "allocate local segment";
+        return ret;
+    }
+
+    ret = startHandshakeDaemon(local_server_name);
+    if (ret) {
+        LOG(ERROR) << "CxiTransport: cannot start handshake daemon";
+        return ret;
+    }
+
+    ret = metadata_->updateLocalSegmentDesc();
+    if (ret) {
+        LOG(ERROR) << "CxiTransport: cannot publish segments";
+        return ret;
+    }
+
+    // Start CQ polling worker threads
+    startWorkerThreads();
+
+    return 0;
+}
+
+int CxiTransport::preTouchMemory(void* addr, size_t length) {
+    if (context_list_.size() == 0) {
+        return 0;
+    }
+
+    unsigned int hwc = std::thread::hardware_concurrency();
+    unsigned int num_threads = hwc > 64 ? 16 : std::min(hwc, 8u);
+    num_threads = std::max(
+        num_threads, 1u);  // guard in case hardware_concurrency() returns 0
+    size_t block_size = length / num_threads;
+    if (block_size == 0) {
+        return 0;
+    }
+
+    std::vector<std::thread> threads;
+    threads.reserve(num_threads);
+    std::vector<int> thread_results(num_threads, 0);
+
+    for (size_t thread_i = 0; thread_i < num_threads; ++thread_i) {
+        void* block_addr = static_cast<char*>(addr) + thread_i * block_size;
+        threads.emplace_back([this, thread_i, block_addr, block_size,
+                              &thread_results]() {
+            int ret = context_list_[0]->preTouchMemory(block_addr, block_size);
+            thread_results[thread_i] = ret;
+        });
+    }
+
+    for (auto& thread : threads) {
+        thread.join();
+    }
+
+    for (size_t i = 0; i < num_threads; ++i) {
+        if (thread_results[i] != 0) {
+            return thread_results[i];
+        }
+    }
+
+    return 0;
+}
+
+int CxiTransport::registerLocalMemory(void* addr, size_t length,
+                                      const std::string& name,
+                                      bool remote_accessible,
+                                      bool update_metadata) {
+    return registerLocalMemoryInternal(addr, length, name, remote_accessible,
+                                       update_metadata, false);
+}
+
+int CxiTransport::registerLocalMemoryInternal(void* addr, size_t length,
+                                              const std::string& name,
+                                              bool remote_accessible,
+                                              bool update_metadata,
+                                              bool force_sequential) {
+    (void)remote_accessible;
+    const int kBaseAccessRights =
+        IBV_ACCESS_LOCAL_WRITE | IBV_ACCESS_REMOTE_WRITE |
+        IBV_ACCESS_REMOTE_READ;  // this is not used, ignore
+
+    int access_rights = kBaseAccessRights;
+    size_t max_mr = (size_t)globalConfig().max_mr_size;
+
+    // max_mr must be set, slingshot handles mempages differently compared to
+    // efa
+    size_t chunk_limit = max_mr;
+    LOG(INFO) << "Auto-split params: max_mr_size=" << max_mr
+              << ", chunk_limit=" << chunk_limit;
+
+    // Determine chunk boundaries
+    std::vector<std::pair<void*, size_t>> chunks;
+    if (length > chunk_limit) {
+        size_t offset = 0;
+        while (offset < length) {
+            size_t chunk_len = std::min(chunk_limit, length - offset);
+            chunks.emplace_back(static_cast<char*>(addr) + offset, chunk_len);
+            offset += chunk_len;
+        }
+        LOG(WARNING) << "Auto-splitting buffer " << addr << " (" << length
+                     << " bytes) into " << chunks.size()
+                     << " chunks of <= " << chunk_limit << " bytes each";
+    } else {
+        chunks.emplace_back(addr, length);
+    }
+
+    // Resolve location name once (based on original buffer)
+    std::string resolved_name;
+    if (name == kWildcardLocation) {
+        bool only_first_page = true;
+        const std::vector<MemoryLocationEntry> entries = getMemoryLocation(
+            addr, length, only_first_page);  // check only first page
+        if (entries.empty()) return ERR_DEVICE_NOT_FOUND;
+        resolved_name = entries[0].location;
+    } else {
+        resolved_name = name;
+    }
+
+    size_t num_nics = context_list_.size();
+    size_t num_chunks = chunks.size();
+
+    auto policy = getReplicaPolicy();
+
+    std::vector<std::vector<size_t>> nic_assignments(num_chunks);
+
+    int selected_device = local_topology_->selectDevice(resolved_name);
+    if (selected_device == ERR_DEVICE_NOT_FOUND) {
+        LOG(ERROR) << "could not select a NIC for data, resolved_name="
+                   << resolved_name;
+        return ERR_DEVICE_NOT_FOUND;
+    }
+    std::string nic = local_topology_->getHcaList().at(
+        local_topology_->selectDevice(resolved_name));
+    LOG(INFO) << "for this allocation, selected NIC " << nic;
+    int id = -1;
+    for (size_t i = 0; i < context_list_.size(); i++) {
+        auto& entry = context_list_[i];
+        if (entry->deviceName() == nic) {
+            id = i;
+            break;
+        }
+    }
+    if (id == -1) {
+        LOG(ERROR) << "possible mismatch between context list devices and "
+                      "topology HCA names!";
+        return -1;
+    }
+    int numa_node = id;
+
+    // implicit assumption here, the address and relative chunks all lie in the
+    // same numa node
+    for (size_t ci = 0; ci < num_chunks; ci++) {
+        // based on replication policy, assign chunk to nic
+        switch (policy) {
+            case NUMA_AWARE:
+                nic_assignments[ci].push_back(numa_node);
+                break;
+            case REPLICATE_ALL:
+                for (size_t nic_idx = 0; nic_idx < num_nics; nic_idx++)
+                    nic_assignments[ci].push_back(nic_idx);
+                break;
+            default:
+                nic_assignments[ci].push_back(
+                    0);  // register everything to nic 0
+                break;
+        }
+    }
+
+    auto rollbackChunks = [&](size_t up_to_ci) {
+        for (size_t ri = 0; ri <= up_to_ci; ++ri) {
+            for (size_t nic_idx : nic_assignments[ri]) {
+                context_list_[nic_idx]->unregisterMemoryRegion(
+                    chunks[ri].first);
+            }
+        }
+    };
+
+    // Register each chunk on its assigned NICs
+    for (size_t ci = 0; ci < chunks.size(); ++ci) {
+        void* chunk_addr = chunks[ci].first;
+        size_t chunk_len = chunks[ci].second;
+
+        const auto& assigned_nics = nic_assignments[ci];
+
+        // preTouchMemory does a CPU-side store to each page, which segfaults
+        // on GPU VRAM (cudaMalloc'd pointers). Restrict it to host memory.
+        bool is_host_mem = resolved_name.rfind("cpu", 0) == 0;
+        bool do_pre_touch = is_host_mem && context_list_.size() > 0 &&
+                            std::thread::hardware_concurrency() >= 4 &&
+                            chunk_len >= (size_t)4 * 1024 * 1024 * 1024;
+        if (do_pre_touch) {
+            int ret = preTouchMemory(chunk_addr, chunk_len);
+            if (ret != 0) {
+                if (ci > 0) rollbackChunks(ci - 1);
+                return ret;
+            }
+        }
+
+        int use_parallel_reg = 0;
+        if (!force_sequential) {
+            use_parallel_reg = globalConfig().parallel_reg_mr;
+            if (use_parallel_reg == -1) {
+                use_parallel_reg = assigned_nics.size() > 1 && do_pre_touch;
+            }
+        }
+
+        auto reg_start = std::chrono::steady_clock::now();
+
+        if (use_parallel_reg) {
+            std::vector<std::thread> reg_threads;
+            reg_threads.reserve(assigned_nics.size());
+            std::vector<int> ret_codes(assigned_nics.size(), 0);
+            const int ar = access_rights;
+
+            for (size_t j = 0; j < assigned_nics.size(); ++j) {
+                size_t nic_idx = assigned_nics[j];
+                reg_threads.emplace_back([this, &ret_codes, j, nic_idx,
+                                          chunk_addr, chunk_len, ar]() {
+                    ret_codes[j] = context_list_[nic_idx]->registerMemoryRegion(
+                        chunk_addr, chunk_len, ar);
+                });
+            }
+
+            for (auto& thread : reg_threads) {
+                thread.join();
+            }
+
+            for (size_t j = 0; j < ret_codes.size(); ++j) {
+                if (ret_codes[j] != 0) {
+                    LOG(ERROR)
+                        << "Failed to register memory region chunk " << ci
+                        << " with CXI context " << assigned_nics[j];
+                    rollbackChunks(ci);
+                    return ret_codes[j];
+                }
+            }
+        } else {
+            for (size_t nic_idx : assigned_nics) {
+                int ret = context_list_[nic_idx]->registerMemoryRegion(
+                    chunk_addr, chunk_len, access_rights);
+                if (ret) {
+                    LOG(ERROR) << "Failed to register memory region chunk "
+                               << ci << " with CXI context " << nic_idx;
+                    rollbackChunks(ci);
+                    return ret;
+                }
+            }
+        }
+
+        auto reg_end = std::chrono::steady_clock::now();
+        auto reg_duration_ms =
+            std::chrono::duration_cast<std::chrono::milliseconds>(reg_end -
+                                                                  reg_start)
+                .count();
+
+        if (globalConfig().trace) {
+            LOG(INFO) << "CXI registerMemoryRegion: chunk " << ci
+                      << ", addr=" << chunk_addr << ", length=" << chunk_len
+                      << ", nics=" << assigned_nics.size() << "/"
+                      << context_list_.size()
+                      << ", parallel=" << (use_parallel_reg ? "true" : "false")
+                      << ", duration=" << reg_duration_ms << "ms";
+        }
+
+        LOG(WARNING) << "Chunk " << ci << "/" << chunks.size()
+                     << " registered on " << assigned_nics.size() << " NICs"
+                     << ", addr=" << chunk_addr << ", length=" << chunk_len
+                     << ", duration=" << reg_duration_ms << "ms";
+
+        // Collect keys: assigned NICs have valid keys, others get 0
+        BufferDesc buffer_desc;
+        for (auto& context : context_list_) {
+            buffer_desc.lkey.push_back(context->lkey(chunk_addr));
+            buffer_desc.rkey.push_back(context->rkey(chunk_addr));
+        }
+
+        buffer_desc.name = resolved_name;
+        buffer_desc.addr = (uint64_t)chunk_addr;
+        buffer_desc.length = chunk_len;
+        int rc = metadata_->addLocalMemoryBuffer(buffer_desc, update_metadata);
+        if (rc) {
+            rollbackChunks(ci);
+            return rc;
+        }
+    }
+
+    // Track chunks and NIC assignments for unregistration
+    if (chunks.size() > 1) {
+        std::lock_guard<std::mutex> lock(chunk_map_mutex_);
+        std::vector<ChunkRegistration> regs;
+        regs.reserve(chunks.size());
+        for (size_t ci = 0; ci < chunks.size(); ++ci) {
+            regs.push_back({(uint64_t)chunks[ci].first, nic_assignments[ci]});
+        }
+        chunk_map_[(uint64_t)addr] = std::move(regs);
+    }
+
+    return 0;
+}
+
+int CxiTransport::unregisterLocalMemory(void* addr, bool update_metadata) {
+    return unregisterLocalMemoryInternal(addr, update_metadata, false);
+}
+
+int CxiTransport::unregisterLocalMemoryInternal(void* addr,
+                                                bool update_metadata,
+                                                bool force_sequential) {
+    // Check if this buffer was split into chunks (per-NIC partition)
+    std::vector<ChunkRegistration> chunk_regs;
+    {
+        std::lock_guard<std::mutex> lock(chunk_map_mutex_);
+        auto it = chunk_map_.find((uint64_t)addr);
+        if (it != chunk_map_.end()) {
+            chunk_regs = std::move(it->second);
+            chunk_map_.erase(it);
+        }
+    }
+
+    if (!chunk_regs.empty()) {
+        // Unregister each chunk from its assigned NICs only
+        for (auto& reg : chunk_regs) {
+            void* ca = (void*)reg.addr;
+            int rc = metadata_->removeLocalMemoryBuffer(ca, update_metadata);
+            if (rc) {
+                LOG(ERROR) << "Failed to remove chunk metadata at " << ca;
+                return rc;
+            }
+
+            for (size_t nic_idx : reg.nic_indices) {
+                int ret = context_list_[nic_idx]->unregisterMemoryRegion(ca);
+                if (ret) {
+                    LOG(ERROR) << "Failed to unregister chunk " << ca
+                               << " with CXI context " << nic_idx;
+                    return ret;
+                }
+            }
+        }
+        return 0;
+    }
+
+    // Non-chunked buffer: original path
+    int rc = metadata_->removeLocalMemoryBuffer(addr, update_metadata);
+    if (rc) return rc;
+
+    int use_parallel_unreg = 0;
+    if (!force_sequential) {
+        use_parallel_unreg = globalConfig().parallel_reg_mr;
+        if (use_parallel_unreg == -1) {
+            use_parallel_unreg = context_list_.size() > 1;
+        }
+    }
+
+    if (use_parallel_unreg) {
+        std::vector<std::thread> unreg_threads;
+        unreg_threads.reserve(context_list_.size());
+        std::vector<int> ret_codes(context_list_.size(), 0);
+
+        for (size_t i = 0; i < context_list_.size(); ++i) {
+            unreg_threads.emplace_back([this, &ret_codes, i, addr]() {
+                ret_codes[i] = context_list_[i]->unregisterMemoryRegion(addr);
+            });
+        }
+
+        for (auto& thread : unreg_threads) {
+            thread.join();
+        }
+
+        for (size_t i = 0; i < ret_codes.size(); ++i) {
+            if (ret_codes[i] != 0) {
+                LOG(ERROR)
+                    << "Failed to unregister memory region with CXI context "
+                    << i;
+                return ret_codes[i];
+            }
+        }
+    } else {
+        for (size_t i = 0; i < context_list_.size(); ++i) {
+            int ret = context_list_[i]->unregisterMemoryRegion(addr);
+            if (ret) {
+                LOG(ERROR)
+                    << "Failed to unregister memory region with CXI context "
+                    << i;
+                return ret;
+            }
+        }
+    }
+
+    return 0;
+}
+
+int CxiTransport::allocateLocalSegmentID() {
+    auto desc = std::make_shared<SegmentDesc>();
+    if (!desc) return ERR_MEMORY;
+    desc->name = local_server_name_;
+    desc->protocol = "cxi";
+    for (auto& entry : context_list_) {
+        TransferMetadata::DeviceDesc device_desc;
+        device_desc.name = entry->deviceName();
+        device_desc.lid = entry->lid();
+        device_desc.gid = entry->gid();
+        desc->devices.push_back(device_desc);
+    }
+    desc->topology = *(local_topology_.get());
+    metadata_->addLocalSegment(LOCAL_SEGMENT_ID, local_server_name_,
+                               std::move(desc));
+    return 0;
+}
+
+int CxiTransport::registerLocalMemoryBatch(
+    const std::vector<CxiTransport::BufferEntry>& buffer_list,
+    const std::string& location) {
+    std::vector<std::future<int>> results;
+    for (auto& buffer : buffer_list) {
+        results.emplace_back(
+            std::async(std::launch::async, [this, buffer, location]() -> int {
+                return registerLocalMemoryInternal(buffer.addr, buffer.length,
+                                                   location, true, false, true);
+            }));
+    }
+
+    for (size_t i = 0; i < buffer_list.size(); ++i) {
+        if (results[i].get()) {
+            LOG(WARNING) << "CxiTransport: Failed to register memory: addr "
+                         << buffer_list[i].addr << " length "
+                         << buffer_list[i].length;
+        }
+    }
+
+    return metadata_->updateLocalSegmentDesc();
+}
+
+int CxiTransport::unregisterLocalMemoryBatch(
+    const std::vector<void*>& addr_list) {
+    std::vector<std::future<int>> results;
+    for (auto& addr : addr_list) {
+        results.emplace_back(
+            std::async(std::launch::async, [this, addr]() -> int {
+                return unregisterLocalMemoryInternal(addr, false, true);
+            }));
+    }
+
+    for (size_t i = 0; i < addr_list.size(); ++i) {
+        if (results[i].get())
+            LOG(WARNING) << "CxiTransport: Failed to unregister memory: addr "
+                         << addr_list[i];
+    }
+
+    return metadata_->updateLocalSegmentDesc();
+}
+
+int CxiTransport::warmupSegment(const std::string& segment_name) {
+    if (!metadata_) {
+        LOG(ERROR) << "CxiTransport::warmupSegment: metadata_ is null";
+        return ERR_INVALID_ARGUMENT;
+    }
+    if (segment_name.empty() || segment_name == local_server_name_) {
+        // Loopback / empty name — nothing to pre-connect.
+        return 0;
+    }
+
+    auto desc = metadata_->getSegmentDescByName(segment_name);
+    if (!desc) {
+        LOG(ERROR) << "CxiTransport::warmupSegment: segment '" << segment_name
+                   << "' not found in metadata (did you openSegment() first?)";
+        return ERR_INVALID_ARGUMENT;
+    }
+    if (desc->devices.empty()) {
+        LOG(WARNING) << "CxiTransport::warmupSegment: segment '" << segment_name
+                     << "' has no devices";
+        return 0;
+    }
+
+    // Build peer_nic_path list: "<segment_name>@<device_name>" for each NIC.
+    std::vector<std::string> peer_paths;
+    peer_paths.reserve(desc->devices.size());
+    for (const auto& dev : desc->devices) {
+        peer_paths.emplace_back(segment_name + "@" + dev.name);
+    }
+
+    auto t0 = std::chrono::steady_clock::now();
+    size_t n_pairs = context_list_.size() * peer_paths.size();
+
+    // Idempotent short-circuit: if every (local_ctx, peer_nic) pair already
+    // has a connected endpoint, skip the whole async dispatch. Matters for
+    // callers that invoke warmupSegment per request loop — without this the
+    // 256-thread fan-out runs every time even though there is no work to do.
+    size_t already_ready = 0;
+    for (auto& ctx : context_list_) {
+        for (const auto& path : peer_paths) {
+            auto ep = ctx->peekEndpoint(path);
+            if (ep && ep->connected()) ++already_ready;
+        }
+    }
+    if (already_ready == n_pairs) {
+        VLOG(1) << "CxiTransport::warmupSegment('" << segment_name << "'): all "
+                << n_pairs << " endpoints already connected, "
+                << "skipping";
+        return 0;
+    }
+
+    // Warm up every (local_ctx, peer_nic) pair concurrently.  Under the
+    // shared-endpoint model each warmup is just a handshake RPC +
+    // fi_av_insert (no fi_endpoint, no fi_enable), so the critical path is
+    // max(handshake RTT), not sum.  We still dispatch with std::async for
+    // concurrency, but total wall time is typically ms-level.
+    std::vector<std::future<int>> futs;
+    futs.reserve(n_pairs);
+    for (auto& ctx : context_list_) {
+        for (const auto& path : peer_paths) {
+            futs.emplace_back(
+                std::async(std::launch::async, [ctx, path]() -> int {
+                    auto ep = ctx->endpoint(path);
+                    if (!ep) {
+                        LOG(WARNING) << "warmupSegment: endpoint() returned "
+                                        "null for "
+                                     << path;
+                        return -1;
+                    }
+                    if (ep->connected()) return 0;
+                    int rc = ep->setupConnectionsByActive();
+                    if (rc != 0) {
+                        // Handshake failed: drop the peer handle so the AV
+                        // slot is freed and the next warmup retry starts
+                        // clean.  Cheap under the shared-endpoint model —
+                        // no fi_endpoint teardown required.
+                        ctx->deleteEndpoint(path);
+                    }
+                    return rc;
+                }));
+        }
+    }
+    int ok = 0, fail = 0;
+    for (auto& f : futs) {
+        int rc = f.get();
+        if (rc == 0)
+            ++ok;
+        else
+            ++fail;
+    }
+    auto elapsed =
+        std::chrono::duration<double>(std::chrono::steady_clock::now() - t0)
+            .count();
+    LOG(INFO) << "CxiTransport::warmupSegment('" << segment_name << "'): " << ok
+              << "/" << n_pairs << " endpoints connected (" << fail
+              << " failed) in " << elapsed << "s (" << context_list_.size()
+              << " local NICs x " << peer_paths.size() << " peer NICs)";
+    return fail == 0 ? 0 : ERR_ENDPOINT;
+}
+
+Status CxiTransport::submitTransfer(
+    BatchID batch_id, const std::vector<TransferRequest>& entries) {
+    auto& batch_desc = *((BatchDesc*)(batch_id));
+    if (batch_desc.task_list.size() + entries.size() > batch_desc.batch_size) {
+        LOG(ERROR) << "CxiTransport: Exceed the limitation of current batch's "
+                      "capacity";
+        return Status::InvalidArgument(
+            "CxiTransport: Exceed the limitation of capacity, batch id: " +
+            std::to_string(batch_id));
+    }
+
+    size_t task_id = batch_desc.task_list.size();
+    batch_desc.task_list.resize(task_id + entries.size());
+    std::vector<TransferTask*> task_list;
+    for (auto& task : batch_desc.task_list) task_list.push_back(&task);
+    return submitTransferTask(task_list);
+}
+
+Status CxiTransport::submitTransferTask(
+    const std::vector<TransferTask*>& task_list) {
+    std::unordered_map<std::shared_ptr<CxiContext>, std::vector<Slice*>>
+        slices_to_post;
+    auto local_segment_desc = metadata_->getSegmentDescByID(LOCAL_SEGMENT_ID);
+    assert(local_segment_desc.get());
+    const int kMaxRetryCount = globalConfig().retry_cnt;
+
+    for (size_t index = 0; index < task_list.size(); ++index) {
+        assert(task_list[index]);
+        auto& task = *task_list[index];
+        assert(task.request);
+        auto& request = *task.request;
+
+        if (request.length == 0) continue;
+
+        // Find which buffer and preferred device covers this request
+        int request_buffer_id = -1, request_device_id = -1;
+        if (selectDevice(local_segment_desc.get(), (uint64_t)request.source,
+                         request.length, request_buffer_id,
+                         request_device_id)) {
+            request_buffer_id = -1;
+            request_device_id = -1;
+        }
+
+        if (request_buffer_id >= 0 && request_device_id >= 0) {
+            // One slice per request.  Round-robin NIC selection is
+            // handled by selectDevice above.
+            auto& context = context_list_[request_device_id];
+            if (!context || !context->active()) {
+                LOG(ERROR) << "CXI Device " << request_device_id
+                           << " is not active";
+                return Status::InvalidArgument(
+                    "CXI Device " + std::to_string(request_device_id) +
+                    " is not active");
+            }
+
+            Slice* slice = getSliceCache().allocate();
+            assert(slice);
+            slice->peer_nic_path.clear();
+            slice->rdma.dest_rkey = 0;  // will be set later
+
+            slice->source_addr = (char*)request.source;
+            slice->length = request.length;
+            slice->opcode = request.opcode;
+            slice->rdma.dest_addr = request.target_offset;
+            slice->rdma.retry_cnt = request.advise_retry_cnt;
+            slice->rdma.max_retry_cnt = kMaxRetryCount;
+            slice->task = &task;
+            slice->target_id = request.target_id;
+            slice->status = Slice::PENDING;
+            slice->ts = 0;
+            task.slice_list.push_back(slice);
+
+            slice->rdma.source_lkey =
+                local_segment_desc->buffers[request_buffer_id]
+                    .lkey[request_device_id];
+            slices_to_post[context].push_back(slice);
+            __sync_fetch_and_add(&task.total_bytes, slice->length);
+            __sync_fetch_and_add(&task.slice_count, 1);
+        } else {
+            // FALLBACK: device not found via initial selectDevice.
+            // Try per-slice retry with increasing retry_cnt to find any
+            // available device (handles edge cases like multi-buffer spans).
+            int buffer_id = -1, device_id = -1;
+            int retry_cnt = request.advise_retry_cnt;
+            bool found_device = false;
+            while (retry_cnt < kMaxRetryCount && !found_device) {
+                if (selectDevice(local_segment_desc.get(),
+                                 (uint64_t)request.source, request.length,
+                                 buffer_id, device_id, retry_cnt++))
+                    continue;
+                if (device_id >= 0 &&
+                    static_cast<size_t>(device_id) < context_list_.size() &&
+                    context_list_[device_id] &&
+                    context_list_[device_id]->active()) {
+                    found_device = true;
+                    break;
+                }
+            }
+            if (!found_device) {
+                LOG(ERROR) << "Memory region not registered by any active CXI "
+                              "device(s): "
+                           << request.source;
+                for (auto& entry : slices_to_post)
+                    for (auto s : entry.second) s->markFailed();
+                return Status::AddressNotRegistered(
+                    "Memory region not registered by any active CXI "
+                    "device(s): " +
+                    std::to_string(
+                        reinterpret_cast<uintptr_t>(request.source)));
+            }
+
+            // Found a device via retry — create single slice
+            Slice* slice = getSliceCache().allocate();
+            assert(slice);
+            slice->peer_nic_path.clear();
+            slice->rdma.dest_rkey = 0;
+
+            slice->source_addr = (char*)request.source;
+            slice->length = request.length;
+            slice->opcode = request.opcode;
+            slice->rdma.dest_addr = request.target_offset;
+            slice->rdma.retry_cnt = request.advise_retry_cnt;
+            slice->rdma.max_retry_cnt = kMaxRetryCount;
+            slice->task = &task;
+            slice->target_id = request.target_id;
+            slice->status = Slice::PENDING;
+            slice->ts = 0;
+            task.slice_list.push_back(slice);
+
+            auto& context = context_list_[device_id];
+            slice->rdma.source_lkey =
+                local_segment_desc->buffers[buffer_id].lkey[device_id];
+            slices_to_post[context].push_back(slice);
+            __sync_fetch_and_add(&task.total_bytes, slice->length);
+            __sync_fetch_and_add(&task.slice_count, 1);
+        }
+    }
+
+    for (auto& entry : slices_to_post)
+        if (!entry.second.empty()) entry.first->submitPostSend(entry.second);
+    return Status::OK();
+}
+
+Status CxiTransport::getTransferStatus(BatchID batch_id,
+                                       std::vector<TransferStatus>& status) {
+    auto& batch_desc = *((BatchDesc*)(batch_id));
+    const size_t task_count = batch_desc.task_list.size();
+    status.resize(task_count);
+    for (size_t task_id = 0; task_id < task_count; task_id++) {
+        auto& task = batch_desc.task_list[task_id];
+        status[task_id].transferred_bytes = task.transferred_bytes;
+        uint64_t success_slice_count = task.success_slice_count;
+        uint64_t failed_slice_count = task.failed_slice_count;
+        if (success_slice_count + failed_slice_count == task.slice_count) {
+            if (failed_slice_count)
+                status[task_id].s = TransferStatusEnum::FAILED;
+            else
+                status[task_id].s = TransferStatusEnum::COMPLETED;
+            task.is_finished = true;
+        } else {
+            status[task_id].s = TransferStatusEnum::WAITING;
+        }
+    }
+    return Status::OK();
+}
+
+Status CxiTransport::getTransferStatus(BatchID batch_id, size_t task_id,
+                                       TransferStatus& status) {
+    auto& batch_desc = *((BatchDesc*)(batch_id));
+    const size_t task_count = batch_desc.task_list.size();
+    if (task_id >= task_count) {
+        return Status::InvalidArgument(
+            "CxiTransport::getTransportStatus invalid argument, batch id: " +
+            std::to_string(batch_id));
+    }
+    auto& task = batch_desc.task_list[task_id];
+    status.transferred_bytes = task.transferred_bytes;
+    uint64_t success_slice_count = task.success_slice_count;
+    uint64_t failed_slice_count = task.failed_slice_count;
+    if (success_slice_count + failed_slice_count == task.slice_count) {
+        if (failed_slice_count)
+            status.s = TransferStatusEnum::FAILED;
+        else
+            status.s = TransferStatusEnum::COMPLETED;
+        task.is_finished = true;
+    } else {
+        status.s = TransferStatusEnum::WAITING;
+    }
+    return Status::OK();
+}
+
+CxiTransport::SegmentID CxiTransport::getSegmentID(
+    const std::string& segment_name) {
+    return metadata_->getSegmentID(segment_name);
+}
+
+int CxiTransport::onSetupCxiConnections(const HandShakeDesc& peer_desc,
+                                        HandShakeDesc& local_desc) {
+    auto local_nic_name = getNicNameFromNicPath(peer_desc.peer_nic_path);
+    if (local_nic_name.empty()) return ERR_INVALID_ARGUMENT;
+
+    // Find context by device name instead of using hca_list index, since
+    // context_list_ only contains CXI devices and may have different
+    // indexing than the full hca_list.
+    std::shared_ptr<CxiContext> context;
+    for (auto& entry : context_list_) {
+        if (entry->deviceName() == local_nic_name) {
+            context = entry;
+            break;
+        }
+    }
+    if (!context) return ERR_INVALID_ARGUMENT;
+
+    auto endpoint = context->endpoint(peer_desc.local_nic_path);
+    if (!endpoint) return ERR_ENDPOINT;
+    return endpoint->setupConnectionsByPassive(peer_desc, local_desc);
+}
+
+int CxiTransport::initializeCxiResources() {
+    auto hca_list = local_topology_->getHcaList();
+
+    std::vector<std::string> cxi_devices;
+    std::vector<std::string> non_cxi_devices;
+    for (auto& device_name : hca_list) {
+        if (device_name.find("cxi") != std::string::npos) {
+            cxi_devices.push_back(device_name);
+        } else {
+            non_cxi_devices.push_back(device_name);
+        }
+    }
+
+    if (cxi_devices.empty()) {
+        LOG(WARNING) << "CxiTransport: No CXI devices found, falling back to "
+                        "all devices";
+        cxi_devices = hca_list;
+        non_cxi_devices.clear();
+    }
+
+    // Disable non-CXI devices
+    for (auto& device_name : non_cxi_devices) {
+        local_topology_->disableDevice(device_name);
+        LOG(INFO) << "CxiTransport: Disabled non-CXI device " << device_name
+                  << " in topology";
+    }
+
+    for (auto& device_name : cxi_devices) {
+        auto context = std::make_shared<CxiContext>(*this, device_name);
+        auto& config = globalConfig();
+        int ret = context->construct(config.num_cq_per_ctx, config.max_cqe,
+                                     config.max_ep_per_ctx);
+        if (ret) {
+            local_topology_->disableDevice(device_name);
+            LOG(WARNING) << "CxiTransport: Disable device " << device_name;
+        } else {
+            context_list_.push_back(context);
+            LOG(INFO) << "CxiTransport: Initialized CXI device " << device_name;
+        }
+    }
+    if (context_list_.empty()) {
+        LOG(ERROR) << "CxiTransport: No available CXI devices";
+        return ERR_DEVICE_NOT_FOUND;
+    }
+
+    return 0;
+}
+
+int CxiTransport::startHandshakeDaemon(std::string& local_server_name) {
+    return metadata_->startHandshakeDaemon(
+        std::bind(&CxiTransport::onSetupCxiConnections, this,
+                  std::placeholders::_1, std::placeholders::_2),
+        metadata_->localRpcMeta().rpc_port, metadata_->localRpcMeta().sockfd);
+}
+
+int CxiTransport::selectDevice(SegmentDesc* desc, uint64_t offset,
+                               size_t length, std::string_view hint,
+                               int& buffer_id, int& device_id,
+                               int retry_count) {
+    if (desc == nullptr) return ERR_ADDRESS_NOT_REGISTERED;
+    const auto& buffers = desc->buffers;
+    for (buffer_id = 0; buffer_id < static_cast<int>(buffers.size());
+         ++buffer_id) {
+        const auto& buffer = buffers[buffer_id];
+
+        if (offset < buffer.addr || length > buffer.length ||
+            offset - buffer.addr > buffer.length - length) {
+            continue;
+        }
+
+        int num_devices = static_cast<int>(desc->devices.size());
+        // this must be <=, because the first attempt will be random, so to
+        // guarantee that the device is found we need num_device+1 attempts
+        for (int attempt = 0; attempt <= num_devices; ++attempt) {
+            int try_count = retry_count + attempt;
+            device_id =
+                hint.empty()
+                    ? desc->topology.selectDevice(buffer.name, try_count)
+                    : desc->topology.selectDevice(buffer.name, hint, try_count);
+            if (device_id >= 0 &&
+                static_cast<size_t>(device_id) < buffer.rkey.size() &&
+                buffer.rkey[device_id] != 0) {
+                return 0;
+            }
+            device_id = hint.empty() ? desc->topology.selectDevice(
+                                           kWildcardLocation, try_count)
+                                     : desc->topology.selectDevice(
+                                           kWildcardLocation, hint, try_count);
+            if (device_id >= 0 &&
+                static_cast<size_t>(device_id) < buffer.rkey.size() &&
+                buffer.rkey[device_id] != 0) {
+                return 0;
+            }
+        }
+    }
+    return ERR_ADDRESS_NOT_REGISTERED;
+}
+
+int CxiTransport::selectDevice(SegmentDesc* desc, uint64_t offset,
+                               size_t length, int& buffer_id, int& device_id,
+                               int retry_count) {
+    return selectDevice(desc, offset, length, "", buffer_id, device_id,
+                        retry_count);
+}
+
+}  // namespace mooncake
diff --git a/mooncake-transfer-engine/src/transport/device/CMakeLists.txt b/mooncake-transfer-engine/src/transport/device/CMakeLists.txt
new file mode 100644
index 00000000..903ec819
--- /dev/null
+++ b/mooncake-transfer-engine/src/transport/device/CMakeLists.txt
@@ -0,0 +1,21 @@
+# Device transport sources — conditional on GPU vendor. mlx5gda.cpp (IBGDA QP
+# lifecycle) is compiled directly into the device_transport OBJECT library, like
+# every other transport module, so its mlx5gda_* symbols flow into
+# transfer_engine without a separate static lib. This keeps them visible to
+# consumers that link transfer_engine via hand- written ldflags (Go p2p-store /
+# mooncake-store), which bypass CMake's target_link_libraries propagation.
+set(DEVICE_TRANSPORT_SOURCES p2p_device_transport.cpp)
+if((USE_CUDA OR USE_MUSA) AND NOT (USE_CXI))
+  list(APPEND DEVICE_TRANSPORT_SOURCES ibgda_device_transport.cpp mlx5gda.cpp)
+endif()
+
+add_library(device_transport OBJECT ${DEVICE_TRANSPORT_SOURCES})
+target_include_directories(device_transport PRIVATE ${CMAKE_SOURCE_DIR}/include)
+
+if(USE_CUDA)
+  target_include_directories(device_transport PRIVATE /usr/local/cuda/include)
+endif()
+if(USE_MUSA)
+  target_include_directories(device_transport PRIVATE /usr/local/musa/include)
+  target_compile_definitions(device_transport PRIVATE USE_MUSA)
+endif()
diff --git a/mooncake-transfer-engine/src/transport/device/ibgda_device_transport.cpp b/mooncake-transfer-engine/src/transport/device/ibgda_device_transport.cpp
new file mode 100644
index 00000000..2dcdac17
--- /dev/null
+++ b/mooncake-transfer-engine/src/transport/device/ibgda_device_transport.cpp
@@ -0,0 +1,430 @@
+// Copyright 2024 KVCache.AI
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+// IBGDA device transport — wraps mlx5gda QP lifecycle behind RdmaTransport.
+//
+// Uses TE's Topology for NIC auto-detection when device_name is empty.
+// Uses ibv_open_device / ibv_alloc_pd directly (same as the original
+// MooncakeEpBuffer::init_ibgda()), but encapsulated in a reusable class.
+
+#include "transport/device/device_transport.h"
+
+#include <arpa/inet.h>
+#include <glog/logging.h>
+#include <infiniband/mlx5dv.h>
+#include <infiniband/verbs.h>
+
+#include <cstring>
+#include <stdexcept>
+
+#include "cuda_alike.h"
+#include "transport/device/ibgda/memheap.h"
+#include "transport/device/ibgda/mlx5gda.h"
+#include "topology.h"
+
+namespace mooncake {
+namespace device {
+
+static constexpr size_t kCtrlBufSize = 1024ULL * 1024 * 1024;  // 1 GiB
+
+// Check if IPv6 address is IPv4-mapped (::ffff:x.x.x.x)
+static bool isIpv4Mapped(const struct in6_addr* a) {
+    return ((a->s6_addr32[0] | a->s6_addr32[1]) == 0 &&
+            a->s6_addr32[2] == htonl(0x0000ffff));
+}
+
+static int findBestGidIndex(ibv_context* ctx, uint8_t port,
+                            const ibv_port_attr& port_attr) {
+    for (int i = 0; i < port_attr.gid_tbl_len; ++i) {
+        ibv_gid_entry entry;
+        if (ibv_query_gid_ex(ctx, port, i, &entry, 0)) continue;
+
+        if (entry.gid_type == IBV_GID_TYPE_ROCE_V2) {
+            bool v4mapped = isIpv4Mapped(
+                reinterpret_cast<const struct in6_addr*>(entry.gid.raw));
+            if (v4mapped) return i;
+        } else if (entry.gid_type == IBV_GID_TYPE_IB) {
+            return i;
+        }
+    }
+
+    return -1;
+}
+
+// Auto-detect the best NIC for the current GPU using TE's Topology.
+// filter: if non-empty, only consider NICs in this list.
+static std::string autoDetectNic(const std::vector<std::string>& filter) {
+    Topology topo;
+    if (topo.discover(filter) != 0) return "";
+    const auto& hca_list = topo.getHcaList();
+    if (hca_list.empty()) return "";
+
+    // Build a location string for the current GPU so Topology picks the
+    // topologically closest NIC.  Fall back to wildcard if cudaGetDevice fails.
+    int device_id = 0;
+    cudaGetDevice(&device_id);
+    std::string location = "cuda:" + std::to_string(device_id);
+
+    int idx = topo.selectDevice(location);
+    if (idx < 0) idx = topo.selectDevice("*");  // wildcard fallback
+    if (idx < 0 || idx >= static_cast<int>(hca_list.size())) return "";
+    return hca_list[idx];
+}
+
+class IbgdaDeviceTransportImpl : public RdmaTransport {
+   public:
+    explicit IbgdaDeviceTransportImpl(std::vector<std::string> filter)
+        : device_filter_(std::move(filter)) {}
+
+    ~IbgdaDeviceTransportImpl() override { teardown(); }
+
+    int initialize(const std::string& device_name, int num_ranks,
+                   int num_qps) override {
+        if (num_qps < num_ranks) {
+            LOG(ERROR) << "[EP IBGDA] num_qps (" << num_qps
+                       << ") must be >= num_ranks (" << num_ranks << ")";
+            return -1;
+        }
+        num_ranks_ = num_ranks;
+        num_qps_ = num_qps;
+
+        std::string nic =
+            device_name.empty() ? autoDetectNic(device_filter_) : device_name;
+        if (nic.empty()) {
+            LOG(WARNING) << "[EP IBGDA] No RDMA NIC found";
+            return -1;
+        }
+
+        int num_devices = 0;
+        ibv_device** dev_list = ibv_get_device_list(&num_devices);
+        if (!dev_list) {
+            LOG(ERROR) << "[EP IBGDA] ibv_get_device_list failed";
+            return -1;
+        }
+
+        ibv_device* dev = nullptr;
+        for (int i = 0; i < num_devices; ++i) {
+            if (nic == ibv_get_device_name(dev_list[i])) {
+                dev = dev_list[i];
+                break;
+            }
+        }
+        if (!dev) {
+            LOG(ERROR) << "[EP IBGDA] NIC '" << nic << "' not found";
+            ibv_free_device_list(dev_list);
+            return -1;
+        }
+
+        ctx_ = ibv_open_device(dev);
+        ibv_free_device_list(dev_list);
+        if (!ctx_) {
+            LOG(ERROR) << "[EP IBGDA] ibv_open_device failed for " << nic;
+            return -1;
+        }
+
+        const uint8_t port = 1;
+        ibv_port_attr port_attr{};
+        if (ibv_query_port(ctx_, port, &port_attr)) {
+            LOG(ERROR) << "[EP IBGDA] ibv_query_port failed";
+            return -1;
+        }
+
+        gid_index_ = findBestGidIndex(ctx_, port, port_attr);
+        if (gid_index_ < 0) {
+            LOG(ERROR) << "[EP IBGDA] No suitable GID on " << nic;
+            return -1;
+        }
+
+        if (ibv_query_gid(ctx_, port, gid_index_, &gid_)) {
+            LOG(ERROR) << "[EP IBGDA] ibv_query_gid failed";
+            return -1;
+        }
+
+        is_roce_ = (port_attr.link_layer == IBV_LINK_LAYER_ETHERNET);
+        lid_ = port_attr.lid;
+        device_name_ = nic;
+
+        pd_ = ibv_alloc_pd(ctx_);
+        if (!pd_) {
+            LOG(ERROR) << "[EP IBGDA] ibv_alloc_pd failed";
+            return -1;
+        }
+
+        mlx5dv_obj dv_obj{};
+        dv_obj.pd.in = pd_;
+        dv_obj.pd.out = &mpd_;
+        if (mlx5dv_init_obj(&dv_obj, MLX5DV_OBJ_PD)) {
+            LOG(ERROR) << "[EP IBGDA] mlx5dv_init_obj failed";
+            return -1;
+        }
+
+        // Allocate device-visible tables
+        if (cudaMalloc(&raddrs_, num_ranks_ * sizeof(uint64_t)) !=
+                cudaSuccess ||
+            cudaMalloc(&rkeys_, num_ranks_ * sizeof(uint32_t)) != cudaSuccess ||
+            cudaMalloc(&qp_devctxs_, num_qps_ * sizeof(mlx5gda_qp_devctx)) !=
+                cudaSuccess) {
+            LOG(ERROR) << "[EP IBGDA] Failed to allocate device-visible tables";
+            return -1;
+        }
+
+        LOG(INFO) << "[EP IBGDA] Initialized on " << nic
+                  << " (gid_index=" << gid_index_ << ", roce=" << is_roce_
+                  << ")";
+        return 0;
+    }
+
+    int registerMemory(void* ptr, size_t bytes) override {
+        mr_ =
+            ibv_reg_mr(pd_, ptr, bytes,
+                       IBV_ACCESS_LOCAL_WRITE | IBV_ACCESS_REMOTE_READ |
+                           IBV_ACCESS_REMOTE_WRITE | IBV_ACCESS_REMOTE_ATOMIC);
+        if (!mr_) {
+            LOG(ERROR) << "[EP IBGDA] ibv_reg_mr failed";
+            return -1;
+        }
+        mr_ptr_ = ptr;
+        return 0;
+    }
+
+    int allocateControlBuffer() override {
+        cudaError_t err = cudaMalloc(&ctrl_buf_, kCtrlBufSize);
+        if (err != cudaSuccess) {
+            LOG(ERROR) << "[EP IBGDA] cudaMalloc ctrl_buf failed: "
+                       << cudaGetErrorString(err);
+            return -1;
+        }
+
+        ctrl_buf_umem_ = mlx5dv_devx_umem_reg(ctx_, ctrl_buf_, kCtrlBufSize,
+                                              IBV_ACCESS_LOCAL_WRITE);
+        if (!ctrl_buf_umem_) {
+            LOG(ERROR) << "[EP IBGDA] mlx5dv_devx_umem_reg failed (errno="
+                       << errno << ")";
+            return -1;
+        }
+        LOG(INFO) << "[EP IBGDA] ctrl_buf UMEM registered via VA path";
+
+        ctrl_buf_heap_ = memheap_create(kCtrlBufSize);
+        if (!ctrl_buf_heap_) {
+            LOG(ERROR) << "[EP IBGDA] memheap_create failed";
+            return -1;
+        }
+        return 0;
+    }
+
+    int createQueuePairs(void* stream_ptr) override {
+        auto stream = static_cast<cudaStream_t>(stream_ptr);
+        for (int i = 0; i < num_qps_; ++i) {
+            mlx5gda_qp* qp =
+                mlx5gda_create_rc_qp(mpd_, ctrl_buf_, ctrl_buf_umem_,
+                                     ctrl_buf_heap_, pd_, 16384, 1, stream);
+            if (!qp) {
+                LOG(ERROR) << "[EP IBGDA] mlx5gda_create_rc_qp failed at " << i;
+                return -1;
+            }
+            if (mlx5gda_modify_rc_qp_rst2init(qp, 0)) {
+                LOG(ERROR) << "[EP IBGDA] rst2init failed at " << i;
+                mlx5gda_destroy_qp(ctrl_buf_heap_, qp);
+                return -1;
+            }
+            cudaStreamSynchronize(stream);
+            mlx5gda_qp_devctx devctx{
+                .qpn = qp->qpn,
+                .wqeid_mask = qp->num_wqebb - 1,
+                .wq = reinterpret_cast<mlx5gda_wqebb*>(
+                    static_cast<char*>(ctrl_buf_) + qp->wq_offset),
+                .cq = reinterpret_cast<mlx5_cqe64*>(
+                    static_cast<char*>(ctrl_buf_) + qp->send_cq->cq_offset),
+                .dbr = reinterpret_cast<mlx5gda_wq_dbr*>(
+                    static_cast<char*>(ctrl_buf_) + qp->dbr_offset),
+                .bf = static_cast<char*>(qp->uar->reg_addr),
+            };
+            cudaMemcpy(
+                static_cast<char*>(qp_devctxs_) + i * sizeof(mlx5gda_qp_devctx),
+                &devctx, sizeof(mlx5gda_qp_devctx), cudaMemcpyHostToDevice);
+            qps_.push_back(qp);
+        }
+        return 0;
+    }
+
+    int recreateQueuePairs(void* stream_ptr) override {
+        auto stream = static_cast<cudaStream_t>(stream_ptr);
+        for (auto* qp : qps_) {
+            if (qp) mlx5gda_destroy_qp(ctrl_buf_heap_, qp);
+        }
+        qps_.clear();
+        return createQueuePairs(stream_ptr);
+    }
+
+    int connectPeers(int local_rank, bool is_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<int32_t>& remote_lids,
+                     const std::vector<int64_t>& subnet_prefixes,
+                     const std::vector<int64_t>& interface_ids,
+                     const std::vector<int>& active_ranks_mask) override {
+        for (int i = 0; i < num_qps_; ++i) {
+            int peer_rank = i * num_ranks_ / num_qps_;
+            if (active_ranks_mask[peer_rank] == 0) continue;
+
+            ibv_ah_attr ah_attr{};
+            if (is_roce) {
+                ibv_gid remote_gid{};
+                remote_gid.global.subnet_prefix = subnet_prefixes[peer_rank];
+                remote_gid.global.interface_id = interface_ids[peer_rank];
+                ah_attr.is_global = 1;
+                ah_attr.grh.dgid = remote_gid;
+                ah_attr.grh.sgid_index = gid_index_;
+                // Match the legacy IBGDA path.  mlx5gda previously hard-coded
+                // QPC hop_limit=255 and ignored this field; after moving QP
+                // setup into the device transport this value reaches hardware.
+                ah_attr.grh.hop_limit = 255;
+                ah_attr.port_num = 1;
+                ah_attr.dlid = qps_[i]->port_attr.lid | 0xC000;
+            } else {
+                ah_attr.dlid = static_cast<uint16_t>(remote_lids[i]);
+                ah_attr.port_num = 0;
+            }
+
+            if (mlx5gda_modify_rc_qp_init2rtr(qps_[i], ah_attr, remote_qpns[i],
+                                              IBV_MTU_4096)) {
+                LOG(ERROR) << "[EP IBGDA] init2rtr failed for QP " << i
+                           << " (roce=" << is_roce << " gid_idx=" << gid_index_
+                           << " remote_qpn=" << remote_qpns[i]
+                           << " udp_sport=" << ah_attr.dlid
+                           << " hop_limit=" << (int)ah_attr.grh.hop_limit
+                           << ")";
+                return -1;
+            }
+            if (mlx5gda_modify_rc_qp_rtr2rts(qps_[i])) {
+                LOG(ERROR) << "[EP IBGDA] rtr2rts failed for QP " << i;
+                return -1;
+            }
+        }
+
+        // Populate device-visible raddrs/rkeys tables
+        for (int i = 0; i < num_ranks_; ++i) {
+            if (active_ranks_mask[i] == 0) continue;
+            uint64_t raddr = static_cast<uint64_t>(remote_addrs[i]);
+            uint32_t rkey = (i == local_rank)
+                                ? static_cast<uint32_t>(mr_->lkey)
+                                : static_cast<uint32_t>(remote_keys[i]);
+            cudaMemcpy(static_cast<char*>(raddrs_) + i * sizeof(uint64_t),
+                       &raddr, sizeof(uint64_t), cudaMemcpyHostToDevice);
+            cudaMemcpy(static_cast<char*>(rkeys_) + i * sizeof(uint32_t), &rkey,
+                       sizeof(uint32_t), cudaMemcpyHostToDevice);
+        }
+        return 0;
+    }
+
+    RdmaLocalMetadata localMetadata() const override {
+        RdmaLocalMetadata meta;
+        meta.raddr = mr_ ? reinterpret_cast<int64_t>(mr_->addr) : 0;
+        meta.rkey = mr_ ? static_cast<int32_t>(mr_->rkey) : 0;
+        meta.subnet_prefix = static_cast<int64_t>(gid_.global.subnet_prefix);
+        meta.interface_id = static_cast<int64_t>(gid_.global.interface_id);
+        for (auto* qp : qps_) {
+            meta.qpns.push_back(static_cast<int32_t>(qp->qpn));
+            meta.lids.push_back(static_cast<int32_t>(lid_));
+        }
+        return meta;
+    }
+
+    void* raddrsPtr() override { return raddrs_; }
+    void* rkeysPtr() override { return rkeys_; }
+    void* qpDevCtxsPtr() override { return qp_devctxs_; }
+    bool isRoce() const override { return is_roce_; }
+    int gidIndex() const override { return gid_index_; }
+
+   private:
+    void teardown() {
+        for (auto* qp : qps_) {
+            if (qp) mlx5gda_destroy_qp(ctrl_buf_heap_, qp);
+        }
+        qps_.clear();
+        if (ctrl_buf_heap_) {
+            memheap_destroy(ctrl_buf_heap_);
+            ctrl_buf_heap_ = nullptr;
+        }
+        if (ctrl_buf_umem_) {
+            mlx5dv_devx_umem_dereg(ctrl_buf_umem_);
+            ctrl_buf_umem_ = nullptr;
+        }
+        if (ctrl_buf_) {
+            cudaFree(ctrl_buf_);
+            ctrl_buf_ = nullptr;
+        }
+        if (mr_) {
+            ibv_dereg_mr(mr_);
+            mr_ = nullptr;
+        }
+        if (raddrs_) {
+            cudaFree(raddrs_);
+            raddrs_ = nullptr;
+        }
+        if (rkeys_) {
+            cudaFree(rkeys_);
+            rkeys_ = nullptr;
+        }
+        if (qp_devctxs_) {
+            cudaFree(qp_devctxs_);
+            qp_devctxs_ = nullptr;
+        }
+        if (pd_) {
+            ibv_dealloc_pd(pd_);
+            pd_ = nullptr;
+        }
+        if (ctx_) {
+            ibv_close_device(ctx_);
+            ctx_ = nullptr;
+        }
+    }
+
+    // IB resources
+    ibv_context* ctx_ = nullptr;
+    ibv_pd* pd_ = nullptr;
+    mlx5dv_pd mpd_{};
+    ibv_mr* mr_ = nullptr;
+    void* mr_ptr_ = nullptr;
+    ibv_gid gid_{};
+    int gid_index_ = -1;
+    uint16_t lid_ = 0;
+    bool is_roce_ = false;
+    std::string device_name_;
+    std::vector<std::string> device_filter_;
+
+    // Control buffer
+    void* ctrl_buf_ = nullptr;  // GPU VA
+    mlx5dv_devx_umem* ctrl_buf_umem_ = nullptr;
+    memheap* ctrl_buf_heap_ = nullptr;
+
+    // QPs
+    std::vector<mlx5gda_qp*> qps_;
+    int num_ranks_ = 0;
+    int num_qps_ = 0;
+
+    // Device-visible tables
+    void* raddrs_ = nullptr;
+    void* rkeys_ = nullptr;
+    void* qp_devctxs_ = nullptr;
+};
+
+std::unique_ptr<RdmaTransport> createIbgdaDeviceTransport(
+    const std::vector<std::string>& device_filter) {
+    return std::make_unique<IbgdaDeviceTransportImpl>(device_filter);
+}
+
+}  // namespace device
+}  // namespace mooncake
diff --git a/mooncake-ep/src/mooncake_ibgda/mlx5gda.cpp b/mooncake-transfer-engine/src/transport/device/mlx5gda.cpp
similarity index 91%
rename from mooncake-ep/src/mooncake_ibgda/mlx5gda.cpp
rename to mooncake-transfer-engine/src/transport/device/mlx5gda.cpp
index 96e8f939..ca601b25 100644
--- a/mooncake-ep/src/mooncake_ibgda/mlx5gda.cpp
+++ b/mooncake-transfer-engine/src/transport/device/mlx5gda.cpp
@@ -1,14 +1,14 @@
 #include <cmath>
 
-#include <cuda_runtime_api.h>
+#include "cuda_alike.h"
 
 #include <infiniband/verbs.h>
 #include <infiniband/mlx5dv.h>
 
-#include <mooncake_ibgda/memheap.h>
-#include <mooncake_ibgda/mlx5gda.h>
-#include <mooncake_ibgda/mlx5_ifc.h>
-#include <mooncake_ibgda/mlx5_prm.h>
+#include <transport/device/ibgda/memheap.h>
+#include <transport/device/ibgda/mlx5gda.h>
+#include <transport/device/ibgda/mlx5_ifc.h>
+#include <transport/device/ibgda/mlx5_prm.h>
 
 template <typename T>
 inline T IBGDA_ILOG2(T _n) {
@@ -30,10 +30,22 @@ 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) {
+#ifdef USE_MUSA
+    const char* err_str = musaGetErrorString(musaGetLastError());
+#else
     const char* err_str = cudaGetErrorString(cudaGetLastError());
+#endif
     fprintf(stderr, "%s: %s\n", msg, err_str);
 }
 
+// Create UAR for BF (Blue Flame) doorbell ringing.
+// On CUDA: registers the BF MMIO region into GPU address space so the
+//          GPU kernel can directly write the doorbell (lowest latency).
+// On MUSA: musaHostRegisterIoMemory is not supported for MMIO addresses,
+//          so we skip BF registration and return a UAR with reg_addr=NULL.
+//          The GPU kernel will use DBR-only mode (write to memory-mapped
+//          doorbell record, NIC polls it) — slightly higher latency but
+//          functionally correct.
 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);
@@ -41,6 +53,15 @@ static struct mlx5dv_devx_uar* create_uar(struct ibv_context* ctx) {
         errno = EIO;
         return NULL;
     }
+#ifdef USE_MUSA
+    // MUSA cannot map MMIO addresses into GPU VA.  Skip the
+    // musaHostRegister(IoMemory) call entirely — attempting it
+    // corrupts the MUSA runtime, causing all subsequent device-side
+    // fill operations to fail with "illegal memory access".
+    // Use DBR-only mode: the kernel writes to the doorbell record
+    // in GPU memory instead of the BF MMIO register.
+    uar->reg_addr = NULL;
+#else
     if (cudaHostRegister(uar->reg_addr, MLX5GDA_BF_SIZE * 2,
                          cudaHostRegisterPortable | cudaHostRegisterMapped |
                              cudaHostRegisterIoMemory) != cudaSuccess) {
@@ -49,13 +70,16 @@ static struct mlx5dv_devx_uar* create_uar(struct ibv_context* ctx) {
         mlx5dv_devx_free_uar(uar);
         return NULL;
     }
+#endif
     return 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");
+    if (uar->reg_addr) {
+        if (cudaHostUnregister(uar->reg_addr) != cudaSuccess) {
+            print_cuda_error("Failed to unregister MMIO memory");
+        }
     }
     mlx5dv_devx_free_uar(uar);
 }
@@ -444,7 +468,8 @@ int mlx5gda_modify_rc_qp_init2rtr(struct mlx5gda_qp* qp,
 
         memcpy(DEVX_ADDR_OF(qpc, qpc, primary_address_path.rmac_47_32),
                &dah.av->rmac, sizeof(dah.av->rmac));
-        DEVX_SET(qpc, qpc, primary_address_path.hop_limit, 255);
+        DEVX_SET(qpc, qpc, primary_address_path.hop_limit,
+                 ah_attr.grh.hop_limit);
         DEVX_SET(qpc, qpc, primary_address_path.src_addr_index,
                  ah_attr.grh.sgid_index);
         DEVX_SET(qpc, qpc, primary_address_path.udp_sport, ah_attr.dlid);
diff --git a/mooncake-transfer-engine/src/transport/device/p2p_device_transport.cpp b/mooncake-transfer-engine/src/transport/device/p2p_device_transport.cpp
new file mode 100644
index 00000000..23e02255
--- /dev/null
+++ b/mooncake-transfer-engine/src/transport/device/p2p_device_transport.cpp
@@ -0,0 +1,268 @@
+// Copyright 2024 KVCache.AI
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+// P2P device transport — unified NVLink (CUDA) + MTLink (MUSA) implementation.
+//
+// Uses cuda_alike.h so all cuda* APIs map to musa* when USE_MUSA is defined.
+// No #ifdef USE_MUSA / MOONCAKE_EP_USE_MUSA in this file.
+
+#include "transport/device/device_transport.h"
+
+#include <glog/logging.h>
+#include <algorithm>
+#include <cstring>
+
+#include "cuda_alike.h"
+
+namespace mooncake {
+namespace device {
+
+class P2pDeviceTransportImpl : public P2pTransport {
+   public:
+    explicit P2pDeviceTransportImpl(int num_ranks) : num_ranks_(num_ranks) {
+        cudaMalloc(&available_table_, num_ranks_ * sizeof(int32_t));
+        cudaMemset(available_table_, 0, num_ranks_ * sizeof(int32_t));
+        cudaMallocHost(&peer_ptrs_host_, num_ranks_ * sizeof(void*));
+        cudaMalloc(&peer_ptrs_dev_, num_ranks_ * sizeof(void*));
+        for (int i = 0; i < num_ranks_; ++i) peer_ptrs_host_[i] = nullptr;
+        cudaMemset(peer_ptrs_dev_, 0, num_ranks_ * sizeof(void*));
+    }
+
+    ~P2pDeviceTransportImpl() override {
+        if (available_table_) cudaFree(available_table_);
+        if (peer_ptrs_dev_) cudaFree(peer_ptrs_dev_);
+        if (peer_ptrs_host_) {
+            for (int i = 0; i < num_ranks_; ++i) {
+                if (peer_ptrs_host_[i] && peer_ptrs_host_[i] != local_ptr_) {
+                    cudaIpcCloseMemHandle(peer_ptrs_host_[i]);
+                }
+            }
+            cudaFreeHost(peer_ptrs_host_);
+        }
+    }
+
+    void* allocateBuffer(size_t bytes) override {
+        void* ptr = nullptr;
+        cudaError_t err = cudaMalloc(&ptr, bytes);
+        if (err != cudaSuccess) {
+            LOG(ERROR) << "[EP P2P] cudaMalloc(" << bytes
+                       << ") failed: " << cudaGetErrorString(err);
+            return nullptr;
+        }
+        return ptr;
+    }
+
+    void freeBuffer(void* ptr) override { cudaFree(ptr); }
+
+    std::vector<int32_t> exportIpcHandle(void* ptr) override {
+        cudaIpcMemHandle_t handle;
+        cudaError_t err = cudaIpcGetMemHandle(&handle, ptr);
+        if (err != cudaSuccess) {
+            LOG(ERROR) << "[EP P2P] cudaIpcGetMemHandle failed: "
+                       << cudaGetErrorString(err);
+            return {};
+        }
+        constexpr size_t kHandleBytes = sizeof(cudaIpcMemHandle_t);
+        constexpr size_t kNumInt32s =
+            (kHandleBytes + sizeof(int32_t) - 1) / sizeof(int32_t);
+        std::vector<int32_t> result(kNumInt32s);
+        memcpy(result.data(), &handle, kHandleBytes);
+        return result;
+    }
+
+    void importPeerHandles(
+        void* local_ptr, int rank, int num_ranks,
+        const std::vector<std::vector<int32_t>>& remote_handles,
+        const std::vector<int>& active_ranks_mask) override {
+        local_ptr_ = local_ptr;
+        int device_id = 0;
+        cudaGetDevice(&device_id);
+        int device_count = 0;
+        cudaGetDeviceCount(&device_count);
+        CHECK_GT(device_count, 0) << "No CUDA/MUSA devices found";
+
+        std::vector<int32_t> available(num_ranks_, 0);
+        available[rank] = 1;
+        peer_ptrs_host_[rank] = local_ptr;
+
+        int node_id = rank / device_count;
+        int group_start = node_id * device_count;
+        int group_end = std::min(group_start + device_count, num_ranks_);
+
+        for (int dst = group_start; dst < group_end; ++dst) {
+            if (active_ranks_mask[dst] == 0) continue;
+            if (dst == rank) continue;
+
+            int dst_device = dst % device_count;
+            int can_access = 0;
+            cudaDeviceCanAccessPeer(&can_access, device_id, dst_device);
+            LOG(INFO) << "[EP P2P] rank " << rank << " (device " << device_id
+                      << ") -> rank " << dst << " (device " << dst_device
+                      << "): canAccessPeer=" << can_access;
+            if (!can_access) continue;
+
+            cudaError_t err = cudaDeviceEnablePeerAccess(dst_device, 0);
+            if (err != cudaSuccess &&
+                err != cudaErrorPeerAccessAlreadyEnabled) {
+                LOG(WARNING) << "[EP P2P] rank " << rank
+                             << " failed to enable peer access to device "
+                             << dst_device << ": " << cudaGetErrorString(err);
+                continue;
+            }
+            if (err == cudaErrorPeerAccessAlreadyEnabled) cudaGetLastError();
+
+            if (dst >= static_cast<int>(remote_handles.size())) continue;
+            const auto& h = remote_handles[dst];
+            if (h.empty()) continue;
+
+            constexpr size_t kHandleBytes = sizeof(cudaIpcMemHandle_t);
+            constexpr size_t kNumInt32s =
+                (kHandleBytes + sizeof(int32_t) - 1) / sizeof(int32_t);
+            if (h.size() < kNumInt32s) continue;
+
+            cudaIpcMemHandle_t handle;
+            memcpy(&handle, h.data(), kHandleBytes);
+
+            void* peer_ptr = nullptr;
+            err = cudaIpcOpenMemHandle(&peer_ptr, handle,
+                                       cudaIpcMemLazyEnablePeerAccess);
+            if (err != cudaSuccess) {
+                LOG(WARNING) << "[EP P2P] rank " << rank
+                             << " failed to open IPC handle for rank " << dst
+                             << ": " << cudaGetErrorString(err);
+                continue;
+            }
+            LOG(INFO) << "[EP P2P] rank " << rank
+                      << " opened IPC handle for rank " << dst
+                      << ": peer_ptr=" << peer_ptr;
+            available[dst] = 1;
+            peer_ptrs_host_[dst] = peer_ptr;
+        }
+
+        // Determine if all active ranks have P2P
+        all_peers_accessible_ = true;
+        for (int i = 0; i < num_ranks_; ++i) {
+            if (active_ranks_mask[i] == 0) continue;
+            if (!available[i] || !peer_ptrs_host_[i]) {
+                all_peers_accessible_ = false;
+                break;
+            }
+        }
+        // Multi-node: P2P only within a node
+        if (all_peers_accessible_ && num_ranks_ > 1) {
+            int first_node = 0 / device_count;
+            int last_node = (num_ranks_ - 1) / device_count;
+            if (first_node != last_node) all_peers_accessible_ = false;
+        }
+
+        cudaMemcpy(available_table_, available.data(),
+                   num_ranks_ * sizeof(int32_t), cudaMemcpyHostToDevice);
+        cudaMemcpy(peer_ptrs_dev_, peer_ptrs_host_, num_ranks_ * sizeof(void*),
+                   cudaMemcpyHostToDevice);
+    }
+
+    int32_t* availableTablePtr() override { return available_table_; }
+    void** peerPtrsTablePtr() override { return peer_ptrs_dev_; }
+    bool allPeersAccessible() const override { return all_peers_accessible_; }
+
+    bool verifyPeerAccess() override {
+        if (!all_peers_accessible_) return false;
+
+        int device_id = 0;
+        cudaGetDevice(&device_id);
+
+        bool all_ok = true;
+        for (int i = 0; i < num_ranks_; ++i) {
+            if (i == device_id) continue;
+            if (!peer_ptrs_host_[i] || peer_ptrs_host_[i] == local_ptr_)
+                continue;
+
+            // Test: write a pattern to the peer buffer via cudaMemcpy, then
+            // read it back.  This verifies the IPC mapping is writable.
+            constexpr int kTestBytes = 256;
+            std::vector<uint8_t> pattern(kTestBytes);
+            for (int j = 0; j < kTestBytes; ++j)
+                pattern[j] = (uint8_t)(j ^ 0xA5);
+
+            cudaError_t err = cudaMemcpy(peer_ptrs_host_[i], pattern.data(),
+                                         kTestBytes, cudaMemcpyHostToDevice);
+            if (err != cudaSuccess) {
+                LOG(WARNING) << "[EP P2P] verifyPeerAccess: rank " << device_id
+                             << " cannot write to peer " << i
+                             << " mapped buffer: " << cudaGetErrorString(err);
+                all_ok = false;
+                continue;
+            }
+
+            // Read back and verify
+            std::vector<uint8_t> readback(kTestBytes, 0);
+            err = cudaMemcpy(readback.data(), peer_ptrs_host_[i], kTestBytes,
+                             cudaMemcpyDeviceToHost);
+            if (err != cudaSuccess) {
+                LOG(WARNING) << "[EP P2P] verifyPeerAccess: rank " << device_id
+                             << " cannot read back from peer " << i
+                             << " mapped buffer: " << cudaGetErrorString(err);
+                all_ok = false;
+                continue;
+            }
+
+            bool match =
+                (memcmp(readback.data(), pattern.data(), kTestBytes) == 0);
+            if (!match) {
+                LOG(WARNING) << "[EP P2P] verifyPeerAccess: rank " << device_id
+                             << " readback mismatch from peer " << i;
+                all_ok = false;
+            } else {
+                LOG(INFO) << "[EP P2P] verifyPeerAccess: rank " << device_id
+                          << " peer " << i << " OK (memcpy write/read)";
+            }
+        }
+
+        if (!all_ok) {
+            all_peers_accessible_ = false;
+            // Update device table to reflect failure
+            std::vector<int32_t> avail_h(num_ranks_, 0);
+            cudaMemcpy(avail_h.data(), available_table_,
+                       num_ranks_ * sizeof(int32_t), cudaMemcpyDeviceToHost);
+            for (int i = 0; i < num_ranks_; ++i) {
+                if (i == device_id) continue;
+                if (!peer_ptrs_host_[i] || peer_ptrs_host_[i] == local_ptr_)
+                    continue;
+                // Leave self as available; only clear failed peers
+            }
+            // Re-upload with all peers marked unavailable except self
+            std::vector<int32_t> cleared(num_ranks_, 0);
+            cleared[device_id] = 1;
+            cudaMemcpy(available_table_, cleared.data(),
+                       num_ranks_ * sizeof(int32_t), cudaMemcpyHostToDevice);
+        }
+
+        return all_ok;
+    }
+
+   private:
+    int num_ranks_;
+    void* local_ptr_ = nullptr;
+    int32_t* available_table_ = nullptr;
+    void** peer_ptrs_host_ = nullptr;
+    void** peer_ptrs_dev_ = nullptr;
+    bool all_peers_accessible_ = false;
+};
+
+std::unique_ptr<P2pTransport> createP2pDeviceTransport(int num_ranks) {
+    return std::make_unique<P2pDeviceTransportImpl>(num_ranks);
+}
+
+}  // namespace device
+}  // namespace mooncake
diff --git a/mooncake-transfer-engine/src/transport/efa_transport/efa_context.cpp b/mooncake-transfer-engine/src/transport/efa_transport/efa_context.cpp
index 4c96b7e1..6f8abe49 100644
--- a/mooncake-transfer-engine/src/transport/efa_transport/efa_context.cpp
+++ b/mooncake-transfer-engine/src/transport/efa_transport/efa_context.cpp
@@ -19,12 +19,12 @@
 
 #include <atomic>
 #include <cassert>
-#include <fstream>
-#include <memory>
-#include <thread>
 #include <cstring>
+#include <fstream>
 #include <iomanip>
+#include <memory>
 #include <sstream>
+#include <thread>
 
 #include "config.h"
 #include "cuda_alike.h"
@@ -34,117 +34,8 @@
 
 namespace mooncake {
 
-// EfaEndpointStore implementation
-
-EfaEndpointStore::EfaEndpointStore(size_t max_endpoints,
-                                   double inactive_timeout_sec)
-    : max_endpoints_(max_endpoints),
-      inactive_timeout_sec_(inactive_timeout_sec) {}
-
-std::shared_ptr<EfaEndPoint> EfaEndpointStore::get(
-    const std::string& peer_nic_path) {
-    RWSpinlock::ReadGuard guard(lock_);
-    auto it = endpoints_.find(peer_nic_path);
-    if (it != endpoints_.end()) {
-        it->second->set_active(true);
-        return it->second;
-    }
-    return nullptr;
-}
-
-std::shared_ptr<EfaEndPoint> EfaEndpointStore::getOrInsert(
-    const std::string& peer_nic_path, std::shared_ptr<EfaEndPoint> new_ep) {
-    RWSpinlock::WriteGuard guard(lock_);
-    auto it = endpoints_.find(peer_nic_path);
-    if (it != endpoints_.end()) {
-        it->second->set_active(true);
-        return it->second;  // Another thread already created it
-    }
-    // Evict stale endpoints if at capacity
-    if (endpoints_.size() >= max_endpoints_) {
-        size_t evicted = evictStaleLocked();
-        if (evicted == 0 && endpoints_.size() >= max_endpoints_) {
-            LOG(WARNING) << "EfaEndpointStore at capacity (" << max_endpoints_
-                         << ") with no stale endpoints to evict";
-        }
-    }
-    endpoints_[peer_nic_path] = new_ep;
-    new_ep->set_active(true);
-    return new_ep;
-}
-
-void EfaEndpointStore::add(const std::string& peer_nic_path,
-                           std::shared_ptr<EfaEndPoint> endpoint) {
-    RWSpinlock::WriteGuard guard(lock_);
-    endpoints_[peer_nic_path] = endpoint;
-}
-
-void EfaEndpointStore::remove(const std::string& peer_nic_path) {
-    RWSpinlock::WriteGuard guard(lock_);
-    endpoints_.erase(peer_nic_path);
-}
-
-int EfaEndpointStore::disconnectAll() {
-    RWSpinlock::WriteGuard guard(lock_);
-    for (auto& entry : endpoints_) {
-        if (entry.second) {
-            entry.second->disconnect();
-        }
-    }
-    return 0;
-}
-
-size_t EfaEndpointStore::size() const {
-    RWSpinlock::ReadGuard guard(lock_);
-    return endpoints_.size();
-}
-
-size_t EfaEndpointStore::evictStale() {
-    RWSpinlock::WriteGuard guard(lock_);
-    return evictStaleLocked();
-}
-
-size_t EfaEndpointStore::evictStaleLocked() {
-    size_t evicted = 0;
-    for (auto it = endpoints_.begin(); it != endpoints_.end();) {
-        auto& ep = it->second;
-        double age = ep ? ep->lastUsedAge() : 0;
-        bool outstanding = ep ? ep->hasOutstandingSlice() : false;
-        if (ep && !outstanding && age > inactive_timeout_sec_) {
-            LOG(INFO) << "Evicting stale EFA endpoint: " << it->first
-                      << " (idle " << age
-                      << "s, timeout=" << inactive_timeout_sec_ << "s)";
-            ep->disconnect();
-            it = endpoints_.erase(it);
-            ++evicted;
-        } else {
-            ++it;
-        }
-    }
-    if (endpoints_.size() > 0) {
-        VLOG(1) << "evictStale: " << endpoints_.size() << " endpoints, "
-                << evicted << " evicted";
-    }
-    return evicted;
-}
-
-size_t EfaEndpointStore::removeDisconnected() {
-    RWSpinlock::WriteGuard guard(lock_);
-    size_t removed = 0;
-    for (auto it = endpoints_.begin(); it != endpoints_.end();) {
-        auto& ep = it->second;
-        if (ep && !ep->connected() && !ep->hasOutstandingSlice()) {
-            LOG(INFO) << "Removing disconnected EFA endpoint: " << it->first;
-            it = endpoints_.erase(it);
-            ++removed;
-        } else {
-            ++it;
-        }
-    }
-    return removed;
-}
-
 // EfaContext implementation
+
 EfaContext::EfaContext(EfaTransport& engine, const std::string& device_name)
     : engine_(engine),
       device_name_(device_name),
@@ -153,17 +44,18 @@ EfaContext::EfaContext(EfaTransport& engine, const std::string& device_name)
       fabric_(nullptr),
       domain_(nullptr),
       av_(nullptr),
-      active_(true) {}
+      active_(true),
+      shared_ep_(nullptr),
+      wr_depth_(0),
+      max_wr_depth_(0),
+      post_lock_(ATOMIC_FLAG_INIT) {}
 
 EfaContext::~EfaContext() {
     if (fabric_) deconstruct();
 }
 
-int EfaContext::construct(size_t num_cq_list, size_t num_comp_channels,
-                          uint8_t port, int gid_index, size_t max_cqe,
+int EfaContext::construct(size_t num_cq_list, size_t max_cqe,
                           int max_endpoints) {
-    endpoint_store_ = std::make_shared<EfaEndpointStore>(max_endpoints);
-
 #if !defined(USE_CUDA) && !defined(USE_HIP)
     // When built without GPU support, prevent libfabric's EFA provider from
     // dlopen-ing libcudart/libcuda at fi_getinfo/fi_domain time. That
@@ -181,7 +73,20 @@ int EfaContext::construct(size_t num_cq_list, size_t num_comp_channels,
     }
 
     hints_->caps =
-        FI_MSG | FI_RMA | FI_READ | FI_WRITE | FI_REMOTE_READ | FI_REMOTE_WRITE;
+        FI_MSG | FI_RMA | FI_READ | FI_WRITE | FI_REMOTE_READ |
+        FI_REMOTE_WRITE
+#if defined(USE_CUDA) || defined(USE_HIP)
+        // Declare FI_HMEM so the provider wires up HMEM-aware copy routines
+        // (cudaMemcpy) on every data path, including the intra-node SHM SAR
+        // segmentation/reassembly path. Registering device memory via
+        // FI_MR_HMEM alone is NOT enough: without FI_HMEM in caps the SHM
+        // sub-provider initializes its copy callbacks in plain-host mode and
+        // does a host memcpy() straight into a CUDA device VA during SAR,
+        // which SIGSEGVs in __memcpy_avx512_unaligned_erms.
+        // See https://github.com/ofiwg/libfabric/issues/12328.
+        | FI_HMEM
+#endif
+        ;
     hints_->mode = FI_CONTEXT;
     hints_->ep_attr->type = FI_EP_RDM;  // EFA uses RDM endpoints
     hints_->fabric_attr->prov_name = strdup("efa");
@@ -197,9 +102,29 @@ int EfaContext::construct(size_t num_cq_list, size_t num_comp_channels,
         ;
     hints_->domain_attr->threading = FI_THREAD_SAFE;
 
-    // Get fabric info
+    // Get fabric info.
+    //
+    // Request libfabric API 1.18+ so the EFA provider's
+    // efa_rdm_get_use_device_rdma() takes the "new API" branch and keys
+    // the default for FI_EFA_USE_DEVICE_RDMA on hardware capability
+    // (hw_support) instead of vendor_part_id.  Under the older 1.14
+    // request, the provider's legacy branch hardcoded
+    //   default_val = (vendor_part_id == 0xefa0 || 0xefa1) ? false : true
+    // which silently disabled device RDMA on Nitro v4 EFA (p5/p5e, part
+    // id 0xefa1) while leaving it enabled on Nitro v5+ (p5en and newer).
+    // With device RDMA disabled, fi_write falls back to libfabric's
+    // emulated RDMA data path, and libfabric 2.4.0 has a thread-safety
+    // regression there between fi_av_insert and concurrent fi_cq_read
+    // that segfaults Mooncake once the handshake wave finishes and the
+    // first real transfers start.  Bumping the requested API to 1.18
+    // restores the same default path we already got on newer hardware,
+    // and applications that still want emulated RDMA can opt out with
+    // FI_EFA_USE_DEVICE_RDMA=0.
+    //
+    // 1.18 is from March 2023 (EFA installer 1.26+ ships 1.18 or later);
+    // all Mooncake deployments today run libfabric >> 1.18.
     int ret =
-        fi_getinfo(FI_VERSION(1, 14), nullptr, nullptr, 0, hints_, &fi_info_);
+        fi_getinfo(FI_VERSION(1, 18), nullptr, nullptr, 0, hints_, &fi_info_);
     if (ret) {
         LOG(ERROR) << "fi_getinfo failed for device " << device_name_ << ": "
                    << fi_strerror(-ret);
@@ -232,7 +157,9 @@ int EfaContext::construct(size_t num_cq_list, size_t num_comp_channels,
         return ERR_CONTEXT;
     }
 
-    // Create address vector
+    // Create address vector.  Capacity sized for the largest peer count we
+    // expect to support in a single process.  AV entries are cheap (no QP
+    // cost), so we over-provision.
     struct fi_av_attr av_attr = {};
     av_attr.type = FI_AV_TABLE;
     av_attr.count = max_endpoints;
@@ -269,18 +196,106 @@ int EfaContext::construct(size_t num_cq_list, size_t num_comp_channels,
         cq_list_[i] = cq;
     }
 
+    // Build the shared endpoint that services every peer through AV lookup.
+    ret = buildSharedEndpoint(globalConfig().max_wr, 64);
+    if (ret) {
+        LOG(ERROR) << "EfaContext::construct: buildSharedEndpoint failed for "
+                   << device_name_;
+        return ret;
+    }
+
     LOG(INFO) << "EFA device (libfabric): " << device_name_
               << ", domain: " << fi_info_->domain_attr->name
-              << ", provider: " << fi_info_->fabric_attr->prov_name;
+              << ", fabric: " << fi_info_->fabric_attr->name
+              << ", provider: " << fi_info_->fabric_attr->prov_name
+              << " (shared endpoint, max_wr=" << max_wr_depth_ << ")";
 
     return 0;
 }
 
+int EfaContext::buildSharedEndpoint(size_t max_wr, size_t max_inline) {
+    (void)max_inline;
+    if (shared_ep_) return 0;
+
+    shared_cq_ = cq_list_.empty() ? nullptr : cq_list_[0];
+    if (!shared_cq_) {
+        LOG(ERROR) << "EfaContext::buildSharedEndpoint: no CQ available";
+        return ERR_CONTEXT;
+    }
+    max_wr_depth_ = static_cast<int>(max_wr);
+
+    int ret = fi_endpoint(domain_, fi_info_, &shared_ep_, nullptr);
+    if (ret) {
+        LOG(ERROR) << "fi_endpoint failed: " << fi_strerror(-ret);
+        shared_ep_ = nullptr;
+        return ERR_ENDPOINT;
+    }
+
+    ret = fi_ep_bind(shared_ep_, &av_->fid, 0);
+    if (ret) {
+        LOG(ERROR) << "fi_ep_bind(av) failed: " << fi_strerror(-ret);
+        fi_close(&shared_ep_->fid);
+        shared_ep_ = nullptr;
+        return ERR_ENDPOINT;
+    }
+
+    ret = fi_ep_bind(shared_ep_, &shared_cq_->cq->fid, FI_TRANSMIT);
+    if (ret) {
+        LOG(ERROR) << "fi_ep_bind(tx_cq) failed: " << fi_strerror(-ret);
+        fi_close(&shared_ep_->fid);
+        shared_ep_ = nullptr;
+        return ERR_ENDPOINT;
+    }
+
+    ret = fi_ep_bind(shared_ep_, &shared_cq_->cq->fid, FI_RECV);
+    if (ret) {
+        LOG(ERROR) << "fi_ep_bind(rx_cq) failed: " << fi_strerror(-ret);
+        fi_close(&shared_ep_->fid);
+        shared_ep_ = nullptr;
+        return ERR_ENDPOINT;
+    }
+
+    ret = fi_enable(shared_ep_);
+    if (ret) {
+        LOG(ERROR) << "fi_enable failed: " << fi_strerror(-ret);
+        fi_close(&shared_ep_->fid);
+        shared_ep_ = nullptr;
+        return ERR_ENDPOINT;
+    }
+
+    // Cache our own libfabric address for handshake advertisement.
+    size_t addr_len = 64;
+    local_ep_addr_.assign(addr_len, 0);
+    ret = fi_getname(&shared_ep_->fid, local_ep_addr_.data(), &addr_len);
+    if (ret) {
+        LOG(ERROR) << "fi_getname failed: " << fi_strerror(-ret);
+        fi_close(&shared_ep_->fid);
+        shared_ep_ = nullptr;
+        local_ep_addr_.clear();
+        return ERR_ENDPOINT;
+    }
+    local_ep_addr_.resize(addr_len);
+    return 0;
+}
+
 int EfaContext::deconstruct() {
-    // Destroy all endpoints before closing domain/fabric/AV.
-    // Endpoints hold fi_ep handles that reference the domain, so they must
-    // be closed first.
-    endpoint_store_.reset();
+    // Teardown order matters for the EFA provider: the shared endpoint must
+    // be closed before the AV it is bound to.  We also cannot call
+    // fi_av_remove() after fi_close(ep) — the provider faults — so we just
+    // clear the peer map (dropping shared_ptrs) and let fi_av_close() below
+    // invalidate every AV slot in one shot.
+    if (shared_ep_) {
+        fi_close(&shared_ep_->fid);
+        shared_ep_ = nullptr;
+    }
+
+    {
+        RWSpinlock::WriteGuard guard(peer_map_lock_);
+        for (auto& entry : peer_map_) {
+            if (entry.second) entry.second->markDetachedForTeardown();
+        }
+        peer_map_.clear();
+    }
 
     {
         RWSpinlock::WriteGuard guard(mr_lock_);
@@ -299,6 +314,7 @@ int EfaContext::deconstruct() {
         }
     }
     cq_list_.clear();
+    shared_cq_.reset();
 
     if (av_) {
         fi_close(&av_->fid);
@@ -331,6 +347,7 @@ int EfaContext::deconstruct() {
 int EfaContext::registerMemoryRegionInternal(void* addr, size_t length,
                                              int access,
                                              EfaMemoryRegionMeta& mrMeta) {
+    (void)access;
     if (length > (size_t)globalConfig().max_mr_size) {
         LOG(ERROR) << "Buffer length " << length
                    << " exceeds device max_mr_size "
@@ -417,24 +434,22 @@ int EfaContext::registerMemoryRegion(void* addr, size_t length, int access) {
 int EfaContext::unregisterMemoryRegion(void* addr) {
     RWSpinlock::WriteGuard guard(mr_lock_);
     auto it = mr_map_.find((uint64_t)addr);
-    if (it != mr_map_.end()) {
-        if (it->second.mr) {
-            int ret = fi_close(&it->second.mr->fid);
-            if (ret) {
-                LOG(ERROR) << "Failed to unregister memory " << addr << ": "
-                           << fi_strerror(-ret);
-                return ERR_CONTEXT;
-            }
-        }
-        mr_map_.erase(it);
+    if (it == mr_map_.end()) {
+        return 0;
     }
+    if (it->second.mr) {
+        fi_close(&it->second.mr->fid);
+    }
+    mr_map_.erase(it);
     return 0;
 }
 
 int EfaContext::preTouchMemory(void* addr, size_t length) {
-    volatile char* ptr = (volatile char*)addr;
-    for (size_t i = 0; i < length; i += 4096) {
-        ptr[i] = ptr[i];
+    volatile char* p = static_cast<char*>(addr);
+    const long sc = sysconf(_SC_PAGESIZE);
+    const size_t page_size = sc > 0 ? static_cast<size_t>(sc) : 4096;
+    for (size_t off = 0; off < length; off += page_size) {
+        p[off] = p[off];
     }
     return 0;
 }
@@ -445,7 +460,7 @@ uint64_t EfaContext::rkey(void* addr) {
     if (it != mr_map_.begin()) {
         --it;
         if ((uint64_t)addr < it->first + it->second.length && it->second.mr) {
-            return it->second.key;
+            return fi_mr_key(it->second.mr);
         }
     }
     return 0;
@@ -477,58 +492,79 @@ void* EfaContext::mrDesc(void* addr) {
 
 std::shared_ptr<EfaEndPoint> EfaContext::endpoint(
     const std::string& peer_nic_path) {
-    if (!endpoint_store_) return nullptr;
+    // Key the peer map by the full "host:port@nic" path, not the
+    // port-stripped form.  Under sglang DP>1 every DP worker on a peer
+    // host is a distinct process with its own Mooncake TransferEngine
+    // and its own P2PHANDSHAKE RPC port.  They share the same host+NIC
+    // but map to different EFA QPNs / memory regions.  If we normalize
+    // the port away, every DP worker on that host collapses onto the
+    // same EfaEndPoint slot: each arriving handshake looks like a
+    // "peer reconnected with new address" to the previous holder,
+    // triggering fi_av_remove + fi_av_insert (and an AH warm-up) on
+    // every KV transfer.  At high DP this devolves into permanent
+    // thrashing and eventually "Remote MR invalid" when an in-flight
+    // fi_write runs against a stale AV slot.
+    //
+    // The port is stable for the lifetime of a sglang worker process
+    // (Mooncake only calls initialize() once), so keying by full path
+    // costs nothing in steady state.  A genuine peer restart (process
+    // re-launch → new RPC port) creates a new map entry and leaks the
+    // old EfaEndPoint — that's at most a few bytes per ex-worker and
+    // far cheaper than the churn the old normalization caused.
+    const std::string& key = peer_nic_path;
 
-    // Use normalized key (strip port) so the same physical peer reuses its
-    // endpoint across reconnections.  Each P2PHANDSHAKE run picks a random
-    // port, producing a different peer_nic_path for the same peer host+NIC.
-    std::string key = normalizeNicPath(peer_nic_path);
-
-    // Fast path: endpoint already exists for this physical peer.
-    // Update peer_nic_path in case the port changed (new initiator run).
-    // setPeerNicPath disconnects old connection (fi_av_remove) if needed.
-    auto ep = endpoint_store_->get(key);
-    if (ep) {
-        ep->setPeerNicPath(peer_nic_path);
-        return ep;
-    }
-
-    // Slow path: create new endpoint, then atomically insert (or get existing
-    // if another thread raced us).  getOrInsert prevents duplicate endpoints
-    // and duplicate AV entries for the same peer.
-    auto new_endpoint = std::make_shared<EfaEndPoint>(*this);
-    auto cq = nextCq();
-    if (cq) {
-        int ret = new_endpoint->construct(cq->cq, &cq->outstanding, 1, 4,
-                                          globalConfig().max_wr, 64);
-        if (ret != 0) {
-            LOG(ERROR) << "Failed to construct EFA endpoint";
-            return nullptr;
+    {
+        RWSpinlock::ReadGuard guard(peer_map_lock_);
+        auto it = peer_map_.find(key);
+        if (it != peer_map_.end()) {
+            return it->second;
         }
     }
-    // Still set the full peer_nic_path (with port) for handshake routing
-    new_endpoint->setPeerNicPath(peer_nic_path);
-    ep = endpoint_store_->getOrInsert(key, new_endpoint);
-    // If another thread won the race, new_endpoint is discarded (RAII cleanup)
-    return ep;
+
+    auto new_ep = std::make_shared<EfaEndPoint>(*this);
+    new_ep->setPeerNicPath(peer_nic_path);
+
+    RWSpinlock::WriteGuard guard(peer_map_lock_);
+    auto it = peer_map_.find(key);
+    if (it != peer_map_.end()) {
+        return it->second;
+    }
+    peer_map_[key] = new_ep;
+    return new_ep;
+}
+
+std::shared_ptr<EfaEndPoint> EfaContext::peekEndpoint(
+    const std::string& peer_nic_path) {
+    RWSpinlock::ReadGuard guard(peer_map_lock_);
+    auto it = peer_map_.find(peer_nic_path);
+    if (it == peer_map_.end()) return nullptr;
+    return it->second;
 }
 
 int EfaContext::deleteEndpoint(const std::string& peer_nic_path) {
-    if (endpoint_store_) {
-        endpoint_store_->remove(peer_nic_path);
+    std::shared_ptr<EfaEndPoint> ep;
+    {
+        RWSpinlock::WriteGuard guard(peer_map_lock_);
+        auto it = peer_map_.find(peer_nic_path);
+        if (it == peer_map_.end()) return 0;
+        ep = it->second;
+        peer_map_.erase(it);
     }
+    if (ep) ep->disconnect();  // runs fi_av_remove
     return 0;
 }
 
 int EfaContext::disconnectAllEndpoints() {
-    if (endpoint_store_) {
-        return endpoint_store_->disconnectAll();
+    RWSpinlock::WriteGuard guard(peer_map_lock_);
+    for (auto& entry : peer_map_) {
+        if (entry.second) entry.second->disconnect();
     }
     return 0;
 }
 
 size_t EfaContext::getTotalQPNumber() const {
-    return endpoint_store_ ? endpoint_store_->size() : 0;
+    RWSpinlock::ReadGuard guard(peer_map_lock_);
+    return peer_map_.size();
 }
 
 std::string EfaContext::nicPath() const {
@@ -536,11 +572,7 @@ std::string EfaContext::nicPath() const {
 }
 
 std::string EfaContext::localAddr() const {
-    // Return a hex string representation of the local address info
-    if (!fi_info_ || !fi_info_->src_addr) {
-        return "";
-    }
-
+    if (!fi_info_ || !fi_info_->src_addr) return "";
     std::ostringstream oss;
     const uint8_t* addr = static_cast<const uint8_t*>(fi_info_->src_addr);
     for (size_t i = 0; i < fi_info_->src_addrlen; ++i) {
@@ -549,27 +581,150 @@ std::string EfaContext::localAddr() const {
     return oss.str();
 }
 
+std::string EfaContext::localEpAddr() const {
+    static constexpr char kHex[] = "0123456789abcdef";
+    std::string out;
+    out.resize(local_ep_addr_.size() * 2);
+    for (size_t i = 0; i < local_ep_addr_.size(); ++i) {
+        out[2 * i] = kHex[(local_ep_addr_[i] >> 4) & 0xF];
+        out[2 * i + 1] = kHex[local_ep_addr_[i] & 0xF];
+    }
+    return out;
+}
+
+// Decode one hex nibble, -1 on invalid input.
+static inline int hexNibble(char c) {
+    if (c >= '0' && c <= '9') return c - '0';
+    if (c >= 'a' && c <= 'f') return c - 'a' + 10;
+    if (c >= 'A' && c <= 'F') return c - 'A' + 10;
+    return -1;
+}
+
+int EfaContext::insertPeerAddr(const std::string& peer_hex_addr,
+                               fi_addr_t& out) {
+    if (peer_hex_addr.empty() || (peer_hex_addr.size() % 2) != 0) {
+        LOG(ERROR) << "insertPeerAddr: invalid hex length "
+                   << peer_hex_addr.size();
+        return ERR_INVALID_ARGUMENT;
+    }
+    const size_t n = peer_hex_addr.size() / 2;
+    std::vector<uint8_t> bin(n);
+    for (size_t i = 0; i < n; ++i) {
+        int hi = hexNibble(peer_hex_addr[2 * i]);
+        int lo = hexNibble(peer_hex_addr[2 * i + 1]);
+        if (hi < 0 || lo < 0) {
+            LOG(ERROR) << "insertPeerAddr: non-hex char at offset " << (2 * i);
+            return ERR_INVALID_ARGUMENT;
+        }
+        bin[i] = static_cast<uint8_t>((hi << 4) | lo);
+    }
+    return insertPeerAddrBytes(bin.data(), bin.size(), out);
+}
+
+int EfaContext::insertPeerAddrBytes(const uint8_t* addr, size_t len,
+                                    fi_addr_t& out) {
+    if (!addr || len == 0) return ERR_INVALID_ARGUMENT;
+    int ret = fi_av_insert(av_, addr, 1, &out, 0, nullptr);
+    if (ret != 1) {
+        LOG(ERROR) << "fi_av_insert failed: " << fi_strerror(-ret);
+        return ERR_ENDPOINT;
+    }
+    return 0;
+}
+
+void EfaContext::removePeerAddr(fi_addr_t fi_addr) {
+    if (fi_addr == FI_ADDR_UNSPEC) return;
+    int ret = fi_av_remove(av_, &fi_addr, 1, 0);
+    if (ret) {
+        LOG(WARNING) << "fi_av_remove failed: " << fi_strerror(-ret);
+    }
+}
+
+bool EfaContext::tryLoopbackCopy(Transport::Slice* slice) {
+    // Only intra-process self-transfers qualify.  We compare the SERVER name
+    // (host:rpc_port), NOT the full nic path: local_server_name embeds this
+    // process's unique RPC port, so a server-name match guarantees the peer
+    // is *this very process* — meaning both source_addr and dest_addr are
+    // pointers we can dereference directly.  We deliberately ignore the
+    // device suffix: with multiple NICs the source slice is routed by its
+    // source buffer's device while peer_nic_path names the destination
+    // buffer's device, so the two device names often differ even for a pure
+    // self-loopback.  Same-host *cross-process* peers carry a different RPC
+    // port and never match, so we never memcpy across address spaces.
+    if (getServerNameFromNicPath(slice->peer_nic_path) !=
+        engine_.local_server_name()) {
+        return false;
+    }
+
+    // Direction depends on opcode, mirroring fi_read / fi_write below:
+    //   WRITE: fi_write(buf=source_addr -> addr=dest_addr)  data src->dst
+    //   READ : fi_read (buf=source_addr <- addr=dest_addr)  data dst->src
+    // source_addr and rdma.dest_addr are two distinct local buffers here, so
+    // the copy is NOT symmetric — we must honor the opcode's direction.
+    void* local_buf = slice->source_addr;
+    void* remote_buf = reinterpret_cast<void*>(slice->rdma.dest_addr);
+    void* dst;
+    void* src;
+    if (slice->opcode == Transport::TransferRequest::READ) {
+        dst = local_buf;   // read INTO local
+        src = remote_buf;  // FROM remote (== local) buffer
+    } else {
+        dst = remote_buf;  // write INTO remote (== local) buffer
+        src = local_buf;   // FROM local
+    }
+    size_t len = slice->length;
+    // GPU-aware copy.  Guard on the SAME backends that register GPU memory
+    // as FI_HMEM (see the FI_MR_HMEM hint and the registration path: only
+    // USE_CUDA / USE_HIP tag MRs with a device iface).  Every other build —
+    // including non-EFA GPU backends such as MUSA/MLU/MACA, which never run
+    // on AWS EFA hardware — registers loopback buffers as host memory, so a
+    // plain memcpy is both correct and the only portable option (those
+    // backends do not expose the cuda* symbols).  *MemcpyDefault picks
+    // H2H/H2D/D2H/D2D from the pointer attributes.
+#if defined(USE_CUDA)
+    auto rc = cudaMemcpy(dst, src, len, cudaMemcpyDefault);
+    if (rc != cudaSuccess) {
+        LOG(ERROR) << "EFA loopback cudaMemcpy failed: "
+                   << cudaGetErrorString(rc) << " (dst=" << dst
+                   << ", src=" << src << ", len=" << len << ")";
+        slice->markFailed();
+        return true;
+    }
+#elif defined(USE_HIP)
+    auto rc = hipMemcpy(dst, src, len, hipMemcpyDefault);
+    if (rc != hipSuccess) {
+        LOG(ERROR) << "EFA loopback hipMemcpy failed: " << hipGetErrorString(rc)
+                   << " (dst=" << dst << ", src=" << src << ", len=" << len
+                   << ")";
+        slice->markFailed();
+        return true;
+    }
+#else
+    memcpy(dst, src, len);
+#endif
+    slice->markSuccess();
+    return true;
+}
+
 int EfaContext::submitPostSend(
     const std::vector<Transport::Slice*>& slice_list) {
-    // Route slices to appropriate endpoints for sending
-    // Group slices by peer NIC path
+    // Route slices to appropriate peer handles.  Group by peer NIC path.
     std::unordered_map<std::string, std::vector<Transport::Slice*>>
         slices_by_peer;
-    std::vector<Transport::Slice*> failed_slices;
 
     for (auto* slice : slice_list) {
         if (!slice) continue;
 
-        // Fast path: peer info already resolved by submitTransferTask's
-        // striping path (dest_rkey and peer_nic_path pre-set on slice).
-        // This eliminates per-slice metadata lookup, selectDevice(), and
-        // string construction — the main bottleneck for multi-NIC striping.
+        // Fast path: peer info already filled in by the caller
+        // (dest_rkey and peer_nic_path set on the slice before dispatch).
         if (!slice->peer_nic_path.empty()) {
+            // Self-loopback: satisfy locally, skip EFA entirely.
+            if (tryLoopbackCopy(slice)) continue;
             slices_by_peer[slice->peer_nic_path].push_back(slice);
             continue;
         }
 
-        // Slow path: resolve peer info per-slice (non-striped transfers)
+        // Slow path: resolve peer info per-slice.
         auto peer_segment_desc =
             engine_.meta()->getSegmentDescByID(slice->target_id);
         if (!peer_segment_desc) {
@@ -579,7 +734,6 @@ int EfaContext::submitPostSend(
             continue;
         }
 
-        // Find the buffer and device for this destination address
         int buffer_id = -1, device_id = -1;
         if (EfaTransport::selectDevice(peer_segment_desc.get(),
                                        slice->rdma.dest_addr, slice->length,
@@ -590,46 +744,235 @@ int EfaContext::submitPostSend(
             continue;
         }
 
-        // Set the remote key from the peer's registered memory region
         slice->rdma.dest_rkey =
             peer_segment_desc->buffers[buffer_id].rkey[device_id];
 
-        // Construct peer NIC path: "server_name@device_name"
-        std::string peer_nic_path = peer_segment_desc->name + "@" +
+        std::string peer_nic_path = peer_segment_desc->nicPathServerName() +
+                                    "@" +
                                     peer_segment_desc->devices[device_id].name;
         slice->peer_nic_path = peer_nic_path;
-
+        // Self-loopback: satisfy locally, skip EFA entirely.
+        if (tryLoopbackCopy(slice)) continue;
         slices_by_peer[peer_nic_path].push_back(slice);
     }
 
-    // Now send to each peer endpoint
     for (auto& entry : slices_by_peer) {
         const std::string& peer_nic_path = entry.first;
         auto& peer_slices = entry.second;
 
-        // Get or create endpoint for this peer
         auto ep = endpoint(peer_nic_path);
         if (!ep) {
-            LOG(ERROR) << "Cannot create endpoint for peer " << peer_nic_path;
-            for (auto* slice : peer_slices) {
-                slice->markFailed();
-            }
+            LOG(ERROR) << "Cannot create peer handle for " << peer_nic_path;
+            for (auto* slice : peer_slices) slice->markFailed();
             continue;
         }
 
-        // Submit to endpoint
         std::vector<Transport::Slice*> failed_slice_list;
-        ep->submitPostSend(peer_slices, failed_slice_list);
+        int rc = ep->submitPostSend(peer_slices, failed_slice_list);
+        for (auto* slice : failed_slice_list) slice->markFailed();
 
-        // Handle any slices that failed to post
-        for (auto* slice : failed_slice_list) {
-            slice->markFailed();
+        // Drop peer handle if it is no longer connected after submit,
+        // freeing its AV entry for reuse.  Under the shared-endpoint model
+        // this is cheap (no fid_ep to destroy).
+        if (rc != 0 && !ep->connected()) {
+            deleteEndpoint(peer_nic_path);
         }
     }
 
     return 0;
 }
 
+int EfaContext::submitSlicesOnPeer(
+    fi_addr_t peer_fi_addr, std::vector<Transport::Slice*>& slice_list,
+    std::vector<Transport::Slice*>& failed_slice_list) {
+    // Batched submission against the shared endpoint.  Mirrors the previous
+    // per-endpoint submit path but uses context-level wr_depth / post_lock.
+    //
+    // 1. Reserve N WR+CQ slots in bulk (single CAS each)
+    // 2. Prepare MR descriptors and op contexts outside the lock
+    // 3. Hold post_lock_ once for the entire batch of fi_write calls
+    const int kMaxBackoffYields = 100000;
+    const int cq_limit = static_cast<int>(globalConfig().max_cqe);
+    std::atomic<int>* cq_outstanding =
+        shared_cq_ ? &shared_cq_->outstanding : nullptr;
+
+    struct BatchEntry {
+        Transport::Slice* slice;
+        void* local_desc;
+        EfaOpContext* op_ctx;
+    };
+
+    // Consume slice_list via a moving index instead of erase-from-front,
+    // which was O(N^2) on large batches.  retry_slices accumulate across
+    // passes and are applied by rewinding the cursor.
+    size_t cursor = 0;
+    std::vector<Transport::Slice*> retry_slices;
+    while (cursor < slice_list.size() || !retry_slices.empty()) {
+        if (!retry_slices.empty()) {
+            // Splice retry slices back in at the current cursor so the next
+            // pass picks them up.  O(retry_slices.size()) per retry wave,
+            // which is bounded by valid_count of the last batch.
+            slice_list.insert(slice_list.begin() + cursor, retry_slices.begin(),
+                              retry_slices.end());
+            retry_slices.clear();
+            std::this_thread::yield();
+        }
+
+        const size_t remaining = slice_list.size() - cursor;
+        int batch_count = 0;
+        int backoff = 0;
+        bool timed_out = false;
+        while (batch_count == 0) {
+            int cur_wr = wr_depth_.load(std::memory_order_relaxed);
+            int wr_avail = max_wr_depth_ - cur_wr;
+            if (wr_avail <= 0) {
+                if (++backoff > kMaxBackoffYields) {
+                    timed_out = true;
+                    break;
+                }
+                std::this_thread::yield();
+                continue;
+            }
+            int want = std::min(wr_avail, (int)remaining);
+            if (cq_outstanding) {
+                int cur_cq = cq_outstanding->load(std::memory_order_relaxed);
+                int cq_avail = cq_limit - cur_cq;
+                if (cq_avail <= 0) {
+                    if (++backoff > kMaxBackoffYields) {
+                        timed_out = true;
+                        break;
+                    }
+                    std::this_thread::yield();
+                    continue;
+                }
+                want = std::min(want, cq_avail);
+                if (!wr_depth_.compare_exchange_weak(
+                        cur_wr, cur_wr + want, std::memory_order_acq_rel,
+                        std::memory_order_relaxed)) {
+                    continue;
+                }
+                cur_cq = cq_outstanding->load(std::memory_order_relaxed);
+                cq_avail = cq_limit - cur_cq;
+                if (cq_avail < want) {
+                    wr_depth_.fetch_sub(want, std::memory_order_acq_rel);
+                    continue;
+                }
+                if (!cq_outstanding->compare_exchange_weak(
+                        cur_cq, cur_cq + want, std::memory_order_acq_rel,
+                        std::memory_order_relaxed)) {
+                    wr_depth_.fetch_sub(want, std::memory_order_acq_rel);
+                    continue;
+                }
+            } else {
+                if (!wr_depth_.compare_exchange_weak(
+                        cur_wr, cur_wr + want, std::memory_order_acq_rel,
+                        std::memory_order_relaxed)) {
+                    continue;
+                }
+            }
+            batch_count = want;
+        }
+
+        if (timed_out) {
+            LOG(WARNING) << "EFA submitSlicesOnPeer: timed out waiting for CQ"
+                         << " drain (wr_depth="
+                         << wr_depth_.load(std::memory_order_relaxed)
+                         << ", max=" << max_wr_depth_ << ", cq_outstanding="
+                         << (cq_outstanding ? cq_outstanding->load(
+                                                  std::memory_order_relaxed)
+                                            : -1)
+                         << ", max_cqe=" << cq_limit << ")";
+            for (size_t i = cursor; i < slice_list.size(); ++i) {
+                failed_slice_list.push_back(slice_list[i]);
+            }
+            slice_list.clear();
+            return 0;
+        }
+
+        std::vector<BatchEntry> batch(batch_count);
+        int valid_count = 0;
+
+        for (int i = 0; i < batch_count; i++) {
+            Transport::Slice* slice = slice_list[cursor + i];
+            void* local_desc = mrDesc(slice->source_addr);
+            if (!local_desc) {
+                LOG(ERROR) << "No MR descriptor found for address "
+                           << slice->source_addr;
+                failed_slice_list.push_back(slice);
+                continue;
+            }
+            EfaOpContext* op_ctx = new EfaOpContext();
+            memset(op_ctx, 0, sizeof(EfaOpContext));
+            op_ctx->slice = slice;
+            op_ctx->wr_depth = &wr_depth_;
+            batch[valid_count++] = {slice, local_desc, op_ctx};
+        }
+
+        int mr_failures = batch_count - valid_count;
+        if (mr_failures > 0) {
+            wr_depth_.fetch_sub(mr_failures, std::memory_order_acq_rel);
+            if (cq_outstanding)
+                cq_outstanding->fetch_sub(mr_failures,
+                                          std::memory_order_acq_rel);
+        }
+
+        if (valid_count > 0) {
+            while (post_lock_.test_and_set(std::memory_order_acquire)) {
+            }
+            for (int i = 0; i < valid_count; i++) {
+                auto& entry = batch[i];
+                ssize_t ret;
+                if (entry.slice->opcode == Transport::TransferRequest::READ) {
+                    ret = fi_read(shared_ep_, (void*)entry.slice->source_addr,
+                                  entry.slice->length, entry.local_desc,
+                                  peer_fi_addr, entry.slice->rdma.dest_addr,
+                                  entry.slice->rdma.dest_rkey,
+                                  &entry.op_ctx->fi_ctx);
+                } else {
+                    ret = fi_write(shared_ep_, (void*)entry.slice->source_addr,
+                                   entry.slice->length, entry.local_desc,
+                                   peer_fi_addr, entry.slice->rdma.dest_addr,
+                                   entry.slice->rdma.dest_rkey,
+                                   &entry.op_ctx->fi_ctx);
+                }
+                if (ret == 0) {
+                    entry.slice->status = Transport::Slice::PENDING;
+                } else if (ret == -FI_EAGAIN) {
+                    delete entry.op_ctx;
+                    int not_posted = valid_count - i;
+                    wr_depth_.fetch_sub(not_posted, std::memory_order_acq_rel);
+                    if (cq_outstanding)
+                        cq_outstanding->fetch_sub(not_posted,
+                                                  std::memory_order_acq_rel);
+                    for (int j = i; j < valid_count; j++) {
+                        if (j > i) delete batch[j].op_ctx;
+                        retry_slices.push_back(batch[j].slice);
+                    }
+                    break;
+                } else {
+                    LOG(ERROR)
+                        << "fi_read/fi_write failed: " << fi_strerror(-ret)
+                        << " (source=" << entry.slice->source_addr
+                        << ", len=" << entry.slice->length
+                        << ", dest=" << (void*)entry.slice->rdma.dest_addr
+                        << ", rkey=" << entry.slice->rdma.dest_rkey << ")";
+                    delete entry.op_ctx;
+                    wr_depth_.fetch_sub(1, std::memory_order_acq_rel);
+                    if (cq_outstanding)
+                        cq_outstanding->fetch_sub(1, std::memory_order_acq_rel);
+                    failed_slice_list.push_back(entry.slice);
+                }
+            }
+            post_lock_.clear(std::memory_order_release);
+        }
+
+        cursor += batch_count;
+    }
+
+    slice_list.clear();
+    return 0;
+}
+
 int EfaContext::pollCq(int max_entries, int cq_index) {
     if (cq_index < 0 || (size_t)cq_index >= cq_list_.size()) {
         return 0;
@@ -638,15 +981,13 @@ int EfaContext::pollCq(int max_entries, int cq_index) {
     struct fid_cq* cq = cq_list_[cq_index]->cq;
     if (!cq) return 0;
 
-    // Use fi_cq_data format for completions
     struct fi_cq_data_entry entries[64];
     int to_poll = std::min(max_entries, 64);
 
     ssize_t ret = fi_cq_read(cq, entries, to_poll);
 
     if (ret > 0) {
-        // Process completions outside the lock (markSuccess / delete are safe)
-        std::unordered_map<volatile int*, int> wr_depth_set;
+        std::unordered_map<std::atomic<int>*, int> wr_depth_set;
         for (ssize_t i = 0; i < ret; i++) {
             EfaOpContext* op_ctx =
                 reinterpret_cast<EfaOpContext*>(entries[i].op_context);
@@ -659,18 +1000,17 @@ int EfaContext::pollCq(int max_entries, int cq_index) {
             }
         }
         for (auto& entry : wr_depth_set) {
-            __sync_fetch_and_sub(entry.first, entry.second);
+            entry.first->fetch_sub(entry.second, std::memory_order_acq_rel);
         }
-        __sync_fetch_and_sub(&cq_list_[cq_index]->outstanding,
-                             static_cast<int>(ret));
+        cq_list_[cq_index]->outstanding.fetch_sub(static_cast<int>(ret),
+                                                  std::memory_order_acq_rel);
         return static_cast<int>(ret);
     } else if (ret == -FI_EAGAIN) {
         return 0;
     } else if (ret < 0) {
-        // CQ error - drain all queued error entries under the domain lock
         int err_count = 0;
         struct fi_cq_err_entry err_entry;
-        std::unordered_map<volatile int*, int> wr_depth_set;
+        std::unordered_map<std::atomic<int>*, int> wr_depth_set;
 
         while ((ret = fi_cq_readerr(cq, &err_entry, 0)) > 0) {
             EfaOpContext* op_ctx =
@@ -690,10 +1030,11 @@ int EfaContext::pollCq(int max_entries, int cq_index) {
         }
 
         for (auto& entry : wr_depth_set) {
-            __sync_fetch_and_sub(entry.first, entry.second);
+            entry.first->fetch_sub(entry.second, std::memory_order_acq_rel);
         }
         if (err_count > 0) {
-            __sync_fetch_and_sub(&cq_list_[cq_index]->outstanding, err_count);
+            cq_list_[cq_index]->outstanding.fetch_sub(
+                err_count, std::memory_order_acq_rel);
         }
         return err_count;
     }
diff --git a/mooncake-transfer-engine/src/transport/efa_transport/efa_endpoint.cpp b/mooncake-transfer-engine/src/transport/efa_transport/efa_endpoint.cpp
index 67a94cb6..7357a22f 100644
--- a/mooncake-transfer-engine/src/transport/efa_transport/efa_endpoint.cpp
+++ b/mooncake-transfer-engine/src/transport/efa_transport/efa_endpoint.cpp
@@ -16,122 +16,22 @@
 
 #include <glog/logging.h>
 
-#include <cassert>
-#include <cstddef>
-#include <cstring>
-#include <iomanip>
+#include <cstdlib>
 #include <sstream>
-#include <thread>
 
 #include "config.h"
 
 namespace mooncake {
 
 EfaEndPoint::EfaEndPoint(EfaContext& context)
-    : context_(context),
-      status_(INITIALIZING),
-      ep_(nullptr),
-      tx_cq_(nullptr),
-      rx_cq_(nullptr),
-      peer_fi_addr_(FI_ADDR_UNSPEC),
-      local_addr_len_(0),
-      wr_depth_(0),
-      max_wr_depth_(0),
-      cq_outstanding_(nullptr),
-      active_(true),
-      inactive_time_(0),
-      last_used_time_(getCurrentTimeInNano()) {}
+    : context_(context), status_(INITIALIZING), peer_fi_addr_(FI_ADDR_UNSPEC) {}
 
-EfaEndPoint::~EfaEndPoint() {
-    if (ep_) deconstruct();
-}
-
-int EfaEndPoint::construct(struct fid_cq* cq, volatile int* cq_outstanding,
-                           size_t num_qp_list, size_t max_sge, size_t max_wr,
-                           size_t max_inline) {
-    if (status_.load(std::memory_order_relaxed) != INITIALIZING) {
-        LOG(ERROR) << "EFA Endpoint has already been constructed";
-        return ERR_ENDPOINT;
-    }
-
-    tx_cq_ = cq;
-    rx_cq_ = cq;  // Use same CQ for TX and RX
-    max_wr_depth_ = max_wr;
-    cq_outstanding_ = cq_outstanding;
-
-    // Create endpoint
-    int ret = fi_endpoint(context_.domain(), context_.info(), &ep_, nullptr);
-    if (ret) {
-        LOG(ERROR) << "fi_endpoint failed: " << fi_strerror(-ret);
-        return ERR_ENDPOINT;
-    }
-
-    // Bind endpoint to AV
-    ret = fi_ep_bind(ep_, &context_.av()->fid, 0);
-    if (ret) {
-        LOG(ERROR) << "fi_ep_bind (av) failed: " << fi_strerror(-ret);
-        fi_close(&ep_->fid);
-        ep_ = nullptr;
-        return ERR_ENDPOINT;
-    }
-
-    // Bind endpoint to TX CQ
-    ret = fi_ep_bind(ep_, &tx_cq_->fid, FI_TRANSMIT);
-    if (ret) {
-        LOG(ERROR) << "fi_ep_bind (tx_cq) failed: " << fi_strerror(-ret);
-        fi_close(&ep_->fid);
-        ep_ = nullptr;
-        return ERR_ENDPOINT;
-    }
-
-    // Bind endpoint to RX CQ
-    ret = fi_ep_bind(ep_, &rx_cq_->fid, FI_RECV);
-    if (ret) {
-        LOG(ERROR) << "fi_ep_bind (rx_cq) failed: " << fi_strerror(-ret);
-        fi_close(&ep_->fid);
-        ep_ = nullptr;
-        return ERR_ENDPOINT;
-    }
-
-    // Enable endpoint
-    ret = fi_enable(ep_);
-    if (ret) {
-        LOG(ERROR) << "fi_enable failed: " << fi_strerror(-ret);
-        fi_close(&ep_->fid);
-        ep_ = nullptr;
-        return ERR_ENDPOINT;
-    }
-
-    // Get local endpoint address
-    local_addr_len_ = 64;  // EFA addresses are typically 32 bytes
-    local_addr_.resize(local_addr_len_);
-    ret = fi_getname(&ep_->fid, local_addr_.data(), &local_addr_len_);
-    if (ret) {
-        LOG(ERROR) << "fi_getname failed: " << fi_strerror(-ret);
-        fi_close(&ep_->fid);
-        ep_ = nullptr;
-        return ERR_ENDPOINT;
-    }
-    local_addr_.resize(local_addr_len_);
-
-    status_.store(UNCONNECTED, std::memory_order_relaxed);
-    return 0;
-}
-
-int EfaEndPoint::deconstruct() {
-    if (ep_) {
-        fi_close(&ep_->fid);
-        ep_ = nullptr;
-    }
-    return 0;
-}
-
-int EfaEndPoint::destroyQP() { return deconstruct(); }
+EfaEndPoint::~EfaEndPoint() { disconnect(); }
 
 void EfaEndPoint::setPeerNicPath(const std::string& peer_nic_path) {
     RWSpinlock::WriteGuard guard(lock_);
     if (peer_nic_path_ == peer_nic_path) return;  // No change
-    if (connected()) {
+    if (status_.load(std::memory_order_relaxed) == CONNECTED) {
         LOG(INFO) << "Peer reconnected with new address, re-establishing: "
                   << peer_nic_path_ << " -> " << peer_nic_path;
         disconnectUnlocked();
@@ -139,62 +39,28 @@ void EfaEndPoint::setPeerNicPath(const std::string& peer_nic_path) {
     peer_nic_path_ = peer_nic_path;
 }
 
-std::string EfaEndPoint::getLocalAddr() const {
-    std::ostringstream oss;
-    for (size_t i = 0; i < local_addr_.size(); ++i) {
-        oss << std::hex << std::setw(2) << std::setfill('0')
-            << (int)local_addr_[i];
-    }
-    return oss.str();
-}
-
-int EfaEndPoint::insertPeerAddr(const std::string& peer_addr) {
-    // Convert hex string to binary address
-    std::vector<uint8_t> addr_bin;
-    addr_bin.reserve(peer_addr.size() / 2);
-
-    for (size_t i = 0; i < peer_addr.size(); i += 2) {
-        std::string byte_str = peer_addr.substr(i, 2);
-        uint8_t byte = (uint8_t)strtol(byte_str.c_str(), nullptr, 16);
-        addr_bin.push_back(byte);
-    }
-
-    // Insert into address vector
-    int ret = fi_av_insert(context_.av(), addr_bin.data(), 1, &peer_fi_addr_, 0,
-                           nullptr);
-    if (ret != 1) {
-        LOG(ERROR) << "fi_av_insert failed: " << fi_strerror(-ret);
-        return ERR_ENDPOINT;
-    }
-
-    return 0;
-}
-
 int EfaEndPoint::setupConnectionsByActive() {
     RWSpinlock::WriteGuard guard(lock_);
-    if (connected()) {
-        LOG(INFO) << "EFA Connection has been established";
-        return 0;
-    }
+    if (status_.load(std::memory_order_relaxed) == CONNECTED) return 0;
 
-    // Loopback mode
+    // Loopback: handshake against ourselves.  Use the binary overload
+    // so we avoid hex-encoding the local address just to have
+    // insertPeerAddr decode it right back.
     if (context_.nicPath() == peer_nic_path_) {
-        // For loopback, insert our own address
-        int ret = insertPeerAddr(getLocalAddr());
-        if (ret != 0) {
-            return ret;
-        }
+        const auto& bytes = context_.localEpAddrBytes();
+        int ret = context_.insertPeerAddrBytes(bytes.data(), bytes.size(),
+                                               peer_fi_addr_);
+        if (ret != 0) return ret;
         status_.store(CONNECTED, std::memory_order_release);
         LOG(INFO) << "EFA loopback connection established: " << toString();
         return 0;
     }
 
-    // Exchange addresses via handshake
+    // Exchange addresses via the transfer-metadata handshake RPC.
     TransferMetadata::HandShakeDesc local_desc, peer_desc;
     local_desc.local_nic_path = context_.nicPath();
     local_desc.peer_nic_path = peer_nic_path_;
-    // Store our EFA endpoint address in efa_addr field (hex encoded)
-    local_desc.efa_addr = getLocalAddr();
+    local_desc.efa_addr = context_.localEpAddr();
 
     auto peer_server_name = getServerNameFromNicPath(peer_nic_path_);
     auto peer_nic_name = getNicNameFromNicPath(peer_nic_path_);
@@ -212,14 +78,11 @@ int EfaEndPoint::setupConnectionsByActive() {
         return ERR_REJECT_HANDSHAKE;
     }
 
-    // Insert peer's address into our AV
-    rc = insertPeerAddr(peer_desc.efa_addr);
-    if (rc != 0) {
-        return rc;
-    }
+    rc = context_.insertPeerAddr(peer_desc.efa_addr, peer_fi_addr_);
+    if (rc != 0) return rc;
+    cached_peer_addr_ = peer_desc.efa_addr;
 
     status_.store(CONNECTED, std::memory_order_release);
-    touchLastUsed();
     VLOG(1) << "EFA connection established: " << toString()
             << " peer_fi_addr=" << peer_fi_addr_;
     return 0;
@@ -228,10 +91,6 @@ int EfaEndPoint::setupConnectionsByActive() {
 int EfaEndPoint::setupConnectionsByPassive(const HandShakeDesc& peer_desc,
                                            HandShakeDesc& local_desc) {
     RWSpinlock::WriteGuard guard(lock_);
-    if (connected()) {
-        LOG(WARNING) << "Re-establish EFA connection: " << toString();
-        disconnectUnlocked();
-    }
 
     if (peer_desc.peer_nic_path != context_.nicPath() ||
         peer_desc.local_nic_path != peer_nic_path_) {
@@ -244,27 +103,50 @@ int EfaEndPoint::setupConnectionsByPassive(const HandShakeDesc& peer_desc,
         return ERR_REJECT_HANDSHAKE;
     }
 
-    // Insert peer's address from handshake
     if (peer_desc.efa_addr.empty()) {
         local_desc.reply_msg = "No EFA address provided";
         LOG(ERROR) << "Peer did not provide EFA address";
         return ERR_REJECT_HANDSHAKE;
     }
 
-    int ret = insertPeerAddr(peer_desc.efa_addr);
+    // Classify this passive handshake so we can emit the right log level
+    // without changing functional behavior: the AV reinsert below must
+    // still happen on every handshake because libfabric's EFA provider
+    // tracks per-peer transport state that depends on a fresh
+    // fi_av_insert (e.g. provider-internal AH activation and RNR state).
+    // Skipping reinsert caused request-level stalls under bilateral
+    // sglang P/D load even when the peer EFA address was unchanged.
+    //
+    // Log semantics:
+    //   * Same cached peer address → benign symmetric handshake under
+    //     bilateral traffic.  Demote to INFO to keep decode logs readable
+    //     without hiding real reconnects.
+    //   * Different cached peer address → genuine reconnect (peer
+    //     restart, port reshuffle that reached disconnect first, etc.).
+    //     Keep at WARNING so it stays visible.
+    if (status_.load(std::memory_order_relaxed) == CONNECTED) {
+        if (!cached_peer_addr_.empty() &&
+            peer_desc.efa_addr == cached_peer_addr_) {
+            VLOG(1) << "EFA passive handshake (same peer addr): " << toString();
+        } else {
+            LOG(WARNING) << "Re-establish EFA connection: " << toString();
+        }
+        disconnectUnlocked();
+    }
+
+    int ret = context_.insertPeerAddr(peer_desc.efa_addr, peer_fi_addr_);
     if (ret != 0) {
         local_desc.reply_msg = "Failed to insert peer address";
         return ret;
     }
+    cached_peer_addr_ = peer_desc.efa_addr;
 
-    // Provide our address to peer (using efa_addr field, not reply_msg)
     local_desc.local_nic_path = context_.nicPath();
     local_desc.peer_nic_path = peer_nic_path_;
-    local_desc.efa_addr = getLocalAddr();
-    // reply_msg should be empty on success
+    local_desc.efa_addr = context_.localEpAddr();
+    // reply_msg empty on success
 
     status_.store(CONNECTED, std::memory_order_release);
-    touchLastUsed();
     VLOG(1) << "EFA connection established (passive): " << toString();
     return 0;
 }
@@ -275,17 +157,18 @@ void EfaEndPoint::disconnect() {
 }
 
 void EfaEndPoint::disconnectUnlocked() {
-    // Remove peer address from AV to free the slot for reuse.
-    // Without this, AV entries leak on every initiator disconnect,
-    // eventually exhausting the AV and causing fi_av_insert/fi_endpoint
-    // failures ("No data available") after 3-5 initiator runs.
     if (peer_fi_addr_ != FI_ADDR_UNSPEC) {
-        int ret = fi_av_remove(context_.av(), &peer_fi_addr_, 1, 0);
-        if (ret) {
-            LOG(WARNING) << "fi_av_remove failed: " << fi_strerror(-ret);
-        }
+        context_.removePeerAddr(peer_fi_addr_);
+        peer_fi_addr_ = FI_ADDR_UNSPEC;
     }
+    cached_peer_addr_.clear();
+    status_.store(UNCONNECTED, std::memory_order_release);
+}
+
+void EfaEndPoint::markDetachedForTeardown() {
+    RWSpinlock::WriteGuard guard(lock_);
     peer_fi_addr_ = FI_ADDR_UNSPEC;
+    cached_peer_addr_.clear();
     status_.store(UNCONNECTED, std::memory_order_release);
 }
 
@@ -293,28 +176,12 @@ const std::string EfaEndPoint::toString() const {
     return "EfaEndPoint[" + context_.nicPath() + " <-> " + peer_nic_path_ + "]";
 }
 
-bool EfaEndPoint::hasOutstandingSlice() const { return wr_depth_ > 0; }
-
-int EfaEndPoint::doSetupConnection(const std::string& peer_addr,
-                                   std::string* reply_msg) {
-    int ret = insertPeerAddr(peer_addr);
-    if (ret != 0) {
-        if (reply_msg) *reply_msg = "Failed to insert peer address into AV";
-        return ret;
-    }
-
-    status_.store(CONNECTED, std::memory_order_release);
-    return 0;
-}
-
 int EfaEndPoint::submitPostSend(
     std::vector<Transport::Slice*>& slice_list,
     std::vector<Transport::Slice*>& failed_slice_list) {
-    if (!connected()) {
-        // Try to establish connection first
+    if (status_.load(std::memory_order_relaxed) != CONNECTED) {
         int ret = setupConnectionsByActive();
         if (ret != 0) {
-            // Move all slices to failed list
             for (auto* slice : slice_list) {
                 failed_slice_list.push_back(slice);
             }
@@ -323,177 +190,33 @@ int EfaEndPoint::submitPostSend(
         }
     }
 
-    // Batched submission — mirrors RDMA's ibv_post_send approach:
-    // 1. Reserve N WR+CQ slots in bulk (single CAS each)
-    // 2. Prepare MR descriptors and op contexts outside the lock
-    // 3. Hold post_lock once for the entire batch of fi_write calls
-    // This reduces per-slice overhead from 2 CAS + 1 spinlock to amortized
-    // 2/N CAS + 1/N spinlock per slice.
-    const int kMaxBackoffYields = 100000;
-    const int cq_limit = static_cast<int>(globalConfig().max_cqe);
+    // Hold a read lock for the entire submit window so setPeerNicPath() /
+    // disconnect() (which take the write lock) cannot fi_av_remove() our
+    // slot while an fi_write() inside submitSlicesOnPeer is still in
+    // flight.  The previous code latched peer_fi_addr_ under the lock and
+    // released it before calling into the context, leaving a race: a
+    // concurrent peer reconnect could remove the AV entry after the latch,
+    // and libfabric would segfault on the stale fi_addr_t.  Read locks
+    // stack, so multiple senders to the same peer still submit in parallel.
+    RWSpinlock::ReadGuard guard(lock_);
 
-    struct BatchEntry {
-        Transport::Slice* slice;
-        void* local_desc;
-        EfaOpContext* op_ctx;
-    };
-
-    while (!slice_list.empty()) {
-        // --- Step 1: Atomically reserve a batch of WR+CQ slots ---
-        int batch_count = 0;
-        int backoff = 0;
-        bool timed_out = false;
-        while (batch_count == 0) {
-            int cur_wr = wr_depth_;
-            int wr_avail = max_wr_depth_ - cur_wr;
-            if (wr_avail <= 0) {
-                if (++backoff > kMaxBackoffYields) {
-                    timed_out = true;
-                    break;
-                }
-                std::this_thread::yield();
-                continue;
-            }
-            int want = std::min(wr_avail, (int)slice_list.size());
-            if (cq_outstanding_) {
-                int cur_cq = *cq_outstanding_;
-                int cq_avail = cq_limit - cur_cq;
-                if (cq_avail <= 0) {
-                    if (++backoff > kMaxBackoffYields) {
-                        timed_out = true;
-                        break;
-                    }
-                    std::this_thread::yield();
-                    continue;
-                }
-                want = std::min(want, cq_avail);
-                if (!__sync_bool_compare_and_swap(&wr_depth_, cur_wr,
-                                                  cur_wr + want)) {
-                    continue;
-                }
-                cur_cq = *cq_outstanding_;
-                cq_avail = cq_limit - cur_cq;
-                if (cq_avail < want) {
-                    __sync_fetch_and_sub(&wr_depth_, want);
-                    continue;
-                }
-                if (!__sync_bool_compare_and_swap(cq_outstanding_, cur_cq,
-                                                  cur_cq + want)) {
-                    __sync_fetch_and_sub(&wr_depth_, want);
-                    continue;
-                }
-            } else {
-                if (!__sync_bool_compare_and_swap(&wr_depth_, cur_wr,
-                                                  cur_wr + want)) {
-                    continue;
-                }
-            }
-            batch_count = want;
-        }
-
-        if (timed_out) {
-            LOG(WARNING) << "EFA submitPostSend: timed out waiting for CQ drain"
-                         << " (wr_depth=" << wr_depth_
-                         << ", max=" << max_wr_depth_ << ", cq_outstanding="
-                         << (cq_outstanding_ ? *cq_outstanding_ : -1)
-                         << ", max_cqe=" << cq_limit << ")";
-            for (auto* slice : slice_list) {
-                failed_slice_list.push_back(slice);
-            }
-            slice_list.clear();
-            return 0;
-        }
-
-        // --- Step 2: Prepare MR descriptors outside the lock ---
-        std::vector<BatchEntry> batch(batch_count);
-        int valid_count = 0;
-
-        for (int i = 0; i < batch_count; i++) {
-            Transport::Slice* slice = slice_list[i];
-            void* local_desc = context_.mrDesc(slice->source_addr);
-            if (!local_desc) {
-                LOG(ERROR) << "No MR descriptor found for address "
-                           << slice->source_addr;
-                failed_slice_list.push_back(slice);
-                continue;
-            }
-            EfaOpContext* op_ctx = new EfaOpContext();
-            memset(op_ctx, 0, sizeof(EfaOpContext));
-            op_ctx->slice = slice;
-            op_ctx->wr_depth = &wr_depth_;
-
-            batch[valid_count++] = {slice, local_desc, op_ctx};
-        }
-
-        int mr_failures = batch_count - valid_count;
-        if (mr_failures > 0) {
-            __sync_fetch_and_sub(&wr_depth_, mr_failures);
-            if (cq_outstanding_)
-                __sync_fetch_and_sub(cq_outstanding_, mr_failures);
-        }
-
-        // --- Step 3: Post entire batch under a single lock hold ---
-        std::vector<Transport::Slice*> retry_slices;
-        if (valid_count > 0) {
-            while (post_lock_.test_and_set(std::memory_order_acquire)) {
-            }
-            for (int i = 0; i < valid_count; i++) {
-                auto& entry = batch[i];
-                ssize_t ret;
-                if (entry.slice->opcode == Transport::TransferRequest::READ) {
-                    ret = fi_read(ep_, (void*)entry.slice->source_addr,
-                                  entry.slice->length, entry.local_desc,
-                                  peer_fi_addr_, entry.slice->rdma.dest_addr,
-                                  entry.slice->rdma.dest_rkey,
-                                  &entry.op_ctx->fi_ctx);
-                } else {
-                    ret = fi_write(ep_, (void*)entry.slice->source_addr,
-                                   entry.slice->length, entry.local_desc,
-                                   peer_fi_addr_, entry.slice->rdma.dest_addr,
-                                   entry.slice->rdma.dest_rkey,
-                                   &entry.op_ctx->fi_ctx);
-                }
-                if (ret == 0) {
-                    entry.slice->status = Transport::Slice::PENDING;
-                } else if (ret == -FI_EAGAIN) {
-                    delete entry.op_ctx;
-                    int not_posted = valid_count - i;
-                    __sync_fetch_and_sub(&wr_depth_, not_posted);
-                    if (cq_outstanding_)
-                        __sync_fetch_and_sub(cq_outstanding_, not_posted);
-                    for (int j = i; j < valid_count; j++) {
-                        if (j > i) delete batch[j].op_ctx;
-                        retry_slices.push_back(batch[j].slice);
-                    }
-                    break;
-                } else {
-                    LOG(ERROR)
-                        << "fi_read/fi_write failed: " << fi_strerror(-ret)
-                        << " (source=" << entry.slice->source_addr
-                        << ", len=" << entry.slice->length
-                        << ", dest=" << (void*)entry.slice->rdma.dest_addr
-                        << ", rkey=" << entry.slice->rdma.dest_rkey << ")";
-                    delete entry.op_ctx;
-                    __sync_fetch_and_sub(&wr_depth_, 1);
-                    if (cq_outstanding_)
-                        __sync_fetch_and_sub(cq_outstanding_, 1);
-                    failed_slice_list.push_back(entry.slice);
-                }
-            }
-            post_lock_.clear(std::memory_order_release);
-            touchLastUsed();
-        }
-
-        // Remove consumed batch from slice_list, prepend any retry slices
-        slice_list.erase(slice_list.begin(), slice_list.begin() + batch_count);
-        if (!retry_slices.empty()) {
-            slice_list.insert(slice_list.begin(), retry_slices.begin(),
-                              retry_slices.end());
-            std::this_thread::yield();
-        }
+    // Re-check status under the lock — a racing disconnect() could have
+    // flipped us to UNCONNECTED between the outer check and acquiring the
+    // lock.  Fail the batch so the caller retries with a fresh setup.
+    if (status_.load(std::memory_order_acquire) != CONNECTED) {
+        for (auto* slice : slice_list) failed_slice_list.push_back(slice);
+        slice_list.clear();
+        return ERR_ENDPOINT;
     }
 
-    return 0;
+    fi_addr_t peer = peer_fi_addr_;
+    if (peer == FI_ADDR_UNSPEC) {
+        for (auto* slice : slice_list) failed_slice_list.push_back(slice);
+        slice_list.clear();
+        return ERR_ENDPOINT;
+    }
+
+    return context_.submitSlicesOnPeer(peer, slice_list, failed_slice_list);
 }
 
 }  // namespace mooncake
diff --git a/mooncake-transfer-engine/src/transport/efa_transport/efa_transport.cpp b/mooncake-transfer-engine/src/transport/efa_transport/efa_transport.cpp
index d3a305c1..617a0ee0 100644
--- a/mooncake-transfer-engine/src/transport/efa_transport/efa_transport.cpp
+++ b/mooncake-transfer-engine/src/transport/efa_transport/efa_transport.cpp
@@ -23,6 +23,7 @@
 #include <cassert>
 #include <chrono>
 #include <cstddef>
+#include <cstdlib>
 #include <fstream>
 #include <future>
 #include <set>
@@ -32,6 +33,7 @@
 
 #include "common.h"
 #include "config.h"
+#include "environ.h"
 #include "memory_location.h"
 #include "topology.h"
 #include "transport/efa_transport/efa_context.h"
@@ -43,7 +45,6 @@ namespace mooncake {
 // roughly 24 million PTEs per NIC, but we use 22M as a conservative default.
 // With 4KB pages: 22M × 4KB ≈ 88GB per NIC.
 // With 2MB hugepages: 22M × 2MB ≈ 44TB per NIC (effectively unlimited).
-// Override via MC_EFA_MAX_PTE_ENTRIES environment variable.
 static constexpr size_t kDefaultMaxPteEntries = 22ULL * 1024 * 1024;  // 22M
 
 // Detect the kernel page size backing the memory at `addr` by reading
@@ -58,8 +59,12 @@ static size_t detectBufferPageSize(void* addr) {
     bool in_range = false;
 
     while (std::getline(smaps, line)) {
-        // VMA header: "start-end perms offset dev inode [pathname]"
-        if (!line.empty() && std::isxdigit(line[0])) {
+        // VMA header: "start-end perms offset dev inode [pathname]".
+        // Cast to unsigned char before std::isxdigit: passing a (possibly
+        // signed) char whose value is > 0x7F is UB, since the argument must be
+        // representable as unsigned char or equal EOF.
+        if (!line.empty() &&
+            std::isxdigit(static_cast<unsigned char>(line[0]))) {
             unsigned long start = 0, end = 0;
             if (sscanf(line.c_str(), "%lx-%lx", &start, &end) == 2) {
                 in_range = (target >= start && target < end);
@@ -75,20 +80,7 @@ static size_t detectBufferPageSize(void* addr) {
     return fallback;
 }
 
-static size_t getMaxPteEntries() {
-    static size_t cached = []() {
-        const char* env = std::getenv("MC_EFA_MAX_PTE_ENTRIES");
-        if (env) {
-            size_t val = std::stoull(env);
-            if (val > 0) {
-                LOG(INFO) << "MC_EFA_MAX_PTE_ENTRIES override: " << val;
-                return val;
-            }
-        }
-        return kDefaultMaxPteEntries;
-    }();
-    return cached;
-}
+static size_t getMaxPteEntries() { return kDefaultMaxPteEntries; }
 
 EfaTransport::EfaTransport() {
     LOG(INFO) << "[EFA] AWS Elastic Fabric Adapter transport initialized";
@@ -105,8 +97,12 @@ void EfaTransport::startWorkerThreads() {
     if (worker_running_) return;
 
     worker_running_ = true;
-    // One poller thread per context for responsive CQ draining under load
+    // MC_EFA_CQ_THREADS caps CQ poller count (default 1). Set 0 to disable cap.
     size_t num_threads = context_list_.size();
+    int cq_cap = Environ::Get().GetEfaCqThreads();
+    if (cq_cap > 0 && static_cast<size_t>(cq_cap) < num_threads) {
+        num_threads = static_cast<size_t>(cq_cap);
+    }
     for (size_t i = 0; i < num_threads; i++) {
         worker_threads_.emplace_back(&EfaTransport::workerThreadFunc, this, i);
     }
@@ -147,12 +143,10 @@ void EfaTransport::workerThreadFunc(int thread_id) {
             }
         }
 
-        // Stale endpoint eviction happens on-demand in
-        // EfaEndpointStore::getOrInsert when the store is at capacity,
-        // not periodically here.  Periodic eviction is unsafe because
-        // target-side endpoints never get touchLastUsed() updates
-        // (the initiator writes remotely via fi_write, bypassing
-        // the target's endpoint code).
+        // Under the shared-endpoint model there is no per-peer QP to evict:
+        // a new peer costs one AV entry (~bytes), not an fi_endpoint slot.
+        // Stale peers are reclaimed when submitPostSend() drops a peer whose
+        // handshake failed.
 
         // If no work was done, yield CPU briefly
         if (!did_work) {
@@ -406,7 +400,10 @@ int EfaTransport::registerLocalMemoryInternal(void* addr, size_t length,
         size_t chunk_len = chunks[ci].second;
         const auto& assigned_nics = nic_assignments[ci];
 
-        bool do_pre_touch = context_list_.size() > 0 &&
+        // preTouchMemory does a CPU-side store to each page, which segfaults
+        // on GPU VRAM (cudaMalloc'd pointers). Restrict it to host memory.
+        bool is_host_mem = resolved_name.rfind("cpu", 0) == 0;
+        bool do_pre_touch = is_host_mem && context_list_.size() > 0 &&
                             std::thread::hardware_concurrency() >= 4 &&
                             chunk_len >= (size_t)4 * 1024 * 1024 * 1024;
         if (do_pre_touch) {
@@ -698,14 +695,32 @@ int EfaTransport::warmupSegment(const std::string& segment_name) {
         peer_paths.emplace_back(segment_name + "@" + dev.name);
     }
 
-    // Warm up every (local_ctx, peer_nic) pair concurrently. Each
-    // EfaContext::endpoint() + setupConnectionsByActive() is idempotent and
-    // takes its own lock, so parallel calls across distinct peer_nic_paths
-    // (and distinct contexts) are safe. We use std::async to get roughly
-    // per-pair parallelism — the critical path is now max(handshake RTT) not
-    // sum(handshake RTT).
     auto t0 = std::chrono::steady_clock::now();
     size_t n_pairs = context_list_.size() * peer_paths.size();
+
+    // Idempotent short-circuit: if every (local_ctx, peer_nic) pair already
+    // has a connected endpoint, skip the whole async dispatch. Matters for
+    // callers that invoke warmupSegment per request loop — without this the
+    // 256-thread fan-out runs every time even though there is no work to do.
+    size_t already_ready = 0;
+    for (auto& ctx : context_list_) {
+        for (const auto& path : peer_paths) {
+            auto ep = ctx->peekEndpoint(path);
+            if (ep && ep->connected()) ++already_ready;
+        }
+    }
+    if (already_ready == n_pairs) {
+        VLOG(1) << "EfaTransport::warmupSegment('" << segment_name << "'): all "
+                << n_pairs << " endpoints already connected, "
+                << "skipping";
+        return 0;
+    }
+
+    // Warm up every (local_ctx, peer_nic) pair concurrently.  Under the
+    // shared-endpoint model each warmup is just a handshake RPC +
+    // fi_av_insert (no fi_endpoint, no fi_enable), so the critical path is
+    // max(handshake RTT), not sum.  We still dispatch with std::async for
+    // concurrency, but total wall time is typically ms-level.
     std::vector<std::future<int>> futs;
     futs.reserve(n_pairs);
     for (auto& ctx : context_list_) {
@@ -720,7 +735,15 @@ int EfaTransport::warmupSegment(const std::string& segment_name) {
                         return -1;
                     }
                     if (ep->connected()) return 0;
-                    return ep->setupConnectionsByActive();
+                    int rc = ep->setupConnectionsByActive();
+                    if (rc != 0) {
+                        // Handshake failed: drop the peer handle so the AV
+                        // slot is freed and the next warmup retry starts
+                        // clean.  Cheap under the shared-endpoint model —
+                        // no fi_endpoint teardown required.
+                        ctx->deleteEndpoint(path);
+                    }
+                    return rc;
                 }));
         }
     }
@@ -756,7 +779,17 @@ Status EfaTransport::submitTransfer(
     size_t task_id = batch_desc.task_list.size();
     batch_desc.task_list.resize(task_id + entries.size());
     std::vector<TransferTask*> task_list;
-    for (auto& task : batch_desc.task_list) task_list.push_back(&task);
+    for (auto& request : entries) {
+        auto& task = batch_desc.task_list[task_id];
+        ++task_id;
+        task.batch_id = batch_id;
+#ifdef USE_ASCEND_HETEROGENEOUS
+        task.request = const_cast<TransferRequest*>(&request);
+#else
+        task.request = &request;
+#endif
+        task_list.push_back(&task);
+    }
     return submitTransferTask(task_list);
 }
 
@@ -767,15 +800,6 @@ Status EfaTransport::submitTransferTask(
     auto local_segment_desc = metadata_->getSegmentDescByID(LOCAL_SEGMENT_ID);
     assert(local_segment_desc.get());
     const int kMaxRetryCount = globalConfig().retry_cnt;
-    // Striping threshold: transfers larger than this stripe across all
-    // active local NICs (one chunk per NIC) instead of creating thousands
-    // of small slices.  This dramatically reduces per-slice overhead
-    // (spinlock acquisition, atomic ops, MR lookup, heap allocation).
-    // Configurable via GlobalConfig (MC_EFA_STRIPING_THRESHOLD env var).
-    // Default 2MB: below this, single-NIC is faster (avoids endpoint overhead
-    // for 16 tiny chunks); above this, multi-NIC striping with pre-resolved
-    // peer info gives up to +54% throughput.
-    const size_t kStripingThreshold = globalConfig().efa_striping_threshold;
 
     for (size_t index = 0; index < task_list.size(); ++index) {
         assert(task_list[index]);
@@ -794,97 +818,9 @@ Status EfaTransport::submitTransferTask(
             request_device_id = -1;
         }
 
-        if (request.length > kStripingThreshold && request_buffer_id >= 0) {
-            // LARGE TRANSFER: stripe across all active local NICs.
-            // Creates one slice per NIC instead of (length / 64KB) slices,
-            // e.g. 240MB transfer: 4-32 slices vs 3840 slices.
-            std::vector<std::pair<int, std::shared_ptr<EfaContext>>>
-                active_devs;
-            for (int d = 0; d < static_cast<int>(context_list_.size()); ++d) {
-                auto& ctx = context_list_[d];
-                if (!ctx || !ctx->active()) continue;
-                if (static_cast<size_t>(request_buffer_id) <
-                        local_segment_desc->buffers.size() &&
-                    local_segment_desc->buffers[request_buffer_id].lkey.size() >
-                        static_cast<size_t>(d) &&
-                    local_segment_desc->buffers[request_buffer_id].lkey[d] !=
-                        0) {
-                    active_devs.push_back({d, ctx});
-                }
-            }
-
-            if (active_devs.empty()) {
-                LOG(ERROR) << "No active EFA device for striping transfer of "
-                           << request.length << " bytes";
-                for (auto& entry : slices_to_post)
-                    for (auto s : entry.second) s->markFailed();
-                return Status::AddressNotRegistered(
-                    "No active EFA device found for striping");
-            }
-
-            size_t num_nics = active_devs.size();
-            size_t chunk_size = request.length / num_nics;
-
-            // Pre-resolve remote peer info ONCE for all striped slices.
-            // This eliminates per-slice overhead in submitPostSend:
-            //   - getSegmentDescByID lookup
-            //   - selectDevice() call
-            //   - string construction for peer_nic_path
-            // Using retry_count=i distributes slices across remote NICs.
-            auto peer_segment_desc =
-                metadata_->getSegmentDescByID(request.target_id);
-
-            for (size_t i = 0; i < num_nics; ++i) {
-                Slice* slice = getSliceCache().allocate();
-                assert(slice);
-                slice->peer_nic_path.clear();
-                slice->rdma.dest_rkey = 0;
-
-                size_t offset = i * chunk_size;
-                size_t len =
-                    (i == num_nics - 1) ? request.length - offset : chunk_size;
-
-                slice->source_addr = (char*)request.source + offset;
-                slice->length = len;
-                slice->opcode = request.opcode;
-                slice->rdma.dest_addr = request.target_offset + offset;
-                slice->rdma.retry_cnt = request.advise_retry_cnt;
-                slice->rdma.max_retry_cnt = kMaxRetryCount;
-                slice->task = &task;
-                slice->target_id = request.target_id;
-                slice->status = Slice::PENDING;
-                slice->ts = 0;
-                task.slice_list.push_back(slice);
-
-                int dev_id = active_devs[i].first;
-                auto& context = active_devs[i].second;
-                slice->rdma.source_lkey =
-                    local_segment_desc->buffers[request_buffer_id].lkey[dev_id];
-
-                // Pre-resolve remote NIC: set dest_rkey and peer_nic_path
-                // so submitPostSend() can skip expensive per-slice lookups.
-                if (peer_segment_desc) {
-                    int remote_buffer_id = -1, remote_device_id = -1;
-                    if (selectDevice(peer_segment_desc.get(),
-                                     request.target_offset + offset, len,
-                                     remote_buffer_id, remote_device_id,
-                                     static_cast<int>(i)) == 0) {
-                        slice->rdma.dest_rkey =
-                            peer_segment_desc->buffers[remote_buffer_id]
-                                .rkey[remote_device_id];
-                        slice->peer_nic_path =
-                            peer_segment_desc->name + "@" +
-                            peer_segment_desc->devices[remote_device_id].name;
-                    }
-                }
-
-                slices_to_post[context].push_back(slice);
-                __sync_fetch_and_add(&task.total_bytes, slice->length);
-                __sync_fetch_and_add(&task.slice_count, 1);
-            }
-        } else if (request_buffer_id >= 0 && request_device_id >= 0) {
-            // SMALL TRANSFER (or single-NIC): one slice, no sub-slicing.
-            // Round-robin NIC selection is handled by selectDevice above.
+        if (request_buffer_id >= 0 && request_device_id >= 0) {
+            // One slice per request.  Round-robin NIC selection is
+            // handled by selectDevice above.
             auto& context = context_list_[request_device_id];
             if (!context || !context->active()) {
                 LOG(ERROR) << "EFA Device " << request_device_id
@@ -1093,10 +1029,8 @@ int EfaTransport::initializeEfaResources() {
     for (auto& device_name : efa_devices) {
         auto context = std::make_shared<EfaContext>(*this, device_name);
         auto& config = globalConfig();
-        int ret = context->construct(config.num_cq_per_ctx,
-                                     config.num_comp_channels_per_ctx,
-                                     config.port, config.gid_index,
-                                     config.max_cqe, config.max_ep_per_ctx);
+        int ret = context->construct(config.num_cq_per_ctx, config.max_cqe,
+                                     config.max_ep_per_ctx);
         if (ret) {
             local_topology_->disableDevice(device_name);
             LOG(WARNING) << "EfaTransport: Disable device " << device_name;
diff --git a/mooncake-transfer-engine/src/transport/hip_transport/hip_transport.cpp b/mooncake-transfer-engine/src/transport/hip_transport/hip_transport.cpp
index 2b672fc2..f91206fa 100644
--- a/mooncake-transfer-engine/src/transport/hip_transport/hip_transport.cpp
+++ b/mooncake-transfer-engine/src/transport/hip_transport/hip_transport.cpp
@@ -27,6 +27,7 @@
 #include "common.h"
 #include "common/serialization.h"
 #include "config.h"
+#include "hip_device_guard.h"
 #include "transfer_metadata.h"
 #include "transport/transport.h"
 
@@ -251,6 +252,9 @@ static int setDeviceContext(void* source_ptr, int& device_id) {
 }
 
 static void setupP2PAccess(int num_devices) {
+    // The loop switches devices; the guard restores the caller's on return.
+    HipDeviceGuard device_guard;
+
     auto clearStickyPeerAccessError = [](int src_device, int dst_device) {
         // hipDeviceEnablePeerAccess may leave hipErrorPeerAccessAlreadyEnabled
         // in the runtime's last-error slot. Clear it so subsequent PyTorch HIP
@@ -440,6 +444,8 @@ Status HipTransport::startAsyncTransfer(const TransferRequest& request,
     hipError_t err;
     int device_id;
 
+    HipDeviceGuard device_guard;
+
     if (setDeviceContext(request.source, device_id) != 0) {
         return Status::InvalidArgument("Failed to set device context");
     }
diff --git a/mooncake-transfer-engine/src/transport/intranode_nvlink_transport/intranode_nvlink_transport.cpp b/mooncake-transfer-engine/src/transport/intranode_nvlink_transport/intranode_nvlink_transport.cpp
index bf9bfb63..710718ad 100644
--- a/mooncake-transfer-engine/src/transport/intranode_nvlink_transport/intranode_nvlink_transport.cpp
+++ b/mooncake-transfer-engine/src/transport/intranode_nvlink_transport/intranode_nvlink_transport.cpp
@@ -24,6 +24,7 @@
 #include <cstdint>
 #include <iomanip>
 #include <memory>
+#include <vector>
 
 #include "common.h"
 #include "common/serialization.h"
@@ -42,6 +43,169 @@ static bool checkCudaErrorReturn(cudaError_t result, const char *message) {
 }
 
 namespace mooncake {
+
+namespace {
+struct CudaStreamNVLinkRAII {
+    cudaStream_t stream_;
+    CudaStreamNVLinkRAII() : stream_(nullptr) {
+        auto err = cudaStreamCreateWithFlags(&stream_, cudaStreamNonBlocking);
+        if (err != cudaSuccess) {
+            LOG(FATAL) << "Failed to create NVLink CUDA stream: " << err
+                       << " - " << cudaGetErrorString(err);
+        }
+    }
+    ~CudaStreamNVLinkRAII() { cudaStreamDestroy(stream_); }
+};
+
+static thread_local CudaStreamNVLinkRAII tl_nvlink_stream;
+
+/// Thread-local CUDA event for GPU-level stream synchronization.
+/// Used to establish a GPU-visible dependency between cudaStreamPerThread
+/// and nvlink_stream, which is required by cudaMemcpyBatchAsync's
+/// srcAccessOrderStream attribute for cross-stream P2P copies.
+struct CudaEventNVLinkRAII {
+    cudaEvent_t event_;
+    CudaEventNVLinkRAII() {
+        auto err = cudaEventCreateWithFlags(&event_, cudaEventDisableTiming);
+        if (err != cudaSuccess) {
+            LOG(FATAL) << "Failed to create NVLink CUDA sync event: " << err
+                       << " - " << cudaGetErrorString(err);
+        }
+    }
+    ~CudaEventNVLinkRAII() { cudaEventDestroy(event_); }
+};
+
+static thread_local CudaEventNVLinkRAII tl_nvlink_sync_event;
+}  // anonymous namespace
+
+using Slice = Transport::Slice;
+
+/// Submit batched memcpy operations using cudaMemcpyBatchAsync when available
+/// (CUDA 12.8+), falling back to per-slice cudaMemcpyAsync otherwise.
+/// Uses cudaMemcpySrcAccessOrderStream to ensure source data visibility for
+/// P2P copies. This attribute is REQUIRED — without it, the GPU does not
+/// insert the necessary memory barriers for P2P access, causing segfaults.
+/// The caller must also establish GPU-level stream synchronization
+/// (via cudaEventRecord + cudaStreamWaitEvent) before calling this function.
+/// Individual slice errors are tracked so that slices whose memcpy failed
+/// are marked as FAILED while successfully submitted ones are POSTED.
+static void submitBatchMemcpy(const std::vector<Slice *> &slices,
+                              const std::vector<void *> &srcs,
+                              const std::vector<void *> &dsts,
+                              const std::vector<size_t> &sizes,
+                              cudaStream_t stream) {
+    if (slices.empty()) return;
+
+    const size_t count = slices.size();
+    cudaError_t err = cudaSuccess;
+
+    // Log the active memcpy path once per process lifetime
+    static const bool logged_once = [] {
+#if CUDART_VERSION >= 13000
+        LOG(INFO) << "IntraNodeNvlinkTransport: using cudaMemcpyBatchAsync "
+                  << "(CUDA >= 13.0 path)";
+#elif CUDART_VERSION >= 12080
+        LOG(INFO) << "IntraNodeNvlinkTransport: using cudaMemcpyBatchAsync "
+                  << "(CUDA >= 12.8 path)";
+#else
+        LOG(INFO)
+            << "IntraNodeNvlinkTransport: using per-slice cudaMemcpyAsync "
+            << "(CUDA < 12.8 fallback path)";
+#endif
+        return true;
+    }();
+    (void)logged_once;
+
+#if CUDART_VERSION >= 12080
+    // srcAccessOrderStream is REQUIRED for P2P copies — without it, the GPU
+    // does not insert necessary memory barriers for cross-device access,
+    // resulting in segmentation faults. The caller also establishes a
+    // GPU-level dependency via cudaEventRecord(cudaStreamPerThread) +
+    // cudaStreamWaitEvent(nvlink_stream) to ensure source data is coherent
+    // before the memcpy starts; these two mechanisms are complementary.
+    cudaMemcpyAttributes attr{};
+    attr.srcAccessOrder = cudaMemcpySrcAccessOrderStream;
+    size_t attrs_idx = 0;
+    // cudaMemcpyBatchAsync in CUDA 12.8 takes non-const size_t* for sizes
+    std::vector<size_t> mutable_sizes(sizes);
+    size_t fail_idx = count;
+#endif
+
+#if CUDART_VERSION >= 13000
+    err = cudaMemcpyBatchAsync(const_cast<const void **>(dsts.data()),
+                               const_cast<const void **>(srcs.data()),
+                               mutable_sizes.data(), static_cast<size_t>(count),
+                               &attr, &attrs_idx, 1, stream);
+    if (err != cudaSuccess) {
+        LOG(ERROR) << "IntraNodeNvlinkTransport: cudaMemcpyBatchAsync "
+                   << "failed: " << cudaGetErrorString(err);
+        // CUDA >= 13.0 does not return fail_idx; conservatively mark all
+        // as FAILED since we cannot determine which copies succeeded.
+        for (size_t i = 0; i < count; ++i) {
+            if (slices[i]->status == Slice::PENDING) {
+                slices[i]->markFailed();
+            }
+        }
+    } else {
+        for (size_t i = 0; i < count; ++i) {
+            slices[i]->status = Slice::POSTED;
+            slices[i]->local.cuda_stream = (void *)stream;
+        }
+    }
+#elif CUDART_VERSION >= 12080
+    err = cudaMemcpyBatchAsync(const_cast<void **>(dsts.data()),
+                               const_cast<void **>(srcs.data()),
+                               mutable_sizes.data(), static_cast<size_t>(count),
+                               &attr, &attrs_idx, 1, &fail_idx, stream);
+    if (err != cudaSuccess) {
+        if (fail_idx < count) {
+            LOG(ERROR) << "IntraNodeNvlinkTransport: cudaMemcpyBatchAsync "
+                       << "failed at index " << fail_idx
+                       << " (src=" << srcs[fail_idx]
+                       << ", dst=" << dsts[fail_idx]
+                       << ", size=" << sizes[fail_idx]
+                       << "): " << cudaGetErrorString(err);
+        } else {
+            LOG(ERROR) << "IntraNodeNvlinkTransport: cudaMemcpyBatchAsync "
+                       << "failed: " << cudaGetErrorString(err);
+        }
+        // Copies [0, fail_idx) were submitted successfully → POSTED.
+        // Copy [fail_idx] failed → FAILED.
+        // Copies (fail_idx, count) were never submitted → FAILED.
+        for (size_t i = 0; i < fail_idx; ++i) {
+            slices[i]->status = Slice::POSTED;
+            slices[i]->local.cuda_stream = (void *)stream;
+        }
+        for (size_t i = fail_idx; i < count; ++i) {
+            if (slices[i]->status == Slice::PENDING) {
+                slices[i]->markFailed();
+            }
+        }
+    } else {
+        for (size_t i = 0; i < count; ++i) {
+            slices[i]->status = Slice::POSTED;
+            slices[i]->local.cuda_stream = (void *)stream;
+        }
+    }
+#else
+    // Fallback for CUDA < 12.8: submit each memcpy individually
+    for (size_t i = 0; i < count; ++i) {
+        auto single_err = cudaMemcpyAsync(dsts[i], srcs[i], sizes[i],
+                                          cudaMemcpyDefault, stream);
+        if (single_err != cudaSuccess) {
+            LOG(ERROR) << "IntraNodeNvlinkTransport: cudaMemcpyAsync failed at "
+                       << "index " << i << ": "
+                       << cudaGetErrorString(single_err);
+            slices[i]->markFailed();
+            continue;
+        }
+        slices[i]->status = Slice::POSTED;
+        slices[i]->local.cuda_stream = (void *)stream;
+    }
+    return;  // Slice states already set above
+#endif
+}
+
 static int getNumDevices() {
     static int cached_num_devices = -1;
     if (cached_num_devices == -1) {
@@ -176,6 +340,41 @@ Status IntraNodeNvlinkTransport::submitTransfer(
     size_t task_id = batch_desc.task_list.size();
     batch_desc.task_list.resize(task_id + entries.size());
 
+    // Synchronize with the caller's CUDA stream before issuing any memcpy.
+    // PyTorch uses cudaStreamPerThread (per-thread default stream), NOT the
+    // legacy default stream (nullptr). Recording an event on
+    // cudaStreamPerThread and making nvlink_stream wait for it ensures that all
+    // PyTorch GPU operations on source/dest buffers complete before the NVLink
+    // memcpy starts. This GPU-level dependency is also required by
+    // cudaMemcpyBatchAsync's srcAccessOrderStream attribute, which needs
+    // the source data to be visible through the nvlink_stream's access order.
+    //
+    // Do NOT use the legacy default stream (nullptr) for cudaEventRecord,
+    // as it would trigger implicit synchronization with blocking streams and
+    // could cause deadlocks.
+    cudaStream_t stream = tl_nvlink_stream.stream_;
+    cudaError_t sync_err =
+        cudaEventRecord(tl_nvlink_sync_event.event_, cudaStreamPerThread);
+    if (sync_err != cudaSuccess) {
+        LOG(ERROR) << "IntraNodeNvlinkTransport: cudaEventRecord on "
+                      "cudaStreamPerThread failed: "
+                   << cudaGetErrorString(sync_err);
+        return Status::Context("cudaEventRecord failed: " +
+                               std::string(cudaGetErrorString(sync_err)));
+    }
+    sync_err = cudaStreamWaitEvent(stream, tl_nvlink_sync_event.event_, 0);
+    if (sync_err != cudaSuccess) {
+        LOG(ERROR) << "IntraNodeNvlinkTransport: cudaStreamWaitEvent failed: "
+                   << cudaGetErrorString(sync_err);
+        return Status::Context("cudaStreamWaitEvent failed: " +
+                               std::string(cudaGetErrorString(sync_err)));
+    }
+
+    // Phase 1: Prepare slices and collect memcpy parameters
+    std::vector<void *> dsts, srcs;
+    std::vector<size_t> sizes;
+    std::vector<Slice *> slices;
+
     for (auto &request : entries) {
         TransferTask &task = batch_desc.task_list[task_id];
         ++task_id;
@@ -194,20 +393,25 @@ Status IntraNodeNvlinkTransport::submitTransfer(
         slice->task = &task;
         slice->target_id = request.target_id;
         slice->status = Slice::PENDING;
+        slice->ts = getCurrentTimeInNano();
+        task.slice_list.push_back(slice);
         __sync_fetch_and_add(&task.slice_count, 1);
-        cudaError_t err;
-        if (slice->opcode == TransferRequest::READ)
-            err = cudaMemcpy(slice->source_addr, (void *)slice->local.dest_addr,
-                             slice->length, cudaMemcpyDefault);
-        else
-            err = cudaMemcpy((void *)slice->local.dest_addr, slice->source_addr,
-                             slice->length, cudaMemcpyDefault);
-        if (err != cudaSuccess)
-            slice->markFailed();
-        else
-            slice->markSuccess();
+
+        void *src = (request.opcode == TransferRequest::READ)
+                        ? (void *)slice->local.dest_addr
+                        : (void *)slice->source_addr;
+        void *dst = (request.opcode == TransferRequest::READ)
+                        ? slice->source_addr
+                        : (void *)slice->local.dest_addr;
+        srcs.push_back(src);
+        dsts.push_back(dst);
+        sizes.push_back(slice->length);
+        slices.push_back(slice);
     }
 
+    // Phase 2: Submit all memcpy operations
+    submitBatchMemcpy(slices, srcs, dsts, sizes, stream);
+
     return Status::OK();
 }
 
@@ -223,6 +427,29 @@ Status IntraNodeNvlinkTransport::getTransferStatus(BatchID batch_id,
             std::to_string(batch_id));
     }
     auto &task = batch_desc.task_list[task_id];
+    // Poll POSTED slices for async completion via cudaStreamQuery.
+    // Cache the query result per stream to avoid redundant driver calls,
+    // since multiple slices typically share the same CUDA stream.
+    std::unordered_map<cudaStream_t, cudaError_t> stream_status_cache;
+    for (auto *slice : task.slice_list) {
+        if (slice && slice->status == Slice::POSTED) {
+            cudaStream_t stream = (cudaStream_t)slice->local.cuda_stream;
+            auto it = stream_status_cache.find(stream);
+            cudaError_t cuda_err;
+            if (it == stream_status_cache.end()) {
+                cuda_err = cudaStreamQuery(stream);
+                stream_status_cache[stream] = cuda_err;
+            } else {
+                cuda_err = it->second;
+            }
+            if (cuda_err == cudaSuccess) {
+                slice->markSuccess();
+            } else if (cuda_err != cudaErrorNotReady) {
+                slice->markFailed();
+            }
+            // cudaErrorNotReady means still in progress, keep POSTED
+        }
+    }
     status.transferred_bytes = task.transferred_bytes;
     uint64_t success_slice_count = task.success_slice_count;
     uint64_t failed_slice_count = task.failed_slice_count;
@@ -241,6 +468,31 @@ Status IntraNodeNvlinkTransport::getTransferStatus(BatchID batch_id,
 
 Status IntraNodeNvlinkTransport::submitTransferTask(
     const std::vector<TransferTask *> &task_list) {
+    // Synchronize with the caller's CUDA stream before issuing any memcpy.
+    // See submitTransfer() for detailed rationale.
+    cudaStream_t stream = tl_nvlink_stream.stream_;
+    cudaError_t sync_err =
+        cudaEventRecord(tl_nvlink_sync_event.event_, cudaStreamPerThread);
+    if (sync_err != cudaSuccess) {
+        LOG(ERROR) << "IntraNodeNvlinkTransport: cudaEventRecord on "
+                      "cudaStreamPerThread failed: "
+                   << cudaGetErrorString(sync_err);
+        return Status::Context("cudaEventRecord failed: " +
+                               std::string(cudaGetErrorString(sync_err)));
+    }
+    sync_err = cudaStreamWaitEvent(stream, tl_nvlink_sync_event.event_, 0);
+    if (sync_err != cudaSuccess) {
+        LOG(ERROR) << "IntraNodeNvlinkTransport: cudaStreamWaitEvent failed: "
+                   << cudaGetErrorString(sync_err);
+        return Status::Context("cudaStreamWaitEvent failed: " +
+                               std::string(cudaGetErrorString(sync_err)));
+    }
+
+    // Phase 1: Prepare slices and collect memcpy parameters
+    std::vector<void *> dsts, srcs;
+    std::vector<size_t> sizes;
+    std::vector<Slice *> slices;
+
     for (size_t index = 0; index < task_list.size(); ++index) {
         assert(task_list[index]);
         auto &task = *task_list[index];
@@ -261,20 +513,25 @@ Status IntraNodeNvlinkTransport::submitTransferTask(
         slice->task = &task;
         slice->target_id = request.target_id;
         slice->status = Slice::PENDING;
+        slice->ts = getCurrentTimeInNano();
         task.slice_list.push_back(slice);
         __sync_fetch_and_add(&task.slice_count, 1);
-        cudaError_t err;
-        if (slice->opcode == TransferRequest::READ)
-            err = cudaMemcpy(slice->source_addr, (void *)slice->local.dest_addr,
-                             slice->length, cudaMemcpyDefault);
-        else
-            err = cudaMemcpy((void *)slice->local.dest_addr, slice->source_addr,
-                             slice->length, cudaMemcpyDefault);
-        if (err != cudaSuccess)
-            slice->markFailed();
-        else
-            slice->markSuccess();
+
+        void *src = (request.opcode == TransferRequest::READ)
+                        ? (void *)slice->local.dest_addr
+                        : (void *)slice->source_addr;
+        void *dst = (request.opcode == TransferRequest::READ)
+                        ? slice->source_addr
+                        : (void *)slice->local.dest_addr;
+        srcs.push_back(src);
+        dsts.push_back(dst);
+        sizes.push_back(slice->length);
+        slices.push_back(slice);
     }
+
+    // Phase 2: Submit all memcpy operations
+    submitBatchMemcpy(slices, srcs, dsts, sizes, stream);
+
     return Status::OK();
 }
 
diff --git a/mooncake-transfer-engine/src/transport/kunpeng_transport/CMakeLists.txt b/mooncake-transfer-engine/src/transport/kunpeng_transport/CMakeLists.txt
index c8616968..5583b7d8 100644
--- a/mooncake-transfer-engine/src/transport/kunpeng_transport/CMakeLists.txt
+++ b/mooncake-transfer-engine/src/transport/kunpeng_transport/CMakeLists.txt
@@ -1,4 +1,4 @@
-file(GLOB UB_SOURCES "*.cpp" "urma/urma_endpoint.cpp")
+file(GLOB UB_SOURCES "*.cpp" "urma/urma_endpoint.cpp" "ub_allocator.cpp")
 
 # Check if liburma.so exists
 find_library(URMA_LIBRARY urma PATHS /usr/lib64)
diff --git a/mooncake-transfer-engine/src/transport/kunpeng_transport/ub_allocator.cpp b/mooncake-transfer-engine/src/transport/kunpeng_transport/ub_allocator.cpp
new file mode 100644
index 00000000..609ecf98
--- /dev/null
+++ b/mooncake-transfer-engine/src/transport/kunpeng_transport/ub_allocator.cpp
@@ -0,0 +1,76 @@
+#include <algorithm>
+#include <cstddef>
+#include <mutex>
+#include <unordered_map>
+#include <vector>
+#include <glog/logging.h>
+#include <numa.h>
+
+#include "ub_allocator.h"
+
+namespace mooncake {
+struct UbStoreMemRange {
+    void* base;
+    size_t size;
+};
+std::mutex g_ub_store_mem_mutex;
+std::vector<UbStoreMemRange> g_ub_store_mem_ranges;
+
+size_t remove_store_memory_range(void* ptr) {
+    std::lock_guard<std::mutex> store_lock(g_ub_store_mem_mutex);
+
+    auto it = std::find_if(
+        g_ub_store_mem_ranges.begin(), g_ub_store_mem_ranges.end(),
+        [ptr](const UbStoreMemRange& range) { return range.base == ptr; });
+
+    if (it == g_ub_store_mem_ranges.end()) {
+        LOG(ERROR) << "failed for UB protocol, addr at " << ptr;
+        return 0;
+    }
+
+    size_t sz = it->size;
+    g_ub_store_mem_ranges.erase(it);
+    return sz;
+}
+
+void* ub_allocate_memory(size_t alignment, size_t total_size) {
+    void* ptr = numa_alloc_local(total_size);
+    if (!ptr) {
+        LOG(ERROR) << "failed for UB protocol, size=" << total_size
+                   << ", alignment : " << alignment;
+        return nullptr;
+    }
+    LOG(INFO) << "UB:  allocated total size : " << total_size
+              << ", alignment : " << alignment << " addr at " << ptr;
+
+    std::lock_guard<std::mutex> store_lock(g_ub_store_mem_mutex);
+    g_ub_store_mem_ranges.push_back({ptr, total_size});
+
+    return ptr;
+}
+
+void ub_free_memory(void* ptr) {
+    if (!ptr) {
+        return;
+    }
+    auto size = remove_store_memory_range(ptr);
+    numa_free(ptr, size);
+    LOG(INFO) << "UB: freed  bytes at " << ptr;
+}
+
+bool ub_is_store_memory(void* addr, size_t length) {
+    if (!addr || length == 0) return false;
+    auto addr_start = reinterpret_cast<uintptr_t>(addr);
+    uintptr_t addr_end = addr_start + length;
+    std::lock_guard<std::mutex> lock(g_ub_store_mem_mutex);
+    for (const auto& range : g_ub_store_mem_ranges) {
+        auto range_start = reinterpret_cast<uintptr_t>(range.base);
+        uintptr_t range_end = range_start + range.size;
+        if (addr_start >= range_start && addr_end <= range_end) {
+            return true;
+        }
+    }
+    return false;
+}
+
+}  // namespace mooncake
\ No newline at end of file
diff --git a/mooncake-transfer-engine/src/transport/kunpeng_transport/ub_context.cpp b/mooncake-transfer-engine/src/transport/kunpeng_transport/ub_context.cpp
index 4814a183..46a3d1e4 100644
--- a/mooncake-transfer-engine/src/transport/kunpeng_transport/ub_context.cpp
+++ b/mooncake-transfer-engine/src/transport/kunpeng_transport/ub_context.cpp
@@ -80,6 +80,29 @@ int UbSIEVEEndpointStore::deleteEndpoint(const std::string& peer_nic_path) {
     return 0;
 }
 
+int UbSIEVEEndpointStore::deleteEndpointByPtr(UbEndPoint* point_ptr) {
+    RWSpinlock::WriteGuard guard(endpoint_map_lock_);
+    for (auto iter = endpoint_map_.begin(); iter != endpoint_map_.end();
+         iter++) {
+        if (iter->second.first.get() == point_ptr) {
+            std::string peer_nic_path = iter->first;
+            iter->second.first->deconstruct();
+            waiting_list_len_++;
+            waiting_list_.insert(iter->second.first);
+            auto fifo_iter = fifo_map_[peer_nic_path];
+            if (hand_.has_value() && hand_.value() == fifo_iter) {
+                fifo_iter == fifo_list_.begin() ? hand_ = std::nullopt
+                                                : hand_ = std::prev(fifo_iter);
+            }
+            fifo_list_.erase(fifo_iter);
+            fifo_map_.erase(peer_nic_path);
+            endpoint_map_.erase(iter);
+            return 0;
+        }
+    }
+    return 0;
+}
+
 void UbSIEVEEndpointStore::evictEndpoint() {
     if (fifo_list_.empty()) {
         return;
@@ -246,8 +269,14 @@ int UbWorkerPool::submitPostSend(
         auto targetSegment =
             peer_segment_desc->buffers[buffer_id].tseg[device_id];
         slice->ub.r_seg = context_.retrieveRemoteSeg(targetSegment);
+        if (!slice->ub.r_seg) {
+            LOG(ERROR) << "[UB] retrieveRemoteSeg failed for target_id="
+                       << slice->target_id << " buffer_id=" << buffer_id
+                       << " device_id" << device_id
+                       << " dest_addr=" << slice->ub.dest_addr;
+        }
         auto peer_nic_path =
-            MakeNicPath(peer_segment_desc->name,
+            MakeNicPath(peer_segment_desc->nicPathServerName(),
                         peer_segment_desc->devices[device_id].name);
         slice->peer_nic_path = peer_nic_path;
         int shard_id = (slice->target_id * 10007 + device_id) % kShardCount;
@@ -333,7 +362,7 @@ void UbWorkerPool::performPostSend(int thread_id) {
         }
         if (!endpoint->active()) {
             if (endpoint->inactiveTime() > 1.0)
-                context_.deleteEndpoint(entry.first);
+                context_.deleteEndpointByPtr(endpoint.get());
             // enable for re-establishation
             for (auto& slice : entry.second) failed_slice_list.push_back(slice);
             entry.second.clear();
@@ -355,6 +384,10 @@ void UbWorkerPool::performPostSend(int thread_id) {
             entry.second.clear();
             continue;
         }
+        // Set endpoint pointer for each slice before submitting
+        for (auto& slice : entry.second) {
+            slice->ub.endpoint = endpoint.get();
+        }
         endpoint->submitPostSend(entry.second, failed_slice_list);
 #endif
     }
@@ -368,44 +401,42 @@ void UbWorkerPool::performPostSend(int thread_id) {
 void UbWorkerPool::performPoll(int thread_id) {
     int processed_slice_count = 0;
     const static size_t kPollCount = 64;
+    // context_.poll() aggregates each completion's jetty_depth here and
+    // calls markSuccess() on successful slices in place. Successful slices
+    // are NOT returned from poll(), so this worker never dereferences them
+    // after they may have been recycled by the submitting thread.
     std::unordered_map<volatile int*, int> jetty_depth_set;
+    std::vector<UbTransport::Slice*> failed_slices;
     for (int jfc_index = thread_id; jfc_index < context_.jfcCount();
          jfc_index += kTransferWorkerCount) {
-        UbTransport::Slice* cr[kPollCount];
-        int nr_poll = context_.poll(kPollCount, cr, jfc_index);
+        UbTransport::Slice* failed[kPollCount];
+        int num_failed = 0;
+        int nr_poll = context_.poll(kPollCount, failed, num_failed,
+                                    jetty_depth_set, jfc_index);
         if (nr_poll < 0) {
             LOG(ERROR) << "Worker: Failed to poll jetty for complete";
             continue;
         }
-        for (int i = 0; i < nr_poll; ++i) {
-            UbTransport::Slice* slice = cr[i];
+        int num_success = nr_poll - num_failed;
+        success_nr_polls += num_success;
+        processed_slice_count += num_success;
+        for (int i = 0; i < num_failed; ++i) {
+            UbTransport::Slice* slice = failed[i];
             assert(slice);
-            if (jetty_depth_set.count(slice->ub.jetty_depth))
-                jetty_depth_set[slice->ub.jetty_depth]++;
-            else
-                jetty_depth_set[slice->ub.jetty_depth] = 1;
-            if (cr[i]->status != Transport::Slice::SUCCESS) {
-                failed_nr_polls++;
-                if (context_.active() && failed_nr_polls > 32 &&
-                    !success_nr_polls) {
-                    LOG(WARNING) << "Too many errors found in local RNIC "
-                                 << context_.nicPath() << ", mark it inactive";
-                    context_.set_active(false);
-                }
-                context_.deleteEndpoint(slice->peer_nic_path);
-                slice->ub.retry_cnt++;
-                if (slice->ub.retry_cnt >= slice->ub.max_retry_cnt) {
-                    slice->markFailed();
-                    processed_slice_count_++;
-                } else {
-                    collective_slice_queue_[thread_id][slice->peer_nic_path]
-                        .push_back(slice);
-                    redispatch_counter_++;
-                }
+            failed_nr_polls++;
+            if (context_.active() && failed_nr_polls > 32 &&
+                !success_nr_polls) {
+                LOG(WARNING) << "Too many errors found in local RNIC "
+                             << context_.nicPath() << ", mark it inactive";
+                context_.set_active(false);
+            }
+            slice->ub.retry_cnt++;
+            if (slice->ub.retry_cnt >= slice->ub.max_retry_cnt) {
+                failed_slices.push_back(slice);
             } else {
-                // slice->markSuccess();
-                processed_slice_count++;
-                success_nr_polls++;
+                collective_slice_queue_[thread_id][slice->peer_nic_path]
+                    .push_back(slice);
+                redispatch_counter_++;
             }
         }
         if (nr_poll)
@@ -415,6 +446,18 @@ void UbWorkerPool::performPoll(int thread_id) {
     for (auto& entry : jetty_depth_set)
         __sync_fetch_and_sub(entry.first, entry.second);
 
+    // Slices that hit max_retry: final markFailed() after all reads (and the
+    // jetty depth returns above) are done. Failed slices were never published
+    // by poll(), so they remained safe to deref up to this point.
+    for (auto& slice : failed_slices) {
+        if (slice->ub.endpoint) {
+            auto ptr = static_cast<UbEndPoint*>(slice->ub.endpoint);
+            context_.deleteEndpointByPtr(ptr);
+        }
+        slice->markFailed();
+        processed_slice_count_++;
+    }
+
     if (processed_slice_count)
         processed_slice_count_.fetch_add(processed_slice_count);
 }
@@ -450,7 +493,7 @@ void UbWorkerPool::redispatch(std::vector<Transport::Slice*>& slice_list,
                 peer_segment_desc->buffers[buffer_id].tseg[device_id];
             slice->ub.r_seg = context_.retrieveRemoteSeg(targetSegment);
             auto peer_nic_path =
-                MakeNicPath(peer_segment_desc->name,
+                MakeNicPath(peer_segment_desc->nicPathServerName(),
                             peer_segment_desc->devices[device_id].name);
             slice->peer_nic_path = peer_nic_path;
             collective_slice_queue_[thread_id][peer_nic_path].push_back(slice);
diff --git a/mooncake-transfer-engine/src/transport/kunpeng_transport/ub_transport.cpp b/mooncake-transfer-engine/src/transport/kunpeng_transport/ub_transport.cpp
index 77b72c14..2dd3dc42 100644
--- a/mooncake-transfer-engine/src/transport/kunpeng_transport/ub_transport.cpp
+++ b/mooncake-transfer-engine/src/transport/kunpeng_transport/ub_transport.cpp
@@ -186,8 +186,18 @@ Status UbTransport::submitTransfer(
     size_t task_id = batch_desc.task_list.size();
     batch_desc.task_list.resize(task_id + entries.size());
     std::vector<TransferTask*> task_list;
-    task_list.reserve(batch_desc.task_list.size());
-    for (auto& task : batch_desc.task_list) task_list.push_back(&task);
+    task_list.reserve(entries.size());
+    for (auto& request : entries) {
+        auto& task = batch_desc.task_list[task_id];
+        ++task_id;
+        task.batch_id = batch_id;
+#ifdef USE_ASCEND_HETEROGENEOUS
+        task.request = const_cast<TransferRequest*>(&request);
+#else
+        task.request = &request;
+#endif
+        task_list.push_back(&task);
+    }
     return submitTransferTask(task_list);
 }
 
diff --git a/mooncake-transfer-engine/src/transport/kunpeng_transport/urma/mock_urma.cpp b/mooncake-transfer-engine/src/transport/kunpeng_transport/urma/mock_urma.cpp
index 020c3a6b..c392577a 100644
--- a/mooncake-transfer-engine/src/transport/kunpeng_transport/urma/mock_urma.cpp
+++ b/mooncake-transfer-engine/src/transport/kunpeng_transport/urma/mock_urma.cpp
@@ -1,16 +1,26 @@
 #include "urma_api.h"
-#include <map>
-#include <vector>
+#include <algorithm>
+#include <atomic>
 #include <cstring>
+#include <deque>
+#include <map>
 #include <mutex>
+#include <shared_mutex>
+#include <vector>
 
 namespace {
-std::mutex mock_mutex;
+
+struct JfcState {
+    std::mutex mutex;
+    std::deque<uint64_t> pending_ctx;
+};
+
+std::shared_mutex g_rw_mutex;
 bool initialized = false;
 std::vector<urma_device_t *> device_list;
 std::map<urma_context_t *, int> context_map;
 std::map<urma_jfce_t *, int> jfce_map;
-std::map<urma_jfc_t *, std::vector<uint64_t>> jfc_user_ctx_map;
+std::map<urma_jfc_t *, JfcState *> jfc_state_map;
 std::map<urma_jfr_t *, int> jfr_map;
 std::map<urma_target_seg_t *, int> seg_map;
 std::map<urma_jetty_t *, int> jetty_map;
@@ -32,10 +42,11 @@ urma_eid_info_t mock_eid_info = {
     .eid = {{0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, 0x09, 0x0A, 0x0B,
              0x0C, 0x0D, 0x0E, 0x0F, 0x10}},
     .eid_index = 0};
+
 }  // namespace
 
 urma_status_t urma_init(urma_init_attr_t *init_attr) {
-    std::lock_guard<std::mutex> lock(mock_mutex);
+    std::unique_lock<std::shared_mutex> lock(g_rw_mutex);
     if (initialized) {
         return URMA_EEXIST;
     }
@@ -44,7 +55,7 @@ urma_status_t urma_init(urma_init_attr_t *init_attr) {
 }
 
 urma_status_t urma_uninit(void) {
-    std::lock_guard<std::mutex> lock(mock_mutex);
+    std::unique_lock<std::shared_mutex> lock(g_rw_mutex);
     initialized = false;
     for (auto device : device_list) {
         delete device;
@@ -52,7 +63,10 @@ urma_status_t urma_uninit(void) {
     device_list.clear();
     context_map.clear();
     jfce_map.clear();
-    jfc_user_ctx_map.clear();
+    for (auto &kv : jfc_state_map) {
+        delete kv.second;
+    }
+    jfc_state_map.clear();
     jfr_map.clear();
     seg_map.clear();
     jetty_map.clear();
@@ -61,53 +75,81 @@ urma_status_t urma_uninit(void) {
 }
 
 urma_device_t **urma_get_device_list(int *num_devices) {
-    std::lock_guard<std::mutex> lock(mock_mutex);
-    if (!initialized) {
-        *num_devices = 0;
-        return nullptr;
+    {
+        std::shared_lock<std::shared_mutex> lock(g_rw_mutex);
+        if (!initialized) {
+            *num_devices = 0;
+            return nullptr;
+        }
+        if (!device_list.empty()) {
+            *num_devices = device_list.size();
+            urma_device_t **devices = new urma_device_t *[device_list.size()];
+            for (size_t i = 0; i < device_list.size(); ++i) {
+                devices[i] = device_list[i];
+            }
+            return devices;
+        }
     }
-
-    if (device_list.empty()) {
-        urma_device_t *device = new urma_device_t;
-        strcpy(device->name, "mock_urma_device");
-        strcpy(device->path, "/sys/class/infiniband/mock_device");
-        device->type = URMA_TRANSPORT_UB;
-        device->ops = nullptr;
-        device->sysfs_dev = nullptr;
-        device_list.push_back(device);
+    {
+        std::unique_lock<std::shared_mutex> write_lock(g_rw_mutex);
+        if (!initialized) {
+            *num_devices = 0;
+            return nullptr;
+        }
+        if (device_list.empty()) {
+            urma_device_t *device = new urma_device_t;
+            strcpy(device->name, "mock_urma_device");
+            strcpy(device->path, "/sys/class/infiniband/mock_device");
+            device->type = URMA_TRANSPORT_UB;
+            device->ops = nullptr;
+            device->sysfs_dev = nullptr;
+            device_list.push_back(device);
+        }
+        *num_devices = device_list.size();
+        urma_device_t **devices = new urma_device_t *[device_list.size()];
+        for (size_t i = 0; i < device_list.size(); ++i) {
+            devices[i] = device_list[i];
+        }
+        return devices;
     }
-
-    *num_devices = device_list.size();
-    urma_device_t **devices = new urma_device_t *[device_list.size()];
-    for (size_t i = 0; i < device_list.size(); ++i) {
-        devices[i] = device_list[i];
-    }
-    return devices;
 }
 
 urma_device_t *urma_get_device_by_name(const char *name) {
-    std::lock_guard<std::mutex> lock(mock_mutex);
-    if (!initialized) {
-        return nullptr;
-    }
-
-    if (device_list.empty()) {
-        auto *device = new urma_device_t;
-        strcpy(device->name, "mock_urma_device");
-        strcpy(device->path, "/sys/class/infiniband/mock_device");
-        device->type = URMA_TRANSPORT_UB;
-        device->ops = nullptr;
-        device->sysfs_dev = nullptr;
-        device_list.push_back(device);
-    }
-
-    for (auto device : device_list) {
-        if (strcmp(device->name, name) == 0) {
-            return device;
+    {
+        std::shared_lock<std::shared_mutex> lock(g_rw_mutex);
+        if (!initialized) {
+            return nullptr;
+        }
+        if (!device_list.empty()) {
+            for (auto device : device_list) {
+                if (strcmp(device->name, name) == 0) {
+                    return device;
+                }
+            }
+            return device_list[0];
         }
     }
-
-    return device_list.empty() ? nullptr : device_list[0];
+    {
+        std::unique_lock<std::shared_mutex> write_lock(g_rw_mutex);
+        if (!initialized) {
+            return nullptr;
+        }
+        if (device_list.empty()) {
+            auto *device = new urma_device_t;
+            strcpy(device->name, "mock_urma_device");
+            strcpy(device->path, "/sys/class/infiniband/mock_device");
+            device->type = URMA_TRANSPORT_UB;
+            device->ops = nullptr;
+            device->sysfs_dev = nullptr;
+            device_list.push_back(device);
+        }
+        for (auto device : device_list) {
+            if (strcmp(device->name, name) == 0) {
+                return device;
+            }
+        }
+        return device_list.empty() ? nullptr : device_list[0];
+    }
 }
 
 void urma_free_device_list(urma_device_t **device_list) {
@@ -118,7 +160,6 @@ void urma_free_device_list(urma_device_t **device_list) {
 
 urma_status_t urma_query_device(urma_device_t *device,
                                 urma_device_attr_t *attr) {
-    std::lock_guard<std::mutex> lock(mock_mutex);
     if (!device || !attr) {
         return URMA_EINVAL;
     }
@@ -129,7 +170,6 @@ urma_status_t urma_query_device(urma_device_t *device,
 }
 
 urma_eid_info_t *urma_get_eid_list(urma_device_t *device, uint32_t *eid_cnt) {
-    std::lock_guard<std::mutex> lock(mock_mutex);
     if (!device || !eid_cnt) {
         return nullptr;
     }
@@ -146,7 +186,7 @@ void urma_free_eid_list(urma_eid_info_t *eid_list) {
 }
 
 urma_context_t *urma_create_context(urma_device_t *device, uint32_t eid_index) {
-    std::lock_guard<std::mutex> lock(mock_mutex);
+    std::unique_lock<std::shared_mutex> lock(g_rw_mutex);
     if (!device) {
         return nullptr;
     }
@@ -158,7 +198,7 @@ urma_context_t *urma_create_context(urma_device_t *device, uint32_t eid_index) {
 }
 
 urma_status_t urma_delete_context(urma_context_t *ctx) {
-    std::lock_guard<std::mutex> lock(mock_mutex);
+    std::unique_lock<std::shared_mutex> lock(g_rw_mutex);
     if (!ctx || context_map.find(ctx) == context_map.end()) {
         return URMA_EINVAL;
     }
@@ -168,7 +208,7 @@ urma_status_t urma_delete_context(urma_context_t *ctx) {
 }
 
 urma_jfce_t *urma_create_jfce(urma_context_t *ctx) {
-    std::lock_guard<std::mutex> lock(mock_mutex);
+    std::unique_lock<std::shared_mutex> lock(g_rw_mutex);
     if (!ctx || context_map.find(ctx) == context_map.end()) {
         return nullptr;
     }
@@ -178,7 +218,7 @@ urma_jfce_t *urma_create_jfce(urma_context_t *ctx) {
 }
 
 urma_status_t urma_delete_jfce(urma_jfce_t *jfce) {
-    std::lock_guard<std::mutex> lock(mock_mutex);
+    std::unique_lock<std::shared_mutex> lock(g_rw_mutex);
     if (!jfce || jfce_map.find(jfce) == jfce_map.end()) {
         return URMA_EINVAL;
     }
@@ -188,7 +228,7 @@ urma_status_t urma_delete_jfce(urma_jfce_t *jfce) {
 }
 
 urma_jfc_t *urma_create_jfc(urma_context_t *ctx, urma_jfc_cfg_t *cfg) {
-    std::lock_guard<std::mutex> lock(mock_mutex);
+    std::unique_lock<std::shared_mutex> lock(g_rw_mutex);
     if (!ctx || !cfg || context_map.find(ctx) == context_map.end()) {
         return nullptr;
     }
@@ -201,22 +241,23 @@ urma_jfc_t *urma_create_jfc(urma_context_t *ctx, urma_jfc_cfg_t *cfg) {
     jfc->comp_events_acked = 0;
     jfc->async_events_acked = 0;
     jfc->jfc_cfg = *cfg;
-    jfc_user_ctx_map[jfc] = std::vector<uint64_t>();
+    jfc_state_map[jfc] = new JfcState();
     return jfc;
 }
 
 urma_status_t urma_delete_jfc(urma_jfc_t *jfc) {
-    std::lock_guard<std::mutex> lock(mock_mutex);
-    if (!jfc || jfc_user_ctx_map.find(jfc) == jfc_user_ctx_map.end()) {
+    std::unique_lock<std::shared_mutex> lock(g_rw_mutex);
+    if (!jfc || jfc_state_map.find(jfc) == jfc_state_map.end()) {
         return URMA_EINVAL;
     }
-    jfc_user_ctx_map.erase(jfc);
+    delete jfc_state_map[jfc];
+    jfc_state_map.erase(jfc);
     delete jfc;
     return URMA_SUCCESS;
 }
 
 urma_jfr_t *urma_create_jfr(urma_context_t *ctx, urma_jfr_cfg_t *cfg) {
-    std::lock_guard<std::mutex> lock(mock_mutex);
+    std::unique_lock<std::shared_mutex> lock(g_rw_mutex);
     if (!ctx || !cfg || context_map.find(ctx) == context_map.end()) {
         return nullptr;
     }
@@ -226,7 +267,7 @@ urma_jfr_t *urma_create_jfr(urma_context_t *ctx, urma_jfr_cfg_t *cfg) {
 }
 
 urma_status_t urma_delete_jfr(urma_jfr_t *jfr) {
-    std::lock_guard<std::mutex> lock(mock_mutex);
+    std::unique_lock<std::shared_mutex> lock(g_rw_mutex);
     if (!jfr || jfr_map.find(jfr) == jfr_map.end()) {
         return URMA_EINVAL;
     }
@@ -236,7 +277,7 @@ urma_status_t urma_delete_jfr(urma_jfr_t *jfr) {
 }
 
 urma_target_seg_t *urma_register_seg(urma_context_t *ctx, urma_seg_cfg_t *cfg) {
-    std::lock_guard<std::mutex> lock(mock_mutex);
+    std::unique_lock<std::shared_mutex> lock(g_rw_mutex);
     if (!ctx || !cfg || context_map.find(ctx) == context_map.end()) {
         return nullptr;
     }
@@ -252,7 +293,7 @@ urma_target_seg_t *urma_register_seg(urma_context_t *ctx, urma_seg_cfg_t *cfg) {
 }
 
 urma_status_t urma_unregister_seg(urma_target_seg_t *seg) {
-    std::lock_guard<std::mutex> lock(mock_mutex);
+    std::unique_lock<std::shared_mutex> lock(g_rw_mutex);
     if (!seg || seg_map.find(seg) == seg_map.end()) {
         return URMA_EINVAL;
     }
@@ -264,7 +305,7 @@ urma_status_t urma_unregister_seg(urma_target_seg_t *seg) {
 urma_target_seg_t *urma_import_seg(urma_context_t *ctx, urma_seg_t *seg,
                                    urma_token_t *token_value, uint64_t addr,
                                    urma_import_seg_flag_t flag) {
-    std::lock_guard<std::mutex> lock(mock_mutex);
+    std::unique_lock<std::shared_mutex> lock(g_rw_mutex);
     if (!ctx || !seg || !token_value ||
         context_map.find(ctx) == context_map.end()) {
         return nullptr;
@@ -277,7 +318,7 @@ urma_target_seg_t *urma_import_seg(urma_context_t *ctx, urma_seg_t *seg,
 }
 
 urma_status_t urma_unimport_seg(urma_target_seg_t *tseg) {
-    std::lock_guard<std::mutex> lock(mock_mutex);
+    std::unique_lock<std::shared_mutex> lock(g_rw_mutex);
     if (!tseg || seg_map.find(tseg) == seg_map.end()) {
         return URMA_EINVAL;
     }
@@ -288,8 +329,11 @@ urma_status_t urma_unimport_seg(urma_target_seg_t *tseg) {
 
 urma_status_t urma_get_async_event(urma_context_t *ctx,
                                    urma_async_event_t *event) {
-    std::lock_guard<std::mutex> lock(mock_mutex);
-    if (!ctx || !event || context_map.find(ctx) == context_map.end()) {
+    if (!ctx || !event) {
+        return URMA_EINVAL;
+    }
+    std::shared_lock<std::shared_mutex> lock(g_rw_mutex);
+    if (context_map.find(ctx) == context_map.end()) {
         return URMA_EINVAL;
     }
     return URMA_ETIMEOUT;
@@ -298,7 +342,7 @@ urma_status_t urma_get_async_event(urma_context_t *ctx,
 void urma_ack_async_event(urma_async_event_t *event) {}
 
 urma_jetty_t *urma_create_jetty(urma_context_t *ctx, urma_jetty_cfg_t *cfg) {
-    std::lock_guard<std::mutex> lock(mock_mutex);
+    std::unique_lock<std::shared_mutex> lock(g_rw_mutex);
     if (!ctx || !cfg || context_map.find(ctx) == context_map.end()) {
         return nullptr;
     }
@@ -314,7 +358,7 @@ urma_jetty_t *urma_create_jetty(urma_context_t *ctx, urma_jetty_cfg_t *cfg) {
 }
 
 urma_status_t urma_delete_jetty(urma_jetty_t *jetty) {
-    std::lock_guard<std::mutex> lock(mock_mutex);
+    std::unique_lock<std::shared_mutex> lock(g_rw_mutex);
     if (!jetty || jetty_map.find(jetty) == jetty_map.end()) {
         return URMA_EINVAL;
     }
@@ -324,7 +368,7 @@ urma_status_t urma_delete_jetty(urma_jetty_t *jetty) {
 }
 
 urma_status_t urma_unbind_jetty(urma_jetty_t *jetty) {
-    std::lock_guard<std::mutex> lock(mock_mutex);
+    std::unique_lock<std::shared_mutex> lock(g_rw_mutex);
     if (!jetty || jetty_map.find(jetty) == jetty_map.end()) {
         return URMA_EINVAL;
     }
@@ -335,7 +379,7 @@ urma_status_t urma_unbind_jetty(urma_jetty_t *jetty) {
 urma_target_jetty_t *urma_import_jetty(urma_context_t *ctx,
                                        urma_rjetty_t *rjetty,
                                        urma_token_t *token_value) {
-    std::lock_guard<std::mutex> lock(mock_mutex);
+    std::unique_lock<std::shared_mutex> lock(g_rw_mutex);
     if (!ctx || !rjetty || !token_value ||
         context_map.find(ctx) == context_map.end()) {
         return nullptr;
@@ -348,7 +392,7 @@ urma_target_jetty_t *urma_import_jetty(urma_context_t *ctx,
 }
 
 urma_status_t urma_unimport_jetty(urma_target_jetty_t *tjetty) {
-    std::lock_guard<std::mutex> lock(mock_mutex);
+    std::unique_lock<std::shared_mutex> lock(g_rw_mutex);
     if (!tjetty || target_jetty_map.find(tjetty) == target_jetty_map.end()) {
         return URMA_EINVAL;
     }
@@ -359,7 +403,7 @@ urma_status_t urma_unimport_jetty(urma_target_jetty_t *tjetty) {
 
 urma_status_t urma_bind_jetty(urma_jetty_t *jetty,
                               urma_target_jetty_t *tjetty) {
-    std::lock_guard<std::mutex> lock(mock_mutex);
+    std::unique_lock<std::shared_mutex> lock(g_rw_mutex);
     if (!jetty || !tjetty || jetty_map.find(jetty) == jetty_map.end() ||
         target_jetty_map.find(tjetty) == target_jetty_map.end()) {
         return URMA_EINVAL;
@@ -369,7 +413,7 @@ urma_status_t urma_bind_jetty(urma_jetty_t *jetty,
 }
 
 urma_status_t urma_modify_jetty(urma_jetty_t *jetty, urma_jetty_attr_t *attr) {
-    std::lock_guard<std::mutex> lock(mock_mutex);
+    std::shared_lock<std::shared_mutex> lock(g_rw_mutex);
     if (!jetty || !attr || jetty_map.find(jetty) == jetty_map.end()) {
         return URMA_EINVAL;
     }
@@ -378,19 +422,37 @@ urma_status_t urma_modify_jetty(urma_jetty_t *jetty, urma_jetty_attr_t *attr) {
 
 urma_status_t urma_post_jetty_send_wr(urma_jetty_t *jetty, urma_jfs_wr_t *wr,
                                       urma_jfs_wr_t **bad_wr) {
-    std::lock_guard<std::mutex> lock(mock_mutex);
-    if (!jetty || !wr || jetty_map.find(jetty) == jetty_map.end()) {
-        if (bad_wr) {
-            *bad_wr = wr;
+    {
+        std::shared_lock<std::shared_mutex> lock(g_rw_mutex);
+        if (!jetty || !wr || jetty_map.find(jetty) == jetty_map.end()) {
+            if (bad_wr) {
+                *bad_wr = wr;
+            }
+            return URMA_EINVAL;
         }
-        return URMA_EINVAL;
     }
 
-    urma_jfs_wr_t *current_wr = wr;
-    while (current_wr) {
-        jfc_user_ctx_map[jetty->jetty_cfg.jfs_cfg.jfc].push_back(
-            current_wr->user_ctx);
-        current_wr = current_wr->next;
+    urma_jfc_t *jfc = jetty->jetty_cfg.jfs_cfg.jfc;
+    JfcState *state = nullptr;
+    {
+        std::shared_lock<std::shared_mutex> lock(g_rw_mutex);
+        auto it = jfc_state_map.find(jfc);
+        if (it == jfc_state_map.end()) {
+            if (bad_wr) {
+                *bad_wr = wr;
+            }
+            return URMA_EINVAL;
+        }
+        state = it->second;
+    }
+
+    {
+        std::lock_guard<std::mutex> jfc_lock(state->mutex);
+        urma_jfs_wr_t *current_wr = wr;
+        while (current_wr) {
+            state->pending_ctx.push_back(current_wr->user_ctx);
+            current_wr = current_wr->next;
+        }
     }
 
     if (bad_wr) {
@@ -400,18 +462,27 @@ urma_status_t urma_post_jetty_send_wr(urma_jetty_t *jetty, urma_jfs_wr_t *wr,
 }
 
 int urma_poll_jfc(urma_jfc_t *jfc, int num_entries, urma_cr_t *cr_list) {
-    std::lock_guard<std::mutex> lock(mock_mutex);
-    if (!jfc || !cr_list ||
-        jfc_user_ctx_map.find(jfc) == jfc_user_ctx_map.end()) {
-        return -1;
+    JfcState *state = nullptr;
+    {
+        std::shared_lock<std::shared_mutex> lock(g_rw_mutex);
+        auto it = jfc_state_map.find(jfc);
+        if (it == jfc_state_map.end()) {
+            return -1;
+        }
+        state = it->second;
     }
-    int available = jfc_user_ctx_map[jfc].size();
-    int num_completed = std::min(num_entries, available);
-    for (int i = 0; i < num_completed; ++i) {
-        cr_list[i].status = URMA_CR_SUCCESS;
-        cr_list[i].user_ctx = jfc_user_ctx_map[jfc][i];
+
+    int num_completed = 0;
+    {
+        std::lock_guard<std::mutex> jfc_lock(state->mutex);
+        int available = static_cast<int>(state->pending_ctx.size());
+        num_completed = std::min(num_entries, available);
+        for (int i = 0; i < num_completed; ++i) {
+            cr_list[i].status = URMA_CR_SUCCESS;
+            cr_list[i].user_ctx = state->pending_ctx[i];
+        }
+        state->pending_ctx.erase(state->pending_ctx.begin(),
+                                 state->pending_ctx.begin() + num_completed);
     }
-    jfc_user_ctx_map[jfc].erase(jfc_user_ctx_map[jfc].begin(),
-                                jfc_user_ctx_map[jfc].begin() + num_completed);
     return num_completed;
 }
diff --git a/mooncake-transfer-engine/src/transport/kunpeng_transport/urma/urma_endpoint.cpp b/mooncake-transfer-engine/src/transport/kunpeng_transport/urma/urma_endpoint.cpp
index f7080a76..06c65d03 100644
--- a/mooncake-transfer-engine/src/transport/kunpeng_transport/urma/urma_endpoint.cpp
+++ b/mooncake-transfer-engine/src/transport/kunpeng_transport/urma/urma_endpoint.cpp
@@ -38,7 +38,9 @@ UrmaContext::~UrmaContext() {
     auto thisString = toString();
     worker_pool_.reset();
     LOG(INFO) << "destroy worker pool done.";
-    endpoint_store_->destroy();
+    if (endpoint_store_) {
+        endpoint_store_->destroy();
+    }
     LOG(INFO) << "destroy endpoint store done.";
     if (urma_context_) deconstruct();
     LOG(WARNING) << "finished destroy context : " << thisString;
@@ -420,13 +422,16 @@ int UrmaContext::openDevice(const std::string& device_name, uint8_t port,
             return ERR_CONTEXT;
         }
         for (int p = 0; p < MAX_PORT_CNT; p++) {
-            if (dev_attr_.port_attr[p].state == URMA_PORT_ACTIVE) {
+            auto port_attr = dev_attr_.port_attr[p];
+            if (port_attr.state == URMA_PORT_ACTIVE ||
+                port_attr.state == URMA_PORT_ACTIVE_DEFER) {
                 port_ = p;
                 break;
             }
         }
         if (dev_attr_.port_cnt != 0 &&
-            dev_attr_.port_attr[port_].state != URMA_PORT_ACTIVE) {
+            dev_attr_.port_attr[port_].state != URMA_PORT_ACTIVE &&
+            dev_attr_.port_attr[port_].state != URMA_PORT_ACTIVE_DEFER) {
             LOG(WARNING) << "Device " << device_name
                          << " not found active port";
             if (urma_delete_context(context)) {
@@ -513,8 +518,11 @@ bool UrmaContext::transEidFromString(const std::string& eid_str,
     return index == URMA_EID_SIZE;
 }
 
-int UrmaContext::poll(int num_entries, Transport::Slice** slices,
+int UrmaContext::poll(int num_entries, Transport::Slice** failed_slices,
+                      int& num_failed,
+                      std::unordered_map<volatile int*, int>& jetty_depth_set,
                       int jfc_index) {
+    num_failed = 0;
     urma_cr_t cr[num_entries];
     int nr_poll = urma_poll_jfc(jfc_list_[jfc_index].native, num_entries, cr);
     if (nr_poll < 0) {
@@ -522,17 +530,29 @@ int UrmaContext::poll(int num_entries, Transport::Slice** slices,
                    << device_name_;
         return ERR_CONTEXT;
     }
-    Transport::Slice s[nr_poll];
     for (int i = 0; i < nr_poll; ++i) {
         auto slice = (Transport::Slice*)cr[i].user_ctx;
         if (!slice) {
             continue;
         }
+
+        // All deref of `slice` (including the jetty_depth aggregation below)
+        // MUST happen before markSuccess(): once that publishes completion,
+        // the submitting thread may recycle the slice immediately.
+        auto* depth = slice->ub.jetty_depth;
+        auto it = jetty_depth_set.find(depth);
+        if (it != jetty_depth_set.end())
+            it->second++;
+        else
+            jetty_depth_set[depth] = 1;
+
         if (cr[i].status == URMA_CR_SUCCESS) {
+            // Safe to publish here — we are done with this slice and do not
+            // return it to the caller, so no one else will deref it.
             slice->markSuccess();
-            slices[i] = slice;
             continue;
         }
+
         if (cr[i].status != URMA_CR_WR_FLUSH_ERR ||
             show_work_request_flushed_error_)
             LOG(ERROR) << "Worker: Process failed for slice (opcode: "
@@ -550,6 +570,11 @@ int UrmaContext::poll(int num_entries, Transport::Slice** slices,
                        << ", comp_events_acked: "
                        << jfc_list_[jfc_index].native->comp_events_acked << " "
                        << jfc_list_[jfc_index].native->async_events_acked;
+
+        // Failed: hand the slice back so the caller can decide retry vs
+        // final markFailed(). Slice is NOT published, so the caller may
+        // safely deref it.
+        failed_slices[num_failed++] = slice;
     }
     return nr_poll;
 }
@@ -886,6 +911,10 @@ int UrmaEndpoint::submitPostSend(
         wr.flag.bs.inline_flag = 0;
         // Check if the jetty is in the imported_jetty_map_
         auto it = imported_jetty_map_.find(jetty_list_[jetty_index]);
+        if (it == imported_jetty_map_.end()) {
+            LOG(ERROR) << "Jetty not imported for endpoint, tjetty is nullptr"
+                       << jetty_index << ", local_nic=";
+        }
         if (it != imported_jetty_map_.end()) {
             wr.tjetty = it->second;
         } else {
@@ -895,6 +924,8 @@ int UrmaEndpoint::submitPostSend(
         slice->ts = getCurrentTimeInNano();
         slice->status = Transport::Slice::POSTED;
         slice->ub.jetty_depth = &wr_depth_list_[jetty_index];
+        // Set endpoint pointer for each slice before submitting
+        slice->ub.endpoint = this;
     }
     __sync_fetch_and_add(&wr_depth_list_[jetty_index], wr_count);
     __sync_fetch_and_add(jfc_outstanding_, wr_count);
@@ -966,6 +997,8 @@ int UrmaEndpoint::doSetupConnection(int jetty_index,
     rjetty.jetty_id.eid = eid;
     rjetty.trans_mode = URMA_TM_RC;
     rjetty.type = URMA_JETTY;
+    rjetty.tp_type = URMA_CTP;
+    rjetty.flag.value = 0;
     LOG(INFO) << "Peer jetty id = " << peer_jetty_num;
     urma_target_jetty_t* imported_jetty =
         urma_import_jetty(context_->urma_context_, &rjetty, &urma_token);
diff --git a/mooncake-transfer-engine/src/transport/maca_transport/CMakeLists.txt b/mooncake-transfer-engine/src/transport/maca_transport/CMakeLists.txt
new file mode 100644
index 00000000..eb1b41ec
--- /dev/null
+++ b/mooncake-transfer-engine/src/transport/maca_transport/CMakeLists.txt
@@ -0,0 +1,5 @@
+file(GLOB MACA_TRANSPORT_SOURCES "*.cpp")
+
+add_library(maca_transport OBJECT ${MACA_TRANSPORT_SOURCES})
+
+target_include_directories(maca_transport PUBLIC ${MACA_INCLUDE_DIR})
diff --git a/mooncake-transfer-engine/src/transport/maca_transport/maca_transport.cpp b/mooncake-transfer-engine/src/transport/maca_transport/maca_transport.cpp
new file mode 100644
index 00000000..517cea88
--- /dev/null
+++ b/mooncake-transfer-engine/src/transport/maca_transport/maca_transport.cpp
@@ -0,0 +1,482 @@
+// Copyright 2024 KVCache.AI
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "transport/maca_transport/maca_transport.h"
+
+#include <bits/stdint-uintn.h>
+#include "cuda_alike.h"
+#include <glog/logging.h>
+
+#include <algorithm>
+#include <cassert>
+#include <cstddef>
+#include <cstdint>
+#include <iomanip>
+#include <memory>
+
+#include "common.h"
+#include "common/serialization.h"
+#include "config.h"
+#include "transfer_engine.h"
+#include "transfer_metadata.h"
+#include "transport/transport.h"
+
+static bool checkCudaErrorReturn(cudaError_t result, const char *message) {
+    if (result != cudaSuccess) {
+        LOG(ERROR) << message << " (Error code: " << result << " - "
+                   << cudaGetErrorString(result) << ")" << std::endl;
+        return false;
+    }
+    return true;
+}
+
+namespace mooncake {
+static int getNumDevices() {
+    static int cached_num_devices = -1;
+    if (cached_num_devices == -1) {
+        if (!checkCudaErrorReturn(cudaGetDeviceCount(&cached_num_devices),
+                                  "MacaTransport: cudaGetDeviceCount failed")) {
+            return 0;
+        }
+    }
+    return cached_num_devices;
+}
+
+static bool enableP2PAccess(int src_device_id, int dst_device_id) {
+    int original_device;
+    cudaGetDevice(&original_device);
+
+    int canAccessPeer = 0;
+    if (!checkCudaErrorReturn(cudaDeviceCanAccessPeer(
+                                  &canAccessPeer, src_device_id, dst_device_id),
+                              "MacaTransport: failed to query peer access")) {
+        cudaSetDevice(original_device);
+        return false;
+    }
+
+    if (!canAccessPeer) {
+        LOG(ERROR) << "MacaTransport: device " << src_device_id
+                   << " cannot p2p access device " << dst_device_id;
+        cudaSetDevice(original_device);
+        return false;
+    }
+
+    // enable src->dst p2p access
+    if (!checkCudaErrorReturn(cudaSetDevice(src_device_id),
+                              "MacaTransport: failed to set device")) {
+        cudaSetDevice(original_device);
+        return false;
+    }
+    cudaError_t result = cudaDeviceEnablePeerAccess(dst_device_id, 0);
+
+    if (result != cudaSuccess && result != cudaErrorPeerAccessAlreadyEnabled) {
+        LOG(ERROR) << "MacaTransport: failed to enable p2p access (Error code: "
+                   << result << " - " << cudaGetErrorString(result) << ")"
+                   << std::endl;
+        cudaSetDevice(original_device);
+        return false;
+    }
+
+    // enable dst->src p2p access
+    if (!checkCudaErrorReturn(cudaSetDevice(dst_device_id),
+                              "MacaTransport: failed to set device")) {
+        cudaSetDevice(original_device);
+        return false;
+    }
+    result = cudaDeviceEnablePeerAccess(src_device_id, 0);
+
+    if (result != cudaSuccess && result != cudaErrorPeerAccessAlreadyEnabled) {
+        LOG(ERROR) << "MacaTransport: failed to enable p2p access (Error code: "
+                   << result << " - " << cudaGetErrorString(result) << ")"
+                   << std::endl;
+        cudaSetDevice(original_device);
+        return false;
+    }
+
+    cudaSetDevice(original_device);
+    return true;
+}
+
+static int getDeviceFromPointer(void *ptr) {
+    if (!ptr) return -1;
+
+    cudaPointerAttributes attr;
+    cudaError_t err = cudaPointerGetAttributes(&attr, ptr);
+    if (err != cudaSuccess) {
+        LOG(ERROR) << "MacaTransport: cudaPointerGetAttributes failed for "
+                   << ptr;
+        return -1;
+    }
+
+    if (attr.type == cudaMemoryTypeDevice) {
+        return attr.device;
+    }
+    return -1;
+}
+
+MacaTransport::MacaTransport() {
+    int num_devices = getNumDevices();
+    if (globalConfig().trace) {
+        LOG(INFO) << "MacaTransport: num_devices: " << num_devices;
+    }
+
+    for (int src_device_id = 0; src_device_id < num_devices; ++src_device_id) {
+        for (int dst_device_id = src_device_id + 1; dst_device_id < num_devices;
+             ++dst_device_id) {
+            if (enableP2PAccess(src_device_id, dst_device_id)) {
+                if (globalConfig().trace) {
+                    LOG(INFO)
+                        << "MacaTransport: enabled p2p access between device "
+                        << src_device_id << " and " << dst_device_id;
+                }
+            } else {
+                LOG(ERROR)
+                    << "MacaTransport: failed to enable p2p access between "
+                       "device "
+                    << src_device_id << " and " << dst_device_id;
+            }
+        }
+    }
+}
+
+MacaTransport::~MacaTransport() {
+    for (auto &entry : remap_entries_) {
+        cudaIpcCloseMemHandle(entry.second.shm_addr);
+    }
+    remap_entries_.clear();
+}
+
+int MacaTransport::install(std::string &local_server_name,
+                           std::shared_ptr<TransferMetadata> metadata,
+                           std::shared_ptr<Topology> topology) {
+    metadata_ = metadata;
+    local_server_name_ = local_server_name;
+
+    auto desc = std::make_shared<SegmentDesc>();
+    if (!desc) return ERR_MEMORY;
+    desc->name = local_server_name_;
+    desc->protocol = "maca";
+    metadata_->addLocalSegment(LOCAL_SEGMENT_ID, local_server_name_,
+                               std::move(desc));
+    return 0;
+}
+
+Status MacaTransport::submitTransfer(
+    BatchID batch_id, const std::vector<TransferRequest> &entries) {
+    auto &batch_desc = *((BatchDesc *)(batch_id));
+    if (batch_desc.task_list.size() + entries.size() > batch_desc.batch_size) {
+        LOG(ERROR) << "MacaTransport: Exceed the limitation of current batch's "
+                      "capacity";
+        return Status::InvalidArgument(
+            "MacaTransport: Exceed the limitation of capacity, batch id: " +
+            std::to_string(batch_id));
+    }
+
+    size_t task_id = batch_desc.task_list.size();
+    batch_desc.task_list.resize(task_id + entries.size());
+
+    for (auto &request : entries) {
+        TransferTask &task = batch_desc.task_list[task_id];
+        ++task_id;
+        uint64_t dest_addr = request.target_offset;
+        if (request.target_id != LOCAL_SEGMENT_ID) {
+            int rc = relocateSharedMemoryAddress(dest_addr, request.length,
+                                                 request.target_id);
+            if (rc) return Status::Memory("device memory not registered");
+        }
+        task.total_bytes = request.length;
+        Slice *slice = getSliceCache().allocate();
+        slice->source_addr = (char *)request.source;
+        slice->local.dest_addr = (char *)dest_addr;
+        slice->length = request.length;
+        slice->opcode = request.opcode;
+        slice->task = &task;
+        slice->target_id = request.target_id;
+        slice->status = Slice::PENDING;
+        __sync_fetch_and_add(&task.slice_count, 1);
+
+        // Set correct device context before memcpy
+        int original_device = -1;
+        cudaGetDevice(&original_device);
+        int target_device = getDeviceFromPointer(request.source);
+        if (target_device < 0)
+            target_device = getDeviceFromPointer((void *)dest_addr);
+        if (target_device >= 0) cudaSetDevice(target_device);
+
+        cudaError_t err;
+        if (slice->opcode == TransferRequest::READ)
+            err = cudaMemcpy(slice->source_addr, (void *)slice->local.dest_addr,
+                             slice->length, cudaMemcpyDefault);
+        else
+            err = cudaMemcpy((void *)slice->local.dest_addr, slice->source_addr,
+                             slice->length, cudaMemcpyDefault);
+        if (err != cudaSuccess)
+            slice->markFailed();
+        else
+            slice->markSuccess();
+
+        if (original_device >= 0) cudaSetDevice(original_device);
+    }
+
+    return Status::OK();
+}
+
+Status MacaTransport::getTransferStatus(BatchID batch_id, size_t task_id,
+                                        TransferStatus &status) {
+    auto &batch_desc = *((BatchDesc *)(batch_id));
+    const size_t task_count = batch_desc.task_list.size();
+    if (task_id >= task_count) {
+        return Status::InvalidArgument(
+            "MacaTransport::getTransportStatus invalid argument, batch id: " +
+            std::to_string(batch_id));
+    }
+    auto &task = batch_desc.task_list[task_id];
+    status.transferred_bytes = task.transferred_bytes;
+    uint64_t success_slice_count = task.success_slice_count;
+    uint64_t failed_slice_count = task.failed_slice_count;
+    if (success_slice_count + failed_slice_count == task.slice_count) {
+        if (failed_slice_count) {
+            status.s = TransferStatusEnum::FAILED;
+        } else {
+            status.s = TransferStatusEnum::COMPLETED;
+        }
+        task.is_finished = true;
+    } else {
+        status.s = TransferStatusEnum::WAITING;
+    }
+    return Status::OK();
+}
+
+Status MacaTransport::submitTransferTask(
+    const std::vector<TransferTask *> &task_list) {
+    for (size_t index = 0; index < task_list.size(); ++index) {
+        assert(task_list[index]);
+        auto &task = *task_list[index];
+        assert(task.request);
+        auto &request = *task.request;
+        uint64_t dest_addr = request.target_offset;
+        if (request.target_id != LOCAL_SEGMENT_ID) {
+            int rc = relocateSharedMemoryAddress(dest_addr, request.length,
+                                                 request.target_id);
+            if (rc) return Status::Memory("device memory not registered");
+        }
+        task.total_bytes = request.length;
+        Slice *slice = getSliceCache().allocate();
+        slice->source_addr = (char *)request.source;
+        slice->local.dest_addr = (char *)dest_addr;
+        slice->length = request.length;
+        slice->opcode = request.opcode;
+        slice->task = &task;
+        slice->target_id = request.target_id;
+        slice->status = Slice::PENDING;
+        task.slice_list.push_back(slice);
+        __sync_fetch_and_add(&task.slice_count, 1);
+
+        // Set correct device context before memcpy
+        int original_device = -1;
+        cudaGetDevice(&original_device);
+        int target_device = getDeviceFromPointer(request.source);
+        if (target_device < 0)
+            target_device = getDeviceFromPointer((void *)dest_addr);
+        if (target_device >= 0) cudaSetDevice(target_device);
+
+        cudaError_t err;
+        if (slice->opcode == TransferRequest::READ)
+            err = cudaMemcpy(slice->source_addr, (void *)slice->local.dest_addr,
+                             slice->length, cudaMemcpyDefault);
+        else
+            err = cudaMemcpy((void *)slice->local.dest_addr, slice->source_addr,
+                             slice->length, cudaMemcpyDefault);
+        if (err != cudaSuccess)
+            slice->markFailed();
+        else
+            slice->markSuccess();
+
+        if (original_device >= 0) cudaSetDevice(original_device);
+    }
+    return Status::OK();
+}
+
+int MacaTransport::registerLocalMemory(void *addr, size_t length,
+                                       const std::string &location,
+                                       bool remote_accessible,
+                                       bool update_metadata) {
+    std::lock_guard<std::mutex> lock(register_mutex_);
+    if (globalConfig().trace) {
+        LOG(INFO) << "register memory: addr " << addr << ", length " << length;
+    }
+    cudaPointerAttributes attr;
+    cudaError_t err = cudaPointerGetAttributes(&attr, addr);
+    if (err != cudaSuccess) {
+        LOG(ERROR) << "MacaTransport: cudaPointerGetAttributes failed";
+        return -1;
+    }
+
+    if (attr.type != cudaMemoryTypeDevice) {
+        LOG(ERROR) << "Unsupported memory type, " << addr << " " << attr.type;
+        return -1;
+    }
+
+    // Resolve the true cudaMalloc base address. Framework caching allocators
+    // (PyTorch, etc.) sub-allocate tensors within larger cudaMalloc segments.
+    // cudaIpcGetMemHandle always returns a handle for the entire segment, so
+    // we must register at segment granularity for correct IPC relocation.
+    CUdeviceptr base_ptr = 0;
+    size_t alloc_size = 0;
+    CUresult cu_err =
+        cuMemGetAddressRange(&base_ptr, &alloc_size, (CUdeviceptr)addr);
+    if (cu_err != CUDA_SUCCESS) {
+        LOG(ERROR) << "MacaTransport: cuMemGetAddressRange failed "
+                   << "for addr " << addr << " (error " << cu_err << ")";
+        return -1;
+    }
+
+    // Skip if this cudaMalloc block is already registered
+    if (registered_base_addrs_.count((uint64_t)base_ptr)) {
+        return 0;
+    }
+
+    cudaIpcMemHandle_t handle;
+    err = cudaIpcGetMemHandle(&handle, (void *)base_ptr);
+    if (err != cudaSuccess) {
+        LOG(ERROR) << "MacaTransport: cudaIpcGetMemHandle failed";
+        return -1;
+    }
+
+    (void)remote_accessible;
+    BufferDesc desc;
+    desc.addr = (uint64_t)base_ptr;
+    desc.length = alloc_size;
+    desc.name = location;
+    desc.shm_name = serializeBinaryData(&handle, sizeof(cudaIpcMemHandle_t));
+    int rc = metadata_->addLocalMemoryBuffer(desc, true);
+    if (rc == 0) {
+        registered_base_addrs_.insert((uint64_t)base_ptr);
+    }
+    return rc;
+}
+
+int MacaTransport::unregisterLocalMemory(void *addr, bool update_metadata) {
+    CUdeviceptr base_ptr = 0;
+    size_t alloc_size = 0;
+    CUresult cu_err =
+        cuMemGetAddressRange(&base_ptr, &alloc_size, (CUdeviceptr)addr);
+
+    void *key_ptr = addr;
+    if (cu_err == CUDA_SUCCESS) {
+        key_ptr = (void *)base_ptr;
+    } else {
+        LOG(WARNING)
+            << "MacaTransport: cuMemGetAddressRange failed for "
+            << "addr " << addr << " during unregister (error " << cu_err
+            << "). Memory may already be freed, using provided address.";
+    }
+
+    {
+        std::lock_guard<std::mutex> lock(register_mutex_);
+        registered_base_addrs_.erase((uint64_t)key_ptr);
+    }
+    return metadata_->removeLocalMemoryBuffer(key_ptr, update_metadata);
+}
+
+int MacaTransport::relocateSharedMemoryAddress(uint64_t &dest_addr,
+                                               uint64_t length,
+                                               uint64_t target_id) {
+    auto desc = metadata_->getSegmentDescByID(target_id);
+    int index = 0;
+    for (auto &entry : desc->buffers) {
+        if (!entry.shm_name.empty() && entry.addr <= dest_addr &&
+            dest_addr + length <= entry.addr + entry.length) {
+            remap_lock_.lockShared();
+            if (remap_entries_.count(std::make_pair(target_id, entry.addr))) {
+                auto shm_addr =
+                    remap_entries_[std::make_pair(target_id, entry.addr)]
+                        .shm_addr;
+                remap_lock_.unlockShared();
+                dest_addr = dest_addr - entry.addr + ((uint64_t)shm_addr);
+                return 0;
+            }
+            remap_lock_.unlockShared();
+            RWSpinlock::WriteGuard lock_guard(remap_lock_);
+            if (!remap_entries_.count(std::make_pair(target_id, entry.addr))) {
+                std::vector<unsigned char> output_buffer;
+                deserializeBinaryData(entry.shm_name, output_buffer);
+                if (output_buffer.size() == sizeof(cudaIpcMemHandle_t)) {
+                    cudaIpcMemHandle_t handle;
+                    memcpy(&handle, output_buffer.data(), sizeof(handle));
+                    void *shm_addr = nullptr;
+                    cudaError_t err = cudaIpcOpenMemHandle(
+                        &shm_addr, handle, cudaIpcMemLazyEnablePeerAccess);
+                    if (err != cudaSuccess) {
+                        LOG(ERROR) << "MacaTransport: "
+                                      "cudaIpcOpenMemHandle failed: "
+                                   << cudaGetErrorString(err);
+                        return -1;
+                    }
+                    OpenedShmEntry shm_entry;
+                    shm_entry.shm_addr = shm_addr;
+                    shm_entry.length = entry.length;
+                    remap_entries_[std::make_pair(target_id, entry.addr)] =
+                        shm_entry;
+                } else {
+                    LOG(ERROR) << "Mismatched MACA data transfer method";
+                    return -1;
+                }
+            }
+            auto shm_addr =
+                remap_entries_[std::make_pair(target_id, entry.addr)].shm_addr;
+            dest_addr = dest_addr - entry.addr + ((uint64_t)shm_addr);
+            return 0;
+        }
+        index++;
+    }
+    LOG(ERROR) << "Requested address " << (void *)dest_addr << " to "
+               << (void *)(dest_addr + length) << " not found!";
+    return ERR_INVALID_ARGUMENT;
+}
+
+int MacaTransport::registerLocalMemoryBatch(
+    const std::vector<Transport::BufferEntry> &buffer_list,
+    const std::string &location) {
+    for (auto &buffer : buffer_list)
+        registerLocalMemory(buffer.addr, buffer.length, location, true, false);
+    return metadata_->updateLocalSegmentDesc();
+}
+
+int MacaTransport::unregisterLocalMemoryBatch(
+    const std::vector<void *> &addr_list) {
+    for (auto &addr : addr_list) unregisterLocalMemory(addr, false);
+    return metadata_->updateLocalSegmentDesc();
+}
+
+void *MacaTransport::allocatePinnedLocalMemory(size_t size) {
+    void *ptr = nullptr;
+    cudaError_t res = cudaMalloc(&ptr, size);
+    if (res == cudaSuccess) {
+        LOG(INFO) << "MacaTransport: allocated device memory " << size
+                  << " bytes";
+        return ptr;
+    } else {
+        LOG(ERROR) << "MacaTransport: cudaMalloc failed: "
+                   << cudaGetErrorString(res);
+        return nullptr;
+    }
+}
+
+void MacaTransport::freePinnedLocalMemory(void *ptr) {
+    cudaFree(ptr);
+    return;
+}
+
+}  // namespace mooncake
diff --git a/mooncake-transfer-engine/src/transport/nvlink_transport/nvlink_transport.cpp b/mooncake-transfer-engine/src/transport/nvlink_transport/nvlink_transport.cpp
index 7cf80652..705b3abe 100644
--- a/mooncake-transfer-engine/src/transport/nvlink_transport/nvlink_transport.cpp
+++ b/mooncake-transfer-engine/src/transport/nvlink_transport/nvlink_transport.cpp
@@ -24,6 +24,7 @@
 #include <cstdint>
 #include <iomanip>
 #include <memory>
+#include <vector>
 
 #include "common.h"
 #include "common/serialization.h"
@@ -42,6 +43,168 @@ static bool checkCudaErrorReturn(cudaError_t result, const char *message) {
 }
 
 namespace mooncake {
+
+namespace {
+struct CudaStreamNVLinkRAII {
+    cudaStream_t stream_;
+    CudaStreamNVLinkRAII() : stream_(nullptr) {
+        auto err = cudaStreamCreateWithFlags(&stream_, cudaStreamNonBlocking);
+        if (err != cudaSuccess) {
+            LOG(FATAL) << "Failed to create NVLink CUDA stream: " << err
+                       << " - " << cudaGetErrorString(err);
+        }
+    }
+    ~CudaStreamNVLinkRAII() { cudaStreamDestroy(stream_); }
+};
+
+static thread_local CudaStreamNVLinkRAII tl_nvlink_stream;
+
+/// Thread-local CUDA event for GPU-level stream synchronization.
+/// Used to establish a GPU-visible dependency between cudaStreamPerThread
+/// and nvlink_stream, which is required by cudaMemcpyBatchAsync's
+/// srcAccessOrderStream attribute for cross-stream P2P copies.
+struct CudaEventNVLinkRAII {
+    cudaEvent_t event_;
+    CudaEventNVLinkRAII() {
+        auto err = cudaEventCreateWithFlags(&event_, cudaEventDisableTiming);
+        if (err != cudaSuccess) {
+            LOG(FATAL) << "Failed to create NVLink CUDA sync event: " << err
+                       << " - " << cudaGetErrorString(err);
+        }
+    }
+    ~CudaEventNVLinkRAII() { cudaEventDestroy(event_); }
+};
+
+static thread_local CudaEventNVLinkRAII tl_nvlink_sync_event;
+}  // anonymous namespace
+
+using Slice = Transport::Slice;
+
+/// Submit batched memcpy operations using cudaMemcpyBatchAsync when available
+/// (CUDA 12.8+), falling back to per-slice cudaMemcpyAsync otherwise.
+/// Uses cudaMemcpySrcAccessOrderStream to ensure source data visibility for
+/// P2P copies. This attribute is REQUIRED — without it, the GPU does not
+/// insert the necessary memory barriers for P2P access, causing segfaults.
+/// The caller must also establish GPU-level stream synchronization
+/// (via cudaEventRecord + cudaStreamWaitEvent) before calling this function.
+/// Individual slice errors are tracked so that slices whose memcpy failed
+/// are marked as FAILED while successfully submitted ones are POSTED.
+static void submitBatchMemcpy(const std::vector<Slice *> &slices,
+                              const std::vector<void *> &srcs,
+                              const std::vector<void *> &dsts,
+                              const std::vector<size_t> &sizes,
+                              cudaStream_t stream) {
+    if (slices.empty()) return;
+
+    const size_t count = slices.size();
+    cudaError_t err = cudaSuccess;
+
+    // Log the active memcpy path once per process lifetime
+    static const bool logged_once = [] {
+#if CUDART_VERSION >= 13000
+        LOG(INFO) << "NvlinkTransport: using cudaMemcpyBatchAsync "
+                  << "(CUDA >= 13.0 path)";
+#elif CUDART_VERSION >= 12080
+        LOG(INFO) << "NvlinkTransport: using cudaMemcpyBatchAsync "
+                  << "(CUDA >= 12.8 path)";
+#else
+        LOG(INFO) << "NvlinkTransport: using per-slice cudaMemcpyAsync "
+                  << "(CUDA < 12.8 fallback path)";
+#endif
+        return true;
+    }();
+    (void)logged_once;
+
+#if CUDART_VERSION >= 12080
+    // srcAccessOrderStream is REQUIRED for P2P copies — without it, the GPU
+    // does not insert necessary memory barriers for cross-device access,
+    // resulting in segmentation faults. The caller also establishes a
+    // GPU-level dependency via cudaEventRecord(cudaStreamPerThread) +
+    // cudaStreamWaitEvent(nvlink_stream) to ensure source data is coherent
+    // before the memcpy starts; these two mechanisms are complementary.
+    cudaMemcpyAttributes attr{};
+    attr.srcAccessOrder = cudaMemcpySrcAccessOrderStream;
+    size_t attrs_idx = 0;
+    // cudaMemcpyBatchAsync in CUDA 12.8 takes non-const size_t* for sizes
+    std::vector<size_t> mutable_sizes(sizes);
+    size_t fail_idx = count;
+#endif
+
+#if CUDART_VERSION >= 13000
+    err = cudaMemcpyBatchAsync(const_cast<const void **>(dsts.data()),
+                               const_cast<const void **>(srcs.data()),
+                               mutable_sizes.data(), static_cast<size_t>(count),
+                               &attr, &attrs_idx, 1, stream);
+    if (err != cudaSuccess) {
+        LOG(ERROR) << "NvlinkTransport: cudaMemcpyBatchAsync "
+                   << "failed: " << cudaGetErrorString(err);
+        // CUDA >= 13.0 does not return fail_idx; conservatively mark all
+        // as FAILED since we cannot determine which copies succeeded.
+        for (size_t i = 0; i < count; ++i) {
+            if (slices[i]->status == Slice::PENDING) {
+                slices[i]->markFailed();
+            }
+        }
+    } else {
+        for (size_t i = 0; i < count; ++i) {
+            slices[i]->status = Slice::POSTED;
+            slices[i]->local.cuda_stream = (void *)stream;
+        }
+    }
+#elif CUDART_VERSION >= 12080
+    err = cudaMemcpyBatchAsync(const_cast<void **>(dsts.data()),
+                               const_cast<void **>(srcs.data()),
+                               mutable_sizes.data(), static_cast<size_t>(count),
+                               &attr, &attrs_idx, 1, &fail_idx, stream);
+    if (err != cudaSuccess) {
+        if (fail_idx < count) {
+            LOG(ERROR) << "NvlinkTransport: cudaMemcpyBatchAsync "
+                       << "failed at index " << fail_idx
+                       << " (src=" << srcs[fail_idx]
+                       << ", dst=" << dsts[fail_idx]
+                       << ", size=" << sizes[fail_idx]
+                       << "): " << cudaGetErrorString(err);
+        } else {
+            LOG(ERROR) << "NvlinkTransport: cudaMemcpyBatchAsync "
+                       << "failed: " << cudaGetErrorString(err);
+        }
+        // Copies [0, fail_idx) were submitted successfully → POSTED.
+        // Copy [fail_idx] failed → FAILED.
+        // Copies (fail_idx, count) were never submitted → FAILED.
+        for (size_t i = 0; i < fail_idx; ++i) {
+            slices[i]->status = Slice::POSTED;
+            slices[i]->local.cuda_stream = (void *)stream;
+        }
+        for (size_t i = fail_idx; i < count; ++i) {
+            if (slices[i]->status == Slice::PENDING) {
+                slices[i]->markFailed();
+            }
+        }
+    } else {
+        for (size_t i = 0; i < count; ++i) {
+            slices[i]->status = Slice::POSTED;
+            slices[i]->local.cuda_stream = (void *)stream;
+        }
+    }
+#else
+    // Fallback for CUDA < 12.8: submit each memcpy individually
+    for (size_t i = 0; i < count; ++i) {
+        auto single_err = cudaMemcpyAsync(dsts[i], srcs[i], sizes[i],
+                                          cudaMemcpyDefault, stream);
+        if (single_err != cudaSuccess) {
+            LOG(ERROR) << "NvlinkTransport: cudaMemcpyAsync failed at "
+                       << "index " << i << ": "
+                       << cudaGetErrorString(single_err);
+            slices[i]->markFailed();
+            continue;
+        }
+        slices[i]->status = Slice::POSTED;
+        slices[i]->local.cuda_stream = (void *)stream;
+    }
+    return;  // Slice states already set above
+#endif
+}
+
 static int getNumDevices() {
     static int cached_num_devices = -1;
     if (cached_num_devices == -1) {
@@ -203,6 +366,41 @@ Status NvlinkTransport::submitTransfer(
     size_t task_id = batch_desc.task_list.size();
     batch_desc.task_list.resize(task_id + entries.size());
 
+    // Synchronize with the caller's CUDA stream before issuing any memcpy.
+    // PyTorch uses cudaStreamPerThread (per-thread default stream), NOT the
+    // legacy default stream (nullptr). Recording an event on
+    // cudaStreamPerThread and making nvlink_stream wait for it ensures that all
+    // PyTorch GPU operations on source/dest buffers complete before the NVLink
+    // memcpy starts. This GPU-level dependency is also required by
+    // cudaMemcpyBatchAsync's srcAccessOrderStream attribute, which needs
+    // the source data to be visible through the nvlink_stream's access order.
+    //
+    // Do NOT use the legacy default stream (nullptr) for cudaEventRecord,
+    // as it would trigger implicit synchronization with blocking streams and
+    // could cause deadlocks.
+    cudaStream_t stream = tl_nvlink_stream.stream_;
+    cudaError_t sync_err =
+        cudaEventRecord(tl_nvlink_sync_event.event_, cudaStreamPerThread);
+    if (sync_err != cudaSuccess) {
+        LOG(ERROR) << "NvlinkTransport: cudaEventRecord on "
+                      "cudaStreamPerThread failed: "
+                   << cudaGetErrorString(sync_err);
+        return Status::Context("cudaEventRecord failed: " +
+                               std::string(cudaGetErrorString(sync_err)));
+    }
+    sync_err = cudaStreamWaitEvent(stream, tl_nvlink_sync_event.event_, 0);
+    if (sync_err != cudaSuccess) {
+        LOG(ERROR) << "NvlinkTransport: cudaStreamWaitEvent failed: "
+                   << cudaGetErrorString(sync_err);
+        return Status::Context("cudaStreamWaitEvent failed: " +
+                               std::string(cudaGetErrorString(sync_err)));
+    }
+
+    // Phase 1: Prepare slices and collect memcpy parameters
+    std::vector<void *> dsts, srcs;
+    std::vector<size_t> sizes;
+    std::vector<Slice *> slices;
+
     for (auto &request : entries) {
         TransferTask &task = batch_desc.task_list[task_id];
         ++task_id;
@@ -221,20 +419,25 @@ Status NvlinkTransport::submitTransfer(
         slice->task = &task;
         slice->target_id = request.target_id;
         slice->status = Slice::PENDING;
+        slice->ts = getCurrentTimeInNano();
+        task.slice_list.push_back(slice);
         __sync_fetch_and_add(&task.slice_count, 1);
-        cudaError_t err;
-        if (slice->opcode == TransferRequest::READ)
-            err = cudaMemcpy(slice->source_addr, (void *)slice->local.dest_addr,
-                             slice->length, cudaMemcpyDefault);
-        else
-            err = cudaMemcpy((void *)slice->local.dest_addr, slice->source_addr,
-                             slice->length, cudaMemcpyDefault);
-        if (err != cudaSuccess)
-            slice->markFailed();
-        else
-            slice->markSuccess();
+
+        void *src = (request.opcode == TransferRequest::READ)
+                        ? (void *)slice->local.dest_addr
+                        : (void *)slice->source_addr;
+        void *dst = (request.opcode == TransferRequest::READ)
+                        ? slice->source_addr
+                        : (void *)slice->local.dest_addr;
+        srcs.push_back(src);
+        dsts.push_back(dst);
+        sizes.push_back(slice->length);
+        slices.push_back(slice);
     }
 
+    // Phase 2: Submit all memcpy operations
+    submitBatchMemcpy(slices, srcs, dsts, sizes, stream);
+
     return Status::OK();
 }
 
@@ -248,6 +451,29 @@ Status NvlinkTransport::getTransferStatus(BatchID batch_id, size_t task_id,
             std::to_string(batch_id));
     }
     auto &task = batch_desc.task_list[task_id];
+    // Poll POSTED slices for async completion via cudaStreamQuery.
+    // Cache the query result per stream to avoid redundant driver calls,
+    // since multiple slices typically share the same CUDA stream.
+    std::unordered_map<cudaStream_t, cudaError_t> stream_status_cache;
+    for (auto *slice : task.slice_list) {
+        if (slice && slice->status == Slice::POSTED) {
+            cudaStream_t stream = (cudaStream_t)slice->local.cuda_stream;
+            auto it = stream_status_cache.find(stream);
+            cudaError_t cuda_err;
+            if (it == stream_status_cache.end()) {
+                cuda_err = cudaStreamQuery(stream);
+                stream_status_cache[stream] = cuda_err;
+            } else {
+                cuda_err = it->second;
+            }
+            if (cuda_err == cudaSuccess) {
+                slice->markSuccess();
+            } else if (cuda_err != cudaErrorNotReady) {
+                slice->markFailed();
+            }
+            // cudaErrorNotReady means still in progress, keep POSTED
+        }
+    }
     status.transferred_bytes = task.transferred_bytes;
     uint64_t success_slice_count = task.success_slice_count;
     uint64_t failed_slice_count = task.failed_slice_count;
@@ -266,6 +492,31 @@ Status NvlinkTransport::getTransferStatus(BatchID batch_id, size_t task_id,
 
 Status NvlinkTransport::submitTransferTask(
     const std::vector<TransferTask *> &task_list) {
+    // Synchronize with the caller's CUDA stream before issuing any memcpy.
+    // See submitTransfer() for detailed rationale.
+    cudaStream_t stream = tl_nvlink_stream.stream_;
+    cudaError_t sync_err =
+        cudaEventRecord(tl_nvlink_sync_event.event_, cudaStreamPerThread);
+    if (sync_err != cudaSuccess) {
+        LOG(ERROR) << "NvlinkTransport: cudaEventRecord on "
+                      "cudaStreamPerThread failed: "
+                   << cudaGetErrorString(sync_err);
+        return Status::Context("cudaEventRecord failed: " +
+                               std::string(cudaGetErrorString(sync_err)));
+    }
+    sync_err = cudaStreamWaitEvent(stream, tl_nvlink_sync_event.event_, 0);
+    if (sync_err != cudaSuccess) {
+        LOG(ERROR) << "NvlinkTransport: cudaStreamWaitEvent failed: "
+                   << cudaGetErrorString(sync_err);
+        return Status::Context("cudaStreamWaitEvent failed: " +
+                               std::string(cudaGetErrorString(sync_err)));
+    }
+
+    // Phase 1: Prepare slices and collect memcpy parameters
+    std::vector<void *> dsts, srcs;
+    std::vector<size_t> sizes;
+    std::vector<Slice *> slices;
+
     for (size_t index = 0; index < task_list.size(); ++index) {
         assert(task_list[index]);
         auto &task = *task_list[index];
@@ -286,20 +537,25 @@ Status NvlinkTransport::submitTransferTask(
         slice->task = &task;
         slice->target_id = request.target_id;
         slice->status = Slice::PENDING;
+        slice->ts = getCurrentTimeInNano();
         task.slice_list.push_back(slice);
         __sync_fetch_and_add(&task.slice_count, 1);
-        cudaError_t err;
-        if (slice->opcode == TransferRequest::READ)
-            err = cudaMemcpy(slice->source_addr, (void *)slice->local.dest_addr,
-                             slice->length, cudaMemcpyDefault);
-        else
-            err = cudaMemcpy((void *)slice->local.dest_addr, slice->source_addr,
-                             slice->length, cudaMemcpyDefault);
-        if (err != cudaSuccess)
-            slice->markFailed();
-        else
-            slice->markSuccess();
+
+        void *src = (request.opcode == TransferRequest::READ)
+                        ? (void *)slice->local.dest_addr
+                        : (void *)slice->source_addr;
+        void *dst = (request.opcode == TransferRequest::READ)
+                        ? slice->source_addr
+                        : (void *)slice->local.dest_addr;
+        srcs.push_back(src);
+        dsts.push_back(dst);
+        sizes.push_back(slice->length);
+        slices.push_back(slice);
     }
+
+    // Phase 2: Submit all memcpy operations
+    submitBatchMemcpy(slices, srcs, dsts, sizes, stream);
+
     return Status::OK();
 }
 
diff --git a/mooncake-transfer-engine/src/transport/nvmeof_transport/nvmeof_transport.cpp b/mooncake-transfer-engine/src/transport/nvmeof_transport/nvmeof_transport.cpp
index 67e174ef..d019abd7 100644
--- a/mooncake-transfer-engine/src/transport/nvmeof_transport/nvmeof_transport.cpp
+++ b/mooncake-transfer-engine/src/transport/nvmeof_transport/nvmeof_transport.cpp
@@ -87,8 +87,7 @@ Status NVMeoFTransport::getTransferStatus(BatchID batch_id, size_t task_id,
     for (size_t i = slice_id; i < slice_id + slice_num; ++i) {
         // LOG(INFO) << "task " << task_id << " i " << i << " upper bound " <<
         // slice_num;
-        auto event =
-            desc_pool_->getTransferStatus(nvmeof_desc.desc_idx_, slice_id);
+        auto event = desc_pool_->getTransferStatus(nvmeof_desc.desc_idx_, i);
         transfer_status.s = from_cufile_transfer_status(event.status);
         // TODO(FIXME): what to do if multi slices have different status?
         if (transfer_status.s == COMPLETED) {
@@ -206,12 +205,13 @@ Status NVMeoFTransport::submitTransfer(
 
 Status NVMeoFTransport::freeBatchID(BatchID batch_id) {
     auto &batch_desc = *((BatchDesc *)(batch_id));
-    auto &nvmeof_desc = *((NVMeoFBatchDesc *)(batch_desc.context));
-    int desc_idx = nvmeof_desc.desc_idx_;
+    auto *nvmeof_desc_ptr = (NVMeoFBatchDesc *)(batch_desc.context);
+    int desc_idx = nvmeof_desc_ptr->desc_idx_;
     Status rc = Transport::freeBatchID(batch_id);
     if (rc != Status::OK()) {
         return rc;
     }
+    delete nvmeof_desc_ptr;
     desc_pool_->freeCUfileDesc(desc_idx);
     return Status::OK();
 }
diff --git a/mooncake-transfer-engine/src/transport/rdma_transport/CMakeLists.txt b/mooncake-transfer-engine/src/transport/rdma_transport/CMakeLists.txt
index e3b211a0..dc2c10ea 100644
--- a/mooncake-transfer-engine/src/transport/rdma_transport/CMakeLists.txt
+++ b/mooncake-transfer-engine/src/transport/rdma_transport/CMakeLists.txt
@@ -1,4 +1,9 @@
 file(GLOB RDMA_SOURCES "*.cpp")
 
 add_library(rdma_transport OBJECT ${RDMA_SOURCES})
-target_link_libraries(rdma_transport PRIVATE JsonCpp::JsonCpp glog::glog pthread)
+target_link_libraries(rdma_transport PRIVATE JsonCpp::JsonCpp glog::glog
+                                             pthread)
+
+if(USE_MLX5DV)
+  target_compile_definitions(rdma_transport PRIVATE USE_MLX5DV)
+endif()
diff --git a/mooncake-transfer-engine/src/transport/rdma_transport/endpoint_store.cpp b/mooncake-transfer-engine/src/transport/rdma_transport/endpoint_store.cpp
index 7dfc550a..9fc47179 100644
--- a/mooncake-transfer-engine/src/transport/rdma_transport/endpoint_store.cpp
+++ b/mooncake-transfer-engine/src/transport/rdma_transport/endpoint_store.cpp
@@ -78,11 +78,14 @@ std::shared_ptr<RdmaEndPoint> FIFOEndpointStore::insertEndpoint(
 int FIFOEndpointStore::deleteEndpoint(const std::string &peer_nic_path) {
     RWSpinlock::WriteGuard guard(endpoint_map_lock_);
     auto iter = endpoint_map_.find(peer_nic_path);
-    // remove endpoint but leaving it status unchanged
-    // in case it is setting up connection or submitting slice
+    // Begin two-phase destruction: mark endpoint as destroying and move QPs
+    // to ERR state so inflight WRs are flushed to CQ. The endpoint is moved
+    // to waiting_list_ and will be fully destroyed by reclaimEndpoint() once
+    // all outstanding WRs have been drained.
     if (iter != endpoint_map_.end()) {
+        waiting_list_len_++;
+        iter->second->beginDestroy();
         waiting_list_.insert(iter->second);
-        iter->second->set_active(false);
         endpoint_map_.erase(iter);
         auto fifo_iter = fifo_map_[peer_nic_path];
         fifo_list_.erase(fifo_iter);
@@ -91,23 +94,49 @@ int FIFOEndpointStore::deleteEndpoint(const std::string &peer_nic_path) {
     return 0;
 }
 
+int FIFOEndpointStore::deleteEndpointByPtr(const RdmaEndPoint *endpoint_ptr) {
+    RWSpinlock::WriteGuard guard(endpoint_map_lock_);
+    // Find endpoint by pointer
+    for (auto iter = endpoint_map_.begin(); iter != endpoint_map_.end();
+         ++iter) {
+        if (iter->second.get() == endpoint_ptr) {
+            std::string peer_nic_path = iter->first;
+            waiting_list_len_++;
+            iter->second->beginDestroy();
+            waiting_list_.insert(iter->second);
+            endpoint_map_.erase(iter);
+            auto fifo_iter = fifo_map_[peer_nic_path];
+            fifo_list_.erase(fifo_iter);
+            fifo_map_.erase(peer_nic_path);
+            return 0;
+        }
+    }
+    return -1;  // Not found
+}
+
 void FIFOEndpointStore::evictEndpoint() {
     if (fifo_list_.empty()) return;
     std::string victim = fifo_list_.front();
     fifo_list_.pop_front();
     fifo_map_.erase(victim);
     LOG(INFO) << victim << " evicted";
-    waiting_list_.insert(endpoint_map_[victim]);
+    waiting_list_len_++;
+    auto victim_endpoint = endpoint_map_[victim];
+    victim_endpoint->beginDestroy();
+    waiting_list_.insert(victim_endpoint);
     endpoint_map_.erase(victim);
     return;
 }
 
 void FIFOEndpointStore::reclaimEndpoint() {
+    if (waiting_list_len_.load(std::memory_order_relaxed) == 0) return;
     RWSpinlock::WriteGuard guard(endpoint_map_lock_);
     std::vector<std::shared_ptr<RdmaEndPoint>> to_delete;
-    for (auto &endpoint : waiting_list_)
-        if (!endpoint->hasOutstandingSlice()) to_delete.push_back(endpoint);
+    for (auto &endpoint : waiting_list_) {
+        if (endpoint->finishDestroy()) to_delete.push_back(endpoint);
+    }
     for (auto &endpoint : to_delete) waiting_list_.erase(endpoint);
+    waiting_list_len_ -= to_delete.size();
 }
 
 size_t FIFOEndpointStore::getSize() { return endpoint_map_.size(); }
@@ -120,9 +149,15 @@ int FIFOEndpointStore::destroyQPs() {
 }
 
 int FIFOEndpointStore::disconnectQPs() {
+    RWSpinlock::WriteGuard guard(endpoint_map_lock_);
     for (auto &kv : endpoint_map_) {
-        kv.second->disconnect();
+        kv.second->beginDestroy();
+        waiting_list_.insert(kv.second);
     }
+    waiting_list_len_ += endpoint_map_.size();
+    endpoint_map_.clear();
+    fifo_list_.clear();
+    fifo_map_.clear();
     return 0;
 }
 
@@ -135,6 +170,13 @@ size_t FIFOEndpointStore::getTotalQPNumber() {
     return total_qps;
 }
 
+void FIFOEndpointStore::testOnlyInsertWaiting(
+    std::shared_ptr<RdmaEndPoint> ep) {
+    RWSpinlock::WriteGuard guard(endpoint_map_lock_);
+    waiting_list_.insert(ep);
+    waiting_list_len_++;
+}
+
 std::shared_ptr<RdmaEndPoint> SIEVEEndpointStore::getEndpoint(
     const std::string &peer_nic_path) {
     RWSpinlock::ReadGuard guard(endpoint_map_lock_);
@@ -192,12 +234,14 @@ std::shared_ptr<RdmaEndPoint> SIEVEEndpointStore::insertEndpoint(
 int SIEVEEndpointStore::deleteEndpoint(const std::string &peer_nic_path) {
     RWSpinlock::WriteGuard guard(endpoint_map_lock_);
     auto iter = endpoint_map_.find(peer_nic_path);
-    // remove endpoint but leaving it status unchanged
-    // in case it is setting up connection or submitting slice
+    // Begin two-phase destruction: mark endpoint as destroying and move QPs
+    // to ERR state so inflight WRs are flushed to CQ. The endpoint is moved
+    // to waiting_list_ and will be fully destroyed by reclaimEndpoint() once
+    // all outstanding WRs have been drained.
     if (iter != endpoint_map_.end()) {
+        iter->second.first->beginDestroy();
         waiting_list_len_++;
         waiting_list_.insert(iter->second.first);
-        iter->second.first->set_active(false);
         endpoint_map_.erase(iter);
         auto fifo_iter = fifo_map_[peer_nic_path];
         if (hand_.has_value() && hand_.value() == fifo_iter) {
@@ -210,6 +254,30 @@ int SIEVEEndpointStore::deleteEndpoint(const std::string &peer_nic_path) {
     return 0;
 }
 
+int SIEVEEndpointStore::deleteEndpointByPtr(const RdmaEndPoint *endpoint_ptr) {
+    RWSpinlock::WriteGuard guard(endpoint_map_lock_);
+    // Find endpoint by pointer
+    for (auto iter = endpoint_map_.begin(); iter != endpoint_map_.end();
+         ++iter) {
+        if (iter->second.first.get() == endpoint_ptr) {
+            std::string peer_nic_path = iter->first;
+            iter->second.first->beginDestroy();
+            waiting_list_len_++;
+            waiting_list_.insert(iter->second.first);
+            auto fifo_iter = fifo_map_[peer_nic_path];
+            if (hand_.has_value() && hand_.value() == fifo_iter) {
+                fifo_iter == fifo_list_.begin() ? hand_ = std::nullopt
+                                                : hand_ = std::prev(fifo_iter);
+            }
+            fifo_list_.erase(fifo_iter);
+            fifo_map_.erase(peer_nic_path);
+            endpoint_map_.erase(iter);
+            return 0;
+        }
+    }
+    return -1;  // Not found
+}
+
 void SIEVEEndpointStore::evictEndpoint() {
     if (fifo_list_.empty()) {
         return;
@@ -231,7 +299,7 @@ void SIEVEEndpointStore::evictEndpoint() {
     fifo_map_.erase(victim);
     LOG(INFO) << victim << " evicted";
     auto victim_instance = endpoint_map_[victim].first;
-    victim_instance->set_active(false);
+    victim_instance->beginDestroy();
     waiting_list_len_++;
     waiting_list_.insert(victim_instance);
     endpoint_map_.erase(victim);
@@ -242,8 +310,9 @@ void SIEVEEndpointStore::reclaimEndpoint() {
     if (waiting_list_len_.load(std::memory_order_relaxed) == 0) return;
     RWSpinlock::WriteGuard guard(endpoint_map_lock_);
     std::vector<std::shared_ptr<RdmaEndPoint>> to_delete;
-    for (auto &endpoint : waiting_list_)
-        if (!endpoint->hasOutstandingSlice()) to_delete.push_back(endpoint);
+    for (auto &endpoint : waiting_list_) {
+        if (endpoint->finishDestroy()) to_delete.push_back(endpoint);
+    }
     for (auto &endpoint : to_delete) waiting_list_.erase(endpoint);
     waiting_list_len_ -= to_delete.size();
 }
@@ -255,13 +324,28 @@ int SIEVEEndpointStore::destroyQPs() {
 }
 
 int SIEVEEndpointStore::disconnectQPs() {
-    for (auto &endpoint : waiting_list_) endpoint->disconnect();
-    for (auto &kv : endpoint_map_) kv.second.first->disconnect();
+    RWSpinlock::WriteGuard guard(endpoint_map_lock_);
+    for (auto &kv : endpoint_map_) {
+        kv.second.first->beginDestroy();
+        waiting_list_.insert(kv.second.first);
+    }
+    waiting_list_len_ += endpoint_map_.size();
+    endpoint_map_.clear();
+    fifo_list_.clear();
+    fifo_map_.clear();
+    hand_ = std::nullopt;
     return 0;
 }
 
 size_t SIEVEEndpointStore::getSize() { return endpoint_map_.size(); }
 
+void SIEVEEndpointStore::testOnlyInsertWaiting(
+    std::shared_ptr<RdmaEndPoint> ep) {
+    RWSpinlock::WriteGuard guard(endpoint_map_lock_);
+    waiting_list_.insert(ep);
+    waiting_list_len_++;
+}
+
 size_t SIEVEEndpointStore::getTotalQPNumber() {
     RWSpinlock::ReadGuard guard(endpoint_map_lock_);
     size_t total_qps = 0;
diff --git a/mooncake-transfer-engine/src/transport/rdma_transport/rdma_context.cpp b/mooncake-transfer-engine/src/transport/rdma_transport/rdma_context.cpp
index 282a4ea6..71fdcc93 100644
--- a/mooncake-transfer-engine/src/transport/rdma_transport/rdma_context.cpp
+++ b/mooncake-transfer-engine/src/transport/rdma_transport/rdma_context.cpp
@@ -14,18 +14,35 @@
 
 #include "transport/rdma_transport/rdma_context.h"
 
+#include <algorithm>
+#include <cerrno>
+#include <cstdio>
 #include <fcntl.h>
+#include <netinet/in.h>
 #include <sys/epoll.h>
+#include <unistd.h>
 
 #include <atomic>
 #include <cassert>
+#include <exception>
 #include <fstream>
 #include <memory>
 #include <thread>
 
 #include "config.h"
 #include "cuda_alike.h"
+#include "environ.h"
+#include "hip_device_guard.h"
+#if defined(USE_HIP_DMABUF)
+#include <sys/utsname.h>
+
+#include <vector>
+
+#include <hsa/hsa.h>
+#include <hsa/hsa_ext_amd.h>
+#endif
 #include "transport/rdma_transport/endpoint_store.h"
+#include "transport/rdma_transport/rdma_gid_probe.h"
 #include "transport/rdma_transport/rdma_endpoint.h"
 #include "transport/rdma_transport/rdma_transport.h"
 #include "transport/rdma_transport/worker_pool.h"
@@ -39,6 +56,108 @@ static int isNullGid(union ibv_gid *gid) {
     return 1;
 }
 
+namespace {
+bool containsAddress(const MemoryRegionMeta &region, uintptr_t addr) {
+    const auto region_start = reinterpret_cast<uintptr_t>(region.addr);
+    const auto region_length = static_cast<uintptr_t>(region.mr->length);
+    return region_start <= addr && addr - region_start < region_length;
+}
+
+#if defined(USE_HIP_DMABUF)
+// Returns true when the kernel has CONFIG_PCI_P2PDMA and
+// CONFIG_DMABUF_MOVE_NOTIFY enabled, which are required for ibv_reg_dmabuf_mr
+// to produce working RDMA transfers (not just successful registration).
+// Result is cached after the first call.
+bool isKernelDmabufSupported() {
+    static const bool supported = []() {
+        if (const char *env = std::getenv("MOONCAKE_DISABLE_HIP_DMABUF")) {
+            if (std::string(env) != "0") {
+                LOG(INFO)
+                    << "HIP dmabuf disabled via MOONCAKE_DISABLE_HIP_DMABUF";
+                return false;
+            }
+        }
+        struct utsname uts{};
+        std::string release;
+        if (uname(&uts) == 0) release = uts.release;
+
+        const char *needles[] = {"CONFIG_PCI_P2PDMA=y",
+                                 "CONFIG_DMABUF_MOVE_NOTIFY=y"};
+        std::vector<bool> found(2, false);
+
+        const std::vector<std::string> candidates = {
+            "/proc/config.gz",  // compressed; checked below via gzip-magic skip
+            "/boot/config-" + release,
+            "/usr/src/linux-" + release + "/.config",
+            "/usr/src/linux/.config",
+            "/usr/lib/modules/" + release + "/config",
+            "/usr/lib/ostree-boot/config-" + release,
+            "/usr/lib/kernel/config-" + release,
+            "/usr/src/linux-headers-" + release + "/.config",
+            "/lib/modules/" + release + "/build/.config",
+        };
+
+        for (const auto &path : candidates) {
+            std::ifstream f(path);
+            if (!f.good()) continue;
+            // /proc/config.gz is gzipped; we skip it here (the kallsyms
+            // fallback below covers that case in practice).
+            if (path.find(".gz") != std::string::npos) continue;
+            std::string line;
+            while (std::getline(f, line)) {
+                for (size_t i = 0; i < 2; ++i) {
+                    if (!found[i] && line.find(needles[i]) != std::string::npos)
+                        found[i] = true;
+                }
+                if (found[0] && found[1]) break;
+            }
+            if (found[0] && found[1]) break;
+        }
+
+        // Fallback: probe /proc/kallsyms for the corresponding kernel symbols.
+        if (!found[0] || !found[1]) {
+            std::ifstream f("/proc/kallsyms");
+            if (f.good()) {
+                std::string line;
+                while (std::getline(f, line)) {
+                    if (!found[0] &&
+                        line.find("pci_p2pdma") != std::string::npos)
+                        found[0] = true;
+                    if (!found[1] &&
+                        line.find("dma_buf_move_notify") != std::string::npos)
+                        found[1] = true;
+                    if (found[0] && found[1]) break;
+                }
+            }
+        }
+
+        bool ok = found[0] && found[1];
+        if (!ok) {
+            LOG(WARNING)
+                << "Kernel lacks CONFIG_PCI_P2PDMA / CONFIG_DMABUF_MOVE_NOTIFY "
+                << "(p2pdma=" << found[0] << " move_notify=" << found[1]
+                << "), HIP dmabuf MR registration disabled, falling back to "
+                << "ibv_reg_mr() (which requires an amdgpu peermem driver). "
+                << "Rebuild kernel with both options for GPU-direct RDMA.";
+        }
+        return ok;
+    }();
+    return supported;
+}
+#endif  // USE_HIP_DMABUF
+
+std::string gidBytesToString(const uint8_t *raw) {
+    std::string gid_str;
+    char buf[16] = {0};
+    const static size_t kGidLength = 16;
+    for (size_t i = 0; i < kGidLength; ++i) {
+        snprintf(buf, sizeof(buf), "%02x", raw[i]);
+        gid_str += i == 0 ? buf : std::string(":") + buf;
+    }
+    return gid_str;
+}
+}  // namespace
+
 RdmaContext::RdmaContext(RdmaTransport &engine, const std::string &device_name)
     : device_name_(device_name),
       engine_(engine),
@@ -91,7 +210,7 @@ int RdmaContext::construct(size_t num_cq_list, size_t num_comp_channels,
     }
 
     num_comp_channel_ = num_comp_channels;
-    comp_channel_ = new ibv_comp_channel *[num_comp_channels];
+    comp_channel_ = new ibv_comp_channel *[num_comp_channels]();
     for (size_t i = 0; i < num_comp_channels; ++i) {
         comp_channel_[i] = ibv_create_comp_channel(context_);
         if (!comp_channel_[i]) {
@@ -110,6 +229,7 @@ int RdmaContext::construct(size_t num_cq_list, size_t num_comp_channels,
     if (joinNonblockingPollList(event_fd_, context_->async_fd)) {
         LOG(ERROR) << "Failed to register context async fd to epoll";
         close(event_fd_);
+        event_fd_ = -1;
         return ERR_CONTEXT;
     }
 
@@ -118,6 +238,7 @@ int RdmaContext::construct(size_t num_cq_list, size_t num_comp_channels,
             LOG(ERROR) << "Failed to register completion channel " << i
                        << " to epoll";
             close(event_fd_);
+            event_fd_ = -1;
             return ERR_CONTEXT;
         }
 
@@ -130,6 +251,7 @@ int RdmaContext::construct(size_t num_cq_list, size_t num_comp_channels,
         if (!cq) {
             PLOG(ERROR) << "Failed to create completion queue";
             close(event_fd_);
+            event_fd_ = -1;
             return ERR_CONTEXT;
         }
         cq_list_[i].native = cq;
@@ -137,8 +259,20 @@ int RdmaContext::construct(size_t num_cq_list, size_t num_comp_channels,
 
     worker_pool_ = std::make_shared<WorkerPool>(*this, socketId());
 
+#ifdef USE_MLX5DV
+    {
+        mlx5dv_context dv_ctx = {};
+        dv_ctx.comp_mask = MLX5DV_CONTEXT_MASK_NUM_LAG_PORTS;
+        if (mlx5dv_query_device(context_, &dv_ctx) == 0)
+            num_lag_ports_ = dv_ctx.num_lag_ports;
+    }
+    LOG(INFO) << "RDMA device: " << context_->device->name << ", LID: " << lid_
+              << ", GID: (GID_Index " << gid_index_ << ") " << gid()
+              << ", num_lag_ports: " << (int)num_lag_ports_;
+#else
     LOG(INFO) << "RDMA device: " << context_->device->name << ", LID: " << lid_
               << ", GID: (GID_Index " << gid_index_ << ") " << gid();
+#endif
 
     return 0;
 }
@@ -162,13 +296,13 @@ int RdmaContext::deconstruct() {
 
     endpoint_store_->destroyQPs();
 
-    for (auto &entry : memory_region_list_) {
+    for (auto &[_, entry] : memory_region_map_) {
         int ret = ibv_dereg_mr(entry.mr);
         if (ret) {
             PLOG(ERROR) << "Failed to unregister memory region";
         }
     }
-    memory_region_list_.clear();
+    memory_region_map_.clear();
 
     for (size_t i = 0; i < cq_list_.size(); ++i) {
         if (!cq_list_[i].native) continue;
@@ -217,7 +351,7 @@ int RdmaContext::registerMemoryRegionInternal(void *addr, size_t length,
                       << "shrink it to " << globalConfig().max_mr_size;
         length = (size_t)globalConfig().max_mr_size;
     }
-#if defined(USE_MLU) || (!defined(WITH_NVIDIA_PEERMEM) && defined(USE_CUDA))
+#if defined(USE_MLU) || defined(USE_MACA) || defined(USE_CUDA)
     // Implement register memory in a way that does not assume the presence of
     // nvidia-peermem. If memory is on CPU call ibv_reg_mr() as usual. If memory
     // is on GPU then use ibv_reg_dmabuf_mr() instead which does not require
@@ -230,32 +364,154 @@ int RdmaContext::registerMemoryRegionInternal(void *addr, size_t length,
     if (result != CUDA_SUCCESS || memType == CU_MEMORYTYPE_HOST) {
         mrMeta.addr = addr;
         mrMeta.mr = ibv_reg_mr(pd_, addr, length, access);
+#if defined(USE_CUDA)
+    } else if (memType == CU_MEMORYTYPE_DEVICE &&
+               Environ::Get().GetWithNvidiaPeermem()) {
+        // WITH_NVIDIA_PEERMEM env var is set: use ibv_reg_mr() directly for
+        // GPU memory (requires the nvidia-peermem kernel module to be loaded).
+        mrMeta.addr = addr;
+        mrMeta.mr = ibv_reg_mr(pd_, addr, length, access);
+#endif
     } else if (memType == CU_MEMORYTYPE_DEVICE) {
+#if defined(USE_CUDA)
+        // Ensure a CUDA context is current — worker threads or callers
+        // from non-CUDA threads may lack one.
+        unsigned int devOrd = 0;
+        cuPointerGetAttribute(&devOrd, CU_POINTER_ATTRIBUTE_DEVICE_ORDINAL,
+                              (CUdeviceptr)addr);
+        CUdevice cuDev;
+        CUcontext cuCtx;
+        cuDeviceGet(&cuDev, devOrd);
+        cuDevicePrimaryCtxRetain(&cuCtx, cuDev);
+        cuCtxSetCurrent(cuCtx);
+
+        // Use cuMemGetAddressRange to get the true allocation base and
+        // size — addr may sit at an offset within a larger cudaMalloc
+        // block (e.g. PyTorch caching allocator packs multiple tensors
+        // into one allocation).  cuMemGetHandleForAddressRange requires
+        // the exact allocation boundaries.
+#endif
+        CUdeviceptr allocBase;
         size_t allocSize;
-        result = cuPointerGetAttribute(
-            &allocSize, CU_POINTER_ATTRIBUTE_RANGE_SIZE, (CUdeviceptr)addr);
+        result =
+            cuMemGetAddressRange(&allocBase, &allocSize, (CUdeviceptr)addr);
         if (result != CUDA_SUCCESS) {
             const char *errStr;
             cuGetErrorString(result, &errStr);
-            LOG(ERROR) << "Failed to call cuPointerGetAttribute for "
+            LOG(ERROR) << "Failed to call cuMemGetAddressRange for "
                        << (uintptr_t)addr << " cuda error=" << errStr;
+#if defined(USE_CUDA)
+            cuDevicePrimaryCtxRelease(cuDev);
+#endif
             return ERR_CONTEXT;
         }
 
         int dmabuf_fd;
         result = cuMemGetHandleForAddressRange(
-            &dmabuf_fd, (CUdeviceptr)addr, allocSize,
+            &dmabuf_fd, allocBase, allocSize,
             CU_MEM_RANGE_HANDLE_TYPE_DMA_BUF_FD, 0);
         if (result != CUDA_SUCCESS) {
             const char *errStr;
             cuGetErrorString(result, &errStr);
             LOG(ERROR) << "Failed to retrieve dmabuf for " << (uintptr_t)addr
-                       << " cuda error=" << errStr;
+                       << " base=" << (uintptr_t)allocBase
+                       << " size=" << allocSize << " cuda error=" << errStr;
+#if defined(USE_CUDA)
+            cuDevicePrimaryCtxRelease(cuDev);
+#endif
             return ERR_CONTEXT;
         }
         mrMeta.addr = addr;
-        mrMeta.mr = ibv_reg_dmabuf_mr(pd_, 0 /* offset */, length,
+        uint64_t dmabuf_offset = (uintptr_t)addr - (uintptr_t)allocBase;
+        mrMeta.mr = ibv_reg_dmabuf_mr(pd_, dmabuf_offset, length,
                                       (uintptr_t)addr, dmabuf_fd, access);
+        const int regErrno = errno;
+        if (close(dmabuf_fd) != 0) {
+            PLOG(WARNING) << "Failed to close dmabuf fd";
+        }
+        if (!mrMeta.mr) {
+            errno = regErrno;
+        }
+#if defined(USE_CUDA)
+        cuDevicePrimaryCtxRelease(cuDev);
+#endif
+    }
+#elif defined(USE_HIP_DMABUF)
+    hipPointerAttribute_t hipAttr{};
+    hipError_t hipRes = hipPointerGetAttributes(&hipAttr, addr);
+
+    if (hipRes != hipSuccess || hipAttr.type == hipMemoryTypeHost ||
+        hipAttr.type == hipMemoryTypeUnregistered) {
+        // Host memory — standard ibv_reg_mr() path.
+        mrMeta.addr = addr;
+        mrMeta.mr = ibv_reg_mr(pd_, addr, length, access);
+    } else if (hipAttr.type == hipMemoryTypeManaged) {
+        // Managed (unified) memory pages can migrate between host and device;
+        // hsa_amd_portable_export_dmabuf captures the device-side handle at
+        // export time only, making the dmabuf fd stale after migration. Fall
+        // back to ibv_reg_mr() for safety.
+        LOG(WARNING) << "HIP managed memory at " << (uintptr_t)addr
+                     << " — dmabuf export skipped (pages may migrate); "
+                        "falling back to ibv_reg_mr";
+        mrMeta.addr = addr;
+        mrMeta.mr = ibv_reg_mr(pd_, addr, length, access);
+    } else if (hipAttr.type == hipMemoryTypeDevice &&
+               !isKernelDmabufSupported()) {
+        // Kernel lacks CONFIG_PCI_P2PDMA / CONFIG_DMABUF_MOVE_NOTIFY —
+        // ibv_reg_dmabuf_mr may succeed but transfers will silently fail.
+        // Fail at registration time instead.
+        mrMeta.addr = addr;
+        mrMeta.mr = ibv_reg_mr(pd_, addr, length, access);
+    } else if (hipAttr.type == hipMemoryTypeDevice) {
+        // Pin to the owning device while exporting the dmabuf fd.
+        HipDeviceGuard dev_guard(hipAttr.device);
+        if (!dev_guard.set_ok()) {
+            LOG(ERROR) << "Failed to set HIP device to " << hipAttr.device
+                       << " for dmabuf export of " << (uintptr_t)addr;
+            return ERR_CONTEXT;
+        }
+
+        // Get the allocation base + size, since `addr` may sit at an offset
+        // within a larger hipMalloc block (caching allocators pack tensors).
+        hipDeviceptr_t allocBase = nullptr;
+        size_t allocSize = 0;
+        hipRes = hipMemGetAddressRange(&allocBase, &allocSize,
+                                       reinterpret_cast<hipDeviceptr_t>(addr));
+        if (hipRes != hipSuccess) {
+            LOG(ERROR) << "Failed to call hipMemGetAddressRange for "
+                       << (uintptr_t)addr
+                       << " hip error=" << hipGetErrorString(hipRes);
+            return ERR_CONTEXT;
+        }
+
+        int dmabuf_fd = -1;
+        uint64_t hsa_dmabuf_offset = 0;
+        hsa_status_t hsaRes = hsa_amd_portable_export_dmabuf(
+            allocBase, allocSize, &dmabuf_fd, &hsa_dmabuf_offset);
+        if (hsaRes != HSA_STATUS_SUCCESS) {
+            const char *hsaErr = nullptr;
+            hsa_status_string(hsaRes, &hsaErr);
+            LOG(ERROR) << "Failed to retrieve dmabuf for " << (uintptr_t)addr
+                       << " base=" << (uintptr_t)allocBase
+                       << " size=" << allocSize
+                       << " hsa error=" << (hsaErr ? hsaErr : "unknown");
+            return ERR_CONTEXT;
+        }
+
+        mrMeta.addr = addr;
+        // Offset within the dmabuf-backed region: distance from the
+        // allocation base, plus any offset hsa returned for the export.
+        uint64_t reg_offset =
+            (uintptr_t)addr - (uintptr_t)allocBase + hsa_dmabuf_offset;
+        mrMeta.mr = ibv_reg_dmabuf_mr(pd_, reg_offset, length, (uintptr_t)addr,
+                                      dmabuf_fd, access);
+        const int regErrno = errno;
+        if (close(dmabuf_fd) != 0) {
+            PLOG(WARNING) << "Failed to close dmabuf fd";
+        }
+        if (!mrMeta.mr) {
+            errno = regErrno;
+        }
     }
 #else
     mrMeta.addr = addr;
@@ -275,29 +531,21 @@ int RdmaContext::registerMemoryRegion(void *addr, size_t length, int access) {
         return ret;
     }
     RWSpinlock::WriteGuard guard(memory_regions_lock_);
-    memory_region_list_.push_back(mrMeta);
+    memory_region_map_[reinterpret_cast<uintptr_t>(mrMeta.addr)] = mrMeta;
     return 0;
 }
 
 int RdmaContext::unregisterMemoryRegion(void *addr) {
     RWSpinlock::WriteGuard guard(memory_regions_lock_);
-    bool has_removed;
-    do {
-        has_removed = false;
-        for (auto iter = memory_region_list_.begin();
-             iter != memory_region_list_.end(); ++iter) {
-            if (iter->addr <= addr &&
-                addr < (char *)(iter->addr) + iter->mr->length) {
-                if (ibv_dereg_mr(iter->mr)) {
-                    LOG(ERROR) << "Failed to unregister memory " << addr;
-                    return ERR_CONTEXT;
-                }
-                memory_region_list_.erase(iter);
-                has_removed = true;
-                break;
-            }
-        }
-    } while (has_removed);
+    auto iter = findMemoryRegionContaining(reinterpret_cast<uintptr_t>(addr));
+    if (iter == memory_region_map_.end()) {
+        return 0;
+    }
+    if (ibv_dereg_mr(iter->second.mr)) {
+        LOG(ERROR) << "Failed to unregister memory " << addr;
+        return ERR_CONTEXT;
+    }
+    memory_region_map_.erase(iter);
     return 0;
 }
 
@@ -313,11 +561,8 @@ int RdmaContext::preTouchMemory(void *addr, size_t length) {
 
 uint32_t RdmaContext::rkey(void *addr) {
     RWSpinlock::ReadGuard guard(memory_regions_lock_);
-    for (auto iter = memory_region_list_.begin();
-         iter != memory_region_list_.end(); ++iter)
-        if (iter->addr <= addr &&
-            addr < (char *)(iter->addr) + iter->mr->length)
-            return iter->mr->rkey;
+    auto iter = findMemoryRegionContaining(reinterpret_cast<uintptr_t>(addr));
+    if (iter != memory_region_map_.end()) return iter->second.mr->rkey;
 
     LOG(ERROR) << "Address " << addr << " rkey not found for " << deviceName();
     return 0;
@@ -325,16 +570,35 @@ uint32_t RdmaContext::rkey(void *addr) {
 
 uint32_t RdmaContext::lkey(void *addr) {
     RWSpinlock::ReadGuard guard(memory_regions_lock_);
-    for (auto iter = memory_region_list_.begin();
-         iter != memory_region_list_.end(); ++iter)
-        if (iter->addr <= addr &&
-            addr < (char *)(iter->addr) + iter->mr->length)
-            return iter->mr->lkey;
+    auto iter = findMemoryRegionContaining(reinterpret_cast<uintptr_t>(addr));
+    if (iter != memory_region_map_.end()) return iter->second.mr->lkey;
 
     LOG(ERROR) << "Address " << addr << " lkey not found for " << deviceName();
     return 0;
 }
 
+RdmaContext::MemoryRegionMap::iterator RdmaContext::findMemoryRegionContaining(
+    uintptr_t addr) {
+    auto upper = memory_region_map_.upper_bound(addr);
+    if (upper == memory_region_map_.begin()) {
+        return memory_region_map_.end();
+    }
+    auto candidate = std::prev(upper);
+    return containsAddress(candidate->second, addr) ? candidate
+                                                    : memory_region_map_.end();
+}
+
+RdmaContext::MemoryRegionMap::const_iterator
+RdmaContext::findMemoryRegionContaining(uintptr_t addr) const {
+    auto upper = memory_region_map_.upper_bound(addr);
+    if (upper == memory_region_map_.begin()) {
+        return memory_region_map_.end();
+    }
+    auto candidate = std::prev(upper);
+    return containsAddress(candidate->second, addr) ? candidate
+                                                    : memory_region_map_.end();
+}
+
 std::shared_ptr<RdmaEndPoint> RdmaContext::endpoint(
     const std::string &peer_nic_path) {
     if (!active_) {
@@ -370,24 +634,38 @@ int RdmaContext::deleteEndpoint(const std::string &peer_nic_path) {
     return endpoint_store_->deleteEndpoint(peer_nic_path);
 }
 
+int RdmaContext::deleteEndpointByPtr(const RdmaEndPoint *endpoint_ptr) {
+    return endpoint_store_->deleteEndpointByPtr(endpoint_ptr);
+}
+
+void RdmaContext::reclaimEndpoints() { endpoint_store_->reclaimEndpoint(); }
+
+size_t RdmaContext::waitingListSize() const {
+    return endpoint_store_->waitingListSize();
+}
+
+void RdmaContext::testOnlyInsertWaiting(std::shared_ptr<RdmaEndPoint> ep) {
+    endpoint_store_->testOnlyInsertWaiting(std::move(ep));
+}
+
 size_t RdmaContext::getTotalQPNumber() const {
     return endpoint_store_->getTotalQPNumber();
 }
 
 std::string RdmaContext::nicPath() const {
-    return MakeNicPath(engine_.local_server_name_, device_name_);
+    return MakeNicPath(engine_.rdma_server_name_, device_name_);
 }
 
-std::string RdmaContext::gid() const {
-    std::string gid_str;
-    char buf[16] = {0};
-    const static size_t kGidLength = 16;
-    for (size_t i = 0; i < kGidLength; ++i) {
-        sprintf(buf, "%02x", gid_.raw[i]);
-        gid_str += i == 0 ? buf : std::string(":") + buf;
-    }
+std::string RdmaContext::gid() const { return gidSelection().gid; }
 
-    return gid_str;
+GidSelectionSnapshot RdmaContext::gidSelection() const {
+    std::lock_guard<std::mutex> guard(gid_lock_);
+    return {gidBytesToString(gid_.raw), gid_index_};
+}
+
+int RdmaContext::gidIndex() const {
+    std::lock_guard<std::mutex> guard(gid_lock_);
+    return gid_index_;
 }
 
 ibv_cq *RdmaContext::cq() {
@@ -428,16 +706,53 @@ static std::string readGidNdev(const std::string &device_name, uint8_t port,
     return ndev;
 }
 
-// Returns 1 if the GID has an associated network device, 0 otherwise.
-static int hasNetworkDevice(const std::string &device_name, uint8_t port,
-                            int gid_index) {
-    return !readGidNdev(device_name, port, gid_index).empty() ? 1 : 0;
+static inline bool isOverlayNetwork(const std::string &ndev) {
+    return ndev.find("flannel") == 0 || ndev.find("cni") == 0 ||
+           ndev.find("calico") == 0 || ndev.find("vxlan") == 0 ||
+           ndev.find("docker") == 0 || ndev == "tunl0";
+}
+
+static inline bool isOverlayIPv4(const struct in6_addr *addr) {
+    if (!ipv6_addr_v4mapped(addr)) return false;
+
+    uint32_t ipv4 = ntohl(addr->s6_addr32[3]);
+    uint8_t octet1 = (ipv4 >> 24) & 0xFF;
+    uint8_t octet2 = (ipv4 >> 16) & 0xFF;
+
+    if (octet1 == 10) return true;
+    if (octet1 == 172 && octet2 >= 16 && octet2 <= 31) return true;
+    if (octet1 == 100 && octet2 >= 64 && octet2 <= 127) return true;
+    return false;
+}
+
+static inline bool isLinkLocalIpv6(const struct in6_addr *addr) {
+    return IN6_IS_ADDR_LINKLOCAL(addr);
 }
 
 static const char *GidNetworkStateToString(GidNetworkState state) {
-    return (state == GidNetworkState::GID_WITH_NETWORK)
-               ? "with network device"
-               : "without network device";
+    switch (state) {
+        case GidNetworkState::GID_WITH_NETWORK:
+            return "with network device";
+        case GidNetworkState::GID_WITHOUT_NETWORK:
+            return "without network device";
+        case GidNetworkState::GID_NOT_FOUND:
+            return "not found";
+    }
+    return "unknown";
+}
+
+static GidNetworkState autoGidStateFromSelection(
+    const AutoGidSelection &selection) {
+    switch (selection.candidate_class) {
+        case AutoGidCandidateClass::kNetworkRoutable:
+        case AutoGidCandidateClass::kNetworkDegraded:
+            return GidNetworkState::GID_WITH_NETWORK;
+        case AutoGidCandidateClass::kNoNetworkRoutable:
+        case AutoGidCandidateClass::kNoNetworkDegraded:
+        case AutoGidCandidateClass::kFallbackNonzero:
+            return GidNetworkState::GID_WITHOUT_NETWORK;
+    }
+    return GidNetworkState::GID_NOT_FOUND;
 }
 
 GidNetworkState RdmaContext::findBestGidIndex(const std::string &device_name,
@@ -445,67 +760,167 @@ GidNetworkState RdmaContext::findBestGidIndex(const std::string &device_name,
                                               ibv_port_attr &port_attr,
                                               uint8_t port, int &gid_index) {
     gid_index = -1;
-    int i;
-    struct ibv_gid_entry gid_entry;
-    int fallback_ipv4_gid_without_network = -1;
-    int fallback_ipv6_gid_with_network = -1;
-    int fallback_ipv6_gid_without_network = -1;
-    GidNetworkState state = GidNetworkState::GID_NOT_FOUND;
+    std::vector<AutoGidCandidate> candidates;
+    candidates.reserve(port_attr.gid_tbl_len);
 
-    for (i = 0; i < port_attr.gid_tbl_len; i++) {
+    for (int i = 0; i < port_attr.gid_tbl_len; i++) {
+        AutoGidCandidate candidate;
+        candidate.gid_index = i;
+
+        struct ibv_gid_entry gid_entry;
         if (ibv_query_gid_ex(context, port, i, &gid_entry, 0)) {
-            // Reached end of valid GID indices
-            break;
-        }
-
-        if (gid_entry.gid_type != IBV_GID_TYPE_ROCE_V2 &&
-            gid_entry.gid_type != IBV_GID_TYPE_IB) {
+            candidate.query_succeeded = false;
+            candidates.push_back(candidate);
             continue;
         }
 
-        const bool is_ipv4_gid =
-            gid_entry.gid_type == IBV_GID_TYPE_ROCE_V2 &&
-            ipv6_addr_v4mapped((struct in6_addr *)gid_entry.gid.raw);
-        const bool has_network_device = hasNetworkDevice(device_name, port, i);
+        const auto *gid_addr =
+            reinterpret_cast<const struct in6_addr *>(gid_entry.gid.raw);
+        std::string ndev = readGidNdev(device_name, port, i);
+        candidate.gid = gidBytesToString(gid_entry.gid.raw);
+        candidate.gid_type = gid_entry.gid_type;
+        candidate.has_network_device = !ndev.empty();
+        candidate.is_ipv4_mapped = ipv6_addr_v4mapped(gid_addr);
+        candidate.is_link_local_ipv6 = isLinkLocalIpv6(gid_addr);
+        candidate.is_overlay_network =
+            candidate.has_network_device && isOverlayNetwork(ndev);
+        candidate.is_overlay_ipv4 =
+            candidate.is_ipv4_mapped && isOverlayIPv4(gid_addr);
+        candidate.is_null_gid = isNullGid(&gid_entry.gid);
+        candidates.push_back(candidate);
+    }
 
-        if (is_ipv4_gid) {
-            if (has_network_device) {
-                gid_index = i;
-                return GidNetworkState::GID_WITH_NETWORK;
-            }
-            if (fallback_ipv4_gid_without_network < 0) {
-                gid_index = i;
-                fallback_ipv4_gid_without_network = i;
-                state = GidNetworkState::GID_WITHOUT_NETWORK;
-            }
+    auto selection = selectBestAutoGidCandidate(candidates);
+    if (!selection.has_value()) {
+        return GidNetworkState::GID_NOT_FOUND;
+    }
+
+    gid_index = selection->gid_index;
+    VLOG(1) << "Selected auto GID[" << gid_index << "] on " << device_name
+            << " with class "
+            << autoGidCandidateClassToString(selection->candidate_class);
+    return autoGidStateFromSelection(*selection);
+}
+
+bool RdmaContext::reprobeAutoGid(
+    const GidSelectionSnapshot &expected_selection,
+    const std::vector<AutoGidSelectionIdentity> &tried_selections,
+    std::string *previous_gid, std::string *next_gid) {
+    std::lock_guard<std::mutex> reprobe_guard(gid_reprobe_lock_);
+    std::string current_gid_string;
+    std::string next_gid_string;
+    int current_gid_index = -1;
+    int next_gid_index = -1;
+    uint16_t current_lid = 0;
+    ibv_context *current_context = nullptr;
+    uint8_t current_port = 0;
+    AutoGidCandidateClass next_candidate_class =
+        AutoGidCandidateClass::kFallbackNonzero;
+    {
+        std::lock_guard<std::mutex> guard(gid_lock_);
+        if (!auto_gid_selection_enabled_ || !context_) {
+            return false;
+        }
+
+        current_gid_index = gid_index_;
+        current_gid_string = gidBytesToString(gid_.raw);
+        current_lid = lid_;
+        current_context = context_;
+        current_port = port_;
+        if (current_gid_index != expected_selection.gid_index ||
+            current_gid_string != expected_selection.gid) {
+            return false;
+        }
+    }
+
+    ibv_port_attr port_attr;
+    if (ibv_query_port(current_context, current_port, &port_attr)) {
+        PLOG(WARNING) << "Failed to reprobe port attributes on " << device_name_
+                      << "/" << static_cast<int>(current_port);
+        return false;
+    }
+
+    std::vector<AutoGidCandidate> candidates;
+    candidates.reserve(port_attr.gid_tbl_len);
+    for (int i = 0; i < port_attr.gid_tbl_len; ++i) {
+        AutoGidCandidate candidate;
+        candidate.gid_index = i;
+
+        struct ibv_gid_entry gid_entry;
+        if (ibv_query_gid_ex(current_context, current_port, i, &gid_entry, 0)) {
+            candidate.query_succeeded = false;
+            candidates.push_back(candidate);
             continue;
         }
 
-        if (has_network_device && fallback_ipv6_gid_with_network < 0) {
-            fallback_ipv6_gid_with_network = i;
-        }
+        const auto *gid_addr =
+            reinterpret_cast<const struct in6_addr *>(gid_entry.gid.raw);
+        std::string ndev = readGidNdev(device_name_, current_port, i);
+        candidate.gid = gidBytesToString(gid_entry.gid.raw);
+        candidate.gid_type = gid_entry.gid_type;
+        candidate.has_network_device = !ndev.empty();
+        candidate.is_ipv4_mapped = ipv6_addr_v4mapped(gid_addr);
+        candidate.is_link_local_ipv6 = isLinkLocalIpv6(gid_addr);
+        candidate.is_overlay_network =
+            candidate.has_network_device && isOverlayNetwork(ndev);
+        candidate.is_overlay_ipv4 =
+            candidate.is_ipv4_mapped && isOverlayIPv4(gid_addr);
+        candidate.is_null_gid = isNullGid(&gid_entry.gid);
+        candidates.push_back(candidate);
+    }
 
-        if (!has_network_device && fallback_ipv6_gid_without_network < 0) {
-            fallback_ipv6_gid_without_network = i;
+    auto selection = reselectAutoGidCandidate(
+        candidates, current_gid_index, current_gid_string, tried_selections);
+    if (!selection.has_value()) {
+        if (next_gid) {
+            *next_gid = current_gid_string;
+        }
+        return false;
+    }
+
+    ibv_gid new_gid = {};
+    if (ibv_query_gid(current_context, current_port, selection->gid_index,
+                      &new_gid)) {
+        return false;
+    }
+    if (isNullGid(&new_gid)) {
+        return false;
+    }
+
+    if (previous_gid) {
+        *previous_gid = current_gid_string;
+    }
+    next_gid_string = gidBytesToString(new_gid.raw);
+
+    int publish_ret = engine_.refreshLocalDeviceDesc(device_name_, current_lid,
+                                                     next_gid_string);
+    if (publish_ret) {
+        LOG(ERROR) << "Failed to refresh local device descriptor for "
+                   << device_name_ << ": " << publish_ret;
+        return false;
+    }
+
+    {
+        std::lock_guard<std::mutex> guard(gid_lock_);
+        gid_ = new_gid;
+        gid_index_ = selection->gid_index;
+        next_gid_index = selection->gid_index;
+        next_candidate_class = selection->candidate_class;
+        if (next_gid) {
+            *next_gid = next_gid_string;
         }
     }
 
-    if (fallback_ipv4_gid_without_network >= 0) {
-        gid_index = fallback_ipv4_gid_without_network;
-        return GidNetworkState::GID_WITHOUT_NETWORK;
+    if (next_gid_string.empty()) {
+        return false;
     }
 
-    if (fallback_ipv6_gid_with_network >= 0) {
-        gid_index = fallback_ipv6_gid_with_network;
-        return GidNetworkState::GID_WITH_NETWORK;
-    }
-
-    if (fallback_ipv6_gid_without_network >= 0) {
-        gid_index = fallback_ipv6_gid_without_network;
-        return GidNetworkState::GID_WITHOUT_NETWORK;
-    }
-
-    return state;
+    LOG(WARNING) << "Auto GID reprobe switched " << device_name_ << "/"
+                 << static_cast<int>(port_) << " from index "
+                 << current_gid_index << " (" << current_gid_string << ") to "
+                 << next_gid_index << " (" << next_gid_string << "), class "
+                 << autoGidCandidateClassToString(next_candidate_class);
+    return true;
 }
 
 int RdmaContext::openRdmaDevice(const std::string &device_name, uint8_t port,
@@ -565,27 +980,92 @@ int RdmaContext::openRdmaDevice(const std::string &device_name, uint8_t port,
             return ERR_CONTEXT;
         }
 
-#if !defined(WITH_NVIDIA_PEERMEM) && defined(USE_CUDA)
-        // Verify dmabuf support which is required if not using nvidia-peermem.
-        // Assume device index matches.
-        CUdevice cuDevice;
-        CUresult result = cuDeviceGet(&cuDevice, i);
-        if (result != CUDA_SUCCESS) {
-            LOG(ERROR) << "Failed to query CUDA device";
-            return ERR_CONTEXT;
-        }
-        int dmaBufSupported;
-        result = cuDeviceGetAttribute(
-            &dmaBufSupported, CU_DEVICE_ATTRIBUTE_DMA_BUF_SUPPORTED, cuDevice);
-        if (result != CUDA_SUCCESS) {
-            LOG(ERROR) << "Failed to query CUDA device attributes";
-            return ERR_CONTEXT;
-        }
-        if (!dmaBufSupported) {
-            LOG(ERROR) << "DMA BUF supported required for GPU RDMA without "
-                          "nvidia-peermem";
-            return ERR_CONTEXT;
-        }
+#if defined(USE_MACA) || defined(USE_CUDA)
+        // Verify DMA-BUF support against the GPU device(s) that the local
+        // topology explicitly maps to this RNIC, rather than assuming the
+        // verbs enumeration order matches GPU enumeration.
+        // Validate DMA-BUF support for every GPU that can reach this RNIC,
+        // not just GPUs listing it as preferred.  Runtime selection falls
+        // back to avail_hca when a preferred NIC is disabled, so we must
+        // validate both lists.
+        if (!Environ::Get().GetWithNvidiaPeermem()) {
+            std::vector<int> mapped_gpu_devices;
+            if (engine_.local_topology_) {
+                const auto topology_matrix =
+                    engine_.local_topology_->getMatrix();
+                for (const auto &entry : topology_matrix) {
+                    if (entry.first.rfind(GPU_PREFIX, 0) != 0) continue;
+                    bool in_preferred =
+                        std::find(entry.second.preferred_hca.begin(),
+                                  entry.second.preferred_hca.end(),
+                                  device_name) !=
+                        entry.second.preferred_hca.end();
+                    bool in_avail =
+                        std::find(entry.second.avail_hca.begin(),
+                                  entry.second.avail_hca.end(),
+                                  device_name) != entry.second.avail_hca.end();
+                    if (!in_preferred && !in_avail) continue;
+
+                    try {
+                        mapped_gpu_devices.push_back(
+                            std::stoi(entry.first.substr(GPU_PREFIX.size())));
+                    } catch (const std::exception &e) {
+                        LOG(WARNING) << "Ignore malformed topology GPU entry "
+                                     << entry.first << ": " << e.what();
+                    }
+                }
+            }
+
+            std::sort(mapped_gpu_devices.begin(), mapped_gpu_devices.end());
+            mapped_gpu_devices.erase(std::unique(mapped_gpu_devices.begin(),
+                                                 mapped_gpu_devices.end()),
+                                     mapped_gpu_devices.end());
+
+            if (mapped_gpu_devices.empty()) {
+                LOG(INFO) << "No GPU device is explicitly mapped to RNIC "
+                          << device_name
+                          << "; skip DMA-BUF affinity validation";
+            } else {
+                // cuInit is process-global and idempotent; call it once before
+                // the per-device loop, not per cuDeviceGet.
+#if defined(USE_CUDA)
+                CUresult result = cuInit(0);
+                if (result != CUDA_SUCCESS) {
+                    LOG(ERROR) << "Failed to initialize CUDA driver for RNIC "
+                               << device_name;
+                    goto cleanup_context_and_devices;
+                }
+#endif
+                for (int gpu_device : mapped_gpu_devices) {
+                    CUdevice cuDevice;
+                    CUresult result = cuDeviceGet(&cuDevice, gpu_device);
+                    if (result != CUDA_SUCCESS) {
+                        LOG(ERROR) << "Failed to query GPU device "
+                                   << gpu_device << " for RNIC " << device_name;
+                        goto cleanup_context_and_devices;
+                    }
+                    int dmaBufSupported;
+                    result = cuDeviceGetAttribute(
+                        &dmaBufSupported, CU_DEVICE_ATTRIBUTE_DMA_BUF_SUPPORTED,
+                        cuDevice);
+                    if (result != CUDA_SUCCESS) {
+                        LOG(ERROR)
+                            << "Failed to query GPU device attributes for "
+                            << "GPU device " << gpu_device << " and RNIC "
+                            << device_name;
+                        goto cleanup_context_and_devices;
+                    }
+                    if (!dmaBufSupported) {
+                        LOG(ERROR)
+                            << "DMA BUF supported required for GPU RDMA "
+                               "without "
+                               "nvidia-peermem on GPU device "
+                            << gpu_device << " mapped to RNIC " << device_name;
+                        goto cleanup_context_and_devices;
+                    }
+                }
+            }
+        }  // !Environ::Get().GetWithNvidiaPeermem()
 #endif
 
         ibv_port_attr port_attr;
@@ -602,6 +1082,7 @@ int RdmaContext::openRdmaDevice(const std::string &device_name, uint8_t port,
 
         updateGlobalConfig(device_attr);
         GidNetworkState gid_state;
+        auto_gid_selection_enabled_ = gid_index < 0;
         if (gid_index < 0) {
             int found_gid_index = -1;
             gid_state = findBestGidIndex(device_name, context, port_attr, port,
@@ -619,7 +1100,7 @@ int RdmaContext::openRdmaDevice(const std::string &device_name, uint8_t port,
             }
         } else {
             // Also check network state for user-specified GID
-            bool has_ndev = hasNetworkDevice(device_name, port, gid_index);
+            bool has_ndev = !readGidNdev(device_name, port, gid_index).empty();
             if (!has_ndev) {
                 LOG(WARNING) << "User-specified GID index " << gid_index
                              << " on " << device_name << "/" << port
@@ -653,7 +1134,10 @@ int RdmaContext::openRdmaDevice(const std::string &device_name, uint8_t port,
         lid_ = attr.lid;
         active_mtu_ = attr.active_mtu;
         active_speed_ = attr.active_speed;
-        gid_index_ = gid_index;
+        {
+            std::lock_guard<std::mutex> guard(gid_lock_);
+            gid_index_ = gid_index;
+        }
 
         ibv_free_device_list(devices);
         return 0;
diff --git a/mooncake-transfer-engine/src/transport/rdma_transport/rdma_endpoint.cpp b/mooncake-transfer-engine/src/transport/rdma_transport/rdma_endpoint.cpp
index 104c0322..f5ada5d3 100644
--- a/mooncake-transfer-engine/src/transport/rdma_transport/rdma_endpoint.cpp
+++ b/mooncake-transfer-engine/src/transport/rdma_transport/rdma_endpoint.cpp
@@ -17,26 +17,80 @@
 #include <glog/logging.h>
 
 #include <cassert>
+#include <cerrno>
 #include <cstddef>
 #include <chrono>
 #include <thread>
+#include <sstream>
+
+#ifdef USE_MLX5DV
+#include <infiniband/mlx5dv.h>
+#endif
 
 #include "common.h"
 #include "config.h"
+#include "transport/rdma_transport/rdma_gid_probe.h"
 
 namespace mooncake {
-const static uint8_t MAX_HOP_LIMIT = 16;
-const static uint8_t TIMEOUT = 14;
-const static uint8_t RETRY_CNT = 7;
+constexpr uint8_t kMaxHopLimit = 16;
+constexpr uint8_t kTimeout = 14;
+constexpr uint8_t kRetryCount = 7;
+
+static GidSelectionSnapshot fillLocalHandshakeDesc(
+    RdmaContext &context, const std::string &peer_nic,
+    const std::vector<uint32_t> &qp_num,
+    RdmaEndPoint::HandShakeDesc &local_desc) {
+    auto gid_selection = context.gidSelection();
+    local_desc.local_nic_path = context.nicPath();
+    local_desc.local_lid = context.lid();
+    local_desc.local_gid = gid_selection.gid;
+    local_desc.peer_nic_path = peer_nic;
+    local_desc.qp_num = qp_num;
+    local_desc.reply_msg.clear();
+    return gid_selection;
+}
+
+static void rememberAutoGidSelection(
+    std::vector<AutoGidSelectionIdentity> &attempted_selections,
+    const GidSelectionSnapshot &selection) {
+    auto already_attempted =
+        std::any_of(attempted_selections.begin(), attempted_selections.end(),
+                    [&](const AutoGidSelectionIdentity &attempted) {
+                        return attempted.gid_index == selection.gid_index &&
+                               attempted.gid == selection.gid;
+                    });
+    if (!already_attempted) {
+        attempted_selections.push_back({selection.gid_index, selection.gid});
+    }
+}
+
+static std::string qpListToString(const std::vector<uint32_t> &qp_num) {
+    std::ostringstream oss;
+    oss << "[";
+    for (size_t i = 0; i < qp_num.size(); ++i) {
+        if (i) oss << ",";
+        oss << qp_num[i];
+    }
+    oss << "]";
+    return oss.str();
+}
 
 RdmaEndPoint::RdmaEndPoint(RdmaContext &context)
     : context_(context),
       status_(INITIALIZING),
+      has_connected_(false),
+      wr_depth_list_(nullptr),
       active_(true),
       cq_outstanding_(nullptr) {}
 
 RdmaEndPoint::~RdmaEndPoint() {
-    if (!qp_list_.empty()) deconstruct();
+    if (!qp_list_.empty()) {
+        // In normal flow, beginDestroy()+finishDestroy() should have been
+        // called already via endpoint_store. This is a fallback for abnormal
+        // shutdown (e.g., process exit).
+        RWSpinlock::WriteGuard guard(lock_);
+        deconstructLocked();
+    }
 }
 
 int RdmaEndPoint::construct(ibv_cq *cq, size_t num_qp_list,
@@ -54,7 +108,7 @@ int RdmaEndPoint::construct(ibv_cq *cq, size_t num_qp_list,
     max_sge_per_wr_ = max_sge_per_wr;
     max_inline_bytes_ = max_inline_bytes;
 
-    wr_depth_list_ = new volatile int[num_qp_list];
+    wr_depth_list_ = new volatile int[num_qp_list]();
     if (!wr_depth_list_) {
         LOG(ERROR) << "Failed to allocate memory for work request depth list";
         return ERR_MEMORY;
@@ -90,7 +144,8 @@ int RdmaEndPoint::reconstruct() {
     auto max_inline_bytes = max_inline_bytes_;
 
     // Deconstruct and reconstruct to get fresh QPs (same as delete+create)
-    int ret = deconstruct();
+    // Use deconstructLocked because callers already hold lock_
+    int ret = deconstructLocked();
     if (ret) {
         LOG(ERROR) << "Failed to deconstruct endpoint: " << ret;
         return ret;
@@ -112,36 +167,146 @@ int RdmaEndPoint::reconstruct() {
 }
 
 int RdmaEndPoint::deconstruct() {
-    for (size_t i = 0; i < qp_list_.size(); ++i) {
-        if (ibv_destroy_qp(qp_list_[i])) {
-            PLOG(ERROR) << "Failed to destroy QP";
-            return ERR_ENDPOINT;
-        }
-        // After destroying QP, the wr_depth_list_ won't change
-        bool displayed = false;
-        if (wr_depth_list_[i] != 0) {
-            if (!displayed) {
-                LOG(WARNING) << "Outstanding work requests found, CQ will not "
-                                "be generated";
-                displayed = true;
+    RWSpinlock::WriteGuard guard(lock_);
+    return deconstructLocked();
+}
+
+int RdmaEndPoint::deconstructLocked() {
+    // Adjust cq_outstanding_ before destroying QPs, so the counter is
+    // always corrected even if ibv_destroy_qp fails and we return early.
+    bool displayed = false;
+    if (wr_depth_list_) {
+        for (size_t i = 0; i < qp_list_.size(); ++i) {
+            if (wr_depth_list_[i] != 0) {
+                if (!displayed) {
+                    LOG(WARNING)
+                        << "Outstanding work requests found, CQ will not "
+                           "be generated";
+                    displayed = true;
+                }
+                __sync_fetch_and_sub(cq_outstanding_, wr_depth_list_[i]);
+                wr_depth_list_[i] = 0;
             }
-            __sync_fetch_and_sub(cq_outstanding_, wr_depth_list_[i]);
-            wr_depth_list_[i] = 0;
         }
     }
+
+    int result = 0;
+    for (size_t i = 0; i < qp_list_.size(); ++i) {
+        if (!qp_list_[i]) continue;  // already destroyed in a previous call
+        if (ibv_destroy_qp(qp_list_[i])) {
+            PLOG(ERROR) << "Failed to destroy QP[" << i << "]";
+            result = ERR_ENDPOINT;
+        } else {
+            qp_list_[i] = nullptr;
+        }
+    }
+
+    if (result) return result;
+
     qp_list_.clear();
     peer_qp_num_list_.clear();
     delete[] wr_depth_list_;
+    wr_depth_list_ = nullptr;
     return 0;
 }
 
 int RdmaEndPoint::destroyQP() { return deconstruct(); }
 
+void RdmaEndPoint::beginDestroy() {
+    RWSpinlock::WriteGuard guard(lock_);
+    beginDestroyLocked();
+}
+
+void RdmaEndPoint::beginDestroyLocked() {
+    auto current_status = status_.load(std::memory_order_relaxed);
+    if (current_status == DESTROYING || current_status == DESTROYED) return;
+
+    active_ = false;
+    inactive_time_ = getCurrentTimeInNano();
+    status_.store(DESTROYING, std::memory_order_release);
+
+    // Transition QPs to ERR state so hardware flushes all inflight WRs to CQ.
+    // This allows performPollCq to drain them naturally.
+    ibv_qp_attr attr;
+    memset(&attr, 0, sizeof(attr));
+    attr.qp_state = IBV_QPS_ERR;
+    for (size_t i = 0; i < qp_list_.size(); ++i) {
+        if (ibv_modify_qp(qp_list_[i], &attr, IBV_QP_STATE)) {
+            PLOG(WARNING) << "Failed to modify QP to ERR during beginDestroy";
+        }
+    }
+}
+
+bool RdmaEndPoint::finishDestroy() {
+    RWSpinlock::WriteGuard guard(lock_);
+    auto current_status = status_.load(std::memory_order_relaxed);
+
+    // Gate 1: already done.
+    if (current_status == DESTROYED) return true;
+
+    // Gate 2: non-two-phase path (status != DESTROYING). The endpoint
+    // reached waiting_list_ without going through beginDestroy(). This is
+    // the contract expected by EndpointStore::testOnlyInsertWaiting() and
+    // serves as a safety net for any future non-two-phase path. Mirror the
+    // pre-two-phase predicate (!hasOutstandingSlice == !active_): only
+    // inactive endpoints are eligible for reclaim; active ones must stay.
+    if (current_status != DESTROYING) {
+        if (active_) return false;
+        // Endpoints that never reached construct() own no RDMA resources
+        // and have wr_depth_list_ uninitialized; deconstructLocked() would
+        // delete[] a wild pointer. Drop them directly.
+        if (qp_list_.empty()) {
+            status_.store(DESTROYED, std::memory_order_relaxed);
+            return true;
+        }
+        LOG(WARNING) << "finishDestroy called in unexpected state: "
+                     << current_status
+                     << ", forcing destruction to avoid waiting_list_ leak";
+        // Fall through to the unified destroy path.
+    } else {
+        // Gate 3: two-phase path. Wait for inflight WRs to drain via CQ
+        // polling. If ibv_modify_qp-to-ERR failed in beginDestroy, WRs may
+        // never be flushed; enforce a timeout to avoid leaking forever.
+        bool has_outstanding = false;
+        for (size_t i = 0; i < qp_list_.size(); ++i) {
+            if (wr_depth_list_[i] != 0) {
+                has_outstanding = true;
+                break;
+            }
+        }
+        if (has_outstanding) {
+            double elapsed = (getCurrentTimeInNano() - inactive_time_) / 1e9;
+            if (elapsed < kFinishDestroyTimeoutSec) return false;
+            LOG(WARNING) << "finishDestroy timed out after " << elapsed
+                         << "s with outstanding WRs, forcing destruction";
+        }
+    }
+
+    // Unified destroy: tear down QPs (deconstructLocked handles
+    // cq_outstanding_ adjustment internally) and bound retries to avoid
+    // log flooding when ibv_destroy_qp fails permanently.
+    int ret = deconstructLocked();
+    if (ret) {
+        finish_destroy_retries_++;
+        LOG(ERROR) << "Failed to finish destroying endpoint (attempt "
+                   << finish_destroy_retries_ << "/" << kFinishDestroyMaxRetries
+                   << "): " << ret;
+        if (finish_destroy_retries_ < kFinishDestroyMaxRetries) return false;
+        LOG(ERROR) << "Giving up after " << finish_destroy_retries_
+                   << " retries (possible resource leak)";
+    }
+    status_.store(DESTROYED, std::memory_order_relaxed);
+    return true;
+}
+
 void RdmaEndPoint::setPeerNicPath(const std::string &peer_nic_path) {
     RWSpinlock::WriteGuard guard(lock_);
-    if (connected()) {
-        LOG(WARNING) << "Previous connection will be discarded";
-        disconnectUnlocked();
+    auto curr_status = status_.load(std::memory_order_relaxed);
+    if (curr_status != INITIALIZING && curr_status != UNCONNECTED) {
+        LOG(ERROR) << "Cannot change peer NIC path after endpoint lifecycle "
+                      "has started: "
+                   << toString();
+        return;
     }
     peer_nic_path_ = peer_nic_path;
 }
@@ -150,6 +315,8 @@ int RdmaEndPoint::setupConnectionsByActive() {
     HandShakeDesc local_desc, peer_desc;
     std::string peer_server_name, peer_nic_name;
     bool do_rpc = false;
+    int auto_gid_retry_count = 0;
+    std::vector<AutoGidSelectionIdentity> attempted_auto_gid_selections;
 
     {
         RWSpinlock::WriteGuard guard(lock_);
@@ -160,16 +327,7 @@ int RdmaEndPoint::setupConnectionsByActive() {
 
         // loopback mode
         if (context_.nicPath() == peer_nic_path_) {
-            auto segment_desc =
-                context_.engine().meta()->getSegmentDescByID(LOCAL_SEGMENT_ID);
-            if (segment_desc) {
-                for (auto &nic : segment_desc->devices)
-                    if (nic.name == context_.deviceName())
-                        return doSetupConnection(nic.gid, nic.lid, qpNum());
-            }
-            LOG(ERROR) << "Peer NIC " << context_.deviceName()
-                       << " not found in localhost";
-            return ERR_DEVICE_NOT_FOUND;
+            return doSetupConnection(context_.gid(), context_.lid(), qpNum());
         }
 
         // Only proceed with RPC if we are the first to transition from
@@ -187,10 +345,6 @@ int RdmaEndPoint::setupConnectionsByActive() {
                 disconnectUnlocked();
                 return ERR_INVALID_ARGUMENT;
             }
-
-            local_desc.local_nic_path = context_.nicPath();
-            local_desc.peer_nic_path = peer_nic_path_;
-            local_desc.qp_num = qpNum();
         }
     }
 
@@ -227,86 +381,181 @@ int RdmaEndPoint::setupConnectionsByActive() {
         return connected() ? 0 : ERR_ENDPOINT;
     }
 
-    // Perform the RPC without holding the lock to avoid deadlock and allow
-    // "simultaneous open" handshake handling.
-    int rc = context_.engine().sendHandshake(peer_server_name, local_desc,
-                                             peer_desc);
+    for (;;) {
+        std::vector<uint32_t> local_qp_num;
+        {
+            RWSpinlock::ReadGuard guard(lock_);
+            local_qp_num = qpNum();
+        }
+        auto local_gid_selection = fillLocalHandshakeDesc(
+            context_, peer_nic_path_, local_qp_num, local_desc);
+        rememberAutoGidSelection(attempted_auto_gid_selections,
+                                 local_gid_selection);
+        peer_desc = HandShakeDesc();
 
-    // We should check the RPC return code before comparing `peer_qp_num_list_`
-    // with `peer_desc.qp_num`, since a failed RPC may result in an
-    // invalid `peer_desc.qp_num`.
-    //
-    // If the RPC is failed, even if the state is CONNECTED, (which means
-    // it is handled by setupConnectionsByPassive in another thread during the
-    // RPC, or "simultaneous open"), we should resetConnection to be safe.
-    // Because we're not sure whether the peer needs a connection
-    // re-establishment. (We don't know `peer_desc.qp_num`)
-    if (rc) {
-        RWSpinlock::WriteGuard write_guard(lock_);
-        resetConnection("handshake RPC failure");
-        return rc;
-    }
+        // Perform the RPC without holding the lock to avoid deadlock and allow
+        // "simultaneous open" handshake handling.
+        int rc = context_.engine().sendHandshake(peer_server_name, local_desc,
+                                                 peer_desc);
 
-    // Re-acquire lock after RPC to finalize state transition
-    RWSpinlock::WriteGuard guard(lock_);
+        // We should check the RPC return code before comparing
+        // `peer_qp_num_list_` with `peer_desc.qp_num`, since a failed RPC may
+        // result in an invalid `peer_desc.qp_num`.
+        //
+        // If the RPC failed but simultaneous-open passive setup has already
+        // made this endpoint CONNECTED with the same local QPs used by this
+        // RPC, reuse that connection. Otherwise reset because
+        // `peer_desc.qp_num` is invalid and we cannot safely infer the peer
+        // state.
+        if (rc) {
+            RWSpinlock::WriteGuard write_guard(lock_);
+            if (connected()) {
+                auto current_qp_num = qpNum();
+                if (current_qp_num == local_desc.qp_num) {
+                    LOG(WARNING)
+                        << "Active handshake RPC failed, but simultaneous-open "
+                           "passive setup already connected this endpoint. "
+                           "Reusing existing connection. rc="
+                        << rc << ", local_desc.qp_num="
+                        << qpListToString(local_desc.qp_num)
+                        << ", current_qp_num=" << qpListToString(current_qp_num)
+                        << ", endpoint=" << toString();
+                    return 0;
+                }
+            }
 
-    // Handle simultaneous open: if the peer initiates a connection during our
-    // RPC and it is passively established in setupConnectionsByPassive, simply
-    // reuse the existing endpoint.
-    if (connected()) {
-        if (peer_qp_num_list_ == peer_desc.qp_num) {
-            LOG(INFO) << "Received same peer QP numbers, reusing connection.";
-            return 0;
+            LOG(ERROR) << "Active handshake RPC failed; resetting endpoint. rc="
+                       << rc << ", local_desc.qp_num="
+                       << qpListToString(local_desc.qp_num)
+                       << ", endpoint=" << toString();
+            resetConnection("handshake RPC failure");
+            return rc;
         }
 
-        // This mismatch scenario should be rare. It may occur when a peer
-        // first sends us an Active RPC and establishes a connection,
-        // then restarts, and eventually accepts and responds to our
-        // Active RPC.
-        LOG(WARNING) << "Peer QP list mismatch on connected endpoint, "
-                        "re-establishing connection: "
-                     << toString();
+        bool retry_with_new_gid = false;
+        {
+            // Re-acquire lock after RPC to finalize state transition
+            RWSpinlock::WriteGuard guard(lock_);
 
-        int ret = resetConnection("re-establishing connection (active)");
-        if (ret) return ret;
-    }
+            // Handle simultaneous open: if the peer initiates a connection
+            // during our RPC and it is passively established in
+            // setupConnectionsByPassive, simply reuse the existing endpoint.
+            if (connected()) {
+                if (peer_qp_num_list_ == peer_desc.qp_num) {
+                    LOG(INFO)
+                        << "Received same peer QP numbers, reusing connection.";
+                    return 0;
+                }
 
-    if (!peer_desc.reply_msg.empty()) {
-        LOG(ERROR) << "Rejected handshake request by peer "
-                   << local_desc.peer_nic_path;
-        disconnectUnlocked();
-        return ERR_REJECT_HANDSHAKE;
-    }
+                // This mismatch scenario should be rare. It may occur when a
+                // peer first sends us an Active RPC and establishes a
+                // connection, then restarts, and eventually accepts and
+                // responds to our Active RPC.
+                LOG(WARNING) << "Peer QP list mismatch on connected endpoint, "
+                                "re-establishing connection: "
+                             << toString();
 
-    if (peer_desc.local_nic_path != peer_nic_path_ ||
-        peer_desc.peer_nic_path != local_desc.local_nic_path) {
-        LOG(ERROR) << "Invalid argument: received packet mismatch, "
-                      "local.local_nic_path: "
-                   << local_desc.local_nic_path
-                   << ", local.peer_nic_path: " << local_desc.peer_nic_path
-                   << ", peer.local_nic_path: " << peer_desc.local_nic_path
-                   << ", peer.peer_nic_path: " << peer_desc.peer_nic_path;
-        disconnectUnlocked();
-        return ERR_REJECT_HANDSHAKE;
-    }
+                int ret =
+                    resetConnection("re-establishing connection (active)");
+                if (ret) return ret;
+            }
 
-    auto segment_desc =
-        context_.engine().meta()->getSegmentDescByName(peer_server_name);
-    if (segment_desc) {
-        for (auto &nic : segment_desc->devices) {
-            if (nic.name == peer_nic_name) {
-                int ret = doSetupConnection(nic.gid, nic.lid, peer_desc.qp_num);
-                if (ret != 0) {
+            if (!peer_desc.reply_msg.empty()) {
+                LOG(ERROR) << "Rejected handshake request by peer "
+                           << local_desc.peer_nic_path;
+                disconnectUnlocked();
+                return ERR_REJECT_HANDSHAKE;
+            }
+
+            if (peer_desc.local_nic_path != peer_nic_path_ ||
+                peer_desc.peer_nic_path != local_desc.local_nic_path) {
+                LOG(ERROR) << "Invalid argument: received packet mismatch, "
+                              "local.local_nic_path: "
+                           << local_desc.local_nic_path
+                           << ", local.peer_nic_path: "
+                           << local_desc.peer_nic_path
+                           << ", peer.local_nic_path: "
+                           << peer_desc.local_nic_path
+                           << ", peer.peer_nic_path: "
+                           << peer_desc.peer_nic_path;
+                disconnectUnlocked();
+                return ERR_REJECT_HANDSHAKE;
+            }
+
+            int ret = ERR_DEVICE_NOT_FOUND;
+            std::string failure_message;
+            SetupConnectionFailureInfo failure_info;
+            if (!peer_desc.local_gid.empty()) {
+                ret = doSetupConnection(peer_desc.local_gid,
+                                        peer_desc.local_lid, peer_desc.qp_num,
+                                        &failure_message, &failure_info);
+            } else {
+                auto segment_desc =
+                    context_.engine().meta()->getSegmentDescByName(
+                        peer_server_name);
+                if (segment_desc) {
+                    for (auto &nic : segment_desc->devices) {
+                        if (nic.name == peer_nic_name) {
+                            ret = doSetupConnection(
+                                nic.gid, nic.lid, peer_desc.qp_num,
+                                &failure_message, &failure_info);
+                            break;
+                        }
+                    }
+                }
+            }
+
+            if (ret == 0) {
+                return 0;
+            }
+
+            if (shouldAttemptAutoGidHandshakeRetry(
+                    context_.autoGidSelectionEnabled(), auto_gid_retry_count,
+                    globalConfig().auto_gid_max_retries,
+                    failure_info.stage == SetupConnectionFailureStage::kRtr,
+                    failure_info.sys_errno)) {
+                std::string previous_gid;
+                std::string next_gid;
+                bool reprobe_changed = context_.reprobeAutoGid(
+                    local_gid_selection, attempted_auto_gid_selections,
+                    &previous_gid, &next_gid);
+                auto current_gid_selection = context_.gidSelection();
+                auto retry_action = decideAutoGidRetryAction(
+                    reprobe_changed, local_gid_selection.gid_index,
+                    local_gid_selection.gid, current_gid_selection.gid_index,
+                    current_gid_selection.gid);
+                if (retry_action != AutoGidRetryAction::kDoNotRetry) {
+                    int reset_ret = resetConnection(
+                        retry_action ==
+                                AutoGidRetryAction::kRetryWithReprobedGid
+                            ? "retry after auto GID reprobe (active)"
+                            : "retry with externally reprobed GID (active)");
+                    if (reset_ret) return reset_ret;
+                    status_.store(CONNECTING, std::memory_order_relaxed);
+                    ++auto_gid_retry_count;
+                    retry_with_new_gid = true;
+                    LOG(WARNING)
+                        << "Retry active handshake with updated local GID on "
+                        << context_.deviceName() << ": "
+                        << local_gid_selection.gid << " -> "
+                        << current_gid_selection.gid << " (attempt "
+                        << auto_gid_retry_count << "/"
+                        << globalConfig().auto_gid_max_retries << ")";
+                }
+            }
+
+            if (!retry_with_new_gid) {
+                if (ret == ERR_DEVICE_NOT_FOUND) {
+                    LOG(ERROR) << "Peer NIC " << peer_nic_name
+                               << " not found in " << peer_server_name;
+                    disconnectUnlocked();
+                } else {
                     resetConnection("failed connection setup (active)");
                 }
                 return ret;
             }
         }
     }
-    LOG(ERROR) << "Peer NIC " << peer_nic_name << " not found in "
-               << peer_server_name;
-    disconnectUnlocked();
-    return ERR_DEVICE_NOT_FOUND;
 }
 
 int RdmaEndPoint::setupConnectionsByPassive(const HandShakeDesc &peer_desc,
@@ -315,9 +564,8 @@ int RdmaEndPoint::setupConnectionsByPassive(const HandShakeDesc &peer_desc,
     if (connected()) {
         // If already connected with the same peer QP info, return success
         if (peer_qp_num_list_ == peer_desc.qp_num) {
-            local_desc.local_nic_path = context_.nicPath();
-            local_desc.peer_nic_path = peer_nic_path_;
-            local_desc.qp_num = qpNum();
+            fillLocalHandshakeDesc(context_, peer_nic_path_, qpNum(),
+                                   local_desc);
             LOG(INFO) << "Received same peer QP numbers, reusing connection.";
             return 0;
         }
@@ -356,26 +604,82 @@ int RdmaEndPoint::setupConnectionsByPassive(const HandShakeDesc &peer_desc,
         return ERR_INVALID_ARGUMENT;
     }
 
-    local_desc.local_nic_path = context_.nicPath();
-    local_desc.peer_nic_path = peer_nic_path_;
-    local_desc.qp_num = qpNum();
+    status_.store(CONNECTING, std::memory_order_relaxed);
 
-    auto segment_desc =
-        context_.engine().meta()->getSegmentDescByName(peer_server_name);
-    if (segment_desc) {
-        for (auto &nic : segment_desc->devices) {
-            if (nic.name == peer_nic_name) {
-                int ret = doSetupConnection(nic.gid, nic.lid, peer_desc.qp_num,
-                                            &local_desc.reply_msg);
-                if (ret != 0) {
-                    resetConnection("failed connection setup (passive)");
-                }
+    auto attempt_setup_with_peer = [&](const std::string &peer_gid,
+                                       uint16_t peer_lid) -> int {
+        int auto_gid_retry_count = 0;
+        std::vector<AutoGidSelectionIdentity> attempted_auto_gid_selections;
+        for (;;) {
+            auto local_gid_selection = fillLocalHandshakeDesc(
+                context_, peer_nic_path_, qpNum(), local_desc);
+            rememberAutoGidSelection(attempted_auto_gid_selections,
+                                     local_gid_selection);
+
+            SetupConnectionFailureInfo failure_info;
+            int ret = doSetupConnection(peer_gid, peer_lid, peer_desc.qp_num,
+                                        &local_desc.reply_msg, &failure_info);
+            if (ret == 0) {
+                return 0;
+            }
+
+            if (!shouldAttemptAutoGidHandshakeRetry(
+                    context_.autoGidSelectionEnabled(), auto_gid_retry_count,
+                    globalConfig().auto_gid_max_retries,
+                    failure_info.stage == SetupConnectionFailureStage::kRtr,
+                    failure_info.sys_errno)) {
+                resetConnection("failed connection setup (passive)");
                 return ret;
             }
+
+            std::string previous_gid;
+            std::string next_gid;
+            bool reprobe_changed = context_.reprobeAutoGid(
+                local_gid_selection, attempted_auto_gid_selections,
+                &previous_gid, &next_gid);
+            auto current_gid_selection = context_.gidSelection();
+            auto retry_action = decideAutoGidRetryAction(
+                reprobe_changed, local_gid_selection.gid_index,
+                local_gid_selection.gid, current_gid_selection.gid_index,
+                current_gid_selection.gid);
+            if (retry_action == AutoGidRetryAction::kDoNotRetry) {
+                resetConnection("failed connection setup (passive)");
+                return ret;
+            }
+
+            int reset_ret = resetConnection(
+                retry_action == AutoGidRetryAction::kRetryWithReprobedGid
+                    ? "retry after auto GID reprobe (passive)"
+                    : "retry with externally reprobed GID (passive)");
+            if (reset_ret) return reset_ret;
+            status_.store(CONNECTING, std::memory_order_relaxed);
+            ++auto_gid_retry_count;
+            LOG(WARNING) << "Retry passive handshake with updated local GID on "
+                         << context_.deviceName() << ": "
+                         << local_gid_selection.gid << " -> "
+                         << current_gid_selection.gid << " (attempt "
+                         << auto_gid_retry_count << "/"
+                         << globalConfig().auto_gid_max_retries << ")";
+        }
+    };
+
+    if (!peer_desc.local_gid.empty()) {
+        return attempt_setup_with_peer(peer_desc.local_gid,
+                                       peer_desc.local_lid);
+    } else {
+        auto segment_desc =
+            context_.engine().meta()->getSegmentDescByName(peer_server_name);
+        if (segment_desc) {
+            for (auto &nic : segment_desc->devices) {
+                if (nic.name == peer_nic_name) {
+                    return attempt_setup_with_peer(nic.gid, nic.lid);
+                }
+            }
         }
     }
     local_desc.reply_msg =
         "Peer nic not found in that server: " + peer_nic_path_;
+    status_.store(UNCONNECTED, std::memory_order_relaxed);
     LOG(ERROR) << local_desc.reply_msg;
     return ERR_DEVICE_NOT_FOUND;
 }
@@ -389,51 +693,60 @@ int RdmaEndPoint::disconnectUnlocked() {
     auto curr_status = status_.load(std::memory_order_acquire);
     if (curr_status != CONNECTED && curr_status != CONNECTING) return 0;
 
-    ibv_qp_attr attr;
-    memset(&attr, 0, sizeof(attr));
-    attr.qp_state = IBV_QPS_RESET;
-    int ret = 0;
-    for (size_t i = 0; i < qp_list_.size(); ++i) {
-        int curr_ret = ibv_modify_qp(qp_list_[i], &attr, IBV_QP_STATE);
-        if (curr_ret) {
-            PLOG(ERROR) << "Failed to modify QP to RESET";
-            ret = ERR_ENDPOINT;
+    if (!has_connected_) {
+        // Pre-connected handshake retries are allowed to reuse this endpoint:
+        // no user WR has been posted yet. eRDMA still needs fresh QPs because
+        // a QP that reached RTS cannot be reliably reset back to RTS.
+#ifdef CONFIG_ERDMA
+        for (size_t i = 0; i < qp_list_.size(); ++i) {
+            CHECK_EQ(wr_depth_list_[i], 0)
+                << "Pre-connected endpoint must not have outstanding WRs";
         }
-        // After resetting QP, the wr_depth_list_ won't change
-        bool displayed = false;
-        if (wr_depth_list_[i] != 0) {
-            if (!displayed) {
-                LOG(WARNING) << "Outstanding work requests found, CQ will not "
-                                "be generated";
-                displayed = true;
+        return reconstruct();
+#else
+        ibv_qp_attr attr;
+        memset(&attr, 0, sizeof(attr));
+        attr.qp_state = IBV_QPS_RESET;
+        int ret = 0;
+        for (size_t i = 0; i < qp_list_.size(); ++i) {
+            int curr_ret = ibv_modify_qp(qp_list_[i], &attr, IBV_QP_STATE);
+            if (curr_ret) {
+                PLOG(ERROR) << "Failed to modify pre-connected QP to RESET";
+                ret = ERR_ENDPOINT;
             }
-            __sync_fetch_and_sub(cq_outstanding_, wr_depth_list_[i]);
-            wr_depth_list_[i] = 0;
+            CHECK_EQ(wr_depth_list_[i], 0)
+                << "Pre-connected endpoint must not have outstanding WRs";
         }
+        peer_qp_num_list_.clear();
+        status_.store(UNCONNECTED, std::memory_order_release);
+        return ret;
+#endif
     }
-    peer_qp_num_list_.clear();
-    status_.store(UNCONNECTED, std::memory_order_release);
-    return ret;
+
+    beginDestroyLocked();
+    return 0;
 }
 
 int RdmaEndPoint::resetConnection(const std::string &reason) {
     auto curr_status = status_.load(std::memory_order_acquire);
     if (curr_status != CONNECTING && curr_status != CONNECTED) return 0;
 
-#ifdef CONFIG_ERDMA
-    int ret = reconstruct();
-#else
-    int ret = disconnectUnlocked();
-#endif
-
-    if (ret) {
-        LOG(ERROR) << "Failed to reset the endpoint (triggered by: " << reason
-                   << "): error=" << ret;
-    } else {
-        LOG(INFO) << "Successfully reset the endpoint (triggered by: " << reason
-                  << ").";
+    if (!has_connected_) {
+        int ret = disconnectUnlocked();
+        if (ret) {
+            LOG(ERROR) << "Failed to reset pre-connected endpoint "
+                       << "(triggered by: " << reason << "): error=" << ret;
+        } else {
+            LOG(INFO) << "Successfully reset pre-connected endpoint "
+                      << "(triggered by: " << reason << ").";
+        }
+        return ret;
     }
-    return ret;
+
+    LOG(WARNING) << "Retiring endpoint instead of resetting it (triggered by: "
+                 << reason << "): " << toString();
+    beginDestroyLocked();
+    return ERR_ENDPOINT;
 }
 
 const std::string RdmaEndPoint::toString() const {
@@ -441,29 +754,38 @@ const std::string RdmaEndPoint::toString() const {
     if (status == CONNECTED)
         return "EndPoint: local " + context_.nicPath() + ", peer " +
                peer_nic_path_;
+    else if (status == DESTROYING)
+        return "EndPoint: local " + context_.nicPath() + ", peer " +
+               peer_nic_path_ + " (destroying)";
+    else if (status == DESTROYED)
+        return "EndPoint: local " + context_.nicPath() + " (destroyed)";
     else
         return "EndPoint: local " + context_.nicPath() + " (unconnected)";
 }
 
-bool RdmaEndPoint::hasOutstandingSlice() const {
-    if (active_) return true;
-    for (size_t i = 0; i < qp_list_.size(); i++)
-        if (wr_depth_list_[i] != 0) return true;
-    return false;
-}
-
 int RdmaEndPoint::submitPostSend(
     std::vector<Transport::Slice *> &slice_list,
     std::vector<Transport::Slice *> &failed_slice_list) {
     RWSpinlock::WriteGuard guard(lock_);
-    if (!active_) return 0;
+    if (!active_ || status_.load(std::memory_order_relaxed) != CONNECTED) {
+        for (auto &slice : slice_list) failed_slice_list.push_back(slice);
+        slice_list.clear();
+        return 0;
+    }
 
     const size_t num_qp = qp_list_.size();
     if (slice_list.empty()) return 0;
     const size_t requested = slice_list.size();
     int cq_remaining = int(globalConfig().max_cqe) - *cq_outstanding_;
-    std::vector<ibv_send_wr> wr_list(requested, ibv_send_wr{});
-    std::vector<ibv_sge> sge_list(requested);
+    if (cq_remaining <= 0) return 0;
+
+    // Only allocate for the max number of WRs we can actually post per QP,
+    // not the entire requested slice count. Each QP iteration reuses the
+    // wr_list/sge_list from index 0, so we only need max_wr_depth_ entries.
+    size_t max_postable_per_qp =
+        std::min({(size_t)max_wr_depth_, (size_t)cq_remaining, requested});
+    std::vector<ibv_send_wr> wr_list(max_postable_per_qp, ibv_send_wr{});
+    std::vector<ibv_sge> sge_list(max_postable_per_qp);
     size_t total_posted = 0;
     size_t cursor = 0;
 
@@ -542,34 +864,95 @@ std::vector<uint32_t> RdmaEndPoint::qpNum() const {
     return ret;
 }
 
+static int parseGidString(const std::string &gid_str, ibv_gid &gid_out) {
+    if (gid_str.empty()) {
+        LOG(ERROR) << "GID string is empty";
+        return ERR_INVALID_ARGUMENT;
+    }
+
+    // Prepend a colon to the GID string to simplify parsing.
+    std::istringstream iss(":" + gid_str);
+    for (size_t i = 0; i < sizeof(gid_out.raw); i++) {
+        if (iss.get() != ':') {
+            LOG(ERROR) << "Invalid GID format at byte " << i
+                       << ", peer_gid=" << gid_str;
+            return ERR_INVALID_ARGUMENT;
+        }
+
+        uint32_t byte = 0;
+        iss >> std::hex >> byte;
+
+        if (iss.fail() || byte > 0xFF) {
+            LOG(ERROR) << "Invalid GID format at byte " << i
+                       << ", peer_gid=" << gid_str;
+            return ERR_INVALID_ARGUMENT;
+        }
+
+        gid_out.raw[i] = static_cast<uint8_t>(byte);
+    }
+
+    // Ensure no trailing data remains after 16 bytes
+    char extra;
+    if (iss.get(extra)) {
+        LOG(ERROR) << "GID string has trailing data after 16 bytes"
+                   << ", peer_gid=" << gid_str;
+        return ERR_INVALID_ARGUMENT;
+    }
+
+    return 0;
+}
+
 int RdmaEndPoint::doSetupConnection(const std::string &peer_gid,
                                     uint16_t peer_lid,
                                     std::vector<uint32_t> peer_qp_num_list,
-                                    std::string *reply_msg) {
+                                    std::string *reply_msg,
+                                    SetupConnectionFailureInfo *failure_info) {
     if (qp_list_.size() != peer_qp_num_list.size()) {
         std::string message =
             "QP count mismatch in peer and local endpoints, check "
             "MC_MAX_EP_PER_CTX";
         LOG(ERROR) << "[Handshake] " << message;
         if (reply_msg) *reply_msg = message;
+        if (failure_info) {
+            failure_info->stage = SetupConnectionFailureStage::kPeerValidation;
+            failure_info->sys_errno = 0;
+        }
         return ERR_INVALID_ARGUMENT;
     }
 
+    // Verify and parse the peer GID before proceeding.
+    ibv_gid peer_gid_raw = {};
+    int ret = parseGidString(peer_gid, peer_gid_raw);
+    if (ret) {
+        std::string message = "Invalid peer GID: " + peer_gid;
+        LOG(ERROR) << "[Handshake] " << message;
+        if (reply_msg) *reply_msg = message;
+        if (failure_info) {
+            failure_info->stage = SetupConnectionFailureStage::kPeerValidation;
+            failure_info->sys_errno = 0;
+        }
+        return ret;
+    }
+
+    int local_gid_index = context_.gidIndex();
     for (int qp_index = 0; qp_index < (int)qp_list_.size(); ++qp_index) {
-        int ret = doSetupConnection(qp_index, peer_gid, peer_lid,
-                                    peer_qp_num_list[qp_index], reply_msg);
+        int ret = doSetupConnection(qp_index, peer_gid_raw, peer_lid,
+                                    peer_qp_num_list[qp_index], local_gid_index,
+                                    reply_msg, failure_info);
         if (ret) return ret;
     }
 
     peer_qp_num_list_ = std::move(peer_qp_num_list);
+    has_connected_ = true;
     status_.store(CONNECTED, std::memory_order_relaxed);
     return 0;
 }
 
-int RdmaEndPoint::doSetupConnection(int qp_index, const std::string &peer_gid,
+int RdmaEndPoint::doSetupConnection(int qp_index, const ibv_gid &peer_gid,
                                     uint16_t peer_lid, uint32_t peer_qp_num,
-                                    std::string *reply_msg) {
-    if (qp_index < 0 || qp_index > (int)qp_list_.size())
+                                    int local_gid_index, std::string *reply_msg,
+                                    SetupConnectionFailureInfo *failure_info) {
+    if (qp_index < 0 || qp_index >= (int)qp_list_.size())
         return ERR_INVALID_ARGUMENT;
     auto &qp = qp_list_[qp_index];
 
@@ -582,6 +965,10 @@ int RdmaEndPoint::doSetupConnection(int qp_index, const std::string &peer_gid,
         std::string message = "Failed to modify QP to RESET";
         PLOG(ERROR) << "[Handshake] " << message;
         if (reply_msg) *reply_msg = message + ": " + strerror(errno);
+        if (failure_info) {
+            failure_info->stage = SetupConnectionFailureStage::kReset;
+            failure_info->sys_errno = errno;
+        }
         return ERR_ENDPOINT;
     }
 
@@ -589,7 +976,7 @@ int RdmaEndPoint::doSetupConnection(int qp_index, const std::string &peer_gid,
     memset(&attr, 0, sizeof(attr));
     attr.qp_state = IBV_QPS_INIT;
     attr.port_num = context_.portNum();
-    attr.pkey_index = 0;
+    attr.pkey_index = globalConfig().pkey_index;
     attr.qp_access_flags = IBV_ACCESS_LOCAL_WRITE | IBV_ACCESS_REMOTE_READ |
                            IBV_ACCESS_REMOTE_WRITE | IBV_ACCESS_REMOTE_ATOMIC;
     ret = ibv_modify_qp(
@@ -600,6 +987,10 @@ int RdmaEndPoint::doSetupConnection(int qp_index, const std::string &peer_gid,
             "Failed to modify QP to INIT, check local context port num";
         PLOG(ERROR) << "[Handshake] " << message;
         if (reply_msg) *reply_msg = message + ": " + strerror(errno);
+        if (failure_info) {
+            failure_info->stage = SetupConnectionFailureStage::kInit;
+            failure_info->sys_errno = errno;
+        }
         return ERR_ENDPOINT;
     }
 
@@ -609,25 +1000,21 @@ int RdmaEndPoint::doSetupConnection(int qp_index, const std::string &peer_gid,
     attr.path_mtu = context_.activeMTU();
     if (globalConfig().mtu_length < attr.path_mtu)
         attr.path_mtu = globalConfig().mtu_length;
-    ibv_gid peer_gid_raw;
-    std::istringstream iss(peer_gid);
-    for (int i = 0; i < 16; ++i) {
-        int value;
-        iss >> std::hex >> value;
-        peer_gid_raw.raw[i] = static_cast<uint8_t>(value);
-        if (i < 15) iss.ignore(1, ':');
-    }
-    attr.ah_attr.grh.dgid = peer_gid_raw;
+    attr.ah_attr.grh.dgid = peer_gid;
     // TODO gidIndex and portNum must fetch from REMOTE
-    attr.ah_attr.grh.sgid_index = context_.gidIndex();
-    attr.ah_attr.grh.hop_limit = MAX_HOP_LIMIT;
+    attr.ah_attr.grh.sgid_index = local_gid_index;
+    attr.ah_attr.grh.hop_limit = kMaxHopLimit;
     // Set traffic class if configured (-1 means use default)
     if (globalConfig().ib_traffic_class >= 0) {
         attr.ah_attr.grh.traffic_class =
             static_cast<uint8_t>(globalConfig().ib_traffic_class);
     }
     attr.ah_attr.dlid = peer_lid;
+    // Set service level if configured (-1 means use default)
     attr.ah_attr.sl = 0;
+    if (globalConfig().ib_service_level >= 0) {
+        attr.ah_attr.sl = static_cast<uint8_t>(globalConfig().ib_service_level);
+    }
     attr.ah_attr.src_path_bits = 0;
     attr.ah_attr.static_rate = 0;
     attr.ah_attr.is_global = 1;
@@ -645,14 +1032,18 @@ int RdmaEndPoint::doSetupConnection(int qp_index, const std::string &peer_gid,
             "Failed to modify QP to RTR, check mtu, gid, peer lid, peer qp num";
         PLOG(ERROR) << "[Handshake] " << message;
         if (reply_msg) *reply_msg = message + ": " + strerror(errno);
+        if (failure_info) {
+            failure_info->stage = SetupConnectionFailureStage::kRtr;
+            failure_info->sys_errno = errno;
+        }
         return ERR_ENDPOINT;
     }
 
     // RTR -> RTS
     memset(&attr, 0, sizeof(attr));
     attr.qp_state = IBV_QPS_RTS;
-    attr.timeout = TIMEOUT;
-    attr.retry_cnt = RETRY_CNT;
+    attr.timeout = kTimeout;
+    attr.retry_cnt = kRetryCount;
     attr.rnr_retry = 7;  // or 7,RNR error
     attr.sq_psn = 0;
     attr.max_rd_atomic = 16;
@@ -664,9 +1055,74 @@ int RdmaEndPoint::doSetupConnection(int qp_index, const std::string &peer_gid,
         std::string message = "Failed to modify QP to RTS";
         PLOG(ERROR) << "[Handshake] " << message;
         if (reply_msg) *reply_msg = message + ": " + strerror(errno);
+        if (failure_info) {
+            failure_info->stage = SetupConnectionFailureStage::kRts;
+            failure_info->sys_errno = errno;
+        }
         return ERR_ENDPOINT;
     }
 
+    // Optional: pin QP to a specific LAG port for even traffic distribution
+    // across bonded physical ports. num_lag_ports is queried from hardware at
+    // context construction time; this block is a no-op on non-LAG devices.
+    if (globalConfig().mlx5_qp_lag_port_balance) {
+#ifdef USE_MLX5DV
+        uint8_t n = context_.numLagPorts();
+        if (n > 1) {
+            uint8_t target = (uint8_t)(qp_index % n) + 1;
+            int lag_ret = mlx5dv_modify_qp_lag_port(qp, target);
+            if (lag_ret) {
+                LOG_FIRST_N(WARNING, 4)
+                    << "[RDMA] mlx5dv_modify_qp_lag_port failed"
+                    << " (qp_index=" << qp_index
+                    << ", target_port=" << (int)target
+                    << "): " << strerror(lag_ret);
+            } else {
+                uint8_t cfg = 0, active = 0;
+                if (mlx5dv_query_qp_lag_port(qp, &cfg, &active) == 0) {
+                    VLOG(1)
+                        << "[RDMA] QP[" << qp_index << "] qpn=" << qp->qp_num
+                        << " lag_port cfg=" << (int)cfg
+                        << " active=" << (int)active;
+                } else {
+                    LOG_FIRST_N(WARNING, 4)
+                        << "[RDMA] mlx5dv_query_qp_lag_port failed"
+                        << " (qp_index=" << qp_index << ")";
+                }
+            }
+        }
+#else
+        LOG_FIRST_N(WARNING, 1)
+            << "MC_MLX5_QP_LAG_PORT_BALANCE is set but binary was not built "
+               "with USE_MLX5DV; ignoring";
+#endif
+    }
+
+    // Optional: override the RoCEv2 UDP source port to spread QPs across
+    // different ECMP/LAG paths. Empty list = leave whatever the driver
+    // picked. Failure is non-fatal so unsupported devices/firmware degrade
+    // gracefully.
+    const auto &sports = globalConfig().mlx5_qp_udp_sports;
+    if (!sports.empty()) {
+#ifdef USE_MLX5DV
+        uint16_t sport = sports[qp_index % sports.size()];
+        int sp_ret = mlx5dv_modify_qp_udp_sport(qp, sport);
+        if (sp_ret) {
+            LOG_FIRST_N(WARNING, 4)
+                << "[RDMA] mlx5dv_modify_qp_udp_sport failed (qp_index="
+                << qp_index << ", sport=" << sport << "): " << strerror(sp_ret);
+        } else {
+            VLOG(1) << "[RDMA] QP[" << qp_index << "] qpn=" << qp->qp_num
+                    << " udp_sport=" << sport;
+        }
+#else
+        LOG_FIRST_N(WARNING, 1)
+            << "MC_MLX5_QP_UDP_SPORTS is set but binary was not built with "
+               "USE_MLX5DV; ignoring";
+#endif
+    }
+
     return 0;
 }
+
 }  // namespace mooncake
diff --git a/mooncake-transfer-engine/src/transport/rdma_transport/rdma_transport.cpp b/mooncake-transfer-engine/src/transport/rdma_transport/rdma_transport.cpp
index ac49970a..e02b4348 100644
--- a/mooncake-transfer-engine/src/transport/rdma_transport/rdma_transport.cpp
+++ b/mooncake-transfer-engine/src/transport/rdma_transport/rdma_transport.cpp
@@ -21,6 +21,7 @@
 #include <cassert>
 #include <chrono>
 #include <cstddef>
+#include <cstdlib>
 #include <future>
 #include <set>
 #include <thread>
@@ -29,6 +30,7 @@
 
 #include "common.h"
 #include "config.h"
+#include "environ.h"
 #include "memory_location.h"
 #include "topology.h"
 #include "transport/rdma_transport/rdma_context.h"
@@ -39,6 +41,17 @@ namespace mooncake {
 static bool MCIbRelaxedOrderingEnabled = false;
 static int MCIbRelaxedOrderingMode = 2;
 
+static std::string resolveBufferLocation(
+    const TransferMetadata::BufferDesc &buffer, uint64_t offset) {
+    std::string location = buffer.name;
+    SegmentsLocationInfo seg_info;
+    if (parseSegmentsLocation(buffer.name, seg_info)) {
+        location = resolveSegmentsLocation(seg_info, buffer.length,
+                                           offset - buffer.addr);
+    }
+    return location;
+}
+
 // Mode definition for MC_IB_PCI_RELAXED_ORDERING env.
 // 0 - disabled, 1 - enabled if supported, 2 - auto (default, same as 1 today).
 static int getIbRelaxedOrderingMode() {
@@ -101,6 +114,21 @@ int RdmaTransport::install(std::string &local_server_name,
     local_server_name_ = local_server_name;
     local_topology_ = topo;
 
+    // In dual-NIC environments (e.g. separate TCP and RDMA interfaces),
+    // MC_RDMA_BIND_ADDRESS allows NIC paths to use an RDMA-reachable IP
+    // while local_server_name_ keeps the TCP-reachable address for P2P.
+    const char *rdma_bind_addr = std::getenv("MC_RDMA_BIND_ADDRESS");
+    if (rdma_bind_addr && rdma_bind_addr[0] != '\0') {
+        auto [host_name, port] = parseHostNameWithPort(local_server_name);
+        rdma_server_name_ =
+            std::string(rdma_bind_addr) + ":" + std::to_string(port);
+        LOG(INFO) << "RdmaTransport: using RDMA bind address "
+                  << rdma_server_name_
+                  << " (TCP address: " << local_server_name_ << ")";
+    } else {
+        rdma_server_name_ = local_server_name_;
+    }
+
     auto ret = initializeRdmaResources();
     if (ret) {
         LOG(ERROR) << "RdmaTransport: cannot initialize RDMA resources";
@@ -190,7 +218,7 @@ int RdmaTransport::registerLocalMemoryInternal(void *addr, size_t length,
                                   IBV_ACCESS_REMOTE_WRITE |
                                   IBV_ACCESS_REMOTE_READ;
 
-    static int access_rights = kBaseAccessRights;
+    int access_rights = kBaseAccessRights;
     if (MCIbRelaxedOrderingEnabled) {
         access_rights |= IBV_ACCESS_RELAXED_ORDERING;
     }
@@ -358,9 +386,14 @@ int RdmaTransport::unregisterLocalMemoryInternal(void *addr,
 }
 
 int RdmaTransport::allocateLocalSegmentID() {
-    auto desc = metadata_->getSegmentDesc(local_server_name_);
+    auto desc = metadata_->getSegmentDescByID(LOCAL_SEGMENT_ID);
     if (!desc) desc = std::make_shared<SegmentDesc>();
     desc->name = local_server_name_;
+    // Store RDMA server name for dual-NIC setups; when it differs from
+    // local_server_name_ the peer will use it for NIC path construction.
+    if (rdma_server_name_ != local_server_name_) {
+        desc->rdma_server_name = rdma_server_name_;
+    }
 #ifdef ENABLE_MULTI_PROTOCOL
     if (!desc->protocol.empty()) desc->protocol += ",";
     desc->protocol += "rdma";
@@ -380,36 +413,71 @@ int RdmaTransport::allocateLocalSegmentID() {
     return 0;
 }
 
+int RdmaTransport::refreshLocalDeviceDesc(const std::string &device_name,
+                                          uint16_t lid,
+                                          const std::string &gid) {
+    std::lock_guard<std::mutex> guard(local_desc_lock_);
+    auto original_desc = metadata_->getSegmentDescByID(LOCAL_SEGMENT_ID);
+    if (!original_desc) {
+        return ERR_ADDRESS_NOT_REGISTERED;
+    }
+
+    auto updated_desc = std::make_shared<SegmentDesc>(*original_desc);
+    for (auto &device : updated_desc->devices) {
+        if (device.name != device_name) continue;
+        device.lid = lid;
+        device.gid = gid;
+        metadata_->addLocalSegment(LOCAL_SEGMENT_ID, local_server_name_,
+                                   std::move(updated_desc));
+        int ret = metadata_->updateLocalSegmentDesc();
+        if (ret) {
+            auto rollback_desc = original_desc;
+            metadata_->addLocalSegment(LOCAL_SEGMENT_ID, local_server_name_,
+                                       std::move(rollback_desc));
+        }
+        return ret;
+    }
+
+    return ERR_DEVICE_NOT_FOUND;
+}
+
 int RdmaTransport::registerLocalMemoryBatch(
     const std::vector<RdmaTransport::BufferEntry> &buffer_list,
     const std::string &location) {
-#if !defined(WITH_NVIDIA_PEERMEM) && defined(USE_CUDA)
-    for (auto &buffer : buffer_list) {
-        int ret = registerLocalMemory(buffer.addr, buffer.length, location,
-                                      true, false);
-        if (ret) {
-            LOG(WARNING) << "RdmaTransport: Failed to register memory: addr "
-                         << buffer.addr << " length " << buffer.length;
+#if defined(USE_CUDA)
+    if (!Environ::Get().GetWithNvidiaPeermem()) {
+        for (auto &buffer : buffer_list) {
+            int ret = registerLocalMemory(buffer.addr, buffer.length, location,
+                                          true, false);
+            if (ret) {
+                LOG(WARNING)
+                    << "RdmaTransport: Failed to register memory: addr "
+                    << buffer.addr << " length " << buffer.length;
+            }
+        }
+    } else {
+#endif
+        std::vector<std::future<int>> results;
+        for (auto &buffer : buffer_list) {
+            results.emplace_back(std::async(
+                std::launch::async, [this, buffer, location]() -> int {
+                    // Use force_sequential=true to avoid nested parallelism
+                    return registerLocalMemoryInternal(buffer.addr,
+                                                       buffer.length, location,
+                                                       true, false, true);
+                }));
         }
-    }
-#else
-    std::vector<std::future<int>> results;
-    for (auto &buffer : buffer_list) {
-        results.emplace_back(
-            std::async(std::launch::async, [this, buffer, location]() -> int {
-                // Use force_sequential=true to avoid nested parallelism
-                return registerLocalMemoryInternal(buffer.addr, buffer.length,
-                                                   location, true, false, true);
-            }));
-    }
 
-    for (size_t i = 0; i < buffer_list.size(); ++i) {
-        if (results[i].get()) {
-            LOG(WARNING) << "RdmaTransport: Failed to register memory: addr "
-                         << buffer_list[i].addr << " length "
-                         << buffer_list[i].length;
+        for (size_t i = 0; i < buffer_list.size(); ++i) {
+            if (results[i].get()) {
+                LOG(WARNING)
+                    << "RdmaTransport: Failed to register memory: addr "
+                    << buffer_list[i].addr << " length "
+                    << buffer_list[i].length;
+            }
         }
-    }
+#if defined(USE_CUDA)
+    }  // Environ::Get().GetWithNvidiaPeermem()
 #endif
 
     return metadata_->updateLocalSegmentDesc();
@@ -449,7 +517,17 @@ Status RdmaTransport::submitTransfer(
     size_t task_id = batch_desc.task_list.size();
     batch_desc.task_list.resize(task_id + entries.size());
     std::vector<TransferTask *> task_list;
-    for (auto &task : batch_desc.task_list) task_list.push_back(&task);
+    for (auto &request : entries) {
+        auto &task = batch_desc.task_list[task_id];
+        ++task_id;
+        task.batch_id = batch_id;
+#ifdef USE_ASCEND_HETEROGENEOUS
+        task.request = const_cast<TransferRequest *>(&request);
+#else
+        task.request = &request;
+#endif
+        task_list.push_back(&task);
+    }
     return submitTransferTask(task_list);
 }
 
@@ -494,6 +572,7 @@ Status RdmaTransport::submitTransferTask(
             slice->source_addr = (char *)request.source + offset;
             slice->length =
                 merge_final_slice ? request.length - offset : kBlockSize;
+            slice->source_location.clear();
             slice->opcode = request.opcode;
             slice->rdma.dest_addr = request.target_offset + offset;
             slice->rdma.retry_cnt = request.advise_retry_cnt;
@@ -550,6 +629,11 @@ Status RdmaTransport::submitTransferTask(
                 }
                 slice->rdma.source_lkey =
                     local_segment_desc->buffers[buffer_id].lkey[device_id];
+                if (globalConfig().log_rdma_slice_affinity) {
+                    slice->source_location = resolveBufferLocation(
+                        local_segment_desc->buffers[buffer_id],
+                        reinterpret_cast<uint64_t>(slice->source_addr));
+                }
                 slices_to_post[context].push_back(slice);
                 task.total_bytes += slice->length;
                 __sync_fetch_and_add(&task.slice_count, 1);
@@ -645,7 +729,9 @@ int RdmaTransport::onSetupRdmaConnections(const HandShakeDesc &peer_desc,
     // Use existing endpoint or create new one.
     auto endpoint = context->endpoint(peer_desc.local_nic_path);
     if (!endpoint) return ERR_ENDPOINT;
-    return endpoint->setupConnectionsByPassive(peer_desc, local_desc);
+    int ret = endpoint->setupConnectionsByPassive(peer_desc, local_desc);
+    if (endpoint->retired()) context->deleteEndpointByPtr(endpoint.get());
+    return ret;
 }
 
 int RdmaTransport::initializeRdmaResources() {
@@ -719,6 +805,33 @@ int RdmaTransport::selectDevice(SegmentDesc *desc, uint64_t offset,
     return ERR_ADDRESS_NOT_REGISTERED;
 }
 
+int RdmaTransport::selectDeviceByLocalHca(SegmentDesc *desc, uint64_t offset,
+                                          size_t length,
+                                          std::string_view local_hca,
+                                          int &buffer_id, int &device_id,
+                                          int retry_count) {
+    if (desc == nullptr) return ERR_ADDRESS_NOT_REGISTERED;
+    const auto &buffers = desc->buffers;
+    for (buffer_id = 0; buffer_id < static_cast<int>(buffers.size());
+         ++buffer_id) {
+        const auto &buffer = buffers[buffer_id];
+
+        if (offset < buffer.addr || length > buffer.length ||
+            offset - buffer.addr > buffer.length - length) {
+            continue;
+        }
+
+        const auto location = resolveBufferLocation(buffer, offset);
+        device_id = desc->topology.selectDeviceByLocalHca(location, local_hca,
+                                                          retry_count);
+        if (device_id >= 0) return 0;
+        device_id = desc->topology.selectDeviceByLocalHca(
+            kWildcardLocation, local_hca, retry_count);
+        if (device_id >= 0) return 0;
+    }
+    return ERR_ADDRESS_NOT_REGISTERED;
+}
+
 int RdmaTransport::selectDevice(SegmentDesc *desc, uint64_t offset,
                                 size_t length, int &buffer_id, int &device_id,
                                 int retry_count) {
diff --git a/mooncake-transfer-engine/src/transport/rdma_transport/worker_pool.cpp b/mooncake-transfer-engine/src/transport/rdma_transport/worker_pool.cpp
index c469ab26..8b9a1bc2 100644
--- a/mooncake-transfer-engine/src/transport/rdma_transport/worker_pool.cpp
+++ b/mooncake-transfer-engine/src/transport/rdma_transport/worker_pool.cpp
@@ -19,6 +19,7 @@
 #include <cassert>
 
 #include "config.h"
+#include "memory_location.h"
 #include "transport/rdma_transport/rdma_context.h"
 #include "transport/rdma_transport/rdma_endpoint.h"
 #include "transport/rdma_transport/rdma_transport.h"
@@ -31,11 +32,44 @@ namespace mooncake {
 
 const static int kTransferWorkerCount = globalConfig().workers_per_ctx;
 
+static std::string resolveBufferLocation(
+    const TransferMetadata::BufferDesc &buffer, uint64_t offset) {
+    std::string location = buffer.name;
+    SegmentsLocationInfo seg_info;
+    if (parseSegmentsLocation(buffer.name, seg_info)) {
+        location = resolveSegmentsLocation(seg_info, buffer.length,
+                                           offset - buffer.addr);
+    }
+    return location;
+}
+
+static const std::string &sourceLocationOrUnknown(Transport::Slice *slice) {
+    static const std::string kUnknown = "<unknown>";
+    return slice->source_location.empty() ? kUnknown : slice->source_location;
+}
+
+static int selectPeerDevice(RdmaTransport::SegmentDesc *peer_segment_desc,
+                            uint64_t offset, size_t length,
+                            const std::string &local_hca, int &buffer_id,
+                            int &device_id, int retry_count = 0) {
+    const auto &config = globalConfig();
+    if (config.enable_hca_peer_affinity) {
+        return RdmaTransport::selectDeviceByLocalHca(
+            peer_segment_desc, offset, length, local_hca, buffer_id, device_id,
+            retry_count);
+    }
+
+    auto hint = config.enable_dest_device_affinity ? std::string_view(local_hca)
+                                                   : std::string_view();
+    return RdmaTransport::selectDevice(peer_segment_desc, offset, length, hint,
+                                       buffer_id, device_id, retry_count);
+}
+
 WorkerPool::WorkerPool(RdmaContext &context, int numa_socket_id)
     : context_(context),
       numa_socket_id_(numa_socket_id),
       workers_running_(true),
-      suspended_flag_(0),
+      parked_worker_count_(0),
       redispatch_counter_(0),
       submitted_slice_count_(0),
       processed_slice_count_(0) {
@@ -97,6 +131,7 @@ int WorkerPool::submitPostSend(
 
     SliceList slice_list_map[kShardCount];
     uint64_t submitted_slice_count = 0;
+    int all_rails_failed_count = 0;
     thread_local std::unordered_map<int, uint64_t> failed_target_ids;
     for (auto &slice : slice_list) {
         if (failed_target_ids.count(slice->target_id)) {
@@ -110,12 +145,9 @@ int WorkerPool::submitPostSend(
         }
         auto &peer_segment_desc = segment_desc_map[slice->target_id];
         int buffer_id, device_id;
-        auto hint = globalConfig().enable_dest_device_affinity
-                        ? context_.deviceName()
-                        : "";
-        if (RdmaTransport::selectDevice(peer_segment_desc.get(),
-                                        slice->rdma.dest_addr, slice->length,
-                                        hint, buffer_id, device_id)) {
+        if (selectPeerDevice(peer_segment_desc.get(), slice->rdma.dest_addr,
+                             slice->length, context_.deviceName(), buffer_id,
+                             device_id)) {
             peer_segment_desc = context_.engine().meta()->getSegmentDescByID(
                 slice->target_id, true);
             if (!peer_segment_desc) {
@@ -126,9 +158,9 @@ int WorkerPool::submitPostSend(
                 continue;
             }
 
-            if (RdmaTransport::selectDevice(
-                    peer_segment_desc.get(), slice->rdma.dest_addr,
-                    slice->length, hint, buffer_id, device_id)) {
+            if (selectPeerDevice(peer_segment_desc.get(), slice->rdma.dest_addr,
+                                 slice->length, context_.deviceName(),
+                                 buffer_id, device_id)) {
                 slice->markFailed();
                 context_.engine().meta()->dumpMetadataContent(
                     peer_segment_desc->name, slice->rdma.dest_addr,
@@ -143,9 +175,49 @@ int WorkerPool::submitPostSend(
         slice->rdma.dest_rkey =
             peer_segment_desc->buffers[buffer_id].rkey[device_id];
         auto peer_nic_path =
-            MakeNicPath(peer_segment_desc->name,
+            MakeNicPath(peer_segment_desc->nicPathServerName(),
                         peer_segment_desc->devices[device_id].name);
+
+        // If selected rail is paused, try alternative devices
+        if (!isRailAvailable(peer_nic_path)) {
+            bool found = false;
+            for (size_t alt_dev_id = 0;
+                 alt_dev_id < peer_segment_desc->devices.size(); ++alt_dev_id) {
+                if (alt_dev_id == (size_t)device_id) continue;
+                auto alt_path =
+                    MakeNicPath(peer_segment_desc->name,
+                                peer_segment_desc->devices[alt_dev_id].name);
+                if (isRailAvailable(alt_path)) {
+                    device_id = alt_dev_id;
+                    slice->rdma.dest_rkey =
+                        peer_segment_desc->buffers[buffer_id].rkey[device_id];
+                    peer_nic_path = alt_path;
+                    found = true;
+                    break;
+                }
+            }
+            if (!found) {
+                slice->markFailed();  // All rails unavailable
+                all_rails_failed_count++;
+                continue;
+            }
+        }
+
         slice->peer_nic_path = peer_nic_path;
+        if (globalConfig().log_rdma_slice_affinity) {
+            VLOG(1) << "RDMA slice affinity: source_location="
+                    << sourceLocationOrUnknown(slice) << ", target_location="
+                    << resolveBufferLocation(
+                           peer_segment_desc->buffers[buffer_id],
+                           slice->rdma.dest_addr)
+                    << ", local_device_name=" << context_.deviceName()
+                    << ", peer_device_name="
+                    << peer_segment_desc->devices[device_id].name
+                    << ", target_id=" << slice->target_id
+                    << ", source_addr=" << slice->source_addr << ", dest_addr="
+                    << reinterpret_cast<void *>(slice->rdma.dest_addr)
+                    << ", length=" << slice->length;
+        }
         int shard_id = (slice->target_id * 10007 + device_id) % kShardCount;
         slice_list_map[shard_id].push_back(slice);
         submitted_slice_count++;
@@ -162,15 +234,43 @@ int WorkerPool::submitPostSend(
     }
 
     submitted_slice_count_.fetch_add(submitted_slice_count);
-    if (suspended_flag_.load()) {
+    if (submitted_slice_count &&
+        parked_worker_count_.load(std::memory_order_acquire) > 0) {
         std::lock_guard<std::mutex> lock(cond_mutex_);
         cond_var_.notify_all();
     }
 
+    // Context-level health tracking: if all slices failed due to no available
+    // rails, increment the context failure counter. This detects catastrophic
+    // local RNIC hardware failure where all paths through the RNIC are down.
+    if (submitted_slice_count == 0 &&
+        all_rails_failed_count == (int)slice_list.size()) {
+        markContextFailure();
+    }
+
     return 0;
 }
 
 void WorkerPool::performPostSend(int thread_id) {
+    // Fast-fail if context is unhealthy due to catastrophic hardware failure
+    if (!contextHealthy()) {
+        for (int shard_id = thread_id; shard_id < kShardCount;
+             shard_id += kTransferWorkerCount) {
+            if (slice_queue_count_[shard_id].load(std::memory_order_relaxed) ==
+                0)
+                continue;
+            slice_queue_lock_[shard_id].lock();
+            for (auto &entry : slice_queue_[shard_id]) {
+                for (auto &slice : entry.second) slice->markFailed();
+                processed_slice_count_ += entry.second.size();
+            }
+            slice_queue_[shard_id].clear();
+            slice_queue_count_[shard_id].store(0, std::memory_order_relaxed);
+            slice_queue_lock_[shard_id].unlock();
+        }
+        return;
+    }
+
     auto &local_slice_queue = collective_slice_queue_[thread_id];
     for (int shard_id = thread_id; shard_id < kShardCount;
          shard_id += kTransferWorkerCount) {
@@ -232,37 +332,36 @@ void WorkerPool::performPostSend(int thread_id) {
             entry.second.clear();
             continue;
         }
-        if (!endpoint->active()) {
-            if (endpoint->inactiveTime() > 1.0)
-                context_.deleteEndpoint(
-                    entry.first);  // enable for re-establishation
+        if (!endpoint->connected() && endpoint->setupConnectionsByActive()) {
+            LOG(ERROR) << "Worker: Cannot make connection for endpoint: "
+                       << entry.first << ", deleting endpoint";
+            // Unified path failure handling
+            handlePathFailure(entry.first, endpoint.get());
             for (auto &slice : entry.second) failed_slice_list.push_back(slice);
             entry.second.clear();
             continue;
         }
-        if (!endpoint->connected() && endpoint->setupConnectionsByActive()) {
-            LOG(ERROR) << "Worker: Cannot make connection for endpoint: "
-                       << entry.first << ", mark it inactive";
-            for (auto &slice : entry.second) failed_slice_list.push_back(slice);
-            endpoint->set_active(false);
-            failed_nr_polls++;
-            if (context_.active() && failed_nr_polls > 32 &&
-                !success_nr_polls) {
-                LOG(WARNING)
-                    << "Failed to establish peer endpoints in local RNIC "
-                    << context_.nicPath() << ", mark it inactive";
-                context_.set_active(false);
-            }
-            entry.second.clear();
-            continue;
+        // Set endpoint pointer for each slice before submitting
+        for (auto &slice : entry.second) {
+            slice->rdma.endpoint = endpoint.get();
         }
         endpoint->submitPostSend(entry.second, failed_slice_list);
 #endif
     }
 
     if (!failed_slice_list.empty()) {
-        for (auto &slice : failed_slice_list) slice->rdma.retry_cnt++;
-        redispatch(failed_slice_list, thread_id);
+        SliceList retry_list;
+        for (auto &slice : failed_slice_list) {
+            if (shouldRetrySlice(slice)) {
+                retry_list.push_back(slice);
+            } else {
+                slice->markFailed();
+                processed_slice_count_++;
+            }
+        }
+        if (!retry_list.empty()) {
+            redispatch(retry_list, thread_id);
+        }
     }
 }
 
@@ -270,6 +369,7 @@ void WorkerPool::performPollCq(int thread_id) {
     int processed_slice_count = 0;
     const static size_t kPollCount = 64;
     std::unordered_map<volatile int *, int> qp_depth_set;
+    SliceList failed_slice_list;  // Unified: collect all slices for redispatch
     for (int cq_index = thread_id; cq_index < context_.cqCount();
          cq_index += kTransferWorkerCount) {
         ibv_wc wc[kPollCount];
@@ -288,45 +388,44 @@ void WorkerPool::performPollCq(int thread_id) {
                 qp_depth_set[slice->rdma.qp_depth] = 1;
             // __sync_fetch_and_sub(slice->rdma.qp_depth, 1);
             if (wc[i].status != IBV_WC_SUCCESS) {
-                bool show_work_request_flushed_error = globalConfig().trace;
-                // After detect an error, subsequent work requests will result
-                // in work_request_flushed_error, we hide this by default
-                if (wc[i].status != IBV_WC_WR_FLUSH_ERR ||
-                    show_work_request_flushed_error)
-                    LOG(ERROR)
-                        << "Worker: Process failed for slice (opcode: "
-                        << slice->opcode
-                        << ", source_addr: " << slice->source_addr
-                        << ", length: " << slice->length
-                        << ", dest_addr: " << (void *)slice->rdma.dest_addr
-                        << ", local_nic: " << context_.deviceName()
-                        << ", peer_nic: " << slice->peer_nic_path
-                        << ", dest_rkey: " << slice->rdma.dest_rkey
-                        << ", retry_cnt: " << slice->rdma.retry_cnt
-                        << "): " << ibv_wc_status_str(wc[i].status);
-                failed_nr_polls++;
-                if (context_.active() && failed_nr_polls > 32 &&
-                    !success_nr_polls) {
-                    LOG(WARNING) << "Too many errors found in local RNIC "
-                                 << context_.nicPath() << ", mark it inactive";
-                    context_.set_active(false);
+                // Flush errors are generated when QPs transition to ERR state
+                // during normal endpoint destruction (beginDestroy). They are
+                // not real network errors and should not trigger rail failure
+                // handling or endpoint deletion.
+                if (wc[i].status == IBV_WC_WR_FLUSH_ERR) {
+                    if (globalConfig().trace)
+                        LOG(INFO) << "Worker: WR flush error (peer_nic: "
+                                  << slice->peer_nic_path
+                                  << "), marking failed without retry";
+                    slice->markFailed();
+                    processed_slice_count++;
+                    continue;
                 }
-                context_.deleteEndpoint(slice->peer_nic_path);
-                slice->rdma.retry_cnt++;
-                if (slice->rdma.retry_cnt >= slice->rdma.max_retry_cnt) {
+
+                // All other WC errors indicate real path/network failures and
+                // should trigger redispatch to an alternate path (or fail if
+                // retry exhausted)
+                LOG(ERROR) << "Worker: Process failed for slice (opcode: "
+                           << slice->opcode
+                           << ", source_addr: " << slice->source_addr
+                           << ", length: " << slice->length
+                           << ", dest_addr: " << (void *)slice->rdma.dest_addr
+                           << ", local_nic: " << context_.deviceName()
+                           << ", peer_nic: " << slice->peer_nic_path
+                           << ", dest_rkey: " << slice->rdma.dest_rkey
+                           << ", retry_cnt: " << slice->rdma.retry_cnt
+                           << "): " << ibv_wc_status_str(wc[i].status);
+                // Unified path failure handling
+                handlePathFailure(slice->peer_nic_path, slice->rdma.endpoint);
+                if (shouldRetrySlice(slice)) {
+                    failed_slice_list.push_back(slice);
+                } else {
                     slice->markFailed();
                     processed_slice_count_++;
-                } else {
-                    collective_slice_queue_[thread_id][slice->peer_nic_path]
-                        .push_back(slice);
-                    redispatch_counter_++;
-                    // std::vector<RdmaTransport::Slice *> slice_list { slice };
-                    // redispatch(slice_list, thread_id);
                 }
             } else {
                 slice->markSuccess();
                 processed_slice_count++;
-                success_nr_polls++;
             }
         }
         if (nr_poll)
@@ -337,8 +436,14 @@ void WorkerPool::performPollCq(int thread_id) {
     for (auto &entry : qp_depth_set)
         __sync_fetch_and_sub(entry.first, entry.second);
 
-    if (processed_slice_count)
+    if (processed_slice_count) {
         processed_slice_count_.fetch_add(processed_slice_count);
+        markContextSuccess();
+    }
+
+    if (!failed_slice_list.empty()) {
+        redispatch(failed_slice_list, thread_id);
+    }
 }
 
 void WorkerPool::redispatch(std::vector<Transport::Slice *> &slice_list,
@@ -361,10 +466,9 @@ void WorkerPool::redispatch(std::vector<Transport::Slice *> &slice_list,
             auto &peer_segment_desc = segment_desc_map[slice->target_id];
             int buffer_id, device_id;
             if (!peer_segment_desc ||
-                RdmaTransport::selectDevice(peer_segment_desc.get(),
-                                            slice->rdma.dest_addr,
-                                            slice->length, buffer_id, device_id,
-                                            slice->rdma.retry_cnt)) {
+                selectPeerDevice(peer_segment_desc.get(), slice->rdma.dest_addr,
+                                 slice->length, context_.deviceName(),
+                                 buffer_id, device_id, slice->rdma.retry_cnt)) {
                 slice->markFailed();
                 processed_slice_count_++;
                 continue;
@@ -372,14 +476,39 @@ void WorkerPool::redispatch(std::vector<Transport::Slice *> &slice_list,
             slice->rdma.dest_rkey =
                 peer_segment_desc->buffers[buffer_id].rkey[device_id];
             auto peer_nic_path =
-                MakeNicPath(peer_segment_desc->name,
+                MakeNicPath(peer_segment_desc->nicPathServerName(),
                             peer_segment_desc->devices[device_id].name);
             slice->peer_nic_path = peer_nic_path;
+            if (globalConfig().log_rdma_slice_affinity) {
+                VLOG(1) << "RDMA slice affinity: source_location="
+                        << sourceLocationOrUnknown(slice)
+                        << ", target_location="
+                        << resolveBufferLocation(
+                               peer_segment_desc->buffers[buffer_id],
+                               slice->rdma.dest_addr)
+                        << ", local_device_name=" << context_.deviceName()
+                        << ", peer_device_name="
+                        << peer_segment_desc->devices[device_id].name
+                        << ", target_id=" << slice->target_id
+                        << ", source_addr=" << slice->source_addr
+                        << ", dest_addr="
+                        << reinterpret_cast<void *>(slice->rdma.dest_addr)
+                        << ", length=" << slice->length
+                        << ", retry_cnt=" << slice->rdma.retry_cnt;
+            }
             collective_slice_queue_[thread_id][peer_nic_path].push_back(slice);
         }
     }
 }
 
+bool WorkerPool::hasOutstandingCq(int thread_id) {
+    for (int cq_index = thread_id; cq_index < context_.cqCount();
+         cq_index += kTransferWorkerCount) {
+        if (*context_.cqOutstandingCount(cq_index) > 0) return true;
+    }
+    return false;
+}
+
 void WorkerPool::transferWorker(int thread_id) {
     bindToSocket(numa_socket_id_);
     const static uint64_t kWaitPeriodInNano = 100000000;  // 100ms
@@ -389,18 +518,21 @@ void WorkerPool::transferWorker(int thread_id) {
             processed_slice_count_.load(std::memory_order_relaxed);
         auto submitted_slice_count =
             submitted_slice_count_.load(std::memory_order_relaxed);
-        if (processed_slice_count == submitted_slice_count) {
+        if (processed_slice_count == submitted_slice_count &&
+            !hasOutstandingCq(thread_id)) {
             uint64_t curr_wait_ts = getCurrentTimeInNano();
             if (curr_wait_ts - last_wait_ts > kWaitPeriodInNano) {
                 std::unique_lock<std::mutex> lock(cond_mutex_);
-                suspended_flag_.fetch_add(1);
+                parked_worker_count_.fetch_add(1, std::memory_order_acq_rel);
                 // Double-check condition after acquiring lock to avoid lost
-                // wakeup
+                // wakeup. parked_worker_count_ is set before this check so
+                // producers that submit after it will notify this worker.
                 if (processed_slice_count_.load(std::memory_order_relaxed) ==
-                    submitted_slice_count_.load()) {
+                        submitted_slice_count_.load() &&
+                    !hasOutstandingCq(thread_id)) {
                     cond_var_.wait_for(lock, std::chrono::seconds(1));
                 }
-                suspended_flag_.fetch_sub(1);
+                parked_worker_count_.fetch_sub(1, std::memory_order_acq_rel);
                 last_wait_ts = curr_wait_ts;
             }
             continue;
@@ -409,6 +541,7 @@ void WorkerPool::transferWorker(int thread_id) {
 #ifndef USE_FAKE_POST_SEND
         performPollCq(thread_id);
 #endif
+        last_wait_ts = getCurrentTimeInNano();
     }
 }
 
@@ -444,10 +577,7 @@ int WorkerPool::doProcessContextEvents() {
          * the endpoint from context_ and use that shared_ptr to access the
          * endpoint.
          */
-        auto endpoint = context_.getEndpointByPtr(endpoint_ptr);
-        if (endpoint) {
-            endpoint->set_active(false);
-        }
+        context_.deleteEndpointByPtr(endpoint_ptr);
     } else if (event.event_type == IBV_EVENT_DEVICE_FATAL ||
                event.event_type == IBV_EVENT_CQ_ERR ||
                event.event_type == IBV_EVENT_WQ_FATAL ||
@@ -475,6 +605,7 @@ int WorkerPool::doProcessContextEvents() {
                   << " is now inactive";
     } else if (event.event_type == IBV_EVENT_PORT_ACTIVE) {
         context_.set_active(true);
+        markContextSuccess();  // Reset failure counter on port recovery
         LOG(INFO) << "Worker: Context " << context_.deviceName()
                   << " is now active";
     }
@@ -492,7 +623,11 @@ void WorkerPool::monitorWorker() {
     while (workers_running_) {
         auto current_ts = getCurrentTimeInNano();
         if (current_ts - last_reset_ts > 1000000000ll) {
-            context_.set_active(true);
+            // Drain endpoint_store_->waiting_list_ even when no new
+            // insertions are happening. Without this, reclaim only runs
+            // from RdmaContext::endpoint() and the waiting list grows
+            // unboundedly under failure load. See issue #1845.
+            context_.reclaimEndpoints();
             last_reset_ts = current_ts;
         }
         struct epoll_event event;
@@ -511,4 +646,50 @@ void WorkerPool::monitorWorker() {
             doProcessContextEvents();
     }
 }
+
+void WorkerPool::markRailFailed(const std::string &peer_nic_path) {
+    std::lock_guard<std::mutex> lock(rail_state_lock_);
+    auto &state = rail_states_[peer_nic_path];
+    state.error_count++;
+    if (state.error_count >= kRailErrorThreshold) {
+        uint64_t now = getCurrentTimeInNano();
+        state.pause_until_ns = now + kRailPauseNs;
+        LOG(WARNING) << "Rail paused: peer=" << peer_nic_path
+                     << " error_count=" << state.error_count;
+    }
+}
+
+bool WorkerPool::isRailAvailable(const std::string &peer_nic_path) {
+    std::lock_guard<std::mutex> lock(rail_state_lock_);
+    auto it = rail_states_.find(peer_nic_path);
+    if (it == rail_states_.end()) return true;
+    auto &state = it->second;
+    if (state.pause_until_ns == 0) return true;
+    uint64_t now = getCurrentTimeInNano();
+    if (now >= state.pause_until_ns) {
+        // Auto-recover: pause expired
+        state.error_count = 0;
+        state.pause_until_ns = 0;
+        return true;
+    }
+    return false;
+}
+
+// Unified retry logic: increment retry count and return whether retry is
+// allowed
+bool WorkerPool::shouldRetrySlice(Transport::Slice *slice) {
+    slice->rdma.retry_cnt++;
+    return slice->rdma.retry_cnt < slice->rdma.max_retry_cnt;
+}
+
+// Unified path failure handler
+void WorkerPool::handlePathFailure(const std::string &peer_nic_path,
+                                   RdmaEndPoint *endpoint) {
+    markRailFailed(peer_nic_path);
+    redispatch_counter_++;  // Notify all workers to redispatch their queues
+    if (endpoint) {
+        context_.deleteEndpointByPtr(endpoint);
+    }
+}
+
 }  // namespace mooncake
diff --git a/mooncake-transfer-engine/src/transport/rpc_communicator/CMakeLists.txt b/mooncake-transfer-engine/src/transport/rpc_communicator/CMakeLists.txt
index 841a36aa..3b6bf0cd 100644
--- a/mooncake-transfer-engine/src/transport/rpc_communicator/CMakeLists.txt
+++ b/mooncake-transfer-engine/src/transport/rpc_communicator/CMakeLists.txt
@@ -1,6 +1,6 @@
 file(GLOB CORO_RPC_SOURCES "*.cpp")
 
-find_package(Python3 COMPONENTS Interpreter Development REQUIRED)
+find_package(Python3 COMPONENTS Interpreter Development.Module REQUIRED)
 
 add_library(rpc_communicator OBJECT ${CORO_RPC_SOURCES})
 
@@ -10,7 +10,6 @@ target_link_libraries(rpc_communicator
     yalantinglibs::yalantinglibs 
     glog::glog 
     pthread 
-    ${Python3_LIBRARIES}
 )
 
 target_include_directories(rpc_communicator 
diff --git a/mooncake-transfer-engine/src/transport/sunrise_link/CMakeLists.txt b/mooncake-transfer-engine/src/transport/sunrise_link/CMakeLists.txt
new file mode 100644
index 00000000..1f9a1c0b
--- /dev/null
+++ b/mooncake-transfer-engine/src/transport/sunrise_link/CMakeLists.txt
@@ -0,0 +1,6 @@
+if(USE_SUNRISE)
+  file(GLOB SUNRISE_LINK_TRANSPORT_SOURCES "*.cpp")
+  add_library(sunrise_link_transport OBJECT ${SUNRISE_LINK_TRANSPORT_SOURCES})
+  target_include_directories(sunrise_link_transport PUBLIC ${MC_TANGRT_ROOT}/include)
+  target_link_libraries(sunrise_link_transport PUBLIC glog::glog dl)
+endif()
diff --git a/mooncake-transfer-engine/src/transport/sunrise_link/sunrise_link_transport.cpp b/mooncake-transfer-engine/src/transport/sunrise_link/sunrise_link_transport.cpp
new file mode 100644
index 00000000..f3c3d7ea
--- /dev/null
+++ b/mooncake-transfer-engine/src/transport/sunrise_link/sunrise_link_transport.cpp
@@ -0,0 +1,687 @@
+// Copyright 2024 KVCache.AI
+
+#include "transport/sunrise_link_transport/sunrise_link_transport.h"
+
+#include <dlfcn.h>
+#include <glog/logging.h>
+#include <ptml.h>
+#include <tang_runtime_api.h>
+
+#include <algorithm>
+#include <cerrno>
+#include <cstdlib>
+#include <cstring>
+#include <limits>
+#include <string>
+
+#include "common.h"
+#include "common/serialization.h"
+
+namespace {
+
+static std::string TangRtSharedObjectPath(const char* soname) {
+    const char* lib_dir = std::getenv("MC_TANGRT_LIB_DIR");
+    if (lib_dir && lib_dir[0]) {
+        std::string d(lib_dir);
+        while (!d.empty() && d.back() == '/') d.pop_back();
+        return d + "/" + soname;
+    }
+    const char* root = std::getenv("MC_TANGRT_ROOT");
+    std::string r = (root && root[0]) ? std::string(root)
+                                      : std::string("/usr/local/tangrt");
+    while (!r.empty() && r.back() == '/') r.pop_back();
+    const char* arch = std::getenv("MC_TANGRT_LIB_ARCH");
+    std::string arch_dir =
+        (arch && arch[0]) ? std::string(arch) : std::string("linux-x86_64");
+    std::string arch_path = r + "/lib/" + arch_dir + "/" + soname;
+    if (access(arch_path.c_str(), R_OK) == 0) return arch_path;
+    return r + "/lib/" + soname;
+}
+
+static const char* kRawAddrPrefix = "RAW_ADDR:";
+
+struct RuntimeState {
+    bool loaded{false};
+};
+
+static std::mutex& processRuntimeMutex() {
+    static std::mutex m;
+    return m;
+}
+
+static RuntimeState& processRuntime() {
+    static RuntimeState state;
+    return state;
+}
+
+constexpr int kMaxTangDevLocks = 64;
+
+static int devSlot(int dev) {
+    if (dev < 0) return -1;
+    return dev % kMaxTangDevLocks;
+}
+
+static std::mutex& devMutex(int slot) {
+    static std::mutex table[kMaxTangDevLocks];
+    static std::mutex fallback;
+    if (slot < 0 || slot >= kMaxTangDevLocks) return fallback;
+    return table[slot];
+}
+
+struct ScopedTangDeviceOps {
+    std::unique_lock<std::mutex> lo_lock;
+    std::unique_lock<std::mutex> hi_lock;
+
+    explicit ScopedTangDeviceOps(int src_dev, int dst_dev) {
+        if (src_dev < 0 && dst_dev < 0) return;
+        if (src_dev < 0 || dst_dev < 0) {
+            int slot = devSlot(src_dev >= 0 ? src_dev : dst_dev);
+            lo_lock = std::unique_lock<std::mutex>(devMutex(slot));
+            return;
+        }
+        int slot_lo = devSlot(src_dev);
+        int slot_hi = devSlot(dst_dev);
+        if (slot_lo == slot_hi) {
+            lo_lock = std::unique_lock<std::mutex>(devMutex(slot_lo));
+            return;
+        }
+        if (slot_lo > slot_hi) std::swap(slot_lo, slot_hi);
+        lo_lock =
+            std::unique_lock<std::mutex>(devMutex(slot_lo), std::defer_lock);
+        hi_lock =
+            std::unique_lock<std::mutex>(devMutex(slot_hi), std::defer_lock);
+        std::lock(lo_lock, hi_lock);
+    }
+};
+
+static tangMemcpyKind MemcpyKindForPointers(tangMemoryType st,
+                                            tangMemoryType dt) {
+    if (st == tangMemoryTypeDevice && dt == tangMemoryTypeHost) {
+        return tangMemcpyDeviceToHost;
+    }
+    if (st == tangMemoryTypeHost && dt == tangMemoryTypeDevice) {
+        return tangMemcpyHostToDevice;
+    }
+    if (st == tangMemoryTypeDevice && dt == tangMemoryTypeDevice) {
+        return tangMemcpyDeviceToDevice;
+    }
+    return tangMemcpyHostToHost;
+}
+
+struct SavedTangDevice {
+    int dev{-1};
+    SavedTangDevice() { tangGetDevice(&dev); }
+    ~SavedTangDevice() {
+        if (dev >= 0) tangSetDevice(dev);
+    }
+    SavedTangDevice(const SavedTangDevice&) = delete;
+    SavedTangDevice& operator=(const SavedTangDevice&) = delete;
+};
+
+static void NormalizeMappedHostPointer(void** ptr,
+                                       tangPointerAttributes* attr) {
+    if (!ptr || !*ptr || !attr) return;
+    if (attr->type != tangMemoryTypeHost) return;
+
+    void* dev_ptr = nullptr;
+    tangError_t ret = tangHostGetDevicePointer(&dev_ptr, *ptr, 0);
+    if (ret != tangSuccess || !dev_ptr) return;
+
+    tangPointerAttributes dev_attr = {};
+    tangPointerGetAttributes(&dev_attr, dev_ptr);
+    if (dev_attr.type == tangMemoryTypeDevice) {
+        *ptr = dev_ptr;
+        *attr = dev_attr;
+    }
+}
+
+static tangError_t QueryPointerAttrsBestEffort(void* ptr,
+                                               tangPointerAttributes* attr,
+                                               int preferred_dev) {
+    if (!ptr || !attr) return tangErrorInvalidValue;
+    SavedTangDevice saved;
+
+    if (preferred_dev >= 0) {
+        tangSetDevice(preferred_dev);
+    }
+
+    tangError_t ret = tangPointerGetAttributes(attr, ptr);
+    if (ret == tangSuccess) return ret;
+
+    int dev_count = 0;
+    if (tangGetDeviceCount(&dev_count) != tangSuccess || dev_count <= 0) {
+        return ret;
+    }
+    for (int d = 0; d < dev_count; ++d) {
+        if (d == preferred_dev) continue;
+        tangSetDevice(d);
+        ret = tangPointerGetAttributes(attr, ptr);
+        if (ret == tangSuccess) return ret;
+    }
+    return ret;
+}
+
+static tangError_t doTangCopy(void* dst, int dst_dev, void* src, int src_dev,
+                              size_t length,
+                              tangMemoryType src_type = tangMemoryTypeDevice,
+                              tangMemoryType dst_type = tangMemoryTypeDevice) {
+    if (src_dev >= 0 && dst_dev >= 0 && src_dev != dst_dev) {
+        tangError_t ret = tangMemcpyPeer(dst, dst_dev, src, src_dev, length);
+        if (ret != tangSuccess) {
+            ret = tangMemcpyPeer_v2(dst, dst_dev, src, src_dev, length);
+        }
+        return ret;
+    }
+    if (src_dev >= 0 && dst_dev >= 0 && src_dev == dst_dev) {
+        ScopedTangDeviceOps ctx_lock(src_dev, dst_dev);
+        int saved_dev = -1;
+        tangGetDevice(&saved_dev);
+        tangSetDevice(src_dev);
+        tangError_t ret =
+            tangMemcpy(dst, src, length, tangMemcpyDeviceToDevice);
+        if (saved_dev >= 0) tangSetDevice(saved_dev);
+        return ret;
+    }
+    ScopedTangDeviceOps ctx_lock(src_dev, dst_dev);
+    int saved_dev = -1;
+    tangGetDevice(&saved_dev);
+    if (src_dev >= 0) {
+        tangSetDevice(src_dev);
+    } else if (dst_dev >= 0) {
+        tangSetDevice(dst_dev);
+    }
+    tangError_t ret =
+        tangMemcpy(dst, src, length, MemcpyKindForPointers(src_type, dst_type));
+    if (saved_dev >= 0) tangSetDevice(saved_dev);
+    return ret;
+}
+
+static int ParseSunriseDeviceId(const std::string& location) {
+    auto pos = location.rfind(':');
+    if (pos == std::string::npos) return -1;
+    const char* value = location.c_str() + pos + 1;
+    if (*value == '\0') return -1;
+
+    char* end = nullptr;
+    errno = 0;
+    long parsed = std::strtol(value, &end, 10);
+    if (errno != 0 || end == value || *end != '\0') return -1;
+    if (parsed < 0 || parsed > std::numeric_limits<int>::max()) return -1;
+    return static_cast<int>(parsed);
+}
+
+static bool RequestRangeFitsInBuffer(uint64_t dest_addr, uint64_t length,
+                                     uint64_t buffer_addr,
+                                     uint64_t buffer_length) {
+    if (dest_addr < buffer_addr) return false;
+
+    uint64_t offset = dest_addr - buffer_addr;
+    if (offset > buffer_length) return false;
+    return length <= buffer_length - offset;
+}
+
+}  // namespace
+
+namespace mooncake {
+
+int SunriseLinkTransport::parseDeviceIdForTest(const std::string& location) {
+    return ParseSunriseDeviceId(location);
+}
+
+bool SunriseLinkTransport::requestRangeFitsInBufferForTest(
+    uint64_t dest_addr, uint64_t length, uint64_t buffer_addr,
+    uint64_t buffer_length) {
+    return RequestRangeFitsInBuffer(dest_addr, length, buffer_addr,
+                                    buffer_length);
+}
+
+SunriseLinkTransport::SunriseLinkTransport() = default;
+
+SunriseLinkTransport::~SunriseLinkTransport() {
+    for (auto& entry : remap_entries_) {
+        if (entry.second.shm_addr && !entry.second.is_raw_addr)
+            tangIpcCloseMemHandle(entry.second.shm_addr);
+    }
+    remap_entries_.clear();
+}
+
+int SunriseLinkTransport::install(std::string& local_server_name,
+                                  std::shared_ptr<TransferMetadata> metadata,
+                                  std::shared_ptr<Topology> topology) {
+    (void)topology;
+    metadata_ = metadata;
+    local_server_name_ = local_server_name;
+
+    if (!loadRuntime()) {
+        LOG(ERROR) << "SunriseLinkTransport: failed to load Tang/PTML runtime";
+        return -1;
+    }
+
+    auto desc = std::make_shared<SegmentDesc>();
+    if (!desc) return ERR_MEMORY;
+    desc->name = local_server_name_;
+    desc->protocol = "sunrise_link";
+    metadata_->addLocalSegment(LOCAL_SEGMENT_ID, local_server_name_,
+                               std::move(desc));
+    return 0;
+}
+
+bool SunriseLinkTransport::loadRuntime() {
+    std::lock_guard<std::mutex> guard(processRuntimeMutex());
+    auto& rt = processRuntime();
+    if (rt.loaded) return true;
+
+    void* ptml_handle =
+        dlopen(TangRtSharedObjectPath("libptml_shared.so").c_str(),
+               RTLD_NOW | RTLD_GLOBAL);
+    if (!ptml_handle) {
+        LOG(ERROR) << "SunriseLinkTransport: failed to load PTML: "
+                   << dlerror();
+        return false;
+    }
+
+    void* tang_handle =
+        dlopen(TangRtSharedObjectPath("libtangrt_shared.so").c_str(),
+               RTLD_NOW | RTLD_GLOBAL);
+    if (!tang_handle) {
+        LOG(ERROR) << "SunriseLinkTransport: failed to load Tang runtime: "
+                   << dlerror();
+        return false;
+    }
+
+    ptmlReturn_t ret = ptmlInit();
+    if (ret != PTML_SUCCESS) {
+        LOG(ERROR) << "SunriseLinkTransport: ptmlInit failed: " << ret;
+        return false;
+    }
+    ret = ptmlPtlinkEnableAll();
+    if (ret != PTML_SUCCESS) {
+        LOG(WARNING) << "SunriseLinkTransport: ptmlPtlinkEnableAll failed: "
+                     << ret;
+    }
+
+    rt.loaded = true;
+    return true;
+}
+
+int SunriseLinkTransport::inferRegisteredDevice(void* ptr) const {
+    uint64_t p = reinterpret_cast<uint64_t>(ptr);
+    std::lock_guard<std::mutex> guard(register_mutex_);
+    for (const auto& entry : registered_regions_) {
+        uint64_t base = reinterpret_cast<uint64_t>(entry.first);
+        if (base <= p && p < base + entry.second.length)
+            return entry.second.gpu_id;
+    }
+    return -1;
+}
+
+int SunriseLinkTransport::inferPointerDeviceBestEffort(
+    void* ptr, int preferred_dev) const {
+    int dev = inferRegisteredDevice(ptr);
+    if (dev >= 0) return dev;
+
+    tangPointerAttributes attr = {};
+    tangError_t ret = QueryPointerAttrsBestEffort(ptr, &attr, preferred_dev);
+    if (ret != tangSuccess) return -1;
+
+    NormalizeMappedHostPointer(&ptr, &attr);
+    return attr.type == tangMemoryTypeDevice ? attr.device : -1;
+}
+
+int SunriseLinkTransport::registerLocalMemory(void* addr, size_t length,
+                                              const std::string& location,
+                                              bool remote_accessible,
+                                              bool update_metadata) {
+    (void)remote_accessible;
+    int preferred_dev = ParseSunriseDeviceId(location);
+    tangPointerAttributes attr = {};
+    tangError_t err = QueryPointerAttrsBestEffort(addr, &attr, preferred_dev);
+    if (err != tangSuccess || attr.type != tangMemoryTypeDevice) {
+        LOG(ERROR) << "SunriseLinkTransport: unsupported memory for register"
+                   << " addr=" << addr << " location=" << location
+                   << " preferred_dev=" << preferred_dev << " err=" << err
+                   << " type=" << attr.type << " device=" << attr.device;
+        return -1;
+    }
+
+    tangIpcMemHandle_t handle;
+    int saved_dev = -1;
+    tangGetDevice(&saved_dev);
+    int handle_dev = preferred_dev >= 0 ? preferred_dev : attr.device;
+    if (handle_dev >= 0) {
+        tangError_t sd = tangSetDevice(handle_dev);
+        if (sd != tangSuccess) {
+            LOG(ERROR) << "SunriseLinkTransport: tangSetDevice before "
+                       << "tangIpcGetMemHandle failed: " << sd << " "
+                       << tangGetErrorString(sd) << " device=" << handle_dev;
+            if (saved_dev >= 0) tangSetDevice(saved_dev);
+            return -1;
+        }
+    }
+    err = tangIpcGetMemHandle(&handle, addr);
+    if (saved_dev >= 0) tangSetDevice(saved_dev);
+
+    std::string shm_name;
+    if (err == tangSuccess) {
+        shm_name = serializeBinaryData(&handle, sizeof(handle));
+    } else {
+        LOG(WARNING) << "SunriseLinkTransport: tangIpcGetMemHandle failed: "
+                     << err << " " << tangGetErrorString(err)
+                     << ", falling back to RAW_ADDR";
+        shm_name = kRawAddrPrefix + std::to_string(attr.device);
+    }
+
+    {
+        std::lock_guard<std::mutex> guard(register_mutex_);
+        if (registered_regions_.count(addr)) {
+            return 0;
+        }
+        registered_regions_[addr] = RegisteredRegion{length, attr.device};
+    }
+
+    BufferDesc desc;
+    desc.addr = reinterpret_cast<uint64_t>(addr);
+    desc.length = length;
+    desc.name = location;
+    desc.shm_name = shm_name;
+    return metadata_->addLocalMemoryBuffer(desc, update_metadata);
+}
+
+int SunriseLinkTransport::unregisterLocalMemory(void* addr,
+                                                bool update_metadata) {
+    {
+        std::lock_guard<std::mutex> guard(register_mutex_);
+        registered_regions_.erase(addr);
+    }
+    return metadata_->removeLocalMemoryBuffer(addr, update_metadata);
+}
+
+int SunriseLinkTransport::registerLocalMemoryBatch(
+    const std::vector<BufferEntry>& buffer_list, const std::string& location) {
+    for (const auto& buffer : buffer_list) {
+        registerLocalMemory(buffer.addr, buffer.length, location, true, false);
+    }
+    return metadata_->updateLocalSegmentDesc();
+}
+
+int SunriseLinkTransport::unregisterLocalMemoryBatch(
+    const std::vector<void*>& addr_list) {
+    for (auto* addr : addr_list) unregisterLocalMemory(addr, false);
+    return metadata_->updateLocalSegmentDesc();
+}
+
+int SunriseLinkTransport::relocateSharedMemoryAddress(uint64_t& dest_addr,
+                                                      uint64_t length,
+                                                      uint64_t target_id,
+                                                      int* target_gpu_id) {
+    auto desc = metadata_->getSegmentDescByID(target_id);
+    for (auto& entry : desc->buffers) {
+        if (!entry.shm_name.empty() &&
+            RequestRangeFitsInBuffer(dest_addr, length, entry.addr,
+                                     entry.length)) {
+            remap_lock_.lockShared();
+            auto key = std::make_pair(target_id, entry.addr);
+            if (remap_entries_.count(key)) {
+                auto opened = remap_entries_[key];
+                remap_lock_.unlockShared();
+                if (target_gpu_id) *target_gpu_id = opened.gpu_id;
+                dest_addr = dest_addr - entry.addr +
+                            reinterpret_cast<uint64_t>(opened.shm_addr);
+                return 0;
+            }
+            remap_lock_.unlockShared();
+
+            RWSpinlock::WriteGuard guard(remap_lock_);
+            if (!remap_entries_.count(key)) {
+                if (entry.shm_name.compare(0, strlen(kRawAddrPrefix),
+                                           kRawAddrPrefix) == 0) {
+                    int remote_gpu_id = std::stoi(
+                        entry.shm_name.substr(strlen(kRawAddrPrefix)));
+                    OpenedShmEntry shm_entry;
+                    shm_entry.shm_addr = reinterpret_cast<void*>(entry.addr);
+                    shm_entry.length = entry.length;
+                    shm_entry.gpu_id = remote_gpu_id;
+                    shm_entry.is_raw_addr = true;
+                    remap_entries_[key] = shm_entry;
+                } else {
+                    std::vector<unsigned char> output_buffer;
+                    deserializeBinaryData(entry.shm_name, output_buffer);
+                    if (output_buffer.size() != sizeof(tangIpcMemHandle_t)) {
+                        LOG(ERROR)
+                            << "SunriseLinkTransport: invalid IPC handle size";
+                        return -1;
+                    }
+                    tangIpcMemHandle_t handle;
+                    memcpy(&handle, output_buffer.data(), sizeof(handle));
+                    void* shm_addr = nullptr;
+                    int remote_gpu_id = ParseSunriseDeviceId(entry.name);
+                    int saved_dev = -1;
+                    tangGetDevice(&saved_dev);
+                    if (remote_gpu_id >= 0) {
+                        tangError_t sd = tangSetDevice(remote_gpu_id);
+                        if (sd != tangSuccess) {
+                            LOG(ERROR)
+                                << "SunriseLinkTransport: tangSetDevice "
+                                << "before tangIpcOpenMemHandle failed: " << sd
+                                << " " << tangGetErrorString(sd)
+                                << " device=" << remote_gpu_id;
+                            if (saved_dev >= 0) tangSetDevice(saved_dev);
+                            return -1;
+                        }
+                    }
+                    tangError_t err = tangIpcOpenMemHandle(
+                        &shm_addr, handle, tangIpcMemLazyEnablePeerAccess);
+                    if (saved_dev >= 0) tangSetDevice(saved_dev);
+                    if (err != tangSuccess) {
+                        LOG(ERROR)
+                            << "SunriseLinkTransport: tangIpcOpenMemHandle "
+                            << "failed: " << err << " "
+                            << tangGetErrorString(err);
+                        return -1;
+                    }
+                    OpenedShmEntry shm_entry;
+                    shm_entry.shm_addr = shm_addr;
+                    shm_entry.length = entry.length;
+                    shm_entry.gpu_id = remote_gpu_id;
+                    remap_entries_[key] = shm_entry;
+                }
+            }
+
+            auto opened = remap_entries_[key];
+            if (target_gpu_id) *target_gpu_id = opened.gpu_id;
+            dest_addr = dest_addr - entry.addr +
+                        reinterpret_cast<uint64_t>(opened.shm_addr);
+            return 0;
+        }
+    }
+    LOG(ERROR)
+        << "SunriseLinkTransport: requested address not found in segment";
+    return ERR_INVALID_ARGUMENT;
+}
+
+int SunriseLinkTransport::startCopy(void* src, void* dst, size_t length,
+                                    int remote_dev, int local_dev) {
+    tangPointerAttributes src_attr = {};
+    tangPointerAttributes dst_attr = {};
+    tangError_t src_ret =
+        QueryPointerAttrsBestEffort(src, &src_attr, local_dev);
+    tangError_t dst_ret =
+        QueryPointerAttrsBestEffort(dst, &dst_attr, remote_dev);
+    NormalizeMappedHostPointer(&src, &src_attr);
+    NormalizeMappedHostPointer(&dst, &dst_attr);
+
+    int src_dev =
+        (src_ret == tangSuccess && src_attr.type == tangMemoryTypeDevice)
+            ? src_attr.device
+            : -1;
+    int dst_dev =
+        (dst_ret == tangSuccess && dst_attr.type == tangMemoryTypeDevice)
+            ? dst_attr.device
+            : -1;
+
+    tangMemoryType src_type =
+        src_ret == tangSuccess ? src_attr.type : tangMemoryTypeHost;
+    tangMemoryType dst_type =
+        dst_ret == tangSuccess ? dst_attr.type : tangMemoryTypeHost;
+
+    tangError_t ret =
+        doTangCopy(dst, dst_dev, src, src_dev, length, src_type, dst_type);
+    if (ret != tangSuccess) {
+        LOG(ERROR) << "SunriseLinkTransport: tang copy failed: " << ret << " "
+                   << tangGetErrorString(ret) << " src_dev=" << src_dev
+                   << " dst_dev=" << dst_dev;
+        return -1;
+    }
+    return 0;
+}
+
+int SunriseLinkTransport::submitRequest(TransferTask& task,
+                                        const TransferRequest& request,
+                                        bool keep_slice_ref) {
+    uint64_t dest_addr = request.target_offset;
+    int target_gpu_id = -1;
+    bool is_ipc = (request.target_id != LOCAL_SEGMENT_ID);
+    if (is_ipc) {
+        int rc = relocateSharedMemoryAddress(dest_addr, request.length,
+                                             request.target_id, &target_gpu_id);
+        if (rc) return rc;
+    }
+
+    int local_gpu_id = inferPointerDeviceBestEffort(request.source, -1);
+
+    task.total_bytes = request.length;
+    Slice* slice = getSliceCache().allocate();
+    slice->source_addr = static_cast<char*>(request.source);
+    slice->local.dest_addr = reinterpret_cast<void*>(dest_addr);
+    slice->length = request.length;
+    slice->opcode = request.opcode;
+    slice->task = &task;
+    slice->target_id = request.target_id;
+    slice->status = Slice::PENDING;
+    if (keep_slice_ref) task.slice_list.push_back(slice);
+    __sync_fetch_and_add(&task.slice_count, 1);
+
+    void* src = nullptr;
+    void* dst = nullptr;
+    if (request.opcode == TransferRequest::READ) {
+        src = reinterpret_cast<void*>(dest_addr);
+        dst = request.source;
+    } else {
+        src = request.source;
+        dst = reinterpret_cast<void*>(dest_addr);
+    }
+
+    if (is_ipc && target_gpu_id >= 0) {
+        int remote_dev = target_gpu_id;
+        void* local_ptr = request.opcode == TransferRequest::WRITE ? src : dst;
+        tangPointerAttributes local_attr = {};
+        tangError_t local_ret =
+            QueryPointerAttrsBestEffort(local_ptr, &local_attr, local_gpu_id);
+        if (local_ret == tangSuccess) {
+            NormalizeMappedHostPointer(&local_ptr, &local_attr);
+        }
+        if (local_ret == tangSuccess &&
+            local_attr.type == tangMemoryTypeDevice) {
+            int local_dev = local_attr.device;
+            if (request.opcode == TransferRequest::WRITE) {
+                src = local_ptr;
+            } else {
+                dst = local_ptr;
+            }
+            int src_dev = (request.opcode == TransferRequest::WRITE)
+                              ? local_dev
+                              : remote_dev;
+            int dst_dev = (request.opcode == TransferRequest::WRITE)
+                              ? remote_dev
+                              : local_dev;
+
+            tangError_t ret =
+                doTangCopy(dst, dst_dev, src, src_dev, request.length);
+            if (ret != tangSuccess) {
+                if (src_dev == dst_dev) {
+                    LOG(ERROR)
+                        << "SunriseLink IPC same-device copy failed "
+                        << "(IPC memory is read-only on same device, "
+                        << "use cross-GPU for WRITE): " << ret << " "
+                        << tangGetErrorString(ret) << " src_dev=" << src_dev
+                        << " dst_dev=" << dst_dev;
+                } else {
+                    LOG(ERROR)
+                        << "SunriseLinkTransport: tang copy failed: " << ret
+                        << " " << tangGetErrorString(ret)
+                        << " src_dev=" << src_dev << " dst_dev=" << dst_dev;
+                }
+                slice->markFailed();
+                return -1;
+            }
+
+            slice->markSuccess();
+            return 0;
+        }
+    }
+
+    if (startCopy(src, dst, request.length, target_gpu_id, local_gpu_id) != 0) {
+        slice->markFailed();
+        return -1;
+    }
+
+    slice->markSuccess();
+    return 0;
+}
+
+Status SunriseLinkTransport::submitTransfer(
+    BatchID batch_id, const std::vector<TransferRequest>& entries) {
+    auto& batch_desc = *((BatchDesc*)(batch_id));
+    if (batch_desc.task_list.size() + entries.size() > batch_desc.batch_size) {
+        return Status::InvalidArgument(
+            "SunriseLinkTransport: exceed batch capacity");
+    }
+
+    size_t task_id = batch_desc.task_list.size();
+    batch_desc.task_list.resize(task_id + entries.size());
+    for (const auto& request : entries) {
+        auto& task = batch_desc.task_list[task_id++];
+        if (submitRequest(task, request, false) != 0) {
+            return Status::Memory("sunrise_link transfer failed");
+        }
+    }
+    return Status::OK();
+}
+
+Status SunriseLinkTransport::submitTransferTask(
+    const std::vector<TransferTask*>& task_list) {
+    for (auto* task_ptr : task_list) {
+        if (!task_ptr || !task_ptr->request) {
+            return Status::InvalidArgument(
+                "SunriseLinkTransport: invalid transfer task");
+        }
+        if (submitRequest(*task_ptr, *task_ptr->request, true) != 0) {
+            return Status::Memory("sunrise_link transfer failed");
+        }
+    }
+    return Status::OK();
+}
+
+Status SunriseLinkTransport::getTransferStatus(BatchID batch_id, size_t task_id,
+                                               TransferStatus& status) {
+    auto& batch_desc = *((BatchDesc*)(batch_id));
+    if (task_id >= batch_desc.task_list.size()) {
+        return Status::InvalidArgument(
+            "SunriseLinkTransport::getTransferStatus invalid task id");
+    }
+
+    auto& task = batch_desc.task_list[task_id];
+    status.transferred_bytes = task.transferred_bytes;
+    uint64_t success_slice_count = task.success_slice_count;
+    uint64_t failed_slice_count = task.failed_slice_count;
+    if (success_slice_count + failed_slice_count == task.slice_count) {
+        status.s = failed_slice_count ? TransferStatusEnum::FAILED
+                                      : TransferStatusEnum::COMPLETED;
+        task.is_finished = true;
+    } else {
+        status.s = TransferStatusEnum::WAITING;
+    }
+    return Status::OK();
+}
+
+}  // namespace mooncake
diff --git a/mooncake-transfer-engine/src/transport/tcp_transport/tcp_transport.cpp b/mooncake-transfer-engine/src/transport/tcp_transport/tcp_transport.cpp
index 44ea42af..2ac73e11 100644
--- a/mooncake-transfer-engine/src/transport/tcp_transport/tcp_transport.cpp
+++ b/mooncake-transfer-engine/src/transport/tcp_transport/tcp_transport.cpp
@@ -20,9 +20,11 @@
 
 #include <algorithm>
 #include <cassert>
+#include <cctype>
 #include <chrono>
 #include <cstddef>
 #include <cstdint>
+#include <cstdlib>
 #include <memory>
 #include <random>
 
@@ -36,7 +38,30 @@
 
 namespace mooncake {
 using tcpsocket = asio::ip::tcp::socket;
-const static size_t kDefaultBufferSize = 65536;
+static size_t getChunkSize() {
+    static const size_t val = [] {
+        const char* env = std::getenv("MC_TCP_SLICE_SIZE");
+        if (env) {
+            try {
+                size_t v = std::stoull(env);
+                if (v > 0) return v;
+                LOG(WARNING)
+                    << "Ignore non-positive MC_TCP_SLICE_SIZE value: " << env
+                    << ", using default 65536";
+            } catch (const std::exception& e) {
+                // A non-numeric or out-of-range value makes std::stoull throw;
+                // fall through to the default instead of letting the exception
+                // propagate out of this static initializer and abort the
+                // transfer that first reads the chunk size.
+                LOG(WARNING)
+                    << "Invalid MC_TCP_SLICE_SIZE value: " << env
+                    << ". Error: " << e.what() << ", using default 65536";
+            }
+        }
+        return size_t(65536);  // 64KB default
+    }();
+    return val;
+}
 
 struct SessionHeader {
     uint64_t size;
@@ -44,26 +69,59 @@ struct SessionHeader {
     uint8_t opcode;
 };
 
-#if defined(USE_CUDA) || defined(USE_MUSA) || defined(USE_HIP) || \
-    defined(USE_MLU) || defined(USE_MACA)
+#if defined(USE_CUDA) || defined(USE_MUSA) || defined(USE_HIP) ||  \
+    defined(USE_MLU) || defined(USE_MACA) || defined(USE_HYGON) || \
+    defined(USE_COREX)
 static bool isCudaMemory(void* addr) {
     cudaPointerAttributes attributes;
     auto status = cudaPointerGetAttributes(&attributes, addr);
     if (status != cudaSuccess) return false;
-    if (attributes.type == cudaMemoryTypeDevice) return true;
-    return false;
+    return attributes.type == cudaMemoryTypeDevice;
 }
+
+// Returns the CUDA device ordinal if addr is device memory, or -1 otherwise.
+// Callers must call cudaSetDevice before any cudaMemcpy to avoid implicit
+// GPU 0 context creation.
+static int getCudaDeviceId(void* addr) {
+    cudaPointerAttributes attributes;
+    auto status = cudaPointerGetAttributes(&attributes, addr);
+    if (status != cudaSuccess) return -1;
+    if (attributes.type == cudaMemoryTypeDevice) return attributes.device;
+    return -1;
+}
+
+#ifdef USE_MACA
+static cudaError_t copyTcpCudaMemory(void* dst, const void* src, size_t size) {
+    cudaStream_t stream;
+    cudaError_t status =
+        cudaStreamCreateWithFlags(&stream, cudaStreamNonBlocking);
+    if (status != cudaSuccess) return status;
+
+    status = cudaMemcpyAsync(dst, src, size, cudaMemcpyDefault, stream);
+    if (status == cudaSuccess) {
+        status = cudaStreamSynchronize(stream);
+    }
+
+    cudaError_t destroy_status = cudaStreamDestroy(stream);
+    return status == cudaSuccess ? destroy_status : status;
+}
+#endif
 #endif
 
 // Forward declaration
 class TcpTransport;
 
+using ValidateAddrFn = std::function<bool(uint64_t, uint64_t)>;
+
 // Server-side session: handles one transfer request on a persistent connection
 struct ServerSession : public std::enable_shared_from_this<ServerSession> {
-    explicit ServerSession(std::shared_ptr<tcpsocket> socket)
-        : socket_(std::move(socket)) {}
+    explicit ServerSession(std::shared_ptr<tcpsocket> socket,
+                           ValidateAddrFn validate_addr)
+        : socket_(std::move(socket)),
+          validate_addr_(std::move(validate_addr)) {}
 
     std::shared_ptr<tcpsocket> socket_;
+    ValidateAddrFn validate_addr_;
     SessionHeader header_;
     uint64_t total_transferred_bytes_;
     char* local_buffer_;
@@ -83,8 +141,6 @@ struct ServerSession : public std::enable_shared_from_this<ServerSession> {
             *socket_, asio::buffer(&header_, sizeof(SessionHeader)),
             [this, self](const asio::error_code& ec, std::size_t len) {
                 if (ec || len != sizeof(SessionHeader)) {
-                    // If client closed connection (EOF), this is normal - don't
-                    // log
                     if (ec.value() != asio::error::eof) {
                         LOG(WARNING)
                             << "ServerSession::readHeader failed. Error: "
@@ -92,10 +148,20 @@ struct ServerSession : public std::enable_shared_from_this<ServerSession> {
                             << ", bytes read: " << len;
                     }
                     session_mutex_.unlock();
-                    return;  // Don't continue, socket will be closed
+                    return;
                 }
 
                 local_buffer_ = (char*)(le64toh(header_.addr));
+                uint64_t size = le64toh(header_.size);
+                if (validate_addr_ &&
+                    !validate_addr_((uint64_t)local_buffer_, size)) {
+                    LOG(ERROR) << "ServerSession: remote-supplied address 0x"
+                               << std::hex << (uint64_t)local_buffer_
+                               << std::dec << " with size " << size
+                               << " is not within any registered buffer";
+                    session_mutex_.unlock();
+                    return;
+                }
                 if (header_.opcode == (uint8_t)TransferRequest::WRITE)
                     readBody();
                 else
@@ -109,7 +175,7 @@ struct ServerSession : public std::enable_shared_from_this<ServerSession> {
         char* addr = local_buffer_;
 
         size_t buffer_size =
-            std::min(kDefaultBufferSize, size - total_transferred_bytes_);
+            std::min(getChunkSize(), size - total_transferred_bytes_);
         if (buffer_size == 0) {
             session_mutex_.unlock();
             // Transfer complete, wait for next request on this connection
@@ -118,14 +184,23 @@ struct ServerSession : public std::enable_shared_from_this<ServerSession> {
         }
 
         char* dram_buffer = addr + total_transferred_bytes_;
+        int cuda_device = -1;
 
-#if defined(USE_CUDA) || defined(USE_MUSA) || defined(USE_HIP) || \
-    defined(USE_MLU) || defined(USE_MACA)
-        if (isCudaMemory(addr)) {
+#if defined(USE_CUDA) || defined(USE_MUSA) || defined(USE_HIP) ||  \
+    defined(USE_MLU) || defined(USE_MACA) || defined(USE_HYGON) || \
+    defined(USE_COREX)
+        cuda_device = getCudaDeviceId(addr);
+        if (cuda_device >= 0) {
             dram_buffer = new char[buffer_size];
+            cudaSetDevice(cuda_device);
+#ifdef USE_MACA
+            cudaError_t cuda_status = copyTcpCudaMemory(
+                dram_buffer, addr + total_transferred_bytes_, buffer_size);
+#else
             cudaError_t cuda_status =
                 cudaMemcpy(dram_buffer, addr + total_transferred_bytes_,
                            buffer_size, cudaMemcpyDefault);
+#endif
             if (cuda_status != cudaSuccess) {
                 LOG(ERROR) << "ServerSession::writeBody failed to copy from "
                               "CUDA memory. "
@@ -139,11 +214,12 @@ struct ServerSession : public std::enable_shared_from_this<ServerSession> {
 
         asio::async_write(
             *socket_, asio::buffer(dram_buffer, buffer_size),
-            [this, addr, dram_buffer, self](const asio::error_code& ec,
-                                            std::size_t transferred_bytes) {
-#if defined(USE_CUDA) || defined(USE_MUSA) || defined(USE_HIP) || \
-    defined(USE_MLU) || defined(USE_MACA)
-                if (isCudaMemory(addr)) {
+            [this, addr, dram_buffer, cuda_device, self](
+                const asio::error_code& ec, std::size_t transferred_bytes) {
+#if defined(USE_CUDA) || defined(USE_MUSA) || defined(USE_HIP) ||  \
+    defined(USE_MLU) || defined(USE_MACA) || defined(USE_HYGON) || \
+    defined(USE_COREX)
+                if (cuda_device >= 0) {
                     delete[] dram_buffer;
                 }
 #endif
@@ -168,7 +244,7 @@ struct ServerSession : public std::enable_shared_from_this<ServerSession> {
         char* addr = local_buffer_;
 
         size_t buffer_size =
-            std::min(kDefaultBufferSize, size - total_transferred_bytes_);
+            std::min(getChunkSize(), size - total_transferred_bytes_);
         if (buffer_size == 0) {
             session_mutex_.unlock();
             // Transfer complete, wait for next request on this connection
@@ -177,20 +253,20 @@ struct ServerSession : public std::enable_shared_from_this<ServerSession> {
         }
 
         char* dram_buffer = addr + total_transferred_bytes_;
+        int cuda_device = -1;
 
-#if defined(USE_CUDA) || defined(USE_MUSA) || defined(USE_HIP) || \
-    defined(USE_MLU) || defined(USE_MACA)
-        bool is_cuda_memory = isCudaMemory(addr);
-        if (is_cuda_memory) {
+#if defined(USE_CUDA) || defined(USE_MUSA) || defined(USE_HIP) ||  \
+    defined(USE_MLU) || defined(USE_MACA) || defined(USE_HYGON) || \
+    defined(USE_COREX)
+        cuda_device = getCudaDeviceId(addr);
+        if (cuda_device >= 0) {
             dram_buffer = new char[buffer_size];
         }
-#else
-        bool is_cuda_memory = false;
 #endif
 
         asio::async_read(
             *socket_, asio::buffer(dram_buffer, buffer_size),
-            [this, addr, dram_buffer, is_cuda_memory, self](
+            [this, addr, dram_buffer, cuda_device, self](
                 const asio::error_code& ec, std::size_t transferred_bytes) {
                 if (ec) {
                     // If client closed connection (EOF), this is normal - don't
@@ -205,19 +281,24 @@ struct ServerSession : public std::enable_shared_from_this<ServerSession> {
                             << " (value: " << ec.value() << ")";
                     }
                     session_mutex_.unlock();
-#if defined(USE_CUDA) || defined(USE_MUSA) || defined(USE_HIP) || \
-    defined(USE_MLU) || defined(USE_MACA)
-                    if (is_cuda_memory) delete[] dram_buffer;
-#endif
+                    if (cuda_device >= 0) delete[] dram_buffer;
                     return;  // Connection will be closed
                 }
 
-#if defined(USE_CUDA) || defined(USE_MUSA) || defined(USE_HIP) || \
-    defined(USE_MLU) || defined(USE_MACA)
-                if (is_cuda_memory) {
+#if defined(USE_CUDA) || defined(USE_MUSA) || defined(USE_HIP) ||  \
+    defined(USE_MLU) || defined(USE_MACA) || defined(USE_HYGON) || \
+    defined(USE_COREX)
+                if (cuda_device >= 0) {
+                    cudaSetDevice(cuda_device);
+#ifdef USE_MACA
+                    cudaError_t cuda_status =
+                        copyTcpCudaMemory(addr + total_transferred_bytes_,
+                                          dram_buffer, transferred_bytes);
+#else
                     cudaError_t cuda_status =
                         cudaMemcpy(addr + total_transferred_bytes_, dram_buffer,
                                    transferred_bytes, cudaMemcpyDefault);
+#endif
                     if (cuda_status != cudaSuccess) {
                         LOG(ERROR)
                             << "ServerSession::readBody failed to copy to CUDA "
@@ -268,12 +349,19 @@ struct ClientSession : public std::enable_shared_from_this<ClientSession> {
             *socket_, asio::buffer(&header_, sizeof(SessionHeader)),
             [this, self](const asio::error_code& ec, std::size_t len) {
                 if (ec || len != sizeof(SessionHeader)) {
-                    LOG(ERROR) << "ClientSession::writeHeader failed. Error: "
-                               << ec.message() << " (value: " << ec.value()
-                               << ")" << ", bytes written: " << len;
-                    if (on_finalize_) on_finalize_(TransferStatusEnum::FAILED);
-                    session_mutex_.unlock();
-                    if (on_complete_) on_complete_();
+                    LOG(ERROR)
+                        << "ClientSession::writeHeader failed. Error: "
+                        << ec.message() << " (value: " << ec.value() << ")"
+                        << ", bytes written: " << len;
+                    asio::post(
+                        socket_->get_executor(),
+                        [this, self, on_finalize = std::move(on_finalize_),
+                         on_complete = std::move(on_complete_)]() {
+                            if (on_finalize)
+                                on_finalize(TransferStatusEnum::FAILED);
+                            session_mutex_.unlock();
+                            if (on_complete) on_complete();
+                        });
                     return;
                 }
                 if (header_.opcode == (uint8_t)TransferRequest::WRITE)
@@ -289,29 +377,34 @@ struct ClientSession : public std::enable_shared_from_this<ClientSession> {
         char* addr = local_buffer_;
 
         size_t buffer_size =
-            std::min(kDefaultBufferSize, size - total_transferred_bytes_);
+            std::min(getChunkSize(), size - total_transferred_bytes_);
         if (buffer_size == 0) {
-            if (on_finalize_) on_finalize_(TransferStatusEnum::COMPLETED);
-            session_mutex_.unlock();
-            if (on_complete_) on_complete_();
+            asio::post(socket_->get_executor(),
+                       [this, self, on_finalize = std::move(on_finalize_),
+                        on_complete = std::move(on_complete_)]() {
+                           if (on_finalize)
+                               on_finalize(TransferStatusEnum::COMPLETED);
+                           session_mutex_.unlock();
+                           if (on_complete) on_complete();
+                       });
             return;
         }
 
         char* dram_buffer = addr + total_transferred_bytes_;
+        int cuda_device = -1;
 
-#if defined(USE_CUDA) || defined(USE_MUSA) || defined(USE_HIP) || \
-    defined(USE_MLU) || defined(USE_MACA)
-        bool is_cuda_memory = isCudaMemory(addr);
-        if (is_cuda_memory) {
+#if defined(USE_CUDA) || defined(USE_MUSA) || defined(USE_HIP) ||  \
+    defined(USE_MLU) || defined(USE_MACA) || defined(USE_HYGON) || \
+    defined(USE_COREX)
+        cuda_device = getCudaDeviceId(addr);
+        if (cuda_device >= 0) {
             dram_buffer = new char[buffer_size];
         }
-#else
-        bool is_cuda_memory = false;
 #endif
 
         asio::async_read(
             *socket_, asio::buffer(dram_buffer, buffer_size),
-            [this, addr, dram_buffer, is_cuda_memory, self](
+            [this, addr, dram_buffer, cuda_device, self](
                 const asio::error_code& ec, std::size_t transferred_bytes) {
                 if (ec) {
                     LOG(ERROR)
@@ -320,32 +413,57 @@ struct ClientSession : public std::enable_shared_from_this<ClientSession> {
                         << " using buffer " << static_cast<void*>(dram_buffer)
                         << ". Error: " << ec.message()
                         << " (value: " << ec.value() << ")";
-                    if (on_finalize_) on_finalize_(TransferStatusEnum::FAILED);
-                    if (on_complete_) on_complete_();
-#if defined(USE_CUDA) || defined(USE_MUSA) || defined(USE_HIP) || \
-    defined(USE_MLU) || defined(USE_MACA)
-                    if (is_cuda_memory) delete[] dram_buffer;
+                    // Post entire cleanup to ensure it runs after callback
+                    // returns
+                    asio::post(socket_->get_executor(),
+                               [this, self, dram_buffer, cuda_device,
+                                on_finalize = std::move(on_finalize_),
+                                on_complete = std::move(on_complete_)]() {
+                                   if (on_finalize)
+                                       on_finalize(TransferStatusEnum::FAILED);
+#if defined(USE_CUDA) || defined(USE_MUSA) || defined(USE_HIP) ||  \
+    defined(USE_MLU) || defined(USE_MACA) || defined(USE_HYGON) || \
+    defined(USE_COREX)
+                                   if (cuda_device >= 0) delete[] dram_buffer;
 #endif
-                    session_mutex_.unlock();
+                                   session_mutex_.unlock();
+                                   if (on_complete) on_complete();
+                               });
                     return;
                 }
 
-#if defined(USE_CUDA) || defined(USE_MUSA) || defined(USE_HIP) || \
-    defined(USE_MLU) || defined(USE_MACA)
-                if (is_cuda_memory) {
+#if defined(USE_CUDA) || defined(USE_MUSA) || defined(USE_HIP) ||  \
+    defined(USE_MLU) || defined(USE_MACA) || defined(USE_HYGON) || \
+    defined(USE_COREX)
+                if (cuda_device >= 0) {
+                    cudaSetDevice(cuda_device);
+#ifdef USE_MACA
+                    cudaError_t cuda_status =
+                        copyTcpCudaMemory(addr + total_transferred_bytes_,
+                                          dram_buffer, transferred_bytes);
+#else
                     cudaError_t cuda_status =
                         cudaMemcpy(addr + total_transferred_bytes_, dram_buffer,
                                    transferred_bytes, cudaMemcpyDefault);
+#endif
                     if (cuda_status != cudaSuccess) {
                         LOG(ERROR)
                             << "ClientSession::readBody failed to copy to CUDA "
                                "memory. "
                             << "Error: " << cudaGetErrorString(cuda_status);
-                        if (on_finalize_)
-                            on_finalize_(TransferStatusEnum::FAILED);
-                        if (on_complete_) on_complete_();
-                        delete[] dram_buffer;
-                        session_mutex_.unlock();
+                        // Post entire cleanup to ensure it runs after callback
+                        // returns
+                        asio::post(
+                            socket_->get_executor(),
+                            [this, self, dram_buffer,
+                             on_finalize = std::move(on_finalize_),
+                             on_complete = std::move(on_complete_)]() {
+                                if (on_finalize)
+                                    on_finalize(TransferStatusEnum::FAILED);
+                                delete[] dram_buffer;
+                                session_mutex_.unlock();
+                                if (on_complete) on_complete();
+                            });
                         return;
                     }
                     delete[] dram_buffer;
@@ -362,31 +480,53 @@ struct ClientSession : public std::enable_shared_from_this<ClientSession> {
         char* addr = local_buffer_;
 
         size_t buffer_size =
-            std::min(kDefaultBufferSize, size - total_transferred_bytes_);
+            std::min(getChunkSize(), size - total_transferred_bytes_);
         if (buffer_size == 0) {
-            if (on_finalize_) on_finalize_(TransferStatusEnum::COMPLETED);
-            session_mutex_.unlock();
-            if (on_complete_) on_complete_();
+            // Post cleanup to ensure it runs after callback returns
+            asio::post(socket_->get_executor(),
+                       [this, self, on_finalize = std::move(on_finalize_),
+                        on_complete = std::move(on_complete_)]() {
+                           if (on_finalize)
+                               on_finalize(TransferStatusEnum::COMPLETED);
+                           session_mutex_.unlock();
+                           if (on_complete) on_complete();
+                       });
             return;
         }
 
         char* dram_buffer = addr + total_transferred_bytes_;
+        int cuda_device = -1;
 
-#if defined(USE_CUDA) || defined(USE_MUSA) || defined(USE_HIP) || \
-    defined(USE_MLU) || defined(USE_MACA)
-        if (isCudaMemory(addr)) {
+#if defined(USE_CUDA) || defined(USE_MUSA) || defined(USE_HIP) ||  \
+    defined(USE_MLU) || defined(USE_MACA) || defined(USE_HYGON) || \
+    defined(USE_COREX)
+        cuda_device = getCudaDeviceId(addr);
+        if (cuda_device >= 0) {
             dram_buffer = new char[buffer_size];
+            cudaSetDevice(cuda_device);
+#ifdef USE_MACA
+            cudaError_t cuda_status = copyTcpCudaMemory(
+                dram_buffer, addr + total_transferred_bytes_, buffer_size);
+#else
             cudaError_t cuda_status =
                 cudaMemcpy(dram_buffer, addr + total_transferred_bytes_,
                            buffer_size, cudaMemcpyDefault);
+#endif
             if (cuda_status != cudaSuccess) {
                 LOG(ERROR) << "ClientSession::writeBody failed to copy from "
                               "CUDA memory. "
                            << "Error: " << cudaGetErrorString(cuda_status);
-                if (on_finalize_) on_finalize_(TransferStatusEnum::FAILED);
-                if (on_complete_) on_complete_();
-                session_mutex_.unlock();
-                delete[] dram_buffer;
+                // Post entire cleanup to ensure it runs after callback returns
+                asio::post(socket_->get_executor(),
+                           [this, self, dram_buffer,
+                            on_finalize = std::move(on_finalize_),
+                            on_complete = std::move(on_complete_)]() {
+                               if (on_finalize)
+                                   on_finalize(TransferStatusEnum::FAILED);
+                               delete[] dram_buffer;
+                               session_mutex_.unlock();
+                               if (on_complete) on_complete();
+                           });
                 return;
             }
         }
@@ -394,14 +534,11 @@ struct ClientSession : public std::enable_shared_from_this<ClientSession> {
 
         asio::async_write(
             *socket_, asio::buffer(dram_buffer, buffer_size),
-            [this, addr, dram_buffer, self](const asio::error_code& ec,
-                                            std::size_t transferred_bytes) {
-#if defined(USE_CUDA) || defined(USE_MUSA) || defined(USE_HIP) || \
-    defined(USE_MLU) || defined(USE_MACA)
-                if (isCudaMemory(addr)) {
+            [this, addr, dram_buffer, cuda_device, self](
+                const asio::error_code& ec, std::size_t transferred_bytes) {
+                if (cuda_device >= 0) {
                     delete[] dram_buffer;
                 }
-#endif
                 if (ec) {
                     LOG(ERROR)
                         << "ClientSession::writeBody failed. "
@@ -409,9 +546,17 @@ struct ClientSession : public std::enable_shared_from_this<ClientSession> {
                         << " using buffer " << static_cast<void*>(dram_buffer)
                         << ". Error: " << ec.message()
                         << " (value: " << ec.value() << ")";
-                    if (on_finalize_) on_finalize_(TransferStatusEnum::FAILED);
-                    if (on_complete_) on_complete_();
-                    session_mutex_.unlock();
+                    // Post entire cleanup to ensure it runs after callback
+                    // returns
+                    asio::post(
+                        socket_->get_executor(),
+                        [this, self, on_finalize = std::move(on_finalize_),
+                         on_complete = std::move(on_complete_)]() {
+                            if (on_finalize)
+                                on_finalize(TransferStatusEnum::FAILED);
+                            session_mutex_.unlock();
+                            if (on_complete) on_complete();
+                        });
                     return;
                 }
                 total_transferred_bytes_ += transferred_bytes;
@@ -421,7 +566,8 @@ struct ClientSession : public std::enable_shared_from_this<ClientSession> {
 };
 
 struct TcpContext {
-    TcpContext(short port) : acceptor(io_context) {
+    TcpContext(short port, ValidateAddrFn validate_addr)
+        : acceptor(io_context), validate_addr_(std::move(validate_addr)) {
         std::error_code ec;
         asio::ip::tcp::endpoint endpoint(asio::ip::tcp::v6(), port);
 
@@ -451,11 +597,13 @@ struct TcpContext {
     void doAccept() {
         acceptor.async_accept([this](asio::error_code ec, tcpsocket socket) {
             if (!ec) {
+                asio::error_code nodelay_ec;
+                socket.set_option(asio::ip::tcp::no_delay(true), nodelay_ec);
                 auto socket_ptr =
                     std::make_shared<tcpsocket>(std::move(socket));
-                auto session = std::make_shared<ServerSession>(socket_ptr);
-                session->start();  // Start processing requests on this
-                                   // persistent connection
+                auto session =
+                    std::make_shared<ServerSession>(socket_ptr, validate_addr_);
+                session->start();
             }
             doAccept();
         });
@@ -463,12 +611,14 @@ struct TcpContext {
 
     asio::io_context io_context;
     asio::ip::tcp::acceptor acceptor;
+    ValidateAddrFn validate_addr_;
 };
 
 TcpTransport::TcpTransport() : context_(nullptr), running_(false) {
     if (getenv("MC_TCP_ENABLE_CONNECTION_POOL") != nullptr) {
         std::string val(getenv("MC_TCP_ENABLE_CONNECTION_POOL"));
-        std::transform(val.begin(), val.end(), val.begin(), ::tolower);
+        std::transform(val.begin(), val.end(), val.begin(),
+                       [](unsigned char c) -> char { return std::tolower(c); });
         if (val == "0" || val == "false" || val == "no") {
             enable_connection_pool_ = false;
         } else {
@@ -539,7 +689,9 @@ int TcpTransport::install(std::string& local_server_name,
 
     close(sockfd);  // the above function has opened a socket
     LOG(INFO) << "TcpTransport: listen on port " << tcp_port;
-    context_ = new TcpContext(tcp_port);
+    context_ = new TcpContext(tcp_port, [this](uint64_t addr, uint64_t size) {
+        return validateAddress(addr, size);
+    });
     running_ = true;
     thread_ = std::thread(&TcpTransport::worker, this);
     return 0;
@@ -688,6 +840,7 @@ void TcpTransport::worker() {
             LOG(ERROR) << "TcpTransport::worker encountered an exception "
                           "during doAccept/run: "
                        << e.what();
+            context_->io_context.restart();
         }
     }
 }
@@ -703,6 +856,7 @@ std::shared_ptr<asio::ip::tcp::socket> TcpTransport::getConnection(
             auto socket_ptr =
                 std::make_shared<asio::ip::tcp::socket>(context_->io_context);
             asio::connect(*socket_ptr, endpoint_iterator);
+            socket_ptr->set_option(asio::ip::tcp::no_delay(true));
             return socket_ptr;
         } catch (std::exception& e) {
             LOG(ERROR)
@@ -755,6 +909,7 @@ std::shared_ptr<asio::ip::tcp::socket> TcpTransport::getConnection(
         new_socket =
             std::make_shared<asio::ip::tcp::socket>(context_->io_context);
         asio::connect(*new_socket, endpoint_iterator);
+        new_socket->set_option(asio::ip::tcp::no_delay(true));
     } catch (std::exception& e) {
         LOG(ERROR)
             << "TcpTransport::getConnection failed to create connection to "
@@ -829,9 +984,8 @@ void TcpTransport::cleanupIdleConnections() {
     for (auto it = connection_pool_.begin(); it != connection_pool_.end();) {
         auto& queue = it->second;
 
-        // Remove idle connections that exceed timeout
-        while (!queue.empty()) {
-            auto& entry = queue.back();
+        for (auto entry_it = queue.begin(); entry_it != queue.end();) {
+            auto& entry = *entry_it;
             if (!entry->in_use) {
                 auto idle_duration =
                     std::chrono::duration_cast<std::chrono::seconds>(
@@ -842,16 +996,13 @@ void TcpTransport::cleanupIdleConnections() {
                         asio::error_code ec;
                         entry->socket->close(ec);
                     }
-                    queue.pop_back();
-                } else {
-                    break;
+                    entry_it = queue.erase(entry_it);
+                    continue;
                 }
-            } else {
-                break;
             }
+            ++entry_it;
         }
 
-        // Remove empty endpoint queues
         if (queue.empty()) {
             it = connection_pool_.erase(it);
         } else {
@@ -860,6 +1011,21 @@ void TcpTransport::cleanupIdleConnections() {
     }
 }
 
+bool TcpTransport::validateAddress(uint64_t addr, uint64_t size) const {
+    if (size == 0) return false;
+    if (addr + size < addr) return false;
+
+    auto desc = metadata_->getSegmentDescByID(LOCAL_SEGMENT_ID);
+    if (!desc) return false;
+
+    for (const auto& buffer : desc->buffers) {
+        if (buffer.addr + buffer.length < buffer.addr) continue;
+        if (buffer.addr <= addr && addr + size <= buffer.addr + buffer.length)
+            return true;
+    }
+    return false;
+}
+
 void TcpTransport::startTransfer(Slice* slice) {
     auto desc = metadata_->getSegmentDescByID(slice->target_id);
     if (!desc) {
@@ -933,7 +1099,8 @@ void TcpTransport::startTransfer(Slice* slice) {
         }
         if (enable_connection_pool_) {
             // Remove the connection from pool if it was pooled
-            ConnectionKey key{meta_entry.ip_or_host_name, desc->tcp_data_port};
+            ConnectionKey key{meta_entry.ip_or_host_name,
+                              static_cast<uint16_t>(desc->tcp_data_port)};
             std::lock_guard<std::mutex> lock(pool_mutex_);
             auto it = connection_pool_.find(key);
             if (it != connection_pool_.end()) {
diff --git a/mooncake-transfer-engine/tent/CMakeLists.txt b/mooncake-transfer-engine/tent/CMakeLists.txt
index 82bced27..b5bcc379 100644
--- a/mooncake-transfer-engine/tent/CMakeLists.txt
+++ b/mooncake-transfer-engine/tent/CMakeLists.txt
@@ -8,4 +8,6 @@ set(CMAKE_POSITION_INDEPENDENT_CODE ON)
 
 add_subdirectory(src)
 add_subdirectory(plugins)
-add_subdirectory(tests)
+if (BUILD_UNIT_TESTS)
+    add_subdirectory(tests)
+endif()
diff --git a/mooncake-transfer-engine/tent/config/transfer-engine.json b/mooncake-transfer-engine/tent/config/transfer-engine.json
index 1fd2a5eb..bfe46b2e 100644
--- a/mooncake-transfer-engine/tent/config/transfer-engine.json
+++ b/mooncake-transfer-engine/tent/config/transfer-engine.json
@@ -2,7 +2,6 @@
     "local_segment_name": "",
     "metadata_type": "p2p",
     "metadata_servers": "127.0.0.1:2379",
-    "redis_db_index": 0,
     "rpc_server_hostname": "127.0.0.1",
     "rpc_server_port": 0,
     "topology": {
@@ -11,6 +10,8 @@
     },
     "log_level": "warning",
     "max_failover_attempts": 3,
+    "enable_auto_failover_on_poll": true,
+    "enable_progress_worker": false,
     "metrics": {
         "enabled": true,
         "http_port": 9100,
@@ -88,5 +89,18 @@
         "mnnvl": {
             "enable" : false
         }
-    }
+    },
+    "policy": [
+        {
+            "name": "default_memory",
+            "segment_type": "memory",
+            "devices": ["mlx5_0", "mlx5_2"],
+            "transports": ["nvlink", "rdma", "shm"]
+        },
+        {
+            "name": "file_storage",
+            "segment_type": "file",
+            "transports": ["gds", "io_uring", "rdma"]
+        }
+    ]
 }
diff --git a/mooncake-transfer-engine/tent/include/tent/common/status.h b/mooncake-transfer-engine/tent/include/tent/common/status.h
index f90b3e84..67049821 100644
--- a/mooncake-transfer-engine/tent/include/tent/common/status.h
+++ b/mooncake-transfer-engine/tent/include/tent/common/status.h
@@ -60,6 +60,16 @@
     } while (0)
 #endif
 
+#ifdef USE_HIP
+#define CHECK_HIP(call)                                                      \
+    do {                                                                     \
+        auto err = call;                                                     \
+        if (err != hipSuccess)                                               \
+            return Status::InternalError(std::string(#call) + ": " +         \
+                                         hipGetErrorString(err) + LOC_MARK); \
+    } while (0)
+#endif
+
 #if defined(USE_ASCEND) || defined(USE_ASCEND_DIRECT)
 #define CHECK_ASCEND(call)                                                \
     do {                                                                  \
diff --git a/mooncake-transfer-engine/tent/include/tent/common/types.h b/mooncake-transfer-engine/tent/include/tent/common/types.h
index da862f42..27a7e98d 100644
--- a/mooncake-transfer-engine/tent/include/tent/common/types.h
+++ b/mooncake-transfer-engine/tent/include/tent/common/types.h
@@ -18,6 +18,7 @@
 #include <cstddef>
 #include <cstdint>
 #include <memory>
+#include <optional>
 #include <string>
 #include <unordered_map>
 #include <vector>
@@ -28,6 +29,11 @@ namespace tent {
 using BatchID = uint64_t;
 using SegmentID = uint64_t;
 
+// QoS priority levels
+static constexpr uint8_t PRIO_HIGH = 0;
+static constexpr uint8_t PRIO_MEDIUM = 1;
+static constexpr uint8_t PRIO_LOW = 2;
+
 struct Notification {
     std::string name;
     std::string msg;
@@ -37,6 +43,28 @@ struct Notification {
 #define LOCAL_SEGMENT_ID (0ull)
 #endif
 
+enum TransportType : int {
+    UNSPEC = 0,
+    RDMA,
+    MNNVL,
+    SHM,
+    NVLINK,
+    GDS,
+    IOURING,
+    TCP,
+    AscendDirect,
+    SUNRISE_LINK,
+    // Sentinel: must remain the last enumerator.
+    kNumTransportTypes,
+};
+
+const static int kSupportedTransportTypes = (int)kNumTransportTypes;
+
+inline TransportType c_to_transport_hint(int v) {
+    if (v < 0 || v >= kSupportedTransportTypes) return UNSPEC;
+    return static_cast<TransportType>(v);
+}
+
 struct Request {
     enum OpCode { READ, WRITE };
     OpCode opcode;
@@ -44,6 +72,19 @@ struct Request {
     SegmentID target_id;
     uint64_t target_offset;
     size_t length;
+    int priority =
+        PRIO_HIGH;  // Request priority (PRIO_HIGH, PRIO_MEDIUM, PRIO_LOW)
+    std::optional<std::string>
+        policy_name;  // Optional: bind to specific policy by name
+    TransportType transport_hint =
+        UNSPEC;  // UNSPEC = follow policy; otherwise pin this request to the
+                 // name transport.
+    // Optional SLO deadline as an absolute steady_clock timestamp in
+    // nanoseconds. 0 = no deadline (default), behaves exactly as today.
+    // When set, the engine emits an observability-only feasibility metric
+    // (MLU = actual transfer time / available window) on completion; it does
+    // not yet drive any admission or scheduling decision. See RFC #2519.
+    uint64_t deadline_ns = 0;
 };
 
 enum TransferStatusEnum {
@@ -70,19 +111,6 @@ enum Permission {
 using Location = std::string;
 const static std::string kWildcardLocation = "*";
 
-enum TransportType {
-    RDMA = 0,
-    MNNVL,
-    SHM,
-    NVLINK,
-    GDS,
-    IOURING,
-    TCP,
-    AscendDirect,
-    UNSPEC
-};
-const static int kSupportedTransportTypes = 8;
-
 struct MemoryOptions {
     Location location = kWildcardLocation;
     Permission perm = kGlobalReadWrite;
@@ -107,4 +135,4 @@ struct SegmentInfo {
 }  // namespace tent
 }  // namespace mooncake
 
-#endif  // TENT_TYPES_H
\ No newline at end of file
+#endif  // TENT_TYPES_H
diff --git a/mooncake-transfer-engine/tent/include/tent/metrics/tent_metrics.h b/mooncake-transfer-engine/tent/include/tent/metrics/tent_metrics.h
index 90162e11..4abc12fa 100644
--- a/mooncake-transfer-engine/tent/include/tent/metrics/tent_metrics.h
+++ b/mooncake-transfer-engine/tent/include/tent/metrics/tent_metrics.h
@@ -81,6 +81,12 @@ class TentMetrics {
     void recordWriteFailed(size_t bytes);
     void recordTransportFailover();
 
+    // Record the deadline feasibility ratio (MLU) for a completed transfer
+    // that carried a deadline. mlu = actual_transfer_seconds / window_seconds,
+    // where window_seconds is (deadline - submit_time). mlu < 1 means the
+    // transfer met its deadline; mlu >= 1 means it missed. Observability only.
+    void recordDeadlineMLU(double mlu);
+
     // Get metrics for HTTP server
     std::string getPrometheusMetrics();
     std::string getJsonMetrics();
@@ -161,6 +167,17 @@ class TentMetrics {
         "tent_write_size_bytes", "Write request size distribution in bytes",
         kSizeBuckets};
 
+    // Deadline feasibility ratio (MLU) distribution for transfers that carried
+    // a deadline. Stored in per-mille (MLU x 1000) so the histogram can use
+    // integer observe() like the others; the 1000 boundary is MLU == 1.0, the
+    // feasible/infeasible line (< 1000 met the deadline, >= 1000 missed it).
+    static inline const std::vector<double> kMluPerMilleBuckets{
+        100, 250, 500, 750, 900, 1000, 1250, 1500, 2000, 5000};
+    ylt::metric::histogram_t deadline_mlu_{
+        "tent_deadline_mlu_permille",
+        "Deadline feasibility ratio (MLU x 1000) distribution",
+        kMluPerMilleBuckets};
+
     // Helper to register all metrics to the vectors
     void registerMetrics();
 #endif  // TENT_METRICS_ENABLED
diff --git a/mooncake-transfer-engine/tent/include/tent/platform/rocm.h b/mooncake-transfer-engine/tent/include/tent/platform/rocm.h
new file mode 100644
index 00000000..5ef0a356
--- /dev/null
+++ b/mooncake-transfer-engine/tent/include/tent/platform/rocm.h
@@ -0,0 +1,131 @@
+// Copyright 2025 KVCache.AI
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#ifndef ROCM_H
+#define ROCM_H
+
+#include "tent/runtime/platform.h"
+#include "tent/common/config.h"
+#include "tent/common/concurrent/rw_spinlock.h"
+
+#include <hip/hip_runtime.h>
+
+namespace mooncake {
+namespace tent {
+
+class HIPStreamPool;
+
+// RAII wrapper for hipStream_t.
+// When this handle destructs, the stream is automatically returned to the pool.
+class HIPStreamHandle {
+   public:
+    HIPStreamHandle() = default;
+
+    HIPStreamHandle(hipStream_t stream, int deviceId, HIPStreamPool* pool)
+        : stream_(stream), deviceId_(deviceId), pool_(pool) {}
+
+    HIPStreamHandle(const HIPStreamHandle&) = delete;
+    HIPStreamHandle& operator=(const HIPStreamHandle&) = delete;
+
+    HIPStreamHandle(HIPStreamHandle&& other) noexcept
+        : stream_(other.stream_),
+          deviceId_(other.deviceId_),
+          pool_(other.pool_) {
+        other.stream_ = nullptr;
+        other.pool_ = nullptr;
+    }
+
+    HIPStreamHandle& operator=(HIPStreamHandle&& other) noexcept;
+    ~HIPStreamHandle();
+
+    [[nodiscard]] hipStream_t get() const { return stream_; }
+
+   private:
+    void releaseToPool();
+
+    hipStream_t stream_ = nullptr;
+    int deviceId_ = -1;
+    HIPStreamPool* pool_ = nullptr;
+};
+
+// HIP Stream Pool managing all devices.
+class HIPStreamPool {
+    friend class HIPStreamHandle;
+
+   public:
+    HIPStreamPool() = default;
+    ~HIPStreamPool() = default;
+
+    HIPStreamPool(const HIPStreamPool&) = delete;
+    HIPStreamPool& operator=(const HIPStreamPool&) = delete;
+
+    static constexpr int kCurrentDevice = -1;
+    Status acquire(HIPStreamHandle& outHandle, int deviceId = kCurrentDevice);
+
+   private:
+    class DevicePool {
+       public:
+        explicit DevicePool(int deviceId);
+        ~DevicePool();
+
+        Status acquire(hipStream_t& outStream);
+        void release(hipStream_t stream);
+
+       private:
+        int deviceId_;
+        RWSpinlock dev_lock_;
+        std::vector<hipStream_t> availableStreams_;
+    };
+
+    void release(int deviceId, hipStream_t stream);
+    DevicePool* getDevicePool(int deviceId);
+
+    RWSpinlock pools_lock_;
+    std::vector<std::unique_ptr<DevicePool>> devicePools_;
+};
+
+class RocmPlatform : public Platform {
+   public:
+    RocmPlatform(std::shared_ptr<Config> config) : conf(std::move(config)) {}
+
+    virtual ~RocmPlatform() {}
+
+    virtual Status probe(std::vector<Topology::NicEntry>& nic_list,
+                         std::vector<Topology::MemEntry>& mem_list);
+
+    virtual Status allocate(void** pptr, size_t size, MemoryOptions& options);
+
+    virtual Status free(void* ptr, size_t size);
+
+    virtual Status copy(void* dst, void* src, size_t length);
+
+    virtual MemoryType getMemoryType(void* addr);
+
+    virtual const std::vector<RangeLocation> getLocation(
+        void* start, size_t len, bool skip_prefault = false);
+
+    virtual const std::string type() const { return "rocm"; }
+
+    Status getStreamFromPool(HIPStreamHandle& outHandle,
+                             int deviceId = HIPStreamPool::kCurrentDevice);
+
+   private:
+    std::shared_ptr<Config> conf;
+    HIPStreamPool stream_pool_;
+};
+
+}  // namespace tent
+}  // namespace mooncake
+
+#endif  // ROCM_H
diff --git a/mooncake-transfer-engine/tent/include/tent/platform/sunrise.h b/mooncake-transfer-engine/tent/include/tent/platform/sunrise.h
new file mode 100644
index 00000000..14df6c06
--- /dev/null
+++ b/mooncake-transfer-engine/tent/include/tent/platform/sunrise.h
@@ -0,0 +1,53 @@
+// Copyright 2026 KVCache.AI
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#ifndef SUNRISE_H
+#define SUNRISE_H
+
+#include "tent/common/config.h"
+#include "tent/runtime/platform.h"
+
+namespace mooncake {
+namespace tent {
+
+class SunrisePlatform : public Platform {
+   public:
+    SunrisePlatform(std::shared_ptr<Config> config) : conf(std::move(config)) {}
+
+    virtual ~SunrisePlatform() {}
+
+    virtual Status probe(std::vector<Topology::NicEntry>& nic_list,
+                         std::vector<Topology::MemEntry>& mem_list);
+
+    virtual Status allocate(void** pptr, size_t size, MemoryOptions& options);
+
+    virtual Status free(void* ptr, size_t size);
+
+    virtual Status copy(void* dst, void* src, size_t length);
+
+    virtual MemoryType getMemoryType(void* addr);
+
+    virtual const std::vector<RangeLocation> getLocation(
+        void* start, size_t len, bool skip_prefault = false);
+
+    virtual const std::string type() const { return "sunrise"; }
+
+   private:
+    std::shared_ptr<Config> conf;
+};
+
+}  // namespace tent
+}  // namespace mooncake
+
+#endif  // SUNRISE_H
diff --git a/mooncake-transfer-engine/tent/include/tent/runtime/admission_queue.h b/mooncake-transfer-engine/tent/include/tent/runtime/admission_queue.h
new file mode 100644
index 00000000..e52f7060
--- /dev/null
+++ b/mooncake-transfer-engine/tent/include/tent/runtime/admission_queue.h
@@ -0,0 +1,125 @@
+// Copyright 2026 KVCache.AI
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#ifndef ADMISSION_QUEUE_H_
+#define ADMISSION_QUEUE_H_
+
+#include <cstddef>
+#include <cstdint>
+#include <deque>
+#include <map>
+#include <utility>
+#include <vector>
+
+#include "tent/common/status.h"
+#include "tent/common/types.h"
+
+namespace mooncake {
+namespace tent {
+
+using QueueOwnerId = uint64_t;
+
+// Owner kind is used only for admission accounting. It is not a dispatch
+// priority.
+enum class QueueOwnerKind {
+    User,
+    StagingInternal,
+};
+
+struct QueueLimits {
+    size_t max_outstanding_owners{0};
+    size_t max_outstanding_bytes{0};
+    size_t staging_owner_reserve{0};
+    size_t staging_byte_reserve{0};
+};
+
+struct QueueOwnerInput {
+    // Absolute task id within the caller's Batch, not relative to this submit.
+    size_t owner_task_id{0};
+    std::vector<size_t> derived_task_ids;
+    Request request{};
+    QueueOwnerKind kind{QueueOwnerKind::User};
+};
+
+struct QueueSubmit {
+    uint64_t batch_token{0};
+    // Caller-computed remaining public task slots for this submit.
+    size_t batch_slots_left{0};
+    std::vector<QueueOwnerInput> owners;
+};
+
+// Runtime-private admission model. It is intentionally single-threaded; the
+// eventual TransferEngineImpl integration owns synchronization.
+class LocalTransferAdmissionQueue {
+   public:
+    explicit LocalTransferAdmissionQueue(QueueLimits limits);
+
+    LocalTransferAdmissionQueue(const LocalTransferAdmissionQueue&) = delete;
+    LocalTransferAdmissionQueue& operator=(const LocalTransferAdmissionQueue&) =
+        delete;
+    LocalTransferAdmissionQueue(LocalTransferAdmissionQueue&&) = delete;
+    LocalTransferAdmissionQueue& operator=(LocalTransferAdmissionQueue&&) =
+        delete;
+
+    Status tryAdmit(const QueueSubmit& submit,
+                    std::vector<QueueOwnerId>& admitted_owner_ids);
+
+    std::vector<QueueOwnerId> pickForDispatch(size_t max_owners,
+                                              size_t max_bytes);
+
+    Status complete(QueueOwnerId owner_id, TransferStatusEnum terminal_status);
+
+    Status retireBatch(uint64_t batch_token);
+
+    Status resolveOwner(uint64_t batch_token, size_t public_task_id,
+                        QueueOwnerId& owner_id) const;
+
+    Status getPublicStatus(uint64_t batch_token, size_t public_task_id,
+                           TransferStatusEnum& status) const;
+
+    size_t outstandingOwners() const;
+
+    size_t outstandingBytes() const;
+
+   private:
+    enum class QueueState {
+        Queued,
+        Dispatching,
+        Terminal,
+    };
+
+    struct QueueOwner {
+        uint64_t batch_token{0};
+        Request request{};
+        QueueOwnerKind kind{QueueOwnerKind::User};
+        QueueState state{QueueState::Queued};
+        TransferStatusEnum terminal_status{TransferStatusEnum::PENDING};
+    };
+
+    QueueLimits limits_;
+    Status limits_status_;
+    QueueOwnerId next_owner_id_{1};
+    std::map<QueueOwnerId, QueueOwner> owners_;
+    std::map<std::pair<uint64_t, size_t>, QueueOwnerId> public_to_owner_;
+    std::deque<QueueOwnerId> fifo_;
+    size_t outstanding_owners_{0};
+    size_t outstanding_bytes_{0};
+    size_t outstanding_user_owners_{0};
+    size_t outstanding_user_bytes_{0};
+};
+
+}  // namespace tent
+}  // namespace mooncake
+
+#endif  // ADMISSION_QUEUE_H_
diff --git a/mooncake-transfer-engine/tent/include/tent/runtime/control_plane.h b/mooncake-transfer-engine/tent/include/tent/runtime/control_plane.h
index e7ce7ffe..a4739f27 100644
--- a/mooncake-transfer-engine/tent/include/tent/runtime/control_plane.h
+++ b/mooncake-transfer-engine/tent/include/tent/runtime/control_plane.h
@@ -114,14 +114,6 @@ class ControlService {
     ControlService(const std::string& type, const std::string& servers,
                    TransferEngineImpl* impl);
 
-    ControlService(const std::string& type, const std::string& servers,
-                   const std::string& password, uint8_t db_index,
-                   TransferEngineImpl* impl);
-
-    ControlService(const std::string& type, const std::string& servers,
-                   const std::string& username, const std::string& password,
-                   uint8_t db_index, TransferEngineImpl* impl);
-
     ~ControlService();
 
     ControlService(const ControlService&) = delete;
diff --git a/mooncake-transfer-engine/tent/include/tent/runtime/metastore.h b/mooncake-transfer-engine/tent/include/tent/runtime/metastore.h
index a70fdd8f..ab383b6d 100644
--- a/mooncake-transfer-engine/tent/include/tent/runtime/metastore.h
+++ b/mooncake-transfer-engine/tent/include/tent/runtime/metastore.h
@@ -26,17 +26,6 @@ struct MetaStore {
     static std::shared_ptr<MetaStore> Create(const std::string &type,
                                              const std::string &servers);
 
-    static std::shared_ptr<MetaStore> Create(const std::string &type,
-                                             const std::string &servers,
-                                             const std::string &password,
-                                             uint8_t db_index);
-
-    static std::shared_ptr<MetaStore> Create(const std::string &type,
-                                             const std::string &servers,
-                                             const std::string &username,
-                                             const std::string &password,
-                                             uint8_t db_index);
-
     MetaStore() {}
 
     virtual ~MetaStore() {}
diff --git a/mooncake-transfer-engine/tent/include/tent/runtime/platform.h b/mooncake-transfer-engine/tent/include/tent/runtime/platform.h
index 6e947463..95a4c7ae 100644
--- a/mooncake-transfer-engine/tent/include/tent/runtime/platform.h
+++ b/mooncake-transfer-engine/tent/include/tent/runtime/platform.h
@@ -20,7 +20,7 @@
 namespace mooncake {
 namespace tent {
 
-enum MemoryType { MTYPE_UNKNOWN, MTYPE_CPU, MTYPE_CUDA };
+enum MemoryType { MTYPE_UNKNOWN, MTYPE_CPU, MTYPE_CUDA, MTYPE_ROCM };
 
 class Platform {
    public:
diff --git a/mooncake-transfer-engine/tent/include/tent/runtime/progress_worker.h b/mooncake-transfer-engine/tent/include/tent/runtime/progress_worker.h
new file mode 100644
index 00000000..30fc01a0
--- /dev/null
+++ b/mooncake-transfer-engine/tent/include/tent/runtime/progress_worker.h
@@ -0,0 +1,82 @@
+// Copyright 2024 KVCache.AI
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#ifndef PROGRESS_WORKER_H_
+#define PROGRESS_WORKER_H_
+
+#include <atomic>
+#include <chrono>
+#include <condition_variable>
+#include <deque>
+#include <mutex>
+#include <thread>
+#include <unordered_set>
+
+#include "tent/common/types.h"
+
+namespace mooncake {
+namespace tent {
+
+class TransferEngineImpl;
+
+// Event-driven progress worker for issue #2116. When the engine is configured
+// with enable_progress_worker=true, transports call
+// notifyBatchMaybeReady to wake this worker, which then drives one
+// progressBatch step per notification. When the runtime queue has active work,
+// the worker also uses a low-frequency fallback tick so transports that do not
+// yet emit completion wakes can still drain the queue.
+class ProgressWorker {
+   public:
+    explicit ProgressWorker(TransferEngineImpl* impl,
+                            std::chrono::microseconds fallback_interval =
+                                std::chrono::microseconds(0));
+    ~ProgressWorker();
+
+    ProgressWorker(const ProgressWorker&) = delete;
+    ProgressWorker& operator=(const ProgressWorker&) = delete;
+
+    void start();
+
+    // Idempotent. Signals the worker thread to exit and joins it. After stop()
+    // returns, notifyBatchMaybeReady becomes a no-op.
+    void stop();
+
+    // Safe from any thread. De-duplicates: enqueueing a batch that is already
+    // queued is a no-op. No-op if the worker has been stopped or never
+    // started.
+    void notifyBatchMaybeReady(BatchID batch_id);
+
+    // Safe from any thread. Coalesces multiple runtime queue wakes into one
+    // bounded refill step.
+    void notifyRuntimeQueueReady();
+
+   private:
+    void runner();
+
+    TransferEngineImpl* impl_;
+    std::chrono::microseconds fallback_interval_;
+    std::atomic<bool> running_{false};
+    std::thread thread_;
+
+    std::mutex mu_;
+    std::condition_variable cv_;
+    std::unordered_set<BatchID> queued_;
+    std::deque<BatchID> order_;
+    bool queue_ready_{false};
+};
+
+}  // namespace tent
+}  // namespace mooncake
+
+#endif  // PROGRESS_WORKER_H_
diff --git a/mooncake-transfer-engine/tent/include/tent/runtime/proxy_manager.h b/mooncake-transfer-engine/tent/include/tent/runtime/proxy_manager.h
index 7f4e47d8..959197ca 100644
--- a/mooncake-transfer-engine/tent/include/tent/runtime/proxy_manager.h
+++ b/mooncake-transfer-engine/tent/include/tent/runtime/proxy_manager.h
@@ -31,6 +31,7 @@ struct StageBufferCache;
 
 struct StagingTask {
     TaskInfo* native{nullptr};
+    BatchID batch{0};
     std::vector<std::string> params;
 };
 
@@ -44,14 +45,15 @@ class ProxyManager {
 
    public:
     explicit ProxyManager(TransferEngineImpl* impl,
-                          size_t chunk_size = 8 * 1024 * 1024,
-                          size_t chunk_count = 32);
+                          size_t chunk_size = 4 * 1024 * 1024,
+                          size_t chunk_count = 64);
 
     ~ProxyManager();
 
     Status deconstruct();
 
-    Status submit(TaskInfo* task, const std::vector<std::string>& params);
+    Status submit(TaskInfo* task, BatchID batch,
+                  const std::vector<std::string>& params);
 
     Status getStatus(TaskInfo* task, TransferStatus& task_status);
 
@@ -113,4 +115,4 @@ class ProxyManager {
 }  // namespace tent
 }  // namespace mooncake
 
-#endif  // PROXY_MANAGER_H_
\ No newline at end of file
+#endif  // PROXY_MANAGER_H_
diff --git a/mooncake-transfer-engine/tent/include/tent/runtime/segment.h b/mooncake-transfer-engine/tent/include/tent/runtime/segment.h
index 21ee86e4..e913abe9 100644
--- a/mooncake-transfer-engine/tent/include/tent/runtime/segment.h
+++ b/mooncake-transfer-engine/tent/include/tent/runtime/segment.h
@@ -18,6 +18,7 @@
 #include <glog/logging.h>
 #include <netdb.h>
 
+#include <algorithm>
 #include <atomic>
 #include <cstdint>
 #include <functional>
@@ -25,7 +26,9 @@
 #include <string>
 #include <thread>
 #include <unordered_map>
+#include <utility>
 #include <variant>
+#include <vector>
 
 #include "tent/common/concurrent/rw_spinlock.h"
 #include "tent/common/concurrent/thread_local_storage.h"
@@ -169,12 +172,25 @@ struct SegmentDesc {
     std::string rpc_server_addr;
     std::variant<MemorySegmentDesc, FileSegmentDesc> detail;
 
+    // In dual-NIC setups (MC_RDMA_BIND_ADDRESS), the RDMA-reachable
+    // address may differ from the TCP-routable segment name.  When
+    // non-empty, NIC paths are constructed using this value instead
+    // of `name`.
+    std::string rdma_server_name;
+
    public:
     BufferDesc* findBuffer(uint64_t base, uint64_t length);
     DeviceDesc* findDevice(const std::string& name);
     const MemorySegmentDesc& getMemory() const {
         return std::get<MemorySegmentDesc>(detail);
     }
+
+    // Returns the server name to use for NIC path construction.
+    // Uses rdma_server_name when available, otherwise falls back
+    // to name.
+    const std::string& nicPathServerName() const {
+        return rdma_server_name.empty() ? name : rdma_server_name;
+    }
 };
 
 inline void to_json(json& j, const SegmentDesc& s) {
@@ -182,6 +198,9 @@ inline void to_json(json& j, const SegmentDesc& s) {
              {"type", s.type},
              {"machine_id", s.machine_id},
              {"rpc_server_addr", s.rpc_server_addr}};
+    if (!s.rdma_server_name.empty()) {
+        j["rdma_server_name"] = s.rdma_server_name;
+    }
     if (s.type == SegmentType::Memory) {
         j["detail"] = std::get<MemorySegmentDesc>(s.detail);
     } else {
@@ -194,6 +213,9 @@ inline void from_json(const json& j, SegmentDesc& s) {
     j.at("type").get_to(s.type);
     j.at("machine_id").get_to(s.machine_id);
     j.at("rpc_server_addr").get_to(s.rpc_server_addr);
+    if (j.contains("rdma_server_name")) {
+        j.at("rdma_server_name").get_to(s.rdma_server_name);
+    }
     if (s.type == SegmentType::Memory) {
         s.detail = j.at("detail").get<MemorySegmentDesc>();
     } else {
@@ -203,6 +225,73 @@ inline void from_json(const json& j, SegmentDesc& s) {
 
 using SegmentDescRef = std::shared_ptr<SegmentDesc>;
 
+// Coalesce a contiguous per-page NUMA probe into a small list of
+// routing-hint Regions. Each bucket emits one Region carrying the bucket's
+// majority NUMA label; receivers consume Region as a max-overlap NIC hint.
+//
+// Precondition: `entries` must tile the buffer as a gap-free, address-ordered
+// sequence (as produced by Platform::getLocation). Passing a non-contiguous or
+// unsorted vector still yields a monotonic byte offset and therefore mislabels
+// the address space; sort/validate at the call site if a producer cannot
+// guarantee this.
+//
+// Invariants: sum(out[i].size) == sum(entries[i].len), and ordering is
+// preserved. The per-bucket NUMA label is the byte-weighted majority, so
+// sub-bucket labels are deliberately lost -- this is a lossy routing hint, not
+// a faithful copy of the fine-grained probe.
+inline std::vector<Region> coalesceRegions(
+    const std::vector<RangeLocation>& entries, size_t max_buckets = 128,
+    size_t min_bucket_bytes = static_cast<size_t>(1) << 20) {
+    std::vector<Region> out;
+    if (entries.empty()) return out;
+
+    size_t total = 0;
+    for (auto& e : entries) total += e.len;
+    size_t computed = max_buckets > 0 ? total / max_buckets +
+                                            (total % max_buckets != 0 ? 1 : 0)
+                                      : 0;
+    // Defensive floor: bucket_bytes == 0 would loop forever / divide by zero.
+    // Hit only when caller passes both min_bucket_bytes=0 and max_buckets=0.
+    size_t bucket_bytes =
+        std::max<size_t>(1, std::max(min_bucket_bytes, computed));
+
+    std::vector<std::pair<std::string, size_t>> votes;
+    size_t bucket = 0;
+
+    auto add_vote = [&](const std::string& loc, size_t n) {
+        for (auto& v : votes) {
+            if (v.first == loc) {
+                v.second += n;
+                return;
+            }
+        }
+        votes.emplace_back(loc, n);
+    };
+    auto flush = [&]() {
+        if (bucket == 0) return;
+        const std::string* best = nullptr;
+        size_t best_n = 0;
+        for (auto& v : votes) {
+            // Strict > preserves "first-seen wins" on ties.
+            if (v.second > best_n) {
+                best_n = v.second;
+                best = &v.first;
+            }
+        }
+        out.push_back(Region{bucket, *best});
+        votes.clear();
+        bucket = 0;
+    };
+
+    for (auto& e : entries) {
+        add_vote(e.location, e.len);
+        bucket += e.len;
+        if (bucket >= bucket_bytes) flush();
+    }
+    flush();
+    return out;
+}
+
 }  // namespace tent
 }  // namespace mooncake
 
diff --git a/mooncake-transfer-engine/tent/include/tent/runtime/segment_manager.h b/mooncake-transfer-engine/tent/include/tent/runtime/segment_manager.h
index e01b2897..98be76bc 100644
--- a/mooncake-transfer-engine/tent/include/tent/runtime/segment_manager.h
+++ b/mooncake-transfer-engine/tent/include/tent/runtime/segment_manager.h
@@ -22,6 +22,7 @@
 #include <cstdint>
 #include <functional>
 #include <memory>
+#include <mutex>
 #include <string>
 #include <thread>
 #include <unordered_map>
@@ -30,6 +31,11 @@
 
 #include "tent/runtime/segment.h"
 
+// Fallback refresh interval for cached remote SegmentDesc. Invalidation
+// normally comes from the best-effort SegmentUpdate push; this TTL only bounds
+// staleness if a push is lost.
+#define TENT_SEGMENT_DESC_TTL_MS (60 * 60 * 1000)  // 1h
+
 namespace mooncake {
 namespace tent {
 class SegmentRegistry;
@@ -110,6 +116,14 @@ class SegmentManager {
    public:
     SegmentDescRef getLocal() { return local_desc_; }
 
+    // Returns a serialized JSON snapshot of local_desc_. The result is cached
+    // and shared across concurrent GetSegmentDesc RPC handlers; the cache is
+    // invalidated by synchronizeLocal() (which is called on every register /
+    // unregisterLocalMemory). This avoids re-dumping the full segment desc on
+    // every peer fetch — the previous behavior multiplied dump cost by the
+    // number of concurrent peer RPCs and dominated remote getRemote latency.
+    std::shared_ptr<const std::string> getLocalDumpedJson();
+
     Status synchronizeLocal();
 
     Status deleteLocal();
@@ -145,11 +159,14 @@ class SegmentManager {
     std::unique_ptr<SegmentRegistry> registry_;
 
     std::string file_desc_basepath_;
-    uint64_t ttl_ms_ = 10 * 1000;  // N.B. Frequent TTL harms p999
 
     // shared_ptr to prevent UAF in async callbacks
     std::shared_ptr<RWSpinlock> subscribers_lock_;
     std::shared_ptr<std::unordered_set<std::string>> subscribers_;
+
+    // Cache for the serialized JSON of local_desc_. Reset by synchronizeLocal.
+    std::mutex local_json_cache_mu_;
+    std::shared_ptr<const std::string> local_json_cache_;
 };
 }  // namespace tent
 }  // namespace mooncake
diff --git a/mooncake-transfer-engine/tent/include/tent/runtime/segment_registry.h b/mooncake-transfer-engine/tent/include/tent/runtime/segment_registry.h
index 97c586d0..46d5dc56 100644
--- a/mooncake-transfer-engine/tent/include/tent/runtime/segment_registry.h
+++ b/mooncake-transfer-engine/tent/include/tent/runtime/segment_registry.h
@@ -53,13 +53,6 @@ class CentralSegmentRegistry : public SegmentRegistry {
    public:
     CentralSegmentRegistry(const std::string &type, const std::string &servers);
 
-    CentralSegmentRegistry(const std::string &type, const std::string &servers,
-                           const std::string &password, uint8_t db_index);
-
-    CentralSegmentRegistry(const std::string &type, const std::string &servers,
-                           const std::string &username,
-                           const std::string &password, uint8_t db_index);
-
     virtual ~CentralSegmentRegistry() {}
 
    public:
diff --git a/mooncake-transfer-engine/tent/include/tent/runtime/topology.h b/mooncake-transfer-engine/tent/include/tent/runtime/topology.h
index d87a1dd5..9360b5e2 100644
--- a/mooncake-transfer-engine/tent/include/tent/runtime/topology.h
+++ b/mooncake-transfer-engine/tent/include/tent/runtime/topology.h
@@ -12,8 +12,11 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-#ifndef TOPOLOGY_H
-#define TOPOLOGY_H
+// NOTE: Guard renamed from TOPOLOGY_H to avoid collision with the classic
+// transfer-engine's mooncake-transfer-engine/include/topology.h, which also
+// uses TOPOLOGY_H.
+#ifndef TENT_TOPOLOGY_H
+#define TENT_TOPOLOGY_H
 
 #include <glog/logging.h>
 #include <netdb.h>
@@ -137,4 +140,4 @@ struct RangeLocation {
 }  // namespace tent
 }  // namespace mooncake
 
-#endif  // TOPOLOGY_H
\ No newline at end of file
+#endif  // TENT_TOPOLOGY_H
\ No newline at end of file
diff --git a/mooncake-transfer-engine/tent/include/tent/runtime/transfer_engine_impl.h b/mooncake-transfer-engine/tent/include/tent/runtime/transfer_engine_impl.h
index 71874af2..41aca93a 100644
--- a/mooncake-transfer-engine/tent/include/tent/runtime/transfer_engine_impl.h
+++ b/mooncake-transfer-engine/tent/include/tent/runtime/transfer_engine_impl.h
@@ -30,13 +30,14 @@
 #include "tent/common/config.h"
 #include "tent/common/status.h"
 #include "tent/common/types.h"
-#include "tent/common/concurrent/thread_local_storage.h"
+#include "tent/runtime/admission_queue.h"
+#include "tent/runtime/transport.h"
+#include "tent/runtime/transport_selector.h"
 
 namespace mooncake {
 namespace tent {
 
 class Batch;
-class BatchSet;
 class Topology;
 class Transport;
 class SegmentDesc;
@@ -45,18 +46,20 @@ class ControlService;
 class SegmentTracker;
 class Platform;
 class ProxyManager;
+class ProgressWorker;
 
 struct TaskInfo {
     TransportType type{UNSPEC};
     int sub_task_id{-1};
-    bool derived{false};  // merged by other tasks
-    int xport_priority{0};
+    bool derived{false};          // merged by other tasks
+    int xport_priority{0};        // transport priority (for fallback)
+    int failover_count{0};        // number of failover attempts
+    uint64_t device_mask{~0ULL};  // Device mask for quota allocation
     Request request;
     bool staging{false};
     TransferStatusEnum status{TransferStatusEnum::PENDING};
     volatile TransferStatusEnum staging_status{TransferStatusEnum::PENDING};
     std::chrono::steady_clock::time_point start_time{};  // For latency tracking
-    int failover_count{0};  // Number of cross-transport failover attempts
 };
 
 class TransferEngineImpl {
@@ -150,6 +153,8 @@ class TransferEngineImpl {
 
     Status getTransferStatus(BatchID batch_id, TransferStatus& overall_status);
 
+    Status progressBatch(BatchID batch_id, TransferStatus& overall_status);
+
     Status waitTransferCompletion(BatchID batch_id);
 
     Status transferSync(const std::vector<Request>& request_list);
@@ -158,7 +163,26 @@ class TransferEngineImpl {
 
     Status unlockStageBuffer(uint64_t addr);
 
+    // Test-only hook: replace the transport in a given slot after construct().
+    // Production code never calls this. Used by failover integration tests to
+    // inject a FaultProxyTransport without bypassing resubmitTransferTask,
+    // resolveTransport, or any other engine state. Not thread-safe with any
+    // in-flight transfer on that slot.
+    void swapTransportForTest(TransportType type,
+                              std::shared_ptr<Transport> xport) {
+        if (type >= 0 && type < (TransportType)kSupportedTransportTypes) {
+            transport_list_[type] = std::move(xport);
+        }
+    }
+
+    // Wake the optional event-driven progress worker for `batch_id`. No-op if
+    // enable_progress_worker is false. Transport completion paths use this as
+    // an idempotent "maybe ready" signal.
+    void notifyBatchMaybeReady(BatchID batch_id);
+
    private:
+    friend class ProgressWorker;
+
     Status construct();
 
     Status deconstruct();
@@ -167,15 +191,77 @@ class TransferEngineImpl {
 
     Status lazyFreeBatch();
 
-    TransportType getTransportType(const Request& request, int priority = 0);
+    SelectionResult getTransportType(const Request& request,
+                                     int transport_index = 0);
 
     std::vector<TransportType> getSupportedTransports(
         TransportType request_type);
 
     Status resubmitTransferTask(Batch* batch, size_t task_id);
 
-    TransportType resolveTransport(const Request& req, int priority,
-                                   bool invalidate_on_fail = true);
+    Status retainBatch(BatchID batch_id, Batch*& batch);
+
+    Status releaseBatch(Batch* batch);
+
+    class BatchRef;
+
+    struct PreparedSubmit;
+
+    Status submitTransfer(BatchID batch_id,
+                          const std::vector<Request>& request_list,
+                          const Notification* notifi,
+                          QueueOwnerKind owner_kind);
+
+    Status submitStagingTransfer(BatchID batch_id,
+                                 const std::vector<Request>& request_list);
+
+    Status enqueuePreparedSubmit(Batch* batch, const PreparedSubmit& prepared,
+                                 QueueOwnerKind owner_kind);
+
+    bool shouldQueueSubmit(const PreparedSubmit& prepared,
+                           QueueOwnerKind owner_kind) const;
+
+    Status prepareSubmit(Batch* batch, const std::vector<Request>& request_list,
+                         PreparedSubmit& prepared);
+
+    Status commitPreparedSubmit(Batch* batch, const PreparedSubmit& prepared);
+
+    void attachProgressNotifier(Batch* batch, Transport::SubBatchRef sub_batch);
+
+    uint64_t nextBatchToken();
+
+    Status refillDispatchWindow();
+
+    Status progressRuntimeQueue();
+
+    bool hasActiveRuntimeQueue();
+
+    void notifyRuntimeQueueReady();
+
+    Status dispatchQueuedOwner(QueueOwnerId owner_id);
+
+    Status markQueuedOwnerSubmitted(QueueOwnerId owner_id);
+
+    Status finishQueuedOwner(QueueOwnerId owner_id,
+                             TransferStatusEnum terminal_status);
+
+    Status retireQueueForBatch(Batch* batch);
+
+    Status pollTaskStatus(Batch* batch, size_t task_id,
+                          TransferStatus& task_status);
+
+    void updateTaskStatusAfterPoll(Batch* batch, size_t task_id,
+                                   TransferStatus& task_status,
+                                   bool allow_failover);
+
+    Status getBatchStatus(BatchID batch_id, TransferStatus& overall_status,
+                          bool allow_failover);
+
+    SelectionResult resolveTransport(const Request& req, int transport_index,
+                                     bool invalidate_on_fail = true);
+
+    // Verify that req.transport_hint is usable for this request
+    Status validateTransportHint(const Request& req, size_t request_index);
 
     Status loadTransports();
 
@@ -184,6 +270,10 @@ class TransferEngineImpl {
 
     Status maybeFireSubmitHooks(Batch* batch, bool check = true);
 
+    void addSubmitHook(Batch* batch, size_t start_task_id,
+                       const std::vector<Request>& request_list,
+                       const Notification& notifi);
+
     void recordTaskCompletionMetrics(TaskInfo& task,
                                      TransferStatusEnum prev_status,
                                      TransferStatusEnum new_status);
@@ -201,17 +291,34 @@ class TransferEngineImpl {
         std::vector<Batch*> freelist;
     };
 
+    struct RuntimeQueueConfig {
+        bool enabled{false};
+        QueueLimits limits{};
+        size_t max_dispatch_owners{0};
+        size_t max_dispatch_bytes{0};
+        std::chrono::microseconds progress_fallback_interval{50000};
+    };
+
+    struct QueuedOwnerState {
+        Batch* batch{nullptr};
+        size_t owner_task_id{0};
+        std::vector<size_t> public_task_ids;
+        size_t byte_charge{0};
+        bool in_dispatch_window{false};
+    };
+
    private:
     std::shared_ptr<Config> conf_;
     std::shared_ptr<ControlService> metadata_;
     std::shared_ptr<Topology> topology_;
+    std::unique_ptr<TransportSelector> transport_selector_;
     bool available_;
 
     std::array<std::shared_ptr<Transport>, kSupportedTransportTypes>
         transport_list_;
     std::unique_ptr<SegmentTracker> local_segment_tracker_;
 
-    ThreadLocalStorage<BatchSet> batch_set_;
+    BatchSet batch_set_;
 
     std::vector<AllocatedMemory> allocated_memory_;
     std::mutex mutex_;
@@ -224,8 +331,24 @@ class TransferEngineImpl {
     std::unique_ptr<ProxyManager> staging_proxy_;
     bool merge_requests_;
     int max_failover_attempts_{3};
+    bool enable_auto_failover_on_poll_{true};
+    bool enable_progress_worker_{false};
+    RuntimeQueueConfig runtime_queue_config_;
+    std::unique_ptr<LocalTransferAdmissionQueue> runtime_queue_;
+    std::unordered_map<QueueOwnerId, QueuedOwnerState> queued_owners_;
+    size_t dispatch_inflight_owners_{0};
+    size_t dispatch_inflight_bytes_{0};
+    uint64_t next_batch_token_{1};
+
+    // Guards alive_batches_ and serializes pollTaskStatus /
+    // updateTaskStatusAfterPoll / lazyFreeBatch against the optional
+    // ProgressWorker thread. Recursive because freeBatch -> lazyFreeBatch ->
+    // getTransferStatus can re-enter on the same thread. See issue #2116.
+    std::recursive_mutex progress_mutex_;
+    std::unordered_set<BatchID> alive_batches_;
+    std::unique_ptr<ProgressWorker> progress_worker_;
 };
 }  // namespace tent
 }  // namespace mooncake
 
-#endif
\ No newline at end of file
+#endif
diff --git a/mooncake-transfer-engine/tent/include/tent/runtime/transport.h b/mooncake-transfer-engine/tent/include/tent/runtime/transport.h
index 36e62a09..b28cf7cf 100644
--- a/mooncake-transfer-engine/tent/include/tent/runtime/transport.h
+++ b/mooncake-transfer-engine/tent/include/tent/runtime/transport.h
@@ -22,12 +22,14 @@
 
 #include <cstddef>
 #include <cstdint>
+#include <functional>
 #include <iostream>
 #include <memory>
 #include <queue>
 #include <string>
 
 #include "tent/common/status.h"
+#include "tent/common/types.h"
 #include "tent/runtime/platform.h"
 #include "tent/runtime/control_plane.h"
 
@@ -45,9 +47,16 @@ struct Capabilities {
 class Transport {
    public:
     struct SubBatch {
-        SubBatch() {}
+        SubBatch() : device_mask(~0ULL) {}
         virtual ~SubBatch() {}
         virtual size_t size() const = 0;
+        void notifyProgress() {
+            if (notify_progress) notify_progress(progress_batch_id);
+        }
+
+        uint64_t device_mask;  // Device mask for transport selection
+        BatchID progress_batch_id{0};
+        std::function<void(BatchID)> notify_progress;
     };
 
     using SubBatchRef = SubBatch *;
diff --git a/mooncake-transfer-engine/tent/include/tent/runtime/transport_selector.h b/mooncake-transfer-engine/tent/include/tent/runtime/transport_selector.h
new file mode 100644
index 00000000..eaa2341f
--- /dev/null
+++ b/mooncake-transfer-engine/tent/include/tent/runtime/transport_selector.h
@@ -0,0 +1,196 @@
+// Copyright 2025 KVCache.AI
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+/**
+ * @file transport_selector.h
+ * @brief Configuration-driven transport selection policy
+ *
+ * Transport selection is driven by configuration with pattern-based rules.
+ *
+ * Example configuration:
+ * {
+ *   "policy": [
+ *     {
+ *       "name": "high_prio_fast",
+ *       "segment_type": "memory",
+ *       "priority": 0,
+ *       "devices": ["mlx5_0", "mlx5_1", "mlx5_2"],
+ *       "transports": ["nvlink", "rdma", "shm"]
+ *     },
+ *     {
+ *       "name": "low_prio_slow",
+ *       "segment_type": "memory",
+ *       "priority": 2,
+ *       "devices": ["mlx5_0"],
+ *       "transports": ["rdma", "tcp"]
+ *     },
+ *     {
+ *       "name": "file_storage",
+ *       "segment_type": "file",
+ *       "transports": ["gds", "io_uring", "rdma"]
+ *     }
+ *   ]
+ * }
+ *
+ * If no "policy" is configured, defaults to original behavior:
+ * - File: GDS -> IOURING -> RDMA
+ * - Memory: uses buffer_transports order
+ */
+
+#ifndef TENT_TRANSPORT_SELECTOR_H
+#define TENT_TRANSPORT_SELECTOR_H
+
+#include "tent/common/config.h"
+#include "tent/common/types.h"
+#include "tent/runtime/segment.h"
+#include "tent/runtime/platform.h"
+
+#include <array>
+#include <cstdint>
+#include <memory>
+#include <optional>
+#include <string>
+#include <vector>
+
+namespace mooncake {
+namespace tent {
+
+class Transport;
+
+/**
+ * @brief Selection context for a single request
+ */
+struct SelectionContext {
+    SegmentType segment_type;       // File or Memory
+    bool same_machine;              // Local or remote
+    MemoryType local_memory_type;   // CPU, CUDA, etc.
+    MemoryType remote_memory_type;  // CPU, CUDA, etc. (for remote)
+    const std::vector<TransportType>*
+        buffer_transports;  // Pointer to transports in buffer
+    size_t transfer_size;   // Transfer size in bytes
+    int priority_level;     // Request priority level (lower = more urgent)
+    std::optional<std::string>
+        policy_name;  // Optional: bind to specific policy by name
+};
+
+/**
+ * @brief Transport selection policy rule
+ */
+struct SelectionPolicy {
+    // Basic identification
+    std::string name;
+
+    // Segment type filter
+    SegmentType segment_type;
+
+    // Location filter
+    std::optional<bool> same_machine;  // nullopt = don't care
+
+    // Memory type filters (supports patterns: "cuda", "cpu", "npu", "*" for
+    // any)
+    std::optional<std::string> local_memory_pattern;
+    std::optional<std::string> remote_memory_pattern;
+
+    // Size filter (nullopt = no limit)
+    std::optional<uint64_t> min_size;  // Minimum transfer size
+    std::optional<uint64_t> max_size;  // Maximum transfer size
+
+    // Priority filter: exact match required (request.priority == priority)
+    // nullopt = match any priority level
+    std::optional<int> priority;
+
+    // Device allocation: list of device names this policy can use
+    // e.g., ["mlx5_0", "mlx5_1", "rocep5s0f0"]
+    // Empty = use all available devices
+    std::vector<std::string> devices;
+
+    // Transport preference list (evaluated in order)
+    std::vector<TransportType> transports;
+};
+
+/**
+ * @brief Result of transport selection
+ */
+struct SelectionResult {
+    TransportType transport = UNSPEC;
+    uint64_t device_mask = ~0ULL;  // Bitmask of allowed devices (~0 = all)
+};
+
+/**
+ * @brief Configuration-driven transport selector
+ */
+class TransportSelector {
+   public:
+    TransportSelector(std::shared_ptr<Config> config);
+
+    /**
+     * @brief Set topology for device name to ID conversion
+     */
+    void setTopology(std::shared_ptr<Topology> topology) {
+        topology_ = topology;
+    }
+
+    /**
+     * @brief Select the best transport for a given context
+     * @param context Selection context
+     * @param available_transports Array of available transports
+     * @param transport_index Transport selection index (0 = first choice, 1 =
+     * second, ...)
+     * @param hint The caller's preferred transport for this request;
+     * UNSPEC means no hint.
+     * @return Selection result with transport type and device mask
+     */
+    SelectionResult select(
+        const SelectionContext& context,
+        const std::array<std::shared_ptr<Transport>, kSupportedTransportTypes>&
+            available_transports,
+        int transport_index = 0, TransportType hint = UNSPEC);
+
+    /**
+     * @brief Enable legacy mode (skip TransportSelector, use original logic)
+     */
+    void setLegacyMode(bool enabled) { legacy_mode_ = enabled; }
+
+    /**
+     * @brief Check if legacy mode is enabled
+     */
+    bool isLegacyMode() const { return legacy_mode_; }
+
+    static std::string transportTypeName(TransportType type);
+    static TransportType parseTransportType(const std::string& str);
+    static std::optional<std::vector<TransportType>> reorderWithHint(
+        const std::vector<TransportType>& raw, TransportType hint);
+
+   private:
+    std::vector<SelectionPolicy> getDefaultPolicies();
+    void loadPolicies();
+    bool matchesMemoryPattern(const std::string& pattern,
+                              MemoryType type) const;
+    bool matchesPolicy(const SelectionPolicy& policy,
+                       const SelectionContext& context) const;
+    bool isTransportAvailable(
+        TransportType type, const SelectionContext& context,
+        const std::array<std::shared_ptr<Transport>, kSupportedTransportTypes>&
+            available_transports) const;
+
+    std::shared_ptr<Config> config_;
+    std::shared_ptr<Topology> topology_;
+    std::vector<SelectionPolicy> policies_;
+    bool legacy_mode_{false};  // If true, skip selector and return empty result
+};
+
+}  // namespace tent
+}  // namespace mooncake
+
+#endif  // TENT_TRANSPORT_SELECTOR_H
diff --git a/mooncake-transfer-engine/tent/include/tent/transfer_engine.h b/mooncake-transfer-engine/tent/include/tent/transfer_engine.h
index c8cd7c5b..ff93e0ed 100644
--- a/mooncake-transfer-engine/tent/include/tent/transfer_engine.h
+++ b/mooncake-transfer-engine/tent/include/tent/transfer_engine.h
@@ -32,12 +32,17 @@ extern "C" {
 #define OPCODE_READ (0)
 #define OPCODE_WRITE (1)
 
+/* IMPORTANT: callers MUST zero-initialize this struct
+ * (e.g. `tent_request_t r = {0};`). transport_hint relies on
+ * UNSPEC == 0 for the no-hint default. */
 struct tent_request {
     int opcode;
     void* source;
     tent_segment_id_t target_id;
     uint64_t target_offset;
     uint64_t length;
+    int priority;       /* Request priority (0=HIGH, 1=MEDIUM, 2=LOW) */
+    int transport_hint; /* TRANSPORT_UNSPEC=follow policy, else pin transport */
 };
 
 typedef struct tent_request tent_request_t;
@@ -94,15 +99,16 @@ typedef struct tent_notifi_info tent_notifi_info;
 #define PERM_GLOBAL_READ_ONLY (1)
 #define PERM_GLOBAL_READ_WRITE (2)
 
-#define TRANSPORT_RDMA (0)
-#define TRANSPORT_MNNVL (1)
-#define TRANSPORT_SHM (2)
-#define TRANSPORT_NVLINK (3)
-#define TRANSPORT_GDS (4)
-#define TRANSPORT_IOURING (5)
-#define TRANSPORT_TCP (6)
-#define TRANSPORT_ASCEND_DIRECT (7)
-#define TRANSPORT_UNSPEC (8)
+#define TRANSPORT_UNSPEC (0)
+#define TRANSPORT_RDMA (1)
+#define TRANSPORT_MNNVL (2)
+#define TRANSPORT_SHM (3)
+#define TRANSPORT_NVLINK (4)
+#define TRANSPORT_GDS (5)
+#define TRANSPORT_IOURING (6)
+#define TRANSPORT_TCP (7)
+#define TRANSPORT_ASCEND_DIRECT (8)
+#define TRANSPORT_SUNRISE_LINK (9)
 
 struct tent_memory_options {
     char location[64];
@@ -307,6 +313,14 @@ class TransferEngine {
 
     Status getTransferStatus(BatchID batch_id, TransferStatus& overall_status);
 
+    // Drive one progress step on a batch and return its aggregated status.
+    // Unlike getTransferStatus, this always allows internal failover/resubmit
+    // regardless of enable_auto_failover_on_poll. The call is non-blocking and
+    // performs at most one state-machine step per task; callers that want to
+    // wait for completion must invoke it in a loop. PENDING means "make
+    // progress later"; terminal states (COMPLETED/FAILED) will not be revived.
+    Status progressBatch(BatchID batch_id, TransferStatus& overall_status);
+
    private:
     std::unique_ptr<TransferEngineImpl> impl_;
 };
diff --git a/mooncake-transfer-engine/tent/include/tent/transport/mnnvl/mnnvl_transport.h b/mooncake-transfer-engine/tent/include/tent/transport/mnnvl/mnnvl_transport.h
index bac1b88f..8986038b 100644
--- a/mooncake-transfer-engine/tent/include/tent/transport/mnnvl/mnnvl_transport.h
+++ b/mooncake-transfer-engine/tent/include/tent/transport/mnnvl/mnnvl_transport.h
@@ -23,6 +23,7 @@
 #include <queue>
 #include <string>
 #include <unordered_set>
+#include <vector>
 
 #include "tent/runtime/control_plane.h"
 #include "tent/runtime/transport.h"
@@ -45,6 +46,13 @@ struct MnnvlSubBatch : public Transport::SubBatch {
     size_t max_size;
     CUDAStreamHandle sync_stream;
     CUDAStreamHandle async_stream;
+    // Completion events created in startTransfer (one per submit). Destroyed by
+    // the destructor (RAII); Slab<T>::deallocate() invokes ~MnnvlSubBatch()
+    // before reusing the storage, so this runs on every free.
+    std::vector<cudaEvent_t> completion_events;
+    ~MnnvlSubBatch() {
+        for (auto event : completion_events) cudaEventDestroy(event);
+    }
     virtual size_t size() const { return task_list.size(); }
 };
 
diff --git a/mooncake-transfer-engine/tent/include/tent/transport/nvlink/nvlink_transport.h b/mooncake-transfer-engine/tent/include/tent/transport/nvlink/nvlink_transport.h
index 9c38162d..38446c64 100644
--- a/mooncake-transfer-engine/tent/include/tent/transport/nvlink/nvlink_transport.h
+++ b/mooncake-transfer-engine/tent/include/tent/transport/nvlink/nvlink_transport.h
@@ -21,6 +21,7 @@
 #include <queue>
 #include <string>
 #include <unordered_set>
+#include <vector>
 
 #include <cuda.h>
 #include <cuda_runtime.h>
@@ -48,6 +49,13 @@ struct NVLinkSubBatch : public Transport::SubBatch {
     size_t max_size;
     CUDAStreamHandle sync_stream;
     CUDAStreamHandle async_stream;
+    // Completion events created in startTransfer (one per submit). Destroyed by
+    // the destructor (RAII); Slab<T>::deallocate() invokes ~NVLinkSubBatch()
+    // before reusing the storage, so this runs on every free.
+    std::vector<cudaEvent_t> completion_events;
+    ~NVLinkSubBatch() {
+        for (auto event : completion_events) cudaEventDestroy(event);
+    }
     virtual size_t size() const { return task_list.size(); }
 };
 
diff --git a/mooncake-transfer-engine/tent/include/tent/transport/rdma/context.h b/mooncake-transfer-engine/tent/include/tent/transport/rdma/context.h
index 141fd9d5..2c010872 100644
--- a/mooncake-transfer-engine/tent/include/tent/transport/rdma/context.h
+++ b/mooncake-transfer-engine/tent/include/tent/transport/rdma/context.h
@@ -114,6 +114,9 @@ class RdmaContext {
 
     RdmaParams &params() const { return *params_.get(); }
 
+    // PCIe Relaxed Ordering support
+    bool isRelaxedOrderingEnabled() const { return relaxed_ordering_enabled_; }
+
     // Notification CQ (dedicated for notification QPs)
     RdmaCQ *notifyCq() { return notify_cq_; }
 
@@ -149,6 +152,9 @@ class RdmaContext {
     // Dedicated CQ for notification QPs (one per device)
     RdmaCQ *notify_cq_ = nullptr;
 
+    // PCIe Relaxed Ordering support
+    bool relaxed_ordering_enabled_ = false;
+
     const IbvSymbols &verbs_;
 };
 
diff --git a/mooncake-transfer-engine/tent/include/tent/transport/rdma/endpoint.h b/mooncake-transfer-engine/tent/include/tent/transport/rdma/endpoint.h
index 169d4744..7ef3588b 100644
--- a/mooncake-transfer-engine/tent/include/tent/transport/rdma/endpoint.h
+++ b/mooncake-transfer-engine/tent/include/tent/transport/rdma/endpoint.h
@@ -81,9 +81,16 @@ class RdmaEndPoint : public std::enable_shared_from_this<RdmaEndPoint> {
     ~RdmaEndPoint();
 
    public:
-    enum EndPointStatus { EP_UNINIT, EP_HANDSHAKING, EP_READY, EP_RESET };
-
-    int reset();
+    // Endpoint lifecycle is unidirectional: once created, endpoints move
+    // forward through states and never return. Failed or discarded endpoints
+    // enter EP_DESTROYING and are reclaimed - they are never reset or reused.
+    enum EndPointStatus {
+        EP_UNINIT,       // Initial state, not yet constructed
+        EP_HANDSHAKING,  // Connection in progress
+        EP_READY,        // Connected and operational
+        EP_DESTROYING,   // Being destroyed (failed or discarded)
+        EP_DESTROYED,    // Fully destroyed
+    };
 
     int construct(RdmaContext* context, EndPointParams* params,
                   const std::string& endpoint_name,
@@ -91,13 +98,27 @@ class RdmaEndPoint : public std::enable_shared_from_this<RdmaEndPoint> {
 
     int deconstruct();
 
+    int resetConnection(const std::string& reason);
+
+    // Two-phase QP destruction to avoid use-after-free in concurrent
+    // submitPostSend. Phase 1 (beginDestroy): sets status_=EP_DESTROYING,
+    // transitions QPs to ERR state so hardware flushes inflight WRs to CQ.
+    // Does not block. Phase 2 (finishDestroy): called after all outstanding
+    // WRs have been drained, actually destroys QPs and frees resources.
+    // Returns true if destruction is complete, false if outstanding WRs remain.
+    void beginDestroy();
+    void beginDestroyNoLock();  // Internal version without locking
+    bool finishDestroy();
+
     Status connect(const std::string& peer_server_name,
                    const std::string& peer_nic_name,
                    const std::string& peer_rpc_server_addr = "");
 
     Status accept(const BootstrapDesc& peer_desc, BootstrapDesc& local_desc);
 
-    EndPointStatus status() const { return status_; }
+    EndPointStatus status() const {
+        return status_.load(std::memory_order_relaxed);
+    }
 
     std::vector<uint32_t> qpNum();
 
@@ -135,8 +156,6 @@ class RdmaEndPoint : public std::enable_shared_from_this<RdmaEndPoint> {
         bool failed;
     };
 
-    int resetUnlocked();
-
     int submitSlices(std::vector<RdmaSlice*>& slice_list, int qp_index);
 
     int submitRecvImmDataRequest(int qp_index, uint64_t id);
@@ -183,6 +202,9 @@ class RdmaEndPoint : public std::enable_shared_from_this<RdmaEndPoint> {
     uint32_t padding_[7];
     RWSpinlock lock_;
 
+    // Two-phase destruction: timestamp when EP entered EP_DESTROYING
+    std::atomic<uint64_t> destroy_start_time_;
+
     std::string peer_server_name_;
     std::string peer_nic_name_;
     std::vector<uint32_t> peer_qp_num_list_;
@@ -203,6 +225,11 @@ class RdmaEndPoint : public std::enable_shared_from_this<RdmaEndPoint> {
     int notify_pending_count_ = 0;    // Number of pending sends
     uint64_t notify_send_wr_id_ = 0;  // Circular counter for wr_id
     bool notify_connected_ = false;
+
+    // Two-phase destruction constants (matching TE)
+    static constexpr double kFinishDestroyTimeoutSec = 30.0;
+    static constexpr int kFinishDestroyMaxRetries = 3;
+    int finish_destroy_retries_ = 0;  // Retry counter for finishDestroy
 };
 }  // namespace tent
 }  // namespace mooncake
diff --git a/mooncake-transfer-engine/tent/include/tent/transport/rdma/quota.h b/mooncake-transfer-engine/tent/include/tent/transport/rdma/quota.h
index 57c70afa..b3a857ec 100644
--- a/mooncake-transfer-engine/tent/include/tent/transport/rdma/quota.h
+++ b/mooncake-transfer-engine/tent/include/tent/transport/rdma/quota.h
@@ -12,8 +12,8 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-#ifndef TENT_QUOTA_H
-#define TENT_QUOTA_H
+#ifndef TENT_SELECTOR_H
+#define TENT_SELECTOR_H
 
 #include <atomic>
 #include <vector>
@@ -30,47 +30,83 @@
 namespace mooncake {
 namespace tent {
 
-class SharedQuotaManager;
+// Bandwidth constants (Gbps)
+static constexpr double kDefaultBwGbps = 400.0;
+static constexpr double kMinBwGbps = 10.0;
+static constexpr double kMaxBwGbps = 800.0;
+
+class SharedSlotManager;
 
 /**
- * @brief DeviceQuota implements NIC selection based on adaptive feedback.
+ * @brief DeviceSelector implements NIC selection with two modes:
  *
- * Each NIC maintains a smoothed estimate of its average service time,
- * updated after each request completes. The allocator predicts the total
- * completion time of each NIC as:
+ * 1. Baseline mode (smart_selection_enabled=false): Simple round-robin
+ *    - Deterministic, no load tracking
+ *    - All devices used equally
  *
- *     predicted_time = (active_bytes / bandwidth) + avg_service_time
+ * 2. Smart mode (smart_selection_enabled=true): EWMA-based selection
+ *    - Tracks global inflight bytes per device
+ *    - Learns effective bandwidth via EWMA
+ *    - Selects device with minimal predicted completion time
+ *    - Supports multi-path for large requests
  *
- * and selects the NIC with the smallest predicted_time.
+ * Selection formula:
+ *     predicted_time = (inflight + slice_bytes) / ewma_bandwidth
  *
- * The estimator is updated using exponential smoothing:
- *
- *     avg_service_time <- (1 - alpha) * avg_service_time + alpha *
- * observed_time
+ * EWMA update:
+ *     ewma_bandwidth <- alpha * ewma_bandwidth + (1 - alpha) *
+ * observed_bandwidth
  */
-class DeviceQuota {
+class DeviceSelector {
    public:
+    // Candidate device for allocation
+    struct Candidate {
+        int dev_id;
+        double score;
+        bool is_cross_numa;
+    };
+
     struct DeviceInfo {
         int dev_id;
         double bw_gbps;
         int numa_id;
         uint64_t padding0[5];
-        std::atomic<uint64_t> active_bytes{0};
+        std::atomic<uint64_t> inflight_bytes{0};
         uint64_t padding1[7];
-        std::atomic<uint64_t> diffusion_active_bytes{0};
+        std::atomic<double> ewma_bandwidth_bps{50e9};
         uint64_t padding2[7];
-        std::atomic<double> beta0{0.0};  // Fixed latency (PCIe, setup)
-        uint64_t padding3[7];
-        std::atomic<double> beta1{1.0};  // Effective bandwidth correction
-        uint64_t padding4[7];
+        std::atomic<uint64_t> total_bytes{0};
+        uint64_t padding3[5];
+
+        uint64_t getInflightBytes() const {
+            return inflight_bytes.load(std::memory_order_relaxed);
+        }
+
+        void addInflight(uint64_t bytes) {
+            inflight_bytes.fetch_add(bytes, std::memory_order_relaxed);
+        }
+
+        void releaseInflight(uint64_t bytes) {
+            inflight_bytes.fetch_sub(bytes, std::memory_order_relaxed);
+        }
+
+        double getEwmaBandwidth() const {
+            return ewma_bandwidth_bps.load(std::memory_order_relaxed);
+        }
+
+        double getTheoreticalBandwidth() const {
+            if (bw_gbps >= kMinBwGbps && bw_gbps <= kMaxBwGbps)
+                return bw_gbps * 1e9 / 8.0;
+            return kDefaultBwGbps * 1e9 / 8.0;
+        }
     };
 
    public:
-    DeviceQuota() = default;
-    ~DeviceQuota() = default;
+    DeviceSelector() = default;
+    ~DeviceSelector() = default;
 
-    DeviceQuota(const DeviceQuota &) = delete;
-    DeviceQuota &operator=(const DeviceQuota &) = delete;
+    DeviceSelector(const DeviceSelector &) = delete;
+    DeviceSelector &operator=(const DeviceSelector &) = delete;
 
     Status loadTopology(std::shared_ptr<Topology> &local_topology);
 
@@ -78,46 +114,111 @@ class DeviceQuota {
 
     Status enableSharedQuota(const std::string &shm_name);
 
+    std::shared_ptr<SharedSlotManager> getSharedSlotManager() const {
+        return slot_manager_;
+    }
+
+    // Allocate devices for a request (new API)
+    // slice_bytes: pre-calculated slice size from rdma_transport to ensure
+    // consistency
+    Status allocate(uint64_t total_length, uint32_t num_slices,
+                    uint64_t slice_bytes, const std::string &location,
+                    std::vector<int> &slice_dev_ids, int priority = PRIO_HIGH,
+                    uint64_t device_mask = ~0ULL);
+
     Status allocate(uint64_t length, const std::string &location,
                     int &chosen_dev_id);
 
     Status release(int dev_id, uint64_t length, double latency);
 
-    void setDiffusionActiveBytes(int dev_id, uint64_t value) {
-        devices_[dev_id].diffusion_active_bytes.store(
-            value, std::memory_order_relaxed);
+    void updateTrafficStats(int dev_id, uint64_t length) {
+        auto it = devices_.find(dev_id);
+        if (it != devices_.end()) {
+            it->second.total_bytes.fetch_add(length, std::memory_order_relaxed);
+        }
     }
 
-    uint64_t getActiveBytes(int dev_id) {
-        return devices_[dev_id].active_bytes.load(std::memory_order_relaxed);
+    void setSmartSelection(bool enable) { smart_selection_enabled_ = enable; }
+    bool getSmartSelection() const { return smart_selection_enabled_; }
+
+    void setLearningRate(double alpha) {
+        sched_params_.bandwidth_learning_rate = std::clamp(alpha, 0.0, 1.0);
     }
 
-    void setLearningRate(double alpha) { alpha_ = std::clamp(alpha, 0.0, 1.0); }
+    int getDeviceRank(const std::string &location, int dev_id) const;
 
-    void setLocalWeight(double local_weight) {
-        local_weight_ = std::clamp(local_weight, 0.0, 1.0);
+    void printTrafficStats();
+
+    void fillDevicePriorities();
+    int getDevicePriority(int dev_id) const;
+
+    struct SchedulingParams {
+        // NUMA tier penalties (rank 0 = local, should be smallest)
+        double numa_tier_weights[Topology::DevicePriorityRanks] = {1.0, 5.0,
+                                                                   10.0};
+
+        // EWMA bandwidth learning rate (0.0 = full adaptation, 1.0 = no
+        // learning)
+        double bandwidth_learning_rate = 0.01;
+
+        // Enable priority-based filtering
+        bool enable_priority_filtering = true;
+
+        // Local device priority rotation interval (microseconds)
+        uint64_t local_rotation_interval_us = 200;
+
+        // Score random jitter range (to avoid deterministic selection)
+        double score_jitter_range = 1e-9;
+
+        // Epsilon for division by zero protection
+        double score_epsilon = 1e-12;
+
+        // EWMA bandwidth bounds (multiplier of theoretical bandwidth)
+        double ewma_min_multiplier = 0.1;   // 10% of theoretical
+        double ewma_max_multiplier = 10.0;  // 1000% of theoretical
+
+        // Default bandwidth (Gbps) when topology info unavailable
+        double default_bandwidth_gbps = 400.0;  // Default NIC bandwidth
+        double min_bandwidth_gbps = 10.0;       // Minimum valid NIC bandwidth
+        double max_bandwidth_gbps = 800.0;      // Maximum valid NIC bandwidth
+
+        // Shared slot rotation interval (milliseconds)
+        int slot_rotation_interval_ms = 2;
+
+        std::vector<int> device_base_priorities;
+    };
+
+    void setSchedulingParams(const SchedulingParams &params) {
+        sched_params_ = params;
     }
 
-    void setDiffusionInterval(uint64_t msec) {
-        diffusion_interval_ = msec * 1000000ull;
+    const SchedulingParams &getSchedulingParams() const {
+        return sched_params_;
     }
 
-    void setCrossNumaAccess(bool enable = true) { allow_cross_numa_ = enable; }
-
    private:
     std::shared_ptr<Topology> local_topology_;
     std::unordered_map<int, DeviceInfo> devices_;
-    mutable std::shared_mutex rwlock_;
-    bool allow_cross_numa_ = false;
-    double alpha_ = 0.01;
-    double local_weight_ = 0.9;
-    uint64_t diffusion_interval_ = 10 * 1000000ull;
-    std::shared_ptr<SharedQuotaManager> shared_quota_;
-    bool enable_quota_ = true;
-    bool update_quota_params_ = true;
+    std::shared_ptr<SharedSlotManager> slot_manager_;
+    bool smart_selection_enabled_ = true;
+    SchedulingParams sched_params_;
+
+    Status buildCandidates(const Topology::MemEntry *entry,
+                           uint64_t slice_bytes, uint64_t device_mask,
+                           std::vector<Candidate> &candidates,
+                           int request_priority = PRIO_HIGH);
+
+    void selectSinglePath(const std::vector<Candidate> &candidates,
+                          uint32_t num_slices, uint64_t total_length,
+                          std::vector<int> &slice_dev_ids);
+
+    void selectMultiPath(const std::vector<Candidate> &candidates,
+                         uint32_t num_slices, uint64_t total_length,
+                         std::vector<int> &slice_dev_ids,
+                         bool probe_mode = false);
 };
 
 }  // namespace tent
 }  // namespace mooncake
 
-#endif  // TENT_QUOTA_H
\ No newline at end of file
+#endif  // TENT_SELECTOR_H
\ No newline at end of file
diff --git a/mooncake-transfer-engine/tent/include/tent/transport/rdma/rail_monitor.h b/mooncake-transfer-engine/tent/include/tent/transport/rdma/rail_monitor.h
index a0dd7cc7..f69e453c 100644
--- a/mooncake-transfer-engine/tent/include/tent/transport/rdma/rail_monitor.h
+++ b/mooncake-transfer-engine/tent/include/tent/transport/rdma/rail_monitor.h
@@ -21,9 +21,27 @@
 
 namespace mooncake {
 namespace tent {
+// Not thread-safe. Each RDMA worker owns its own RailMonitor via
+// WorkerContext::rails, so markFailed / markRecovered / available / load
+// all execute on a single worker thread. Do not share across threads
+// without adding external synchronization.
 class RailMonitor {
     const static size_t kMaxNuma = 16;
 
+    // Upper bound on the exponential-backoff cooldown; prevents the
+    // per-rail pause from growing without bound under repeated failure.
+    static constexpr std::chrono::seconds kMaxCooldown{300};
+
+   public:
+    // Config keys. Exposed as constants so callers (and docs) reference
+    // a single source of truth instead of duplicating the string path.
+    static constexpr const char *kCfgErrorThreshold =
+        "transports/rdma/rail_error_threshold";
+    static constexpr const char *kCfgErrorWindowSecs =
+        "transports/rdma/rail_error_window_secs";
+    static constexpr const char *kCfgCooldownSecs =
+        "transports/rdma/rail_cooldown_secs";
+
    public:
     RailMonitor() = default;
 
@@ -33,8 +51,13 @@ class RailMonitor {
     RailMonitor &operator=(const RailMonitor &) = delete;
 
    public:
+    // rail_topo_json: optional JSON string describing rail topology.
+    // conf: optional Config pointer; when non-null, overrides the default
+    //   error_threshold / error_window_secs / cooldown_secs values via
+    //   kCfgErrorThreshold / kCfgErrorWindowSecs / kCfgCooldownSecs.
     Status load(const Topology *local, const Topology *remote,
-                const std::string &rail_topo_json = "");
+                const std::string &rail_topo_json = "",
+                const Config *conf = nullptr);
 
     bool ready() { return ready_; }
 
@@ -71,8 +94,12 @@ class RailMonitor {
         int error_count = 0;
         std::chrono::seconds cooldown{0};
         std::chrono::steady_clock::time_point last_error{};
-        bool paused = false;
         std::chrono::steady_clock::time_point resume_time{};
+
+        // Derived: a rail is paused iff a resume_time has been armed.
+        bool paused() const {
+            return resume_time != std::chrono::steady_clock::time_point{};
+        }
     };
 
     std::unordered_map<std::pair<int, int>, RailState, PairHash> rail_states_;
diff --git a/mooncake-transfer-engine/tent/include/tent/transport/rdma/rdma_transport.h b/mooncake-transfer-engine/tent/include/tent/transport/rdma/rdma_transport.h
index cabbaac4..2a739579 100644
--- a/mooncake-transfer-engine/tent/include/tent/transport/rdma/rdma_transport.h
+++ b/mooncake-transfer-engine/tent/include/tent/transport/rdma/rdma_transport.h
@@ -47,7 +47,7 @@ class LocalBuffers;
 using RdmaContextSet = std::vector<std::shared_ptr<RdmaContext>>;
 
 struct RdmaSubBatch : public Transport::SubBatch {
-    std::vector<RdmaTask> task_list;
+    std::vector<RdmaTask*> task_list;
     std::vector<RdmaSlice*> slice_chain;
     size_t max_size;
     virtual size_t size() const { return task_list.size(); }
@@ -114,10 +114,17 @@ class RdmaTransport : public Transport {
    public:
     Status setupLocalSegment();
 
+    std::shared_ptr<Config> config() const { return conf_; }
+
    private:
     bool installed_;
     std::shared_ptr<Config> conf_;
     std::string local_segment_name_;
+    // When MC_RDMA_BIND_ADDRESS is set in a dual-NIC environment,
+    // rdma_server_name_ holds the RDMA-reachable address for NIC path
+    // construction, while local_segment_name_ keeps the TCP-reachable
+    // address for P2P routing.
+    std::string rdma_server_name_;
     std::shared_ptr<Topology> local_topology_;
     std::shared_ptr<ControlService> metadata_;
     LocalBufferManager local_buffer_manager_;
diff --git a/mooncake-transfer-engine/tent/include/tent/transport/rdma/shared_quota.h b/mooncake-transfer-engine/tent/include/tent/transport/rdma/shared_quota.h
index e58af14b..7760fee3 100644
--- a/mooncake-transfer-engine/tent/include/tent/transport/rdma/shared_quota.h
+++ b/mooncake-transfer-engine/tent/include/tent/transport/rdma/shared_quota.h
@@ -12,8 +12,8 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-#ifndef TENT_SHARED_QUOTA_H
-#define TENT_SHARED_QUOTA_H
+#ifndef TENT_SHARED_SLOT_H
+#define TENT_SHARED_SLOT_H
 
 #include "quota.h"
 #include "tent/common/status.h"
@@ -31,60 +31,54 @@
 #include <vector>
 #include <iostream>
 #include <errno.h>
+#include <thread>
 
 namespace mooncake {
 namespace tent {
 
-static constexpr int MAX_DEVICES = 64;
-static constexpr int MAX_PID_SLOTS = 256;
-static constexpr uint64_t SHM_MAGIC = 0x2025082772805202ULL;
-static constexpr int SHM_VERSION = 1;
+static constexpr uint64_t SHM_MAGIC = 0x2025082772805203ULL;
+static constexpr int SHM_VERSION = 10;
 
-struct PidUsage {
-    pid_t pid;                     // 0 == free slot
-    volatile uint64_t used_bytes;  // local used bytes reported by this pid
-    uint8_t reserved[56];          // padding -> total 64B
-};
-
-struct SharedDeviceEntry {
-    char dev_name[56];  // NUL-terminated device name, empty means unused
-    volatile uint64_t active_bytes;
-    PidUsage pid_usages[MAX_PID_SLOTS];
-};
+// Time slice configuration for process-level coordination
+// Slot 0:            HIGH only
+// Slot 1:            MEDIUM + HIGH
+// Slot 2:            ALL (LOW + MEDIUM + HIGH)
+// Then repeat
+static constexpr int NUM_SLOTS = PRIO_LOW + 1;
 
 struct SharedHeader {
     uint64_t magic;
     int32_t version;
-    int32_t num_devices;
+    std::atomic<int> current_slot;
     pthread_mutex_t global_mutex;
-    SharedDeviceEntry devices[MAX_DEVICES];
 };
 
-class DeviceQuota;
-class SharedQuotaManager {
+class DeviceSelector;
+class SharedSlotManager {
    public:
-    explicit SharedQuotaManager(DeviceQuota* local_quota);
-    ~SharedQuotaManager();
+    explicit SharedSlotManager(DeviceSelector* local_quota);
+    ~SharedSlotManager();
 
     Status attach(const std::string& shm_name);
     Status detach();
 
-    Status diffusion();
+    // Check if current process can send (global slot)
+    // Returns true if given priority is allowed in current global slot
+    bool canSend(int priority = PRIO_HIGH);
+
+    // Set slot duration in milliseconds (must be > 0)
+    void setRotationIntervalMs(int ms) { rotation_interval_ms_ = ms; }
+    int getRotationIntervalMs() const { return rotation_interval_ms_; }
 
    private:
-    Status attachProcess();
-    Status detachProcess();
-
-   private:
-    PidUsage* findOrCreatePidSlotLocked(int dev_id, pid_t pid);
-    PidUsage* findPidSlotLocked(int dev_id, pid_t pid);
-    int findDeviceIdByNameLocked(const std::string& dev_name);
+    void startBackgroundThread();
+    void stopBackgroundThread();
+    void backgroundThreadLoop();
     Status initializeHeader();
     Status initMutex(pthread_mutex_t* m);
-    void reclaimDeadPidsInternal();
-    static bool isPidAlive(pid_t pid);
-    int lock();
-    int unlock();
+
+    // Check if a priority is allowed in the given slot
+    bool isPriorityAllowedInSlot(int priority, int slot) const;
 
    private:
     std::string name_;
@@ -92,10 +86,15 @@ class SharedQuotaManager {
     int fd_;
     size_t size_;
     bool created_;
-    DeviceQuota* local_quota_;
+    DeviceSelector* device_selector_;
+    int rotation_interval_ms_ = 2;  // Default: 2ms per slot
+
+    // Background thread
+    std::thread background_thread_;
+    std::atomic<bool> background_running_;
 };
 
 }  // namespace tent
 }  // namespace mooncake
 
-#endif  // TENT_SHARED_QUOTA_H
\ No newline at end of file
+#endif  // TENT_SHARED_SLOT_H
\ No newline at end of file
diff --git a/mooncake-transfer-engine/tent/include/tent/transport/rdma/slice.h b/mooncake-transfer-engine/tent/include/tent/transport/rdma/slice.h
index c4bfe1fb..f681c8d2 100644
--- a/mooncake-transfer-engine/tent/include/tent/transport/rdma/slice.h
+++ b/mooncake-transfer-engine/tent/include/tent/transport/rdma/slice.h
@@ -33,12 +33,20 @@
 namespace mooncake {
 namespace tent {
 struct RdmaSlice;
+class RailMonitor;
 
 struct RdmaSliceList {
     RdmaSlice* first = nullptr;
     int num_slices = 0;
 };
 
+// Forward declarations
+class RdmaEndPoint;
+struct RdmaTask;
+
+using RdmaSliceStorage = Slab<RdmaSlice>;
+using RdmaTaskStorage = Slab<RdmaTask>;
+
 struct RdmaTask {
     int num_slices;
     Request request;
@@ -47,6 +55,16 @@ struct RdmaTask {
     volatile int success_slices;
     volatile int resolved_slices;
     volatile TransferStatusEnum first_error = PENDING;
+
+    // Reference counting for UAF protection
+    std::atomic<int> ref_count{0};
+
+    void ref() { ref_count.fetch_add(1, std::memory_order_relaxed); }
+    void deref() {
+        if (ref_count.fetch_sub(1, std::memory_order_acq_rel) == 1) {
+            RdmaTaskStorage::Get().deallocate(this);
+        }
+    }
 };
 
 class RdmaEndPoint;
@@ -71,10 +89,16 @@ struct RdmaSlice {
     bool failed = false;
     uint64_t enqueue_ts = 0;
     uint64_t submit_ts = 0;
+    // Non-owning pointer to the per-worker RailMonitor for this slice's
+    // target machine, resolved once in generatePostPath. Lets asyncPollCq
+    // and disableEndpoint call markRecovered / markFailed without a
+    // string-keyed map lookup on the RDMA hot path. Stable because
+    // WorkerContext::rails stores values via unique_ptr, so rehashes do
+    // not invalidate the pointee.
+    RailMonitor* rail_monitor = nullptr;
+    int priority = PRIO_HIGH;
 };
 
-using RdmaSliceStorage = Slab<RdmaSlice>;
-
 static inline void updateSliceStatus(RdmaSlice* slice,
                                      TransferStatusEnum status) {
     if (status == PENDING) return;
@@ -94,6 +118,7 @@ static inline void updateSliceStatus(RdmaSlice* slice,
         if (final_st == PENDING) final_st = FAILED;
         __sync_bool_compare_and_swap(&task->status_word, PENDING, final_st);
     }
+    task->deref();
 }
 
 }  // namespace tent
diff --git a/mooncake-transfer-engine/tent/include/tent/transport/rdma/workers.h b/mooncake-transfer-engine/tent/include/tent/transport/rdma/workers.h
index 6688ab7d..83b3da8e 100644
--- a/mooncake-transfer-engine/tent/include/tent/transport/rdma/workers.h
+++ b/mooncake-transfer-engine/tent/include/tent/transport/rdma/workers.h
@@ -16,6 +16,7 @@
 #define TENT_WORKERS_H
 
 #include <future>
+#include <memory>
 #include <queue>
 #include <thread>
 #include <unordered_set>
@@ -28,11 +29,14 @@
 #include "rail_monitor.h"
 #include "tent/common/utils/os.h"
 #include "tent/common/concurrent/bounded_mpsc_queue.h"
+#include "tent/common/types.h"
 
 namespace mooncake {
 namespace tent {
 
 class RdmaTransport;
+class DeviceSelector;
+
 class Workers {
    public:
     static constexpr size_t kCapacity = 1024 * 8;
@@ -53,6 +57,8 @@ class Workers {
 
     Status cancel(RdmaSliceList &slice_list);
 
+    DeviceSelector *getDeviceSelector() const { return device_selector_.get(); }
+
    private:
     using Task = std::function<void()>;
 
@@ -175,9 +181,11 @@ class Workers {
         PerfMetric inflight_lat;
     };
 
+    static constexpr int kNumPriorityLevels = PRIO_LOW + 1;
+
     struct WorkerContext {
         std::thread thread;
-        BoundedSliceQueue queue;
+        BoundedSliceQueue queues[kNumPriorityLevels];  // Priority queues
         GroupedRequests requests;
         std::unordered_set<RdmaSlice *> inflight_slice_set;
         std::atomic<int64_t> inflight_slices = 0;
@@ -186,15 +194,25 @@ class Workers {
         std::condition_variable cv;
         volatile bool in_suspend = false;
 
-        std::unordered_map<std::string, RailMonitor> rails;
+        // Next time to check for priority promotions (nanoseconds)
+        uint64_t next_promotion_check_ns = 0;
+
+        // Values are held via unique_ptr so that map rehashing does not
+        // invalidate pointers into RailMonitor stored on in-flight slices
+        // (see RdmaSlice::rail_monitor).
+        std::unordered_map<std::string, std::unique_ptr<RailMonitor>> rails;
         PerfMetricSummary perf;
-        uint64_t padding[16];
+        uint64_t padding[15];
     };
 
+    // Promote timed-out low priority requests to higher priority queues
+    void promoteTimedOutRequests(WorkerContext &worker);
+
     WorkerContext *worker_context_;
     uint64_t slice_timeout_ns_;
+    uint64_t priority_promotion_timeout_ns_;  // Timeout for priority promotion
 
-    std::unique_ptr<DeviceQuota> device_quota_;
+    std::unique_ptr<DeviceSelector> device_selector_;
     bool always_tier1_ = false;
 };
 }  // namespace tent
diff --git a/mooncake-transfer-engine/tent/include/tent/transport/sunrise_link/sunrise_link_transport.h b/mooncake-transfer-engine/tent/include/tent/transport/sunrise_link/sunrise_link_transport.h
new file mode 100644
index 00000000..6fdae540
--- /dev/null
+++ b/mooncake-transfer-engine/tent/include/tent/transport/sunrise_link/sunrise_link_transport.h
@@ -0,0 +1,148 @@
+// Copyright 2024 KVCache.AI
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#pragma once
+
+#include <map>
+#include <mutex>
+#include <utility>
+#include <vector>
+#include <unordered_map>
+
+#include "tent/runtime/control_plane.h"
+#include "tent/runtime/transport.h"
+
+namespace mooncake {
+namespace tent {
+
+struct SunriseLinkTask {
+    Request request;
+    volatile TransferStatusEnum status_word;
+    volatile size_t transferred_bytes;
+    uint64_t target_addr = 0;
+    bool is_tang_ipc = false;
+    int tang_gpu_id = -1;
+};
+
+struct SunriseLinkSubBatch : public Transport::SubBatch {
+    std::vector<SunriseLinkTask> task_list;
+    size_t max_size = 0;
+    /// Opaque `tangStream_t` when built with Tang; mirrors
+    /// NVLinkSubBatch::stream.
+    void* stream = nullptr;
+    virtual size_t size() const { return task_list.size(); }
+};
+
+struct SunriseLinkDeviceInfo {
+    int device_id;
+    std::vector<int> active_ports;
+};
+
+struct SunriseLinkPeerInfo {
+    int peer_device_id;
+    int peer_port;
+    bool is_active;
+};
+
+class SunriseLinkTransport : public Transport {
+   public:
+    SunriseLinkTransport();
+    ~SunriseLinkTransport() override;
+
+    Status install(std::string& local_segment_name,
+                   std::shared_ptr<ControlService> metadata,
+                   std::shared_ptr<Topology> local_topology,
+                   std::shared_ptr<Config> conf = nullptr) override;
+
+    Status uninstall() override;
+
+    Status allocateSubBatch(SubBatchRef& batch, size_t max_size) override;
+
+    Status freeSubBatch(SubBatchRef& batch) override;
+
+    Status submitTransferTasks(
+        SubBatchRef batch, const std::vector<Request>& request_list) override;
+
+    Status getTransferStatus(SubBatchRef batch, int task_id,
+                             TransferStatus& status) override;
+
+    Status addMemoryBuffer(BufferDesc& desc,
+                           const MemoryOptions& options) override;
+
+    Status removeMemoryBuffer(BufferDesc& desc) override;
+
+    const char* getName() const override { return "sunrise_link"; }
+
+    bool supportNotification() const override { return true; }
+
+   private:
+    Status startTransfer(SunriseLinkTask* task, SunriseLinkSubBatch* batch);
+
+    // SunriseLink initialization
+    Status initSunriseLink();
+
+    // Memory registration
+    Status registerMemory(void* addr, size_t size);
+
+    // Memory deregistration
+    Status deregisterMemory(void* addr);
+
+    // Topology detection
+    Status detectTopology();
+    Status relocateRemoteAddress(uint64_t& dest_addr, uint64_t length,
+                                 uint64_t target_id,
+                                 int* target_gpu_id = nullptr);
+
+    // Runtime library handle
+    void* runtime_lib_handle_ = nullptr;
+
+    void* ptml_handle_ = nullptr;
+
+    // Local device information
+    std::vector<SunriseLinkDeviceInfo> local_devices_;
+
+    // Topology map: device_id -> peer info
+    std::map<int, std::vector<SunriseLinkPeerInfo>> topology_map_;
+
+    // (local_chipid, remote_chipid) -> local PTLink port (S2 C2C)
+    std::map<std::pair<int, int>, int> c2c_port_by_pair_;
+
+    // Registered memory regions
+    std::map<void*, size_t> registered_memory_;
+    std::map<void*, int> registered_memory_gpu_id_;
+    mutable std::mutex registered_memory_mutex_;
+
+    struct OpenedIpcEntry {
+        void* dev_ptr{nullptr};
+        uint64_t length{0};
+        int gpu_id{-1};
+    };
+    using IpcRelocateMap =
+        std::unordered_map<SegmentID,
+                           std::unordered_map<uint64_t, OpenedIpcEntry>>;
+    RWSpinlock relocate_lock_;
+    IpcRelocateMap relocate_map_;
+
+    // Installed flag
+    bool installed_ = false;
+
+    // Config
+    std::shared_ptr<Config> conf_;
+    std::shared_ptr<ControlService> metadata_;
+
+    uint64_t async_memcpy_threshold_ = 0;
+};
+
+}  // namespace tent
+}  // namespace mooncake
diff --git a/mooncake-transfer-engine/tent/include/tent/transport/tcp/tcp_transport.h b/mooncake-transfer-engine/tent/include/tent/transport/tcp/tcp_transport.h
index acdf117d..3d5f2ebe 100644
--- a/mooncake-transfer-engine/tent/include/tent/transport/tcp/tcp_transport.h
+++ b/mooncake-transfer-engine/tent/include/tent/transport/tcp/tcp_transport.h
@@ -38,6 +38,8 @@ struct TcpParams {
 
 struct TcpTask {
     Request request;
+    BatchID progress_batch_id{0};
+    std::function<void(BatchID)> notify_progress;
     std::atomic<TransferStatusEnum> status_word{TransferStatusEnum::PENDING};
     std::atomic<size_t> transferred_bytes{0};
     uint64_t target_addr = 0;
@@ -45,6 +47,8 @@ struct TcpTask {
     TcpTask() = default;
     TcpTask(TcpTask &&other) noexcept
         : request(std::move(other.request)),
+          progress_batch_id(other.progress_batch_id),
+          notify_progress(std::move(other.notify_progress)),
           status_word(other.status_word.load(std::memory_order_relaxed)),
           transferred_bytes(
               other.transferred_bytes.load(std::memory_order_relaxed)),
diff --git a/mooncake-transfer-engine/tent/plugins/CMakeLists.txt b/mooncake-transfer-engine/tent/plugins/CMakeLists.txt
index 7b564750..e88e35a9 100644
--- a/mooncake-transfer-engine/tent/plugins/CMakeLists.txt
+++ b/mooncake-transfer-engine/tent/plugins/CMakeLists.txt
@@ -11,3 +11,7 @@ target_compile_features(tent_plugin INTERFACE cxx_std_20)
 if (USE_CUDA)
     add_subdirectory(cuda)
 endif()
+
+if (USE_HIP)
+    add_subdirectory(rocm)
+endif()
diff --git a/mooncake-transfer-engine/tent/plugins/rocm/CMakeLists.txt b/mooncake-transfer-engine/tent/plugins/rocm/CMakeLists.txt
new file mode 100644
index 00000000..1aec2786
--- /dev/null
+++ b/mooncake-transfer-engine/tent/plugins/rocm/CMakeLists.txt
@@ -0,0 +1,22 @@
+set(CMAKE_POSITION_INDEPENDENT_CODE ON)
+
+find_package(HIP REQUIRED)
+
+set(TENT_ROCM_PLUGIN_SOURCES
+    ${CMAKE_CURRENT_SOURCE_DIR}/rocm_plugin.cpp
+)
+
+add_library(tent_rocm_plugin SHARED ${TENT_ROCM_PLUGIN_SOURCES})
+
+target_link_libraries(tent_rocm_plugin
+    PRIVATE hip::host tent_plugin
+)
+
+set_target_properties(tent_rocm_plugin PROPERTIES
+    OUTPUT_NAME "tent_rocm"
+)
+
+install(TARGETS tent_rocm_plugin
+    LIBRARY DESTINATION lib/tent_plugins
+    RUNTIME DESTINATION lib/tent_plugins
+)
diff --git a/mooncake-transfer-engine/tent/plugins/rocm/rocm_plugin.cpp b/mooncake-transfer-engine/tent/plugins/rocm/rocm_plugin.cpp
new file mode 100644
index 00000000..63d26490
--- /dev/null
+++ b/mooncake-transfer-engine/tent/plugins/rocm/rocm_plugin.cpp
@@ -0,0 +1,146 @@
+// Copyright 2025 KVCache.AI
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "tent/device_plugin.h"
+
+#include <hip/hip_runtime.h>
+#include <string.h>
+#include <stdio.h>
+#include <string>
+#include <glog/logging.h>
+
+struct rocm_plugin_ctx_t {
+    // reserved
+};
+
+#define CHECK_HIP(call)                                                       \
+    do {                                                                      \
+        auto err = call;                                                      \
+        if (err != hipSuccess) {                                              \
+            LOG(ERROR) << std::string(#call) + ": " + hipGetErrorString(err); \
+            return -1;                                                        \
+        }                                                                     \
+    } while (0)
+
+class LocationParser {
+   public:
+    LocationParser(const std::string& location) {
+        size_t colonPos = location.find(':');
+        if (colonPos == std::string::npos) {
+            index_ = -1;
+            return;
+        }
+        std::string type = location.substr(0, colonPos);
+        std::string indexStr = location.substr(colonPos + 1);
+        try {
+            type_ = type;
+            index_ = std::stoi(indexStr);
+        } catch (const std::exception& e) {
+            index_ = -1;
+        }
+    }
+
+    std::string type() const { return type_; }
+    int index() const { return index_; }
+
+   private:
+    std::string type_;
+    int index_;
+};
+
+static void* rocm_create_plugin() {
+    rocm_plugin_ctx_t* ctx = new rocm_plugin_ctx_t;
+    return ctx;
+}
+
+static int rocm_destroy_plugin(void* handle) {
+    if (!handle) return 0;
+    delete reinterpret_cast<rocm_plugin_ctx_t*>(handle);
+    return 0;
+}
+
+static int rocm_alloc(void* ctx_, void** pptr, size_t size, const char* loc) {
+    (void)ctx_;
+    LocationParser location(loc);
+    if (location.type() != "rocm") return -1;
+    int hip_dev = 0;
+    CHECK_HIP(hipGetDevice(&hip_dev));
+    CHECK_HIP(hipSetDevice(location.index()));
+    CHECK_HIP(hipMalloc(pptr, size));
+    CHECK_HIP(hipSetDevice(hip_dev));
+    return 0;
+}
+
+static int rocm_free(void* ctx_, void* ptr, size_t size) {
+    (void)ctx_;
+    (void)size;
+    hipPointerAttribute_t attrs;
+    CHECK_HIP(hipPointerGetAttributes(&attrs, ptr));
+    if (attrs.type == hipMemoryTypeDevice) {
+        CHECK_HIP(hipFree(ptr));
+        return 0;
+    }
+    return -2;
+}
+
+static int rocm_memcpy_sync(void* ctx_, void* dst, void* src, size_t length) {
+    (void)ctx_;
+    CHECK_HIP(hipMemcpy(dst, src, length, hipMemcpyDefault));
+    return 0;
+}
+
+static int rocm_query_location(void* ctx_, void* addr, size_t size,
+                               location_t* buf, size_t buf_count) {
+    (void)ctx_;
+    if (buf_count == 0) return -1;
+    hipPointerAttribute_t attr{};
+    hipError_t err = hipPointerGetAttributes(&attr, addr);
+    if (err != hipSuccess || attr.type != hipMemoryTypeDevice) return 0;
+    buf[0].start = addr;
+    buf[0].length = size;
+    snprintf(buf[0].location, LOCATION_LEN, "rocm:%d", attr.device);
+    return 1;
+}
+
+static int rocm_get_device_count(void* ctx_) {
+    (void)ctx_;
+    int count = 0;
+    CHECK_HIP(hipGetDeviceCount(&count));
+    return count;
+}
+
+static int rocm_get_device_pci_bus_id(void* ctx_, int device_index,
+                                      char* bus_id, size_t bus_id_len) {
+    (void)ctx_;
+    hipDeviceProp_t prop;
+    CHECK_HIP(hipGetDeviceProperties(&prop, device_index));
+    snprintf(bus_id, bus_id_len, "%04x:%02x:%02x.0", prop.pciDomainID,
+             prop.pciBusID, prop.pciDeviceID);
+    return 0;
+}
+
+extern "C" int tent_register_device_plugin(device_plugin_t* out) {
+    if (!out) return -1;
+    memset(out, 0, sizeof(*out));
+    out->class_name = "rocm";
+    out->create_plugin = rocm_create_plugin;
+    out->destroy_plugin = rocm_destroy_plugin;
+    out->alloc = rocm_alloc;
+    out->free = rocm_free;
+    out->memcpy_sync = rocm_memcpy_sync;
+    out->query_location = rocm_query_location;
+    out->get_device_count = rocm_get_device_count;
+    out->get_device_pci_bus_id = rocm_get_device_pci_bus_id;
+    return 0;
+}
diff --git a/mooncake-transfer-engine/tent/src/CMakeLists.txt b/mooncake-transfer-engine/tent/src/CMakeLists.txt
index 7fb3174b..73766ce1 100644
--- a/mooncake-transfer-engine/tent/src/CMakeLists.txt
+++ b/mooncake-transfer-engine/tent/src/CMakeLists.txt
@@ -23,6 +23,18 @@ else()
   message(STATUS "CUDA: Disabled")
 endif()
 
+# ROCm / HIP
+if(USE_HIP)
+  list(APPEND CMAKE_PREFIX_PATH "/opt/rocm/lib/cmake")
+  find_package(HIP REQUIRED)
+  message(STATUS "ROCm/HIP: Enabled")
+  target_compile_definitions(tent_interface INTERFACE USE_HIP __HIP_PLATFORM_AMD__)
+  target_include_directories(tent_interface INTERFACE ${HIP_INCLUDE_DIRS})
+  target_link_libraries(tent_interface INTERFACE hip::host)
+else()
+  message(STATUS "ROCm/HIP: Disabled")
+endif()
+
 # GDS
 find_library(CUFILE_LIB cufile PATHS /usr/local/cuda/lib64)
 find_path(CUFILE_INCLUDE cufile.h PATHS /usr/local/cuda/include)
@@ -65,6 +77,15 @@ else()
   message(STATUS "Uring: Disabled")
 endif()
 
+if(USE_SUNRISE)
+  target_compile_definitions(tent_interface INTERFACE USE_SUNRISE)
+  target_include_directories(tent_interface INTERFACE ${MC_TANGRT_ROOT}/include)
+  target_link_directories(tent_interface INTERFACE
+                          ${MC_TANGRT_ROOT}/lib/linux-x86_64
+                          ${MC_TANGRT_ROOT}/lib)
+  target_link_libraries(tent_interface INTERFACE tangrt_shared ptml_shared dl)
+endif()
+
 set(YALANTING_TARGET yalantinglibs::yalantinglibs)
 if(NOT TARGET ${YALANTING_TARGET})
   if(CMAKE_SOURCE_DIR STREQUAL CMAKE_CURRENT_SOURCE_DIR)
@@ -81,7 +102,14 @@ add_subdirectory(runtime)
 add_subdirectory(platform)
 add_subdirectory(transport)
 add_subdirectory(metrics)
-add_subdirectory(python)
+if(NOT TARGET pybind11::module)
+  find_package(pybind11 QUIET)
+endif()
+if(TARGET pybind11::module)
+  add_subdirectory(python)
+else()
+  message(STATUS "pybind11 not found, skipping Tent Python bindings")
+endif()
 
 file(GLOB TENT_ENGINE_SOURCES "${CMAKE_CURRENT_SOURCE_DIR}/*.cpp")
 
@@ -100,7 +128,9 @@ foreach(
   metastore_etcd
   tent_platform_all
   platform_cuda
+  platform_rocm
   platform_ascend
+  platform_sunrise
   tent_xport_gds
   tent_xport_uring
   tent_xport_bufio
@@ -110,13 +140,13 @@ foreach(
   tent_xport_shm
   tent_xport_tcp
   tent_xport_ascend_direct
+  tent_xport_sunrise_link
   tent_metrics)
   if(TARGET ${tgt})
     target_link_libraries(tent_link_group INTERFACE ${tgt})
   endif()
 endforeach()
 target_link_libraries(tent_link_group INTERFACE "-Wl,--end-group")
-
 install(
   TARGETS tent_shared
   ARCHIVE DESTINATION lib
diff --git a/mooncake-transfer-engine/tent/src/common/config.cpp b/mooncake-transfer-engine/tent/src/common/config.cpp
index 0d725dc7..f390333c 100644
--- a/mooncake-transfer-engine/tent/src/common/config.cpp
+++ b/mooncake-transfer-engine/tent/src/common/config.cpp
@@ -16,6 +16,8 @@
 
 #include <glog/logging.h>
 
+#include <algorithm>
+#include <cctype>
 #include <sstream>
 
 namespace mooncake {
@@ -84,16 +86,44 @@ Status ConfigHelper::loadFromEnv(Config& config) {
         }
     }
 
-    // Legacy keys for backward compatibility
+    // Legacy keys for backward compatibility (MC_* env vars)
+    setConfig(config, "MC_RDMA_BIND_ADDRESS", "transports/rdma/bind_address");
+    setConfig(config, "MC_NUM_CQ_PER_CTX",
+              "transports/rdma/device/num_cq_list");
+    setConfig(config, "MC_NUM_COMP_CHANNELS_PER_CTX",
+              "transports/rdma/device/num_comp_channels");
     setConfig(config, "MC_IB_PORT", "transports/rdma/device/port");
     setConfig(config, "MC_GID_INDEX", "transports/rdma/device/gid_index");
+    setConfig(config, "NCCL_IB_GID_INDEX", "transports/rdma/device/gid_index");
+    setConfig(config, "MC_MAX_CQE_PER_CTX", "transports/rdma/device/max_cqe");
+    setConfig(config, "MC_MAX_EP_PER_CTX",
+              "transports/rdma/endpoint/endpoint_store_cap");
+    setConfig(config, "MC_NUM_QP_PER_EP",
+              "transports/rdma/endpoint/qp_mul_factor");
+    setConfig(config, "MC_MAX_SGE", "transports/rdma/endpoint/max_sge");
+    setConfig(config, "MC_MAX_WR", "transports/rdma/endpoint/max_qp_wr");
+    setConfig(config, "MC_MAX_INLINE",
+              "transports/rdma/endpoint/max_inline_bytes");
+    setConfig(config, "MC_PKEY_INDEX", "transports/rdma/endpoint/pkey_index");
+    setConfig(config, "MC_MTU", "transports/rdma/endpoint/path_mtu");
+    setConfig(config, "MC_IB_TC", "transports/rdma/endpoint/traffic_class");
+    setConfig(config, "MC_IB_SL", "transports/rdma/endpoint/service_level");
+    setConfig(config, "MC_IB_PCI_RELAXED_ORDERING",
+              "transports/rdma/pci_relaxed_ordering");
+    setConfig(config, "MC_WORKERS_PER_CTX",
+              "transports/rdma/workers/num_workers");
+    setConfig(config, "MC_SLICE_SIZE", "transports/rdma/workers/block_size");
+    setConfig(config, "MC_RETRY_CNT",
+              "transports/rdma/workers/max_retry_count");
+    setConfig(config, "MC_DISABLE_GPU_DIRECT_RDMA",
+              "transports/rdma/disable_gpu_direct_rdma");
     return status;
 }
 
 bool ConfigHelper::parseBool(const std::string& str, bool default_value) {
     std::string lower_str = str;
     std::transform(lower_str.begin(), lower_str.end(), lower_str.begin(),
-                   ::tolower);
+                   [](unsigned char c) { return std::tolower(c); });
 
     if (lower_str == "true" || lower_str == "1" || lower_str == "yes" ||
         lower_str == "on") {
diff --git a/mooncake-transfer-engine/tent/src/common/ip.cpp b/mooncake-transfer-engine/tent/src/common/ip.cpp
index 299c0fa6..a1323806 100644
--- a/mooncake-transfer-engine/tent/src/common/ip.cpp
+++ b/mooncake-transfer-engine/tent/src/common/ip.cpp
@@ -13,6 +13,8 @@
 // limitations under the License.
 
 #include "tent/common/utils/ip.h"
+#include "common.h"
+
 #include <glog/logging.h>
 
 namespace mooncake {
@@ -92,42 +94,60 @@ Status checkLocalIpAddress(std::string &hostname, bool &ipv6) {
 std::pair<std::string, uint16_t> parseHostNameWithPort(const std::string &url,
                                                        uint16_t default_port) {
     uint16_t port = default_port;
+    if (url.empty()) return std::make_pair(url, port);
 
-    // Check if url is valid IPv6 address
-    in6_addr addr;
-    if (inet_pton(AF_INET6, url.c_str(), &addr) == 1) return {url, port};
-
-    // Check if url is valid IPv6 address with port
-    size_t start_pos = 0;
-    size_t end_pos = url.find_last_of(']');
-    size_t port_pos = url.find_last_of(':');
-    if (url.front() == '[' && end_pos != std::string::npos &&
-        port_pos > end_pos) {
-        auto ip = url.substr(start_pos + 1, end_pos - start_pos - 1);
-        std::string port_str = url.substr(port_pos + 1);
-        int val = std::atoi(port_str.c_str());
-        if (val <= 0 || val > 65535) {
-            LOG(WARNING) << "Illegal port number in " << url
-                         << ". Use default port " << port << " instead";
-        } else {
-            port = static_cast<uint16_t>(val);
+    // Helper: check if a string is a valid numeric IP (IPv4 or IPv6),
+    // automatically stripping scope IDs (e.g., %eth0) before validation
+    auto isNumericAddress = [](const std::string &host) -> bool {
+        std::string check = host;
+        auto pct = host.find('%');
+        if (pct != std::string::npos) {
+            if (host.find(':', pct) != std::string::npos) {
+                return false;
+            }
+            check = host.substr(0, pct);
         }
-        return std::make_pair(port_str, port);
+        struct addrinfo hints, *res;
+        memset(&hints, 0, sizeof(hints));
+        hints.ai_flags = AI_NUMERICHOST;
+        hints.ai_family = AF_UNSPEC;
+        hints.ai_socktype = SOCK_STREAM;
+        int ret = getaddrinfo(check.c_str(), nullptr, &hints, &res);
+        if (ret == 0) freeaddrinfo(res);
+        return ret == 0;
+    };
+
+    if (isNumericAddress(url)) return {url, port};
+
+    size_t end_pos = url.find_last_of(']');
+    if (url.front() == '[' && end_pos != std::string::npos) {
+        std::string ip = url.substr(1, end_pos - 1);
+        if (!isNumericAddress(ip)) {
+            LOG(WARNING) << "Illegal bracketed IPv6 address: " << ip;
+            return std::make_pair(std::string(), port);
+        }
+        if (end_pos + 1 < url.size() && url[end_pos + 1] == ':') {
+            port = getPortFromString(url.substr(end_pos + 2), default_port);
+        }
+        return std::make_pair(ip, port);
     }
 
-    // Check if url has port field
-    auto pos = url.find(':');
-    if (pos == url.npos) return std::make_pair(url, port);
-    auto trimmed_server_name = url.substr(0, pos);
-    auto port_str = url.substr(pos + 1);
-    int val = std::atoi(port_str.c_str());
-    if (val <= 0 || val > 65535)
-        LOG(WARNING) << "Illegal port number in " << url
-                     << ". Use default port " << port << " instead";
-    else
-        port = (uint16_t)val;
+    auto pos = url.rfind(':');
+    if (pos != std::string::npos) {
+        std::string host_part = url.substr(0, pos);
+        std::string port_str = url.substr(pos + 1);
 
-    return std::make_pair(trimmed_server_name, port);
+        if (!isNumericAddress(host_part) &&
+            host_part.find(':') != std::string::npos) {
+            LOG(WARNING) << "Illegal IPv6 address: " << host_part;
+            return std::make_pair(std::string(), port);
+        }
+
+        return std::make_pair(host_part,
+                              getPortFromString(port_str, default_port));
+    }
+
+    return std::make_pair(url, port);
 }
 
 std::string buildIpAddrWithPort(const std::string &hostname, uint16_t port,
@@ -138,4 +158,4 @@ std::string buildIpAddrWithPort(const std::string &hostname, uint16_t port,
         return hostname + ":" + std::to_string(port);
 }
 }  // namespace tent
-}  // namespace mooncake
\ No newline at end of file
+}  // namespace mooncake
diff --git a/mooncake-transfer-engine/tent/src/metrics/tent_metrics.cpp b/mooncake-transfer-engine/tent/src/metrics/tent_metrics.cpp
index f678eb9f..4eb9f28e 100644
--- a/mooncake-transfer-engine/tent/src/metrics/tent_metrics.cpp
+++ b/mooncake-transfer-engine/tent/src/metrics/tent_metrics.cpp
@@ -173,8 +173,8 @@ void TentMetrics::shutdown() {
 void TentMetrics::registerMetrics() {
     // Pre-allocate vectors to avoid reallocation
     counters_.reserve(7);
-    histograms_.reserve(4);
-    histogram_boundaries_.reserve(4);
+    histograms_.reserve(5);
+    histogram_boundaries_.reserve(5);
 
     // Register all counters - add new counters here
     counters_ = {
@@ -186,16 +186,12 @@ void TentMetrics::registerMetrics() {
     // Register all histograms - add new histograms here
     // Note: histogram_boundaries_ must match the order of histograms_
     histograms_ = {
-        &read_latency_,
-        &write_latency_,
-        &read_size_,
-        &write_size_,
+        &read_latency_, &write_latency_, &read_size_,
+        &write_size_,   &deadline_mlu_,
     };
     histogram_boundaries_ = {
-        kLatencyBuckets,
-        kLatencyBuckets,
-        kSizeBuckets,
-        kSizeBuckets,
+        kLatencyBuckets, kLatencyBuckets,     kSizeBuckets,
+        kSizeBuckets,    kMluPerMilleBuckets,
     };
 }
 
@@ -229,6 +225,15 @@ void TentMetrics::recordWriteCompleted(size_t bytes, double latency_seconds) {
     }
 }
 
+void TentMetrics::recordDeadlineMLU(double mlu) {
+    // Fast path: check runtime switch first
+    if (!initialized_ || !runtime_enabled_.load(std::memory_order_relaxed))
+        return;
+    if (mlu < 0.0) return;  // defensive: ignore invalid (e.g. window <= 0)
+    // Store in per-mille so the integer histogram can bucket fractional ratios.
+    deadline_mlu_.observe(static_cast<int64_t>(mlu * 1000.0));
+}
+
 void TentMetrics::recordReadFailed(size_t bytes) {
     // Fast path: check runtime switch first
     if (!initialized_ || !runtime_enabled_.load(std::memory_order_relaxed))
@@ -384,6 +389,7 @@ void TentMetrics::recordWriteCompleted(size_t, double) {}
 void TentMetrics::recordReadFailed(size_t) {}
 void TentMetrics::recordWriteFailed(size_t) {}
 void TentMetrics::recordTransportFailover() {}
+void TentMetrics::recordDeadlineMLU(double) {}
 
 std::string TentMetrics::getPrometheusMetrics() {
     return "# TENT metrics disabled at compile time\n";
diff --git a/mooncake-transfer-engine/tent/src/platform/CMakeLists.txt b/mooncake-transfer-engine/tent/src/platform/CMakeLists.txt
index d17f6436..2f48b686 100644
--- a/mooncake-transfer-engine/tent/src/platform/CMakeLists.txt
+++ b/mooncake-transfer-engine/tent/src/platform/CMakeLists.txt
@@ -1,12 +1,20 @@
 add_subdirectory(cuda)
+add_subdirectory(rocm)
 add_subdirectory(ascend)
+add_subdirectory(sunrise)
 
 file(GLOB PLATFORM_SOURCES "*.cpp")
 add_library(tent_platform_all STATIC ${PLATFORM_SOURCES})
 target_link_libraries(tent_platform_all PUBLIC tent_common)
-if (TARGET platform_cuda)
-    target_link_libraries(tent_platform_all PUBLIC platform_cuda)
+if(TARGET platform_cuda)
+  target_link_libraries(tent_platform_all PUBLIC platform_cuda)
 endif()
-if (TARGET platform_ascend)
-    target_link_libraries(tent_platform_all PUBLIC platform_ascend)
+if(TARGET platform_rocm)
+  target_link_libraries(tent_platform_all PUBLIC platform_rocm)
+endif()
+if(TARGET platform_ascend)
+  target_link_libraries(tent_platform_all PUBLIC platform_ascend)
+endif()
+if(TARGET platform_sunrise)
+  target_link_libraries(tent_platform_all PUBLIC platform_sunrise)
 endif()
diff --git a/mooncake-transfer-engine/tent/src/platform/cuda/cuda_probe.cpp b/mooncake-transfer-engine/tent/src/platform/cuda/cuda_probe.cpp
index 78327f16..09345e98 100644
--- a/mooncake-transfer-engine/tent/src/platform/cuda/cuda_probe.cpp
+++ b/mooncake-transfer-engine/tent/src/platform/cuda/cuda_probe.cpp
@@ -16,6 +16,7 @@
 #include "tent/common/status.h"
 #include "tent/common/utils/prefault.h"
 #include "tent/common/utils/random.h"
+#include "char_util.h"
 
 #include <glog/logging.h>
 #include <fstream>
@@ -26,7 +27,6 @@
 #include <utility>
 #include <vector>
 #include <cuda_runtime.h>
-#include <ctype.h>
 #include <dirent.h>
 #include <infiniband/verbs.h>
 #include <limits.h>
@@ -191,7 +191,7 @@ static void discoverCudaTopology(std::vector<Topology::NicEntry>& nic_list,
                          << cudaGetErrorString(err);
             continue;
         }
-        for (char* ch = pci_bus_id; (*ch = tolower(*ch)); ch++);
+        for (char* ch = pci_bus_id; (*ch = to_lower(*ch)); ch++);
         int numa_node = getNumaNodeFromPciDevice(pci_bus_id);
         int min_distance = INT_MAX;
         std::unordered_map<int, std::vector<int>> distance_map;
diff --git a/mooncake-transfer-engine/tent/src/platform/rocm/CMakeLists.txt b/mooncake-transfer-engine/tent/src/platform/rocm/CMakeLists.txt
new file mode 100644
index 00000000..bdb2747c
--- /dev/null
+++ b/mooncake-transfer-engine/tent/src/platform/rocm/CMakeLists.txt
@@ -0,0 +1,6 @@
+if (USE_HIP)
+    find_package(HIP REQUIRED)
+    file(GLOB TENT_PLATFORM_ROCM_SOURCES "*.cpp")
+    add_library(platform_rocm STATIC ${TENT_PLATFORM_ROCM_SOURCES})
+    target_link_libraries(platform_rocm PUBLIC tent_common hip::host)
+endif()
diff --git a/mooncake-transfer-engine/tent/src/platform/rocm/rocm_allocator.cpp b/mooncake-transfer-engine/tent/src/platform/rocm/rocm_allocator.cpp
new file mode 100644
index 00000000..11aa5bc2
--- /dev/null
+++ b/mooncake-transfer-engine/tent/src/platform/rocm/rocm_allocator.cpp
@@ -0,0 +1,67 @@
+// Copyright 2025 KVCache.AI
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "tent/platform/rocm.h"
+#include "tent/common/status.h"
+
+#include <hip/hip_runtime.h>
+#include <numa.h>
+#include <glog/logging.h>
+
+namespace mooncake {
+namespace tent {
+
+Status RocmPlatform::allocate(void** pptr, size_t size,
+                              MemoryOptions& options) {
+    LocationParser location(options.location);
+    if (location.type() == "rocm") {
+        int hip_dev = 0;
+        CHECK_HIP(hipGetDevice(&hip_dev));
+        CHECK_HIP(hipSetDevice(location.index()));
+        CHECK_HIP(hipMalloc(pptr, size));
+        hipSetDevice(hip_dev);
+        return Status::OK();
+    }
+    int socket_id = 0;
+    if (location.type() == "cpu") socket_id = location.index();
+    *pptr = numa_alloc_onnode(size, socket_id);
+    if (!(*pptr))
+        return Status::InternalError("Unable to allocate DRAM memory");
+    return Status::OK();
+}
+
+Status RocmPlatform::free(void* ptr, size_t size) {
+    hipPointerAttribute_t attributes;
+    CHECK_HIP(hipPointerGetAttributes(&attributes, ptr));
+    if (attributes.type == hipMemoryTypeDevice) {
+        CHECK_HIP(hipFree(ptr));
+    } else if (attributes.type == hipMemoryTypeHost ||
+               attributes.type == hipMemoryTypeUnregistered) {
+        numa_free(ptr, size);
+    } else {
+        LOG(ERROR) << "Unknown memory type, " << ptr << " " << attributes.type;
+    }
+    return Status::OK();
+}
+
+Status RocmPlatform::copy(void* dst, void* src, size_t length) {
+    HIPStreamHandle stream;
+    CHECK_STATUS(getStreamFromPool(stream));
+    CHECK_HIP(hipMemcpyAsync(dst, src, length, hipMemcpyDefault, stream.get()));
+    CHECK_HIP(hipStreamSynchronize(stream.get()));
+    return Status::OK();
+}
+
+}  // namespace tent
+}  // namespace mooncake
diff --git a/mooncake-transfer-engine/tent/src/platform/rocm/rocm_probe.cpp b/mooncake-transfer-engine/tent/src/platform/rocm/rocm_probe.cpp
new file mode 100644
index 00000000..f1a2918b
--- /dev/null
+++ b/mooncake-transfer-engine/tent/src/platform/rocm/rocm_probe.cpp
@@ -0,0 +1,344 @@
+// Copyright 2025 KVCache.AI
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "tent/platform/rocm.h"
+#include "tent/common/status.h"
+#include "tent/common/utils/prefault.h"
+#include "char_util.h"
+
+#include <glog/logging.h>
+#include <fstream>
+#include <map>
+#include <set>
+#include <string>
+#include <utility>
+#include <vector>
+#include <hip/hip_runtime.h>
+#include <dirent.h>
+#include <infiniband/verbs.h>
+#include <limits.h>
+#include <stdlib.h>
+#include <string.h>
+#include <sys/types.h>
+#include <unordered_set>
+#include <algorithm>
+#include <numa.h>
+
+namespace mooncake {
+namespace tent {
+
+static std::vector<Topology::NicEntry> listInfiniBandDevices() {
+    int num_devices = 0;
+    std::vector<Topology::NicEntry> devices;
+
+    struct ibv_device** device_list = ibv_get_device_list(&num_devices);
+    if (!device_list || num_devices <= 0) {
+        LOG(WARNING) << "No RDMA devices found, check your device installation";
+        return {};
+    }
+
+    for (int i = 0; i < num_devices; ++i) {
+        std::string device_name = ibv_get_device_name(device_list[i]);
+        char path[PATH_MAX + 32];
+        char resolved_path[PATH_MAX];
+        snprintf(path, sizeof(path), "/sys/class/infiniband/%s/../..",
+                 device_name.c_str());
+        if (realpath(path, resolved_path) == NULL) {
+            PLOG(ERROR) << "realpath " << path << ":";
+            continue;
+        }
+        std::string pci_bus_id = basename(resolved_path);
+
+        int numa_node = -1;
+        snprintf(path, sizeof(path), "%s/numa_node", resolved_path);
+        std::ifstream(path) >> numa_node;
+
+        devices.push_back(
+            Topology::NicEntry{.name = std::move(device_name),
+                               .pci_bus_id = std::move(pci_bus_id),
+                               .type = Topology::NIC_RDMA,
+                               .numa_node = numa_node});
+    }
+    ibv_free_device_list(device_list);
+    return devices;
+}
+
+static void filterInfiniBandDevices(std::vector<Topology::NicEntry>& devices,
+                                    std::shared_ptr<Config> conf) {
+    auto whitelist = conf->getArray<std::string>("topology/rdma_whitelist");
+    auto blacklist = conf->getArray<std::string>("topology/rdma_blacklist");
+    std::vector<Topology::NicEntry> new_devices;
+    if (!whitelist.empty()) {
+        for (auto& entry : devices) {
+            if (std::find(whitelist.begin(), whitelist.end(), entry.name) !=
+                whitelist.end())
+                new_devices.push_back(entry);
+        }
+        devices.swap(new_devices);
+        return;
+    }
+    if (!blacklist.empty()) {
+        for (auto& entry : devices) {
+            if (std::find(blacklist.begin(), blacklist.end(), entry.name) ==
+                blacklist.end())
+                new_devices.push_back(entry);
+        }
+        devices.swap(new_devices);
+        return;
+    }
+}
+
+static void discoverCpuTopology(std::vector<Topology::NicEntry>& nic_list,
+                                std::vector<Topology::MemEntry>& mem_list) {
+    DIR* dir = opendir("/sys/devices/system/node");
+    struct dirent* entry;
+    if (dir == NULL) {
+        PLOG(WARNING) << "open /sys/devices/system/node failed";
+        return;
+    }
+    while ((entry = readdir(dir))) {
+        const char* prefix = "node";
+        if (entry->d_type != DT_DIR ||
+            strncmp(entry->d_name, prefix, strlen(prefix)) != 0) {
+            continue;
+        }
+        int numa_node = atoi(entry->d_name + strlen(prefix));
+        Topology::MemEntry mem_entry;
+        mem_entry.name = "cpu:" + std::to_string(numa_node);
+        mem_entry.numa_node = numa_node;
+        mem_entry.type = Topology::MEM_HOST;
+        int nic_id = 0;
+        for (const auto& device : nic_list) {
+            if (device.numa_node == numa_node) {
+                mem_entry.device_list[0].push_back(nic_id++);
+            } else {
+                mem_entry.device_list[2].push_back(nic_id++);
+            }
+        }
+        mem_list.push_back(std::move(mem_entry));
+    }
+    (void)closedir(dir);
+}
+
+static int getNumaNodeFromPciDevice(const std::string& pci_bdf) {
+    std::string sysfs_path = "/sys/bus/pci/devices/" + pci_bdf + "/numa_node";
+    std::ifstream numa_file(sysfs_path);
+    if (!numa_file.is_open()) return -1;
+    int numa_node = -1;
+    numa_file >> numa_node;
+    if (numa_file.fail()) return -1;
+    return numa_node;
+}
+
+static int getPciDistance(const char* bus1, const char* bus2) {
+    char buf[PATH_MAX];
+    char path1[PATH_MAX];
+    char path2[PATH_MAX];
+    snprintf(buf, sizeof(buf), "/sys/bus/pci/devices/%s", bus1);
+    if (realpath(buf, path1) == NULL) return -1;
+    snprintf(buf, sizeof(buf), "/sys/bus/pci/devices/%s", bus2);
+    if (realpath(buf, path2) == NULL) return -1;
+
+    char* ptr1 = path1;
+    char* ptr2 = path2;
+    while (*ptr1 && *ptr1 == *ptr2) {
+        ptr1++;
+        ptr2++;
+    }
+    int distance = 0;
+    for (; *ptr1; ptr1++) distance += (*ptr1 == '/');
+    for (; *ptr2; ptr2++) distance += (*ptr2 == '/');
+    return distance;
+}
+
+static void discoverRocmTopology(std::vector<Topology::NicEntry>& nic_list,
+                                 std::vector<Topology::MemEntry>& mem_list) {
+    int device_count = 0;
+    auto err = hipGetDeviceCount(&device_count);
+    if (err != hipSuccess) {
+        LOG(WARNING) << "hipGetDeviceCount: " << hipGetErrorString(err);
+        device_count = 0;
+    }
+    for (int i = 0; i < device_count; i++) {
+        hipDeviceProp_t prop;
+        err = hipGetDeviceProperties(&prop, i);
+        if (err != hipSuccess) {
+            LOG(WARNING) << "hipGetDeviceProperties: "
+                         << hipGetErrorString(err);
+            continue;
+        }
+        // Format PCI bus ID as lowercase: domain:bus:device.function
+        char pci_bus_id[20];
+        snprintf(pci_bus_id, sizeof(pci_bus_id), "%04x:%02x:%02x.%x",
+                 prop.pciDomainID, prop.pciBusID, prop.pciDeviceID, 0);
+        for (char* ch = pci_bus_id; *ch; ch++) *ch = to_lower(*ch);
+
+        int numa_node = getNumaNodeFromPciDevice(pci_bus_id);
+        int min_distance = INT_MAX;
+        std::unordered_map<int, std::vector<int>> distance_map;
+        for (const auto& device : nic_list) {
+            int dist = getPciDistance(device.pci_bus_id.c_str(), pci_bus_id);
+            distance_map[dist].push_back(&device - &nic_list[0]);
+            min_distance = std::min(min_distance, dist);
+        }
+
+        Topology::MemEntry entry;
+        entry.name = "rocm:" + std::to_string(i);
+        entry.numa_node = numa_node;
+        entry.pci_bus_id = pci_bus_id;
+        entry.type = Topology::MEM_ROCM;
+        if (distance_map.count(0)) {
+            entry.device_list[0] = std::move(distance_map[0]);
+        } else if (distance_map.count(min_distance)) {
+            entry.device_list[0] = std::move(distance_map[min_distance]);
+        }
+        std::unordered_set<int> preferred_set;
+        for (const auto& dev_id : entry.device_list[0]) {
+            preferred_set.insert(dev_id);
+        }
+        int dev_id = 0;
+        for (const auto& device : nic_list) {
+            if (!preferred_set.count(dev_id)) {
+                if (numa_node >= 0 && device.numa_node == numa_node)
+                    entry.device_list[1].push_back(dev_id);
+                else
+                    entry.device_list[2].push_back(dev_id);
+            }
+            dev_id++;
+        }
+        mem_list.push_back(std::move(entry));
+    }
+}
+
+static void insertFallbackMemEntry(int nic_list_count,
+                                   std::vector<Topology::MemEntry>& mem_list) {
+    for (auto& entry : mem_list) {
+        if (entry.name == kWildcardLocation) {
+            entry.device_list[2].clear();
+            for (int i = 0; i < nic_list_count; ++i)
+                entry.device_list[2].push_back(i);
+            return;
+        }
+    }
+    Topology::MemEntry new_entry;
+    new_entry.name = kWildcardLocation;
+    new_entry.numa_node = -1;
+    new_entry.type = Topology::MEM_HOST;
+    for (int i = 0; i < nic_list_count; ++i)
+        new_entry.device_list[2].push_back(i);
+    mem_list.push_back(new_entry);
+}
+
+Status RocmPlatform::probe(std::vector<Topology::NicEntry>& nic_list,
+                           std::vector<Topology::MemEntry>& mem_list) {
+    auto new_nic_list = listInfiniBandDevices();
+    filterInfiniBandDevices(new_nic_list, conf);
+    for (auto& entry : new_nic_list) nic_list.push_back(entry);
+    insertFallbackMemEntry((int)nic_list.size(), mem_list);
+    discoverCpuTopology(nic_list, mem_list);
+    discoverRocmTopology(nic_list, mem_list);
+    return Status::OK();
+}
+
+MemoryType RocmPlatform::getMemoryType(void* addr) {
+    hipPointerAttribute_t attributes;
+    hipError_t result = hipPointerGetAttributes(&attributes, addr);
+    if (result != hipSuccess) {
+        LOG(WARNING) << "hipPointerGetAttributes: "
+                     << hipGetErrorString(result);
+        return MTYPE_UNKNOWN;
+    }
+    if (attributes.type == hipMemoryTypeDevice) return MTYPE_ROCM;
+    return MTYPE_CPU;
+}
+
+static inline uintptr_t alignPage(uintptr_t address) {
+    const static size_t kPageSize = 4096;
+    return address & ~(kPageSize - 1);
+}
+
+static inline std::string genCpuNodeName(int node) {
+    if (node >= 0) return "cpu:" + std::to_string(node);
+    return kWildcardLocation;
+}
+
+static inline std::string genRocmNodeName(int node) {
+    if (node >= 0) return "rocm:" + std::to_string(node);
+    return kWildcardLocation;
+}
+
+const std::vector<RangeLocation> RocmPlatform::getLocation(void* start,
+                                                           size_t len,
+                                                           bool skip_prefault) {
+    const static size_t kPageSize = 4096;
+    std::vector<RangeLocation> entries;
+
+    hipPointerAttribute_t attributes;
+    hipError_t result = hipPointerGetAttributes(&attributes, start);
+    if (result != hipSuccess) {
+        LOG(WARNING) << "hipPointerGetAttributes: "
+                     << hipGetErrorString(result);
+        entries.push_back({(uint64_t)start, len, kWildcardLocation});
+        return entries;
+    }
+
+    if (attributes.type == hipMemoryTypeDevice) {
+        entries.push_back(
+            {(uint64_t)start, len, genRocmNodeName(attributes.device)});
+        return entries;
+    }
+
+    uintptr_t aligned_start = alignPage((uintptr_t)start);
+    int n =
+        (uintptr_t(start) - aligned_start + len + kPageSize - 1) / kPageSize;
+    void** pages = (void**)malloc(sizeof(void*) * n);
+    int* status = (int*)malloc(sizeof(int) * n);
+
+    for (int i = 0; i < n; i++) {
+        pages[i] = (void*)((char*)aligned_start + i * kPageSize);
+    }
+
+    if (!skip_prefault) {
+        prefaultBeforeProbe(pages, n, aligned_start, "RocmPlatform");
+    }
+
+    int rc = numa_move_pages(0, n, pages, nullptr, status, 0);
+    if (rc != 0) {
+        entries.push_back({(uint64_t)start, len, kWildcardLocation});
+        ::free(pages);
+        ::free(status);
+        return entries;
+    }
+
+    int node = status[0];
+    uint64_t start_addr = (uint64_t)start;
+    uint64_t new_start_addr;
+    for (int i = 1; i < n; i++) {
+        if (status[i] != node) {
+            new_start_addr = alignPage((uint64_t)start) + i * kPageSize;
+            entries.push_back({start_addr, size_t(new_start_addr - start_addr),
+                               genCpuNodeName(node)});
+            start_addr = new_start_addr;
+            node = status[i];
+        }
+    }
+    entries.push_back(
+        {start_addr, (uint64_t)start + len - start_addr, genCpuNodeName(node)});
+    ::free(pages);
+    ::free(status);
+    return entries;
+}
+
+}  // namespace tent
+}  // namespace mooncake
diff --git a/mooncake-transfer-engine/tent/src/platform/rocm/rocm_stream_pool.cpp b/mooncake-transfer-engine/tent/src/platform/rocm/rocm_stream_pool.cpp
new file mode 100644
index 00000000..65dbbad8
--- /dev/null
+++ b/mooncake-transfer-engine/tent/src/platform/rocm/rocm_stream_pool.cpp
@@ -0,0 +1,150 @@
+// Copyright 2025 KVCache.AI
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "tent/platform/rocm.h"
+
+namespace mooncake {
+namespace tent {
+
+HIPStreamHandle& HIPStreamHandle::operator=(HIPStreamHandle&& other) noexcept {
+    if (this != &other) {
+        releaseToPool();
+        stream_ = other.stream_;
+        deviceId_ = other.deviceId_;
+        pool_ = other.pool_;
+
+        other.stream_ = nullptr;
+        other.pool_ = nullptr;
+    }
+    return *this;
+}
+
+HIPStreamHandle::~HIPStreamHandle() { releaseToPool(); }
+
+void HIPStreamHandle::releaseToPool() {
+    if (stream_ != nullptr && pool_ != nullptr) {
+        pool_->release(deviceId_, stream_);
+        stream_ = nullptr;
+        pool_ = nullptr;
+    }
+}
+
+HIPStreamPool::DevicePool::DevicePool(int deviceId) : deviceId_(deviceId) {}
+
+HIPStreamPool::DevicePool::~DevicePool() {
+    int currentDevice;
+    if (hipGetDevice(&currentDevice) == hipSuccess) {
+        hipSetDevice(deviceId_);
+        for (hipStream_t stream : availableStreams_) {
+            hipStreamDestroy(stream);
+        }
+        hipSetDevice(currentDevice);
+    }
+}
+
+Status HIPStreamPool::DevicePool::acquire(hipStream_t& outStream) {
+    {
+        RWSpinlock::WriteGuard guard(dev_lock_);
+        if (!availableStreams_.empty()) {
+            outStream = availableStreams_.back();
+            availableStreams_.pop_back();
+            return Status::OK();
+        }
+    }
+
+    int currentDevice;
+    CHECK_HIP(hipGetDevice(&currentDevice));
+
+    if (currentDevice != deviceId_) {
+        CHECK_HIP(hipSetDevice(deviceId_));
+    }
+
+    CHECK_HIP(hipStreamCreateWithFlags(&outStream, hipStreamNonBlocking));
+
+    if (currentDevice != deviceId_) {
+        CHECK_HIP(hipSetDevice(currentDevice));
+    }
+
+    return Status::OK();
+}
+
+void HIPStreamPool::DevicePool::release(hipStream_t stream) {
+    RWSpinlock::WriteGuard guard(dev_lock_);
+    availableStreams_.push_back(stream);
+}
+
+Status HIPStreamPool::acquire(HIPStreamHandle& outHandle, int deviceId) {
+    if (deviceId == kCurrentDevice) {
+        if (hipGetDevice(&deviceId) != hipSuccess) {
+            return Status::InternalError("Failed to get current HIP device ID");
+        }
+    } else if (deviceId < 0) {
+        return Status::InternalError("Invalid HIP device ID");
+    }
+
+    DevicePool* devicePool = getDevicePool(deviceId);
+    if (!devicePool) {
+        return Status::InternalError("Failed to get HIP device pool");
+    }
+
+    hipStream_t rawStream;
+    CHECK_STATUS(devicePool->acquire(rawStream));
+
+    outHandle = HIPStreamHandle(rawStream, deviceId, this);
+    return Status::OK();
+}
+
+void HIPStreamPool::release(int deviceId, hipStream_t stream) {
+    DevicePool* devicePool = getDevicePool(deviceId);
+    if (devicePool) devicePool->release(stream);
+}
+
+HIPStreamPool::DevicePool* HIPStreamPool::getDevicePool(int deviceId) {
+    {
+        RWSpinlock::ReadGuard readGuard(pools_lock_);
+        if (static_cast<size_t>(deviceId) < devicePools_.size() &&
+            devicePools_[deviceId]) {
+            return devicePools_[deviceId].get();
+        }
+    }
+
+    RWSpinlock::WriteGuard writeGuard(pools_lock_);
+
+    if (static_cast<size_t>(deviceId) < devicePools_.size() &&
+        devicePools_[deviceId]) {
+        return devicePools_[deviceId].get();
+    }
+
+    int actualDeviceCount = 0;
+    if (hipGetDeviceCount(&actualDeviceCount) != hipSuccess ||
+        deviceId >= actualDeviceCount) {
+        LOG(ERROR) << "Invalid HIP device id " << deviceId;
+        return nullptr;
+    }
+
+    if (static_cast<size_t>(deviceId) >= devicePools_.size()) {
+        devicePools_.resize(deviceId + 1);
+    }
+
+    devicePools_[deviceId] = std::make_unique<DevicePool>(deviceId);
+    return devicePools_[deviceId].get();
+}
+
+Status RocmPlatform::getStreamFromPool(HIPStreamHandle& outHandle,
+                                       int deviceId) {
+    return stream_pool_.acquire(outHandle, deviceId);
+}
+
+}  // namespace tent
+}  // namespace mooncake
diff --git a/mooncake-transfer-engine/tent/src/platform/sunrise/CMakeLists.txt b/mooncake-transfer-engine/tent/src/platform/sunrise/CMakeLists.txt
new file mode 100644
index 00000000..078314b6
--- /dev/null
+++ b/mooncake-transfer-engine/tent/src/platform/sunrise/CMakeLists.txt
@@ -0,0 +1,10 @@
+if(USE_SUNRISE)
+  file(GLOB TENT_PLATFORM_SUNRISE_SOURCES "*.cpp")
+  add_library(platform_sunrise STATIC ${TENT_PLATFORM_SUNRISE_SOURCES})
+  target_include_directories(platform_sunrise PUBLIC ${MC_TANGRT_ROOT}/include)
+  target_link_directories(platform_sunrise PUBLIC
+                          ${MC_TANGRT_ROOT}/lib/linux-x86_64
+                          ${MC_TANGRT_ROOT}/lib)
+  target_link_libraries(platform_sunrise PUBLIC tent_common tangrt_shared
+                                                ptml_shared dl)
+endif()
diff --git a/mooncake-transfer-engine/tent/src/platform/sunrise/sunrise_allocator.cpp b/mooncake-transfer-engine/tent/src/platform/sunrise/sunrise_allocator.cpp
new file mode 100644
index 00000000..56569161
--- /dev/null
+++ b/mooncake-transfer-engine/tent/src/platform/sunrise/sunrise_allocator.cpp
@@ -0,0 +1,92 @@
+// Copyright 2026 KVCache.AI
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "tent/platform/sunrise.h"
+
+#include <numa.h>
+
+#include "tent/runtime/topology.h"
+#include <tang_runtime_api.h>
+
+namespace mooncake {
+namespace tent {
+
+Status SunrisePlatform::allocate(void** pptr, size_t size,
+                                 MemoryOptions& options) {
+    LocationParser location(options.location);
+    if (location.type() == "cuda") {
+        int saved_dev = -1;
+        const tangError_t get_dev_ret = tangGetDevice(&saved_dev);
+        if (tangSetDevice(location.index()) != tangSuccess) {
+            return Status::InternalError("Unable to switch sunrise device");
+        }
+        tangError_t ret = tangMalloc(pptr, size);
+        if (get_dev_ret == tangSuccess && saved_dev >= 0) {
+            tangSetDevice(saved_dev);
+        }
+        if (ret != tangSuccess)
+            return Status::InternalError(
+                "Unable to allocate sunrise device memory");
+        return Status::OK();
+    }
+    int socket_id = 0;
+    if (location.type() == "cpu") socket_id = location.index();
+    *pptr = numa_alloc_onnode(size, socket_id);
+    if (!(*pptr))
+        return Status::InternalError("Unable to allocate DRAM memory");
+    return Status::OK();
+}
+
+Status SunrisePlatform::free(void* ptr, size_t size) {
+    tangPointerAttributes attributes{};
+    tangError_t ret = tangPointerGetAttributes(&attributes, ptr);
+    if (ret == tangSuccess && attributes.type == tangMemoryTypeDevice) {
+        if (attributes.device >= 0) {
+            int saved_dev = 0;
+            tangGetDevice(&saved_dev);
+            tangSetDevice(attributes.device);
+            tangFree(ptr);
+            tangSetDevice(saved_dev);
+        } else {
+            tangFree(ptr);
+        }
+    } else {
+        numa_free(ptr, size);
+    }
+    return Status::OK();
+}
+
+Status SunrisePlatform::copy(void* dst, void* src, size_t length) {
+    tangPointerAttributes sa{};
+    tangPointerAttributes da{};
+    (void)tangPointerGetAttributes(&sa, src);
+    (void)tangPointerGetAttributes(&da, dst);
+    tangMemcpyKind kind = tangMemcpyHostToHost;
+    if (sa.type == tangMemoryTypeDevice && da.type == tangMemoryTypeDevice) {
+        kind = tangMemcpyDeviceToDevice;
+    } else if (sa.type == tangMemoryTypeDevice &&
+               da.type != tangMemoryTypeDevice) {
+        kind = tangMemcpyDeviceToHost;
+    } else if (sa.type != tangMemoryTypeDevice &&
+               da.type == tangMemoryTypeDevice) {
+        kind = tangMemcpyHostToDevice;
+    }
+    tangError_t ret = tangMemcpy(dst, src, length, kind);
+    if (ret != tangSuccess)
+        return Status::InternalError("Sunrise memcpy failed");
+    return Status::OK();
+}
+
+}  // namespace tent
+}  // namespace mooncake
diff --git a/mooncake-transfer-engine/tent/src/platform/sunrise/sunrise_probe.cpp b/mooncake-transfer-engine/tent/src/platform/sunrise/sunrise_probe.cpp
new file mode 100644
index 00000000..db9e4cce
--- /dev/null
+++ b/mooncake-transfer-engine/tent/src/platform/sunrise/sunrise_probe.cpp
@@ -0,0 +1,306 @@
+// Copyright 2026 KVCache.AI
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "tent/platform/sunrise.h"
+
+#include <algorithm>
+#include <dirent.h>
+#include <fstream>
+#include <infiniband/verbs.h>
+#include <limits.h>
+#include <numa.h>
+#include <stdlib.h>
+#include <string.h>
+#include <sys/stat.h>
+#include <unistd.h>
+#include <vector>
+#include "tent/common/utils/prefault.h"
+#include <glog/logging.h>
+#include <tang_runtime_api.h>
+
+namespace mooncake {
+namespace tent {
+
+static inline uintptr_t alignPage(uintptr_t address) {
+    const static size_t kPageSize = 4096;
+    return address & ~(kPageSize - 1);
+}
+
+static inline std::string genCpuNodeName(int node) {
+    if (node >= 0) return "cpu:" + std::to_string(node);
+    return kWildcardLocation;
+}
+
+static inline std::string genCudaNodeName(int node) {
+    if (node >= 0) return "cuda:" + std::to_string(node);
+    return kWildcardLocation;
+}
+
+static bool isIbDeviceAccessible(struct ibv_device* device) {
+    char device_path[PATH_MAX];
+    struct stat st;
+
+    snprintf(device_path, sizeof(device_path), "/dev/infiniband/%s",
+             device->dev_name);
+    if (stat(device_path, &st) != 0) return false;
+    if (!S_ISCHR(st.st_mode)) return false;
+    if (access(device_path, R_OK | W_OK) != 0) return false;
+    return true;
+}
+
+static bool checkIbDevicePort(struct ibv_context* context, uint8_t port_num) {
+    struct ibv_port_attr port_attr;
+    if (ibv_query_port(context, port_num, &port_attr) != 0) return false;
+    if (port_attr.gid_tbl_len == 0) return false;
+    if (port_attr.state != IBV_PORT_ACTIVE) return false;
+    return true;
+}
+
+static bool isIbDeviceAvailable(struct ibv_device* device) {
+    if (!isIbDeviceAccessible(device)) return false;
+    struct ibv_context* context = ibv_open_device(device);
+    if (!context) return false;
+    struct ibv_device_attr device_attr;
+    if (ibv_query_device(context, &device_attr) != 0) {
+        ibv_close_device(context);
+        return false;
+    }
+    bool has_active_port = false;
+    for (uint8_t port = 1; port <= device_attr.phys_port_cnt; ++port) {
+        if (checkIbDevicePort(context, port)) {
+            has_active_port = true;
+            break;
+        }
+    }
+    ibv_close_device(context);
+    return has_active_port;
+}
+
+static std::vector<Topology::NicEntry> listInfiniBandDevices() {
+    int num_devices = 0;
+    std::vector<Topology::NicEntry> devices;
+    struct ibv_device** device_list = ibv_get_device_list(&num_devices);
+    if (!device_list || num_devices <= 0) return {};
+
+    for (int i = 0; i < num_devices; ++i) {
+        std::string device_name = ibv_get_device_name(device_list[i]);
+        if (!isIbDeviceAvailable(device_list[i])) continue;
+
+        char path[PATH_MAX + 32];
+        char resolved_path[PATH_MAX];
+        snprintf(path, sizeof(path), "/sys/class/infiniband/%s/../..",
+                 device_name.c_str());
+        if (realpath(path, resolved_path) == NULL) continue;
+        std::string pci_bus_id = basename(resolved_path);
+
+        int numa_node = -1;
+        snprintf(path, sizeof(path), "%s/numa_node", resolved_path);
+        std::ifstream(path) >> numa_node;
+
+        devices.push_back(
+            Topology::NicEntry{.name = std::move(device_name),
+                               .pci_bus_id = std::move(pci_bus_id),
+                               .type = Topology::NIC_RDMA,
+                               .numa_node = numa_node});
+    }
+    ibv_free_device_list(device_list);
+    return devices;
+}
+
+static void filterInfiniBandDevices(std::vector<Topology::NicEntry>& devices,
+                                    std::shared_ptr<Config> conf) {
+    auto whitelist = conf->getArray<std::string>("topology/rdma_whitelist");
+    auto blacklist = conf->getArray<std::string>("topology/rdma_blacklist");
+    std::vector<Topology::NicEntry> new_devices;
+    if (!whitelist.empty()) {
+        for (auto& entry : devices) {
+            if (std::find(whitelist.begin(), whitelist.end(), entry.name) !=
+                whitelist.end())
+                new_devices.push_back(entry);
+        }
+        devices.swap(new_devices);
+        return;
+    }
+    if (!blacklist.empty()) {
+        for (auto& entry : devices) {
+            if (std::find(blacklist.begin(), blacklist.end(), entry.name) ==
+                blacklist.end())
+                new_devices.push_back(entry);
+        }
+        devices.swap(new_devices);
+    }
+}
+
+static void discoverCpuTopology(std::vector<Topology::NicEntry>& nic_list,
+                                std::vector<Topology::MemEntry>& mem_list) {
+    DIR* dir = opendir("/sys/devices/system/node");
+    struct dirent* entry;
+    if (dir == NULL) return;
+    while ((entry = readdir(dir))) {
+        const char* prefix = "node";
+        if (entry->d_type != DT_DIR ||
+            strncmp(entry->d_name, prefix, strlen(prefix)) != 0) {
+            continue;
+        }
+        int numa_node = atoi(entry->d_name + strlen(prefix));
+        Topology::MemEntry mem_entry;
+        mem_entry.name = "cpu:" + std::to_string(numa_node);
+        mem_entry.numa_node = numa_node;
+        mem_entry.type = Topology::MEM_HOST;
+        int nic_id = 0;
+        for (const auto& device : nic_list) {
+            if (device.numa_node == numa_node) {
+                mem_entry.device_list[0].push_back(nic_id++);
+            } else {
+                mem_entry.device_list[2].push_back(nic_id++);
+            }
+        }
+        mem_list.push_back(std::move(mem_entry));
+    }
+    (void)closedir(dir);
+}
+
+static void insertFallbackMemEntry(int nic_list_count,
+                                   std::vector<Topology::MemEntry>& mem_list) {
+    for (auto& entry : mem_list) {
+        if (entry.name == kWildcardLocation) {
+            entry.device_list[2].clear();
+            for (int i = 0; i < nic_list_count; ++i)
+                entry.device_list[2].push_back(i);
+            return;
+        }
+    }
+    Topology::MemEntry new_entry;
+    new_entry.name = kWildcardLocation;
+    new_entry.numa_node = -1;
+    new_entry.type = Topology::MEM_HOST;
+    for (int i = 0; i < nic_list_count; ++i)
+        new_entry.device_list[2].push_back(i);
+    mem_list.push_back(new_entry);
+}
+
+Status SunrisePlatform::probe(std::vector<Topology::NicEntry>& nic_list,
+                              std::vector<Topology::MemEntry>& mem_list) {
+    auto detected_nic = listInfiniBandDevices();
+    filterInfiniBandDevices(detected_nic, conf);
+    for (auto& entry : detected_nic) nic_list.push_back(entry);
+    insertFallbackMemEntry((int)nic_list.size(), mem_list);
+    discoverCpuTopology(nic_list, mem_list);
+
+    int device_count = 0;
+    if (tangGetDeviceCount(&device_count) != tangSuccess || device_count <= 0) {
+        return Status::OK();
+    }
+
+    for (int i = 0; i < device_count; ++i) {
+        Topology::MemEntry entry;
+        entry.name = "cuda:" + std::to_string(i);
+        entry.numa_node = -1;
+        entry.type = Topology::MEM_CUDA;
+        for (int nic = 0; nic < (int)nic_list.size(); ++nic) {
+            entry.device_list[2].push_back(nic);
+        }
+        mem_list.push_back(std::move(entry));
+    }
+    return Status::OK();
+}
+
+MemoryType SunrisePlatform::getMemoryType(void* addr) {
+    tangPointerAttributes attributes{};
+    tangError_t ret = tangPointerGetAttributes(&attributes, addr);
+    if (ret != tangSuccess) {
+        int saved_dev = 0;
+        if (tangGetDevice(&saved_dev) != tangSuccess) saved_dev = -1;
+        int device_count = 0;
+        if (tangGetDeviceCount(&device_count) == tangSuccess &&
+            device_count > 0) {
+            for (int d = 0; d < device_count; ++d) {
+                if (tangSetDevice(d) != tangSuccess) continue;
+                ret = tangPointerGetAttributes(&attributes, addr);
+                if (ret == tangSuccess) break;
+            }
+            if (saved_dev >= 0) tangSetDevice(saved_dev);
+        }
+    }
+    if (ret != tangSuccess) return MTYPE_CPU;
+    if (attributes.type == tangMemoryTypeDevice) return MTYPE_CUDA;
+    return MTYPE_CPU;
+}
+
+const std::vector<RangeLocation> SunrisePlatform::getLocation(
+    void* start, size_t len, bool skip_prefault) {
+    std::vector<RangeLocation> entries;
+
+    tangPointerAttributes attributes{};
+    tangError_t ret = tangPointerGetAttributes(&attributes, start);
+    if (ret != tangSuccess) {
+        int saved_dev = 0;
+        if (tangGetDevice(&saved_dev) != tangSuccess) saved_dev = -1;
+        int device_count = 0;
+        if (tangGetDeviceCount(&device_count) == tangSuccess &&
+            device_count > 0) {
+            for (int d = 0; d < device_count; ++d) {
+                if (tangSetDevice(d) != tangSuccess) continue;
+                ret = tangPointerGetAttributes(&attributes, start);
+                if (ret == tangSuccess) break;
+            }
+            if (saved_dev >= 0) tangSetDevice(saved_dev);
+        }
+    }
+
+    if (ret == tangSuccess && attributes.type == tangMemoryTypeDevice) {
+        entries.push_back(
+            {(uint64_t)start, len, genCudaNodeName(attributes.device)});
+        return entries;
+    }
+
+    const static size_t kPageSize = 4096;
+    uintptr_t aligned_start = alignPage((uintptr_t)start);
+    int n =
+        (uintptr_t(start) - aligned_start + len + kPageSize - 1) / kPageSize;
+    std::vector<void*> pages(n);
+    std::vector<int> status(n);
+    for (int i = 0; i < n; i++) {
+        pages[i] = reinterpret_cast<void*>(aligned_start + i * kPageSize);
+    }
+
+    if (!skip_prefault) {
+        prefaultBeforeProbe(pages.data(), n, aligned_start, "SunrisePlatform");
+    }
+
+    int rc = numa_move_pages(0, n, pages.data(), nullptr, status.data(), 0);
+    if (rc != 0) {
+        entries.push_back({(uint64_t)start, len, kWildcardLocation});
+        return entries;
+    }
+
+    int node = status[0];
+    uint64_t start_addr = (uint64_t)start;
+    uint64_t new_start_addr;
+    for (int i = 1; i < n; i++) {
+        if (status[i] != node) {
+            new_start_addr = alignPage((uint64_t)start) + i * kPageSize;
+            entries.push_back({start_addr, size_t(new_start_addr - start_addr),
+                               genCpuNodeName(node)});
+            start_addr = new_start_addr;
+            node = status[i];
+        }
+    }
+    entries.push_back(
+        {start_addr, (uint64_t)start + len - start_addr, genCpuNodeName(node)});
+    return entries;
+}
+
+}  // namespace tent
+}  // namespace mooncake
diff --git a/mooncake-transfer-engine/tent/src/python/CMakeLists.txt b/mooncake-transfer-engine/tent/src/python/CMakeLists.txt
index 5ef9b101..5d8ea378 100644
--- a/mooncake-transfer-engine/tent/src/python/CMakeLists.txt
+++ b/mooncake-transfer-engine/tent/src/python/CMakeLists.txt
@@ -23,8 +23,7 @@ set_target_properties(tentpy PROPERTIES OUTPUT_NAME "tent" INSTALL_RPATH
 
 find_package(
   Python3
-  COMPONENTS Interpreter Development
+  COMPONENTS Interpreter Development.Module
   REQUIRED)
-target_link_libraries(tentpy PRIVATE ${Python3_LIBRARIES})
 target_include_directories(tentpy PRIVATE ${Python3_INCLUDE_DIRS})
 target_link_libraries(tentpy PUBLIC tent_link_group)
diff --git a/mooncake-transfer-engine/tent/src/python/pybind.cpp b/mooncake-transfer-engine/tent/src/python/pybind.cpp
index af219231..bc1d3a23 100644
--- a/mooncake-transfer-engine/tent/src/python/pybind.cpp
+++ b/mooncake-transfer-engine/tent/src/python/pybind.cpp
@@ -259,6 +259,11 @@ PYBIND11_MODULE(tent, m) {
     m.attr("LOCAL_SEGMENT_ID") = py::int_(LOCAL_SEGMENT_ID);
     m.attr("kWildcardLocation") = py::str(kWildcardLocation);
 
+    // Priority constants
+    m.attr("PRIO_HIGH") = py::int_(PRIO_HIGH);
+    m.attr("PRIO_MEDIUM") = py::int_(PRIO_MEDIUM);
+    m.attr("PRIO_LOW") = py::int_(PRIO_LOW);
+
     // -------------------------------------------------------------------------
     // Enums
     // -------------------------------------------------------------------------
@@ -284,6 +289,7 @@ PYBIND11_MODULE(tent, m) {
         .export_values();
 
     py::enum_<TransportType>(m, "TransportType")
+        .value("UNSPEC", TransportType::UNSPEC)
         .value("RDMA", TransportType::RDMA)
         .value("MNNVL", TransportType::MNNVL)
         .value("SHM", TransportType::SHM)
@@ -292,7 +298,7 @@ PYBIND11_MODULE(tent, m) {
         .value("IOURING", TransportType::IOURING)
         .value("TCP", TransportType::TCP)
         .value("AscendDirect", TransportType::AscendDirect)
-        .value("UNSPEC", TransportType::UNSPEC)
+        .value("SUNRISE_LINK", TransportType::SUNRISE_LINK)
         .export_values();
 
     py::enum_<SegmentInfo::Type>(m, "SegmentInfoType")
@@ -314,17 +320,22 @@ PYBIND11_MODULE(tent, m) {
         .def(py::init<>())
         .def(py::init([](Request::OpCode opcode, uint64_t source,
                          uint64_t target_id, uint64_t target_offset,
-                         size_t length) {
+                         size_t length, int priority,
+                         TransportType transport_hint) {
                  Request r;
                  r.opcode = opcode;
                  r.source = U64ToPtr(source);
                  r.target_id = target_id;
                  r.target_offset = target_offset;
                  r.length = length;
+                 r.priority = priority;
+                 r.transport_hint = transport_hint;
                  return r;
              }),
              py::arg("opcode"), py::arg("source"), py::arg("target_id"),
-             py::arg("target_offset"), py::arg("length"))
+             py::arg("target_offset"), py::arg("length"),
+             py::arg("priority") = PRIO_HIGH,
+             py::arg("transport_hint") = TransportType::UNSPEC)
         .def_property(
             "opcode", [](const Request& r) { return r.opcode; },
             [](Request& r, Request::OpCode op) { r.opcode = op; })
@@ -333,7 +344,9 @@ PYBIND11_MODULE(tent, m) {
             [](Request& r, uint64_t addr) { r.source = U64ToPtr(addr); })
         .def_readwrite("target_id", &Request::target_id)
         .def_readwrite("target_offset", &Request::target_offset)
-        .def_readwrite("length", &Request::length);
+        .def_readwrite("length", &Request::length)
+        .def_readwrite("priority", &Request::priority)
+        .def_readwrite("transport_hint", &Request::transport_hint);
 
     py::class_<TransferStatus>(m, "TransferStatus")
         .def(py::init<>())
diff --git a/mooncake-transfer-engine/tent/src/rpc/rpc.cpp b/mooncake-transfer-engine/tent/src/rpc/rpc.cpp
index 40fe2077..9091543c 100644
--- a/mooncake-transfer-engine/tent/src/rpc/rpc.cpp
+++ b/mooncake-transfer-engine/tent/src/rpc/rpc.cpp
@@ -104,7 +104,8 @@ Status CoroRpcAgent::start(uint16_t& port, bool ipv6) {
         try {
             if (port == 0)
                 port = kStartPort + SimpleRandom::Get().next(kPortRange);
-            server_ = new coro_rpc::coro_rpc_server(kRpcThreads, port);
+            server_ = new coro_rpc::coro_rpc_server(kRpcThreads, port,
+                                                    ipv6 ? "::" : "0.0.0.0");
             server_->register_handler<&CoroRpcAgent::process>(this);
             server_->async_start();
             const auto err = server_->get_errc();
diff --git a/mooncake-transfer-engine/tent/src/runtime/admission_queue.cpp b/mooncake-transfer-engine/tent/src/runtime/admission_queue.cpp
new file mode 100644
index 00000000..3d30747a
--- /dev/null
+++ b/mooncake-transfer-engine/tent/src/runtime/admission_queue.cpp
@@ -0,0 +1,333 @@
+// Copyright 2026 KVCache.AI
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "tent/runtime/admission_queue.h"
+
+#include <limits>
+#include <set>
+
+namespace mooncake {
+namespace tent {
+namespace {
+
+using PublicTaskKey = std::pair<uint64_t, size_t>;
+
+bool isSupportedTerminalStatus(TransferStatusEnum status) {
+    return status == TransferStatusEnum::COMPLETED ||
+           status == TransferStatusEnum::INVALID ||
+           status == TransferStatusEnum::CANCELED ||
+           status == TransferStatusEnum::TIMEOUT ||
+           status == TransferStatusEnum::FAILED;
+}
+
+bool isSupportedOwnerKind(QueueOwnerKind kind) {
+    switch (kind) {
+        case QueueOwnerKind::User:
+        case QueueOwnerKind::StagingInternal:
+            return true;
+    }
+    return false;
+}
+
+Status checkedAdd(size_t lhs, size_t rhs, size_t& out) {
+    if (rhs > std::numeric_limits<size_t>::max() - lhs) {
+        return Status::InvalidArgument(
+            "admission queue charge overflow" LOC_MARK);
+    }
+    out = lhs + rhs;
+    return Status::OK();
+}
+
+Status validateLimits(const QueueLimits& limits) {
+    if (limits.staging_owner_reserve > limits.max_outstanding_owners) {
+        return Status::InvalidArgument(
+            "staging owner reserve exceeds owner limit" LOC_MARK);
+    }
+    if (limits.staging_byte_reserve > limits.max_outstanding_bytes) {
+        return Status::InvalidArgument(
+            "staging byte reserve exceeds byte limit" LOC_MARK);
+    }
+    return Status::OK();
+}
+
+}  // namespace
+
+LocalTransferAdmissionQueue::LocalTransferAdmissionQueue(QueueLimits limits)
+    : limits_(limits), limits_status_(validateLimits(limits)) {}
+
+Status LocalTransferAdmissionQueue::tryAdmit(
+    const QueueSubmit& submit, std::vector<QueueOwnerId>& admitted_owner_ids) {
+    admitted_owner_ids.clear();
+    CHECK_STATUS(limits_status_);
+    if (submit.batch_token == 0) {
+        return Status::InvalidArgument("invalid batch token" LOC_MARK);
+    }
+    if (submit.owners.empty()) return Status::OK();
+
+    std::set<PublicTaskKey> public_keys;
+    size_t byte_charge = 0;
+    size_t user_owner_charge = 0;
+    size_t user_byte_charge = 0;
+
+    for (const auto& owner : submit.owners) {
+        if (!isSupportedOwnerKind(owner.kind)) {
+            return Status::InvalidArgument(
+                "unsupported queue owner kind" LOC_MARK);
+        }
+        if (owner.request.length == 0) {
+            return Status::InvalidArgument("empty transfer request" LOC_MARK);
+        }
+
+        const PublicTaskKey owner_key{submit.batch_token, owner.owner_task_id};
+        if (!public_keys.insert(owner_key).second) {
+            return Status::InvalidArgument("duplicate public task id" LOC_MARK);
+        }
+        for (const auto derived_task_id : owner.derived_task_ids) {
+            if (derived_task_id == owner.owner_task_id) {
+                return Status::InvalidArgument(
+                    "owner task id appears in derived task ids" LOC_MARK);
+            }
+            const PublicTaskKey derived_key{submit.batch_token,
+                                            derived_task_id};
+            if (!public_keys.insert(derived_key).second) {
+                return Status::InvalidArgument(
+                    "duplicate public task id" LOC_MARK);
+            }
+        }
+
+        CHECK_STATUS(
+            checkedAdd(byte_charge, owner.request.length, byte_charge));
+        if (owner.kind == QueueOwnerKind::User) {
+            CHECK_STATUS(checkedAdd(user_owner_charge, 1, user_owner_charge));
+            CHECK_STATUS(checkedAdd(user_byte_charge, owner.request.length,
+                                    user_byte_charge));
+        }
+    }
+
+    if (public_keys.size() > submit.batch_slots_left) {
+        return Status::TooManyRequests(
+            "batch public task capacity exceeded" LOC_MARK);
+    }
+
+    for (const auto& key : public_keys) {
+        if (public_to_owner_.count(key)) {
+            return Status::InvalidEntry(
+                "public task id already admitted" LOC_MARK);
+        }
+    }
+
+    const size_t owner_charge = submit.owners.size();
+    size_t next_outstanding_owners = 0;
+    size_t next_outstanding_bytes = 0;
+    size_t next_user_owners = 0;
+    size_t next_user_bytes = 0;
+    CHECK_STATUS(
+        checkedAdd(outstanding_owners_, owner_charge, next_outstanding_owners));
+    CHECK_STATUS(
+        checkedAdd(outstanding_bytes_, byte_charge, next_outstanding_bytes));
+    CHECK_STATUS(checkedAdd(outstanding_user_owners_, user_owner_charge,
+                            next_user_owners));
+    CHECK_STATUS(
+        checkedAdd(outstanding_user_bytes_, user_byte_charge, next_user_bytes));
+
+    const size_t user_owner_limit =
+        limits_.max_outstanding_owners - limits_.staging_owner_reserve;
+    const size_t user_byte_limit =
+        limits_.max_outstanding_bytes - limits_.staging_byte_reserve;
+
+    if (next_outstanding_owners > limits_.max_outstanding_owners) {
+        return Status::TooManyRequests(
+            "queue owner capacity exceeded" LOC_MARK);
+    }
+    if (next_outstanding_bytes > limits_.max_outstanding_bytes) {
+        return Status::TooManyRequests("queue byte capacity exceeded" LOC_MARK);
+    }
+    if (next_user_owners > user_owner_limit) {
+        return Status::TooManyRequests("user owner capacity exceeded" LOC_MARK);
+    }
+    if (next_user_bytes > user_byte_limit) {
+        return Status::TooManyRequests("user byte capacity exceeded" LOC_MARK);
+    }
+
+    admitted_owner_ids.reserve(submit.owners.size());
+    for (const auto& owner_input : submit.owners) {
+        const QueueOwnerId owner_id = next_owner_id_++;
+        QueueOwner owner;
+        owner.batch_token = submit.batch_token;
+        owner.request = owner_input.request;
+        owner.kind = owner_input.kind;
+        owners_.emplace(owner_id, owner);
+
+        public_to_owner_[{submit.batch_token, owner_input.owner_task_id}] =
+            owner_id;
+        for (const auto derived_task_id : owner_input.derived_task_ids) {
+            public_to_owner_[{submit.batch_token, derived_task_id}] = owner_id;
+        }
+        fifo_.push_back(owner_id);
+        admitted_owner_ids.push_back(owner_id);
+    }
+
+    outstanding_owners_ = next_outstanding_owners;
+    outstanding_bytes_ = next_outstanding_bytes;
+    outstanding_user_owners_ = next_user_owners;
+    outstanding_user_bytes_ = next_user_bytes;
+    return Status::OK();
+}
+
+std::vector<QueueOwnerId> LocalTransferAdmissionQueue::pickForDispatch(
+    size_t max_owners, size_t max_bytes) {
+    std::vector<QueueOwnerId> picked;
+    if (max_owners == 0 || max_bytes == 0) return picked;
+
+    size_t used_owners = 0;
+    size_t used_bytes = 0;
+    while (!fifo_.empty() && used_owners < max_owners) {
+        auto owner_id = fifo_.front();
+        auto owner_it = owners_.find(owner_id);
+        // Non-queued entries should not normally remain in fifo_, but stale
+        // entries are skipped defensively so retireBatch() does not need to
+        // scan the dispatch queue.
+        if (owner_it == owners_.end() ||
+            owner_it->second.state != QueueState::Queued) {
+            fifo_.pop_front();
+            continue;
+        }
+
+        const auto& owner = owner_it->second;
+        const size_t remaining_bytes = max_bytes - used_bytes;
+        if (owner.request.length > remaining_bytes) break;
+
+        fifo_.pop_front();
+        owner_it->second.state = QueueState::Dispatching;
+        picked.push_back(owner_id);
+        ++used_owners;
+        used_bytes += owner.request.length;
+    }
+    return picked;
+}
+
+Status LocalTransferAdmissionQueue::complete(
+    QueueOwnerId owner_id, TransferStatusEnum terminal_status) {
+    if (owner_id == 0) {
+        return Status::InvalidArgument("invalid queue owner id" LOC_MARK);
+    }
+    if (!isSupportedTerminalStatus(terminal_status)) {
+        return Status::InvalidArgument("unsupported terminal status" LOC_MARK);
+    }
+
+    auto owner_it = owners_.find(owner_id);
+    if (owner_it == owners_.end()) {
+        return Status::InvalidEntry("queue owner not found" LOC_MARK);
+    }
+    auto& owner = owner_it->second;
+    if (owner.state != QueueState::Dispatching) {
+        return Status::InvalidEntry("queue owner is not dispatching" LOC_MARK);
+    }
+
+    owner.state = QueueState::Terminal;
+    owner.terminal_status = terminal_status;
+    --outstanding_owners_;
+    outstanding_bytes_ -= owner.request.length;
+    if (owner.kind == QueueOwnerKind::User) {
+        --outstanding_user_owners_;
+        outstanding_user_bytes_ -= owner.request.length;
+    }
+    return Status::OK();
+}
+
+Status LocalTransferAdmissionQueue::retireBatch(uint64_t batch_token) {
+    if (batch_token == 0) {
+        return Status::InvalidArgument("invalid batch token" LOC_MARK);
+    }
+
+    const PublicTaskKey batch_begin{batch_token, 0};
+    auto public_begin = public_to_owner_.lower_bound(batch_begin);
+    auto public_end = public_to_owner_.upper_bound(
+        {batch_token, std::numeric_limits<size_t>::max()});
+
+    std::set<QueueOwnerId> owner_ids;
+    for (auto it = public_begin; it != public_end; ++it) {
+        owner_ids.insert(it->second);
+    }
+
+    for (const auto owner_id : owner_ids) {
+        auto owner_it = owners_.find(owner_id);
+        if (owner_it == owners_.end()) {
+            return Status::InternalError(
+                "queue owner mapping is stale" LOC_MARK);
+        }
+
+        const auto& owner = owner_it->second;
+        if (owner.batch_token != batch_token) {
+            return Status::InternalError(
+                "queue owner batch token mismatch" LOC_MARK);
+        }
+        if (owner.state != QueueState::Terminal) {
+            return Status::InvalidEntry(
+                "batch has non-terminal queue owners" LOC_MARK);
+        }
+    }
+
+    for (const auto owner_id : owner_ids) {
+        owners_.erase(owner_id);
+    }
+    public_to_owner_.erase(public_begin, public_end);
+    return Status::OK();
+}
+
+Status LocalTransferAdmissionQueue::resolveOwner(uint64_t batch_token,
+                                                 size_t public_task_id,
+                                                 QueueOwnerId& owner_id) const {
+    if (batch_token == 0) {
+        return Status::InvalidArgument("invalid batch token" LOC_MARK);
+    }
+    auto it = public_to_owner_.find({batch_token, public_task_id});
+    if (it == public_to_owner_.end()) {
+        return Status::InvalidEntry("public task id not found" LOC_MARK);
+    }
+    owner_id = it->second;
+    return Status::OK();
+}
+
+Status LocalTransferAdmissionQueue::getPublicStatus(
+    uint64_t batch_token, size_t public_task_id,
+    TransferStatusEnum& status) const {
+    QueueOwnerId owner_id = 0;
+    CHECK_STATUS(resolveOwner(batch_token, public_task_id, owner_id));
+    auto owner_it = owners_.find(owner_id);
+    if (owner_it == owners_.end()) {
+        return Status::InternalError("queue owner mapping is stale" LOC_MARK);
+    }
+    switch (owner_it->second.state) {
+        case QueueState::Queued:
+        case QueueState::Dispatching:
+            status = TransferStatusEnum::PENDING;
+            break;
+        case QueueState::Terminal:
+            status = owner_it->second.terminal_status;
+            break;
+    }
+    return Status::OK();
+}
+
+size_t LocalTransferAdmissionQueue::outstandingOwners() const {
+    return outstanding_owners_;
+}
+
+size_t LocalTransferAdmissionQueue::outstandingBytes() const {
+    return outstanding_bytes_;
+}
+
+}  // namespace tent
+}  // namespace mooncake
diff --git a/mooncake-transfer-engine/tent/src/runtime/control_plane.cpp b/mooncake-transfer-engine/tent/src/runtime/control_plane.cpp
index 384828e4..d82fbb54 100644
--- a/mooncake-transfer-engine/tent/src/runtime/control_plane.cpp
+++ b/mooncake-transfer-engine/tent/src/runtime/control_plane.cpp
@@ -156,26 +156,12 @@ Status ControlClient::unpinStageBuffer(const std::string& server_addr,
 ControlService::ControlService(const std::string& type,
                                const std::string& servers,
                                TransferEngineImpl* impl)
-    : ControlService(type, servers, "", "", 0, impl) {}
-
-ControlService::ControlService(const std::string& type,
-                               const std::string& servers,
-                               const std::string& password, uint8_t db_index,
-                               TransferEngineImpl* impl)
-    : ControlService(type, servers, "", password, db_index, impl) {}
-
-ControlService::ControlService(const std::string& type,
-                               const std::string& servers,
-                               const std::string& username,
-                               const std::string& password, uint8_t db_index,
-                               TransferEngineImpl* impl)
     : bootstrap_callback_(nullptr), notify_callback_(nullptr), impl_(impl) {
     if (type == "p2p") {
         auto agent = std::make_unique<PeerSegmentRegistry>();
         manager_ = std::make_unique<SegmentManager>(std::move(agent));
     } else {
-        auto agent = std::make_unique<CentralSegmentRegistry>(
-            type, servers, username, password, db_index);
+        auto agent = std::make_unique<CentralSegmentRegistry>(type, servers);
         manager_ = std::make_unique<SegmentManager>(std::move(agent));
     }
     rpc_server_ = std::make_shared<CoroRpcAgent>();
@@ -240,8 +226,9 @@ Status ControlService::start(uint16_t& port, bool ipv6_) {
 
 void ControlService::onGetSegmentDesc(const std::string_view& request,
                                       std::string& response) {
-    json j = *manager_->getLocal();
-    response = j.dump();
+    // Re-use the cached dump shared across concurrent peer fetches.
+    auto cached = manager_->getLocalDumpedJson();
+    response = *cached;
 }
 
 void ControlService::onBootstrapRdma(const std::string_view& request,
@@ -257,6 +244,10 @@ void ControlService::onBootstrapRdma(const std::string_view& request,
 
 void ControlService::onSendData(const std::string_view& request,
                                 std::string& response) {
+    if (request.size() < sizeof(XferDataDesc)) {
+        response = "SendData failed: request too short";
+        return;
+    }
     XferDataDesc* desc = (XferDataDesc*)request.data();
     auto local_desc = manager_->getLocal().get();
     auto peer_mem_addr = le64toh(desc->peer_mem_addr);
@@ -277,6 +268,10 @@ void ControlService::onSendData(const std::string_view& request,
 
 void ControlService::onRecvData(const std::string_view& request,
                                 std::string& response) {
+    if (request.size() < sizeof(XferDataDesc)) {
+        response = "RecvData failed: request too short";
+        return;
+    }
     XferDataDesc* desc = (XferDataDesc*)request.data();
     auto local_desc = manager_->getLocal().get();
     auto peer_mem_addr = le64toh(desc->peer_mem_addr);
diff --git a/mooncake-transfer-engine/tent/src/runtime/memory_prober.cpp b/mooncake-transfer-engine/tent/src/runtime/memory_prober.cpp
index 7effdacb..c4da931c 100644
--- a/mooncake-transfer-engine/tent/src/runtime/memory_prober.cpp
+++ b/mooncake-transfer-engine/tent/src/runtime/memory_prober.cpp
@@ -15,6 +15,7 @@
 #include "tent/runtime/memory_prober.h"
 #include "tent/device_plugin.h"
 #include "tent/common/utils/prefault.h"
+#include "char_util.h"
 
 #include <filesystem>
 #include <dlfcn.h>
@@ -374,7 +375,7 @@ void MemoryProber::probeDeviceMemory(
                                                      sizeof(pci_bus_id));
         if (err) continue;
 
-        for (char* ch = pci_bus_id; (*ch = tolower(*ch)); ch++);
+        for (char* ch = pci_bus_id; (*ch = to_lower(*ch)); ch++);
         int numa_node = getNumaNodeFromPciDevice(pci_bus_id);
         int min_distance = INT_MAX;
         std::unordered_map<int, std::vector<int>> distance_map;
diff --git a/mooncake-transfer-engine/tent/src/runtime/metastore.cpp b/mooncake-transfer-engine/tent/src/runtime/metastore.cpp
index f797b4d8..75b92d62 100644
--- a/mooncake-transfer-engine/tent/src/runtime/metastore.cpp
+++ b/mooncake-transfer-engine/tent/src/runtime/metastore.cpp
@@ -27,23 +27,9 @@
 
 namespace mooncake {
 namespace tent {
-std::shared_ptr<MetaStore> MetaStore::Create(const std::string &type,
-                                             const std::string &servers) {
-    return Create(type, servers, "", "", 0);
-}
 
-std::shared_ptr<MetaStore> MetaStore::Create(const std::string &type,
-                                             const std::string &servers,
-                                             const std::string &password,
-                                             uint8_t db_index) {
-    return Create(type, servers, "", password, db_index);
-}
-
-std::shared_ptr<MetaStore> MetaStore::Create(const std::string &type,
-                                             const std::string &servers,
-                                             const std::string &username,
-                                             const std::string &password,
-                                             uint8_t db_index) {
+std::shared_ptr<MetaStore> MetaStore::Create(const std::string& type,
+                                             const std::string& servers) {
     std::shared_ptr<MetaStore> plugin;
 #ifdef USE_ETCD
     if (type == "etcd") {
@@ -52,6 +38,43 @@ std::shared_ptr<MetaStore> MetaStore::Create(const std::string &type,
 #endif  // USE_ETCD
 #ifdef USE_REDIS
     if (type == "redis") {
+        // Get Redis password from environment variable for security
+        std::string password;
+        const char* env_password = std::getenv("MC_REDIS_PASSWORD");
+        if (env_password && *env_password) {
+            password = env_password;
+        }
+
+        std::string username;
+        const char* env_username = std::getenv("MC_REDIS_USERNAME");
+        if (env_username && *env_username) {
+            username = env_username;
+        }
+
+        // Get Redis DB index from environment variable
+        int redis_db_index = REDIS_DEFAULT_DB_INDEX;
+        const char* env_db_index = std::getenv("MC_REDIS_DB_INDEX");
+        if (env_db_index && *env_db_index) {
+            try {
+                redis_db_index = std::stoi(env_db_index);
+            } catch (const std::exception& e) {
+                LOG(WARNING)
+                    << "Invalid MC_REDIS_DB_INDEX environment variable: "
+                    << env_db_index << ", using default "
+                    << static_cast<int>(REDIS_DEFAULT_DB_INDEX);
+            }
+        }
+
+        // Validate redis_db_index range (0-255)
+        uint8_t db_index = REDIS_DEFAULT_DB_INDEX;
+        if (redis_db_index >= 0 && redis_db_index <= REDIS_MAX_DB_INDEX) {
+            db_index = static_cast<uint8_t>(redis_db_index);
+        } else {
+            LOG(WARNING) << "Invalid Redis DB index: " << redis_db_index
+                         << ", using default "
+                         << static_cast<int>(REDIS_DEFAULT_DB_INDEX);
+        }
+
         auto redis_plugin = std::make_shared<RedisMetaStore>();
         auto status =
             redis_plugin->connect(servers, username, password, db_index);
diff --git a/mooncake-transfer-engine/tent/src/runtime/platform.cpp b/mooncake-transfer-engine/tent/src/runtime/platform.cpp
index 9f78b811..72a9d096 100644
--- a/mooncake-transfer-engine/tent/src/runtime/platform.cpp
+++ b/mooncake-transfer-engine/tent/src/runtime/platform.cpp
@@ -16,6 +16,10 @@
 
 #ifdef USE_CUDA
 #include "tent/platform/cuda.h"
+#elif defined(USE_HIP)
+#include "tent/platform/rocm.h"
+#elif defined(USE_SUNRISE)
+#include "tent/platform/sunrise.h"
 #elif defined(USE_ASCEND) || defined(USE_ASCEND_DIRECT)
 #include "tent/platform/ascend.h"
 #else
@@ -31,6 +35,10 @@ Platform& Platform::getLoader(std::shared_ptr<Config> conf) {
     std::call_once(flag, [&]() {
 #ifdef USE_CUDA
         g_instance = std::make_shared<CudaPlatform>(conf);
+#elif defined(USE_HIP)
+        g_instance = std::make_shared<RocmPlatform>(conf);
+#elif defined(USE_SUNRISE)
+        g_instance = std::make_shared<SunrisePlatform>(conf);
 #elif defined(USE_ASCEND) || defined(USE_ASCEND_DIRECT)
         g_instance = std::make_shared<AscendPlatform>(conf);
 #else
diff --git a/mooncake-transfer-engine/tent/src/runtime/progress_worker.cpp b/mooncake-transfer-engine/tent/src/runtime/progress_worker.cpp
new file mode 100644
index 00000000..49a6e470
--- /dev/null
+++ b/mooncake-transfer-engine/tent/src/runtime/progress_worker.cpp
@@ -0,0 +1,122 @@
+// Copyright 2024 KVCache.AI
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "tent/runtime/progress_worker.h"
+
+#include <chrono>
+
+#include "tent/common/status.h"
+#include "tent/runtime/transfer_engine_impl.h"
+
+namespace mooncake {
+namespace tent {
+
+ProgressWorker::ProgressWorker(TransferEngineImpl* impl,
+                               std::chrono::microseconds fallback_interval)
+    : impl_(impl), fallback_interval_(fallback_interval) {}
+
+ProgressWorker::~ProgressWorker() { stop(); }
+
+void ProgressWorker::start() {
+    if (running_.exchange(true, std::memory_order_acq_rel)) return;
+    thread_ = std::thread(&ProgressWorker::runner, this);
+}
+
+void ProgressWorker::stop() {
+    if (!running_.exchange(false, std::memory_order_acq_rel)) return;
+    {
+        std::lock_guard<std::mutex> lk(mu_);
+        // Drop pending work; outstanding batches will be reaped via the
+        // user thread's freeBatch path.
+        order_.clear();
+        queued_.clear();
+        queue_ready_ = false;
+    }
+    cv_.notify_all();
+    if (thread_.joinable()) thread_.join();
+}
+
+void ProgressWorker::notifyBatchMaybeReady(BatchID batch_id) {
+    if (!batch_id) return;
+    if (!running_.load(std::memory_order_acquire)) return;
+    {
+        std::lock_guard<std::mutex> lk(mu_);
+        if (!queued_.insert(batch_id).second) return;
+        order_.push_back(batch_id);
+    }
+    cv_.notify_one();
+}
+
+void ProgressWorker::notifyRuntimeQueueReady() {
+    if (!running_.load(std::memory_order_acquire)) return;
+    {
+        std::lock_guard<std::mutex> lk(mu_);
+        queue_ready_ = true;
+    }
+    cv_.notify_one();
+}
+
+void ProgressWorker::runner() {
+    while (true) {
+        BatchID batch_id = 0;
+        bool queue_ready = false;
+        bool fallback_due = false;
+        bool queue_active = impl_->hasActiveRuntimeQueue();
+        {
+            std::unique_lock<std::mutex> lk(mu_);
+            if (queue_active && fallback_interval_.count() > 0) {
+                const bool woke = cv_.wait_for(lk, fallback_interval_, [&] {
+                    return !running_.load(std::memory_order_acquire) ||
+                           queue_ready_ || !order_.empty();
+                });
+                fallback_due = !woke;
+            } else {
+                cv_.wait(lk, [&] {
+                    return !running_.load(std::memory_order_acquire) ||
+                           queue_ready_ || !order_.empty();
+                });
+            }
+            if (!running_.load(std::memory_order_acquire)) return;
+            queue_ready = queue_ready_;
+            queue_ready_ = false;
+            if (!order_.empty()) {
+                batch_id = order_.front();
+                order_.pop_front();
+                queued_.erase(batch_id);
+            }
+        }
+        if (queue_ready) {
+            (void)impl_->progressRuntimeQueue();
+            (void)impl_->lazyFreeBatch();
+        }
+        // progressBatch acquires the engine's progress_mutex_ and silently
+        // returns InvalidArgument if the batch was freed before we got here.
+        // PENDING means "kick again later"; the next notify wakes us up.
+        // Terminal states are observed here; lazyFreeBatch reclaims a batch
+        // only if freeBatch has already marked it free_requested.
+        if (batch_id) {
+            TransferStatus s;
+            (void)impl_->progressBatch(batch_id, s);
+            (void)impl_->progressRuntimeQueue();
+            (void)impl_->lazyFreeBatch();
+        }
+        if (fallback_due && !queue_ready && !batch_id && queue_active) {
+            (void)impl_->progressRuntimeQueue();
+            (void)impl_->lazyFreeBatch();
+        }
+    }
+}
+
+}  // namespace tent
+}  // namespace mooncake
diff --git a/mooncake-transfer-engine/tent/src/runtime/proxy_manager.cpp b/mooncake-transfer-engine/tent/src/runtime/proxy_manager.cpp
index 559018e7..53a79537 100644
--- a/mooncake-transfer-engine/tent/src/runtime/proxy_manager.cpp
+++ b/mooncake-transfer-engine/tent/src/runtime/proxy_manager.cpp
@@ -58,7 +58,13 @@ BatchID ProxyManager::submitCrossStage(const Request& request,
     inter_stage.target_id = request.target_id;
     inter_stage.target_offset = remote_stage_buffer;
     auto batch = impl_->allocateBatch(1);
-    impl_->submitTransfer(batch, {inter_stage});
+    auto status = impl_->submitStagingTransfer(batch, {inter_stage});
+    if (!status.ok()) {
+        LOG(WARNING) << "failed to submit cross-stage transfer: "
+                     << status.ToString();
+        if (batch) impl_->freeBatch(batch);
+        return 0;
+    }
     return batch;
 }
 
@@ -72,7 +78,13 @@ BatchID ProxyManager::submitLocalStage(const Request& request,
     local_stage.target_id = LOCAL_SEGMENT_ID;
     local_stage.target_offset = local_stage_buffer;
     auto batch = impl_->allocateBatch(1);
-    impl_->submitTransfer(batch, {local_stage});
+    auto status = impl_->submitStagingTransfer(batch, {local_stage});
+    if (!status.ok()) {
+        LOG(WARNING) << "failed to submit local-stage transfer: "
+                     << status.ToString();
+        if (batch) impl_->freeBatch(batch);
+        return 0;
+    }
     return batch;
 }
 
@@ -81,6 +93,7 @@ Status ProxyManager::waitLocalStage(const Request& request,
                                     uint64_t chunk_length, uint64_t offset) {
     auto batch =
         submitLocalStage(request, local_stage_buffer, chunk_length, offset);
+    if (!batch) return Status::TooManyRequests("submit local stage failed");
     return impl_->waitTransferCompletion(batch);
 }
 
@@ -119,13 +132,15 @@ Status ProxyManager::waitCrossStage(const Request& request,
                                     uint64_t chunk_length) {
     auto batch = submitCrossStage(request, local_stage_buffer,
                                   remote_stage_buffer, chunk_length);
+    if (!batch) return Status::TooManyRequests("submit cross stage failed");
     return impl_->waitTransferCompletion(batch);
 }
 
-Status ProxyManager::submit(TaskInfo* task,
+Status ProxyManager::submit(TaskInfo* task, BatchID batch,
                             const std::vector<std::string>& params) {
     StagingTask staging_task;
     staging_task.native = task;
+    staging_task.batch = batch;
     staging_task.params = params;
     task->staging_status = PENDING;
     static std::atomic<size_t> next_queue_index(0);
@@ -228,6 +243,7 @@ void ProxyManager::runner(size_t id) {
         auto staging_status = status.ok() ? COMPLETED : FAILED;
         __atomic_store(&task.native->staging_status, &staging_status,
                        __ATOMIC_RELEASE);
+        impl_->notifyBatchMaybeReady(task.batch);
     }
     cache.reset();
 }
@@ -238,7 +254,8 @@ Status ProxyManager::transferEventLoop(StagingTask& task,
     auto server_addr = task.params[0];
     bool local_staging = !task.params[1].empty();
     bool remote_staging = !task.params[2].empty();
-    const static size_t kStageBuffers = 4;
+    const size_t kStageBuffers =
+        std::min(chunk_count_, static_cast<size_t>(16));
     uint64_t local_stage_buffer[kStageBuffers],
         remote_stage_buffer[kStageBuffers];
     if (local_staging) {
@@ -316,8 +333,14 @@ Status ProxyManager::transferEventLoop(StagingTask& task,
                     local_locked.insert(chunk.local_buf);
                     chunk.batch = submitLocalStage(request, chunk.local_buf,
                                                    chunk.length, chunk.offset);
+                    if (!chunk.batch) {
+                        chunk.state = StageState::FAILED;
+                        event_queue.push(id);
+                        break;
+                    }
                     chunk.prev_state = chunk.state;
                     chunk.state = StageState::INFLIGHT;
+                    event_queue.push(id);
                 } else if (request.opcode == Request::READ && remote_staging) {
                     if (remote_locked.count(chunk.remote_buf)) {
                         event_queue.push(id);
@@ -329,10 +352,11 @@ Status ProxyManager::transferEventLoop(StagingTask& task,
                                       remote_futures[id]);
                     chunk.prev_state = chunk.state;
                     chunk.state = StageState::INFLIGHT_REMOTE;
+                    event_queue.push(id);
                 } else {
                     chunk.state = StageState::CROSS;
+                    event_queue.push(id);
                 }
-                event_queue.push(id);
                 break;
             }
 
@@ -353,6 +377,11 @@ Status ProxyManager::transferEventLoop(StagingTask& task,
                 }
                 chunk.batch = submitCrossStage(request, chunk.local_buf,
                                                chunk.remote_buf, chunk.length);
+                if (!chunk.batch) {
+                    chunk.state = StageState::FAILED;
+                    event_queue.push(id);
+                    break;
+                }
                 chunk.prev_state = chunk.state;
                 chunk.state = StageState::INFLIGHT;
                 event_queue.push(id);
@@ -366,20 +395,28 @@ Status ProxyManager::transferEventLoop(StagingTask& task,
                                       remote_futures[id]);
                     chunk.prev_state = chunk.state;
                     chunk.state = StageState::INFLIGHT_REMOTE;
+                    event_queue.push(id);
                 } else if (request.opcode == Request::READ && local_staging) {
                     chunk.batch = submitLocalStage(request, chunk.local_buf,
                                                    chunk.length, chunk.offset);
+                    if (!chunk.batch) {
+                        chunk.state = StageState::FAILED;
+                        event_queue.push(id);
+                        break;
+                    }
                     chunk.prev_state = chunk.state;
                     chunk.state = StageState::INFLIGHT;
                     event_queue.push(id);
+                } else {
+                    // No staging needed, mark as finished
+                    chunk.state = StageState::FINISH;
                 }
                 break;
             }
 
             case StageState::INFLIGHT: {
                 TransferStatus xfer_status;
-                CHECK_STATUS(
-                    impl_->getTransferStatus(chunk.batch, xfer_status));
+                CHECK_STATUS(impl_->progressBatch(chunk.batch, xfer_status));
                 if (xfer_status.s == PENDING) {
                     event_queue.push(id);
                     break;
@@ -416,6 +453,10 @@ Status ProxyManager::transferEventLoop(StagingTask& task,
             }
 
             case StageState::FAILED: {
+                // Drain the queue to avoid losing chunks
+                while (!event_queue.empty()) {
+                    event_queue.pop();
+                }
                 return Status::InternalError(
                     "Proxy event loop in failed state");
             }
@@ -424,6 +465,7 @@ Status ProxyManager::transferEventLoop(StagingTask& task,
                 auto& fut = remote_futures[id];
                 if (!fut.valid()) {
                     chunk.state = StageState::FAILED;
+                    event_queue.push(id);
                     break;
                 }
                 if (fut.wait_for(std::chrono::seconds(0)) ==
@@ -431,6 +473,7 @@ Status ProxyManager::transferEventLoop(StagingTask& task,
                     Status rs = fut.get();
                     if (!rs.ok()) {
                         chunk.state = StageState::FAILED;
+                        event_queue.push(id);
                         break;
                     }
                     if (chunk.prev_state == StageState::PRE) {
@@ -442,6 +485,7 @@ Status ProxyManager::transferEventLoop(StagingTask& task,
                             remote_staging) {
                             remote_locked.erase(chunk.remote_buf);
                         }
+                        event_queue.push(id);
                     }
                 } else {
                     event_queue.push(id);
@@ -576,4 +620,4 @@ Status ProxyManager::unpinStageBuffer(uint64_t addr) {
 }
 
 }  // namespace tent
-}  // namespace mooncake
\ No newline at end of file
+}  // namespace mooncake
diff --git a/mooncake-transfer-engine/tent/src/runtime/segment_manager.cpp b/mooncake-transfer-engine/tent/src/runtime/segment_manager.cpp
index 9d91dd73..98f6331a 100644
--- a/mooncake-transfer-engine/tent/src/runtime/segment_manager.cpp
+++ b/mooncake-transfer-engine/tent/src/runtime/segment_manager.cpp
@@ -64,7 +64,8 @@ Status SegmentManager::getRemoteCached(SegmentDesc *&desc, SegmentID handle) {
     auto &cache = tl_remote_cache_.get();
     auto current_ts = getCurrentTimeInNano();
     auto current_version = version_.load(std::memory_order_relaxed);
-    if (current_ts - cache.last_refresh > ttl_ms_ * 1000000 ||
+    if (current_ts - cache.last_refresh >
+            static_cast<uint64_t>(TENT_SEGMENT_DESC_TTL_MS) * 1000000 ||
         cache.version != current_version) {
         cache.id_to_desc_map.clear();
         cache.last_refresh = current_ts;
@@ -97,11 +98,16 @@ Status SegmentManager::getRemoteCached(SegmentDesc *&desc, SegmentID handle) {
 }
 
 Status SegmentManager::getRemote(SegmentDescRef &desc, SegmentID handle) {
-    RWSpinlock::WriteGuard guard(lock_);
-    if (!id_to_name_map_.count(handle)) {
-        return Status::InvalidArgument("Invalid segment handle" LOC_MARK);
+    // Only the id_to_name lookup needs the lock.
+    std::string segment_name;
+    {
+        RWSpinlock::ReadGuard guard(lock_);
+        auto it = id_to_name_map_.find(handle);
+        if (it == id_to_name_map_.end()) {
+            return Status::InvalidArgument("Invalid segment handle" LOC_MARK);
+        }
+        segment_name = it->second;
     }
-    auto segment_name = id_to_name_map_[handle];
     if (segment_name.starts_with(kLocalFileSegmentPrefix)) {
         CHECK_STATUS(makeFileRemote(desc, segment_name));
     } else {
@@ -169,7 +175,26 @@ Status SegmentManager::makeFileRemote(SegmentDescRef &desc,
     return Status::OK();
 }
 
+std::shared_ptr<const std::string> SegmentManager::getLocalDumpedJson() {
+    {
+        std::lock_guard<std::mutex> g(local_json_cache_mu_);
+        if (local_json_cache_) return local_json_cache_;
+    }
+    json j = *local_desc_;
+    auto computed = std::make_shared<const std::string>(j.dump());
+
+    std::lock_guard<std::mutex> g(local_json_cache_mu_);
+    if (!local_json_cache_) {
+        local_json_cache_ = computed;
+    }
+    return local_json_cache_;
+}
+
 Status SegmentManager::synchronizeLocal() {
+    {
+        std::lock_guard<std::mutex> g(local_json_cache_mu_);
+        local_json_cache_.reset();
+    }
     CHECK_STATUS(registry_->putSegmentDesc(local_desc_));
 
     std::vector<std::string> subscribers_snapshot;
diff --git a/mooncake-transfer-engine/tent/src/runtime/segment_registry.cpp b/mooncake-transfer-engine/tent/src/runtime/segment_registry.cpp
index f6f85971..3f932b4c 100644
--- a/mooncake-transfer-engine/tent/src/runtime/segment_registry.cpp
+++ b/mooncake-transfer-engine/tent/src/runtime/segment_registry.cpp
@@ -34,21 +34,6 @@ CentralSegmentRegistry::CentralSegmentRegistry(const std::string &type,
     plugin_ = MetaStore::Create(type, servers);
 }
 
-CentralSegmentRegistry::CentralSegmentRegistry(const std::string &type,
-                                               const std::string &servers,
-                                               const std::string &password,
-                                               uint8_t db_index) {
-    plugin_ = MetaStore::Create(type, servers, "", password, db_index);
-}
-
-CentralSegmentRegistry::CentralSegmentRegistry(const std::string &type,
-                                               const std::string &servers,
-                                               const std::string &username,
-                                               const std::string &password,
-                                               uint8_t db_index) {
-    plugin_ = MetaStore::Create(type, servers, username, password, db_index);
-}
-
 Status CentralSegmentRegistry::getSegmentDesc(SegmentDescRef &desc,
                                               const std::string &segment_name) {
     if (!plugin_)
diff --git a/mooncake-transfer-engine/tent/src/runtime/segment_tracker.cpp b/mooncake-transfer-engine/tent/src/runtime/segment_tracker.cpp
index 8cd6a518..2f84486d 100644
--- a/mooncake-transfer-engine/tent/src/runtime/segment_tracker.cpp
+++ b/mooncake-transfer-engine/tent/src/runtime/segment_tracker.cpp
@@ -66,8 +66,7 @@ Status SegmentTracker::add(uint64_t base, size_t length,
         new_desc.location = entries[0].location;
     else {
         new_desc.location = entries[0].location;
-        for (auto& entry : entries)
-            new_desc.regions.push_back(Region{entry.len, entry.location});
+        new_desc.regions = coalesceRegions(entries);
     }
     new_desc.ref_count = 1;
     auto status = callback(new_desc);
diff --git a/mooncake-transfer-engine/tent/src/runtime/transfer_engine_impl.cpp b/mooncake-transfer-engine/tent/src/runtime/transfer_engine_impl.cpp
index 453d9c44..2e3f3455 100644
--- a/mooncake-transfer-engine/tent/src/runtime/transfer_engine_impl.cpp
+++ b/mooncake-transfer-engine/tent/src/runtime/transfer_engine_impl.cpp
@@ -26,10 +26,10 @@
 
 #include "tent/common/config.h"
 #include "tent/common/status.h"
-#include "tent/metastore/redis.h"
 #include "tent/runtime/control_plane.h"
 #include "tent/runtime/segment.h"
 #include "tent/runtime/segment_tracker.h"
+#include "tent/runtime/progress_worker.h"
 #include "tent/runtime/proxy_manager.h"
 #include "tent/runtime/transport.h"
 #include "tent/runtime/topology.h"
@@ -43,6 +43,11 @@
 namespace mooncake {
 namespace tent {
 
+namespace {
+constexpr uint8_t kRedisMaxDbIndex = 255;
+constexpr uint8_t kRedisDefaultDbIndex = 0;
+}  // namespace
+
 struct Batch {
     Batch() : max_size(0) { sub_batch.fill(nullptr); }
 
@@ -51,6 +56,9 @@ struct Batch {
     std::array<Transport::SubBatchRef, kSupportedTransportTypes> sub_batch;
     std::vector<TaskInfo> task_list;
     size_t max_size;
+    size_t runtime_refs{0};
+    bool free_requested{false};
+    uint64_t queue_token{0};
 
     struct SubmitHook {
         size_t start_task_id{0};
@@ -272,38 +280,40 @@ Status TransferEngineImpl::construct() {
     auto metadata_type = conf_->get("metadata_type", "p2p");
     auto metadata_servers = conf_->get("metadata_servers", "");
 
-    // Get Redis password from environment variable for security
-    std::string redis_password;
-    const char* env_password = std::getenv("MC_REDIS_PASSWORD");
-    if (env_password && *env_password) {
-        redis_password = env_password;
-    }
-
-    std::string redis_username;
-    const char* env_username = std::getenv("MC_REDIS_USERNAME");
-    if (env_username && *env_username) {
-        redis_username = env_username;
-    }
-
-    // Get Redis DB index from environment variable or config
-    int redis_db_index_config = conf_->get("redis_db_index", 0);
-    const char* env_db_index = std::getenv("MC_REDIS_DB_INDEX");
-    if (env_db_index && *env_db_index) {
-        try {
-            redis_db_index_config = std::stoi(env_db_index);
-        } catch (const std::exception& e) {
-            LOG(WARNING) << "Invalid REDIS_DB_INDEX environment variable: "
-                         << env_db_index
-                         << ", using config value: " << redis_db_index_config;
-        }
-    }
-
     setLogLevel(conf_->get("log_level", "info"));
     hostname_ = conf_->get("rpc_server_hostname", "");
     local_segment_name_ = conf_->get("local_segment_name", "");
     CHECK_STATUS(getRpcServerPortFromConfig(*conf_, 0, port_));
     merge_requests_ = conf_->get("merge_requests", true);
     max_failover_attempts_ = conf_->get("max_failover_attempts", 3);
+    enable_auto_failover_on_poll_ =
+        conf_->get("enable_auto_failover_on_poll", true);
+    enable_progress_worker_ = conf_->get("enable_progress_worker", false);
+    runtime_queue_config_.enabled = conf_->get("enable_runtime_queue", false);
+    if (runtime_queue_config_.enabled) enable_progress_worker_ = true;
+    runtime_queue_config_.limits.max_outstanding_owners =
+        conf_->get("runtime_queue/max_outstanding_owners", 1024UL);
+    runtime_queue_config_.limits.max_outstanding_bytes =
+        conf_->get("runtime_queue/max_outstanding_bytes", 1UL << 30);
+    runtime_queue_config_.limits.staging_owner_reserve =
+        conf_->get("runtime_queue/staging_owner_reserve", 0UL);
+    runtime_queue_config_.limits.staging_byte_reserve =
+        conf_->get("runtime_queue/staging_byte_reserve", 0UL);
+    runtime_queue_config_.max_dispatch_owners =
+        conf_->get("runtime_queue/max_dispatch_owners", 64UL);
+    runtime_queue_config_.max_dispatch_bytes =
+        conf_->get("runtime_queue/max_dispatch_bytes", 64UL << 20);
+    runtime_queue_config_.progress_fallback_interval =
+        std::chrono::microseconds(
+            conf_->get("runtime_queue/progress_fallback_interval_us", 50000UL));
+    if (runtime_queue_config_.enabled &&
+        (runtime_queue_config_.max_dispatch_owners == 0 ||
+         runtime_queue_config_.max_dispatch_bytes == 0)) {
+        return Status::InvalidArgument(
+            "runtime queue dispatch window must be non-zero" LOC_MARK);
+    }
+    runtime_queue_ = std::make_unique<LocalTransferAdmissionQueue>(
+        runtime_queue_config_.limits);
     if (!hostname_.empty())
         CHECK_STATUS(checkLocalIpAddress(hostname_, ipv6_));
     else
@@ -313,20 +323,8 @@ Status TransferEngineImpl::construct() {
     auto loader = &Platform::getLoader(conf_);
     CHECK_STATUS(topology_->discover({loader}));
 
-    // Validate redis_db_index range (0-255)
-    uint8_t db_index = REDIS_DEFAULT_DB_INDEX;
-    if (redis_db_index_config >= 0 &&
-        redis_db_index_config <= REDIS_MAX_DB_INDEX) {
-        db_index = static_cast<uint8_t>(redis_db_index_config);
-    } else {
-        LOG(WARNING) << "Invalid Redis DB index: " << redis_db_index_config
-                     << ", using default "
-                     << static_cast<int>(REDIS_DEFAULT_DB_INDEX);
-    }
-
-    metadata_ = std::make_shared<ControlService>(
-        metadata_type, metadata_servers, redis_username, redis_password,
-        db_index, this);
+    metadata_ =
+        std::make_shared<ControlService>(metadata_type, metadata_servers, this);
 
     CHECK_STATUS(metadata_->start(port_, ipv6_));
 
@@ -336,6 +334,18 @@ Status TransferEngineImpl::construct() {
         local_segment_name_ = randomSegmentName();
 
     CHECK_STATUS(setupLocalSegment());
+
+    // Initialize transport selector
+    transport_selector_ = std::make_unique<TransportSelector>(conf_);
+    transport_selector_->setTopology(topology_);
+
+    // Check if legacy mode is enabled (use original getTransportType logic)
+    bool legacy_mode = conf_->get("use_legacy_transport_selection", false);
+    transport_selector_->setLegacyMode(legacy_mode);
+    if (legacy_mode) {
+        LOG(INFO) << "Using legacy transport selection (original logic)";
+    }
+
     CHECK_STATUS(loadTransports());
 
     std::string transport_string;
@@ -356,6 +366,14 @@ Status TransferEngineImpl::construct() {
 
     staging_proxy_ = std::make_unique<ProxyManager>(this);
 
+    if (enable_progress_worker_) {
+        progress_worker_ = std::make_unique<ProgressWorker>(
+            this, runtime_queue_config_.enabled
+                      ? runtime_queue_config_.progress_fallback_interval
+                      : std::chrono::microseconds(0));
+        progress_worker_->start();
+    }
+
     // Initialize and start Metrics system
     auto metrics_config = MetricsConfigLoader::loadWithDefaults(conf_.get());
     if (metrics_config.enabled) {
@@ -398,6 +416,14 @@ Status TransferEngineImpl::construct() {
 Status TransferEngineImpl::deconstruct() {
     // Metrics cleanup is handled automatically by TentMetrics destructor
 
+    // Stop the progress worker first so it cannot race with batch teardown
+    // below (it dereferences BatchID into Batch* via progressBatch). Keep the
+    // object alive until transports are destroyed: completion paths may still
+    // issue a final no-op wake while their workers are joining.
+    if (progress_worker_) {
+        progress_worker_->stop();
+    }
+
     // Destroy staging_proxy_ first: its destructor calls back into
     // unregisterLocalMemory/freeLocalMemory, which require
     // local_segment_tracker_ and metadata_ to be alive.
@@ -421,8 +447,11 @@ Status TransferEngineImpl::deconstruct() {
     // does not access transport-internal state (workers, connections).
     // Callers must ensure no transfers are in-flight before calling
     // deconstruct().
-    batch_set_.forEach([&](BatchSet& entry) {
-        for (auto& batch : entry.active) {
+    {
+        std::lock_guard<std::recursive_mutex> lk(progress_mutex_);
+        std::unordered_set<Batch*> released_batches;
+        auto release_batch = [&](Batch* batch) {
+            if (!released_batches.insert(batch).second) return;
             for (size_t type = 0; type < kSupportedTransportTypes; ++type) {
                 auto& transport = transport_list_[type];
                 auto& sub_batch = batch->sub_batch[type];
@@ -430,22 +459,17 @@ Status TransferEngineImpl::deconstruct() {
                 transport->freeSubBatch(sub_batch);
             }
             Slab<Batch>::Get().deallocate(batch);
-        }
-        for (auto& batch : entry.freelist) {
-            for (size_t type = 0; type < kSupportedTransportTypes; ++type) {
-                auto& transport = transport_list_[type];
-                auto& sub_batch = batch->sub_batch[type];
-                if (!transport || !sub_batch) continue;
-                transport->freeSubBatch(sub_batch);
-            }
-            Slab<Batch>::Get().deallocate(batch);
-        }
-        entry.active.clear();
-        entry.freelist.clear();
-    });
+        };
+        for (auto& batch : batch_set_.active) release_batch(batch);
+        for (auto& batch : batch_set_.freelist) release_batch(batch);
+        batch_set_.active.clear();
+        batch_set_.freelist.clear();
+        alive_batches_.clear();
+    }
 
     // Now safe to destroy transports (workers join here)
     for (auto& transport : transport_list_) transport.reset();
+    progress_worker_.reset();
     local_segment_tracker_.reset();
     if (metadata_) {
         metadata_->segmentManager().deleteLocal();
@@ -616,6 +640,7 @@ std::vector<TransportType> TransferEngineImpl::getSupportedTransports(
     if (transport_list_[MNNVL]) result.push_back(MNNVL);
     if (transport_list_[NVLINK]) result.push_back(NVLINK);
     if (transport_list_[RDMA]) result.push_back(RDMA);
+    if (transport_list_[SUNRISE_LINK]) result.push_back(SUNRISE_LINK);
     if (transport_list_[AscendDirect]) result.push_back(AscendDirect);
     if (transport_list_[SHM]) result.push_back(SHM);
     if (transport_list_[TCP]) result.push_back(TCP);
@@ -662,8 +687,7 @@ Status TransferEngineImpl::registerLocalMemory(std::vector<void*> addr_list,
             desc.location = entries[0].location;
         } else {
             desc.location = entries[0].location;
-            for (auto& entry : entries)
-                desc.regions.push_back(Region{entry.len, entry.location});
+            desc.regions = coalesceRegions(entries);
         }
         desc.ref_count = 1;
         if (options.location != kWildcardLocation)
@@ -734,45 +758,132 @@ BatchID TransferEngineImpl::allocateBatch(size_t batch_size) {
     Batch* batch = Slab<Batch>::Get().allocate();
     if (!batch) return (BatchID)0;
     batch->max_size = batch_size;
-    batch_set_.get().active.insert(batch);
-    return (BatchID)batch;
+    BatchID batch_id = (BatchID)batch;
+    std::lock_guard<std::recursive_mutex> lk(progress_mutex_);
+    batch_set_.active.insert(batch);
+    alive_batches_.insert(batch_id);
+    return batch_id;
 }
 
 Status TransferEngineImpl::freeBatch(BatchID batch_id) {
     if (!batch_id) return Status::InvalidArgument("Invalid batch ID" LOC_MARK);
     Batch* batch = (Batch*)(batch_id);
-    batch_set_.get().freelist.push_back(batch);
+    std::lock_guard<std::recursive_mutex> lk(progress_mutex_);
+    if (!alive_batches_.count(batch_id))
+        return Status::InvalidArgument("Batch is not alive" LOC_MARK);
+    if (runtime_queue_config_.enabled && batch->queue_token != 0) {
+        auto retire_status = retireQueueForBatch(batch);
+        if (!retire_status.ok() && !retire_status.IsInvalidEntry()) {
+            return retire_status;
+        }
+    }
+    if (batch->free_requested) {
+        CHECK_STATUS(lazyFreeBatch());
+        return Status::OK();
+    }
+    batch->free_requested = true;
+    batch_set_.freelist.push_back(batch);
     lazyFreeBatch();
     return Status::OK();
 }
 
 Status TransferEngineImpl::lazyFreeBatch() {
-    auto& batch_set = batch_set_.get();
-    for (auto it = batch_set.freelist.begin();
-         it != batch_set.freelist.end();) {
+    std::lock_guard<std::recursive_mutex> lk(progress_mutex_);
+    for (auto it = batch_set_.freelist.begin();
+         it != batch_set_.freelist.end();) {
         auto& batch = *it;
+        if (batch->runtime_refs > 0) {
+            it++;
+            continue;
+        }
         TransferStatus overall_status;
         CHECK_STATUS(getTransferStatus((BatchID)batch, overall_status));
         if (overall_status.s == PENDING) {
             it++;
             continue;
         }
+        if (runtime_queue_config_.enabled && batch->queue_token != 0) {
+            CHECK_STATUS(retireQueueForBatch(batch));
+        }
         for (size_t type = 0; type < kSupportedTransportTypes; ++type) {
             auto& transport = transport_list_[type];
             auto& sub_batch = batch->sub_batch[type];
             if (transport && sub_batch) transport->freeSubBatch(sub_batch);
         }
-        batch_set.active.erase(batch);
+        batch_set_.active.erase(batch);
+        alive_batches_.erase((BatchID)batch);
         Slab<Batch>::Get().deallocate(batch);
-        it = batch_set.freelist.erase(it);
+        it = batch_set_.freelist.erase(it);
     }
     return Status::OK();
 }
 
+Status TransferEngineImpl::retainBatch(BatchID batch_id, Batch*& batch) {
+    if (!batch_id) return Status::InvalidArgument("Invalid batch ID" LOC_MARK);
+    std::lock_guard<std::recursive_mutex> lk(progress_mutex_);
+    if (!alive_batches_.count(batch_id)) {
+        return Status::InvalidArgument("Batch is not alive" LOC_MARK);
+    }
+    batch = (Batch*)batch_id;
+    if (batch->free_requested) {
+        return Status::InvalidArgument("Batch is being freed" LOC_MARK);
+    }
+    ++batch->runtime_refs;
+    return Status::OK();
+}
+
+Status TransferEngineImpl::releaseBatch(Batch* batch) {
+    if (!batch) return Status::InvalidArgument("Invalid batch" LOC_MARK);
+    std::lock_guard<std::recursive_mutex> lk(progress_mutex_);
+    if (batch->runtime_refs == 0) {
+        return Status::InternalError("Batch runtime ref underflow" LOC_MARK);
+    }
+    --batch->runtime_refs;
+    if (batch->runtime_refs == 0 && batch->free_requested) {
+        CHECK_STATUS(lazyFreeBatch());
+    }
+    return Status::OK();
+}
+
+class TransferEngineImpl::BatchRef {
+   public:
+    BatchRef(TransferEngineImpl& engine, Batch* batch)
+        : engine_(engine), batch_(batch) {}
+
+    ~BatchRef() {
+        if (!batch_) return;
+        auto status = engine_.releaseBatch(batch_);
+        if (!status.ok()) {
+            LOG(WARNING) << "failed to release batch ref: "
+                         << status.ToString();
+        }
+    }
+
+    BatchRef(const BatchRef&) = delete;
+    BatchRef& operator=(const BatchRef&) = delete;
+
+    Batch* get() const { return batch_; }
+
+    Status release() {
+        if (!batch_) return Status::OK();
+        auto status = engine_.releaseBatch(batch_);
+        batch_ = nullptr;
+        return status;
+    }
+
+   private:
+    TransferEngineImpl& engine_;
+    Batch* batch_{nullptr};
+};
+
+static bool isGpuType(MemoryType t) {
+    return t == MTYPE_CUDA || t == MTYPE_ROCM;
+}
+
 static bool checkAvailability(const std::shared_ptr<Transport>& xport,
                               MemoryType local) {
     if (local == MTYPE_CPU) return xport && xport->capabilities().dram_to_file;
-    if (local == MTYPE_CUDA) return xport && xport->capabilities().gpu_to_file;
+    if (isGpuType(local)) return xport && xport->capabilities().gpu_to_file;
     return false;
 }
 
@@ -780,11 +891,11 @@ static bool checkAvailability(const std::shared_ptr<Transport>& xport,
                               MemoryType local, MemoryType remote) {
     if (local == MTYPE_CPU && remote == MTYPE_CPU)
         return xport && xport->capabilities().dram_to_dram;
-    if (local == MTYPE_CUDA && remote == MTYPE_CUDA)
+    if (isGpuType(local) && isGpuType(remote))
         return xport && xport->capabilities().gpu_to_gpu;
-    if (local == MTYPE_CPU && remote == MTYPE_CUDA)
+    if (local == MTYPE_CPU && isGpuType(remote))
         return xport && xport->capabilities().dram_to_gpu;
-    if (local == MTYPE_CUDA && remote == MTYPE_CPU)
+    if (isGpuType(local) && remote == MTYPE_CPU)
         return xport && xport->capabilities().gpu_to_dram;
     return false;
 }
@@ -793,52 +904,126 @@ static MemoryType getTypeEnum(const std::string& type) {
     if (type == "cpu" || type == "*") return MTYPE_CPU;
     if (type == "cuda") return MTYPE_CUDA;
     if (type == "npu") return MTYPE_CUDA;
+    if (type == "rocm") return MTYPE_ROCM;
     return MTYPE_UNKNOWN;
 }
 
-TransportType TransferEngineImpl::getTransportType(const Request& request,
-                                                   int priority) {
+Status TransferEngineImpl::validateTransportHint(const Request& req,
+                                                 size_t request_index) {
+    if (req.transport_hint == UNSPEC) return Status::OK();
+    if ((int)req.transport_hint < 0 ||
+        (int)req.transport_hint >= kSupportedTransportTypes) {
+        return Status::InvalidArgument(
+            "transport_hint out of range for request[" +
+            std::to_string(request_index) + "]" LOC_MARK);
+    }
+    if (!transport_list_[req.transport_hint]) {
+        return Status::InvalidArgument(
+            "transport_hint=" +
+            TransportSelector::transportTypeName(req.transport_hint) +
+            " is not enabled in config (request[" +
+            std::to_string(request_index) + "])" LOC_MARK);
+    }
+    return Status::OK();
+}
+
+SelectionResult TransferEngineImpl::getTransportType(const Request& request,
+                                                     int transport_index) {
     SegmentDesc* desc;
     if (request.target_id == LOCAL_SEGMENT_ID) {
         desc = metadata_->segmentManager().getLocal().get();
     } else {
         auto status = metadata_->segmentManager().getRemoteCached(
             desc, request.target_id);
-        if (!status.ok()) return UNSPEC;
+        if (!status.ok()) return SelectionResult{};
     }
     auto local_mtype = Platform::getLoader().getMemoryType(request.source);
-    if (desc->type == SegmentType::File) {
-        if (checkAvailability(transport_list_[GDS], local_mtype)) {
-            if (priority-- == 0) return GDS;
-        }
-        if (checkAvailability(transport_list_[IOURING], local_mtype)) {
-            if (priority-- == 0) return IOURING;
-        }
-        return UNSPEC;
-    } else {
-        auto entry = desc->findBuffer(request.target_offset, request.length);
-        if (!entry) return UNSPEC;
-        bool same_machine = (request.target_id == LOCAL_SEGMENT_ID);
-        if (!same_machine) {
-            auto local_desc = metadata_->segmentManager().getLocal();
-            same_machine = local_desc && !desc->machine_id.empty() &&
-                           !local_desc->machine_id.empty() &&
-                           desc->machine_id == local_desc->machine_id;
-        }
-        auto remote_mtype = getTypeEnum(LocationParser(entry->location).type());
-        for (auto type : entry->transports) {
-            if ((type == NVLINK || type == SHM) && !same_machine) continue;
-            if (checkAvailability(transport_list_[type], local_mtype,
-                                  remote_mtype)) {
-                if (priority-- == 0) return type;
+
+    const TransportType hint = request.transport_hint;
+
+    // Legacy mode: use original logic (before TransportSelector)
+    if (transport_selector_ && transport_selector_->isLegacyMode()) {
+        SelectionResult result;
+        std::vector<TransportType> raw;
+        if (desc->type == SegmentType::File) {
+            if (checkAvailability(transport_list_[GDS], local_mtype))
+                raw.push_back(GDS);
+            if (checkAvailability(transport_list_[IOURING], local_mtype))
+                raw.push_back(IOURING);
+        } else {
+            auto entry =
+                desc->findBuffer(request.target_offset, request.length);
+            if (entry) {
+                bool same_machine = (request.target_id == LOCAL_SEGMENT_ID);
+                if (!same_machine) {
+                    auto local_desc = metadata_->segmentManager().getLocal();
+                    same_machine = local_desc && !desc->machine_id.empty() &&
+                                   !local_desc->machine_id.empty() &&
+                                   desc->machine_id == local_desc->machine_id;
+                }
+                auto remote_mtype =
+                    getTypeEnum(LocationParser(entry->location).type());
+                for (auto type : entry->transports) {
+                    if ((type == NVLINK || type == SHM) && !same_machine)
+                        continue;
+                    if (checkAvailability(transport_list_[type], local_mtype,
+                                          remote_mtype)) {
+                        raw.push_back(type);
+                    }
+                }
             }
         }
-        return UNSPEC;
+
+        auto candidates = TransportSelector::reorderWithHint(raw, hint);
+        if (!candidates) {
+            return result;  // UNSPEC: hint not authorized for this req
+        }
+        if (transport_index >= 0 &&
+            (size_t)transport_index < candidates->size()) {
+            result.transport = (*candidates)[transport_index];
+        }
+        return result;
     }
+
+    // Selector mode: build ctx, then defer everything to
+    // TransportSelector::select().
+    SelectionContext ctx;
+    ctx.transfer_size = request.length;
+    ctx.priority_level =
+        request.priority;  // Use request priority for selection
+    ctx.policy_name = request.policy_name;  // Optional: bind to specific policy
+
+    if (desc->type == SegmentType::File) {
+        // File segment: use selector with empty buffer_transports
+        ctx.segment_type = SegmentType::File;
+        ctx.same_machine = true;  // File is always local
+        ctx.local_memory_type = local_mtype;
+        ctx.remote_memory_type = MTYPE_CPU;
+        ctx.buffer_transports = nullptr;  // Empty - use policy priority
+    } else {
+        // Memory segment
+        auto entry = desc->findBuffer(request.target_offset, request.length);
+        if (!entry) return SelectionResult{};
+        bool same_machine =
+            (desc->machine_id ==
+             metadata_->segmentManager().getLocal()->machine_id);
+        auto remote_mtype = getTypeEnum(LocationParser(entry->location).type());
+
+        ctx.segment_type = SegmentType::Memory;
+        ctx.same_machine = same_machine;
+        ctx.local_memory_type = local_mtype;
+        ctx.remote_memory_type = remote_mtype;
+        ctx.buffer_transports = &entry->transports;
+    }
+
+    return transport_selector_->select(ctx, transport_list_, transport_index,
+                                       hint);
 }
 
 static const char* transportTypeName(TransportType type) {
     switch (type) {
+        case UNSPEC:
+            return "UNSPEC";
         case RDMA:
             return "RDMA";
         case MNNVL:
@@ -855,16 +1040,18 @@ static const char* transportTypeName(TransportType type) {
             return "TCP";
         case AscendDirect:
             return "AscendDirect";
-        default:
-            return "UNSPEC";
+        case SUNRISE_LINK:
+            return "SUNRISE_LINK";
     }
+    return "UNKNOWN";
 }
 
 std::string printRequest(const Request& request) {
     std::stringstream ss;
     ss << "opcode " << request.opcode << " source " << request.source
        << " target_id " << request.target_id << " target_offset "
-       << (void*)request.target_offset << " length " << request.length;
+       << (void*)request.target_offset << " length " << request.length
+       << " transport_hint " << transportTypeName(request.transport_hint);
     return ss.str();
 }
 
@@ -885,6 +1072,27 @@ struct MergeResult {
     std::map<size_t, size_t> task_lookup;
 };
 
+struct TransferEngineImpl::PreparedSubmit {
+    struct Task {
+        size_t merged_task_index{0};
+        size_t task_id{0};
+    };
+
+    struct Owner {
+        size_t owner_task_id{0};
+        bool has_owner_task_id{false};
+        std::vector<size_t> derived_task_ids;
+        Request request{};
+        SelectionResult route{};
+        bool staging{false};
+        std::vector<std::string> staging_params;
+    };
+
+    std::chrono::steady_clock::time_point submit_time{};
+    std::vector<Task> tasks;
+    std::vector<Owner> owners;
+};
+
 namespace {
 
 bool tryAddUint64(uint64_t lhs, uint64_t rhs, uint64_t& out) {
@@ -933,6 +1141,8 @@ MergeResult mergeRequests(const std::vector<Request>& requests,
         if (a.req.opcode != b.req.opcode) return a.req.opcode < b.req.opcode;
         if (a.req.target_id != b.req.target_id)
             return a.req.target_id < b.req.target_id;
+        if (a.req.transport_hint != b.req.transport_hint)
+            return a.req.transport_hint < b.req.transport_hint;
         if (a.req.target_offset != b.req.target_offset)
             return a.req.target_offset < b.req.target_offset;
         return requestSourceAddr(a.req) < requestSourceAddr(b.req);
@@ -943,6 +1153,10 @@ MergeResult mergeRequests(const std::vector<Request>& requests,
             last.req.target_id != curr.req.target_id) {
             return false;
         }
+        // Mixed transport_hint inside one batch must not be merged.
+        if (last.req.transport_hint != curr.req.transport_hint) {
+            return false;
+        }
         if (!last.boundary.source_key || !curr.boundary.source_key ||
             !last.boundary.target_key || !curr.boundary.target_key) {
             return false;
@@ -993,43 +1207,61 @@ std::optional<BufferKey> toBufferKey(BufferDesc* buffer) {
     return BufferKey{buffer->addr, buffer->length};
 }
 
-RequestBoundaryInfo resolveRequestBoundary(ControlService* metadata,
-                                           SegmentDesc* local_desc,
-                                           const Request& request) {
-    RequestBoundaryInfo boundary;
-
-    if (local_desc) {
-        auto source_addr =
-            static_cast<uint64_t>(reinterpret_cast<uintptr_t>(request.source));
-        boundary.source_key =
-            toBufferKey(local_desc->findBuffer(source_addr, request.length));
-    }
-
-    metadata->segmentManager().withCachedSegment(
-        request.target_id, [&](SegmentDesc* target_desc) {
-            auto buffer =
-                target_desc->findBuffer(request.target_offset, request.length);
-
-            if (!buffer) {
-                boundary.target_key = std::nullopt;
-                return Status::NeedsRefreshCache(
-                    "Requested address is not in registered buffer" LOC_MARK);
-            }
-
-            boundary.target_key = toBufferKey(buffer);
-            return Status::OK();
-        });
-    return boundary;
-}
-
 std::vector<RequestBoundaryInfo> resolveRequestBoundaries(
     ControlService* metadata, const std::vector<Request>& requests) {
-    std::vector<RequestBoundaryInfo> boundaries;
-    boundaries.reserve(requests.size());
+    // Group requests by target_id so withCachedSegment fires at most once per
+    // peer.
+    std::vector<RequestBoundaryInfo> boundaries(requests.size());
     auto* local_desc = metadata->segmentManager().getLocal().get();
-    for (const auto& request : requests) {
-        boundaries.push_back(
-            resolveRequestBoundary(metadata, local_desc, request));
+
+    if (local_desc) {
+        for (size_t i = 0; i < requests.size(); ++i) {
+            auto source_addr = static_cast<uint64_t>(
+                reinterpret_cast<uintptr_t>(requests[i].source));
+            boundaries[i].source_key = toBufferKey(
+                local_desc->findBuffer(source_addr, requests[i].length));
+        }
+    }
+
+    std::unordered_map<SegmentID, std::vector<size_t>> by_target;
+    for (size_t i = 0; i < requests.size(); ++i) {
+        by_target[requests[i].target_id].push_back(i);
+    }
+
+    for (auto& [target_id, idxs] : by_target) {
+        metadata->segmentManager().withCachedSegment(
+            target_id, [&](SegmentDesc* target_desc) {
+                bool any_missing = false;
+                for (size_t i : idxs) {
+                    const auto& r = requests[i];
+                    auto* buffer =
+                        target_desc->findBuffer(r.target_offset, r.length);
+                    if (!buffer) {
+                        any_missing = true;
+                        boundaries[i].target_key = std::nullopt;
+                    } else {
+                        boundaries[i].target_key = toBufferKey(buffer);
+                    }
+                }
+                // Invariant: when this lambda returns NeedsRefreshCache, all
+                // writes it made in this pass are wiped before it returns.
+                // Reason: withCachedSegment will invalidate the cache and try
+                // ONE refetch; if that refetch fails (e.g. peer RPC down) the
+                // retry pass never runs, and any tentative writes from this
+                // (stale) pass would leak downstream into mergeRequests. By
+                // clearing here we leave a clean nullopt state for the group,
+                // and the retry pass (if it does run) repopulates from the
+                // fresh desc so the wipe is harmless.
+                if (any_missing) {
+                    for (size_t i : idxs) {
+                        boundaries[i].target_key = std::nullopt;
+                    }
+                    return Status::NeedsRefreshCache(
+                        "Requested address is not in registered "
+                        "buffer" LOC_MARK);
+                }
+                return Status::OK();
+            });
     }
     return boundaries;
 }
@@ -1098,45 +1330,94 @@ void TransferEngineImpl::findStagingPolicy(const Request& request,
     }
 }
 
-TransportType TransferEngineImpl::resolveTransport(const Request& req,
-                                                   int priority,
-                                                   bool invalidate_on_fail) {
-    auto type = getTransportType(req, priority);
-    if (type == UNSPEC && invalidate_on_fail) {
+SelectionResult TransferEngineImpl::resolveTransport(const Request& req,
+                                                     int transport_index,
+                                                     bool invalidate_on_fail) {
+    auto result = getTransportType(req, transport_index);
+    if (result.transport == UNSPEC && invalidate_on_fail) {
         metadata_->segmentManager().invalidateRemote(req.target_id);
-        type = getTransportType(req, priority);
+        result = getTransportType(req, transport_index);
     }
-    return type;
+    return result;
 }
 
-Status TransferEngineImpl::submitTransfer(
-    BatchID batch_id, const std::vector<Request>& request_list) {
-    if (!batch_id) return Status::InvalidArgument("Invalid batch ID" LOC_MARK);
-    Batch* batch = (Batch*)(batch_id);
-
-    std::vector<Request> classified_request_list[kSupportedTransportTypes];
-    std::vector<size_t> task_id_list[kSupportedTransportTypes];
-    std::unordered_map<size_t, TaskInfo> merged_task_id_map;
-
-    size_t start_task_id = batch->task_list.size();
-    batch->task_list.insert(batch->task_list.end(), request_list.size(),
-                            TaskInfo{});
-
-    // Record start time for metrics tracking
-    auto submit_time = std::chrono::steady_clock::now();
+Status TransferEngineImpl::prepareSubmit(
+    Batch* batch, const std::vector<Request>& request_list,
+    PreparedSubmit& prepared) {
+    if (!batch) return Status::InvalidArgument("Invalid batch" LOC_MARK);
+    for (size_t i = 0; i < request_list.size(); ++i) {
+        auto st = validateTransportHint(request_list[i], i);
+        if (!st.ok()) return st;
+    }
 
+    prepared = PreparedSubmit{};
+    const size_t start_task_id = batch->task_list.size();
+    prepared.submit_time = std::chrono::steady_clock::now();
     auto merge_boundaries =
         merge_requests_
             ? resolveRequestBoundaries(metadata_.get(), request_list)
             : std::vector<RequestBoundaryInfo>{};
     auto merged =
         mergeRequests(request_list, merge_boundaries, merge_requests_);
+
+    prepared.owners.reserve(merged.request_list.size());
+    for (const auto& request : merged.request_list) {
+        PreparedSubmit::Owner owner;
+        owner.request = request;
+        owner.route = resolveTransport(owner.request, 0);
+        if (owner.route.transport == TCP) {
+            findStagingPolicy(owner.request, owner.staging_params);
+            owner.staging = !owner.staging_params.empty() && staging_proxy_;
+        }
+        prepared.owners.push_back(std::move(owner));
+    }
+
+    prepared.tasks.reserve(merged.task_lookup.size());
+    for (const auto& kv : merged.task_lookup) {
+        const size_t public_task_index = kv.first;
+        const size_t merged_task_index = kv.second;
+        const size_t task_id = start_task_id + public_task_index;
+        auto& owner = prepared.owners[merged_task_index];
+        if (!owner.has_owner_task_id) {
+            owner.owner_task_id = task_id;
+            owner.has_owner_task_id = true;
+        } else {
+            owner.derived_task_ids.push_back(task_id);
+        }
+        prepared.tasks.push_back({merged_task_index, task_id});
+    }
+    return Status::OK();
+}
+
+uint64_t TransferEngineImpl::nextBatchToken() { return next_batch_token_++; }
+
+void TransferEngineImpl::attachProgressNotifier(
+    Batch* batch, Transport::SubBatchRef sub_batch) {
+    if (!batch || !sub_batch) return;
+    sub_batch->progress_batch_id = (BatchID)batch;
+    sub_batch->notify_progress = [this](BatchID batch_id) {
+        notifyBatchMaybeReady(batch_id);
+    };
+}
+
+Status TransferEngineImpl::commitPreparedSubmit(
+    Batch* batch, const PreparedSubmit& prepared) {
+    if (!batch) return Status::InvalidArgument("Invalid batch" LOC_MARK);
+
+    std::vector<Request> classified_request_list[kSupportedTransportTypes];
+    std::vector<size_t> task_id_list[kSupportedTransportTypes];
+    std::unordered_map<size_t, TaskInfo> merged_task_id_map;
+
+    batch->task_list.insert(batch->task_list.end(), prepared.tasks.size(),
+                            TaskInfo{});
+
     std::unordered_map<TransportType, size_t> next_sub_task_id;
-    for (auto& kv : merged.task_lookup) {
-        size_t task_id = start_task_id + kv.first;
-        size_t merged_task_id = kv.second;
+    for (const auto& task_plan : prepared.tasks) {
+        size_t task_id = task_plan.task_id;
+        size_t merged_task_id = task_plan.merged_task_index;
         auto& task = batch->task_list[task_id];
-        auto& merged_request = merged.request_list[merged_task_id];
+        const auto& owner = prepared.owners[merged_task_id];
+        auto& merged_request = owner.request;
         if (merged_task_id_map.count(merged_task_id)) {
             task = merged_task_id_map[merged_task_id];
             task.derived = true;
@@ -1144,13 +1425,15 @@ Status TransferEngineImpl::submitTransfer(
             continue;
         }
 
+        task.failover_count = 0;
         task.xport_priority = 0;
         task.status = PENDING;
         task.request = merged_request;
         task.staging = false;
         task.start_time =
-            submit_time;  // Record start time for latency tracking
-        task.type = resolveTransport(merged_request, 0);
+            prepared.submit_time;  // Record start time for latency tracking
+        task.type = owner.route.transport;
+        task.device_mask = owner.route.device_mask;
         if (task.type == UNSPEC) {
             LOG(WARNING) << "Unable to find registered buffer for request: "
                          << printRequest(merged_request);
@@ -1158,14 +1441,10 @@ Status TransferEngineImpl::submitTransfer(
             continue;
         }
 
-        if (task.type == TCP) {
-            std::vector<std::string> staging_params;
-            findStagingPolicy(merged_request, staging_params);
-            if (!staging_params.empty() && staging_proxy_) {
-                task.staging = true;
-                staging_proxy_->submit(&task, staging_params);
-                continue;
-            }
+        if (owner.staging) {
+            task.staging = true;
+            staging_proxy_->submit(&task, (BatchID)batch, owner.staging_params);
+            continue;
         }
 
         if (!batch->sub_batch[task.type]) {
@@ -1178,6 +1457,7 @@ Status TransferEngineImpl::submitTransfer(
                 merged_task_id_map[merged_task_id] = task;
                 continue;
             }
+            attachProgressNotifier(batch, batch->sub_batch[task.type]);
         }
 
         if (!next_sub_task_id.count(task.type))
@@ -1196,6 +1476,15 @@ Status TransferEngineImpl::submitTransfer(
         if (classified_request_list[type].empty()) continue;
         auto& transport = transport_list_[type];
         auto& sub_batch = batch->sub_batch[type];
+
+        // Set device_mask on SubBatch for RDMA transport
+        if (type == RDMA && !task_id_list[type].empty()) {
+            // Use the device_mask from the first task (we assume all tasks in
+            // this batch should have the same policy)
+            sub_batch->device_mask =
+                batch->task_list[task_id_list[type][0]].device_mask;
+        }
+
         auto status = transport->submitTransferTasks(
             sub_batch, classified_request_list[type]);
         if (!status.ok()) {
@@ -1209,6 +1498,303 @@ Status TransferEngineImpl::submitTransfer(
     return Status::OK();
 }
 
+Status TransferEngineImpl::enqueuePreparedSubmit(Batch* batch,
+                                                 const PreparedSubmit& prepared,
+                                                 QueueOwnerKind owner_kind) {
+    std::lock_guard<std::recursive_mutex> lk(progress_mutex_);
+    if (prepared.tasks.empty()) return Status::OK();
+    if (prepared.tasks.size() > batch->max_size - batch->task_list.size()) {
+        return Status::TooManyRequests(
+            "batch public task capacity exceeded" LOC_MARK);
+    }
+
+    const uint64_t batch_token =
+        batch->queue_token != 0 ? batch->queue_token : nextBatchToken();
+    QueueSubmit submit;
+    submit.batch_token = batch_token;
+    submit.batch_slots_left = batch->max_size - batch->task_list.size();
+    submit.owners.reserve(prepared.owners.size());
+    for (const auto& owner : prepared.owners) {
+        if (owner.request.length > runtime_queue_config_.max_dispatch_bytes) {
+            return Status::TooManyRequests(
+                "request exceeds runtime queue dispatch byte window" LOC_MARK);
+        }
+        QueueOwnerInput input;
+        input.owner_task_id = owner.owner_task_id;
+        input.derived_task_ids = owner.derived_task_ids;
+        input.request = owner.request;
+        input.kind = owner_kind;
+        submit.owners.push_back(std::move(input));
+    }
+
+    std::vector<QueueOwnerId> admitted_owner_ids;
+    CHECK_STATUS(runtime_queue_->tryAdmit(submit, admitted_owner_ids));
+    batch->queue_token = batch_token;
+
+    batch->task_list.insert(batch->task_list.end(), prepared.tasks.size(),
+                            TaskInfo{});
+    for (const auto& task_plan : prepared.tasks) {
+        auto& task = batch->task_list[task_plan.task_id];
+        const auto& owner = prepared.owners[task_plan.merged_task_index];
+        task.failover_count = 0;
+        task.xport_priority = 0;
+        task.status = PENDING;
+        task.request = owner.request;
+        task.staging = false;
+        task.start_time = prepared.submit_time;
+        task.type = UNSPEC;
+        task.sub_task_id = -1;
+        task.device_mask = owner.route.device_mask;
+        task.derived = task_plan.task_id != owner.owner_task_id;
+    }
+
+    for (size_t i = 0; i < admitted_owner_ids.size(); ++i) {
+        QueuedOwnerState queued;
+        queued.batch = batch;
+        queued.owner_task_id = prepared.owners[i].owner_task_id;
+        queued.byte_charge = prepared.owners[i].request.length;
+        queued.public_task_ids.push_back(prepared.owners[i].owner_task_id);
+        queued.public_task_ids.insert(
+            queued.public_task_ids.end(),
+            prepared.owners[i].derived_task_ids.begin(),
+            prepared.owners[i].derived_task_ids.end());
+        queued_owners_.emplace(admitted_owner_ids[i], queued);
+    }
+    return Status::OK();
+}
+
+Status TransferEngineImpl::finishQueuedOwner(
+    QueueOwnerId owner_id, TransferStatusEnum terminal_status) {
+    auto queued_it = queued_owners_.find(owner_id);
+    if (queued_it == queued_owners_.end()) {
+        return Status::InvalidEntry("queued owner not found" LOC_MARK);
+    }
+    auto& queued = queued_it->second;
+    if (queued.in_dispatch_window) {
+        if (dispatch_inflight_owners_ == 0 ||
+            dispatch_inflight_bytes_ < queued.byte_charge) {
+            return Status::InternalError(
+                "runtime dispatch window accounting underflow" LOC_MARK);
+        }
+    }
+    CHECK_STATUS(runtime_queue_->complete(owner_id, terminal_status));
+    if (queued.in_dispatch_window) {
+        --dispatch_inflight_owners_;
+        dispatch_inflight_bytes_ -= queued.byte_charge;
+        queued.in_dispatch_window = false;
+    }
+    for (const auto task_id : queued.public_task_ids) {
+        queued.batch->task_list[task_id].status = terminal_status;
+    }
+    queued_owners_.erase(queued_it);
+    return Status::OK();
+}
+
+Status TransferEngineImpl::retireQueueForBatch(Batch* batch) {
+    if (!batch || batch->queue_token == 0) return Status::OK();
+    auto status = runtime_queue_->retireBatch(batch->queue_token);
+    if (!status.ok()) return status;
+    batch->queue_token = 0;
+    return Status::OK();
+}
+
+Status TransferEngineImpl::markQueuedOwnerSubmitted(QueueOwnerId owner_id) {
+    auto queued_it = queued_owners_.find(owner_id);
+    if (queued_it == queued_owners_.end()) {
+        return Status::InternalError("queued owner metadata missing" LOC_MARK);
+    }
+    auto& queued = queued_it->second;
+    if (!queued.in_dispatch_window) {
+        queued.in_dispatch_window = true;
+        ++dispatch_inflight_owners_;
+        dispatch_inflight_bytes_ += queued.byte_charge;
+    }
+    return Status::OK();
+}
+
+Status TransferEngineImpl::dispatchQueuedOwner(QueueOwnerId owner_id) {
+    auto queued_it = queued_owners_.find(owner_id);
+    if (queued_it == queued_owners_.end()) {
+        return Status::InternalError("queued owner metadata missing" LOC_MARK);
+    }
+    const auto queued = queued_it->second;
+    auto* batch = queued.batch;
+    auto& task = batch->task_list[queued.owner_task_id];
+    auto route = resolveTransport(task.request, 0);
+    task.type = route.transport;
+    task.device_mask = route.device_mask;
+    if (task.type == UNSPEC) {
+        return finishQueuedOwner(owner_id, FAILED);
+    }
+
+    if (task.type == TCP) {
+        std::vector<std::string> staging_params;
+        findStagingPolicy(task.request, staging_params);
+        if (!staging_params.empty() && staging_proxy_) {
+            task.staging = true;
+            auto status =
+                staging_proxy_->submit(&task, (BatchID)batch, staging_params);
+            if (!status.ok()) return finishQueuedOwner(owner_id, FAILED);
+            return markQueuedOwnerSubmitted(owner_id);
+        }
+    }
+
+    if (!batch->sub_batch[task.type]) {
+        auto& transport = transport_list_[task.type];
+        if (!transport) return finishQueuedOwner(owner_id, FAILED);
+        auto status = transport->allocateSubBatch(batch->sub_batch[task.type],
+                                                  batch->max_size);
+        if (!status.ok()) return finishQueuedOwner(owner_id, FAILED);
+        attachProgressNotifier(batch, batch->sub_batch[task.type]);
+    }
+
+    auto& transport = transport_list_[task.type];
+    if (!transport) return finishQueuedOwner(owner_id, FAILED);
+    auto& sub_batch = batch->sub_batch[task.type];
+    if (task.type == RDMA) sub_batch->device_mask = task.device_mask;
+    task.sub_task_id = sub_batch->size();
+    auto status = transport->submitTransferTasks(sub_batch, {task.request});
+    if (!status.ok()) {
+        task.type = UNSPEC;
+        return finishQueuedOwner(owner_id, FAILED);
+    }
+    return markQueuedOwnerSubmitted(owner_id);
+}
+
+Status TransferEngineImpl::refillDispatchWindow() {
+    std::lock_guard<std::recursive_mutex> lk(progress_mutex_);
+    if (!runtime_queue_config_.enabled) return Status::OK();
+    if (dispatch_inflight_owners_ >=
+            runtime_queue_config_.max_dispatch_owners ||
+        dispatch_inflight_bytes_ >= runtime_queue_config_.max_dispatch_bytes) {
+        return Status::OK();
+    }
+
+    const size_t owner_budget =
+        runtime_queue_config_.max_dispatch_owners - dispatch_inflight_owners_;
+    const size_t byte_budget =
+        runtime_queue_config_.max_dispatch_bytes - dispatch_inflight_bytes_;
+    auto picked = runtime_queue_->pickForDispatch(owner_budget, byte_budget);
+    for (const auto owner_id : picked) {
+        CHECK_STATUS(dispatchQueuedOwner(owner_id));
+    }
+    return Status::OK();
+}
+
+Status TransferEngineImpl::progressRuntimeQueue() {
+    std::lock_guard<std::recursive_mutex> lk(progress_mutex_);
+    if (!runtime_queue_config_.enabled) return Status::OK();
+
+    CHECK_STATUS(refillDispatchWindow());
+
+    std::vector<QueueOwnerId> owner_ids;
+    owner_ids.reserve(queued_owners_.size());
+    for (const auto& entry : queued_owners_) {
+        if (entry.second.in_dispatch_window) owner_ids.push_back(entry.first);
+    }
+
+    bool released_window = false;
+    for (const auto owner_id : owner_ids) {
+        auto queued_it = queued_owners_.find(owner_id);
+        if (queued_it == queued_owners_.end()) continue;
+
+        auto& queued = queued_it->second;
+        if (!queued.in_dispatch_window) continue;
+        auto* batch = queued.batch;
+        if (!batch || !alive_batches_.count((BatchID)batch)) continue;
+        if (queued.owner_task_id >= batch->task_list.size()) {
+            return Status::InternalError(
+                "queued owner task id out of range" LOC_MARK);
+        }
+
+        auto& task = batch->task_list[queued.owner_task_id];
+        auto prev_status = task.status;
+        TransferStatus task_status;
+        CHECK_STATUS(pollTaskStatus(batch, queued.owner_task_id, task_status));
+        updateTaskStatusAfterPoll(batch, queued.owner_task_id, task_status,
+                                  true);
+        recordTaskCompletionMetrics(task, prev_status, task_status.s);
+
+        if (task_status.s == PENDING) continue;
+
+        CHECK_STATUS(finishQueuedOwner(owner_id, task_status.s));
+        if (task_status.s == COMPLETED)
+            CHECK_STATUS(maybeFireSubmitHooks(batch));
+        released_window = true;
+    }
+
+    if (released_window) CHECK_STATUS(refillDispatchWindow());
+    return Status::OK();
+}
+
+bool TransferEngineImpl::hasActiveRuntimeQueue() {
+    std::lock_guard<std::recursive_mutex> lk(progress_mutex_);
+    return runtime_queue_config_.enabled && !queued_owners_.empty();
+}
+
+bool TransferEngineImpl::shouldQueueSubmit(const PreparedSubmit& prepared,
+                                           QueueOwnerKind owner_kind) const {
+    if (!runtime_queue_config_.enabled) return false;
+    if (owner_kind == QueueOwnerKind::StagingInternal) return true;
+    return std::none_of(
+        prepared.owners.begin(), prepared.owners.end(),
+        [](const PreparedSubmit::Owner& owner) { return owner.staging; });
+}
+
+Status TransferEngineImpl::submitTransfer(
+    BatchID batch_id, const std::vector<Request>& request_list,
+    const Notification* notifi, QueueOwnerKind owner_kind) {
+    Batch* batch = nullptr;
+    CHECK_STATUS(retainBatch(batch_id, batch));
+    BatchRef batch_ref(*this, batch);
+    const size_t start_task_id = batch_ref.get()->task_list.size();
+    PreparedSubmit prepared;
+    CHECK_STATUS(prepareSubmit(batch_ref.get(), request_list, prepared));
+
+    if (shouldQueueSubmit(prepared, owner_kind)) {
+        CHECK_STATUS(
+            enqueuePreparedSubmit(batch_ref.get(), prepared, owner_kind));
+        auto dispatch_status = refillDispatchWindow();
+        if (!dispatch_status.ok()) {
+            LOG(WARNING) << "runtime queue dispatch failed after admission: "
+                         << dispatch_status.ToString();
+        }
+        notifyRuntimeQueueReady();
+    } else {
+        CHECK_STATUS(commitPreparedSubmit(batch_ref.get(), prepared));
+    }
+
+    if (notifi) {
+        addSubmitHook(batch_ref.get(), start_task_id, request_list, *notifi);
+    }
+    return batch_ref.release();
+}
+
+Status TransferEngineImpl::submitTransfer(
+    BatchID batch_id, const std::vector<Request>& request_list) {
+    return submitTransfer(batch_id, request_list, nullptr,
+                          QueueOwnerKind::User);
+}
+
+Status TransferEngineImpl::submitStagingTransfer(
+    BatchID batch_id, const std::vector<Request>& request_list) {
+    return submitTransfer(batch_id, request_list, nullptr,
+                          QueueOwnerKind::StagingInternal);
+}
+
+void TransferEngineImpl::addSubmitHook(Batch* batch, size_t start_task_id,
+                                       const std::vector<Request>& request_list,
+                                       const Notification& notifi) {
+    Batch::SubmitHook hook;
+    hook.start_task_id = start_task_id;
+    hook.end_task_id = start_task_id + request_list.size();
+    hook.notifi = notifi;
+    hook.fired = false;
+    for (const auto& request : request_list)
+        hook.targets.insert(request.target_id);
+    batch->submit_hooks.emplace_back(std::move(hook));
+}
+
 Status TransferEngineImpl::maybeFireSubmitHooks(Batch* batch, bool check) {
     for (auto& hook : batch->submit_hooks) {
         if (hook.fired) continue;
@@ -1244,20 +1830,8 @@ Status TransferEngineImpl::maybeFireSubmitHooks(Batch* batch, bool check) {
 Status TransferEngineImpl::submitTransfer(
     BatchID batch_id, const std::vector<Request>& request_list,
     const Notification& notifi) {
-    if (!batch_id) return Status::InvalidArgument("Invalid batch ID" LOC_MARK);
-    Batch* batch = (Batch*)(batch_id);
-    const size_t start_task_id = batch->task_list.size();
-    CHECK_STATUS(submitTransfer(batch_id, request_list));
-    const size_t end_task_id = start_task_id + request_list.size();
-    Batch::SubmitHook hook;
-    hook.start_task_id = start_task_id;
-    hook.end_task_id = end_task_id;
-    hook.notifi = notifi;
-    hook.fired = false;
-    for (const auto& request : request_list)
-        hook.targets.insert(request.target_id);
-    batch->submit_hooks.emplace_back(std::move(hook));
-    return Status::OK();
+    return submitTransfer(batch_id, request_list, &notifi,
+                          QueueOwnerKind::User);
 }
 
 Status TransferEngineImpl::resubmitTransferTask(Batch* batch, size_t task_id) {
@@ -1275,8 +1849,10 @@ Status TransferEngineImpl::resubmitTransferTask(Batch* batch, size_t task_id) {
     if (task.staging)
         task.staging = false;
     else
-        task.xport_priority++;
-    auto type = resolveTransport(task.request, task.xport_priority);
+        task.xport_priority = task.failover_count;
+
+    auto result = resolveTransport(task.request, task.xport_priority);
+    auto type = result.transport;
     if (type == UNSPEC) {
         LOG(WARNING) << "No more transports available after "
                      << transportTypeName(prev_type) << " failed";
@@ -1289,15 +1865,53 @@ Status TransferEngineImpl::resubmitTransferTask(Batch* batch, size_t task_id) {
     TENT_RECORD_TRANSPORT_FAILOVER();
 
     auto& transport = transport_list_[type];
-    if (!batch->sub_batch[type])
+    if (!batch->sub_batch[type]) {
         CHECK_STATUS(transport->allocateSubBatch(batch->sub_batch[type],
                                                  batch->max_size));
+        attachProgressNotifier(batch, batch->sub_batch[type]);
+    }
     auto& sub_batch = batch->sub_batch[type];
     task.sub_task_id = sub_batch->size();
     task.type = type;
     return transport->submitTransferTasks(sub_batch, {task.request});
 }
 
+Status TransferEngineImpl::pollTaskStatus(Batch* batch, size_t task_id,
+                                          TransferStatus& task_status) {
+    auto& task = batch->task_list[task_id];
+    if (task.staging) {
+        return staging_proxy_->getStatus(&task, task_status);
+    }
+
+    if (task.type == UNSPEC) {
+        task_status.s = FAILED;
+        task_status.transferred_bytes = 0;
+        return Status::OK();
+    }
+
+    auto& transport = transport_list_[task.type];
+    auto& sub_batch = batch->sub_batch[task.type];
+    if (!transport || !sub_batch) {
+        return Status::InvalidArgument("Transport not available" LOC_MARK);
+    }
+    return transport->getTransferStatus(sub_batch, task.sub_task_id,
+                                        task_status);
+}
+
+void TransferEngineImpl::updateTaskStatusAfterPoll(Batch* batch, size_t task_id,
+                                                   TransferStatus& task_status,
+                                                   bool allow_failover) {
+    auto& task = batch->task_list[task_id];
+    task.status = task_status.s;
+    if (!allow_failover || task_status.s != FAILED || task.type == UNSPEC)
+        return;
+
+    if (resubmitTransferTask(batch, task_id).ok()) {
+        task_status.s = PENDING;
+        task.status = PENDING;
+    }
+}
+
 Status TransferEngineImpl::sendNotification(SegmentID target_id,
                                             const Notification& notifi) {
     for (size_t type = 0; type < kSupportedTransportTypes; ++type) {
@@ -1340,37 +1954,58 @@ Status TransferEngineImpl::receiveNotification(
 Status TransferEngineImpl::getTransferStatus(BatchID batch_id, size_t task_id,
                                              TransferStatus& task_status) {
     if (!batch_id) return Status::InvalidArgument("Invalid batch ID" LOC_MARK);
+    std::lock_guard<std::recursive_mutex> lk(progress_mutex_);
+    if (!alive_batches_.count(batch_id))
+        return Status::InvalidArgument("Batch is not alive" LOC_MARK);
     Batch* batch = (Batch*)(batch_id);
     if (task_id >= batch->task_list.size())
         return Status::InvalidArgument("Invalid task ID" LOC_MARK);
-    auto& task = batch->task_list[task_id];
-    auto prev_status = task.status;
-    if (task.staging) {
-        CHECK_STATUS(staging_proxy_->getStatus(&task, task_status));
-    } else {
-        if (task.type == UNSPEC) {
-            task_status.s = FAILED;
-            task_status.transferred_bytes = 0;
-            batch->task_list[task_id].status = task_status.s;
-            return Status::OK();
+    const size_t public_task_id = task_id;
+    size_t poll_task_id = task_id;
+    CHECK_STATUS(refillDispatchWindow());
+    if (runtime_queue_config_.enabled && batch->queue_token != 0) {
+        QueueOwnerId owner_id = 0;
+        auto resolve_status = runtime_queue_->resolveOwner(
+            batch->queue_token, public_task_id, owner_id);
+        if (resolve_status.ok()) {
+            TransferStatusEnum public_status = PENDING;
+            CHECK_STATUS(runtime_queue_->getPublicStatus(
+                batch->queue_token, public_task_id, public_status));
+            auto queued_it = queued_owners_.find(owner_id);
+            if (public_status != PENDING ||
+                (queued_it != queued_owners_.end() &&
+                 !queued_it->second.in_dispatch_window)) {
+                task_status.s = public_status;
+                task_status.transferred_bytes =
+                    public_status == COMPLETED
+                        ? batch->task_list[public_task_id].request.length
+                        : 0;
+                return Status::OK();
+            }
+            if (batch->task_list[public_task_id].derived &&
+                queued_it != queued_owners_.end()) {
+                poll_task_id = queued_it->second.owner_task_id;
+            }
         }
-        auto& transport = transport_list_[task.type];
-        auto& sub_batch = batch->sub_batch[task.type];
-        if (!transport || !sub_batch) {
-            return Status::InvalidArgument("Transport not available" LOC_MARK);
-        }
-        CHECK_STATUS(transport->getTransferStatus(sub_batch, task.sub_task_id,
-                                                  task_status));
     }
-    batch->task_list[task_id].status = task_status.s;
-
-    if (task_status.s == FAILED && resubmitTransferTask(batch, task_id).ok()) {
-        task_status.s = PENDING;
-        batch->task_list[task_id].status = PENDING;
+    auto& task = batch->task_list[poll_task_id];
+    auto prev_status = task.status;
+    CHECK_STATUS(pollTaskStatus(batch, poll_task_id, task_status));
+    updateTaskStatusAfterPoll(batch, poll_task_id, task_status,
+                              enable_auto_failover_on_poll_);
+    if (runtime_queue_config_.enabled && batch->queue_token != 0 &&
+        task_status.s != PENDING) {
+        QueueOwnerId owner_id = 0;
+        auto resolve_status = runtime_queue_->resolveOwner(
+            batch->queue_token, public_task_id, owner_id);
+        if (resolve_status.ok()) {
+            CHECK_STATUS(finishQueuedOwner(owner_id, task_status.s));
+            CHECK_STATUS(refillDispatchWindow());
+        }
     }
 
     // Record metrics when task transitions to terminal state
-    recordTaskCompletionMetrics(batch->task_list[task_id], prev_status,
+    recordTaskCompletionMetrics(batch->task_list[poll_task_id], prev_status,
                                 task_status.s);
 
     if (task_status.s == COMPLETED) CHECK_STATUS(maybeFireSubmitHooks(batch));
@@ -1380,6 +2015,9 @@ Status TransferEngineImpl::getTransferStatus(BatchID batch_id, size_t task_id,
 Status TransferEngineImpl::getTransferStatus(
     BatchID batch_id, std::vector<TransferStatus>& status_list) {
     if (!batch_id) return Status::InvalidArgument("Invalid batch ID" LOC_MARK);
+    std::lock_guard<std::recursive_mutex> lk(progress_mutex_);
+    if (!alive_batches_.count(batch_id))
+        return Status::InvalidArgument("Batch is not alive" LOC_MARK);
     Batch* batch = (Batch*)(batch_id);
     status_list.clear();
     for (size_t task_id = 0; task_id < batch->task_list.size(); ++task_id) {
@@ -1390,78 +2028,135 @@ Status TransferEngineImpl::getTransferStatus(
     return Status::OK();
 }
 
-Status TransferEngineImpl::getTransferStatus(BatchID batch_id,
-                                             TransferStatus& overall_status) {
+Status TransferEngineImpl::getBatchStatus(BatchID batch_id,
+                                          TransferStatus& overall_status,
+                                          bool allow_failover) {
     if (!batch_id) return Status::InvalidArgument("Invalid batch ID" LOC_MARK);
+    std::lock_guard<std::recursive_mutex> lk(progress_mutex_);
+    if (!alive_batches_.count(batch_id))
+        return Status::InvalidArgument("Batch is not alive" LOC_MARK);
+    CHECK_STATUS(refillDispatchWindow());
     Batch* batch = (Batch*)(batch_id);
     overall_status.s = PENDING;
     overall_status.transferred_bytes = 0;
     size_t success_tasks = 0;
+    size_t failed_tasks = 0;
     size_t total_tasks = 0;
+    TransferStatusEnum worst_failure = PENDING;
+    auto isWorse = [](TransferStatusEnum cur, TransferStatusEnum best) {
+        static const std::unordered_map<TransferStatusEnum, int> severity = {
+            {INITIAL, 0},  {PENDING, 0}, {COMPLETED, 0}, {INVALID, 1},
+            {CANCELED, 2}, {TIMEOUT, 3}, {FAILED, 4},
+        };
+        return severity.at(cur) > severity.at(best);
+    };
     for (size_t task_id = 0; task_id < batch->task_list.size(); ++task_id) {
         auto& task = batch->task_list[task_id];
         if (task.derived) continue;  // This task is performed by other tasks
         total_tasks++;
+        if (runtime_queue_config_.enabled && batch->queue_token != 0) {
+            QueueOwnerId owner_id = 0;
+            auto resolve_status = runtime_queue_->resolveOwner(
+                batch->queue_token, task_id, owner_id);
+            if (resolve_status.ok()) {
+                TransferStatusEnum public_status = PENDING;
+                CHECK_STATUS(runtime_queue_->getPublicStatus(
+                    batch->queue_token, task_id, public_status));
+                auto queued_it = queued_owners_.find(owner_id);
+                if (public_status == PENDING) {
+                    if (queued_it != queued_owners_.end() &&
+                        !queued_it->second.in_dispatch_window) {
+                        continue;
+                    }
+                }
+                if (public_status == COMPLETED) {
+                    success_tasks++;
+                    overall_status.transferred_bytes += task.request.length;
+                    continue;
+                }
+                if (public_status != PENDING) {
+                    failed_tasks++;
+                    if (isWorse(public_status, worst_failure))
+                        worst_failure = public_status;
+                    continue;
+                }
+            }
+        }
         TransferStatus task_status;
         if (task.status != PENDING) {
             if (task.status == COMPLETED) {
                 success_tasks++;
                 overall_status.transferred_bytes += task.request.length;
             } else {
-                overall_status.s = task.status;
+                failed_tasks++;
+                if (isWorse(task.status, worst_failure))
+                    worst_failure = task.status;
             }
             continue;
         }
         auto prev_status = task.status;
-        if (task.staging) {
-            CHECK_STATUS(staging_proxy_->getStatus(&task, task_status));
-        } else {
-            if (task.type == UNSPEC) {
-                task.status = FAILED;
-                overall_status.s = FAILED;
-                continue;
+        CHECK_STATUS(pollTaskStatus(batch, task_id, task_status));
+        updateTaskStatusAfterPoll(batch, task_id, task_status, allow_failover);
+        if (runtime_queue_config_.enabled && batch->queue_token != 0 &&
+            task_status.s != PENDING) {
+            QueueOwnerId owner_id = 0;
+            auto resolve_status = runtime_queue_->resolveOwner(
+                batch->queue_token, task_id, owner_id);
+            if (resolve_status.ok()) {
+                CHECK_STATUS(finishQueuedOwner(owner_id, task_status.s));
+                CHECK_STATUS(refillDispatchWindow());
             }
-            auto& transport = transport_list_[task.type];
-            auto& sub_batch = batch->sub_batch[task.type];
-            if (!transport || !sub_batch) {
-                return Status::InvalidArgument(
-                    "Transport not available" LOC_MARK);
-            }
-            CHECK_STATUS(transport->getTransferStatus(
-                sub_batch, task.sub_task_id, task_status));
-        }
-        // memorize task result
-        task.status = task_status.s;
-
-        // Attempt failover before status aggregation so that a
-        // successfully resubmitted task appears as PENDING and does not
-        // overwrite a permanent FAILED from another task in the batch.
-        if (task_status.s == FAILED &&
-            resubmitTransferTask(batch, task_id).ok()) {
-            task.status = PENDING;
-            task_status.s = PENDING;
         }
 
         if (task_status.s == COMPLETED) {
             success_tasks++;
             overall_status.transferred_bytes += task_status.transferred_bytes;
         } else if (task_status.s != PENDING) {
-            overall_status.s = task_status.s;
+            failed_tasks++;
+            if (isWorse(task_status.s, worst_failure))
+                worst_failure = task_status.s;
         }
 
         // Record metrics when task transitions to terminal state
         recordTaskCompletionMetrics(batch->task_list[task_id], prev_status,
                                     task_status.s);
     }
-    if (success_tasks == total_tasks) overall_status.s = COMPLETED;
+    // Determine overall status: COMPLETED only when all succeed; FAILED only
+    // when all tasks are terminal (no in-flight work) and at least one failed;
+    // otherwise PENDING (some tasks still running).
+    if (success_tasks == total_tasks) {
+        overall_status.s = COMPLETED;
+    } else if (success_tasks + failed_tasks == total_tasks) {
+        overall_status.s = worst_failure;
+    }
+    // else: some tasks still PENDING → overall_status.s stays PENDING
     CHECK_STATUS(maybeFireSubmitHooks(batch, overall_status.s == COMPLETED));
     return Status::OK();
 }
 
+Status TransferEngineImpl::getTransferStatus(BatchID batch_id,
+                                             TransferStatus& overall_status) {
+    return getBatchStatus(batch_id, overall_status,
+                          enable_auto_failover_on_poll_);
+}
+
+Status TransferEngineImpl::progressBatch(BatchID batch_id,
+                                         TransferStatus& overall_status) {
+    return getBatchStatus(batch_id, overall_status, true);
+}
+
+void TransferEngineImpl::notifyBatchMaybeReady(BatchID batch_id) {
+    if (progress_worker_) progress_worker_->notifyBatchMaybeReady(batch_id);
+}
+
+void TransferEngineImpl::notifyRuntimeQueueReady() {
+    if (progress_worker_) progress_worker_->notifyRuntimeQueueReady();
+}
+
 Status TransferEngineImpl::waitTransferCompletion(BatchID batch_id) {
     TransferStatus xfer_status;
     while (true) {
-        CHECK_STATUS(getTransferStatus(batch_id, xfer_status));
+        CHECK_STATUS(progressBatch(batch_id, xfer_status));
         if (xfer_status.s != PENDING) {
             freeBatch(batch_id);
             return xfer_status.s == COMPLETED
@@ -1479,7 +2174,7 @@ Status TransferEngineImpl::transferSync(
     CHECK_STATUS(submitTransfer(batch_id, request_list));
     while (true) {
         TransferStatus xfer_status;
-        CHECK_STATUS(getTransferStatus(batch_id, xfer_status));
+        CHECK_STATUS(progressBatch(batch_id, xfer_status));
         if (xfer_status.s == COMPLETED) break;
         if (xfer_status.s != PENDING) {
             CHECK_STATUS(freeBatch(batch_id));
@@ -1519,6 +2214,26 @@ void TransferEngineImpl::recordTaskCompletionMetrics(
                     TentMetrics::instance().recordWriteCompleted(
                         task.request.length, latency_seconds);
                 }
+                // Observability only (RFC #2519): if this transfer carried a
+                // deadline, emit the post-hoc feasibility ratio MLU =
+                // actual_transfer_time / available_window, where the window is
+                // (deadline - submit_time). MLU < 1 met the deadline; >= 1
+                // missed it. This does not drive any admission/scheduling yet.
+                if (task.request.deadline_ns != 0) {
+                    uint64_t start_ns = static_cast<uint64_t>(
+                        std::chrono::duration_cast<std::chrono::nanoseconds>(
+                            start_time.time_since_epoch())
+                            .count());
+                    if (task.request.deadline_ns > start_ns) {
+                        double window_seconds =
+                            (task.request.deadline_ns - start_ns) / 1e9;
+                        TentMetrics::instance().recordDeadlineMLU(
+                            latency_seconds / window_seconds);
+                    } else {
+                        // Deadline already in the past at submit: infeasible.
+                        TentMetrics::instance().recordDeadlineMLU(5.0);
+                    }
+                }
             } else if (new_status == FAILED) {
                 if (task.request.opcode == Request::READ) {
                     TentMetrics::instance().recordReadFailed(
diff --git a/mooncake-transfer-engine/tent/src/runtime/transport_loader.cpp b/mooncake-transfer-engine/tent/src/runtime/transport_loader.cpp
index 4d9d4072..7ecd5127 100644
--- a/mooncake-transfer-engine/tent/src/runtime/transport_loader.cpp
+++ b/mooncake-transfer-engine/tent/src/runtime/transport_loader.cpp
@@ -37,6 +37,10 @@
 #include "tent/transport/ascend/ascend_direct_transport.h"
 #endif
 
+#ifdef USE_SUNRISE
+#include "tent/transport/sunrise_link/sunrise_link_transport.h"
+#endif
+
 namespace mooncake {
 namespace tent {
 
@@ -83,6 +87,13 @@ Status TransferEngineImpl::loadTransports() {
     }
 #endif
 
+#ifdef USE_SUNRISE
+    if (conf_->get("transports/sunrise_link/enable", true)) {
+        transport_list_[SUNRISE_LINK] =
+            std::make_shared<SunriseLinkTransport>();
+    }
+#endif
+
     return Status::OK();
 }
 
diff --git a/mooncake-transfer-engine/tent/src/runtime/transport_selector.cpp b/mooncake-transfer-engine/tent/src/runtime/transport_selector.cpp
new file mode 100644
index 00000000..e26f1e78
--- /dev/null
+++ b/mooncake-transfer-engine/tent/src/runtime/transport_selector.cpp
@@ -0,0 +1,460 @@
+// Copyright 2025 KVCache.AI
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "tent/runtime/transport_selector.h"
+#include "tent/runtime/transport.h"
+#include "tent/runtime/platform.h"
+#include "tent/thirdparty/nlohmann/json.h"
+
+#include <algorithm>
+#include <glog/logging.h>
+
+#include <algorithm>
+#include <cctype>
+
+namespace mooncake {
+namespace tent {
+
+// Transport type name mapping
+static const std::unordered_map<std::string, TransportType> kTransportNameMap =
+    {
+        {"unspec", UNSPEC},       {"rdma", RDMA},
+        {"mnnvl", MNNVL},         {"shm", SHM},
+        {"nvlink", NVLINK},       {"gds", GDS},
+        {"io_uring", IOURING},    {"tcp", TCP},
+        {"ascend", AscendDirect}, {"sunrise_link", SUNRISE_LINK},
+};
+
+static const std::unordered_map<TransportType, std::string>
+    kTransportTypeNames = {
+        {UNSPEC, "unspec"},       {RDMA, "rdma"},
+        {MNNVL, "mnnvl"},         {SHM, "shm"},
+        {NVLINK, "nvlink"},       {GDS, "gds"},
+        {IOURING, "io_uring"},    {TCP, "tcp"},
+        {AscendDirect, "ascend"}, {SUNRISE_LINK, "sunrise_link"},
+};
+
+// Memory type name mapping for pattern matching
+static const std::string kMemoryTypeCpu = "cpu";
+static const std::string kMemoryTypeCuda = "cuda";
+static const std::string kMemoryTypeNpu = "npu";
+static const std::string kMemoryTypeWildcard = "*";
+
+std::string TransportSelector::transportTypeName(TransportType type) {
+    auto it = kTransportTypeNames.find(type);
+    if (it != kTransportTypeNames.end()) {
+        return it->second;
+    }
+    return "unknown";
+}
+
+TransportType TransportSelector::parseTransportType(const std::string& str) {
+    auto it = kTransportNameMap.find(str);
+    if (it != kTransportNameMap.end()) {
+        return it->second;
+    }
+    LOG(WARNING) << "Unknown transport type: " << str;
+    return UNSPEC;
+}
+
+std::vector<SelectionPolicy> TransportSelector::getDefaultPolicies() {
+    // Default policies that match the original hardcoded behavior
+    return {
+        {
+            "file_storage",
+            SegmentType::File,
+            std::nullopt,   // same_machine doesn't matter for file
+            std::nullopt,   // local_memory_pattern
+            std::nullopt,   // remote_memory_pattern
+            std::nullopt,   // min_size
+            std::nullopt,   // max_size
+            std::nullopt,   // priority
+            {},             // devices (empty = all devices)
+            {GDS, IOURING}  // File segment priority (original: GDS -> IOURING)
+        },
+        {
+            "memory_default",
+            SegmentType::Memory,
+            std::nullopt,  // any machine
+            std::nullopt,  // any local memory
+            std::nullopt,  // any remote memory
+            std::nullopt,  // any size
+            std::nullopt,  // min_priority
+            {},            // devices (empty = all devices)
+            {}  // Empty priority = use buffer_transports order (original
+                // behavior)
+        },
+    };
+}
+
+void TransportSelector::loadPolicies() {
+    policies_.clear();
+
+    // Try to load from configuration - read from "policy" subtree
+    auto policies_array = config_->getArray<json>("policy");
+
+    if (policies_array.empty()) {
+        LOG(INFO)
+            << "No 'policy' configured, using default transport selection";
+        policies_ = getDefaultPolicies();
+        return;
+    }
+
+    LOG(INFO) << "Loading transport selection policies from config";
+
+    for (const auto& policy_json : policies_array) {
+        if (!policy_json.is_object()) continue;
+
+        SelectionPolicy policy;
+        policy.name = policy_json.value("name", "unnamed");
+
+        // Parse segment_type
+        std::string segment_type_str = policy_json.value("segment_type", "");
+        if (segment_type_str == "file") {
+            policy.segment_type = SegmentType::File;
+        } else if (segment_type_str == "memory") {
+            policy.segment_type = SegmentType::Memory;
+        } else {
+            LOG(WARNING) << "Invalid segment_type in policy " << policy.name;
+            continue;
+        }
+
+        // Parse same_machine (optional)
+        if (policy_json.contains("same_machine")) {
+            policy.same_machine = policy_json["same_machine"].get<bool>();
+        } else {
+            policy.same_machine = std::nullopt;
+        }
+
+        // Parse memory patterns (optional)
+        if (policy_json.contains("local_memory")) {
+            policy.local_memory_pattern =
+                policy_json["local_memory"].get<std::string>();
+        }
+        if (policy_json.contains("remote_memory")) {
+            policy.remote_memory_pattern =
+                policy_json["remote_memory"].get<std::string>();
+        }
+
+        // Parse size filters (optional)
+        if (policy_json.contains("min_size")) {
+            policy.min_size = policy_json["min_size"].get<uint64_t>();
+        }
+        if (policy_json.contains("max_size")) {
+            policy.max_size = policy_json["max_size"].get<uint64_t>();
+        }
+
+        // Parse priority filter (optional)
+        if (policy_json.contains("priority")) {
+            if (policy_json["priority"].is_string()) {
+                // Parse string: "high", "medium", "low"
+                std::string prio_str =
+                    policy_json["priority"].get<std::string>();
+                // Convert to lowercase for case-insensitive matching
+                std::transform(prio_str.begin(), prio_str.end(),
+                               prio_str.begin(),
+                               [](unsigned char c) { return std::tolower(c); });
+                if (prio_str == "high" || prio_str == "0") {
+                    policy.priority = PRIO_HIGH;
+                } else if (prio_str == "medium" || prio_str == "1") {
+                    policy.priority = PRIO_MEDIUM;
+                } else if (prio_str == "low" || prio_str == "2") {
+                    policy.priority = PRIO_LOW;
+                } else {
+                    LOG(WARNING) << "Invalid priority string: " << prio_str
+                                 << ", using PRIO_LOW";
+                    policy.priority = PRIO_LOW;
+                }
+            } else {
+                // Parse integer for backward compatibility
+                policy.priority = policy_json["priority"].get<int>();
+            }
+        } else {
+            policy.priority = std::nullopt;
+        }
+
+        // Parse devices (optional)
+        if (policy_json.contains("devices")) {
+            for (const auto& device_name : policy_json["devices"]) {
+                if (device_name.is_string()) {
+                    policy.devices.push_back(device_name.get<std::string>());
+                }
+            }
+        }
+
+        // Parse transports list
+        policy.transports.clear();
+        if (policy_json.contains("transports")) {
+            for (const auto& transport_str : policy_json["transports"]) {
+                if (!transport_str.is_string()) continue;
+                TransportType type =
+                    parseTransportType(transport_str.get<std::string>());
+                if (type != UNSPEC) {
+                    policy.transports.push_back(type);
+                }
+            }
+        }
+
+        policies_.push_back(std::move(policy));
+        LOG(INFO) << "Loaded transport policy: " << policy.name
+                  << " (segment_type=" << segment_type_str
+                  << ", transports_count=" << policy.transports.size() << ")";
+    }
+}
+
+TransportSelector::TransportSelector(std::shared_ptr<Config> config)
+    : config_(config) {
+    loadPolicies();
+}
+
+bool TransportSelector::matchesMemoryPattern(const std::string& pattern,
+                                             MemoryType type) const {
+    if (pattern == kMemoryTypeWildcard) {
+        return true;
+    }
+
+    // Convert MemoryType to string for comparison
+    std::string type_str;
+    switch (type) {
+        case MTYPE_CPU:
+            type_str = kMemoryTypeCpu;
+            break;
+        case MTYPE_CUDA:
+            type_str = kMemoryTypeCuda;
+            break;
+        case MTYPE_ROCM:
+            type_str = "rocm";
+            break;
+        default:
+            type_str = "unknown";
+            break;
+    }
+
+    return pattern == type_str;
+}
+
+bool TransportSelector::matchesPolicy(const SelectionPolicy& policy,
+                                      const SelectionContext& context) const {
+    // If context specifies a policy_name, only match that exact policy with
+    // matching segment type
+    if (context.policy_name.has_value()) {
+        return context.policy_name.value() == policy.name &&
+               policy.segment_type == context.segment_type;
+    }
+
+    // Check segment type
+    if (policy.segment_type != context.segment_type) {
+        return false;
+    }
+
+    // Check same_machine constraint
+    if (policy.same_machine.has_value()) {
+        if (policy.same_machine.value() != context.same_machine) {
+            return false;
+        }
+    }
+
+    // Check local memory pattern
+    if (policy.local_memory_pattern.has_value()) {
+        if (!matchesMemoryPattern(policy.local_memory_pattern.value(),
+                                  context.local_memory_type)) {
+            return false;
+        }
+    }
+
+    // Check remote memory pattern
+    if (policy.remote_memory_pattern.has_value()) {
+        if (!matchesMemoryPattern(policy.remote_memory_pattern.value(),
+                                  context.remote_memory_type)) {
+            return false;
+        }
+    }
+
+    // Check size constraints
+    if (policy.min_size.has_value()) {
+        if (context.transfer_size < policy.min_size.value()) {
+            return false;
+        }
+    }
+    if (policy.max_size.has_value()) {
+        if (context.transfer_size > policy.max_size.value()) {
+            return false;
+        }
+    }
+
+    // Check priority constraint: exact match required
+    if (policy.priority.has_value()) {
+        if (context.priority_level != policy.priority.value()) {
+            return false;
+        }
+    }
+
+    return true;
+}
+
+bool TransportSelector::isTransportAvailable(
+    TransportType type, const SelectionContext& context,
+    const std::array<std::shared_ptr<Transport>, kSupportedTransportTypes>&
+        available_transports) const {
+    // Check if transport type is valid
+    if (type < 0 || type >= kSupportedTransportTypes) {
+        return false;
+    }
+
+    // Check if transport exists
+    const auto& transport = available_transports[type];
+    if (!transport) {
+        return false;
+    }
+
+    // Special constraints
+    if ((type == NVLINK || type == SHM) && !context.same_machine) {
+        return false;  // NVLINK and SHM only work on same machine
+    }
+
+    const auto& caps = transport->capabilities();
+
+    // Helper to check if memory type is GPU/NPU
+    auto is_gpu = [](MemoryType t) {
+        return t == MTYPE_CUDA || t == MTYPE_ROCM;
+    };
+
+    // For file segments, check file-specific capabilities (original logic)
+    if (context.segment_type == SegmentType::File) {
+        if (context.local_memory_type == MTYPE_CPU) return caps.dram_to_file;
+        if (is_gpu(context.local_memory_type)) return caps.gpu_to_file;
+        return false;
+    }
+
+    // Check transport capabilities for the memory type combination
+    MemoryType remote_mtype = context.remote_memory_type;
+
+    // CPU to CPU
+    if (context.local_memory_type == MTYPE_CPU && remote_mtype == MTYPE_CPU)
+        return caps.dram_to_dram;
+
+    // GPU/NPU to GPU/NPU (same type)
+    if (is_gpu(context.local_memory_type) && is_gpu(remote_mtype))
+        return caps.gpu_to_gpu;
+
+    // CPU to GPU/NPU
+    if (context.local_memory_type == MTYPE_CPU && is_gpu(remote_mtype))
+        return caps.dram_to_gpu;
+
+    // GPU/NPU to CPU
+    if (is_gpu(context.local_memory_type) && remote_mtype == MTYPE_CPU)
+        return caps.gpu_to_dram;
+
+    return false;
+}
+
+SelectionResult TransportSelector::select(
+    const SelectionContext& context,
+    const std::array<std::shared_ptr<Transport>, kSupportedTransportTypes>&
+        available_transports,
+    int transport_index, TransportType hint) {
+    SelectionResult result;
+
+    // Find the first matching policy (JSON order wins)
+    const SelectionPolicy* matching_policy = nullptr;
+    for (const auto& policy : policies_) {
+        if (matchesPolicy(policy, context)) {
+            matching_policy = &policy;
+            break;  // First match wins
+        }
+    }
+
+    if (!matching_policy) {
+        LOG(WARNING) << "No matching transport policy for segment_type="
+                     << (context.segment_type == SegmentType::File ? "file"
+                                                                   : "memory")
+                     << ", size=" << context.transfer_size
+                     << ", priority_level=" << context.priority_level;
+        return result;  // UNSPEC, all devices
+    }
+
+    // Convert device names to mask
+    result.device_mask = ~0ULL;  // Default: all devices
+    if (!matching_policy->devices.empty() && topology_) {
+        result.device_mask = 0;
+        for (const auto& name : matching_policy->devices) {
+            int dev_id = topology_->getNicId(name);
+            if (dev_id >= 0 && dev_id < 64) {
+                result.device_mask |= (1ULL << dev_id);
+            } else {
+                LOG(WARNING) << "RDMA device not found or ID >= 64: " << name;
+            }
+        }
+        if (result.device_mask == 0) {
+            result.device_mask = ~0ULL;  // Fallback to all if none found
+        }
+    }
+
+    // The "raw" candidate set is whatever the matching policy authorizes:
+    // the policy's explicit transports list, or the buffer's registered
+    // transports as a fallback.
+    static const std::vector<TransportType> kEmpty;
+    const auto& raw = !matching_policy->transports.empty()
+                          ? matching_policy->transports
+                      : context.buffer_transports ? *context.buffer_transports
+                                                  : kEmpty;
+
+    if (transport_index < 0) return result;
+    const int original_index = transport_index;
+
+    auto candidates = reorderWithHint(raw, hint);
+    if (!candidates) return result;  // UNSPEC -> task FAILED downstream
+
+    const bool has_hint = (hint != UNSPEC);
+    for (size_t i = 0; i < candidates->size(); ++i) {
+        TransportType type = (*candidates)[i];
+        if (!isTransportAvailable(type, context, available_transports)) {
+            // The pinned hint at slot 0 must be available; if it isn't,
+            // reject the request rather than falling through to another
+            // transport the caller did not ask for.
+            if (has_hint && i == 0) return result;
+            continue;
+        }
+        if (transport_index == 0) {
+            result.transport = type;
+            break;
+        }
+        --transport_index;
+    }
+    if (result.transport == UNSPEC) return result;
+    VLOG(1) << "Selected transport " << transportTypeName(result.transport)
+            << " for policy " << matching_policy->name << " at index "
+            << original_index << " (hint=" << transportTypeName(hint)
+            << "), device_mask=0x" << std::hex << result.device_mask
+            << std::dec;
+    return result;
+}
+
+std::optional<std::vector<TransportType>> TransportSelector::reorderWithHint(
+    const std::vector<TransportType>& raw, TransportType hint) {
+    if (hint == UNSPEC) return raw;
+    if (std::find(raw.begin(), raw.end(), hint) == raw.end()) {
+        return std::nullopt;
+    }
+    std::vector<TransportType> out;
+    out.reserve(raw.size());
+    out.push_back(hint);
+    for (TransportType t : raw) {
+        if (t != hint) out.push_back(t);
+    }
+    return out;
+}
+
+}  // namespace tent
+}  // namespace mooncake
diff --git a/mooncake-transfer-engine/tent/src/transfer_engine.cpp b/mooncake-transfer-engine/tent/src/transfer_engine.cpp
index 509ea195..5628a45c 100644
--- a/mooncake-transfer-engine/tent/src/transfer_engine.cpp
+++ b/mooncake-transfer-engine/tent/src/transfer_engine.cpp
@@ -169,5 +169,10 @@ Status TransferEngine::getTransferStatus(BatchID batch_id,
     return impl_->getTransferStatus(batch_id, overall_status);
 }
 
+Status TransferEngine::progressBatch(BatchID batch_id,
+                                     TransferStatus& overall_status) {
+    return impl_->progressBatch(batch_id, overall_status);
+}
+
 }  // namespace tent
 }  // namespace mooncake
diff --git a/mooncake-transfer-engine/tent/src/transfer_engine_c.cpp b/mooncake-transfer-engine/tent/src/transfer_engine_c.cpp
index bf8797f8..150df8b8 100644
--- a/mooncake-transfer-engine/tent/src/transfer_engine_c.cpp
+++ b/mooncake-transfer-engine/tent/src/transfer_engine_c.cpp
@@ -216,6 +216,9 @@ int tent_submit(tent_engine_t engine, tent_batch_id_t batch_id,
         req_list[index].target_id = entries[index].target_id;
         req_list[index].target_offset = entries[index].target_offset;
         req_list[index].length = entries[index].length;
+        req_list[index].priority = entries[index].priority;
+        req_list[index].transport_hint =
+            mooncake::tent::c_to_transport_hint(entries[index].transport_hint);
     }
     auto status = CAST(engine)->submitTransfer(batch_id, req_list);
     if (!status.ok()) {
@@ -241,6 +244,9 @@ int tent_submit_notif(tent_engine_t engine, tent_batch_id_t batch_id,
         req_list[index].target_id = entries[index].target_id;
         req_list[index].target_offset = entries[index].target_offset;
         req_list[index].length = entries[index].length;
+        req_list[index].priority = entries[index].priority;
+        req_list[index].transport_hint =
+            mooncake::tent::c_to_transport_hint(entries[index].transport_hint);
     }
     mooncake::tent::Notification notifi;
     notifi.name = name;
diff --git a/mooncake-transfer-engine/tent/src/transport/CMakeLists.txt b/mooncake-transfer-engine/tent/src/transport/CMakeLists.txt
index d6037183..e73f7cb0 100644
--- a/mooncake-transfer-engine/tent/src/transport/CMakeLists.txt
+++ b/mooncake-transfer-engine/tent/src/transport/CMakeLists.txt
@@ -7,18 +7,22 @@ add_subdirectory(gds)
 add_subdirectory(io_uring)
 add_subdirectory(bufio)
 add_subdirectory(ascend)
+add_subdirectory(sunrise_link)
 
 add_library(tent_transport_all INTERFACE)
-foreach(tgt tent_xport_gds
-            tent_xport_uring
-            tent_xport_bufio
-            tent_xport_mnnvl
-            tent_xport_nvlink
-            tent_xport_rdma
-            tent_xport_shm
-            tent_xport_tcp
-            tent_xport_ascend_direct)
-    if (TARGET ${tgt})
-        target_link_libraries(tent_transport_all INTERFACE ${tgt})
-    endif()
+foreach(
+  tgt
+  tent_xport_gds
+  tent_xport_uring
+  tent_xport_bufio
+  tent_xport_mnnvl
+  tent_xport_nvlink
+  tent_xport_rdma
+  tent_xport_shm
+  tent_xport_tcp
+  tent_xport_ascend_direct
+  tent_xport_sunrise_link)
+  if(TARGET ${tgt})
+    target_link_libraries(tent_transport_all INTERFACE ${tgt})
+  endif()
 endforeach()
diff --git a/mooncake-transfer-engine/tent/src/transport/ascend/ascend_direct_transport.cpp b/mooncake-transfer-engine/tent/src/transport/ascend/ascend_direct_transport.cpp
index dd18c4ee..d50fed7f 100644
--- a/mooncake-transfer-engine/tent/src/transport/ascend/ascend_direct_transport.cpp
+++ b/mooncake-transfer-engine/tent/src/transport/ascend/ascend_direct_transport.cpp
@@ -380,9 +380,9 @@ Status AscendDirectTransport::getTransferStatus(SubBatchRef batch, int task_id,
         return Status::InvalidArgument("Invalid task id" LOC_MARK);
     }
     auto &task = hixl_batch->task_list[task_id];
-    status = TransferStatus{task.status_word, task.transferred_bytes};
     if (task.status_word == TransferStatusEnum::PENDING) {
         if (task.req_handle == nullptr) {
+            status = TransferStatus{task.status_word, task.transferred_bytes};
             return Status::OK();
         }
         std::lock_guard<std::mutex> lock(req_mutex_);
@@ -396,6 +396,7 @@ Status AscendDirectTransport::getTransferStatus(SubBatchRef batch, int task_id,
             if (--req_map_[task.req_handle].second == 0) {
                 req_map_.erase(task.req_handle);
             }
+            status = TransferStatus{task.status_word, task.transferred_bytes};
             return Status::OK();
         }
         uint64_t current_ts = getCurrentTimeInNano();
@@ -406,6 +407,7 @@ Status AscendDirectTransport::getTransferStatus(SubBatchRef batch, int task_id,
                 req_map_[task.req_handle] =
                     std::make_pair(task.status_word, task.batch_size - 1);
             }
+            status = TransferStatus{task.status_word, task.transferred_bytes};
             return Status::OK();
         }
         hixl::TransferStatus xfer_status;
@@ -415,6 +417,7 @@ Status AscendDirectTransport::getTransferStatus(SubBatchRef batch, int task_id,
             xfer_status = hixl::TransferStatus::FAILED;
         }
         if (xfer_status == hixl::TransferStatus::WAITING) {
+            status = TransferStatus{task.status_word, task.transferred_bytes};
             return Status::OK();
         }
         if (xfer_status == hixl::TransferStatus::COMPLETED) {
@@ -430,6 +433,9 @@ Status AscendDirectTransport::getTransferStatus(SubBatchRef batch, int task_id,
         req_map_[task.req_handle] =
             std::make_pair(task.status_word, task.batch_size - 1);
     }
+    // Read status AFTER the poll so a just-observed completion/failure is
+    // reported on this call rather than one poll cycle late.
+    status = TransferStatus{task.status_word, task.transferred_bytes};
     return Status::OK();
 }
 
diff --git a/mooncake-transfer-engine/tent/src/transport/bufio/bufio_transport.cpp b/mooncake-transfer-engine/tent/src/transport/bufio/bufio_transport.cpp
index 5bd76f07..6b6ea252 100644
--- a/mooncake-transfer-engine/tent/src/transport/bufio/bufio_transport.cpp
+++ b/mooncake-transfer-engine/tent/src/transport/bufio/bufio_transport.cpp
@@ -80,7 +80,7 @@ Status BufIoTransport::install(std::string& local_segment_name,
 
     installed_ = true;
     caps.dram_to_file = true;
-    if (Platform::getLoader().type() == "cuda") {
+    if (Platform::getLoader().type() != "cpu") {
         caps.gpu_to_file = true;
     }
 
diff --git a/mooncake-transfer-engine/tent/src/transport/io_uring/io_uring_transport.cpp b/mooncake-transfer-engine/tent/src/transport/io_uring/io_uring_transport.cpp
index b594cc38..85e59515 100644
--- a/mooncake-transfer-engine/tent/src/transport/io_uring/io_uring_transport.cpp
+++ b/mooncake-transfer-engine/tent/src/transport/io_uring/io_uring_transport.cpp
@@ -87,7 +87,7 @@ Status IOUringTransport::install(std::string& local_segment_name,
     async_memcpy_threshold_ =
         conf_->get("transports/nvlink/async_memcpy_threshold", 1024) * 1024;
     caps.dram_to_file = true;
-    if (Platform::getLoader().type() == "cuda") {
+    if (Platform::getLoader().type() != "cpu") {
         caps.gpu_to_file = true;
     }
     return Status::OK();
@@ -117,14 +117,16 @@ Status IOUringTransport::allocateSubBatch(SubBatchRef& batch, size_t max_size) {
     auto io_uring_batch = Slab<IOUringSubBatch>::Get().allocate();
     if (!io_uring_batch)
         return Status::InternalError("Unable to allocate IO Uring sub-batch");
-    batch = io_uring_batch;
     io_uring_batch->max_size = max_size;
     io_uring_batch->task_list.reserve(max_size);
     int rc = io_uring_queue_init(max_size, &io_uring_batch->ring, 0);
-    if (rc)
+    if (rc) {
+        Slab<IOUringSubBatch>::Get().deallocate(io_uring_batch);
         return Status::InternalError(
             std::string("io_uring_queue_init failed: ") + strerror(-rc) +
             LOC_MARK);
+    }
+    batch = io_uring_batch;
     return Status::OK();
 }
 
@@ -265,6 +267,7 @@ Status IOUringTransport::getTransferStatus(SubBatchRef batch, int task_id,
             }
         }
         io_uring_cqe_seen(&io_uring_batch->ring, cqe);
+        batch->notifyProgress();
     }
     return Status::OK();
 }
diff --git a/mooncake-transfer-engine/tent/src/transport/mnnvl/mnnvl_transport.cpp b/mooncake-transfer-engine/tent/src/transport/mnnvl/mnnvl_transport.cpp
index e53df7de..75931789 100644
--- a/mooncake-transfer-engine/tent/src/transport/mnnvl/mnnvl_transport.cpp
+++ b/mooncake-transfer-engine/tent/src/transport/mnnvl/mnnvl_transport.cpp
@@ -114,7 +114,7 @@ Status MnnvlTransport::install(std::string &local_segment_name,
         128;
 
     caps.dram_to_gpu = true;
-    if (Platform::getLoader().type() == "cuda") caps.gpu_to_gpu = true;
+    if (Platform::getLoader().type() != "cpu") caps.gpu_to_gpu = true;
 
     installed_ = true;
     supported_ = true;
@@ -145,8 +145,8 @@ Status MnnvlTransport::allocateSubBatch(SubBatchRef &batch, size_t max_size) {
     batch = mnnvl_batch;
     mnnvl_batch->task_list.reserve(max_size);
     mnnvl_batch->max_size = max_size;
-    CHECK_STATUS(platform_->getStreamFromPool(mnnvl_batch->sync_stream));
-    CHECK_STATUS(platform_->getStreamFromPool(mnnvl_batch->async_stream));
+    // Streams are created lazily in submitTransferTasks where the correct
+    // GPU device can be inferred from request source pointers.
     return Status::OK();
 }
 
@@ -154,17 +154,43 @@ Status MnnvlTransport::freeSubBatch(SubBatchRef &batch) {
     auto mnnvl_batch = dynamic_cast<MnnvlSubBatch *>(batch);
     if (!mnnvl_batch)
         return Status::InvalidArgument("Invalid MNNVL sub-batch" LOC_MARK);
-    // CHECK_CUDA(cudaStreamDestroy(mnnvl_batch->stream));
+    // Completion events are destroyed by ~MnnvlSubBatch(), which Slab's
+    // deallocate() invokes before returning the storage. Streams are
+    // pool-managed and not destroyed here.
     Slab<MnnvlSubBatch>::Get().deallocate(mnnvl_batch);
     batch = nullptr;
     return Status::OK();
 }
 
+static int detectDeviceFromPointer(void *ptr) {
+    if (!ptr) return -1;
+    cudaPointerAttributes attrs = {};
+    auto err = cudaPointerGetAttributes(&attrs, ptr);
+    if (err == cudaSuccess && attrs.type == cudaMemoryTypeDevice)
+        return attrs.device;
+    cudaGetLastError();
+    return -1;
+}
+
 Status MnnvlTransport::submitTransferTasks(
     SubBatchRef batch, const std::vector<Request> &request_list) {
     auto mnnvl_batch = dynamic_cast<MnnvlSubBatch *>(batch);
     if (!mnnvl_batch)
         return Status::InvalidArgument("Invalid MNNVL sub-batch" LOC_MARK);
+
+    if (!mnnvl_batch->async_stream.get()) {
+        int device_id = -1;
+        for (auto &req : request_list) {
+            device_id = detectDeviceFromPointer(req.source);
+            if (device_id >= 0) break;
+        }
+        if (device_id < 0) cudaGetDevice(&device_id);
+        CHECK_STATUS(
+            platform_->getStreamFromPool(mnnvl_batch->sync_stream, device_id));
+        CHECK_STATUS(
+            platform_->getStreamFromPool(mnnvl_batch->async_stream, device_id));
+    }
+
     if (request_list.size() + mnnvl_batch->task_list.size() >
         mnnvl_batch->max_size)
         return Status::TooManyRequests("Exceed batch capacity" LOC_MARK);
@@ -260,8 +286,22 @@ void MnnvlTransport::startTransfer(std::vector<MnnvlTask *> &tasks,
     }
 
     cudaEvent_t event;
-    cudaEventCreateWithFlags(&event, cudaEventDisableTiming);
-    cudaEventRecord(event, batch->async_stream.get());
+    auto event_err = cudaEventCreateWithFlags(&event, cudaEventDisableTiming);
+    if (event_err != cudaSuccess) {
+        LOG(ERROR) << "MnnvlTransport: cudaEventCreateWithFlags failed: "
+                   << cudaGetErrorString(event_err);
+        for (auto *task : tasks) task->status_word = TransferStatusEnum::FAILED;
+        return;
+    }
+    auto record_err = cudaEventRecord(event, batch->async_stream.get());
+    if (record_err != cudaSuccess) {
+        LOG(ERROR) << "MnnvlTransport: cudaEventRecord failed: "
+                   << cudaGetErrorString(record_err);
+        cudaEventDestroy(event);
+        for (auto *task : tasks) task->status_word = TransferStatusEnum::FAILED;
+        return;
+    }
+    batch->completion_events.push_back(event);
     for (auto *task : tasks) task->completion_event = event;
 }
 
@@ -272,7 +312,6 @@ Status MnnvlTransport::getTransferStatus(SubBatchRef batch, int task_id,
         return Status::InvalidArgument("Invalid task id" LOC_MARK);
     }
     auto &task = mnnvl_batch->task_list[task_id];
-    status = TransferStatus{task.status_word, task.transferred_bytes};
     if (task.status_word == TransferStatusEnum::PENDING) {
         auto err = cudaEventQuery(task.completion_event);
         if (err == cudaSuccess) {
@@ -282,6 +321,9 @@ Status MnnvlTransport::getTransferStatus(SubBatchRef batch, int task_id,
             task.status_word = TransferStatusEnum::FAILED;
         }
     }
+    // Read status AFTER the poll so a just-observed completion/failure is
+    // reported on this call rather than one poll cycle late.
+    status = TransferStatus{task.status_word, task.transferred_bytes};
     return Status::OK();
 }
 
diff --git a/mooncake-transfer-engine/tent/src/transport/nvlink/nvlink_transport.cpp b/mooncake-transfer-engine/tent/src/transport/nvlink/nvlink_transport.cpp
index 8c16662b..4fecaa9e 100644
--- a/mooncake-transfer-engine/tent/src/transport/nvlink/nvlink_transport.cpp
+++ b/mooncake-transfer-engine/tent/src/transport/nvlink/nvlink_transport.cpp
@@ -33,6 +33,38 @@
 
 namespace mooncake {
 namespace tent {
+namespace {
+
+Status setCudaDeviceForLocation(const LocationParser& location,
+                                int& saved_dev) {
+    saved_dev = -1;
+    CHECK_CUDA(cudaGetDevice(&saved_dev));
+    if (location.index() >= 0 && saved_dev != location.index()) {
+        CHECK_CUDA(cudaSetDevice(location.index()));
+    }
+    return Status::OK();
+}
+
+Status restoreCudaDeviceForLocation(const LocationParser& location,
+                                    int saved_dev) {
+    if (saved_dev >= 0 && location.index() >= 0 &&
+        saved_dev != location.index()) {
+        CHECK_CUDA(cudaSetDevice(saved_dev));
+    }
+    return Status::OK();
+}
+
+int detectDeviceFromPointer(void* ptr) {
+    if (!ptr) return -1;
+    cudaPointerAttributes attrs = {};
+    auto err = cudaPointerGetAttributes(&attrs, ptr);
+    if (err == cudaSuccess && attrs.type == cudaMemoryTypeDevice)
+        return attrs.device;
+    cudaGetLastError();
+    return -1;
+}
+
+}  // namespace
 
 NVLinkTransport::NVLinkTransport() : installed_(false) {}
 
@@ -84,8 +116,8 @@ Status NVLinkTransport::allocateSubBatch(SubBatchRef& batch, size_t max_size) {
     batch = shm_batch;
     shm_batch->task_list.reserve(max_size);
     shm_batch->max_size = max_size;
-    CHECK_STATUS(platform_->getStreamFromPool(shm_batch->sync_stream));
-    CHECK_STATUS(platform_->getStreamFromPool(shm_batch->async_stream));
+    // Streams are created lazily in submitTransferTasks where the correct
+    // GPU device can be inferred from request source pointers.
     return Status::OK();
 }
 
@@ -93,7 +125,9 @@ Status NVLinkTransport::freeSubBatch(SubBatchRef& batch) {
     auto shm_batch = dynamic_cast<NVLinkSubBatch*>(batch);
     if (!shm_batch)
         return Status::InvalidArgument("Invalid NVLink sub-batch" LOC_MARK);
-    // CHECK_CUDA(cudaStreamDestroy(shm_batch->stream));
+    // Completion events are destroyed by ~NVLinkSubBatch(), which Slab's
+    // deallocate() invokes before returning the storage. Streams are
+    // pool-managed and not destroyed here.
     Slab<NVLinkSubBatch>::Get().deallocate(shm_batch);
     batch = nullptr;
     return Status::OK();
@@ -104,6 +138,20 @@ Status NVLinkTransport::submitTransferTasks(
     auto shm_batch = dynamic_cast<NVLinkSubBatch*>(batch);
     if (!shm_batch)
         return Status::InvalidArgument("Invalid NVLink sub-batch" LOC_MARK);
+
+    if (!shm_batch->async_stream.get()) {
+        int device_id = -1;
+        for (auto& req : request_list) {
+            device_id = detectDeviceFromPointer(req.source);
+            if (device_id >= 0) break;
+        }
+        if (device_id < 0) cudaGetDevice(&device_id);
+        CHECK_STATUS(
+            platform_->getStreamFromPool(shm_batch->sync_stream, device_id));
+        CHECK_STATUS(
+            platform_->getStreamFromPool(shm_batch->async_stream, device_id));
+    }
+
     if (request_list.size() + shm_batch->task_list.size() > shm_batch->max_size)
         return Status::TooManyRequests("Exceed batch capacity" LOC_MARK);
     std::vector<NVLinkTask*> new_tasks;
@@ -164,7 +212,8 @@ void NVLinkTransport::startTransfer(std::vector<NVLinkTask*>& tasks,
     err = cudaMemcpyBatchAsync(dsts.data(), srcs.data(), sizes.data(),
                                srcs.size(), &attr, &attrs_idx, 1, &fail_idx,
                                batch->async_stream.get());
-    if (err != cudaSuccess && fail_idx < tasks.size()) {
+    if (err != cudaSuccess && err != cudaErrorCallRequiresNewerDriver &&
+        fail_idx < tasks.size()) {
         LOG(ERROR) << "NVLinkTransport::startTransfer internal error: "
                    << "cudaMemcpyBatchAsync failed at task index " << fail_idx
                    << " (src=" << srcs[fail_idx] << ", dst=" << dsts[fail_idx]
@@ -173,17 +222,22 @@ void NVLinkTransport::startTransfer(std::vector<NVLinkTask*>& tasks,
         tasks[fail_idx]->status_word = TransferStatusEnum::FAILED;
     }
 #else
-    err = cudaSuccess;
-    for (size_t i = 0; i < tasks.size(); ++i) {
-        auto single_err =
-            cudaMemcpyAsync(dsts[i], srcs[i], sizes[i], cudaMemcpyDefault,
-                            batch->async_stream.get());
-        if (single_err != cudaSuccess) {
-            tasks[i]->status_word = TransferStatusEnum::FAILED;
-            err = single_err;
+    err = cudaErrorCallRequiresNewerDriver;
+#endif
+
+    if (err == cudaErrorCallRequiresNewerDriver) {
+        cudaGetLastError();
+        err = cudaSuccess;
+        for (size_t i = 0; i < tasks.size(); ++i) {
+            auto single_err =
+                cudaMemcpyAsync(dsts[i], srcs[i], sizes[i], cudaMemcpyDefault,
+                                batch->async_stream.get());
+            if (single_err != cudaSuccess) {
+                tasks[i]->status_word = TransferStatusEnum::FAILED;
+                err = single_err;
+            }
         }
     }
-#endif
 
     if (err != cudaSuccess) {
         for (auto* task : tasks) {
@@ -193,8 +247,22 @@ void NVLinkTransport::startTransfer(std::vector<NVLinkTask*>& tasks,
     }
 
     cudaEvent_t event;
-    cudaEventCreateWithFlags(&event, cudaEventDisableTiming);
-    cudaEventRecord(event, batch->async_stream.get());
+    auto event_err = cudaEventCreateWithFlags(&event, cudaEventDisableTiming);
+    if (event_err != cudaSuccess) {
+        LOG(ERROR) << "NVLinkTransport: cudaEventCreateWithFlags failed: "
+                   << cudaGetErrorString(event_err);
+        for (auto* task : tasks) task->status_word = TransferStatusEnum::FAILED;
+        return;
+    }
+    auto record_err = cudaEventRecord(event, batch->async_stream.get());
+    if (record_err != cudaSuccess) {
+        LOG(ERROR) << "NVLinkTransport: cudaEventRecord failed: "
+                   << cudaGetErrorString(record_err);
+        cudaEventDestroy(event);
+        for (auto* task : tasks) task->status_word = TransferStatusEnum::FAILED;
+        return;
+    }
+    batch->completion_events.push_back(event);
     for (auto* task : tasks) task->completion_event = event;
 }
 
@@ -205,7 +273,6 @@ Status NVLinkTransport::getTransferStatus(SubBatchRef batch, int task_id,
         return Status::InvalidArgument("Invalid task id" LOC_MARK);
     }
     auto& task = shm_batch->task_list[task_id];
-    status = TransferStatus{task.status_word, task.transferred_bytes};
     if (task.status_word == TransferStatusEnum::PENDING) {
         auto err = cudaEventQuery(task.completion_event);
         if (err == cudaSuccess) {
@@ -215,6 +282,9 @@ Status NVLinkTransport::getTransferStatus(SubBatchRef batch, int task_id,
             task.status_word = TransferStatusEnum::FAILED;
         }
     }
+    // Read status AFTER the poll so a just-observed completion/failure is
+    // reported on this call rather than one poll cycle late.
+    status = TransferStatus{task.status_word, task.transferred_bytes};
     return Status::OK();
 }
 
@@ -222,10 +292,25 @@ Status NVLinkTransport::addMemoryBuffer(BufferDesc& desc,
                                         const MemoryOptions& options) {
     LocationParser location(desc.location);
     if (location.type() == "cuda") {
-        // If the memory region is allocated using cuMemAlloc,
-        // we cannot use cudaIpcGetMemHandle, so skip it
+        // MNNVL allocations are exported by MnnvlTransport instead of CUDA IPC.
         if (options.type == MNNVL) return Status::OK();
 
+        int saved_dev = -1;
+        CHECK_STATUS(setCudaDeviceForLocation(location, saved_dev));
+
+        // VMM allocations have driver allocation handles, but
+        // cudaIpcGetMemHandle only supports cudaMalloc-backed pointers.
+        CUmemGenericAllocationHandle generic_handle;
+        CUresult retain_result =
+            cuMemRetainAllocationHandle(&generic_handle, (void*)desc.addr);
+        if (retain_result == CUDA_SUCCESS) {
+            cuMemRelease(generic_handle);
+            CHECK_STATUS(restoreCudaDeviceForLocation(location, saved_dev));
+            LOG(INFO) << "NVLinkTransport: memory region " << (void*)desc.addr
+                      << " is not cudaMalloc-backed; skip CUDA IPC export.";
+            return Status::OK();
+        }
+
         // Resolve the true cudaMalloc base address. Caching allocators
         // (e.g. PyTorch) sub-allocate tensors within larger cudaMalloc
         // segments. cudaIpcGetMemHandle returns a handle for the whole
@@ -238,6 +323,7 @@ Status NVLinkTransport::addMemoryBuffer(BufferDesc& desc,
             LOG(ERROR) << "NVLinkTransport: cuMemGetAddressRange failed for "
                        << "addr 0x" << std::hex << desc.addr << std::dec
                        << " (error " << cu_err << ")";
+            CHECK_STATUS(restoreCudaDeviceForLocation(location, saved_dev));
             return Status::InternalError(
                 "cuMemGetAddressRange failed" LOC_MARK);
         }
@@ -249,12 +335,24 @@ Status NVLinkTransport::addMemoryBuffer(BufferDesc& desc,
                 desc.addr = (uint64_t)base_ptr;
                 desc.length = alloc_size;
                 desc.transports.push_back(TransportType::NVLINK);
+                CHECK_STATUS(restoreCudaDeviceForLocation(location, saved_dev));
                 return Status::OK();
             }
         }
 
         cudaIpcMemHandle_t handle;
-        CHECK_CUDA(cudaIpcGetMemHandle(&handle, (void*)base_ptr));
+        auto cuda_err = cudaIpcGetMemHandle(&handle, (void*)base_ptr);
+        CHECK_STATUS(restoreCudaDeviceForLocation(location, saved_dev));
+        if (cuda_err != cudaSuccess) {
+            LOG(ERROR) << "NVLinkTransport: cudaIpcGetMemHandle failed for "
+                       << "addr 0x" << std::hex << desc.addr << ", base 0x"
+                       << (uint64_t)base_ptr << std::dec << ", device "
+                       << location.index() << ": "
+                       << cudaGetErrorString(cuda_err);
+            return Status::InternalError(
+                std::string("cudaIpcGetMemHandle(&handle, (void*)base_ptr): ") +
+                cudaGetErrorString(cuda_err) + LOC_MARK);
+        }
         desc.addr = (uint64_t)base_ptr;
         desc.length = alloc_size;
         desc.shm_path =
@@ -281,10 +379,14 @@ Status NVLinkTransport::removeMemoryBuffer(BufferDesc& desc) {
     if (location.type() == "cuda") {
         // Resolve base the same way we did in addMemoryBuffer, so we
         // remove the right entry even for sub-allocated addresses.
+        int saved_dev = -1;
+        CHECK_STATUS(setCudaDeviceForLocation(location, saved_dev));
+
         CUdeviceptr base_ptr = 0;
         size_t alloc_size = 0;
         CUresult cu_err = cuMemGetAddressRange(&base_ptr, &alloc_size,
                                                (CUdeviceptr)desc.addr);
+        CHECK_STATUS(restoreCudaDeviceForLocation(location, saved_dev));
 
         uint64_t key = desc.addr;
         if (cu_err == CUDA_SUCCESS) {
diff --git a/mooncake-transfer-engine/tent/src/transport/rdma/buffers.cpp b/mooncake-transfer-engine/tent/src/transport/rdma/buffers.cpp
index 20f1d216..0c350890 100644
--- a/mooncake-transfer-engine/tent/src/transport/rdma/buffers.cpp
+++ b/mooncake-transfer-engine/tent/src/transport/rdma/buffers.cpp
@@ -78,13 +78,17 @@ LocalBufferManager::LocalBufferManager() {}
 
 LocalBufferManager::~LocalBufferManager() { clear(); }
 
-static inline int getAccessFlags(Permission perm) {
+static inline int getAccessFlags(Permission perm,
+                                 bool relaxed_ordering = false) {
     int access = IBV_ACCESS_LOCAL_WRITE;
     if (perm == kGlobalReadWrite) {
         access |= IBV_ACCESS_REMOTE_WRITE | IBV_ACCESS_REMOTE_READ;
     } else if (perm == kGlobalReadOnly) {
         access |= IBV_ACCESS_REMOTE_READ;
     }
+    if (relaxed_ordering) {
+        access |= IBV_ACCESS_RELAXED_ORDERING;
+    }
     return access;
 }
 
@@ -98,7 +102,6 @@ Status LocalBufferManager::addBufferInternal(BufferDesc& desc,
                                              bool force_sequential) {
     AddressRange range((void*)desc.addr, desc.length);
     BufferEntryForRdma staging;
-    auto access = getAccessFlags(options.perm);
     assert(desc.rkey.empty());
     size_t context_count = 0;
     for (auto* context : context_list_) {
@@ -116,6 +119,9 @@ Status LocalBufferManager::addBufferInternal(BufferDesc& desc,
         for (size_t id = 0; id < context_list_.size(); ++id) {
             auto* context = context_list_[id];
             if (!context) continue;
+            // Calculate access flags per context (for relaxed ordering support)
+            int access = getAccessFlags(options.perm,
+                                        context->isRelaxedOrderingEnabled());
             tasks.emplace_back(std::async(
                 std::launch::async,
                 [context, &mem_reg_list, id, addr, length, access]() {
@@ -128,6 +134,9 @@ Status LocalBufferManager::addBufferInternal(BufferDesc& desc,
         for (size_t id = 0; id < context_list_.size(); ++id) {
             auto* context = context_list_[id];
             if (!context) continue;
+            // Calculate access flags per context (for relaxed ordering support)
+            int access = getAccessFlags(options.perm,
+                                        context->isRelaxedOrderingEnabled());
             mem_reg_list[id] =
                 context->registerMemReg((void*)desc.addr, desc.length, access);
         }
@@ -203,7 +212,8 @@ Status LocalBufferManager::addDevice(RdmaContext* context) {
     for (auto& buffer : buffer_list_) {
         auto range = buffer.first;
         auto& options = buffer.second.options;
-        auto access = getAccessFlags(options.perm);
+        auto access =
+            getAccessFlags(options.perm, context->isRelaxedOrderingEnabled());
         if (buffer.second.mem_reg_map.count(context)) continue;
         auto mem_reg =
             context->registerMemReg(range.addr, range.length, access);
diff --git a/mooncake-transfer-engine/tent/src/transport/rdma/context.cpp b/mooncake-transfer-engine/tent/src/transport/rdma/context.cpp
index a012bc23..31ea13d4 100644
--- a/mooncake-transfer-engine/tent/src/transport/rdma/context.cpp
+++ b/mooncake-transfer-engine/tent/src/transport/rdma/context.cpp
@@ -26,6 +26,13 @@
 #include <fstream>
 #include <memory>
 #include <thread>
+#include <cstdlib>
+#include <dlfcn.h>
+
+#ifdef USE_CUDA
+#include <cuda.h>
+#include <cuda_runtime.h>
+#endif
 
 #include "tent/common/status.h"
 #include "tent/transport/rdma/endpoint_store.h"
@@ -381,6 +388,31 @@ int RdmaContext::enable() {
         return notify_ret;
     }
 
+    // Check PCIe Relaxed Ordering support from config
+    // Mode: 0 = disabled, 1 = enabled if supported, 2 = auto (default)
+    auto mode =
+        transport_.config()->get("transports/rdma/pci_relaxed_ordering", 1);
+    if (mode != 0) {
+        // Check if ibv_reg_mr_iova2 symbol is available (IBVERBS_1.8+)
+        void* sym = dlsym(RTLD_DEFAULT, "ibv_reg_mr_iova2");
+        if (sym) {
+            relaxed_ordering_enabled_ = true;
+            LOG(INFO) << "[RDMA] Relaxed ordering is supported and enabled for "
+                      << device_name_;
+        } else {
+            if (mode == 1) {
+                LOG(WARNING) << "[RDMA] Relaxed ordering requested but NOT "
+                             << "supported (ibv_reg_mr_iova2 missing). "
+                             << "Falling back to strict ordering.";
+            }
+            relaxed_ordering_enabled_ = false;
+        }
+    } else {
+        LOG(INFO) << "[RDMA] Relaxed ordering disabled via config for "
+                  << device_name_;
+        relaxed_ordering_enabled_ = false;
+    }
+
     ibv_port_attr port_attr;
     int ret = verbs_.ibv_query_port_default(native_context_,
                                             params_->device.port, &port_attr);
@@ -473,11 +505,71 @@ RdmaContext::MemReg RdmaContext::registerMemReg(void* addr, size_t length,
         LOG(FATAL) << "RDMA context " << name() << " not constructed";
         return nullptr;
     }
+
+#ifdef USE_CUDA
+    // Ensure CUDA context is current for GPU memory registration
+    // This is needed for worker threads or callers from non-CUDA threads
+    CUmemorytype memType;
+    CUresult result = cuPointerGetAttribute(
+        &memType, CU_POINTER_ATTRIBUTE_MEMORY_TYPE, (CUdeviceptr)addr);
+
+    if (result == CUDA_SUCCESS && memType == CU_MEMORYTYPE_DEVICE) {
+        // Get device ordinal and set primary context current
+        unsigned int devOrd = 0;
+        result = cuPointerGetAttribute(
+            &devOrd, CU_POINTER_ATTRIBUTE_DEVICE_ORDINAL, (CUdeviceptr)addr);
+        if (result != CUDA_SUCCESS) {
+            LOG(ERROR) << "Failed to get CUDA device ordinal: " << result;
+            return nullptr;
+        }
+
+        CUdevice cuDev;
+        result = cuDeviceGet(&cuDev, devOrd);
+        if (result != CUDA_SUCCESS) {
+            LOG(ERROR) << "Failed to get CUDA device: " << result;
+            return nullptr;
+        }
+
+        CUcontext cuCtx;
+        result = cuDevicePrimaryCtxRetain(&cuCtx, cuDev);
+        if (result != CUDA_SUCCESS) {
+            LOG(ERROR) << "Failed to retain CUDA primary context: " << result;
+            return nullptr;
+        }
+
+        result = cuCtxSetCurrent(cuCtx);
+        if (result != CUDA_SUCCESS) {
+            LOG(ERROR) << "Failed to set CUDA context current: " << result;
+            cuDevicePrimaryCtxRelease(cuDev);
+            return nullptr;
+        }
+
+        // Register GPU memory
+        ibv_mr* entry =
+            verbs_.ibv_reg_mr_default(native_pd_, addr, length, access);
+
+        // Release primary context reference
+        cuDevicePrimaryCtxRelease(cuDev);
+
+        if (!entry) {
+            const void* end = static_cast<const char*>(addr) + length;
+            PLOG(ERROR) << "Failed to register GPU memory from " << addr
+                        << " to " << end << " in RDMA device " << device_name_;
+            return nullptr;
+        }
+        mr_set_mutex_.lock();
+        mr_set_.insert(entry);
+        mr_set_mutex_.unlock();
+        return entry;
+    }
+#endif
+
+    // Standard CPU memory registration
     ibv_mr* entry = verbs_.ibv_reg_mr_default(native_pd_, addr, length, access);
     if (!entry) {
+        const void* end = static_cast<const char*>(addr) + length;
         PLOG(ERROR) << "Failed to register memory from " << addr << " to "
-                    << (char*)addr + length << " in RDMA device "
-                    << device_name_;
+                    << end << " in RDMA device " << device_name_;
         return nullptr;
     }
     mr_set_mutex_.lock();
@@ -491,8 +583,10 @@ int RdmaContext::warmupMrRegistration(void* addr, size_t length) {
         LOG(FATAL) << "RDMA context " << name() << " not constructed";
         return -1;
     }
-    ibv_mr* entry = verbs_.ibv_reg_mr_default(native_pd_, addr, length,
-                                              IBV_ACCESS_LOCAL_WRITE);
+    int access_flags = IBV_ACCESS_LOCAL_WRITE;
+    if (relaxed_ordering_enabled_) access_flags |= IBV_ACCESS_RELAXED_ORDERING;
+    ibv_mr* entry =
+        verbs_.ibv_reg_mr_default(native_pd_, addr, length, access_flags);
     if (!entry) {
         PLOG(WARNING) << "ibv_reg_mr warm-up failed on " << device_name_
                       << " for [" << addr << ", " << length << " bytes]";
@@ -519,9 +613,9 @@ int RdmaContext::unregisterMemReg(MemReg id) {
     mr_set_mutex_.unlock();
 
     if (verbs_.ibv_dereg_mr(entry)) {
+        const void* end = static_cast<const char*>(entry->addr) + entry->length;
         LOG(ERROR) << "Failed to unregister memory from " << entry->addr
-                   << " to " << (char*)entry->addr + entry->length
-                   << " in RDMA device " << device_name_;
+                   << " to " << end << " in RDMA device " << device_name_;
     }
 
     return 0;
diff --git a/mooncake-transfer-engine/tent/src/transport/rdma/endpoint.cpp b/mooncake-transfer-engine/tent/src/transport/rdma/endpoint.cpp
index a0326de2..fe94a570 100644
--- a/mooncake-transfer-engine/tent/src/transport/rdma/endpoint.cpp
+++ b/mooncake-transfer-engine/tent/src/transport/rdma/endpoint.cpp
@@ -26,6 +26,7 @@
 #include "tent/common/status.h"
 #include "tent/common/types.h"
 #include "tent/transport/rdma/context.h"
+#include "tent/transport/rdma/endpoint_store.h"
 #include "tent/common/utils/os.h"
 #include "tent/common/utils/string_builder.h"
 #include "tent/thirdparty/nlohmann/json.h"
@@ -41,8 +42,10 @@ static inline const std::string statusToString(
             return "EP_HANDSHAKING";
         case RdmaEndPoint::EP_READY:
             return "EP_READY";
-        case RdmaEndPoint::EP_RESET:
-            return "EP_RESET";
+        case RdmaEndPoint::EP_DESTROYING:
+            return "EP_DESTROYING";
+        case RdmaEndPoint::EP_DESTROYED:
+            return "EP_DESTROYED";
     }
     return "UNKNOWN";
 }
@@ -50,12 +53,15 @@ static inline const std::string statusToString(
 // Forward declaration for notification QP setup
 static int setupNotifyQpConnection(ibv_qp* qp, RdmaContext* ctx,
                                    const std::string& peer_gid_str,
-                                   uint16_t peer_lid, uint32_t peer_qp_num);
+                                   uint16_t peer_lid, uint32_t peer_qp_num,
+                                   uint16_t pkey_index,
+                                   uint8_t service_level = 0,
+                                   uint8_t traffic_class = 0);
 
 RdmaEndPoint::RdmaEndPoint() : status_(EP_UNINIT) {}
 
 RdmaEndPoint::~RdmaEndPoint() {
-    if (status_ != EP_UNINIT) deconstruct();
+    if (status_.load(std::memory_order_relaxed) != EP_UNINIT) deconstruct();
     if (endpoints_count_)
         endpoints_count_->fetch_sub(1, std::memory_order_relaxed);
 }
@@ -118,7 +124,7 @@ int RdmaEndPoint::construct(RdmaContext* context, EndPointParams* params,
     ibv_qp_attr qp_attr;
     memset(&qp_attr, 0, sizeof(qp_attr));
     qp_attr.qp_state = IBV_QPS_INIT;
-    qp_attr.pkey_index = 0;
+    qp_attr.pkey_index = params_->pkey_index;
     qp_attr.port_num = context_->portNum();
     qp_attr.qp_access_flags = IBV_ACCESS_REMOTE_WRITE;
 
@@ -160,7 +166,7 @@ int RdmaEndPoint::construct(RdmaContext* context, EndPointParams* params,
         }
     }
 
-    status_ = EP_HANDSHAKING;
+    status_.store(EP_HANDSHAKING, std::memory_order_relaxed);
     return 0;
 }
 
@@ -170,8 +176,10 @@ int RdmaEndPoint::deconstruct() {
 }
 
 int RdmaEndPoint::deconstructUnlocked() {
-    if (status_ == EP_UNINIT) return 0;
-    status_ = EP_RESET;
+    auto current_status = status_.load(std::memory_order_relaxed);
+    // Idempotent: if already destroyed or never initialized, skip cleanup
+    if (current_status == EP_DESTROYED || current_status == EP_UNINIT) return 0;
+    status_.store(EP_DESTROYED, std::memory_order_relaxed);
     resetInflightSlices();
     peer_qp_num_list_.clear();
 
@@ -214,10 +222,98 @@ int RdmaEndPoint::deconstructUnlocked() {
     wr_depth_list_ = nullptr;
     peer_server_name_.clear();
     peer_nic_name_.clear();
-    status_ = EP_UNINIT;
+    // Status remains EP_DESTROYED (unidirectional lifecycle)
     return 0;
 }
 
+void RdmaEndPoint::beginDestroy() {
+    RWSpinlock::WriteGuard guard(lock_);
+    beginDestroyNoLock();
+}
+
+void RdmaEndPoint::beginDestroyNoLock() {
+    auto current_status = status_.load(std::memory_order_relaxed);
+    if (current_status == EP_DESTROYING || current_status == EP_DESTROYED)
+        return;
+
+    destroy_start_time_ = getCurrentTimeInNano();
+    status_.store(EP_DESTROYING, std::memory_order_release);
+
+    // Transition QPs to ERR state so hardware flushes inflight WRs to CQ
+    ibv_qp_attr attr;
+    memset(&attr, 0, sizeof(attr));
+    attr.qp_state = IBV_QPS_ERR;
+
+    for (size_t i = 0; i < qp_list_.size(); ++i) {
+        int ret =
+            context_->verbs_.ibv_modify_qp(qp_list_[i], &attr, IBV_QP_STATE);
+        if (ret) {
+            PLOG(ERROR) << "Failed to modify QP to ERR in beginDestroy";
+        }
+    }
+}
+
+bool RdmaEndPoint::finishDestroy() {
+    RWSpinlock::WriteGuard guard(lock_);
+    auto current_status = status_.load(std::memory_order_relaxed);
+
+    // Gate 1: already done
+    if (current_status == EP_DESTROYED) return true;
+
+    // Gate 2: non-two-phase path. Endpoint reached waiting_list_ without
+    // going through beginDestroy(). This handles edge cases and serves as
+    // a safety net. Endpoints that never reached construct() own no RDMA
+    // resources; drop them directly.
+    if (current_status != EP_DESTROYING) {
+        if (qp_list_.empty()) {
+            status_.store(EP_DESTROYED, std::memory_order_relaxed);
+            return true;
+        }
+        LOG(WARNING) << "finishDestroy called in unexpected state: "
+                     << statusToString(current_status)
+                     << ", forcing destruction to avoid waiting_list_ leak";
+        // Fall through to the unified destroy path
+    } else {
+        // Gate 3: two-phase path. Wait for inflight WRs to drain via CQ
+        // polling. If ibv_modify_qp-to-ERR failed in beginDestroy, WRs may
+        // never be flushed; enforce a timeout to avoid leaking forever.
+        bool has_outstanding = false;
+        for (size_t i = 0; i < qp_list_.size(); ++i) {
+            if (wr_depth_list_[i].value != 0) {
+                has_outstanding = true;
+                break;
+            }
+        }
+        if (has_outstanding) {
+            double elapsed =
+                (getCurrentTimeInNano() - destroy_start_time_) / 1e9;
+            if (elapsed < kFinishDestroyTimeoutSec) {
+                return false;  // Still waiting for WRs to drain
+            }
+            LOG(WARNING) << "finishDestroy timed out after " << elapsed
+                         << "s with outstanding WRs, forcing destruction";
+        }
+    }
+
+    // Unified destroy: tear down QPs and bound retries to avoid
+    // log flooding when ibv_destroy_qp fails permanently.
+    int ret = deconstructUnlocked();
+    if (ret) {
+        finish_destroy_retries_++;
+        LOG(ERROR) << "Failed to finish destroying endpoint (attempt "
+                   << finish_destroy_retries_ << "/" << kFinishDestroyMaxRetries
+                   << "): " << ret;
+        if (finish_destroy_retries_ < kFinishDestroyMaxRetries) {
+            return false;  // Retry later
+        }
+        LOG(ERROR) << "Giving up after " << finish_destroy_retries_
+                   << " retries (possible resource leak)";
+    }
+
+    status_.store(EP_DESTROYED, std::memory_order_relaxed);
+    return true;
+}
+
 Status RdmaEndPoint::connect(const std::string& peer_server_name,
                              const std::string& peer_nic_name,
                              const std::string& peer_rpc_server_addr) {
@@ -234,18 +330,19 @@ Status RdmaEndPoint::connect(const std::string& peer_server_name,
     {
         RWSpinlock::WriteGuard guard(lock_);
         if (peer_server_name.empty() || peer_nic_name.empty()) {
-            return Status::InvalidArgument("Invalid peer path" LOC_MARK);
+            return mooncake::tent::Status::InvalidArgument(
+                "Invalid peer path" LOC_MARK);
         }
-        if (status_ == EP_READY) {
-            return Status::OK();
+        if (status_.load(std::memory_order_relaxed) == EP_READY) {
+            return mooncake::tent::Status::OK();
         }
-        if (status_ != EP_HANDSHAKING) {
-            return Status::InvalidArgument(
+        if (status_.load(std::memory_order_relaxed) != EP_HANDSHAKING) {
+            return mooncake::tent::Status::InvalidArgument(
                 "Endpoint not in handshaking state" LOC_MARK);
         }
 
         local_desc.local_nic_path =
-            MakeNicPath(transport.local_segment_name_, context_->name());
+            MakeNicPath(transport.rdma_server_name_, context_->name());
         local_desc.peer_nic_path = MakeNicPath(peer_server_name, peer_nic_name);
         qp_num = qpNum();
         local_desc.qp_num = qp_num;
@@ -255,7 +352,8 @@ Status RdmaEndPoint::connect(const std::string& peer_server_name,
 
         same_nic = (local_desc.local_nic_path == local_desc.peer_nic_path);
         is_self =
-            (!same_nic && peer_server_name == transport.local_segment_name_);
+            (!same_nic && (peer_server_name == transport.local_segment_name_ ||
+                           peer_server_name == transport.rdma_server_name_));
     }
 
     // ===== Phase 2: Bootstrap (unlocked) =====
@@ -270,7 +368,7 @@ Status RdmaEndPoint::connect(const std::string& peer_server_name,
         // Same server, different NIC: call handler directly, bypass RPC
         int rc = transport.onSetupRdmaConnections(local_desc, peer_desc);
         if (rc != 0) {
-            return Status::InternalError(
+            return mooncake::tent::Status::InternalError(
                 "Local bootstrap failed: " + peer_desc.reply_msg + LOC_MARK);
         }
         qp_num = peer_desc.qp_num;
@@ -282,22 +380,24 @@ Status RdmaEndPoint::connect(const std::string& peer_server_name,
         if (rpc_server_addr.empty()) {
             SegmentDescRef segment_desc;
             auto& manager = transport.metadata_->segmentManager();
-            CHECK_STATUS(manager.getRemote(segment_desc, peer_server_name));
+            auto status = manager.getRemote(segment_desc, peer_server_name);
+            if (!status.ok()) return status;
             rpc_server_addr = segment_desc->rpc_server_addr;
         }
         if (rpc_server_addr.empty()) {
-            return Status::InvalidArgument(
+            return mooncake::tent::Status::InvalidArgument(
                 "Missing peer RPC server address" LOC_MARK);
         }
-        CHECK_STATUS(
-            ControlClient::bootstrap(rpc_server_addr, local_desc, peer_desc));
+        auto bootstrap_status =
+            ControlClient::bootstrap(rpc_server_addr, local_desc, peer_desc);
+        if (!bootstrap_status.ok()) return bootstrap_status;
         qp_num = peer_desc.qp_num;
         peer_gid = peer_desc.local_gid;
         peer_lid = peer_desc.local_lid;
     }
 
     if (peer_gid.empty()) {
-        return Status::InvalidArgument(
+        return mooncake::tent::Status::InvalidArgument(
             "Missing peer GID in bootstrap" LOC_MARK);
     }
 
@@ -306,11 +406,11 @@ Status RdmaEndPoint::connect(const std::string& peer_server_name,
     // connection while we were doing bootstrap without the lock.
     {
         RWSpinlock::WriteGuard guard(lock_);
-        if (status_ == EP_READY) {
-            return Status::OK();
+        if (status_.load(std::memory_order_relaxed) == EP_READY) {
+            return mooncake::tent::Status::OK();
         }
-        if (status_ != EP_HANDSHAKING) {
-            return Status::InvalidArgument(
+        if (status_.load(std::memory_order_relaxed) != EP_HANDSHAKING) {
+            return mooncake::tent::Status::InvalidArgument(
                 "Endpoint state changed during bootstrap" LOC_MARK);
         }
 
@@ -319,15 +419,17 @@ Status RdmaEndPoint::connect(const std::string& peer_server_name,
         peer_nic_name_ = peer_nic_name;
         int rc = setupAllQPs(peer_gid, peer_lid, qp_num);
         if (rc) {
-            return Status::InternalError(
+            return mooncake::tent::Status::InternalError(
                 "Failed to configure RDMA endpoint" LOC_MARK);
         }
         peer_qp_num_list_ = qp_num;
 
         // Setup notification QP connection if peer supports it
         if (peer_desc.notify_qp_num != 0 && notify_qp_) {
-            rc = setupNotifyQpConnection(notify_qp_, context_, peer_gid,
-                                         peer_lid, peer_desc.notify_qp_num);
+            rc = setupNotifyQpConnection(
+                notify_qp_, context_, peer_gid, peer_lid,
+                peer_desc.notify_qp_num, params_->pkey_index,
+                params_->service_level, params_->traffic_class);
             if (rc) {
                 LOG(WARNING)
                     << "Failed to setup notification QP, notification disabled";
@@ -340,13 +442,13 @@ Status RdmaEndPoint::connect(const std::string& peer_server_name,
         }
     }
 
-    return Status::OK();
+    return mooncake::tent::Status::OK();
 }
 
 Status RdmaEndPoint::accept(const BootstrapDesc& peer_desc,
                             BootstrapDesc& local_desc) {
     RWSpinlock::WriteGuard guard(lock_);
-    if (status_ == EP_READY) {
+    if (status_.load(std::memory_order_relaxed) == EP_READY) {
         // Idempotency check: if the incoming bootstrap is from the same peer
         // and carries the same peer QP list as the established connection,
         // this is a duplicate request (e.g. from a concurrent connect() on
@@ -365,25 +467,28 @@ Status RdmaEndPoint::accept(const BootstrapDesc& peer_desc,
                       << " (duplicate bootstrap, reusing connection)";
             auto& transport = context_->transport_;
             local_desc.local_nic_path =
-                MakeNicPath(transport.local_segment_name_, context_->name());
+                MakeNicPath(transport.rdma_server_name_, context_->name());
             local_desc.peer_nic_path = peer_desc.local_nic_path;
             local_desc.qp_num = qpNum();
             local_desc.local_lid = context_->lid();
             local_desc.local_gid = context_->gid();
             local_desc.notify_qp_num = notifyQpNum();
-            return Status::OK();
+            return mooncake::tent::Status::OK();
         }
-        LOG(WARNING) << "Endpoint is established with " << peer_nic_name_
-                     << " of " << peer_server_name_ << ", resetting first";
-        resetUnlocked();
-        status_ = EP_HANDSHAKING;
-    } else if (status_ == EP_RESET) {
-        resetUnlocked();
-        status_ = EP_HANDSHAKING;
-    } else if (status_ != EP_HANDSHAKING) {
+        // Endpoint already connected to a different peer - reject the request
+        // instead of resetting. Endpoints have unidirectional lifecycle and
+        // are never reset or reused. The caller should create a new endpoint.
+        LOG(ERROR)
+            << "Endpoint already established with " << peer_nic_name_ << " of "
+            << peer_server_name_
+            << ", cannot accept new connection (unidirectional lifecycle)";
+        return mooncake::tent::Status::InternalError(
+            "Endpoint already connected to different peer" LOC_MARK);
+    }
+    if (status_.load(std::memory_order_relaxed) != EP_HANDSHAKING) {
         LOG(ERROR) << "Endpoint not in handshaking state: "
-                   << statusToString(status_);
-        return Status::InvalidArgument(
+                   << statusToString(status_.load(std::memory_order_relaxed));
+        return mooncake::tent::Status::InvalidArgument(
             "Endpoint not in handshaking state" LOC_MARK);
     }
     auto& transport = context_->transport_;
@@ -391,16 +496,17 @@ Status RdmaEndPoint::accept(const BootstrapDesc& peer_desc,
     auto peer_server_name = getServerNameFromNicPath(peer_nic_path);
     auto peer_nic_name = getNicNameFromNicPath(peer_nic_path);
     if (peer_server_name.empty() || peer_nic_name.empty())
-        return Status::InvalidArgument("Invalid peer path" LOC_MARK);
+        return mooncake::tent::Status::InvalidArgument(
+            "Invalid peer path" LOC_MARK);
     local_desc.local_nic_path =
-        MakeNicPath(transport.local_segment_name_, context_->name());
+        MakeNicPath(transport.rdma_server_name_, context_->name());
     local_desc.peer_nic_path = peer_nic_path;
     local_desc.qp_num = qpNum();
     local_desc.local_lid = context_->lid();
     local_desc.local_gid = context_->gid();
     local_desc.notify_qp_num = notifyQpNum();  // Pass notification QP number
     if (peer_desc.local_gid.empty()) {
-        return Status::InvalidArgument(
+        return mooncake::tent::Status::InvalidArgument(
             "Missing peer GID in bootstrap" LOC_MARK);
     }
     peer_server_name_ = peer_server_name;
@@ -408,16 +514,17 @@ Status RdmaEndPoint::accept(const BootstrapDesc& peer_desc,
     int rc =
         setupAllQPs(peer_desc.local_gid, peer_desc.local_lid, peer_desc.qp_num);
     if (rc) {
-        return Status::InternalError(
+        return mooncake::tent::Status::InternalError(
             "Failed to configure RDMA endpoint" LOC_MARK);
     }
     peer_qp_num_list_ = peer_desc.qp_num;
 
     // Setup notification QP connection if peer supports it
     if (peer_desc.notify_qp_num != 0 && notify_qp_) {
-        rc = setupNotifyQpConnection(notify_qp_, context_, peer_desc.local_gid,
-                                     peer_desc.local_lid,
-                                     peer_desc.notify_qp_num);
+        rc = setupNotifyQpConnection(
+            notify_qp_, context_, peer_desc.local_gid, peer_desc.local_lid,
+            peer_desc.notify_qp_num, params_->pkey_index,
+            params_->service_level, params_->traffic_class);
         if (rc) {
             notify_connected_ = false;
         } else {
@@ -427,54 +534,40 @@ Status RdmaEndPoint::accept(const BootstrapDesc& peer_desc,
         }
     }
 
-    return Status::OK();
+    return mooncake::tent::Status::OK();
 }
 
-int RdmaEndPoint::reset() {
-    RWSpinlock::WriteGuard guard(lock_);
-    return resetUnlocked();
-}
+int RdmaEndPoint::resetConnection(const std::string& reason) {
+    // Mark endpoint for destruction due to failure or error.
+    // Endpoints have unidirectional lifecycle - once marked for destruction,
+    // they are never reused. The endpoint store will create a new endpoint
+    // for future connections.
+    RdmaEndPoint* endpoint_ptr = this;
 
-int RdmaEndPoint::resetUnlocked() {
-    if (status_ != EP_READY) return 0;
-    status_ = EP_RESET;
-    peer_qp_num_list_.clear();
-    resetInflightSlices();
+    {
+        RWSpinlock::WriteGuard guard(lock_);
+        auto curr_status = status_.load(std::memory_order_acquire);
+        if (curr_status == EP_DESTROYING || curr_status == EP_DESTROYED)
+            return 0;
+        if (curr_status != EP_HANDSHAKING && curr_status != EP_READY) return 0;
 
-    if (notify_qp_) {
-        // Keep the QP registered in the transport map — the qp_num doesn't
-        // change across resets, and the notify worker may still poll
-        // completions (flush errors) from the CQ between reset and reconnect.
-        // Unregistering here would cause "unknown QP" warnings.
-        notify_connected_ = false;
-        {
-            std::lock_guard<std::mutex> lock(notify_send_mutex_);
-            notify_pending_count_ = 0;
-        }
-        notify_send_cv_.notify_all();
+        destroy_start_time_ = getCurrentTimeInNano();
+        status_.store(EP_DESTROYING, std::memory_order_release);
+        LOG(INFO) << "Endpoint marked for destruction: " << reason;
     }
 
-    ibv_qp_attr attr;
-    memset(&attr, 0, sizeof(attr));
-    attr.qp_state = IBV_QPS_RESET;
-    for (size_t i = 0; i < qp_list_.size(); ++i) {
-        int ret =
-            context_->verbs_.ibv_modify_qp(qp_list_[i], &attr, IBV_QP_STATE);
-        if (ret) {
-            PLOG(ERROR) << "ibv_modify_qp(RESET)";
-            deconstructUnlocked();
-            return -1;
-        }
-        cancelQuota(i, wr_depth_list_[i].value);
-    }
+    // Delete from endpoint store so endpoint() won't return this endpoint.
+    // remove() calls beginDestroyNoLock() to avoid deadlocking when
+    // caller already holds lock_.
+    context_->endpointStore()->remove(endpoint_ptr);
     return 0;
 }
 
 int RdmaEndPoint::setupAllQPs(const std::string& peer_gid, uint16_t peer_lid,
                               std::vector<uint32_t> peer_qp_num_list,
                               std::string* reply_msg) {
-    if (status_ == EP_READY) {
-        status_ = EP_RESET;
+    if (status_.load(std::memory_order_relaxed) == EP_READY) {
+        status_.store(EP_DESTROYING, std::memory_order_relaxed);
         return -1;
     }
 
@@ -485,7 +578,7 @@ int RdmaEndPoint::setupAllQPs(const std::string& peer_gid, uint16_t peer_lid,
            << peer_nic_name_ << " of " << peer_server_name_;
         LOG(ERROR) << ss.str();
         if (reply_msg) *reply_msg = ss.str();
-        status_ = EP_RESET;
+        status_.store(EP_DESTROYING, std::memory_order_relaxed);
         return -1;
     }
 
@@ -493,12 +586,12 @@ int RdmaEndPoint::setupAllQPs(const std::string& peer_gid, uint16_t peer_lid,
         int ret = setupOneQP(qp_index, peer_gid, peer_lid,
                              peer_qp_num_list[qp_index], reply_msg);
         if (ret) {
-            status_ = EP_RESET;
+            status_.store(EP_DESTROYING, std::memory_order_relaxed);
             return ret;
         }
     }
 
-    status_ = EP_READY;
+    status_.store(EP_READY, std::memory_order_relaxed);
     return 0;
 }
 
@@ -520,7 +613,8 @@ int RdmaEndPoint::submitSlices(std::vector<RdmaSlice*>& slice_list,
     if (qp_list_.empty()) return 0;
     if (qp_index < 0) qp_index = 0;
     qp_index %= qp_list_.size();
-    if (status_ != EP_READY) return 0;
+    // Check endpoint status before submitting
+    if (status_.load(std::memory_order_relaxed) != EP_READY) return 0;
     auto cq = context_->cq(qp_index % context_->cqCount());
     int wr_count =
         std::min(cq->maxCqe() - cq->getQuota(),
@@ -586,7 +680,8 @@ int RdmaEndPoint::submitSlices(std::vector<RdmaSlice*>& slice_list,
 
 int RdmaEndPoint::submitRecvImmDataRequest(int qp_index, uint64_t id) {
     RWSpinlock::ReadGuard guard(lock_);
-    if (status_ != EP_READY) return 0;
+    // Check endpoint status before submitting
+    if (status_.load(std::memory_order_relaxed) != EP_READY) return 0;
     if (qp_index < 0 || qp_index >= (int)qp_list_.size()) return 0;
     ibv_recv_wr wr, *bad_wr;
     memset(&wr, 0, sizeof(ibv_recv_wr));
@@ -788,7 +883,9 @@ void RdmaEndPoint::repostAllNotifyRecvs() {
 
 static int setupNotifyQpConnection(ibv_qp* qp, RdmaContext* ctx,
                                    const std::string& peer_gid_str,
-                                   uint16_t peer_lid, uint32_t peer_qp_num) {
+                                   uint16_t peer_lid, uint32_t peer_qp_num,
+                                   uint16_t pkey_index, uint8_t service_level,
+                                   uint8_t traffic_class) {
     // Reconnect path may call this when QP is already in RTS; force a clean
     // state machine: RESET -> INIT -> RTR -> RTS.
     ibv_qp_attr qp_attr = {};
@@ -801,7 +898,7 @@ static int setupNotifyQpConnection(ibv_qp* qp, RdmaContext* ctx,
 
     memset(&qp_attr, 0, sizeof(qp_attr));
     qp_attr.qp_state = IBV_QPS_INIT;
-    qp_attr.pkey_index = 0;
+    qp_attr.pkey_index = pkey_index;
     qp_attr.port_num = ctx->portNum();
     qp_attr.qp_access_flags = IBV_ACCESS_REMOTE_WRITE;
     ret = ibv_modify_qp(
@@ -832,14 +929,14 @@ static int setupNotifyQpConnection(ibv_qp* qp, RdmaContext* ctx,
     qp_attr.min_rnr_timer = 0x12;
     qp_attr.ah_attr.is_global = 1;
     qp_attr.ah_attr.dlid = peer_lid;
-    qp_attr.ah_attr.sl = 0;
+    qp_attr.ah_attr.sl = service_level;
     qp_attr.ah_attr.src_path_bits = 0;
     qp_attr.ah_attr.port_num = ctx->portNum();
     memcpy(&qp_attr.ah_attr.grh.dgid, &peer_gid, 16);
     qp_attr.ah_attr.grh.flow_label = 0;
     qp_attr.ah_attr.grh.sgid_index = ctx->gidIndex();
     qp_attr.ah_attr.grh.hop_limit = 255;
-    qp_attr.ah_attr.grh.traffic_class = 0;
+    qp_attr.ah_attr.grh.traffic_class = traffic_class;
 
     ret = ibv_modify_qp(qp, &qp_attr,
                         IBV_QP_STATE | IBV_QP_PATH_MTU | IBV_QP_DEST_QPN |
diff --git a/mooncake-transfer-engine/tent/src/transport/rdma/endpoint_store.cpp b/mooncake-transfer-engine/tent/src/transport/rdma/endpoint_store.cpp
index 10dc22ba..d1b5ae73 100644
--- a/mooncake-transfer-engine/tent/src/transport/rdma/endpoint_store.cpp
+++ b/mooncake-transfer-engine/tent/src/transport/rdma/endpoint_store.cpp
@@ -28,7 +28,7 @@
 
 namespace mooncake {
 namespace tent {
-std::shared_ptr<RdmaEndPoint> FIFOEndpointStore::get(const std::string &key) {
+std::shared_ptr<RdmaEndPoint> FIFOEndpointStore::get(const std::string& key) {
     RWSpinlock::ReadGuard guard(endpoint_map_lock_);
     auto iter = endpoint_map_.find(key);
     if (iter != endpoint_map_.end()) return iter->second;
@@ -36,12 +36,38 @@ std::shared_ptr<RdmaEndPoint> FIFOEndpointStore::get(const std::string &key) {
 }
 
 std::shared_ptr<RdmaEndPoint> FIFOEndpointStore::getOrInsert(
-    const std::string &key) {
+    const std::string& key) {
     auto endpoint = get(key);
-    if (endpoint) return endpoint;
+    // Endpoints have unidirectional lifecycle - if existing endpoint is
+    // in terminal state (destroying/destroyed), remove it and create new.
+    if (endpoint) {
+        auto status = endpoint->status();
+        if (status == RdmaEndPoint::EP_DESTROYING ||
+            status == RdmaEndPoint::EP_DESTROYED) {
+            remove(endpoint.get());
+            endpoint = nullptr;
+        } else {
+            return endpoint;
+        }
+    }
     RWSpinlock::WriteGuard guard(endpoint_map_lock_);
-    if (endpoint_map_.find(key) != endpoint_map_.end())
-        return endpoint_map_[key];
+    // Double-check after acquiring write lock
+    auto iter = endpoint_map_.find(key);
+    if (iter != endpoint_map_.end()) {
+        auto ep = iter->second;
+        auto status = ep->status();
+        if (status == RdmaEndPoint::EP_DESTROYING ||
+            status == RdmaEndPoint::EP_DESTROYED) {
+            waiting_list_.insert(ep);
+            endpoint_map_.erase(iter);
+            auto fifo_iter = fifo_map_[key];
+            fifo_list_.erase(fifo_iter);
+            fifo_map_.erase(key);
+            // Fall through to create new endpoint
+        } else {
+            return ep;
+        }
+    }
     endpoint = std::make_shared<RdmaEndPoint>();
     int ret = endpoint->construct(&context_, &context_.params().endpoint, key);
     if (ret) {
@@ -56,18 +82,29 @@ std::shared_ptr<RdmaEndPoint> FIFOEndpointStore::getOrInsert(
     return endpoint;
 }
 
-int FIFOEndpointStore::remove(RdmaEndPoint *ep) {
+int FIFOEndpointStore::remove(RdmaEndPoint* ep) {
     RWSpinlock::WriteGuard guard(endpoint_map_lock_);
-    auto key = ep->name();
-    auto iter = endpoint_map_.find(key);
-    if (iter != endpoint_map_.end() && iter->second.get() == ep) {
-        waiting_list_.insert(iter->second);
-        endpoint_map_.erase(iter);
-        auto fifo_iter = fifo_map_[key];
-        fifo_list_.erase(fifo_iter);
-        fifo_map_.erase(key);
+    // Search for and remove the endpoint by pointer comparison
+    for (auto iter = endpoint_map_.begin(); iter != endpoint_map_.end();
+         ++iter) {
+        if (iter->second.get() == ep) {
+            waiting_list_.insert(iter->second);
+            iter->second->beginDestroyNoLock();
+            auto fifo_iter = fifo_map_[iter->first];
+            fifo_list_.erase(fifo_iter);
+            fifo_map_.erase(iter->first);
+            endpoint_map_.erase(iter);
+            return 0;
+        }
     }
-    return 0;
+    // If not found in endpoint_map, check if it's in waiting_list
+    for (const auto& waiting_ep : waiting_list_) {
+        if (waiting_ep.get() == ep) {
+            // Already in waiting list, no action needed
+            return 0;
+        }
+    }
+    return -1;  // Endpoint not found
 }
 
 void FIFOEndpointStore::evictOne() {
@@ -75,17 +112,20 @@ void FIFOEndpointStore::evictOne() {
     std::string victim = fifo_list_.front();
     fifo_list_.pop_front();
     fifo_map_.erase(victim);
-    waiting_list_.insert(endpoint_map_[victim]);
+    auto victim_endpoint = endpoint_map_[victim];
+    victim_endpoint->beginDestroy();
+    waiting_list_.insert(victim_endpoint);
     endpoint_map_.erase(victim);
+    LOG(INFO) << victim << " evicted from FIFOEndpointStore";
 }
 
 void FIFOEndpointStore::reclaim() {
     RWSpinlock::WriteGuard guard(endpoint_map_lock_);
     std::vector<std::shared_ptr<RdmaEndPoint>> to_delete;
-    for (auto &endpoint : waiting_list_) {
-        if (!endpoint->getInflightSlices()) to_delete.push_back(endpoint);
+    for (auto& endpoint : waiting_list_) {
+        if (endpoint->finishDestroy()) to_delete.push_back(endpoint);
     }
-    for (auto &endpoint : to_delete) waiting_list_.erase(endpoint);
+    for (auto& endpoint : to_delete) waiting_list_.erase(endpoint);
 }
 
 size_t FIFOEndpointStore::size() { return endpoint_map_.size(); }
@@ -93,8 +133,8 @@ size_t FIFOEndpointStore::size() { return endpoint_map_.size(); }
 void FIFOEndpointStore::clear() {
     RWSpinlock::WriteGuard guard(endpoint_map_lock_);
     std::vector<std::string> to_delete;
-    for (auto &entry : endpoint_map_) to_delete.push_back(entry.first);
-    for (auto &key : to_delete) {
+    for (auto& entry : endpoint_map_) to_delete.push_back(entry.first);
+    for (auto& key : to_delete) {
         endpoint_map_.erase(key);
         auto fifo_iter = fifo_map_[key];
         fifo_list_.erase(fifo_iter);
@@ -102,7 +142,7 @@ void FIFOEndpointStore::clear() {
     }
 }
 
-std::shared_ptr<RdmaEndPoint> SIEVEEndpointStore::get(const std::string &key) {
+std::shared_ptr<RdmaEndPoint> SIEVEEndpointStore::get(const std::string& key) {
     RWSpinlock::ReadGuard guard(endpoint_map_lock_);
     auto iter = endpoint_map_.find(key);
     if (iter != endpoint_map_.end()) {
@@ -113,12 +153,42 @@ std::shared_ptr<RdmaEndPoint> SIEVEEndpointStore::get(const std::string &key) {
 }
 
 std::shared_ptr<RdmaEndPoint> SIEVEEndpointStore::getOrInsert(
-    const std::string &key) {
+    const std::string& key) {
     auto endpoint = get(key);
-    if (endpoint) return endpoint;
+    // Endpoints have unidirectional lifecycle - if existing endpoint is
+    // in terminal state (destroying/destroyed), remove it and create new.
+    if (endpoint) {
+        auto status = endpoint->status();
+        if (status == RdmaEndPoint::EP_DESTROYING ||
+            status == RdmaEndPoint::EP_DESTROYED) {
+            remove(endpoint.get());
+            endpoint = nullptr;
+        } else {
+            return endpoint;
+        }
+    }
     RWSpinlock::WriteGuard guard(endpoint_map_lock_);
-    if (endpoint_map_.find(key) != endpoint_map_.end()) {
-        return endpoint_map_[key].first;
+    // Double-check after acquiring write lock
+    auto iter = endpoint_map_.find(key);
+    if (iter != endpoint_map_.end()) {
+        auto ep = iter->second.first;
+        auto status = ep->status();
+        if (status == RdmaEndPoint::EP_DESTROYING ||
+            status == RdmaEndPoint::EP_DESTROYED) {
+            waiting_list_len_++;
+            waiting_list_.insert(ep);
+            auto fifo_iter = fifo_map_[key];
+            if (hand_.has_value() && hand_.value() == fifo_iter) {
+                fifo_iter == fifo_list_.begin() ? hand_ = std::nullopt
+                                                : hand_ = std::prev(fifo_iter);
+            }
+            fifo_list_.erase(fifo_iter);
+            fifo_map_.erase(key);
+            endpoint_map_.erase(iter);
+            // Fall through to create new endpoint
+        } else {
+            return ep;
+        }
     }
     endpoint = std::make_shared<RdmaEndPoint>();
     int ret = endpoint->construct(&context_, &context_.params().endpoint, key,
@@ -135,23 +205,34 @@ std::shared_ptr<RdmaEndPoint> SIEVEEndpointStore::getOrInsert(
     return endpoint;
 }
 
-int SIEVEEndpointStore::remove(RdmaEndPoint *ep) {
+int SIEVEEndpointStore::remove(RdmaEndPoint* ep) {
     RWSpinlock::WriteGuard guard(endpoint_map_lock_);
-    auto key = ep->name();
-    auto iter = endpoint_map_.find(key);
-    if (iter != endpoint_map_.end() && iter->second.first.get() == ep) {
-        waiting_list_len_++;
-        waiting_list_.insert(iter->second.first);
-        endpoint_map_.erase(iter);
-        auto fifo_iter = fifo_map_[key];
-        if (hand_.has_value() && hand_.value() == fifo_iter) {
-            fifo_iter == fifo_list_.begin() ? hand_ = std::nullopt
-                                            : hand_ = std::prev(fifo_iter);
+    // Search for and remove the endpoint by pointer comparison
+    for (auto iter = endpoint_map_.begin(); iter != endpoint_map_.end();
+         ++iter) {
+        if (iter->second.first.get() == ep) {
+            waiting_list_len_++;
+            waiting_list_.insert(iter->second.first);
+            iter->second.first->beginDestroyNoLock();
+            auto fifo_iter = fifo_map_[iter->first];
+            if (hand_.has_value() && hand_.value() == fifo_iter) {
+                fifo_iter == fifo_list_.begin() ? hand_ = std::nullopt
+                                                : hand_ = std::prev(fifo_iter);
+            }
+            fifo_list_.erase(fifo_iter);
+            fifo_map_.erase(iter->first);
+            endpoint_map_.erase(iter);
+            return 0;
         }
-        fifo_list_.erase(fifo_iter);
-        fifo_map_.erase(key);
     }
-    return 0;
+    // If not found in endpoint_map, check if it's in waiting_list
+    for (const auto& waiting_ep : waiting_list_) {
+        if (waiting_ep.get() == ep) {
+            // Already in waiting list, no action needed
+            return 0;
+        }
+    }
+    return -1;  // Endpoint not found
 }
 
 void SIEVEEndpointStore::evictOne() {
@@ -174,6 +255,7 @@ void SIEVEEndpointStore::evictOne() {
     fifo_list_.erase(o);
     fifo_map_.erase(victim);
     auto victim_instance = endpoint_map_[victim].first;
+    victim_instance->beginDestroy();
     waiting_list_len_++;
     waiting_list_.insert(victim_instance);
     endpoint_map_.erase(victim);
@@ -185,10 +267,10 @@ void SIEVEEndpointStore::reclaim() {
     if (waiting_list_len_.load(std::memory_order_relaxed) == 0) return;
     RWSpinlock::WriteGuard guard(endpoint_map_lock_);
     std::vector<std::shared_ptr<RdmaEndPoint>> to_delete;
-    for (auto &endpoint : waiting_list_) {
-        if (!endpoint->getInflightSlices()) to_delete.push_back(endpoint);
+    for (auto& endpoint : waiting_list_) {
+        if (endpoint->finishDestroy()) to_delete.push_back(endpoint);
     }
-    for (auto &endpoint : to_delete) waiting_list_.erase(endpoint);
+    for (auto& endpoint : to_delete) waiting_list_.erase(endpoint);
     waiting_list_len_ -= to_delete.size();
 }
 
@@ -197,8 +279,8 @@ size_t SIEVEEndpointStore::size() { return endpoint_map_.size(); }
 void SIEVEEndpointStore::clear() {
     RWSpinlock::WriteGuard guard(endpoint_map_lock_);
     std::vector<std::string> to_delete;
-    for (auto &entry : endpoint_map_) to_delete.push_back(entry.first);
-    for (auto &key : to_delete) {
+    for (auto& entry : endpoint_map_) to_delete.push_back(entry.first);
+    for (auto& key : to_delete) {
         endpoint_map_.erase(key);
         auto fifo_iter = fifo_map_[key];
         if (hand_.has_value() && hand_.value() == fifo_iter) {
diff --git a/mooncake-transfer-engine/tent/src/transport/rdma/quota.cpp b/mooncake-transfer-engine/tent/src/transport/rdma/quota.cpp
index 202c19ad..53a8bb42 100644
--- a/mooncake-transfer-engine/tent/src/transport/rdma/quota.cpp
+++ b/mooncake-transfer-engine/tent/src/transport/rdma/quota.cpp
@@ -15,175 +15,346 @@
 #include "tent/transport/rdma/quota.h"
 #include "tent/transport/rdma/shared_quota.h"
 #include "tent/common/utils/random.h"
+#include "tent/common/utils/os.h"
 
-#include <assert.h>
-#include <unordered_set>
+#include <algorithm>
+#include <iostream>
+#include <iomanip>
 
 namespace mooncake {
 namespace tent {
-Status DeviceQuota::loadTopology(std::shared_ptr<Topology>& local_topology) {
+Status DeviceSelector::loadTopology(std::shared_ptr<Topology>& local_topology) {
     local_topology_ = local_topology;
-    std::unordered_set<int> used_numa_id;
     for (size_t dev_id = 0; dev_id < local_topology->getNicCount(); ++dev_id) {
         auto entry = local_topology->getNicEntry(dev_id);
-        if (entry->type != Topology::NIC_RDMA) continue;
+        if (!entry || entry->type != Topology::NIC_RDMA) continue;
         DeviceInfo& info = devices_[dev_id];
         info.dev_id = dev_id;
-        info.bw_gbps = 200.0;
+        info.bw_gbps = kDefaultBwGbps;
         info.numa_id = entry->numa_node;
-        used_numa_id.insert(entry->numa_node);
+        info.ewma_bandwidth_bps.store(info.getTheoreticalBandwidth(),
+                                      std::memory_order_relaxed);
     }
-    if (used_numa_id.size() == 1) allow_cross_numa_ = true;
+    // Initialize device base priorities after all devices are loaded
+    fillDevicePriorities();
     return Status::OK();
 }
 
-Status DeviceQuota::enableSharedQuota(const std::string& shm_name) {
-    shared_quota_ = std::make_shared<SharedQuotaManager>(this);
-    auto status = shared_quota_->attach(shm_name);
-    if (!status.ok()) shared_quota_.reset();
+Status DeviceSelector::enableSharedQuota(const std::string& shm_name) {
+    slot_manager_ = std::make_shared<SharedSlotManager>(this);
+    slot_manager_->setRotationIntervalMs(
+        sched_params_.slot_rotation_interval_ms);
+    auto status = slot_manager_->attach(shm_name);
+    if (!status.ok()) slot_manager_.reset();
     return status;
 }
 
-struct TlsDeviceInfo {
-    uint64_t active_bytes{0};
-    double beta0{0.0};
-    double beta1{1.0};
-};
-
-thread_local std::unordered_map<int, TlsDeviceInfo> tl_device_info;
-
-Status DeviceQuota::allocate(uint64_t length, const std::string& location,
-                             int& chosen_dev_id) {
+Status DeviceSelector::allocate(uint64_t total_length, uint32_t num_slices,
+                                uint64_t slice_bytes,
+                                const std::string& location,
+                                std::vector<int>& slice_dev_ids, int priority,
+                                uint64_t device_mask) {
+    slice_dev_ids.clear();
+    slice_dev_ids.reserve(num_slices);
     auto entry = local_topology_->getMemEntry(location);
     if (!entry) return Status::InvalidArgument("Unknown location" LOC_MARK);
 
-    if (!enable_quota_) {
-        thread_local int id = 0;
+    if (!smart_selection_enabled_) {
+        // Baseline mode: consistent with original TE behavior
+        // Use devices from the first non-empty rank only
+        thread_local uint64_t tl_rr_counter = 0;
         for (size_t rank = 0; rank < Topology::DevicePriorityRanks; ++rank) {
-            auto& list = entry->device_list[rank];
-            if (list.empty()) continue;
-            chosen_dev_id = list[id % list.size()];
-            id++;
+            thread_local std::vector<int> tl_eligible;
+            tl_eligible.clear();
+            for (int dev_id : entry->device_list[rank]) {
+                if (!devices_.count(dev_id)) continue;
+                if ((device_mask & (1ULL << dev_id)) == 0) continue;
+                tl_eligible.push_back(dev_id);
+            }
+            if (tl_eligible.empty()) continue;
+
+            // Found first non-empty rank, do round-robin within this rank
+            uint64_t offset = 0;
+            for (uint32_t i = 0; i < num_slices; ++i) {
+                int dev_id = tl_eligible[tl_rr_counter % tl_eligible.size()];
+                tl_rr_counter++;
+                slice_dev_ids.push_back(dev_id);
+                uint64_t this_slice_bytes =
+                    std::min(slice_bytes, total_length - offset);
+                offset += this_slice_bytes;
+                devices_[dev_id].total_bytes.fetch_add(
+                    this_slice_bytes, std::memory_order_relaxed);
+            }
             return Status::OK();
         }
-        return Status::DeviceNotFound("no eligible devices for " + location);
+        return Status::DeviceNotFound("no eligible devices");
     }
 
-    static constexpr double penalty[] = {1.0, 3.0, 10.0};
-    const double w = local_weight_;
-    std::unordered_map<int, double> score_map;
-    bool found_device = false;
-    double best_score = std::numeric_limits<double>::infinity();
-    for (size_t rank = 0; rank < Topology::DevicePriorityRanks; ++rank) {
-        if (rank == Topology::DevicePriorityRanks - 1 && !allow_cross_numa_ &&
-            found_device)
-            continue;
-        for (int dev_id : entry->device_list[rank]) {
-            if (!devices_.count(dev_id)) continue;
-            auto& dev = devices_[dev_id];
-            auto& tl_dev = tl_device_info[dev_id];
-            uint64_t overall_active_bytes =
-                dev.diffusion_active_bytes.load(std::memory_order_relaxed) +
-                dev.active_bytes.load(std::memory_order_relaxed);
-            double weighted_active = w * tl_dev.active_bytes +
-                                     (1.0 - w) * overall_active_bytes + length;
-            double beta0_g = dev.beta0.load(std::memory_order_relaxed);
-            double beta1_g = dev.beta1.load(std::memory_order_relaxed);
-            double beta0 = w * tl_dev.beta0 + (1.0 - w) * beta0_g;
-            double beta1 = w * tl_dev.beta1 + (1.0 - w) * beta1_g;
-            double bw = dev.bw_gbps * 1e9 / 8;
-            double predicted_time = (weighted_active / bw) * beta1 + beta0;
-            score_map[dev_id] = penalty[rank] * predicted_time;
-            best_score = std::min(best_score, score_map[dev_id]);
-            found_device = true;
-        }
+    std::vector<DeviceSelector::Candidate> tl_candidates;
+    Status status = buildCandidates(entry, slice_bytes, device_mask,
+                                    tl_candidates, priority);
+    if (!status.ok()) return status;
+    if (num_slices == 1) {
+        selectSinglePath(tl_candidates, num_slices, total_length,
+                         slice_dev_ids);
+    } else {
+        // Probe mode: every 100th call uses round-robin distribution
+        // to ensure all devices are sampled for EWMA updates
+        thread_local uint64_t tl_call_count = 0;
+        bool probe_mode = ((++tl_call_count % 100) == 0);
+        selectMultiPath(tl_candidates, num_slices, total_length, slice_dev_ids,
+                        probe_mode);
     }
-
-    if (!found_device) {
-        return Status::DeviceNotFound("no eligible devices for " + location);
-    }
-
-    std::vector<int> filtered;
-    for (const auto& [dev_id, score] : score_map) {
-        if (score <= best_score * 1.05) filtered.push_back(dev_id);
-    }
-
-    std::sort(filtered.begin(), filtered.end(), [&](int a, int b) {
-        if (std::abs(score_map[a] - score_map[b]) > 1e-9)
-            return score_map[a] < score_map[b];
-        return a < b;
-    });
-
-    thread_local size_t rr_index = 0;
-    chosen_dev_id = filtered[rr_index % filtered.size()];
-    rr_index++;
-
-    tl_device_info[chosen_dev_id].active_bytes += length;
-    if (local_weight_ < 1 - 1e-6)
-        devices_[chosen_dev_id].active_bytes.fetch_add(
-            length, std::memory_order_relaxed);
     return Status::OK();
 }
 
-Status DeviceQuota::release(int dev_id, uint64_t length, double latency) {
-    if (!enable_quota_) return Status::OK();
+int DeviceSelector::getDeviceRank(const std::string& location,
+                                  int dev_id) const {
+    auto entry = local_topology_->getMemEntry(location);
+    if (!entry) return 0;
+    for (size_t rank = 0; rank < Topology::DevicePriorityRanks; ++rank) {
+        for (int id : entry->device_list[rank]) {
+            if (id == dev_id) return static_cast<int>(rank);
+        }
+    }
+    return 0;
+}
+
+Status DeviceSelector::buildCandidates(const Topology::MemEntry* entry,
+                                       uint64_t slice_bytes,
+                                       uint64_t device_mask,
+                                       std::vector<Candidate>& candidates,
+                                       int request_priority) {
+    // Helper lambda to add candidate device
+    // Score formula: predicted_time × numa_penalty + random_jitter
+    // Lower score = better candidate
+    auto add_candidate = [&](int dev_id, size_t rank) {
+        auto& dev = devices_[dev_id];
+        uint64_t inflight = dev.getInflightBytes();
+        double ewma_bw = dev.getEwmaBandwidth();
+        double predicted_time =
+            static_cast<double>(inflight + slice_bytes) / ewma_bw;
+        double rank_penalty = sched_params_.numa_tier_weights[rank];
+        double score = predicted_time * rank_penalty;
+        score +=
+            (SimpleRandom::Get().next(10) * sched_params_.score_jitter_range);
+        bool is_cross_numa = (rank > 0);
+        Candidate c;
+        c.dev_id = dev_id;
+        c.score = score;
+        c.is_cross_numa = is_cross_numa;
+        candidates.push_back(c);
+    };
+
+    // First pass: filter by device priority (QoS filtering)
+    for (size_t rank = 0; rank < Topology::DevicePriorityRanks; ++rank) {
+        for (int dev_id : entry->device_list[rank]) {
+            if (!devices_.count(dev_id)) continue;
+            if ((device_mask & (1ULL << dev_id)) == 0) continue;
+            // QoS: Get device's current priority slot (local, per-process)
+            // Device accepts request if dev_priority >= request_priority
+            int dev_priority = PRIO_LOW;  // Default: accept all
+            if (sched_params_.enable_priority_filtering) {
+                dev_priority = getDevicePriority(dev_id);
+            }
+            if (dev_priority < request_priority) continue;
+            add_candidate(dev_id, rank);
+        }
+    }
+
+    // If no devices after filtering, fallback to all devices
+    if (candidates.empty()) {
+        for (size_t rank = 0; rank < Topology::DevicePriorityRanks; ++rank) {
+            for (int dev_id : entry->device_list[rank]) {
+                if (!devices_.count(dev_id)) continue;
+                if ((device_mask & (1ULL << dev_id)) == 0) continue;
+                add_candidate(dev_id, rank);
+            }
+        }
+    }
+
+    if (candidates.empty()) {
+        return Status::DeviceNotFound("no eligible devices");
+    }
+
+    std::sort(
+        candidates.begin(), candidates.end(),
+        [this](const Candidate& a, const Candidate& b) {
+            if (std::abs(a.score - b.score) > sched_params_.score_jitter_range)
+                return a.score < b.score;
+            return a.dev_id < b.dev_id;
+        });
+    return Status::OK();
+}
+
+void DeviceSelector::selectSinglePath(const std::vector<Candidate>& candidates,
+                                      uint32_t num_slices,
+                                      uint64_t total_length,
+                                      std::vector<int>& slice_dev_ids) {
+    if (candidates.empty()) return;
+
+    const Candidate& best = candidates[0];
+    int dev_id = best.dev_id;
+    auto& dev = devices_[dev_id];
+
+    dev.addInflight(total_length);
+    dev.total_bytes.fetch_add(total_length, std::memory_order_relaxed);
+
+    for (uint32_t i = 0; i < num_slices; ++i) {
+        slice_dev_ids.push_back(dev_id);
+    }
+}
+
+void DeviceSelector::selectMultiPath(const std::vector<Candidate>& candidates,
+                                     uint32_t num_slices, uint64_t total_length,
+                                     std::vector<int>& slice_dev_ids,
+                                     bool probe_mode) {
+    if (candidates.empty()) return;
+    uint64_t slice_bytes = (total_length + num_slices - 1) / num_slices;
+    if (probe_mode) {
+        // Probe mode: round-robin distribution to ensure all devices are
+        // sampled Activates every 100th call to prevent EWMA starvation
+        for (uint32_t i = 0; i < num_slices; ++i) {
+            const Candidate& c = candidates[i % candidates.size()];
+            slice_dev_ids.push_back(c.dev_id);
+            devices_[c.dev_id].addInflight(slice_bytes);
+            devices_[c.dev_id].total_bytes.fetch_add(slice_bytes,
+                                                     std::memory_order_relaxed);
+        }
+    } else {
+        // Normal mode: weighted distribution based on inverse score
+        // Lower score → higher weight → more slices
+        double total_weight = 0.0;
+        double max_weight = -1.0;
+        int best_dev_idx = -1;
+        for (size_t i = 0; i < candidates.size(); ++i) {
+            double w =
+                1.0 / (candidates[i].score + sched_params_.score_epsilon);
+            total_weight += w;
+            if (w > max_weight) {
+                max_weight = w;
+                best_dev_idx = static_cast<int>(i);
+            }
+        }
+        if (best_dev_idx == -1 || num_slices == 0 || total_weight <= 0.0)
+            return;
+        uint32_t remaining_slices = num_slices;
+        for (size_t i = 0; i < candidates.size(); ++i) {
+            double w =
+                1.0 / (candidates[i].score + sched_params_.score_epsilon);
+            uint32_t assigned =
+                static_cast<uint32_t>((w / total_weight) * num_slices);
+            if (assigned > 0) {
+                if (assigned > remaining_slices) assigned = remaining_slices;
+                remaining_slices -= assigned;
+                const Candidate& c = candidates[i];
+                for (uint32_t s = 0; s < assigned; ++s) {
+                    slice_dev_ids.push_back(c.dev_id);
+                }
+                uint64_t total_assigned_bytes =
+                    static_cast<uint64_t>(slice_bytes) * assigned;
+                devices_[c.dev_id].addInflight(total_assigned_bytes);
+                devices_[c.dev_id].total_bytes.fetch_add(
+                    total_assigned_bytes, std::memory_order_relaxed);
+            }
+        }
+        if (remaining_slices > 0) {
+            const Candidate& c = candidates[best_dev_idx];
+            for (uint32_t s = 0; s < remaining_slices; ++s) {
+                slice_dev_ids.push_back(c.dev_id);
+            }
+            uint64_t total_assigned_bytes =
+                static_cast<uint64_t>(slice_bytes) * remaining_slices;
+            devices_[c.dev_id].addInflight(total_assigned_bytes);
+            devices_[c.dev_id].total_bytes.fetch_add(total_assigned_bytes,
+                                                     std::memory_order_relaxed);
+        }
+    }
+}
+
+Status DeviceSelector::allocate(uint64_t length, const std::string& location,
+                                int& chosen_dev_id) {
+    std::vector<int> slice_dev_ids;
+    Status status = allocate(length, 1, length, location, slice_dev_ids, ~0ULL);
+    if (!status.ok()) return status;
+    if (slice_dev_ids.empty()) {
+        return Status::DeviceNotFound("allocation failed");
+    }
+    chosen_dev_id = slice_dev_ids[0];
+    return Status::OK();
+}
+
+Status DeviceSelector::release(int dev_id, uint64_t length, double latency) {
     auto it = devices_.find(dev_id);
     if (it == devices_.end())
         return Status::InvalidArgument("device not found");
 
     auto& dev = it->second;
-    auto& tl_dev = tl_device_info[dev_id];
+    dev.releaseInflight(length);
 
-    if (local_weight_ < 1 - 1e-6)
-        dev.active_bytes.fetch_sub(length, std::memory_order_relaxed);
-    tl_dev.active_bytes -= length;
-
-    if (!update_quota_params_) return Status::OK();
-
-    double bw = dev.bw_gbps * 1e9 / 8;
-    double theory_time = static_cast<double>(length) / bw;
-    double obs_time = latency;
-
-    const double w = local_weight_;
-    double beta0_g = dev.beta0.load(std::memory_order_relaxed);
-    double beta1_g = dev.beta1.load(std::memory_order_relaxed);
-    double beta0 = w * tl_dev.beta0 + (1.0 - w) * beta0_g;
-    double beta1 = w * tl_dev.beta1 + (1.0 - w) * beta1_g;
-
-    double pred_time = beta0 + beta1 * theory_time;
-    double err = obs_time - pred_time;
-    double rel_err = (pred_time > 1e-9) ? (err / pred_time) : 0.0;
-
-    double adapt_alpha = alpha_;
-    if (std::abs(err) > 0.05 * pred_time)
-        adapt_alpha = std::min(1.0, alpha_ * 5.0);
-
-    double delta0 = adapt_alpha * err;
-    double delta1 = adapt_alpha * rel_err;
-
-    double new_beta0_l = tl_dev.beta0 + w * delta0;
-    double new_beta1_l = tl_dev.beta1 * (1.0 + w * delta1);
-    tl_dev.beta0 = std::clamp(new_beta0_l, 0.0, 5e-4);
-    tl_dev.beta1 = std::clamp(new_beta1_l, 0.5, 20.0);
-
-    if (local_weight_ < 1 - 1e-6) {
-        double new_beta0_g = beta0_g + (1.0 - w) * delta0;
-        double new_beta1_g = beta1_g * (1.0 + (1.0 - w) * delta1);
-        dev.beta0.store(std::clamp(new_beta0_g, 0.0, 5e-4),
-                        std::memory_order_relaxed);
-        dev.beta1.store(std::clamp(new_beta1_g, 0.5, 20.0),
-                        std::memory_order_relaxed);
-        if (shared_quota_) {
-            thread_local uint64_t tl_last_ts = 0;
-            uint64_t now = getCurrentTimeInNano();
-            if (now - tl_last_ts > diffusion_interval_) {
-                tl_last_ts = now;
-                return shared_quota_->diffusion();
-            }
-        }
+    if (!smart_selection_enabled_) {
+        return Status::OK();
     }
+
+    // Update EWMA bandwidth: new = α × old + (1-α) × observed
+    // α = 0: always use observed (full adaptation)
+    // α = 1: never update (no learning)
+    double observed_bw = static_cast<double>(length) / latency;
+    double current_ewma = dev.getEwmaBandwidth();
+
+    double alpha = sched_params_.bandwidth_learning_rate;
+    double new_ewma = alpha * current_ewma + (1.0 - alpha) * observed_bw;
+
+    // Clamp to [min_multiplier, max_multiplier] of theoretical bandwidth
+    double theoretical_bw = dev.getTheoreticalBandwidth();
+    new_ewma = std::max(
+        sched_params_.ewma_min_multiplier * theoretical_bw,
+        std::min(sched_params_.ewma_max_multiplier * theoretical_bw, new_ewma));
+
+    dev.ewma_bandwidth_bps.store(new_ewma, std::memory_order_relaxed);
+
     return Status::OK();
 }
 
+void DeviceSelector::printTrafficStats() {
+    std::cout << "=== Device Traffic Statistics ===" << std::endl;
+    for (const auto& [dev_id, dev] : devices_) {
+        uint64_t total = dev.total_bytes.load(std::memory_order_relaxed);
+        double ewma_bw_gbps = dev.getEwmaBandwidth() / 1e9 * 8.0;
+        uint64_t inflight = dev.getInflightBytes();
+        std::cout << "Dev " << dev_id << ": "
+                  << "Total=" << (total / 1024.0 / 1024.0 / 1024.0) << " GB, "
+                  << "EWMA BW=" << std::fixed << std::setprecision(2)
+                  << ewma_bw_gbps << " Gbps, "
+                  << "Inflight=" << inflight << " bytes" << std::endl;
+    }
+}
+
+void DeviceSelector::fillDevicePriorities() {
+    sched_params_.device_base_priorities.clear();
+    for (const auto& [dev_id, dev] : devices_) {
+        sched_params_.device_base_priorities.push_back(dev_id);
+    }
+}
+
+int DeviceSelector::getDevicePriority(int dev_id) const {
+    if (!sched_params_.enable_priority_filtering) return 0;
+    auto it = std::find(sched_params_.device_base_priorities.begin(),
+                        sched_params_.device_base_priorities.end(), dev_id);
+    if (it == sched_params_.device_base_priorities.end()) return 0;
+    size_t base_index =
+        std::distance(sched_params_.device_base_priorities.begin(), it);
+    size_t num_devices = sched_params_.device_base_priorities.size();
+    if (sched_params_.local_rotation_interval_us > 0 && num_devices > 0) {
+        uint64_t now = getCurrentTimeInNano();
+        uint64_t offset_us = now / 1000;
+        size_t rotation_offset =
+            (offset_us / sched_params_.local_rotation_interval_us) %
+            num_devices;
+        base_index = (base_index + rotation_offset) % num_devices;
+    }
+    return static_cast<int>(base_index);
+}
+
 }  // namespace tent
 }  // namespace mooncake
\ No newline at end of file
diff --git a/mooncake-transfer-engine/tent/src/transport/rdma/rail_monitor.cpp b/mooncake-transfer-engine/tent/src/transport/rdma/rail_monitor.cpp
index 5bbb10a1..438b211d 100644
--- a/mooncake-transfer-engine/tent/src/transport/rdma/rail_monitor.cpp
+++ b/mooncake-transfer-engine/tent/src/transport/rdma/rail_monitor.cpp
@@ -17,10 +17,21 @@
 namespace mooncake {
 namespace tent {
 
-Status RailMonitor::load(const Topology *local, const Topology *remote,
-                         const std::string &rail_topo_json) {
+Status RailMonitor::load(const Topology* local, const Topology* remote,
+                         const std::string& rail_topo_json,
+                         const Config* conf) {
     local_ = local;
     remote_ = remote;
+    if (conf) {
+        error_threshold_ = conf->get(kCfgErrorThreshold, error_threshold_);
+        error_window_ = std::chrono::seconds(
+            conf->get(kCfgErrorWindowSecs, (int)error_window_.count()));
+        cooldown_ = std::chrono::seconds(
+            conf->get(kCfgCooldownSecs, (int)cooldown_.count()));
+        LOG(INFO) << "RailMonitor: error_threshold=" << error_threshold_
+                  << " error_window=" << error_window_.count() << "s"
+                  << " cooldown=" << cooldown_.count() << "s";
+    }
     if (!rail_topo_json.empty()) {
         auto status = loadFromJson(rail_topo_json);
         if (status.ok()) return status;
@@ -32,24 +43,24 @@ Status RailMonitor::load(const Topology *local, const Topology *remote,
 bool RailMonitor::available(int local_nic, int remote_nic) {
     auto it = rail_states_.find(std::make_pair(local_nic, remote_nic));
     if (it == rail_states_.end()) return false;
-    auto &st = it->second;
-    if (st.paused) {
-        auto now = std::chrono::steady_clock::now();
-        if (now >= st.resume_time) {
-            st.paused = false;
-            st.error_count = 0;
-            updateBestMapping();
-            return true;
-        }
-        return false;
-    }
+    auto& st = it->second;
+    if (!st.paused()) return true;
+    if (std::chrono::steady_clock::now() < st.resume_time) return false;
+    // Cooldown expired: clear all exponential-backoff memory so a fresh
+    // failure cycle starts from the initial cooldown_, not a doubled value.
+    st.resume_time = {};
+    st.error_count = 0;
+    st.cooldown = std::chrono::seconds(0);
+    updateBestMapping();
+    LOG(INFO) << "Rail recovered: local_nic=" << local_nic
+              << " remote_nic=" << remote_nic << " (cooldown expired)";
     return true;
 }
 
 void RailMonitor::markFailed(int local_nic, int remote_nic) {
     auto it = rail_states_.find(std::make_pair(local_nic, remote_nic));
     if (it == rail_states_.end()) return;
-    auto &st = it->second;
+    auto& st = it->second;
     auto now = std::chrono::steady_clock::now();
     if (st.error_count == 0 || now - st.last_error > error_window_) {
         st.error_count = 1;
@@ -61,11 +72,9 @@ void RailMonitor::markFailed(int local_nic, int remote_nic) {
         st.cooldown = cooldown_;
     } else {
         st.cooldown *= 2;
-        if (st.cooldown > std::chrono::seconds(300))
-            st.cooldown = std::chrono::seconds(300);
+        if (st.cooldown > kMaxCooldown) st.cooldown = kMaxCooldown;
     }
     if (st.error_count >= error_threshold_) {
-        st.paused = true;
         st.resume_time = now + st.cooldown;
         updateBestMapping();
     }
@@ -74,11 +83,23 @@ void RailMonitor::markFailed(int local_nic, int remote_nic) {
 void RailMonitor::markRecovered(int local_nic, int remote_nic) {
     auto it = rail_states_.find(std::make_pair(local_nic, remote_nic));
     if (it == rail_states_.end()) return;
-    auto &st = it->second;
+    auto& st = it->second;
+    // Fast path: a healthy rail stays healthy. 99%+ of completions land
+    // here, so we must not touch best_mapping_ or write any field.
+    if (!st.paused() && st.error_count == 0 && st.cooldown.count() == 0) return;
+    bool was_paused = st.paused();
+    // Clear all exponential-backoff memory: the next failure cycle must
+    // start from the initial cooldown_, not a doubled value left over
+    // from the previous cycle.
     st.error_count = 0;
-    st.paused = false;
     st.resume_time = {};
-    updateBestMapping();
+    st.cooldown = std::chrono::seconds(0);
+    if (was_paused) {
+        LOG(INFO) << "Rail recovered: local_nic=" << local_nic
+                  << " remote_nic=" << remote_nic
+                  << " (un-paused by successful transfer)";
+        updateBestMapping();
+    }
 }
 
 int RailMonitor::findBestRemoteDevice(int local_nic, int remote_numa) {
@@ -108,7 +129,7 @@ int RailMonitor::findBestRemoteDevice(int local_nic, int remote_numa) {
  *   ]
  * }
  */
-Status RailMonitor::loadFromJson(const std::string &rail_topo_json) {
+Status RailMonitor::loadFromJson(const std::string& rail_topo_json) {
     try {
         auto root = json::parse(rail_topo_json);
 
@@ -116,7 +137,7 @@ Status RailMonitor::loadFromJson(const std::string &rail_topo_json) {
         direct_rails_.clear();
 
         if (root.contains("all")) {
-            for (const auto &path_entry : root["all"]) {
+            for (const auto& path_entry : root["all"]) {
                 std::string local_nic_name = path_entry.value("local", "");
                 std::string remote_nic_name = path_entry.value("remote", "");
                 int local_nic_id = local_->getNicId(local_nic_name);
@@ -132,7 +153,7 @@ Status RailMonitor::loadFromJson(const std::string &rail_topo_json) {
         }
 
         if (root.contains("direct")) {
-            for (const auto &path_entry : root["direct"]) {
+            for (const auto& path_entry : root["direct"]) {
                 std::string local_nic_name = path_entry.value("local", "");
                 std::string remote_nic_name = path_entry.value("remote", "");
                 int local_nic_id = local_->getNicId(local_nic_name);
@@ -146,7 +167,7 @@ Status RailMonitor::loadFromJson(const std::string &rail_topo_json) {
                 }
             }
         }
-    } catch (const std::exception &ex) {
+    } catch (const std::exception& ex) {
         LOG(ERROR) << "Failed to parse rail_topo_json: " << ex.what();
         return Status::InvalidArgument("Failed to parse JSON" LOC_MARK);
     }
@@ -156,12 +177,12 @@ Status RailMonitor::loadFromJson(const std::string &rail_topo_json) {
     return Status::OK();
 }
 
-static int matchRemoteNicId(const Topology *local, const Topology *remote,
+static int matchRemoteNicId(const Topology* local, const Topology* remote,
                             int local_nic) {
     std::string mem_name;
     for (size_t i = 0; i < local->getMemCount(); ++i) {
         auto entry = local->getMemEntry(i);
-        auto &prior_devices = entry->device_list[0];
+        auto& prior_devices = entry->device_list[0];
         if (entry->type == Topology::MEM_CUDA && !prior_devices.empty() &&
             prior_devices[0] == local_nic) {
             mem_name = entry->name;
@@ -172,7 +193,7 @@ static int matchRemoteNicId(const Topology *local, const Topology *remote,
     auto mem_id = remote->getMemId(mem_name);
     if (mem_id < 0) return -1;
     auto entry = remote->getMemEntry(mem_id);
-    auto &prior_devices = entry->device_list[0];
+    auto& prior_devices = entry->device_list[0];
     if (entry->type == Topology::MEM_CUDA && !prior_devices.empty())
         return prior_devices[0];
     return -1;
@@ -193,6 +214,23 @@ Status RailMonitor::loadDefault() {
         auto local_entry = local_->getNicEntry(local_nic);
         if (local_entry->type != Topology::NIC_RDMA) continue;
         int numa_id = local_entry->numa_node;
+
+        // Priority 1: Same-name device matching (mlx5_0 -> mlx5_0)
+        bool matched = false;
+        auto local_nic_name = local_entry->name;
+        for (int remote_nic = 0; remote_nic < remote_nic_count; ++remote_nic) {
+            auto remote_entry = remote_->getNicEntry(remote_nic);
+            if (remote_entry && remote_entry->type == Topology::NIC_RDMA &&
+                remote_entry->name == local_nic_name) {
+                remote_load[remote_nic]++;
+                direct_rails_[local_nic] = remote_nic;
+                matched = true;
+                break;
+            }
+        }
+        if (matched) continue;
+
+        // Priority 2: CUDA memory topology matching (GPU-direct NIC)
         int remote_nic = matchRemoteNicId(local_, remote_, local_nic);
         if (remote_nic >= 0) {
             remote_load[remote_nic]++;
@@ -278,7 +316,7 @@ void RailMonitor::updateBestMapping() {
 
     for (size_t local_numa = 0; local_numa < kMaxNuma; ++local_numa) {
         for (size_t remote_numa = 0; remote_numa < kMaxNuma; ++remote_numa) {
-            auto &mapping = best_mapping_[remote_numa];
+            auto& mapping = best_mapping_[remote_numa];
             size_t local_cnt = local_devices[local_numa].size();
             size_t remote_cnt = remote_devices[remote_numa].size();
             if (!local_cnt || !remote_cnt) continue;
diff --git a/mooncake-transfer-engine/tent/src/transport/rdma/rdma_transport.cpp b/mooncake-transfer-engine/tent/src/transport/rdma/rdma_transport.cpp
index dad739f4..4a80f357 100644
--- a/mooncake-transfer-engine/tent/src/transport/rdma/rdma_transport.cpp
+++ b/mooncake-transfer-engine/tent/src/transport/rdma/rdma_transport.cpp
@@ -14,18 +14,23 @@
 
 #include "tent/transport/rdma/rdma_transport.h"
 #include "tent/transport/rdma/ibv_loader.h"
+#include "tent/transport/rdma/quota.h"
 
 #include <glog/logging.h>
 #include <sys/mman.h>
 #include <sys/time.h>
 
 #include <cassert>
+#include <cerrno>
 #include <cstddef>
+#include <cstdlib>
 #include <future>
+#include <limits>
 #include <set>
 #include <sstream>
 
 #include "tent/common/status.h"
+#include "tent/common/utils/ip.h"
 #include "tent/transport/rdma/buffers.h"
 #include "tent/transport/rdma/endpoint_store.h"
 #include "tent/transport/rdma/workers.h"
@@ -44,6 +49,39 @@
 
 namespace mooncake {
 namespace tent {
+
+namespace {
+
+uint16_t getRdmaBindDefaultPort(const Config& config) {
+    constexpr const char* kKey = "rpc_server_port";
+    if (!config.contains(kKey)) return 0;
+
+    json raw_value = config.get<json>(kKey, json());
+    if (raw_value.is_number_integer() || raw_value.is_number_unsigned()) {
+        long long value = raw_value.get<long long>();
+        if (value >= 0 && value <= static_cast<long long>(
+                                       std::numeric_limits<uint16_t>::max())) {
+            return static_cast<uint16_t>(value);
+        }
+        return 0;
+    }
+
+    if (raw_value.is_string()) {
+        const std::string string_value = raw_value.get<std::string>();
+        char* end = nullptr;
+        errno = 0;
+        unsigned long value = std::strtoul(string_value.c_str(), &end, 10);
+        if (errno == 0 && end != string_value.c_str() && *end == '\0' &&
+            value <= std::numeric_limits<uint16_t>::max()) {
+            return static_cast<uint16_t>(value);
+        }
+    }
+
+    return 0;
+}
+
+}  // namespace
+
 static Status configureLaneCount(std::shared_ptr<Config> conf,
                                  std::shared_ptr<RdmaParams> params) {
     constexpr int kUnset = -1;
@@ -163,8 +201,10 @@ static Status convertConfToRdmaParams(std::shared_ptr<Config> conf,
     return Status::OK();
 }
 
-static bool isGpuDirectRdmaSupported() {
-    if (getenv("MC_DISABLE_GPU_DIRECT_RDMA")) {
+static bool isGpuDirectRdmaSupported(std::shared_ptr<Config> conf) {
+    auto disable_gpu_direct =
+        conf->get("transports/rdma/disable_gpu_direct_rdma", false);
+    if (disable_gpu_direct) {
         return false;
     }
     std::ifstream modules("/proc/modules");
@@ -210,6 +250,24 @@ Status RdmaTransport::install(std::string& local_segment_name,
     metadata_ = metadata;
     local_segment_name_ = local_segment_name;
     local_topology_ = local_topology;
+
+    // In dual-NIC environments (e.g. separate TCP and RDMA interfaces),
+    // transports/rdma/bind_address allows NIC paths to use an
+    // RDMA-reachable IP while local_segment_name_ keeps the
+    // TCP-reachable address for P2P.
+    const auto rdma_bind_addr = conf_->get("transports/rdma/bind_address", "");
+    if (!rdma_bind_addr.empty()) {
+        const uint16_t default_port = getRdmaBindDefaultPort(*conf_);
+        auto [host_name, port] =
+            parseHostNameWithPort(local_segment_name, default_port);
+        rdma_server_name_ = rdma_bind_addr + ":" + std::to_string(port);
+        LOG(INFO) << "RdmaTransport(TENT): using RDMA bind address "
+                  << rdma_server_name_
+                  << " (TCP address: " << local_segment_name_ << ")";
+    } else {
+        rdma_server_name_ = local_segment_name_;
+    }
+
     local_buffer_manager_.setTopology(local_topology);
     context_set_.clear();
     for (size_t i = 0; i < local_topology_->getNicCount(); ++i) {
@@ -252,7 +310,7 @@ Status RdmaTransport::install(std::string& local_segment_name,
 
     installed_ = true;
     caps.dram_to_dram = true;
-    if (isGpuDirectRdmaSupported()) {
+    if (isGpuDirectRdmaSupported(conf_)) {
         caps.dram_to_gpu = true;
         caps.gpu_to_dram = true;
         caps.gpu_to_gpu = true;
@@ -293,6 +351,10 @@ Status RdmaTransport::freeSubBatch(SubBatchRef& batch) {
     auto rdma_batch = dynamic_cast<RdmaSubBatch*>(batch);
     if (!rdma_batch)
         return Status::InvalidArgument("Invalid RDMA sub-batch" LOC_MARK);
+    for (auto* task : rdma_batch->task_list) {
+        task->deref();  // Release batch's reference to the task
+    }
+    rdma_batch->task_list.clear();
     for (auto slice : rdma_batch->slice_chain) {
         while (slice) {
             auto next = slice->next;
@@ -320,28 +382,27 @@ Status RdmaTransport::submitTransferTasks(
 
     const size_t default_block_size = params_->workers.block_size;
     const int num_workers = params_->workers.num_workers;
-    const int num_devices = (size_t)local_topology_->getNicCount();
     std::vector<RdmaSliceList> slice_lists(num_workers);
     std::vector<RdmaSlice*> slice_tails(num_workers, nullptr);
     auto enqueue_ts = getCurrentTimeInNano();
 
+    // Distribute starting worker across threads to avoid contention
     static std::atomic<int> g_caller_threads(0);
     thread_local int tl_caller_id = g_caller_threads.fetch_add(1);
-    bool enable_spray =
-        g_caller_threads.load(std::memory_order_relaxed) <= num_workers;
-    int submit_slices = 0;
+    int next_worker_idx = tl_caller_id;
     for (auto& request : request_list) {
         auto opcode = request.opcode;
         auto type = Platform::getLoader().getMemoryType(request.source);
         size_t max_slice_count = 64;
         if (type == MTYPE_CUDA || opcode == Request::WRITE)
             max_slice_count = 32;
-        rdma_batch->task_list.push_back(RdmaTask{});
-        auto& task = rdma_batch->task_list.back();
-        task.request = request;
-        task.num_slices = 0;
-        task.status_word = PENDING;
-        task.transferred_bytes = 0;
+        auto* task = RdmaTaskStorage::Get().allocate();
+        rdma_batch->task_list.push_back(task);
+        task->request = request;
+        task->num_slices = 0;
+        task->status_word = PENDING;
+        task->transferred_bytes = 0;
+        task->ref();  // Batch holds a reference to the task
 
         const double merge_ratio = 0.25;
         uint64_t base_block = default_block_size;
@@ -360,8 +421,27 @@ Status RdmaTransport::submitTransferTasks(
         uint64_t block_size = roundup(
             (request.length + num_slices - 1) / num_slices, default_block_size);
 
-        num_slices = std::max<uint64_t>(
-            1, std::min<uint64_t>(num_slices, max_slice_count));
+        std::vector<int> slice_dev_ids;
+        // Only if a single request is enough, we perform aggregated allocation
+        if (num_slices >= max_slice_count / 2) {
+            std::string source_location = kWildcardLocation;
+            auto source_locations =
+                Platform::getLoader().getLocation(request.source, 1, true);
+            if (!source_locations.empty()) {
+                source_location = source_locations[0].location;
+            }
+            auto device_selector = workers_->getDeviceSelector();
+            if (device_selector) {
+                auto status = device_selector->allocate(
+                    request.length, static_cast<uint32_t>(num_slices),
+                    block_size, source_location, slice_dev_ids,
+                    request.priority, batch->device_mask);
+                if (!status.ok() || slice_dev_ids.empty()) {
+                    LOG(WARNING) << "Device quota allocation failed: "
+                                 << status.message();
+                }
+            }
+        }
 
         uint64_t offset = 0;
         for (uint64_t slice_idx = 0; slice_idx < num_slices; ++slice_idx) {
@@ -371,22 +451,23 @@ Status RdmaTransport::submitTransferTasks(
             slice->source_addr = (char*)request.source + offset;
             slice->target_addr = request.target_offset + offset;
             slice->length = length;
-            slice->task = &task;
+            slice->task = task;
             slice->retry_count = 0;
             slice->ep_weak_ptr.reset();
             slice->word = PENDING;
             slice->next = nullptr;
             slice->enqueue_ts = enqueue_ts;
-            task.num_slices++;
+            slice->priority = request.priority;  // Copy priority from request
+            task->num_slices++;
+            task->ref();  // Each slice holds a reference to the task
+            if (slice_idx < slice_dev_ids.size())
+                slice->source_dev_id = slice_dev_ids[slice_idx];
             offset += length;
-            int part_id =
-                ((enable_spray ? submit_slices : static_cast<int>(slice_idx)) /
-                 num_devices) %
-                num_workers;
+            int part_id = next_worker_idx % num_workers;
             auto& list = slice_lists[part_id];
             auto& tail = slice_tails[part_id];
             list.num_slices++;
-            submit_slices++;
+            next_worker_idx++;
             if (list.first) {
                 tail->next = slice;
                 tail = slice;
@@ -399,7 +480,7 @@ Status RdmaTransport::submitTransferTasks(
     for (int i = 0; i < num_workers; ++i) {
         if (slice_lists[i].first) {
             rdma_batch->slice_chain.push_back(slice_lists[i].first);
-            workers_->submit(slice_lists[i], (tl_caller_id + i) % num_workers);
+            workers_->submit(slice_lists[i], i);
         }
     }
     return Status::OK();
@@ -411,8 +492,8 @@ Status RdmaTransport::getTransferStatus(SubBatchRef batch, int task_id,
     if (task_id < 0 || task_id >= (int)rdma_batch->task_list.size()) {
         return Status::InvalidArgument("Invalid task ID" LOC_MARK);
     }
-    auto& task = rdma_batch->task_list[task_id];
-    status = TransferStatus{task.status_word, task.transferred_bytes};
+    auto* task = rdma_batch->task_list[task_id];
+    status = TransferStatus{task->status_word, task->transferred_bytes};
     return Status::OK();
 }
 
@@ -462,6 +543,11 @@ Status RdmaTransport::setupLocalSegment() {
     auto& manager = metadata_->segmentManager();
     auto segment = manager.getLocal();
     assert(segment);
+    // Store RDMA server name for dual-NIC setups; when it differs from
+    // local_segment_name_ the peer will use it for NIC path construction.
+    if (rdma_server_name_ != local_segment_name_) {
+        segment->rdma_server_name = rdma_server_name_;
+    }
     auto& detail = std::get<MemorySegmentDesc>(segment->detail);
     for (auto& context : context_set_) {
         if (context->status() != RdmaContext::DEVICE_ENABLED) continue;
@@ -550,12 +636,9 @@ std::shared_ptr<RdmaEndPoint> RdmaTransport::getEndpoint(SegmentID target_id,
         return nullptr;
     }
     std::shared_ptr<RdmaEndPoint> endpoint;
-    std::string peer_name = MakeNicPath(segment_desc->name, target_dev_name);
+    std::string peer_name =
+        MakeNicPath(segment_desc->nicPathServerName(), target_dev_name);
     endpoint = context->endpointStore()->getOrInsert(peer_name);
-    if (endpoint && endpoint->status() == RdmaEndPoint::EP_RESET) {
-        context->endpointStore()->remove(endpoint.get());
-        endpoint = context->endpointStore()->getOrInsert(peer_name);
-    }
     if (!endpoint) {
         LOG(ERROR) << "Cannot allocate endpoint " << peer_name;
         return nullptr;
diff --git a/mooncake-transfer-engine/tent/src/transport/rdma/shared_quota.cpp b/mooncake-transfer-engine/tent/src/transport/rdma/shared_quota.cpp
index 5aff9c18..f0b5a4ec 100644
--- a/mooncake-transfer-engine/tent/src/transport/rdma/shared_quota.cpp
+++ b/mooncake-transfer-engine/tent/src/transport/rdma/shared_quota.cpp
@@ -14,94 +14,43 @@
 
 #include "tent/transport/rdma/shared_quota.h"
 #include "tent/common/utils/os.h"
+#include "tent/common/types.h"
+
+#include <glog/logging.h>
+#include <unistd.h>
 
 namespace mooncake {
 namespace tent {
-SharedQuotaManager::SharedQuotaManager(DeviceQuota* local_quota)
+SharedSlotManager::SharedSlotManager(DeviceSelector* device_selector)
     : hdr_(nullptr),
       fd_(-1),
       size_(sizeof(SharedHeader)),
       created_(false),
-      local_quota_(local_quota) {}
+      device_selector_(device_selector),
+      background_running_(false) {}
 
-SharedQuotaManager::~SharedQuotaManager() { detach(); }
+SharedSlotManager::~SharedSlotManager() { detach(); }
 
-Status SharedQuotaManager::attach(const std::string& shm_name) {
-    name_ = shm_name;
-
-    // Open or create shared memory (mode 0666)
-    fd_ = shm_open(name_.c_str(), O_RDWR | O_CREAT, 0666);
-    if (fd_ < 0) {
-        return Status::InternalError("shm_open failed: " +
-                                     std::string(strerror(errno)));
-    }
-
-    // Ensure size
-    if (ftruncate(fd_, static_cast<off_t>(size_)) != 0) {
-        int e = errno;
-        close(fd_);
-        fd_ = -1;
-        return Status::InternalError("ftruncate failed: " +
-                                     std::string(strerror(e)));
-    }
-
-    // mmap
-    void* ptr =
-        mmap(nullptr, size_, PROT_READ | PROT_WRITE, MAP_SHARED, fd_, 0);
-    if (ptr == MAP_FAILED) {
-        int e = errno;
-        close(fd_);
-        fd_ = -1;
-        return Status::InternalError("mmap failed: " +
-                                     std::string(strerror(e)));
-    }
-
-    hdr_ = reinterpret_cast<SharedHeader*>(ptr);
-
-    if (hdr_->magic != SHM_MAGIC || hdr_->version != SHM_VERSION) {
-        created_ = true;
-        Status s = initializeHeader();
-        if (!s.ok()) {
-            munmap(ptr, size_);
-            close(fd_);
-            hdr_ = nullptr;
-            fd_ = -1;
-            return s;
-        }
-    } else {
-        created_ = false;
-    }
-
-    Status s = attachProcess();
-    if (!s.ok()) return s;
-    return Status::OK();
+bool SharedSlotManager::isPriorityAllowedInSlot(int priority, int slot) const {
+    return priority <= slot;
 }
 
-Status SharedQuotaManager::detach() {
-    if (hdr_) {
-        detachProcess();
-        munmap(hdr_, size_);
-        hdr_ = nullptr;
-    }
-    if (fd_ >= 0) {
-        close(fd_);
-        fd_ = -1;
-    }
-    return Status::OK();
-}
+Status SharedSlotManager::initializeHeader() {
+    hdr_->magic = 0;
+    hdr_->version = 0;
+    hdr_->current_slot.store(0, std::memory_order_relaxed);
 
-Status SharedQuotaManager::initializeHeader() {
-    memset(hdr_, 0, size_);
     Status s = initMutex(&hdr_->global_mutex);
-    if (!s.ok()) {
-        return s;
-    }
+    if (!s.ok()) return s;
+
     hdr_->version = SHM_VERSION;
     hdr_->magic = SHM_MAGIC;
+    hdr_->current_slot.store(0, std::memory_order_release);
+
     return Status::OK();
 }
 
-Status SharedQuotaManager::initMutex(pthread_mutex_t* m) {
+Status SharedSlotManager::initMutex(pthread_mutex_t* m) {
     pthread_mutexattr_t attr;
     if (pthread_mutexattr_init(&attr) != 0) {
         return Status::InternalError("pthread_mutexattr_init failed");
@@ -124,162 +73,115 @@ Status SharedQuotaManager::initMutex(pthread_mutex_t* m) {
     return Status::OK();
 }
 
-// attempt to acquire global lock; handle EOWNERDEAD
-int SharedQuotaManager::lock() {
-    if (!hdr_) return EINVAL;
-    int rc = pthread_mutex_lock(&hdr_->global_mutex);
-    if (rc == 0) return 0;
-    if (rc == EOWNERDEAD) {
-        // make consistent so others can continue
-#if defined(PTHREAD_MUTEX_ROBUST)
-        int rc2 = pthread_mutex_consistent(&hdr_->global_mutex);
-        if (rc2 != 0) {
-            return rc2;
+Status SharedSlotManager::attach(const std::string& shm_name) {
+    name_ = shm_name;
+
+    fd_ = shm_open(name_.c_str(), O_RDWR | O_CREAT, 0666);
+    if (fd_ < 0) {
+        return Status::InternalError("shm_open failed: " +
+                                     std::string(std::strerror(errno)));
+    }
+
+    if (ftruncate(fd_, static_cast<off_t>(size_)) != 0) {
+        int e = errno;
+        close(fd_);
+        fd_ = -1;
+        return Status::InternalError("ftruncate failed: " +
+                                     std::string(std::strerror(e)));
+    }
+
+    void* ptr =
+        mmap(nullptr, size_, PROT_READ | PROT_WRITE, MAP_SHARED, fd_, 0);
+    if (ptr == MAP_FAILED) {
+        int e = errno;
+        close(fd_);
+        fd_ = -1;
+        return Status::InternalError("mmap failed: " +
+                                     std::string(std::strerror(e)));
+    }
+
+    hdr_ = reinterpret_cast<SharedHeader*>(ptr);
+
+    if (hdr_->magic != SHM_MAGIC || hdr_->version != SHM_VERSION) {
+        created_ = true;
+        Status s = initializeHeader();
+        if (!s.ok()) {
+            munmap(ptr, size_);
+            close(fd_);
+            hdr_ = nullptr;
+            fd_ = -1;
+            return s;
         }
-#endif
-        // We hold the lock now — repair data if needed
-        reclaimDeadPidsInternal();
-        return 0;
-    }
-    return rc;
-}
-
-int SharedQuotaManager::unlock() {
-    if (!hdr_) return EINVAL;
-    return pthread_mutex_unlock(&hdr_->global_mutex);
-}
-
-void SharedQuotaManager::reclaimDeadPidsInternal() {
-    if (!hdr_) return;
-    // Assumes caller holds lock (but we also call from lock() on EOWNERDEAD)
-    for (int i = 0; i < hdr_->num_devices; ++i) {
-        // skip unused slots (dev_name empty)
-        if (hdr_->devices[i].dev_name[0] == '\0') continue;
-        SharedDeviceEntry& dev = hdr_->devices[i];
-        for (int s = 0; s < MAX_PID_SLOTS; ++s) {
-            pid_t p = dev.pid_usages[s].pid;
-            if (p == 0) continue;
-            if (!isPidAlive(p)) {
-                // zero out slot — we'll recompute active_bytes in diffusion
-                dev.pid_usages[s].pid = 0;
-                dev.pid_usages[s].used_bytes = 0;
-            }
-        }
-    }
-}
-
-bool SharedQuotaManager::isPidAlive(pid_t pid) {
-    if (pid <= 0) return false;
-    int r = kill(pid, 0);
-    if (r == 0) return true;
-    if (errno == ESRCH) return false;
-    return true;  // other errors (EPERM) -> treat as alive
-}
-
-int SharedQuotaManager::findDeviceIdByNameLocked(const std::string& dev_name) {
-    if (!hdr_) return -1;
-    for (int i = 0; i < hdr_->num_devices; ++i) {
-        if (hdr_->devices[i].dev_name[0] == '\0') continue;
-        if (strncmp(hdr_->devices[i].dev_name, dev_name.c_str(),
-                    sizeof(hdr_->devices[i].dev_name)) == 0)
-            return i;
-    }
-    return -1;
-}
-
-PidUsage* SharedQuotaManager::findOrCreatePidSlotLocked(int dev_id, pid_t pid) {
-    if (!hdr_) return nullptr;
-    if (dev_id < 0 || dev_id >= hdr_->num_devices) return nullptr;
-    SharedDeviceEntry& dev = hdr_->devices[dev_id];
-    PidUsage* empty = nullptr;
-    for (int s = 0; s < MAX_PID_SLOTS; ++s) {
-        if (dev.pid_usages[s].pid == pid) return &dev.pid_usages[s];
-        if (dev.pid_usages[s].pid == 0 && empty == nullptr)
-            empty = &dev.pid_usages[s];
-    }
-    if (empty) {
-        empty->pid = pid;
-        empty->used_bytes = 0;
-    }
-    return empty;
-}
-
-PidUsage* SharedQuotaManager::findPidSlotLocked(int dev_id, pid_t pid) {
-    if (!hdr_) return nullptr;
-    if (dev_id < 0 || dev_id >= hdr_->num_devices) return nullptr;
-    SharedDeviceEntry& dev = hdr_->devices[dev_id];
-    for (int s = 0; s < MAX_PID_SLOTS; ++s) {
-        if (dev.pid_usages[s].pid == pid) return &dev.pid_usages[s];
-    }
-    return nullptr;
-}
-
-Status SharedQuotaManager::attachProcess() {
-    if (!hdr_) return Status::InvalidArgument("not attached");
-
-    int rc = lock();
-    if (rc != 0) {
-        return Status::InternalError("failed to lock shared mutex: " +
-                                     std::string(strerror(rc)));
+    } else {
+        created_ = false;
     }
 
-    auto topo = local_quota_->getTopology();
-    for (size_t i = 0; i < topo->getNicCount(); ++i) {
-        if (topo->getNicType(i) != Topology::NIC_RDMA) continue;
-        auto dev_name = topo->getNicName(i);
-        if (findDeviceIdByNameLocked(dev_name) >= 0) continue;
-        int empty_idx = -1;
-        for (int j = 0; j < MAX_DEVICES; ++j) {
-            if (hdr_->devices[j].dev_name[0] == '\0') {
-                empty_idx = j;
-                break;
-            }
-        }
-        if (empty_idx < 0) continue;
-        strncpy(hdr_->devices[empty_idx].dev_name, dev_name.c_str(), 56);
-        hdr_->devices[empty_idx].active_bytes = 0;
-    }
+    startBackgroundThread();
 
-    int count = 0;
-    for (int i = 0; i < MAX_DEVICES; ++i)
-        if (hdr_->devices[i].dev_name[0] != '\0') ++count;
-    hdr_->num_devices = count;
-
-    unlock();
     return Status::OK();
 }
 
-Status SharedQuotaManager::detachProcess() { return Status::OK(); }
+Status SharedSlotManager::detach() {
+    stopBackgroundThread();
 
-Status SharedQuotaManager::diffusion() {
-    if (!hdr_) return Status::InvalidArgument("not attached");
-    pid_t pid = getpid();
-    int rc = lock();
-    if (rc != 0)
-        return Status::InternalError("lock failed: " +
-                                     std::string(strerror(rc)));
-    for (int d = 0; d < hdr_->num_devices; ++d) {
-        std::string dev_name = hdr_->devices[d].dev_name;
-        auto dev_id = local_quota_->getTopology()->getNicId(dev_name);
-        if (dev_name.empty() || dev_id < 0) continue;
-        PidUsage* slot = findOrCreatePidSlotLocked(dev_id, pid);
-        if (!slot) {
-            unlock();
-            return Status::InternalError("no free pid slot for device");
-        }
-        auto used_bytes = local_quota_->getActiveBytes(dev_id);
-        slot->used_bytes = used_bytes;
-        uint64_t sum = 0;
-        for (int s = 0; s < MAX_PID_SLOTS; ++s)
-            sum += hdr_->devices[d].pid_usages[s].used_bytes;
-        uint64_t diffusion_active_bytes =
-            sum < used_bytes ? 0 : sum - used_bytes;
-        hdr_->devices[d].active_bytes = sum;
-        local_quota_->setDiffusionActiveBytes(dev_id, diffusion_active_bytes);
+    if (hdr_) {
+        munmap(hdr_, size_);
+        hdr_ = nullptr;
     }
-    unlock();
+
+    if (fd_ >= 0) {
+        close(fd_);
+        fd_ = -1;
+    }
+
     return Status::OK();
 }
 
+bool SharedSlotManager::canSend(int priority) {
+    if (!hdr_) return true;
+
+    // Get current global slot
+    int current_slot = hdr_->current_slot.load(std::memory_order_acquire);
+
+    // Check if the given priority is allowed in current global slot
+    // Slot 0: HIGH only (priority 0 <= 0)
+    // Slot 1: HIGH + MEDIUM (priority 0 or 1 <= 1)
+    // Slot 2: ALL (priority 0, 1, or 2 <= 2)
+    return isPriorityAllowedInSlot(priority, current_slot);
+}
+
+void SharedSlotManager::startBackgroundThread() {
+    if (background_running_.exchange(true)) return;
+
+    background_thread_ = std::thread([this]() { backgroundThreadLoop(); });
+}
+
+void SharedSlotManager::stopBackgroundThread() {
+    if (!background_running_.exchange(false)) return;
+
+    if (background_thread_.joinable()) {
+        background_thread_.join();
+    }
+}
+
+// Background thread: advance global slot periodically
+void SharedSlotManager::backgroundThreadLoop() {
+    const uint64_t SLEEP_INTERVAL_US = 1000;  // 1ms
+
+    while (background_running_.load(std::memory_order_relaxed)) {
+        // Calculate base slot from time
+        uint64_t now = getCurrentTimeInNano();
+        uint64_t base_slot = now / (rotation_interval_ms_ * 1000000ull);
+
+        pthread_mutex_lock(&hdr_->global_mutex);
+        // Update global slot
+        int global_slot = static_cast<int>(base_slot % NUM_SLOTS);
+        hdr_->current_slot.store(global_slot, std::memory_order_release);
+        pthread_mutex_unlock(&hdr_->global_mutex);
+
+        usleep(SLEEP_INTERVAL_US);
+    }
+}
+
 }  // namespace tent
 }  // namespace mooncake
diff --git a/mooncake-transfer-engine/tent/src/transport/rdma/workers.cpp b/mooncake-transfer-engine/tent/src/transport/rdma/workers.cpp
index 9191b41b..9b9d2bc3 100644
--- a/mooncake-transfer-engine/tent/src/transport/rdma/workers.cpp
+++ b/mooncake-transfer-engine/tent/src/transport/rdma/workers.cpp
@@ -19,6 +19,7 @@
 #include <cassert>
 
 #include "tent/transport/rdma/endpoint_store.h"
+#include "tent/transport/rdma/shared_quota.h"
 #include "tent/common/utils/ip.h"
 #include "tent/common/utils/string_builder.h"
 #include "tent/common/utils/os.h"
@@ -27,25 +28,125 @@
 namespace mooncake {
 namespace tent {
 thread_local int tl_wid = -1;
+
+namespace {
+// Look up (or create) the RailMonitor for `machine_id` on this worker's
+// map. Returning a stable reference is safe because the map stores values
+// via unique_ptr -- rehashes move the pointer slot, not the RailMonitor.
+RailMonitor& getOrCreateRail(
+    std::unordered_map<std::string, std::unique_ptr<RailMonitor>>& rails,
+    const std::string& machine_id) {
+    auto it = rails.find(machine_id);
+    if (it != rails.end()) return *it->second;
+    auto [ins, _] = rails.emplace(machine_id, std::make_unique<RailMonitor>());
+    return *ins->second;
+}
+}  // namespace
+
 Workers::Workers(RdmaTransport* transport)
     : transport_(transport), num_workers_(0), running_(false) {
-    device_quota_ = std::make_unique<DeviceQuota>();
-    device_quota_->loadTopology(transport_->local_topology_);
+    device_selector_ = std::make_unique<DeviceSelector>();
+    device_selector_->loadTopology(transport_->local_topology_);
     auto& conf = transport_->conf_;
+
+    // ============================================================
+    // Core Scheduling Configuration
+    // ============================================================
+
+    // Enable/disable smart scheduling (false = simple round-robin)
+    bool enable_smart_scheduling =
+        conf->get("transports/rdma/enable_smart_scheduling", true);
+    device_selector_->setSmartSelection(enable_smart_scheduling);
+
+    // ============================================================
+    // NUMA Distance Penalties
+    // Higher values = higher penalty for cross-NUMA access
+    // Format: [local_numa, remote_numa1, remote_numa2, ...]
+    // ============================================================
+    DeviceSelector::SchedulingParams params;
+
+    auto numa_penalties =
+        conf->get("transports/rdma/numa_penalties", std::vector<double>{});
+    if (numa_penalties.size() == Topology::DevicePriorityRanks) {
+        for (size_t i = 0; i < Topology::DevicePriorityRanks; ++i) {
+            params.numa_tier_weights[i] = numa_penalties[i];
+        }
+    }
+
+    // ============================================================
+    // Bandwidth Estimation (EWMA)
+    // ============================================================
+
+    // Learning rate: 0.0 = full adaptation, 1.0 = no adaptation
+    params.bandwidth_learning_rate =
+        conf->get("transports/rdma/bandwidth_learning_rate", 0.01);
+
+    // EWMA bounds as multipliers of theoretical bandwidth
+    params.ewma_min_multiplier =
+        conf->get("transports/rdma/ewma_min_bandwidth_multiplier", 0.1);
+    params.ewma_max_multiplier =
+        conf->get("transports/rdma/ewma_max_bandwidth_multiplier", 10.0);
+
+    // ============================================================
+    // Device Selection Scoring
+    // ============================================================
+
+    // Random jitter to avoid deterministic selection
+    params.score_jitter_range =
+        conf->get("transports/rdma/score_jitter_range", 1e-9);
+
+    // Small value to prevent division by zero
+    params.score_epsilon = conf->get("transports/rdma/score_epsilon", 1e-12);
+
+    // ============================================================
+    // Priority-Based Filtering
+    // ============================================================
+
+    params.enable_priority_filtering =
+        conf->get("transports/rdma/enable_priority_filtering", true);
+
+    // Local device priority rotation interval (microseconds)
+    params.local_rotation_interval_us =
+        conf->get("transports/rdma/local_rotation_interval_us", 200);
+
+    // ============================================================
+    // Priority Promotion (Anti-Starvation)
+    // ============================================================
+
+    // Timeout after which low-priority requests get promoted (nanoseconds)
+    // Default: 10ms (10000000 ns)
+    priority_promotion_timeout_ns_ =
+        conf->get("transports/rdma/priority_promotion_timeout_us", 10000) *
+        1000ull;
+
+    // ============================================================
+    // Global Slot Coordination (Multi-Process)
+    // ============================================================
+
+    params.slot_rotation_interval_ms =
+        conf->get("transports/rdma/slot_rotation_interval_ms", 2);
+
+    // ============================================================
+    // Bandwidth Constants (Gbps)
+    // ============================================================
+
+    params.default_bandwidth_gbps =
+        conf->get("transports/rdma/default_bandwidth_gbps", 400.0);
+    params.min_bandwidth_gbps =
+        conf->get("transports/rdma/min_bandwidth_gbps", 10.0);
+    params.max_bandwidth_gbps =
+        conf->get("transports/rdma/max_bandwidth_gbps", 800.0);
+
+    device_selector_->setSchedulingParams(params);
+
+    // ============================================================
+    // Shared Memory Configuration
+    // ============================================================
+
     auto shared_quota_shm_path =
         conf->get("transports/rdma/shared_quota_shm_path", "");
     if (!shared_quota_shm_path.empty())
-        device_quota_->enableSharedQuota(shared_quota_shm_path);
-    auto cross_numa_access =
-        conf->get("transports/rdma/cross_numa_access", false);
-    device_quota_->setCrossNumaAccess(cross_numa_access);
-    auto local_weight = conf->get("transports/rdma/local_weight", 1.0);
-    device_quota_->setLocalWeight(local_weight);
-    auto learning_rate = conf->get("transports/rdma/learning_rate", 0.1);
-    device_quota_->setLearningRate(learning_rate);
-    auto diffusion_interval =
-        conf->get("transports/rdma/diffusion_interval", 10);
-    device_quota_->setDiffusionInterval(diffusion_interval);
+        device_selector_->enableSharedQuota(shared_quota_shm_path);
 }
 
 Workers::~Workers() {
@@ -102,7 +203,14 @@ Status Workers::submit(RdmaSliceList& slice_list, int worker_id) {
         }
     }
     auto& worker = worker_context_[worker_id];
-    worker.queue.push(slice_list);
+
+    // Get priority from first slice (all slices in list have same priority)
+    int priority = PRIO_HIGH;
+    if (slice_list.first && slice_list.first->task) {
+        priority = slice_list.first->priority;
+    }
+
+    worker.queues[priority].push(slice_list);
     if (!worker.inflight_slices.fetch_add(slice_list.num_slices)) {
         std::lock_guard<std::mutex> lock(worker.mutex);
         if (worker.in_suspend) worker.cv.notify_all();
@@ -122,7 +230,8 @@ Status Workers::cancel(RdmaSliceList& slice_list) {
 }
 
 std::shared_ptr<RdmaEndPoint> Workers::getEndpoint(Workers::PostPath path) {
-    std::string rpc_server_addr, target_seg_name, target_dev_name;
+    std::string rpc_server_addr, target_seg_name, target_dev_name,
+        target_nic_path_name;
     RouteHint hint;
     auto& segment_manager = transport_->metadata_->segmentManager();
     auto target_id = path.remote_segment_id;
@@ -140,6 +249,7 @@ std::shared_ptr<RdmaEndPoint> Workers::getEndpoint(Workers::PostPath path) {
                 rpc_server_addr = segment->rpc_server_addr;
             }
             target_seg_name = segment->name;
+            target_nic_path_name = segment->nicPathServerName();
             target_dev_name = hint.topo->getNicName(device_id);
             if (target_seg_name.empty() || target_dev_name.empty()) {
                 return Status::NeedsRefreshCache(
@@ -160,12 +270,8 @@ std::shared_ptr<RdmaEndPoint> Workers::getEndpoint(Workers::PostPath path) {
                          // connection unavailable
     }
     std::shared_ptr<RdmaEndPoint> endpoint;
-    auto peer_name = MakeNicPath(target_seg_name, target_dev_name);
+    auto peer_name = MakeNicPath(target_nic_path_name, target_dev_name);
     endpoint = context->endpointStore()->getOrInsert(peer_name);
-    if (endpoint && endpoint->status() == RdmaEndPoint::EP_RESET) {
-        context->endpointStore()->remove(endpoint.get());
-        endpoint = context->endpointStore()->getOrInsert(peer_name);
-    }
     if (!endpoint) {
         LOG(ERROR) << "Cannot allocate endpoint " << peer_name;
         return nullptr;
@@ -188,30 +294,32 @@ std::shared_ptr<RdmaEndPoint> Workers::getEndpoint(Workers::PostPath path) {
 }
 
 void Workers::disableEndpoint(RdmaSlice* slice) {
-    SegmentDesc* desc = nullptr;
-    auto& segment_manager = transport_->metadata_->segmentManager();
-    auto target_id = slice->task->request.target_id;
-    if (target_id == LOCAL_SEGMENT_ID) {
-        desc = segment_manager.getLocal().get();
-    } else {
-        auto status = segment_manager.getRemoteCached(desc, target_id);
-        if (!status.ok()) return;
-    }
-    if (desc) {
-        auto& worker = worker_context_[tl_wid];
-        auto& rail = worker.rails[desc->machine_id];
-        rail.markFailed(slice->source_dev_id, slice->target_dev_id);
+    if (auto* rail = slice->rail_monitor) {
+        rail->markFailed(slice->source_dev_id, slice->target_dev_id);
     }
     if (auto ep = slice->ep_weak_ptr.lock()) {
         ep->acknowledge(slice, FAILED);
-        ep->reset();
+        ep->resetConnection("Endpoint failed");
     }
 }
 
 void Workers::asyncPostSend() {
     auto& worker = worker_context_[tl_wid];
     std::vector<RdmaSliceList> result;
-    worker.queue.pop(result);
+
+    auto shared_quota =
+        device_selector_ ? device_selector_->getSharedSlotManager() : nullptr;
+
+    // Promote timed-out low priority requests
+    promoteTimedOutRequests(worker);
+
+    // Priority selection: HIGH -> MEDIUM -> LOW
+    for (int prio = PRIO_HIGH; prio < kNumPriorityLevels; ++prio) {
+        if (shared_quota && !shared_quota->canSend(prio)) continue;
+        worker.queues[prio].pop(result);
+        if (!result.empty()) break;
+    }
+
     for (auto& slice_list : result) {
         if (slice_list.num_slices == 0) continue;
         auto slice = slice_list.first;
@@ -282,6 +390,49 @@ void Workers::asyncPostSend() {
     }
 }
 
+void Workers::promoteTimedOutRequests(WorkerContext& worker) {
+    uint64_t current_ts = getCurrentTimeInNano();
+    if (current_ts < worker.next_promotion_check_ns) return;
+
+    // Set next check time (1ms from now)
+    worker.next_promotion_check_ns = current_ts + 1000000ull;
+
+    // Check MEDIUM -> HIGH promotion
+    std::vector<RdmaSliceList> promoted;
+    worker.queues[PRIO_MEDIUM].pop(promoted);
+    if (!promoted.empty()) {
+        auto* slice = promoted.front().first;
+        if (slice && slice->enqueue_ts > 0 &&
+            (current_ts - slice->enqueue_ts) >=
+                priority_promotion_timeout_ns_) {
+            for (auto& slice_list : promoted) {
+                worker.queues[PRIO_HIGH].push(slice_list);
+            }
+            return;
+        }
+        for (auto& slice_list : promoted) {
+            worker.queues[PRIO_MEDIUM].push(slice_list);
+        }
+    }
+
+    // Check LOW -> MEDIUM promotion
+    worker.queues[PRIO_LOW].pop(promoted);
+    if (!promoted.empty()) {
+        auto* slice = promoted.front().first;
+        if (slice && slice->enqueue_ts > 0 &&
+            (current_ts - slice->enqueue_ts) >=
+                priority_promotion_timeout_ns_) {
+            for (auto& slice_list : promoted) {
+                worker.queues[PRIO_MEDIUM].push(slice_list);
+            }
+            return;
+        }
+        for (auto& slice_list : promoted) {
+            worker.queues[PRIO_LOW].push(slice_list);
+        }
+    }
+}
+
 void Workers::asyncPollCq() {
     auto& worker = worker_context_[tl_wid];
     const static size_t kPollCount = 64;
@@ -327,8 +478,8 @@ void Workers::asyncPollCq() {
             double inflight_lat = (poll_ts - slice->submit_ts) / 1000.0;
             double overall_lat_sec = (poll_ts - slice->enqueue_ts) / 1e9;
             if (slice->retry_count == 0) {
-                device_quota_->release(slice->source_dev_id, slice->length,
-                                       overall_lat_sec);
+                device_selector_->release(slice->source_dev_id, slice->length,
+                                          overall_lat_sec);
             }
             if (slice->word != PENDING) continue;
             if (!ep) {
@@ -361,6 +512,14 @@ void Workers::asyncPollCq() {
                 }
             } else {
                 num_slices += ep->acknowledge(slice, COMPLETED);
+                // A successful transfer proves this rail is healthy; clear
+                // any accumulated error count so a previously-cooled-down
+                // rail can be used again without waiting for the full
+                // cooldown to expire. The monitor pointer is resolved once
+                // in generatePostPath, so no map lookup is needed here.
+                if (auto* rail = slice->rail_monitor; rail && rail->ready())
+                    rail->markRecovered(slice->source_dev_id,
+                                        slice->target_dev_id);
                 worker.perf.inflight_lat.add(inflight_lat);
                 worker.perf.enqueue_lat.add(enqueue_lat);
             }
@@ -449,7 +608,25 @@ int Workers::handleContextEvents(std::shared_ptr<RdmaContext>& context) {
 }
 
 void Workers::monitorThread() {
+    // Track time for periodic endpoint reclaim (1 Hz heartbeat)
+    auto last_reclaim_time = std::chrono::steady_clock::now();
+
     while (running_) {
+        // Periodic endpoint reclaim: runs every 1 second to drain waiting_list_
+        // Under failure load, insertions stall but endpoints still need cleanup
+        auto current_time = std::chrono::steady_clock::now();
+        auto time_since_last_reclaim =
+            std::chrono::duration_cast<std::chrono::milliseconds>(
+                current_time - last_reclaim_time)
+                .count();
+
+        if (time_since_last_reclaim >= 1000) {  // 1 second = 1000 ms
+            for (auto& context : transport_->context_set_) {
+                context->endpointStore()->reclaim();
+            }
+            last_reclaim_time = current_time;
+        }
+
         for (auto& context : transport_->context_set_) {
             struct epoll_event event;
             if (context->eventFd() < 0) continue;
@@ -524,7 +701,7 @@ Status Workers::selectOptimalDevice(RouteHint& source, RouteHint& target,
                                     RdmaSlice* slice) {
     auto& worker = worker_context_[tl_wid];
     if (slice->source_dev_id < 0) {
-        CHECK_STATUS(device_quota_->allocate(
+        CHECK_STATUS(device_selector_->allocate(
             slice->length, source.buffer->location, slice->source_dev_id));
     }
 
@@ -532,9 +709,10 @@ Status Workers::selectOptimalDevice(RouteHint& source, RouteHint& target,
         return Status::DeviceNotFound(
             "No device could access the slice memory region" LOC_MARK);
 
-    auto& rail = worker.rails[target.segment->machine_id];
+    auto& rail = getOrCreateRail(worker.rails, target.segment->machine_id);
     if (!rail.ready() || target.topo != rail.remote())
-        rail.load(source.topo, target.topo);
+        rail.load(source.topo, target.topo, /*rail_topo_json=*/"",
+                  transport_->conf_.get());
     if (slice->target_dev_id < 0) {
         int mapped_dev_id = rail.findBestRemoteDevice(
             slice->source_dev_id, target.topo_entry->numa_node);
@@ -630,7 +808,8 @@ Status Workers::selectFallbackDevice(RouteHint& source, RouteHint& target,
             reachable = (sdev == tdev);  // loopback is safe
         } else {
             auto& worker = worker_context_[tl_wid];
-            auto& rail = worker.rails[target.segment->machine_id];
+            auto& rail =
+                getOrCreateRail(worker.rails, target.segment->machine_id);
             reachable = rail.available(sdev, tdev);
         }
 
@@ -663,6 +842,11 @@ Status Workers::generatePostPath(RdmaSlice* slice) {
         CHECK_STATUS(selectFallbackDevice(source, target, slice));
     slice->source_lkey = source.buffer->lkey[slice->source_dev_id];
     slice->target_rkey = target.buffer->rkey[slice->target_dev_id];
+    // Cache the RailMonitor pointer so asyncPollCq / disableEndpoint can
+    // update rail state without a segment lookup or string-keyed map
+    // lookup on the hot path.
+    slice->rail_monitor = &getOrCreateRail(worker_context_[tl_wid].rails,
+                                           target.segment->machine_id);
     return Status::OK();
 }
 }  // namespace tent
diff --git a/mooncake-transfer-engine/tent/src/transport/shm/shm_transport.cpp b/mooncake-transfer-engine/tent/src/transport/shm/shm_transport.cpp
index 930df078..ccd737ee 100644
--- a/mooncake-transfer-engine/tent/src/transport/shm/shm_transport.cpp
+++ b/mooncake-transfer-engine/tent/src/transport/shm/shm_transport.cpp
@@ -132,6 +132,7 @@ void ShmTransport::startTransfer(ShmTask *task, ShmSubBatch *batch) {
     } else {
         task->status_word = TransferStatusEnum::FAILED;
     }
+    batch->notifyProgress();
 }
 
 Status ShmTransport::getTransferStatus(SubBatchRef batch, int task_id,
diff --git a/mooncake-transfer-engine/tent/src/transport/sunrise_link/CMakeLists.txt b/mooncake-transfer-engine/tent/src/transport/sunrise_link/CMakeLists.txt
new file mode 100644
index 00000000..c0bfd594
--- /dev/null
+++ b/mooncake-transfer-engine/tent/src/transport/sunrise_link/CMakeLists.txt
@@ -0,0 +1,21 @@
+# Copyright 2024 KVCache.AI
+#
+# Licensed under the Apache License, Version 2.0 (the "License"); you may not
+# use this file except in compliance with the License. You may obtain a copy of
+# the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+# License for the specific language governing permissions and limitations under
+# the License.
+
+if(USE_SUNRISE)
+  add_library(tent_xport_sunrise_link STATIC sunrise_link_transport.cpp)
+  target_include_directories(tent_xport_sunrise_link
+                              PUBLIC ${MC_TANGRT_ROOT}/include/)
+  target_link_libraries(tent_xport_sunrise_link
+                        PUBLIC tent_common glog glog::glog ${SUNRISE_LIBRARIES})
+endif()
diff --git a/mooncake-transfer-engine/tent/src/transport/sunrise_link/sunrise_link_transport.cpp b/mooncake-transfer-engine/tent/src/transport/sunrise_link/sunrise_link_transport.cpp
new file mode 100644
index 00000000..0b5b959b
--- /dev/null
+++ b/mooncake-transfer-engine/tent/src/transport/sunrise_link/sunrise_link_transport.cpp
@@ -0,0 +1,1098 @@
+// Copyright 2024 KVCache.AI
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "tent/transport/sunrise_link/sunrise_link_transport.h"
+
+#include <dlfcn.h>
+#include <glog/logging.h>
+#include <unistd.h>
+
+#include <cassert>
+#include <cstddef>
+#include <cstdint>
+#include <cstdlib>
+#include <algorithm>
+#include <cstring>
+#include <memory>
+#include <mutex>
+#include <string>
+#include <vector>
+
+#include "tent/common/status.h"
+#include "tent/runtime/slab.h"
+#include "tent/runtime/control_plane.h"
+#include "tent/runtime/topology.h"
+#include "tent/common/utils/string_builder.h"
+
+#include <tang_runtime_api.h>
+#include <ptml.h>
+
+namespace mooncake {
+namespace tent {
+
+namespace {
+
+// Default matches typical Tang RT install; override root and/or arch via env.
+static std::string TangRtSharedObjectPath(const char* soname) {
+    const char* lib_dir = std::getenv("MC_TANGRT_LIB_DIR");
+    if (lib_dir && lib_dir[0]) {
+        std::string d(lib_dir);
+        while (!d.empty() && d.back() == '/') d.pop_back();
+        return d + "/" + soname;
+    }
+    const char* root = std::getenv("MC_TANGRT_ROOT");
+    std::string r = (root && root[0]) ? std::string(root)
+                                      : std::string("/usr/local/tangrt");
+    while (!r.empty() && r.back() == '/') r.pop_back();
+    const char* arch = std::getenv("MC_TANGRT_LIB_ARCH");
+    std::string arch_dir =
+        (arch && arch[0]) ? std::string(arch) : std::string("linux-x86_64");
+    std::string arch_path = r + "/lib/" + arch_dir + "/" + soname;
+    if (access(arch_path.c_str(), R_OK) == 0) return arch_path;
+    return r + "/lib/" + soname;
+}
+
+constexpr int kMaxPhytopoPorts = 10;
+
+static tangMemcpyKind MemcpyKindForPointers(tangMemoryType st,
+                                            tangMemoryType dt) {
+    if (st == tangMemoryTypeDevice && dt == tangMemoryTypeHost) {
+        return tangMemcpyDeviceToHost;
+    }
+    if (st == tangMemoryTypeHost && dt == tangMemoryTypeDevice) {
+        return tangMemcpyHostToDevice;
+    }
+    if (st == tangMemoryTypeDevice && dt == tangMemoryTypeDevice) {
+        return tangMemcpyDeviceToDevice;
+    }
+    return tangMemcpyHostToHost;
+}
+
+// Sunrise S2 path may allocate mapped GPU memory via tangHostAlloc.
+// For kernel/device-copy style paths, convert mapped host pointer into its
+// corresponding device pointer when available.
+static void NormalizeMappedHostPointer(void** ptr,
+                                       tangPointerAttributes* attr) {
+    if (!ptr || !*ptr || !attr) return;
+    if (attr->type != tangMemoryTypeHost) return;
+
+    void* dev_ptr = nullptr;
+    tangError_t ret = tangHostGetDevicePointer(&dev_ptr, *ptr, 0);
+    if (ret != tangSuccess || !dev_ptr) return;
+
+    tangPointerAttributes dev_attr = {};
+    tangPointerGetAttributes(&dev_attr, dev_ptr);
+    if (dev_attr.type == tangMemoryTypeDevice) {
+        *ptr = dev_ptr;
+        *attr = dev_attr;
+    }
+}
+
+struct SavedTangDevice {
+    int dev{-1};
+    SavedTangDevice() { tangGetDevice(&dev); }
+    ~SavedTangDevice() {
+        if (dev >= 0) tangSetDevice(dev);
+    }
+    SavedTangDevice(const SavedTangDevice&) = delete;
+    SavedTangDevice& operator=(const SavedTangDevice&) = delete;
+};
+
+static tangError_t QueryPointerAttrsBestEffort(void* ptr,
+                                               tangPointerAttributes* attr,
+                                               int preferred_dev) {
+    if (!ptr || !attr) return tangErrorInvalidValue;
+    SavedTangDevice saved;
+
+    if (preferred_dev >= 0) {
+        tangSetDevice(preferred_dev);
+    }
+
+    tangError_t ret = tangPointerGetAttributes(attr, ptr);
+    if (ret == tangSuccess) return ret;
+
+    int dev_count = 0;
+    if (tangGetDeviceCount(&dev_count) != tangSuccess || dev_count <= 0) {
+        return ret;
+    }
+    for (int d = 0; d < dev_count; ++d) {
+        if (d == preferred_dev) continue;
+        tangSetDevice(d);
+        ret = tangPointerGetAttributes(attr, ptr);
+        if (ret == tangSuccess) return ret;
+    }
+    return ret;
+}
+
+// Fallback when tangMemcpyPeer_v2 fails: map peer device memory through PTLink.
+static tangError_t C2cMemcpyWithPeerMap(int src_dev, int dst_dev,
+                                        const void* src, void* dst, size_t len,
+                                        int port_src_to_dst,
+                                        int port_dst_to_src) {
+    tangError_t ret = tangErrorUnknown;
+    void* mapped = nullptr;
+
+    if (port_src_to_dst >= 0) {
+        ret = tangDeviceGetPeerPointer(src_dev, port_src_to_dst, dst, &mapped);
+        if (ret != tangSuccess) {
+            ret = tangSetDevice(src_dev);
+            if (ret != tangSuccess) return ret;
+            ret = tangDeviceGetPeerPointer(src_dev, port_src_to_dst, dst,
+                                           &mapped);
+        }
+        if (ret == tangSuccess) {
+            ret = tangMemcpy(mapped, src, len, tangMemcpyDeviceToDevice);
+        }
+        return ret;
+    }
+    if (port_dst_to_src >= 0) {
+        ret = tangDeviceGetPeerPointer(dst_dev, port_dst_to_src,
+                                       const_cast<void*>(src), &mapped);
+        if (ret != tangSuccess) {
+            ret = tangSetDevice(dst_dev);
+            if (ret != tangSuccess) return ret;
+            ret = tangDeviceGetPeerPointer(dst_dev, port_dst_to_src,
+                                           const_cast<void*>(src), &mapped);
+        }
+        if (ret == tangSuccess) {
+            ret = tangMemcpy(dst, mapped, len, tangMemcpyDeviceToDevice);
+        }
+        return ret;
+    }
+    return tangErrorInvalidValue;
+}
+
+}  // namespace
+
+static ptmlReturn_t (*ptmlInit)() = nullptr;
+static ptmlReturn_t (*ptmlPtlinkEnableAll)() = nullptr;
+static ptmlReturn_t (*ptmlPtlinkPhytopoDetect)(ptmlDevice_t dev, int len,
+                                               void* data) = nullptr;
+static ptmlReturn_t (*ptmlDeviceGetCount)(int* devCount) = nullptr;
+
+// Tang memcpy must NOT use one process-wide lock: it serialized every
+// tangMemcpyPeer_v2 and made multi-threaded C2C benches ~3 orders of magnitude
+// slower than PCCL (which runs many concurrent device copies). We only guard
+// tangSetDevice + non-peer tangMemcpy with per-device mutexes; peer copies use
+// device ids passed into tangMemcpyPeer_v2 and run without that serialization.
+
+namespace {
+
+constexpr int kMaxTangDevLocks = 32;
+
+static std::mutex& devMutex(int dev) {
+    static std::mutex table[kMaxTangDevLocks];
+    if (dev < 0 || dev >= kMaxTangDevLocks) {
+        static std::mutex fallback;
+        return fallback;
+    }
+    return table[dev];
+}
+
+// Serialize tangSetDevice + context-bound ops for a device; for peer pairs,
+// lock both devices in a fixed order to avoid deadlock.
+struct ScopedTangDeviceOps {
+    int lo{-1};
+    int hi{-1};
+
+    explicit ScopedTangDeviceOps(int src_dev, int dst_dev) {
+        if (src_dev < 0 && dst_dev < 0) {
+            return;
+        }
+        if (src_dev < 0 || dst_dev < 0) {
+            int d = src_dev >= 0 ? src_dev : dst_dev;
+            devMutex(d).lock();
+            lo = d;
+            return;
+        }
+        if (src_dev == dst_dev) {
+            devMutex(src_dev).lock();
+            lo = src_dev;
+            return;
+        }
+        lo = std::min(src_dev, dst_dev);
+        hi = std::max(src_dev, dst_dev);
+        devMutex(lo).lock();
+        devMutex(hi).lock();
+    }
+
+    ~ScopedTangDeviceOps() {
+        if (hi >= 0) {
+            devMutex(hi).unlock();
+            devMutex(lo).unlock();
+        } else if (lo >= 0) {
+            devMutex(lo).unlock();
+        }
+    }
+};
+
+}  // namespace
+
+class ThreadLocalTangStreamPool {
+   public:
+    tangStream_t getOrCreate(int device_id) {
+        if (device_id < 0) return nullptr;
+        auto it = streams_by_device_.find(device_id);
+        if (it != streams_by_device_.end()) return it->second;
+
+        ScopedTangDeviceOps ctx_lock(device_id, device_id);
+        int saved_dev = -1;
+        tangGetDevice(&saved_dev);
+        tangError_t set_ret = tangSetDevice(device_id);
+        if (set_ret != tangSuccess) {
+            LOG(ERROR) << "tangSetDevice(" << device_id
+                       << ") failed for stream creation: " << set_ret << " "
+                       << tangGetErrorString(set_ret);
+            return nullptr;
+        }
+
+        tangStream_t stream = nullptr;
+        tangError_t create_ret =
+            tangStreamCreateWithFlags(&stream, tangStreamNonBlocking);
+        if (saved_dev >= 0) tangSetDevice(saved_dev);
+        if (create_ret != tangSuccess) {
+            LOG(ERROR) << "tangStreamCreateWithFlags failed: " << create_ret
+                       << " " << tangGetErrorString(create_ret)
+                       << ", device=" << device_id;
+            return nullptr;
+        }
+        streams_by_device_[device_id] = stream;
+        return stream;
+    }
+
+    ~ThreadLocalTangStreamPool() {
+        for (const auto& kv : streams_by_device_) {
+            if (kv.second) tangStreamDestroy(kv.second);
+        }
+    }
+
+   private:
+    std::unordered_map<int, tangStream_t> streams_by_device_;
+};
+
+thread_local ThreadLocalTangStreamPool tl_stream_pool;
+
+SunriseLinkTransport::SunriseLinkTransport() : installed_(false) {
+    // Match NVLink-style caps: routing may see MTYPE_CPU if pointer probe fails
+    // while the remote segment is VRAM; dram_to_gpu must still be available.
+    caps.dram_to_dram = true;
+    caps.dram_to_gpu = true;
+    caps.gpu_to_dram = true;
+    caps.gpu_to_gpu = true;
+}
+
+SunriseLinkTransport::~SunriseLinkTransport() { uninstall(); }
+
+Status SunriseLinkTransport::install(std::string& local_segment_name,
+                                     std::shared_ptr<ControlService> metadata,
+                                     std::shared_ptr<Topology> local_topology,
+                                     std::shared_ptr<Config> conf) {
+    if (installed_) {
+        return Status::InvalidArgument(
+            "SunriseLink transport has been installed" LOC_MARK);
+    }
+
+    conf_ = conf;
+    metadata_ = metadata;
+
+    auto status = initSunriseLink();
+    if (!status.ok()) {
+        LOG(WARNING) << "SunriseLink transport initSunriseLink fail";
+        return status;
+    }
+
+    status = detectTopology();
+    if (!status.ok()) {
+        LOG(WARNING) << "SunriseLink transport detectTopology fail";
+        return status;
+    }
+
+    installed_ = true;
+    async_memcpy_threshold_ =
+        conf_
+            ? conf_->get("transports/sunrise_link/async_memcpy_threshold", 0) *
+                  1024
+            : 0;
+    LOG(INFO) << "SunriseLink transport installed successfully";
+
+    return Status::OK();
+}
+
+Status SunriseLinkTransport::initSunriseLink() {
+    ptml_handle_ =
+        dlopen(TangRtSharedObjectPath("libptml_shared.so").c_str(), RTLD_NOW);
+    if (!ptml_handle_) {
+        char* error = dlerror();
+        LOG(ERROR) << "Failed to load libptml_shared.so: "
+                   << (error ? error : "unknown error");
+        return Status::InternalError("Failed to load PTML library");
+    }
+
+#define LOAD_PTML_SYM(handle, symbol, funcptr)                       \
+    do {                                                             \
+        void* tmp = dlsym(handle, symbol);                           \
+        if (tmp == nullptr) {                                        \
+            LOG(ERROR) << "dlsym failed on " << symbol;              \
+            return Status::InternalError("Failed to load " #symbol); \
+        }                                                            \
+        *(void**)&funcptr = tmp;                                     \
+    } while (0)
+
+    LOAD_PTML_SYM(ptml_handle_, "ptmlInit", ptmlInit);
+    LOAD_PTML_SYM(ptml_handle_, "ptmlPtlinkEnableAll", ptmlPtlinkEnableAll);
+    LOAD_PTML_SYM(ptml_handle_, "ptmlPtlinkPhytopoDetect",
+                  ptmlPtlinkPhytopoDetect);
+    LOAD_PTML_SYM(ptml_handle_, "ptmlDeviceGetCount", ptmlDeviceGetCount);
+
+    ptmlReturn_t ret = ptmlInit();
+    if (ret != PTML_SUCCESS) {
+        LOG(ERROR) << "ptmlInit failed, error code: " << ret;
+        return Status::InternalError("Failed to initialize PTML");
+    }
+
+    ret = ptmlPtlinkEnableAll();
+    if (ret != PTML_SUCCESS) {
+        LOG(WARNING) << "ptmlPtlinkEnableAll failed, error code: " << ret;
+    }
+
+    runtime_lib_handle_ =
+        dlopen(TangRtSharedObjectPath("libtangrt_shared.so").c_str(), RTLD_NOW);
+    if (!runtime_lib_handle_) {
+        char* error = dlerror();
+        LOG(ERROR) << "Failed to load libtangrt_shared.so: "
+                   << (error ? error : "unknown error");
+        return Status::InternalError("Failed to load Tang library");
+    }
+
+    LOG(INFO) << "SunriseLink runtime initialized successfully (PTML + Tang)";
+    return Status::OK();
+}
+
+Status SunriseLinkTransport::detectTopology() {
+    topology_map_.clear();
+    c2c_port_by_pair_.clear();
+    local_devices_.clear();
+
+    int device_count = 0;
+    ptmlReturn_t pret = ptmlDeviceGetCount(&device_count);
+    if (pret != PTML_SUCCESS || device_count <= 0) {
+        LOG(ERROR) << "ptmlDeviceGetCount failed, ret=" << pret
+                   << " count=" << device_count;
+        return Status::InternalError("ptmlDeviceGetCount failed");
+    }
+
+    std::vector<ptPhyTopo_t> phytopo(static_cast<size_t>(device_count) *
+                                     kMaxPhytopoPorts);
+
+    for (int dev = 0; dev < device_count; ++dev) {
+        memset(phytopo.data() + static_cast<size_t>(dev) * kMaxPhytopoPorts, 0,
+               kMaxPhytopoPorts * sizeof(ptPhyTopo_t));
+
+        pret = ptmlPtlinkPhytopoDetect(
+            dev, kMaxPhytopoPorts * sizeof(ptPhyTopo_t),
+            phytopo.data() + static_cast<size_t>(dev) * kMaxPhytopoPorts);
+        if (pret != PTML_SUCCESS) {
+            LOG(WARNING) << "ptmlPtlinkPhytopoDetect failed for device " << dev
+                         << " ret=" << pret;
+            continue;
+        }
+
+        SunriseLinkDeviceInfo info;
+        info.device_id = dev;
+
+        ptPhyTopo_t* row =
+            phytopo.data() + static_cast<size_t>(dev) * kMaxPhytopoPorts;
+        for (int p = 0; p < kMaxPhytopoPorts; ++p) {
+            if (!row[p].link_status) {
+                continue;
+            }
+            int lc = static_cast<int>(row[p].local_chipid);
+            int rc = static_cast<int>(row[p].remote_chipid);
+            int lp = static_cast<int>(row[p].local_port);
+            auto key = std::make_pair(lc, rc);
+            if (c2c_port_by_pair_.find(key) == c2c_port_by_pair_.end()) {
+                c2c_port_by_pair_[key] = lp;
+            }
+
+            SunriseLinkPeerInfo peer;
+            peer.peer_device_id = rc;
+            peer.peer_port = lp;
+            peer.is_active = true;
+            topology_map_[dev].push_back(peer);
+            info.active_ports.push_back(lp);
+
+            LOG(INFO) << "PTLink topo dev=" << dev << " port=" << lp << " link "
+                      << lc << " -> " << rc << " (remote_port="
+                      << static_cast<int>(row[p].remote_port) << ")";
+        }
+
+        local_devices_.push_back(std::move(info));
+    }
+
+    LOG(INFO) << "SunriseLink topology: " << device_count << " device(s), "
+              << c2c_port_by_pair_.size() << " C2C port mapping(s)";
+    return Status::OK();
+}
+
+Status SunriseLinkTransport::allocateSubBatch(SubBatchRef& batch,
+                                              size_t max_size) {
+    auto shm_batch = Slab<SunriseLinkSubBatch>::Get().allocate();
+    if (!shm_batch) {
+        return Status::InternalError(
+            "Unable to allocate SunriseLink sub-batch");
+    }
+    batch = shm_batch;
+    // Slab reuse can leave stale tasks; sub_task_id must index a fresh list.
+    shm_batch->task_list.clear();
+    shm_batch->task_list.reserve(max_size);
+    shm_batch->max_size = max_size;
+    // Synchronous copy path (can_async=false in startTransfer) does not rely on
+    // stream completion; avoid creating thread-local Tang stream here to reduce
+    // device-context coupling in worker threads.
+    shm_batch->stream = nullptr;
+    return Status::OK();
+}
+
+Status SunriseLinkTransport::freeSubBatch(SubBatchRef& batch) {
+    auto shm_batch = dynamic_cast<SunriseLinkSubBatch*>(batch);
+    if (!shm_batch) {
+        return Status::InvalidArgument(
+            "Invalid SunriseLink sub-batch" LOC_MARK);
+    }
+    shm_batch->task_list.clear();
+    Slab<SunriseLinkSubBatch>::Get().deallocate(shm_batch);
+    batch = nullptr;
+    return Status::OK();
+}
+
+Status SunriseLinkTransport::submitTransferTasks(
+    SubBatchRef batch, const std::vector<Request>& request_list) {
+    auto shm_batch = dynamic_cast<SunriseLinkSubBatch*>(batch);
+    if (!shm_batch) {
+        return Status::InvalidArgument(
+            "Invalid SunriseLink sub-batch" LOC_MARK);
+    }
+
+    if (request_list.size() + shm_batch->task_list.size() >
+        shm_batch->max_size) {
+        return Status::TooManyRequests("Exceed batch capacity" LOC_MARK);
+    }
+
+    auto infer_local_dev = [&](void* ptr) -> int {
+        uint64_t p = reinterpret_cast<uint64_t>(ptr);
+        std::lock_guard<std::mutex> guard(registered_memory_mutex_);
+        for (const auto& entry : registered_memory_) {
+            uint64_t base = reinterpret_cast<uint64_t>(entry.first);
+            if (base <= p && p < base + entry.second) {
+                auto it = registered_memory_gpu_id_.find(entry.first);
+                if (it != registered_memory_gpu_id_.end()) return it->second;
+            }
+        }
+        return -1;
+    };
+    auto infer_local_dev_best_effort = [&](void* ptr) -> int {
+        int dev = infer_local_dev(ptr);
+        if (dev >= 0) return dev;
+
+        tangPointerAttributes attr = {};
+        int saved_dev = -1;
+        tangGetDevice(&saved_dev);
+        tangError_t ret = QueryPointerAttrsBestEffort(ptr, &attr, -1);
+        if (saved_dev >= 0) tangSetDevice(saved_dev);
+        if (ret != tangSuccess) return -1;
+
+        NormalizeMappedHostPointer(&ptr, &attr);
+        if (attr.type == tangMemoryTypeDevice) return attr.device;
+        return -1;
+    };
+
+    for (auto& request : request_list) {
+        shm_batch->task_list.push_back(SunriseLinkTask{});
+        auto& task = shm_batch->task_list[shm_batch->task_list.size() - 1];
+
+        uint64_t target_addr = request.target_offset;
+        int target_gpu_id = -1;
+        int local_gpu_id = infer_local_dev_best_effort(request.source);
+        if (local_gpu_id >= 0) {
+            tangError_t sd = tangSetDevice(local_gpu_id);
+            if (sd != tangSuccess) {
+                return Status::InternalError(
+                    "Failed to set local Sunrise device before transfer");
+            }
+        }
+        if (request.target_id != LOCAL_SEGMENT_ID) {
+            auto status = relocateRemoteAddress(
+                target_addr, request.length, request.target_id, &target_gpu_id);
+            if (!status.ok()) return status;
+        }
+        task.target_addr = target_addr;
+        task.request = request;
+        task.status_word = TransferStatusEnum::PENDING;
+        task.transferred_bytes = 0;
+        task.is_tang_ipc = (request.target_id != LOCAL_SEGMENT_ID);
+        task.tang_gpu_id = target_gpu_id;
+
+        auto status = startTransfer(&task, shm_batch);
+        if (!status.ok()) return status;
+    }
+
+    return Status::OK();
+}
+
+Status SunriseLinkTransport::startTransfer(SunriseLinkTask* task,
+                                           SunriseLinkSubBatch* batch) {
+    void* src = nullptr;
+    void* dst = nullptr;
+
+    const uint64_t resolved_target = task->target_addr;
+
+    if (task->request.opcode == Request::READ) {
+        dst = task->request.source;
+        src = reinterpret_cast<void*>(resolved_target);
+    } else {
+        src = task->request.source;
+        dst = reinterpret_cast<void*>(resolved_target);
+    }
+
+    const size_t len = task->request.length;
+    tangStream_t stream = nullptr;
+    int current_dev = -1;
+    tangGetDevice(&current_dev);
+    auto infer_local_dev = [&](void* ptr) -> int {
+        uint64_t p = reinterpret_cast<uint64_t>(ptr);
+        std::lock_guard<std::mutex> guard(registered_memory_mutex_);
+        for (const auto& entry : registered_memory_) {
+            uint64_t base = reinterpret_cast<uint64_t>(entry.first);
+            if (base <= p && p < base + entry.second) {
+                auto it = registered_memory_gpu_id_.find(entry.first);
+                if (it != registered_memory_gpu_id_.end()) return it->second;
+            }
+        }
+        return -1;
+    };
+    auto infer_local_dev_best_effort = [&](void* ptr) -> int {
+        int dev = infer_local_dev(ptr);
+        if (dev >= 0) return dev;
+
+        tangPointerAttributes attr = {};
+        int saved_dev = -1;
+        tangGetDevice(&saved_dev);
+        tangError_t ret = QueryPointerAttrsBestEffort(ptr, &attr, -1);
+        if (saved_dev >= 0) tangSetDevice(saved_dev);
+        if (ret != tangSuccess) return -1;
+
+        NormalizeMappedHostPointer(&ptr, &attr);
+        if (attr.type == tangMemoryTypeDevice) return attr.device;
+        return -1;
+    };
+
+    // Follow ipcC2cTest style: for IPC path, prefer explicit gpu ids and
+    // tangMemcpyPeer directly, instead of relying on pointer attributes.
+    if (task->is_tang_ipc && task->tang_gpu_id >= 0) {
+        int remote_dev = task->tang_gpu_id;
+        int local_dev = -1;
+        if (task->request.opcode == Request::WRITE) {
+            local_dev = infer_local_dev_best_effort(src);
+        } else {
+            local_dev = infer_local_dev_best_effort(dst);
+        }
+        if (local_dev < 0) local_dev = current_dev;
+
+        int src_dev =
+            (task->request.opcode == Request::WRITE) ? local_dev : remote_dev;
+        int dst_dev =
+            (task->request.opcode == Request::WRITE) ? remote_dev : local_dev;
+
+        // Same GPU id on initiator and target (typical single-card P2P bench):
+        // pointers are IPC handles into the same device context; tangMemcpyPeer
+        // is not valid when src_dev == dst_dev.
+        if (src_dev == dst_dev) {
+            ScopedTangDeviceOps ctx_lock(src_dev, dst_dev);
+            tangSetDevice(src_dev);
+            tangError_t ret =
+                tangMemcpy(dst, src, len, tangMemcpyDeviceToDevice);
+            if (ret != tangSuccess) {
+                LOG(ERROR)
+                    << "SunriseLink IPC same-device memcpy failed, tang ret="
+                    << ret << " " << tangGetErrorString(ret)
+                    << " dev=" << src_dev;
+                task->status_word = TransferStatusEnum::FAILED;
+                return Status::InternalError("SunriseLink transfer failed");
+            }
+            task->transferred_bytes = len;
+            task->status_word = TransferStatusEnum::COMPLETED;
+            return Status::OK();
+        }
+
+        tangError_t ret = tangSuccess;
+        // Prefer ipcC2cTest ordering first (current=dst_dev), then retry with
+        // current=src_dev since some Tang stacks are context-sensitive for peer
+        // copy direction.
+        auto try_peer_copy = [&]() -> tangError_t {
+            tangError_t r = tangMemcpyPeer(dst, dst_dev, src, src_dev, len);
+            if (r != tangSuccess) {
+                r = tangMemcpyPeer_v2(dst, dst_dev, src, src_dev, len);
+            }
+            return r;
+        };
+        // ipcC2cTest uses synchronous tangMemcpyPeer only.
+        ret = try_peer_copy();
+        if (ret != tangSuccess) {
+            int p_sd = -1;
+            int p_ds = -1;
+            auto it_sd =
+                c2c_port_by_pair_.find(std::make_pair(src_dev, dst_dev));
+            if (it_sd != c2c_port_by_pair_.end()) p_sd = it_sd->second;
+            auto it_ds =
+                c2c_port_by_pair_.find(std::make_pair(dst_dev, src_dev));
+            if (it_ds != c2c_port_by_pair_.end()) p_ds = it_ds->second;
+            ret = C2cMemcpyWithPeerMap(src_dev, dst_dev, src, dst, len, p_sd,
+                                       p_ds);
+        }
+        if (ret != tangSuccess) {
+            LOG(ERROR) << "SunriseLink IPC peer copy failed, tang ret=" << ret
+                       << " " << tangGetErrorString(ret)
+                       << " src_dev=" << src_dev << " dst_dev=" << dst_dev
+                       << " remote_dev=" << remote_dev
+                       << " local_dev=" << local_dev;
+            task->status_word = TransferStatusEnum::FAILED;
+            return Status::InternalError("SunriseLink transfer failed");
+        }
+        task->transferred_bytes = len;
+        task->status_word = TransferStatusEnum::COMPLETED;
+        return Status::OK();
+    }
+
+    tangPointerAttributes sa = {};
+    tangPointerAttributes da = {};
+    if (task->tang_gpu_id >= 0) tangSetDevice(task->tang_gpu_id);
+
+    // Intentionally no mutex here: a global lock on pointer-attribute queries
+    // serialized all peer copies across threads and capped throughput at a few
+    // MiB/s in multi-threaded benches. Tang is expected to tolerate concurrent
+    // attribute reads for different pointers.
+    tangError_t src_attr_ret =
+        QueryPointerAttrsBestEffort(src, &sa, task->tang_gpu_id);
+    tangError_t dst_attr_ret =
+        QueryPointerAttrsBestEffort(dst, &da, task->tang_gpu_id);
+    NormalizeMappedHostPointer(&src, &sa);
+    NormalizeMappedHostPointer(&dst, &da);
+
+    int src_dev = (sa.type == tangMemoryTypeDevice) ? sa.device : -1;
+    int dst_dev = (da.type == tangMemoryTypeDevice) ? da.device : -1;
+    int remote_dev_hint = task->tang_gpu_id;
+    if (task->request.opcode == Request::WRITE) {
+        if (dst_dev < 0 && remote_dev_hint >= 0) dst_dev = remote_dev_hint;
+        if (src_dev < 0 && current_dev >= 0) src_dev = current_dev;
+    } else {
+        if (src_dev < 0 && remote_dev_hint >= 0) src_dev = remote_dev_hint;
+        if (dst_dev < 0 && current_dev >= 0) dst_dev = current_dev;
+    }
+    if (task->request.opcode == Request::WRITE && src_dev < 0) {
+        int hinted = infer_local_dev(src);
+        if (hinted >= 0) src_dev = hinted;
+    } else if (task->request.opcode == Request::READ && dst_dev < 0) {
+        int hinted = infer_local_dev(dst);
+        if (hinted >= 0) dst_dev = hinted;
+    }
+
+    tangError_t ret = tangSuccess;
+    const int stream_dev = (dst_dev >= 0) ? dst_dev : src_dev;
+    if (stream_dev >= 0) {
+        stream = tl_stream_pool.getOrCreate(stream_dev);
+    }
+    const bool can_async =
+        stream && async_memcpy_threshold_ > 0 && len >= async_memcpy_threshold_;
+    batch->stream = reinterpret_cast<void*>(stream);
+
+    if (task->is_tang_ipc && src_dev >= 0 && dst_dev >= 0 &&
+        src_dev == dst_dev) {
+        ScopedTangDeviceOps ctx_lock(src_dev, dst_dev);
+        tangSetDevice(src_dev);
+        if (can_async) {
+            ret = tangMemcpyAsync(dst, src, len, tangMemcpyDeviceToDevice,
+                                  stream);
+            if (ret == tangSuccess) {
+                return Status::OK();
+            }
+        }
+        ret = tangMemcpy(dst, src, len, tangMemcpyDeviceToDevice);
+        if (ret != tangSuccess) {
+            LOG(ERROR)
+                << "SunriseLink IPC local-device memcpy failed, tang ret="
+                << ret << " " << tangGetErrorString(ret) << " dev=" << src_dev;
+            task->status_word = TransferStatusEnum::FAILED;
+            return Status::InternalError("SunriseLink transfer failed");
+        }
+        task->transferred_bytes = len;
+        task->status_word = TransferStatusEnum::COMPLETED;
+        return Status::OK();
+    }
+
+    if (src_dev >= 0 && dst_dev >= 0 &&
+        (src_dev != dst_dev || task->is_tang_ipc)) {
+        auto collect_ports = [&](int local, int peer) {
+            std::vector<int> ports;
+            auto it = topology_map_.find(local);
+            if (it == topology_map_.end()) return ports;
+            for (const auto& p : it->second) {
+                if (p.peer_device_id == peer && p.is_active) {
+                    ports.push_back(p.peer_port);
+                }
+            }
+            std::sort(ports.begin(), ports.end());
+            ports.erase(std::unique(ports.begin(), ports.end()), ports.end());
+            return ports;
+        };
+
+        auto ports_sd = collect_ports(src_dev, dst_dev);
+        auto ports_ds = collect_ports(dst_dev, src_dev);
+        if (ports_sd.empty()) {
+            auto it = c2c_port_by_pair_.find(std::make_pair(src_dev, dst_dev));
+            if (it != c2c_port_by_pair_.end()) ports_sd.push_back(it->second);
+        }
+        if (ports_ds.empty()) {
+            auto it = c2c_port_by_pair_.find(std::make_pair(dst_dev, src_dev));
+            if (it != c2c_port_by_pair_.end()) ports_ds.push_back(it->second);
+        }
+        constexpr size_t kMultiPortThreshold = 4ull << 20;  // 4 MiB
+        const bool enable_port_striping =
+            conf_ ? conf_->get("transports/sunrise_link/enable_port_striping",
+                               false)
+                  : false;
+        if (enable_port_striping && !ports_sd.empty() && ports_sd.size() > 1 &&
+            len >= kMultiPortThreshold) {
+            size_t lanes = ports_sd.size();
+            size_t offset = 0;
+            const size_t base = len / lanes;
+            const size_t rem = len % lanes;
+            for (size_t lane = 0; lane < lanes; ++lane) {
+                size_t piece = base + (lane < rem ? 1 : 0);
+                size_t lane_offset = offset;
+                offset += piece;
+                int p_sd = ports_sd[lane];
+                ret = C2cMemcpyWithPeerMap(
+                    src_dev, dst_dev, static_cast<char*>(src) + lane_offset,
+                    static_cast<char*>(dst) + lane_offset, piece, p_sd, -1);
+                if (ret != tangSuccess) {
+                    break;
+                }
+            }
+            if (ret == tangSuccess) {
+                task->transferred_bytes = len;
+                task->status_word = TransferStatusEnum::COMPLETED;
+                return Status::OK();
+            }
+        }
+
+        auto peer_sync = [&]() -> tangError_t {
+            tangError_t r = tangMemcpyPeer(dst, dst_dev, src, src_dev, len);
+            if (r != tangSuccess) {
+                r = tangMemcpyPeer_v2(dst, dst_dev, src, src_dev, len);
+            }
+            if (r != tangSuccess) {
+                int p_sd = !ports_sd.empty() ? ports_sd.front() : -1;
+                int p_ds = !ports_ds.empty() ? ports_ds.front() : -1;
+                r = C2cMemcpyWithPeerMap(src_dev, dst_dev, src, dst, len, p_sd,
+                                         p_ds);
+            }
+            return r;
+        };
+
+        // Synchronous tangMemcpyPeer only (see can_async above).
+        (void)can_async;
+        ret = peer_sync();
+    } else if (src_dev >= 0 && dst_dev >= 0 && src_dev == dst_dev) {
+        ScopedTangDeviceOps ctx_lock(src_dev, dst_dev);
+        tangSetDevice(src_dev);
+        if (can_async) {
+            ret = tangMemcpyAsync(dst, src, len, tangMemcpyDeviceToDevice,
+                                  stream);
+            if (ret == tangSuccess) {
+                return Status::OK();
+            }
+        }
+        ret = tangMemcpy(dst, src, len, tangMemcpyDeviceToDevice);
+    } else {
+        ScopedTangDeviceOps ctx_lock(src_dev, dst_dev);
+        if (src_dev >= 0) {
+            tangSetDevice(src_dev);
+        } else if (dst_dev >= 0) {
+            tangSetDevice(dst_dev);
+        }
+        tangMemcpyKind kind = MemcpyKindForPointers(sa.type, da.type);
+        if (can_async) {
+            ret = tangMemcpyAsync(dst, src, len, kind, stream);
+            if (ret == tangSuccess) {
+                return Status::OK();
+            }
+        }
+        ret = tangMemcpy(dst, src, len, kind);
+    }
+
+    if (ret != tangSuccess) {
+        LOG(ERROR) << "SunriseLink transfer failed, tang ret=" << ret << " "
+                   << tangGetErrorString(ret) << " src_dev=" << src_dev
+                   << " dst_dev=" << dst_dev << " src_attr_ret=" << src_attr_ret
+                   << " dst_attr_ret=" << dst_attr_ret
+                   << " task.tang_gpu_id=" << task->tang_gpu_id;
+        task->status_word = TransferStatusEnum::FAILED;
+        return Status::InternalError("SunriseLink transfer failed");
+    }
+
+    task->transferred_bytes = len;
+    task->status_word = TransferStatusEnum::COMPLETED;
+    return Status::OK();
+}
+
+Status SunriseLinkTransport::getTransferStatus(SubBatchRef batch, int task_id,
+                                               TransferStatus& status) {
+    auto shm_batch = dynamic_cast<SunriseLinkSubBatch*>(batch);
+    if (!shm_batch) {
+        return Status::InvalidArgument(
+            "Invalid SunriseLink sub-batch" LOC_MARK);
+    }
+    if (task_id < 0 || task_id >= (int)shm_batch->task_list.size()) {
+        return Status::InvalidArgument("Invalid task id" LOC_MARK);
+    }
+    auto& task = shm_batch->task_list[task_id];
+    if (task.status_word == TransferStatusEnum::PENDING) {
+        tangStream_t stream = reinterpret_cast<tangStream_t>(shm_batch->stream);
+        if (stream) {
+            tangError_t err = tangStreamQuery(stream);
+            if (err == tangSuccess) {
+                tangStreamSynchronize(stream);
+                task.transferred_bytes = task.request.length;
+                task.status_word = TransferStatusEnum::COMPLETED;
+            } else if (err != tangErrorNotReady) {
+                task.status_word = TransferStatusEnum::FAILED;
+            }
+        } else {
+            LOG(ERROR)
+                << "SunriseLink getTransferStatus: PENDING with null stream; "
+                   "cannot complete async bookkeeping";
+            task.status_word = TransferStatusEnum::FAILED;
+        }
+    }
+    status = TransferStatus{task.status_word, task.transferred_bytes};
+    return Status::OK();
+}
+
+Status SunriseLinkTransport::addMemoryBuffer(BufferDesc& desc,
+                                             const MemoryOptions& options) {
+    LocationParser location(desc.location);
+    auto status =
+        registerMemory(reinterpret_cast<void*>(desc.addr), desc.length);
+    if (!status.ok()) return status;
+    if (location.type() == "cuda" || location.type() == "npu") {
+        {
+            std::lock_guard<std::mutex> guard(registered_memory_mutex_);
+            registered_memory_gpu_id_[reinterpret_cast<void*>(desc.addr)] =
+                location.index();
+        }
+        int saved_dev = -1;
+        tangGetDevice(&saved_dev);
+        if (location.index() >= 0) {
+            tangError_t sd = tangSetDevice(location.index());
+            if (sd != tangSuccess) {
+                LOG(ERROR)
+                    << "tangSetDevice before tangIpcGetMemHandle failed: " << sd
+                    << " " << tangGetErrorString(sd);
+                return Status::InternalError("tangSetDevice failed");
+            }
+        }
+        tangIpcMemHandle_t handle;
+        auto ret =
+            tangIpcGetMemHandle(&handle, reinterpret_cast<void*>(desc.addr));
+        if (saved_dev >= 0) {
+            tangSetDevice(saved_dev);
+        }
+        if (ret != tangSuccess) {
+            LOG(ERROR) << "tangIpcGetMemHandle failed: " << ret << " "
+                       << tangGetErrorString(ret);
+            return Status::InternalError("tangIpcGetMemHandle failed");
+        }
+        desc.shm_path =
+            serializeBinaryData(&handle, sizeof(tangIpcMemHandle_t));
+    } else if (location.type() == "cpu" ||
+               location.type() == kWildcardLocation) {
+        std::lock_guard<std::mutex> guard(registered_memory_mutex_);
+        registered_memory_gpu_id_[reinterpret_cast<void*>(desc.addr)] = -1;
+    }
+    desc.transports.push_back(TransportType::SUNRISE_LINK);
+    return Status::OK();
+}
+
+Status SunriseLinkTransport::registerMemory(void* addr, size_t size) {
+    std::lock_guard<std::mutex> guard(registered_memory_mutex_);
+    registered_memory_[addr] = size;
+    return Status::OK();
+}
+
+Status SunriseLinkTransport::deregisterMemory(void* addr) {
+    std::lock_guard<std::mutex> guard(registered_memory_mutex_);
+    auto it = registered_memory_.find(addr);
+    if (it != registered_memory_.end()) {
+        registered_memory_.erase(it);
+        registered_memory_gpu_id_.erase(addr);
+    }
+    return Status::OK();
+}
+
+Status SunriseLinkTransport::removeMemoryBuffer(BufferDesc& desc) {
+    desc.shm_path.clear();
+    return deregisterMemory(reinterpret_cast<void*>(desc.addr));
+}
+
+Status SunriseLinkTransport::relocateRemoteAddress(uint64_t& dest_addr,
+                                                   uint64_t length,
+                                                   uint64_t target_id,
+                                                   int* target_gpu_id) {
+    thread_local IpcRelocateMap tl_relocate_map;
+    if (tl_relocate_map.empty()) {
+        RWSpinlock::ReadGuard guard(relocate_lock_);
+        tl_relocate_map = relocate_map_;
+    }
+
+    for (auto& entry : tl_relocate_map[target_id]) {
+        if (entry.first <= dest_addr &&
+            dest_addr + length <= entry.first + entry.second.length) {
+            dest_addr = dest_addr - entry.first +
+                        reinterpret_cast<uint64_t>(entry.second.dev_ptr);
+            if (target_gpu_id) *target_gpu_id = entry.second.gpu_id;
+            return Status::OK();
+        }
+    }
+
+    RWSpinlock::WriteGuard guard(relocate_lock_);
+    SegmentDesc* desc = nullptr;
+    auto status = metadata_->segmentManager().getRemoteCached(desc, target_id);
+    if (!status.ok()) return status;
+
+    auto buffer = desc->findBuffer(dest_addr, length);
+    if (!buffer || buffer->shm_path.empty()) {
+        // Cache may be stale right after remote registration; refresh once.
+        metadata_->segmentManager().invalidateRemote(target_id);
+        status = metadata_->segmentManager().getRemoteCached(desc, target_id);
+        if (!status.ok()) return status;
+        buffer = desc->findBuffer(dest_addr, length);
+    }
+    if (!buffer || buffer->shm_path.empty()) {
+        return Status::InvalidArgument(
+            "Requested address is not in registered Sunrise buffer" LOC_MARK);
+    }
+
+    if (!relocate_map_[target_id].count(buffer->addr)) {
+        LocationParser location(buffer->location);
+        if (location.type() != "cuda" && location.type() != "npu") {
+            return Status::InvalidArgument(
+                "Requested address is not in registered CUDA/NPU "
+                "buffer" LOC_MARK);
+        }
+        std::vector<unsigned char> decoded;
+        deserializeBinaryData(buffer->shm_path, decoded);
+        if (decoded.size() != sizeof(tangIpcMemHandle_t)) {
+            return Status::InvalidArgument(
+                "Invalid Sunrise IPC handle size" LOC_MARK);
+        }
+        tangIpcMemHandle_t handle;
+        memcpy(&handle, decoded.data(), sizeof(handle));
+
+        void* opened_ptr = nullptr;
+        int saved_dev = -1;
+        tangGetDevice(&saved_dev);
+        tangError_t sd = tangSetDevice(location.index());
+        if (sd != tangSuccess) {
+            LOG(ERROR) << "tangSetDevice before tangIpcOpenMemHandle failed: "
+                       << sd << " " << tangGetErrorString(sd)
+                       << ", target_gpu=" << location.index()
+                       << ", saved_dev=" << saved_dev;
+            return Status::InternalError("tangSetDevice failed");
+        }
+        tangError_t ret = tangIpcOpenMemHandle(&opened_ptr, handle,
+                                               tangIpcMemLazyEnablePeerAccess);
+        if (saved_dev >= 0) {
+            tangSetDevice(saved_dev);
+        }
+        if (ret != tangSuccess) {
+            LOG(ERROR) << "tangIpcOpenMemHandle failed: " << ret << " "
+                       << tangGetErrorString(ret);
+            return Status::InternalError("tangIpcOpenMemHandle failed");
+        }
+
+        OpenedIpcEntry ipc_entry;
+        ipc_entry.dev_ptr = opened_ptr;
+        ipc_entry.length = buffer->length;
+        ipc_entry.gpu_id = location.index();
+        relocate_map_[target_id][buffer->addr] = ipc_entry;
+        tl_relocate_map = relocate_map_;
+    }
+
+    auto opened_ptr = relocate_map_[target_id][buffer->addr].dev_ptr;
+    if (target_gpu_id) {
+        *target_gpu_id = relocate_map_[target_id][buffer->addr].gpu_id;
+    }
+    dest_addr =
+        dest_addr - buffer->addr + reinterpret_cast<uint64_t>(opened_ptr);
+    return Status::OK();
+}
+
+Status SunriseLinkTransport::uninstall() {
+    if (installed_) {
+        std::vector<void*> reg_addrs;
+        {
+            std::lock_guard<std::mutex> guard(registered_memory_mutex_);
+            reg_addrs.reserve(registered_memory_.size());
+            for (auto& entry : registered_memory_) {
+                reg_addrs.push_back(entry.first);
+            }
+        }
+        for (void* addr : reg_addrs) {
+            auto status = deregisterMemory(addr);
+            if (!status.ok()) return status;
+        }
+        {
+            std::lock_guard<std::mutex> guard(registered_memory_mutex_);
+            registered_memory_.clear();
+            registered_memory_gpu_id_.clear();
+        }
+
+        if (runtime_lib_handle_) {
+            dlclose(runtime_lib_handle_);
+            runtime_lib_handle_ = nullptr;
+        }
+
+        for (auto& seg_entry : relocate_map_) {
+            for (auto& entry : seg_entry.second) {
+                if (entry.second.dev_ptr) {
+                    tangIpcCloseMemHandle(entry.second.dev_ptr);
+                }
+            }
+        }
+        relocate_map_.clear();
+
+        if (ptml_handle_) {
+            dlclose(ptml_handle_);
+            ptml_handle_ = nullptr;
+        }
+
+        metadata_.reset();
+        installed_ = false;
+        LOG(INFO) << "SunriseLink transport uninstalled";
+    }
+    return Status::OK();
+}
+
+}  // namespace tent
+}  // namespace mooncake
diff --git a/mooncake-transfer-engine/tent/src/transport/tcp/tcp_transport.cpp b/mooncake-transfer-engine/tent/src/transport/tcp/tcp_transport.cpp
index 10465015..973173ca 100644
--- a/mooncake-transfer-engine/tent/src/transport/tcp/tcp_transport.cpp
+++ b/mooncake-transfer-engine/tent/src/transport/tcp/tcp_transport.cpp
@@ -95,8 +95,7 @@ Status TcpTransport::install(std::string &local_segment_name,
         return 0;
     });
     caps.dram_to_dram = true;
-    if (Platform::getLoader().type() == "cuda" ||
-        Platform::getLoader().type() == "cann") {
+    if (Platform::getLoader().type() != "cpu") {
         caps.dram_to_gpu = true;
         caps.gpu_to_dram = true;
         caps.gpu_to_gpu = true;
@@ -153,6 +152,8 @@ Status TcpTransport::submitTransferTasks(
         tcp_batch->task_list.emplace_back();
         auto &task = tcp_batch->task_list.back();
         task.request = request;
+        task.progress_batch_id = batch->progress_batch_id;
+        task.notify_progress = batch->notify_progress;
         task.status_word.store(TransferStatusEnum::PENDING,
                                std::memory_order_release);
     }
@@ -212,6 +213,7 @@ void TcpTransport::startTransfer(TcpTask *task) {
         task->status_word.store(TransferStatusEnum::FAILED,
                                 std::memory_order_release);
     }
+    if (task->notify_progress) task->notify_progress(task->progress_batch_id);
 }
 
 Status TcpTransport::doTransferWithRetry(TcpTask *task) {
diff --git a/mooncake-transfer-engine/tent/tests/CMakeLists.txt b/mooncake-transfer-engine/tent/tests/CMakeLists.txt
index bd13e5ef..c460e416 100644
--- a/mooncake-transfer-engine/tent/tests/CMakeLists.txt
+++ b/mooncake-transfer-engine/tent/tests/CMakeLists.txt
@@ -20,6 +20,26 @@ target_include_directories(metrics_config_loader_test
                            PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/../include)
 add_test(NAME metrics_config_loader_test COMMAND metrics_config_loader_test)
 
+add_executable(admission_queue_test admission_queue_test.cpp
+                                    ../src/runtime/admission_queue.cpp)
+target_link_libraries(admission_queue_test PRIVATE tent_common gtest gtest_main)
+target_include_directories(admission_queue_test
+                           PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/../include)
+add_test(NAME admission_queue_test COMMAND admission_queue_test)
+
+add_executable(tent_ip_utils_test ip_utils_test.cpp)
+target_link_libraries(tent_ip_utils_test PRIVATE tent_common gtest gtest_main)
+target_include_directories(tent_ip_utils_test
+                           PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/../include)
+add_test(NAME tent_ip_utils_test COMMAND tent_ip_utils_test)
+
+add_executable(tent_coalesce_regions_test coalesce_regions_test.cpp)
+target_link_libraries(tent_coalesce_regions_test PRIVATE tent_common gtest
+                                                         gtest_main)
+target_include_directories(tent_coalesce_regions_test
+                           PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/../include)
+add_test(NAME tent_coalesce_regions_test COMMAND tent_coalesce_regions_test)
+
 add_executable(request_merge_test request_merge_test.cpp)
 target_link_libraries(request_merge_test PRIVATE gtest gtest_main
                                                  tent_link_group)
@@ -58,9 +78,98 @@ add_executable(tent_endpoint_lifecycle_test endpoint_lifecycle_test.cpp)
 target_link_libraries(tent_endpoint_lifecycle_test PRIVATE gtest gtest_main)
 add_test(NAME tent_endpoint_lifecycle_test COMMAND tent_endpoint_lifecycle_test)
 
+if(USE_HIP)
+  find_package(HIP REQUIRED)
+  add_executable(tent_rocm_platform_test rocm_platform_test.cpp)
+  target_link_libraries(tent_rocm_platform_test PRIVATE gtest gtest_main
+                                                        tent_link_group
+                                                        hip::host)
+  target_include_directories(tent_rocm_platform_test
+                             PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/../include)
+  add_test(NAME tent_rocm_platform_test COMMAND tent_rocm_platform_test)
+endif()
+
+if(USE_SUNRISE)
+  add_executable(tent_sunrise_link_transport_test
+                 sunrise_link_transport_test.cpp)
+  target_link_libraries(tent_sunrise_link_transport_test PRIVATE gtest
+                                                                 gtest_main
+                                                                 tent_link_group)
+  target_include_directories(
+    tent_sunrise_link_transport_test
+    PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/../include
+            ${CMAKE_CURRENT_SOURCE_DIR}/../../include)
+  add_test(NAME tent_sunrise_link_transport_test
+           COMMAND tent_sunrise_link_transport_test)
+endif()
 add_executable(tent_fault_proxy_test fault_proxy_test.cpp)
 target_link_libraries(tent_fault_proxy_test PRIVATE gtest gtest_main
                                                      tent_link_group)
 target_include_directories(tent_fault_proxy_test
                            PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/../include)
 add_test(NAME tent_fault_proxy_test COMMAND tent_fault_proxy_test)
+
+add_executable(tent_rail_monitor_test rail_monitor_test.cpp)
+target_link_libraries(tent_rail_monitor_test PRIVATE gtest gtest_main
+                                                      tent_link_group)
+target_include_directories(tent_rail_monitor_test
+                           PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/../include)
+add_test(NAME tent_rail_monitor_test COMMAND tent_rail_monitor_test)
+
+# Transport Selector Unit Test
+add_executable(tent_transport_selector_test transport_selector_test.cpp)
+target_link_libraries(tent_transport_selector_test PRIVATE gtest gtest_main
+                                                            tent_link_group)
+target_include_directories(tent_transport_selector_test
+                           PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/../include)
+add_test(NAME tent_transport_selector_test COMMAND tent_transport_selector_test)
+
+# End-to-end failover test: drives real TransferEngineImpl with
+# FaultProxyTransport-wrapped fakes to exercise resubmitTransferTask.
+add_executable(tent_engine_failover_e2e_test engine_failover_e2e_test.cpp)
+target_link_libraries(tent_engine_failover_e2e_test
+                      PRIVATE gtest gtest_main tent_link_group)
+if(TARGET asio_shared)
+  target_link_libraries(tent_engine_failover_e2e_test PRIVATE asio_shared)
+endif()
+target_include_directories(tent_engine_failover_e2e_test
+                           PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/../include)
+add_test(NAME tent_engine_failover_e2e_test
+         COMMAND tent_engine_failover_e2e_test)
+
+# Per-request transport_hint: validates submitTransfer parameter, routing,
+# disabled-transport rejection, out-of-range rejection, mixed-hint batches.
+add_executable(tent_transport_hint_test transport_hint_test.cpp)
+target_link_libraries(tent_transport_hint_test
+                      PRIVATE gtest gtest_main tent_link_group)
+if(TARGET asio_shared)
+  target_link_libraries(tent_transport_hint_test PRIVATE asio_shared)
+endif()
+target_include_directories(tent_transport_hint_test
+                           PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/../include)
+add_test(NAME tent_transport_hint_test COMMAND tent_transport_hint_test)
+
+# ProgressWorker skeleton test: covers default-off behavior, event-driven
+# progress without poll-failover, and freeBatch races (issue #2116).
+add_executable(tent_progress_worker_test progress_worker_test.cpp)
+target_link_libraries(tent_progress_worker_test
+                      PRIVATE gtest gtest_main tent_link_group)
+if(TARGET asio_shared)
+  target_link_libraries(tent_progress_worker_test PRIVATE asio_shared)
+endif()
+target_include_directories(tent_progress_worker_test
+                           PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/../include)
+add_test(NAME tent_progress_worker_test
+         COMMAND tent_progress_worker_test)
+
+add_executable(tent_runtime_queue_dispatch_test
+               runtime_queue_dispatch_test.cpp)
+target_link_libraries(tent_runtime_queue_dispatch_test
+                      PRIVATE gtest gtest_main tent_link_group)
+if(TARGET asio_shared)
+  target_link_libraries(tent_runtime_queue_dispatch_test PRIVATE asio_shared)
+endif()
+target_include_directories(tent_runtime_queue_dispatch_test
+                           PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/../include)
+add_test(NAME tent_runtime_queue_dispatch_test
+         COMMAND tent_runtime_queue_dispatch_test)
diff --git a/mooncake-transfer-engine/tent/tests/admission_queue_test.cpp b/mooncake-transfer-engine/tent/tests/admission_queue_test.cpp
new file mode 100644
index 00000000..b26ff889
--- /dev/null
+++ b/mooncake-transfer-engine/tent/tests/admission_queue_test.cpp
@@ -0,0 +1,393 @@
+// Copyright 2026 KVCache.AI
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "tent/runtime/admission_queue.h"
+
+#include <utility>
+#include <vector>
+
+#include <gtest/gtest.h>
+
+namespace mooncake {
+namespace tent {
+namespace {
+
+QueueOwnerInput makeOwner(
+    size_t public_task_id, size_t length,
+    QueueOwnerKind kind = QueueOwnerKind::User,
+    std::vector<size_t> derived_task_ids = std::vector<size_t>()) {
+    QueueOwnerInput owner;
+    owner.owner_task_id = public_task_id;
+    owner.derived_task_ids = std::move(derived_task_ids);
+    owner.request.opcode = Request::WRITE;
+    owner.request.source = nullptr;
+    owner.request.target_id = 1;
+    owner.request.target_offset = public_task_id * 4096;
+    owner.request.length = length;
+    owner.kind = kind;
+    return owner;
+}
+
+QueueSubmit makeSubmit(uint64_t batch_token, size_t batch_slots_left,
+                       std::vector<QueueOwnerInput> owners) {
+    QueueSubmit submit;
+    submit.batch_token = batch_token;
+    submit.batch_slots_left = batch_slots_left;
+    submit.owners = std::move(owners);
+    return submit;
+}
+
+TEST(AdmissionQueueTest, AllowsEmptySubmitAsNoOp) {
+    LocalTransferAdmissionQueue queue({2, 128, 0, 0});
+    std::vector<QueueOwnerId> admitted_ids{99};
+
+    auto status = queue.tryAdmit(makeSubmit(1, 0, {}), admitted_ids);
+
+    EXPECT_EQ(status.code(), Status::Code::kOk);
+    EXPECT_TRUE(admitted_ids.empty());
+    EXPECT_EQ(queue.outstandingOwners(), 0u);
+    EXPECT_EQ(queue.outstandingBytes(), 0u);
+}
+
+TEST(AdmissionQueueTest, RejectsSubmitWhenQueueLimitsAreInvalid) {
+    LocalTransferAdmissionQueue queue({1, 128, 2, 0});
+    std::vector<QueueOwnerId> admitted_ids{99};
+
+    auto status =
+        queue.tryAdmit(makeSubmit(1, 1, {makeOwner(0, 16)}), admitted_ids);
+
+    EXPECT_EQ(status.code(), Status::Code::kInvalidArgument);
+    EXPECT_TRUE(admitted_ids.empty());
+    EXPECT_EQ(queue.outstandingOwners(), 0u);
+    EXPECT_EQ(queue.outstandingBytes(), 0u);
+}
+
+TEST(AdmissionQueueTest, RejectsInvalidInputsWithoutPartialAdmission) {
+    LocalTransferAdmissionQueue queue({4, 128, 0, 0});
+    std::vector<QueueOwnerId> admitted_ids{99};
+
+    auto status = queue.tryAdmit(
+        makeSubmit(
+            1, 2,
+            {makeOwner(0, 16, QueueOwnerKind::User, {1}), makeOwner(1, 16)}),
+        admitted_ids);
+
+    EXPECT_EQ(status.code(), Status::Code::kInvalidArgument);
+    EXPECT_TRUE(admitted_ids.empty());
+    EXPECT_EQ(queue.outstandingOwners(), 0u);
+    EXPECT_EQ(queue.outstandingBytes(), 0u);
+
+    status = queue.tryAdmit(makeSubmit(1, 1, {makeOwner(2, 16)}), admitted_ids);
+
+    ASSERT_EQ(status.code(), Status::Code::kOk);
+    ASSERT_EQ(admitted_ids.size(), 1u);
+    EXPECT_EQ(admitted_ids[0], 1u);
+}
+
+TEST(AdmissionQueueTest, RejectsUnsupportedOwnerKindWithoutPartialAdmission) {
+    LocalTransferAdmissionQueue queue({4, 128, 0, 0});
+    std::vector<QueueOwnerId> admitted_ids{99};
+
+    auto invalid_owner = makeOwner(0, 16, static_cast<QueueOwnerKind>(99), {1});
+    auto status = queue.tryAdmit(makeSubmit(1, 2, {std::move(invalid_owner)}),
+                                 admitted_ids);
+
+    EXPECT_EQ(status.code(), Status::Code::kInvalidArgument);
+    EXPECT_TRUE(admitted_ids.empty());
+    EXPECT_EQ(queue.outstandingOwners(), 0u);
+    EXPECT_EQ(queue.outstandingBytes(), 0u);
+
+    status = queue.tryAdmit(makeSubmit(1, 1, {makeOwner(2, 16)}), admitted_ids);
+
+    ASSERT_EQ(status.code(), Status::Code::kOk);
+    ASSERT_EQ(admitted_ids.size(), 1u);
+    EXPECT_EQ(admitted_ids[0], 1u);
+}
+
+TEST(AdmissionQueueTest, RejectsCapacityExceededWithoutPartialAdmission) {
+    LocalTransferAdmissionQueue queue({1, 64, 0, 0});
+    std::vector<QueueOwnerId> admitted_ids;
+
+    auto status = queue.tryAdmit(
+        makeSubmit(1, 2, {makeOwner(0, 16), makeOwner(1, 16)}), admitted_ids);
+
+    EXPECT_EQ(status.code(), Status::Code::kTooManyRequests);
+    EXPECT_TRUE(admitted_ids.empty());
+    EXPECT_EQ(queue.outstandingOwners(), 0u);
+    EXPECT_EQ(queue.outstandingBytes(), 0u);
+
+    status = queue.tryAdmit(makeSubmit(1, 1, {makeOwner(0, 16)}), admitted_ids);
+
+    ASSERT_EQ(status.code(), Status::Code::kOk);
+    ASSERT_EQ(admitted_ids.size(), 1u);
+    EXPECT_EQ(admitted_ids[0], 1u);
+}
+
+TEST(AdmissionQueueTest, RejectsExistingPublicTaskConflictWithoutMutation) {
+    LocalTransferAdmissionQueue queue({4, 128, 0, 0});
+    std::vector<QueueOwnerId> admitted_ids;
+
+    auto status =
+        queue.tryAdmit(makeSubmit(1, 1, {makeOwner(0, 16)}), admitted_ids);
+    ASSERT_EQ(status.code(), Status::Code::kOk);
+    ASSERT_EQ(admitted_ids.size(), 1u);
+    EXPECT_EQ(admitted_ids[0], 1u);
+
+    status = queue.tryAdmit(
+        makeSubmit(1, 2, {makeOwner(1, 16), makeOwner(0, 16)}), admitted_ids);
+
+    EXPECT_EQ(status.code(), Status::Code::kInvalidEntry);
+    EXPECT_TRUE(admitted_ids.empty());
+    EXPECT_EQ(queue.outstandingOwners(), 1u);
+    EXPECT_EQ(queue.outstandingBytes(), 16u);
+
+    QueueOwnerId owner_id = 0;
+    status = queue.resolveOwner(1, 1, owner_id);
+    EXPECT_EQ(status.code(), Status::Code::kInvalidEntry);
+
+    auto picked = queue.pickForDispatch(1, 16);
+    ASSERT_EQ(picked.size(), 1u);
+    status = queue.complete(picked[0], TransferStatusEnum::COMPLETED);
+    ASSERT_EQ(status.code(), Status::Code::kOk);
+    status = queue.retireBatch(1);
+    ASSERT_EQ(status.code(), Status::Code::kOk);
+
+    status = queue.tryAdmit(makeSubmit(2, 1, {makeOwner(0, 16)}), admitted_ids);
+
+    ASSERT_EQ(status.code(), Status::Code::kOk);
+    ASSERT_EQ(admitted_ids.size(), 1u);
+    EXPECT_EQ(admitted_ids[0], 2u);
+}
+
+TEST(AdmissionQueueTest, AccountsPublicSlotsSeparatelyFromQueueOwners) {
+    LocalTransferAdmissionQueue queue({2, 128, 0, 0});
+    std::vector<QueueOwnerId> admitted_ids;
+
+    auto status = queue.tryAdmit(
+        makeSubmit(1, 2, {makeOwner(7, 32, QueueOwnerKind::User, {8, 9})}),
+        admitted_ids);
+
+    EXPECT_EQ(status.code(), Status::Code::kTooManyRequests);
+    EXPECT_TRUE(admitted_ids.empty());
+    EXPECT_EQ(queue.outstandingOwners(), 0u);
+
+    status = queue.tryAdmit(
+        makeSubmit(1, 3, {makeOwner(7, 32, QueueOwnerKind::User, {8, 9})}),
+        admitted_ids);
+
+    ASSERT_EQ(status.code(), Status::Code::kOk);
+    ASSERT_EQ(admitted_ids.size(), 1u);
+    EXPECT_EQ(queue.outstandingOwners(), 1u);
+    EXPECT_EQ(queue.outstandingBytes(), 32u);
+
+    QueueOwnerId resolved_owner = 0;
+    status = queue.resolveOwner(1, 7, resolved_owner);
+    EXPECT_EQ(status.code(), Status::Code::kOk);
+    EXPECT_EQ(resolved_owner, admitted_ids[0]);
+    status = queue.resolveOwner(1, 8, resolved_owner);
+    EXPECT_EQ(status.code(), Status::Code::kOk);
+    EXPECT_EQ(resolved_owner, admitted_ids[0]);
+    status = queue.resolveOwner(1, 9, resolved_owner);
+    EXPECT_EQ(status.code(), Status::Code::kOk);
+    EXPECT_EQ(resolved_owner, admitted_ids[0]);
+    status = queue.resolveOwner(1, 0, resolved_owner);
+    EXPECT_EQ(status.code(), Status::Code::kInvalidEntry);
+}
+
+TEST(AdmissionQueueTest, PreservesStagingReserveForStagingInternalOwners) {
+    LocalTransferAdmissionQueue queue({2, 100, 1, 40});
+    std::vector<QueueOwnerId> admitted_ids;
+
+    auto status =
+        queue.tryAdmit(makeSubmit(1, 1, {makeOwner(0, 60)}), admitted_ids);
+    ASSERT_EQ(status.code(), Status::Code::kOk);
+
+    status = queue.tryAdmit(makeSubmit(2, 1, {makeOwner(0, 1)}), admitted_ids);
+    EXPECT_EQ(status.code(), Status::Code::kTooManyRequests);
+    EXPECT_TRUE(admitted_ids.empty());
+    EXPECT_EQ(queue.outstandingOwners(), 1u);
+    EXPECT_EQ(queue.outstandingBytes(), 60u);
+
+    status = queue.tryAdmit(
+        makeSubmit(3, 1, {makeOwner(0, 40, QueueOwnerKind::StagingInternal)}),
+        admitted_ids);
+    ASSERT_EQ(status.code(), Status::Code::kOk);
+    ASSERT_EQ(admitted_ids.size(), 1u);
+    EXPECT_EQ(admitted_ids[0], 2u);
+    EXPECT_EQ(queue.outstandingOwners(), 2u);
+    EXPECT_EQ(queue.outstandingBytes(), 100u);
+}
+
+TEST(AdmissionQueueTest, KeepsAdmissionOrderForDispatch) {
+    LocalTransferAdmissionQueue queue({4, 128, 0, 0});
+    std::vector<QueueOwnerId> admitted_ids;
+
+    auto status = queue.tryAdmit(
+        makeSubmit(1, 2,
+                   {makeOwner(0, 60),
+                    makeOwner(1, 10, QueueOwnerKind::StagingInternal)}),
+        admitted_ids);
+
+    ASSERT_EQ(status.code(), Status::Code::kOk);
+    ASSERT_EQ(admitted_ids.size(), 2u);
+    const std::vector<QueueOwnerId> expected_ids{1, 2};
+    EXPECT_EQ(admitted_ids, expected_ids);
+
+    EXPECT_TRUE(queue.pickForDispatch(2, 50).empty());
+
+    auto picked = queue.pickForDispatch(2, 70);
+
+    EXPECT_EQ(picked, expected_ids);
+}
+
+TEST(AdmissionQueueTest, RequiresDispatchBeforeTerminalCompletion) {
+    LocalTransferAdmissionQueue queue({2, 128, 0, 0});
+    std::vector<QueueOwnerId> admitted_ids;
+
+    auto status =
+        queue.tryAdmit(makeSubmit(1, 1, {makeOwner(0, 16)}), admitted_ids);
+    ASSERT_EQ(status.code(), Status::Code::kOk);
+    ASSERT_EQ(admitted_ids.size(), 1u);
+
+    status = queue.complete(admitted_ids[0], TransferStatusEnum::COMPLETED);
+    EXPECT_EQ(status.code(), Status::Code::kInvalidEntry);
+    status = queue.complete(admitted_ids[0], TransferStatusEnum::PENDING);
+    EXPECT_EQ(status.code(), Status::Code::kInvalidArgument);
+    EXPECT_EQ(queue.outstandingOwners(), 1u);
+    EXPECT_EQ(queue.outstandingBytes(), 16u);
+
+    auto picked = queue.pickForDispatch(1, 16);
+    ASSERT_EQ(picked.size(), 1u);
+
+    status = queue.complete(picked[0], TransferStatusEnum::COMPLETED);
+    EXPECT_EQ(status.code(), Status::Code::kOk);
+    EXPECT_EQ(queue.outstandingOwners(), 0u);
+    EXPECT_EQ(queue.outstandingBytes(), 0u);
+
+    status = queue.complete(picked[0], TransferStatusEnum::COMPLETED);
+    EXPECT_EQ(status.code(), Status::Code::kInvalidEntry);
+}
+
+TEST(AdmissionQueueTest, RetainsTerminalStatusUntilBatchRetire) {
+    LocalTransferAdmissionQueue queue({2, 128, 0, 0});
+    std::vector<QueueOwnerId> admitted_ids;
+
+    auto status = queue.tryAdmit(
+        makeSubmit(1, 2, {makeOwner(0, 16, QueueOwnerKind::User, {1})}),
+        admitted_ids);
+    ASSERT_EQ(status.code(), Status::Code::kOk);
+
+    TransferStatusEnum public_status = TransferStatusEnum::INVALID;
+    status = queue.getPublicStatus(1, 1, public_status);
+    EXPECT_EQ(status.code(), Status::Code::kOk);
+    EXPECT_EQ(public_status, TransferStatusEnum::PENDING);
+
+    auto picked = queue.pickForDispatch(1, 16);
+    ASSERT_EQ(picked.size(), 1u);
+    status = queue.complete(picked[0], TransferStatusEnum::FAILED);
+    ASSERT_EQ(status.code(), Status::Code::kOk);
+
+    status = queue.getPublicStatus(1, 0, public_status);
+    EXPECT_EQ(status.code(), Status::Code::kOk);
+    EXPECT_EQ(public_status, TransferStatusEnum::FAILED);
+    status = queue.getPublicStatus(1, 1, public_status);
+    EXPECT_EQ(status.code(), Status::Code::kOk);
+    EXPECT_EQ(public_status, TransferStatusEnum::FAILED);
+
+    status = queue.retireBatch(1);
+    ASSERT_EQ(status.code(), Status::Code::kOk);
+
+    QueueOwnerId resolved_owner = 0;
+    status = queue.resolveOwner(1, 0, resolved_owner);
+    EXPECT_EQ(status.code(), Status::Code::kInvalidEntry);
+    status = queue.getPublicStatus(1, 1, public_status);
+    EXPECT_EQ(status.code(), Status::Code::kInvalidEntry);
+}
+
+TEST(AdmissionQueueTest, RetainsSpecificTerminalStatus) {
+    LocalTransferAdmissionQueue queue({1, 128, 0, 0});
+    std::vector<QueueOwnerId> admitted_ids;
+
+    auto status =
+        queue.tryAdmit(makeSubmit(1, 1, {makeOwner(0, 16)}), admitted_ids);
+    ASSERT_EQ(status.code(), Status::Code::kOk);
+
+    auto picked = queue.pickForDispatch(1, 16);
+    ASSERT_EQ(picked.size(), 1u);
+    status = queue.complete(picked[0], TransferStatusEnum::TIMEOUT);
+    ASSERT_EQ(status.code(), Status::Code::kOk);
+
+    TransferStatusEnum public_status = TransferStatusEnum::PENDING;
+    status = queue.getPublicStatus(1, 0, public_status);
+    ASSERT_EQ(status.code(), Status::Code::kOk);
+    EXPECT_EQ(public_status, TransferStatusEnum::TIMEOUT);
+}
+
+TEST(AdmissionQueueTest, RejectsRetireWithNonTerminalOwners) {
+    LocalTransferAdmissionQueue queue({2, 128, 0, 0});
+    std::vector<QueueOwnerId> admitted_ids;
+
+    auto status = queue.tryAdmit(
+        makeSubmit(1, 2, {makeOwner(0, 16), makeOwner(1, 16)}), admitted_ids);
+    ASSERT_EQ(status.code(), Status::Code::kOk);
+
+    auto picked = queue.pickForDispatch(1, 16);
+    ASSERT_EQ(picked.size(), 1u);
+    status = queue.complete(picked[0], TransferStatusEnum::COMPLETED);
+    ASSERT_EQ(status.code(), Status::Code::kOk);
+
+    status = queue.retireBatch(1);
+    EXPECT_EQ(status.code(), Status::Code::kInvalidEntry);
+
+    picked = queue.pickForDispatch(1, 16);
+    ASSERT_EQ(picked.size(), 1u);
+    status = queue.complete(picked[0], TransferStatusEnum::COMPLETED);
+    ASSERT_EQ(status.code(), Status::Code::kOk);
+
+    status = queue.retireBatch(1);
+    EXPECT_EQ(status.code(), Status::Code::kOk);
+}
+
+TEST(AdmissionQueueTest, AllowsBatchTokenReuseAfterRetire) {
+    LocalTransferAdmissionQueue queue({1, 128, 0, 0});
+    std::vector<QueueOwnerId> admitted_ids;
+
+    auto status =
+        queue.tryAdmit(makeSubmit(1, 1, {makeOwner(0, 16)}), admitted_ids);
+    ASSERT_EQ(status.code(), Status::Code::kOk);
+    ASSERT_EQ(admitted_ids.size(), 1u);
+    EXPECT_EQ(admitted_ids[0], 1u);
+
+    auto picked = queue.pickForDispatch(1, 16);
+    ASSERT_EQ(picked.size(), 1u);
+    status = queue.complete(picked[0], TransferStatusEnum::COMPLETED);
+    ASSERT_EQ(status.code(), Status::Code::kOk);
+    status = queue.retireBatch(1);
+    ASSERT_EQ(status.code(), Status::Code::kOk);
+
+    status = queue.tryAdmit(makeSubmit(1, 1, {makeOwner(0, 16)}), admitted_ids);
+
+    ASSERT_EQ(status.code(), Status::Code::kOk);
+    ASSERT_EQ(admitted_ids.size(), 1u);
+    EXPECT_EQ(admitted_ids[0], 2u);
+
+    QueueOwnerId resolved_owner = 0;
+    status = queue.resolveOwner(1, 0, resolved_owner);
+    EXPECT_EQ(status.code(), Status::Code::kOk);
+    EXPECT_EQ(resolved_owner, 2u);
+}
+
+}  // namespace
+}  // namespace tent
+}  // namespace mooncake
diff --git a/mooncake-transfer-engine/tent/tests/coalesce_regions_test.cpp b/mooncake-transfer-engine/tent/tests/coalesce_regions_test.cpp
new file mode 100644
index 00000000..be26161c
--- /dev/null
+++ b/mooncake-transfer-engine/tent/tests/coalesce_regions_test.cpp
@@ -0,0 +1,180 @@
+// Copyright 2026 KVCache.AI
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include <gtest/gtest.h>
+
+#include <cstdint>
+#include <string>
+#include <vector>
+
+#include "tent/runtime/segment.h"
+#include "tent/runtime/topology.h"
+
+namespace mooncake {
+namespace tent {
+namespace {
+
+constexpr size_t kKiB = static_cast<size_t>(1) << 10;
+constexpr size_t kMiB = static_cast<size_t>(1) << 20;
+constexpr size_t kGiB = static_cast<size_t>(1) << 30;
+
+// Build an entries vector where each RangeLocation is address-adjacent to
+// the previous one. Mirrors what Platform::getLoader().getLocation produces.
+std::vector<RangeLocation> makeEntries(
+    const std::vector<std::pair<size_t, std::string>>& segs) {
+    std::vector<RangeLocation> out;
+    uint64_t cursor = 0;
+    out.reserve(segs.size());
+    for (auto& [len, loc] : segs) {
+        out.push_back(RangeLocation{cursor, len, loc});
+        cursor += len;
+    }
+    return out;
+}
+
+size_t sumLens(const std::vector<RangeLocation>& entries) {
+    size_t s = 0;
+    for (auto& e : entries) s += e.len;
+    return s;
+}
+
+size_t sumSizes(const std::vector<Region>& regions) {
+    size_t s = 0;
+    for (auto& r : regions) s += r.size;
+    return s;
+}
+
+// All these inputs are well under (128 buckets × 1 MiB floor), so the
+// effective bucket size is the 1 MiB default min_bucket_bytes floor.
+TEST(CoalesceRegionsTest, TieBreakIsDeterministic) {
+    auto entries = makeEntries({{256 * kKiB, "cpu:0"},
+                                {256 * kKiB, "cpu:1"},
+                                {256 * kKiB, "cpu:0"},
+                                {256 * kKiB, "cpu:1"}});
+    auto first = coalesceRegions(entries);
+    ASSERT_EQ(first.size(), 1u);
+    EXPECT_EQ(first[0].size, 1 * kMiB);
+    EXPECT_EQ(first[0].location, "cpu:0");
+
+    for (int i = 0; i < 100; ++i) {
+        auto run = coalesceRegions(entries);
+        ASSERT_EQ(run.size(), 1u);
+        EXPECT_EQ(run[0].location, "cpu:0") << "run " << i;
+    }
+}
+
+TEST(CoalesceRegionsTest, MajorityWins) {
+    auto entries = makeEntries({{900 * kKiB, "cpu:0"}, {100 * kKiB, "cpu:1"}});
+    auto out = coalesceRegions(entries);
+    ASSERT_EQ(out.size(), 1u);
+    EXPECT_EQ(out[0].size, 1000 * kKiB);
+    EXPECT_EQ(out[0].location, "cpu:0");
+}
+
+TEST(CoalesceRegionsTest, SameLocationCoalesces) {
+    auto entries = makeEntries({{256 * kKiB, "cpu:0"},
+                                {256 * kKiB, "cpu:0"},
+                                {256 * kKiB, "cpu:0"},
+                                {256 * kKiB, "cpu:0"}});
+    auto out = coalesceRegions(entries);
+    ASSERT_EQ(out.size(), 1u);
+    EXPECT_EQ(out[0].size, 1 * kMiB);
+    EXPECT_EQ(out[0].location, "cpu:0");
+}
+
+TEST(CoalesceRegionsTest, EmptyInput) {
+    auto out = coalesceRegions({});
+    EXPECT_TRUE(out.empty());
+}
+
+TEST(CoalesceRegionsTest, TotalBytesPreserved) {
+    auto entries = makeEntries({{700 * kKiB, "cpu:0"},
+                                {300 * kKiB, "cpu:1"},
+                                {2 * kMiB, "cpu:0"},
+                                {500 * kKiB, "cpu:1"},
+                                {1500 * kKiB, "cpu:0"}});
+    auto out = coalesceRegions(entries);
+    EXPECT_EQ(sumSizes(out), sumLens(entries));
+}
+
+TEST(CoalesceRegionsTest, AutoBucketBoundedBy128) {
+    // 256 segments × 16 MiB alternating cpu:0 / cpu:1 = 4 GiB total.
+    // Auto bucket should land near 4 GiB / 128 = 32 MiB → at most 128
+    // regions.
+    std::vector<std::pair<size_t, std::string>> segs;
+    segs.reserve(256);
+    for (int i = 0; i < 256; ++i) {
+        segs.emplace_back(16 * kMiB, i % 2 == 0 ? "cpu:0" : "cpu:1");
+    }
+    auto entries = makeEntries(segs);
+    auto out = coalesceRegions(entries);  // all defaults: auto, 128, 1 MiB
+    EXPECT_LE(out.size(), 128u);
+    EXPECT_EQ(sumSizes(out), sumLens(entries));
+}
+
+TEST(CoalesceRegionsTest, AutoBucketRespectsMinFloor) {
+    auto entries = makeEntries({{50 * kKiB, "cpu:0"}, {50 * kKiB, "cpu:0"}});
+    auto out = coalesceRegions(entries);  // auto: 1 MiB floor applies
+    ASSERT_EQ(out.size(), 1u);
+    EXPECT_EQ(out[0].size, 100 * kKiB);
+    EXPECT_EQ(out[0].location, "cpu:0");
+}
+
+// `max_buckets=0` and `min_bucket_bytes=0` together would historically
+// yield bucket_bytes=0 → divide-by-zero / take=0 infinite loop. Defensive
+// floor inside the helper must keep the call well-defined.
+TEST(CoalesceRegionsTest, DefensiveZeroParams) {
+    auto entries = makeEntries({{256 * kKiB, "cpu:0"}, {256 * kKiB, "cpu:1"}});
+    auto out = coalesceRegions(entries,
+                               /*max_buckets=*/0,
+                               /*min_bucket_bytes=*/0);
+    EXPECT_FALSE(out.empty());
+    EXPECT_EQ(sumSizes(out), sumLens(entries));
+}
+
+// Documents the intentional precision/size tradeoff: a tiny minority
+// region inside a much larger bucket gets the majority label. This is the
+// designed behavior, not a bug — kept here so any future change that
+// "fixes" this case has to consciously remove the test.
+TEST(CoalesceRegionsTest, AcceptsApproximation) {
+    auto entries = makeEntries({{1 * kKiB, "cpu:0"}, {9 * kMiB, "cpu:1"}});
+    auto out = coalesceRegions(entries);  // auto bucket >> 1 KiB
+    ASSERT_FALSE(out.empty());
+    EXPECT_EQ(out.front().location, "cpu:1");
+    EXPECT_EQ(sumSizes(out), sumLens(entries));
+}
+
+// Verifies only spatially adjacent entries get merged into the same
+// Region: a 1 MiB bucket size + 6 alternating 256 KiB segments yields
+// (a) Regions in input order, (b) the second Region cannot pull cpu:0
+// votes from the first bucket back across the boundary.
+TEST(CoalesceRegionsTest, OnlyAdjacentMerged) {
+    auto entries = makeEntries({{256 * kKiB, "cpu:0"},
+                                {256 * kKiB, "cpu:1"},
+                                {256 * kKiB, "cpu:0"},
+                                {256 * kKiB, "cpu:1"},
+                                {256 * kKiB, "cpu:0"},
+                                {256 * kKiB, "cpu:1"}});
+    auto out = coalesceRegions(entries);
+    ASSERT_EQ(out.size(), 2u);
+    EXPECT_EQ(out[0].size, 1 * kMiB);
+    EXPECT_EQ(out[0].location, "cpu:0");  // 4-way tie, first-seen wins
+    EXPECT_EQ(out[1].size, 512 * kKiB);
+    EXPECT_EQ(out[1].location, "cpu:0");  // 2-way tie, first-seen wins
+    EXPECT_EQ(out[0].size + out[1].size, sumLens(entries));
+}
+
+}  // namespace
+}  // namespace tent
+}  // namespace mooncake
diff --git a/mooncake-transfer-engine/tent/tests/engine_failover_e2e_test.cpp b/mooncake-transfer-engine/tent/tests/engine_failover_e2e_test.cpp
new file mode 100644
index 00000000..7100d991
--- /dev/null
+++ b/mooncake-transfer-engine/tent/tests/engine_failover_e2e_test.cpp
@@ -0,0 +1,1041 @@
+// Copyright 2026 KVCache.AI
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+//
+// End-to-end failover test: drives the real TransferEngineImpl through
+// submitTransfer() / getTransferStatus() / resubmitTransferTask() by swapping
+// two FakeTransports into the engine's transport_list_ and wrapping the
+// primary one in a FaultProxyTransport.
+//
+// ---------------------------------------------------------------------------
+// Scenarios:
+//   P0: Primary succeeds at submit but getTransferStatus reports FAILED
+//       (simulates WC error / QP error / peer drop mid-transfer).
+//       -> engine must failover to the secondary and succeed.
+//
+// Submit-stage failures are NOT tested here. Today submitTransferTasks
+// failures are marked UNSPEC and surface as FAILED without a failover
+// attempt -- see the "Known gaps" section of docs/source/design/tent/
+// failover.md for why. Tests for that path belong in a future change
+// that makes submit-stage recovery safe (see known gaps).
+
+#include <gtest/gtest.h>
+
+#include <atomic>
+#include <chrono>
+#include <cstdlib>
+#include <cstring>
+#include <functional>
+#include <memory>
+#include <set>
+#include <string>
+#include <thread>
+#include <utility>
+#include <vector>
+
+#include "tent/common/config.h"
+#include "tent/common/types.h"
+#include "tent/runtime/segment.h"
+#include "tent/runtime/transfer_engine_impl.h"
+#include "tent/runtime/transport.h"
+#include "tent/transport/fault_proxy/fault_proxy_transport.h"
+
+namespace mooncake {
+namespace tent {
+namespace {
+
+// ---------------------------------------------------------------------------
+// FakeTransport: declares itself capable of dram_to_dram, records itself in
+// BufferDesc::transports under a configurable slot, and always completes
+// transfers. Enough to satisfy checkAvailability() + resolveTransport() in
+// TransferEngineImpl.
+// ---------------------------------------------------------------------------
+
+class FakeSubBatch : public Transport::SubBatch {
+   public:
+    size_t size() const override { return task_count; }
+    size_t task_count = 0;
+    std::vector<Request> requests;
+    std::vector<TransferStatus> statuses;
+    std::vector<int> poll_counts;
+};
+
+class FakeTransport : public Transport {
+   public:
+    using StatusFactory = std::function<TransferStatus(const Request&)>;
+    using PollStatusFactory =
+        std::function<TransferStatus(const Request&, int)>;
+
+    explicit FakeTransport(TransportType self_type,
+                           StatusFactory status_factory = {},
+                           PollStatusFactory poll_status_factory = {})
+        : self_type_(self_type),
+          status_factory_(std::move(status_factory)),
+          poll_status_factory_(std::move(poll_status_factory)) {
+        caps.dram_to_dram = true;  // so checkAvailability returns true
+    }
+
+    std::atomic<int> install_calls{0};
+    std::atomic<int> submit_calls{0};
+    std::atomic<int> status_calls{0};
+    std::atomic<int> add_mem_calls{0};
+
+    Status install(std::string& /*local_segment_name*/,
+                   std::shared_ptr<ControlService> /*metadata*/,
+                   std::shared_ptr<Topology> /*local_topology*/,
+                   std::shared_ptr<Config> /*conf*/ = nullptr) override {
+        ++install_calls;
+        return Status::OK();
+    }
+
+    Status allocateSubBatch(SubBatchRef& batch, size_t /*max_size*/) override {
+        batch = new FakeSubBatch();
+        return Status::OK();
+    }
+
+    Status freeSubBatch(SubBatchRef& batch) override {
+        delete batch;
+        batch = nullptr;
+        return Status::OK();
+    }
+
+    Status submitTransferTasks(
+        SubBatchRef batch, const std::vector<Request>& request_list) override {
+        ++submit_calls;
+        auto* fb = static_cast<FakeSubBatch*>(batch);
+        for (const auto& req : request_list) {
+            if (status_factory_) {
+                fb->statuses.push_back(status_factory_(req));
+            } else {
+                fb->statuses.push_back(
+                    {TransferStatusEnum::COMPLETED, req.length});
+            }
+            fb->requests.push_back(req);
+            fb->poll_counts.push_back(0);
+            fb->task_count++;
+        }
+        return Status::OK();
+    }
+
+    Status getTransferStatus(SubBatchRef batch, int task_id,
+                             TransferStatus& status) override {
+        ++status_calls;
+        auto* fb = static_cast<FakeSubBatch*>(batch);
+        if (task_id < 0 || task_id >= (int)fb->statuses.size()) {
+            return Status::InvalidArgument("bad task_id" LOC_MARK);
+        }
+        ++fb->poll_counts[task_id];
+        if (poll_status_factory_) {
+            status = poll_status_factory_(fb->requests[task_id],
+                                          fb->poll_counts[task_id]);
+        } else {
+            status = fb->statuses[task_id];
+        }
+        return Status::OK();
+    }
+
+    // Tag ourselves into the BufferDesc so resolveTransport() considers us.
+    Status addMemoryBuffer(BufferDesc& desc,
+                           const MemoryOptions& /*options*/) override {
+        ++add_mem_calls;
+        desc.transports.push_back(self_type_);
+        return Status::OK();
+    }
+
+    Status addMemoryBuffer(std::vector<BufferDesc>& desc_list,
+                           const MemoryOptions& options) override {
+        for (auto& d : desc_list) {
+            auto s = addMemoryBuffer(d, options);
+            if (!s.ok()) return s;
+        }
+        return Status::OK();
+    }
+
+    Status removeMemoryBuffer(BufferDesc& /*desc*/) override {
+        return Status::OK();
+    }
+
+    Status allocateLocalMemory(void** addr, size_t size,
+                               MemoryOptions& /*options*/) override {
+        *addr = std::malloc(size);
+        if (!*addr) return Status::InternalError("malloc failed" LOC_MARK);
+        return Status::OK();
+    }
+
+    Status freeLocalMemory(void* addr, size_t /*size*/) override {
+        std::free(addr);
+        return Status::OK();
+    }
+
+    bool warmupMemory(void* /*addr*/, size_t /*length*/) override {
+        return false;  // no pinning; engine will fall back to its own path
+    }
+
+    const char* getName() const override {
+        return self_type_ == RDMA ? "<fake-rdma>" : "<fake-tcp>";
+    }
+
+   private:
+    TransportType self_type_;
+    StatusFactory status_factory_;
+    PollStatusFactory poll_status_factory_;
+};
+
+// ---------------------------------------------------------------------------
+// Helpers
+// ---------------------------------------------------------------------------
+
+std::shared_ptr<Config> makeMinimalP2PConfig() {
+    auto cfg = std::make_shared<Config>();
+    // p2p metadata avoids needing an external redis/etcd/http server.
+    cfg->set("metadata_type", "p2p");
+    cfg->set("metadata_servers", "");
+    cfg->set("rpc_server_hostname", "127.0.0.1");
+    cfg->set("rpc_server_port", "0");
+    cfg->set("log_level", "warning");
+    cfg->set("merge_requests", false);
+
+    // Disable every real transport. We'll inject fakes into the slots we care
+    // about via swapTransportForTest.
+    cfg->set("transports/tcp/enable", false);
+    cfg->set("transports/shm/enable", false);
+    cfg->set("transports/rdma/enable", false);
+    cfg->set("transports/io_uring/enable", false);
+    cfg->set("transports/nvlink/enable", false);
+    cfg->set("transports/mnnvl/enable", false);
+    cfg->set("transports/gds/enable", false);
+    cfg->set("transports/ascend_direct/enable", false);
+
+    // Keep failover limit at default (3) but make it explicit.
+    cfg->set("max_failover_attempts", 3);
+    return cfg;
+}
+
+// Wait for a task to leave PENDING. Bounded so tests fail fast.
+TransferStatus pollUntilDone(
+    TransferEngineImpl& engine, BatchID batch_id, size_t task_id,
+    std::chrono::milliseconds timeout = std::chrono::milliseconds(2000)) {
+    TransferStatus ts{};
+    const auto deadline = std::chrono::steady_clock::now() + timeout;
+    while (std::chrono::steady_clock::now() < deadline) {
+        ts = {};
+        auto s = engine.getTransferStatus(batch_id, task_id, ts);
+        if (!s.ok()) {
+            ADD_FAILURE() << "getTransferStatus returned error: "
+                          << s.ToString();
+            return ts;
+        }
+        if (ts.s == TransferStatusEnum::COMPLETED ||
+            ts.s == TransferStatusEnum::FAILED) {
+            return ts;
+        }
+        std::this_thread::sleep_for(std::chrono::milliseconds(2));
+    }
+    return ts;
+}
+
+struct CorruptedRdmaBatch {
+    std::shared_ptr<FakeTransport> fake_rdma;
+    std::shared_ptr<FakeTransport> fake_tcp;
+    std::vector<uint8_t> buf;
+    BatchID batch_id{0};
+};
+
+void submitCorruptedRdmaBatch(TransferEngineImpl& engine,
+                              CorruptedRdmaBatch& batch, uint8_t fill) {
+    batch.fake_rdma = std::make_shared<FakeTransport>(RDMA);
+    batch.fake_tcp = std::make_shared<FakeTransport>(TCP);
+
+    FaultPolicy rdma_policy;
+    rdma_policy.status_corrupt_rate = 1.0;
+    auto proxied_rdma =
+        std::make_shared<FaultProxyTransport>(batch.fake_rdma, rdma_policy);
+
+    std::string seg_name = engine.getSegmentName();
+    ASSERT_TRUE(proxied_rdma->install(seg_name, nullptr, nullptr).ok());
+    ASSERT_TRUE(batch.fake_tcp->install(seg_name, nullptr, nullptr).ok());
+
+    engine.swapTransportForTest(RDMA, proxied_rdma);
+    engine.swapTransportForTest(TCP, batch.fake_tcp);
+
+    constexpr size_t kBufLen = 4096;
+    batch.buf.assign(kBufLen, fill);
+    ASSERT_TRUE(engine.registerLocalMemory(batch.buf.data(), kBufLen).ok());
+
+    batch.batch_id = engine.allocateBatch(8);
+    ASSERT_NE(batch.batch_id, (BatchID)0);
+
+    Request req;
+    req.opcode = Request::WRITE;
+    req.source = batch.buf.data();
+    req.target_id = LOCAL_SEGMENT_ID;
+    req.target_offset = reinterpret_cast<uint64_t>(batch.buf.data());
+    req.length = kBufLen;
+
+    ASSERT_TRUE(engine.submitTransfer(batch.batch_id, {req}).ok());
+}
+
+// ---------------------------------------------------------------------------
+// P0: Completion reports FAILED (simulates WC error / QP error / peer drop
+// mid-transfer). Engine must failover.
+// ---------------------------------------------------------------------------
+
+TEST(EngineFailoverE2E, StatusCorruptionTriggersFailoverToSecondary) {
+    auto cfg = makeMinimalP2PConfig();
+    TransferEngineImpl engine(cfg);
+    ASSERT_TRUE(engine.available());
+
+    auto fake_rdma = std::make_shared<FakeTransport>(RDMA);
+    auto fake_tcp = std::make_shared<FakeTransport>(TCP);
+
+    FaultPolicy rdma_policy;
+    rdma_policy.status_corrupt_rate = 1.0;  // every COMPLETED flipped to FAILED
+    auto proxied_rdma =
+        std::make_shared<FaultProxyTransport>(fake_rdma, rdma_policy);
+
+    std::string seg_name = engine.getSegmentName();
+    ASSERT_TRUE(proxied_rdma->install(seg_name, nullptr, nullptr).ok());
+    ASSERT_TRUE(fake_tcp->install(seg_name, nullptr, nullptr).ok());
+
+    engine.swapTransportForTest(RDMA, proxied_rdma);
+    engine.swapTransportForTest(TCP, fake_tcp);
+
+    constexpr size_t kBufLen = 4096;
+    std::vector<uint8_t> buf(kBufLen, 0xCD);
+    ASSERT_TRUE(engine.registerLocalMemory(buf.data(), kBufLen).ok());
+
+    BatchID batch_id = engine.allocateBatch(8);
+    ASSERT_NE(batch_id, (BatchID)0);
+
+    Request req;
+    req.opcode = Request::WRITE;
+    req.source = buf.data();
+    req.target_id = LOCAL_SEGMENT_ID;
+    req.target_offset = reinterpret_cast<uint64_t>(buf.data());
+    req.length = kBufLen;
+
+    ASSERT_TRUE(engine.submitTransfer(batch_id, {req}).ok());
+
+    auto final_status = pollUntilDone(engine, batch_id, 0);
+    EXPECT_EQ(final_status.s, TransferStatusEnum::COMPLETED);
+
+    // RDMA saw exactly one submit (succeeded on the wire), then its status
+    // was corrupted; engine failed over.
+    EXPECT_EQ(fake_rdma->submit_calls.load(), 1);
+    EXPECT_GE(fake_tcp->submit_calls.load(), 1);
+
+    EXPECT_TRUE(engine.freeBatch(batch_id).ok());
+    EXPECT_TRUE(engine.unregisterLocalMemory(buf.data(), kBufLen).ok());
+}
+
+TEST(EngineFailoverE2E, AutoFailoverOnPollDisabledLeavesTaskFailed) {
+    auto cfg = makeMinimalP2PConfig();
+    cfg->set("enable_auto_failover_on_poll", false);
+    TransferEngineImpl engine(cfg);
+    ASSERT_TRUE(engine.available());
+
+    CorruptedRdmaBatch batch;
+    submitCorruptedRdmaBatch(engine, batch, 0xA1);
+
+    TransferStatus status{};
+    ASSERT_TRUE(engine.getTransferStatus(batch.batch_id, 0, status).ok());
+    EXPECT_EQ(status.s, TransferStatusEnum::FAILED);
+    EXPECT_EQ(batch.fake_rdma->submit_calls.load(), 1);
+    EXPECT_EQ(batch.fake_tcp->submit_calls.load(), 0);
+
+    EXPECT_TRUE(engine.freeBatch(batch.batch_id).ok());
+    EXPECT_TRUE(
+        engine.unregisterLocalMemory(batch.buf.data(), batch.buf.size()).ok());
+}
+
+TEST(EngineFailoverE2E, AutoFailoverOnPollDisabledAppliesToVectorStatus) {
+    auto cfg = makeMinimalP2PConfig();
+    cfg->set("enable_auto_failover_on_poll", false);
+    TransferEngineImpl engine(cfg);
+    ASSERT_TRUE(engine.available());
+
+    CorruptedRdmaBatch batch;
+    submitCorruptedRdmaBatch(engine, batch, 0xA2);
+
+    std::vector<TransferStatus> status_list;
+    ASSERT_TRUE(engine.getTransferStatus(batch.batch_id, status_list).ok());
+    ASSERT_EQ(status_list.size(), 1);
+    EXPECT_EQ(status_list[0].s, TransferStatusEnum::FAILED);
+    EXPECT_EQ(batch.fake_rdma->submit_calls.load(), 1);
+    EXPECT_EQ(batch.fake_tcp->submit_calls.load(), 0);
+
+    EXPECT_TRUE(engine.freeBatch(batch.batch_id).ok());
+    EXPECT_TRUE(
+        engine.unregisterLocalMemory(batch.buf.data(), batch.buf.size()).ok());
+}
+
+TEST(EngineFailoverE2E, AutoFailoverOnPollDisabledAppliesToOverallStatus) {
+    auto cfg = makeMinimalP2PConfig();
+    cfg->set("enable_auto_failover_on_poll", false);
+    TransferEngineImpl engine(cfg);
+    ASSERT_TRUE(engine.available());
+
+    CorruptedRdmaBatch batch;
+    submitCorruptedRdmaBatch(engine, batch, 0xA3);
+
+    TransferStatus overall_status{};
+    ASSERT_TRUE(engine.getTransferStatus(batch.batch_id, overall_status).ok());
+    EXPECT_EQ(overall_status.s, TransferStatusEnum::FAILED);
+    EXPECT_EQ(batch.fake_rdma->submit_calls.load(), 1);
+    EXPECT_EQ(batch.fake_tcp->submit_calls.load(), 0);
+
+    EXPECT_TRUE(engine.freeBatch(batch.batch_id).ok());
+    EXPECT_TRUE(
+        engine.unregisterLocalMemory(batch.buf.data(), batch.buf.size()).ok());
+}
+
+// ---------------------------------------------------------------------------
+// P0c: Explicit progressBatch() drives one progress step and always allows
+// failover/resubmit, regardless of enable_auto_failover_on_poll. Internal
+// sync paths (waitTransferCompletion, transferSync) and the proxy event loop
+// are wired through it so observation-only callers stay decoupled from
+// progress-driving callers.
+// ---------------------------------------------------------------------------
+
+TEST(EngineFailoverE2E, ProgressBatchRetriesWhenPollAutoFailoverDisabled) {
+    auto cfg = makeMinimalP2PConfig();
+    cfg->set("enable_auto_failover_on_poll", false);
+    TransferEngineImpl engine(cfg);
+    ASSERT_TRUE(engine.available());
+
+    CorruptedRdmaBatch batch;
+    submitCorruptedRdmaBatch(engine, batch, 0xA4);
+
+    TransferStatus overall_status{};
+    ASSERT_TRUE(engine.progressBatch(batch.batch_id, overall_status).ok());
+    EXPECT_EQ(overall_status.s, TransferStatusEnum::PENDING);
+    EXPECT_EQ(batch.fake_rdma->submit_calls.load(), 1);
+    EXPECT_EQ(batch.fake_tcp->submit_calls.load(), 1);
+    EXPECT_EQ(batch.fake_tcp->status_calls.load(), 0)
+        << "progressBatch should perform one progress step, not poll the "
+           "fallback submission immediately";
+
+    ASSERT_TRUE(engine.progressBatch(batch.batch_id, overall_status).ok());
+    EXPECT_EQ(overall_status.s, TransferStatusEnum::COMPLETED);
+    EXPECT_EQ(batch.fake_tcp->status_calls.load(), 1);
+
+    EXPECT_TRUE(engine.freeBatch(batch.batch_id).ok());
+    EXPECT_TRUE(
+        engine.unregisterLocalMemory(batch.buf.data(), batch.buf.size()).ok());
+}
+
+TEST(EngineFailoverE2E, ProgressBatchDoesNotReviveObservedFailedTask) {
+    auto cfg = makeMinimalP2PConfig();
+    cfg->set("enable_auto_failover_on_poll", false);
+    TransferEngineImpl engine(cfg);
+    ASSERT_TRUE(engine.available());
+
+    CorruptedRdmaBatch batch;
+    submitCorruptedRdmaBatch(engine, batch, 0xA5);
+
+    TransferStatus overall_status{};
+    ASSERT_TRUE(engine.getTransferStatus(batch.batch_id, overall_status).ok());
+    EXPECT_EQ(overall_status.s, TransferStatusEnum::FAILED);
+    EXPECT_EQ(batch.fake_tcp->submit_calls.load(), 0);
+
+    ASSERT_TRUE(engine.progressBatch(batch.batch_id, overall_status).ok());
+    EXPECT_EQ(overall_status.s, TransferStatusEnum::FAILED);
+    EXPECT_EQ(batch.fake_tcp->submit_calls.load(), 0);
+
+    EXPECT_TRUE(engine.freeBatch(batch.batch_id).ok());
+    EXPECT_TRUE(
+        engine.unregisterLocalMemory(batch.buf.data(), batch.buf.size()).ok());
+}
+
+TEST(EngineFailoverE2E, ProgressBatchHonorsMaxFailoverAttemptsZero) {
+    auto cfg = makeMinimalP2PConfig();
+    cfg->set("enable_auto_failover_on_poll", false);
+    cfg->set("max_failover_attempts", 0);
+    TransferEngineImpl engine(cfg);
+    ASSERT_TRUE(engine.available());
+
+    CorruptedRdmaBatch batch;
+    submitCorruptedRdmaBatch(engine, batch, 0xA6);
+
+    TransferStatus overall_status{};
+    ASSERT_TRUE(engine.progressBatch(batch.batch_id, overall_status).ok());
+    EXPECT_EQ(overall_status.s, TransferStatusEnum::FAILED);
+    EXPECT_EQ(batch.fake_rdma->submit_calls.load(), 1);
+    EXPECT_EQ(batch.fake_tcp->submit_calls.load(), 0);
+
+    EXPECT_TRUE(engine.freeBatch(batch.batch_id).ok());
+    EXPECT_TRUE(
+        engine.unregisterLocalMemory(batch.buf.data(), batch.buf.size()).ok());
+}
+
+TEST(EngineFailoverE2E, ProgressBatchKeepsOverallPendingWithMixedOutcomes) {
+    auto cfg = makeMinimalP2PConfig();
+    cfg->set("enable_auto_failover_on_poll", false);
+    cfg->set("max_failover_attempts", 0);
+    TransferEngineImpl engine(cfg);
+    ASSERT_TRUE(engine.available());
+
+    constexpr size_t kBufLen = 4096;
+    std::vector<uint8_t> failing_buf(kBufLen, 0xA7);
+    std::vector<uint8_t> pending_buf(kBufLen, 0xA8);
+    const uint64_t failing_addr =
+        reinterpret_cast<uint64_t>(failing_buf.data());
+
+    auto fake_rdma = std::make_shared<FakeTransport>(
+        RDMA, FakeTransport::StatusFactory{},
+        [failing_addr](const Request& req, int poll_count) {
+            if (req.target_offset == failing_addr) {
+                return TransferStatus{TransferStatusEnum::FAILED, 0};
+            }
+            if (poll_count == 1) {
+                return TransferStatus{TransferStatusEnum::PENDING, 0};
+            }
+            return TransferStatus{TransferStatusEnum::COMPLETED, req.length};
+        });
+    auto fake_tcp = std::make_shared<FakeTransport>(TCP);
+
+    std::string seg_name = engine.getSegmentName();
+    ASSERT_TRUE(fake_rdma->install(seg_name, nullptr, nullptr).ok());
+    ASSERT_TRUE(fake_tcp->install(seg_name, nullptr, nullptr).ok());
+    engine.swapTransportForTest(RDMA, fake_rdma);
+    engine.swapTransportForTest(TCP, fake_tcp);
+
+    ASSERT_TRUE(engine.registerLocalMemory(failing_buf.data(), kBufLen).ok());
+    ASSERT_TRUE(engine.registerLocalMemory(pending_buf.data(), kBufLen).ok());
+
+    BatchID batch_id = engine.allocateBatch(2);
+    ASSERT_NE(batch_id, (BatchID)0);
+
+    Request failing_req;
+    failing_req.opcode = Request::WRITE;
+    failing_req.source = failing_buf.data();
+    failing_req.target_id = LOCAL_SEGMENT_ID;
+    failing_req.target_offset = failing_addr;
+    failing_req.length = kBufLen;
+
+    Request pending_req;
+    pending_req.opcode = Request::WRITE;
+    pending_req.source = pending_buf.data();
+    pending_req.target_id = LOCAL_SEGMENT_ID;
+    pending_req.target_offset = reinterpret_cast<uint64_t>(pending_buf.data());
+    pending_req.length = kBufLen;
+
+    ASSERT_TRUE(
+        engine.submitTransfer(batch_id, {failing_req, pending_req}).ok());
+
+    TransferStatus overall_status{};
+    ASSERT_TRUE(engine.progressBatch(batch_id, overall_status).ok());
+    EXPECT_EQ(overall_status.s, TransferStatusEnum::PENDING);
+    EXPECT_EQ(fake_tcp->submit_calls.load(), 0);
+
+    ASSERT_TRUE(engine.progressBatch(batch_id, overall_status).ok());
+    EXPECT_EQ(overall_status.s, TransferStatusEnum::FAILED);
+
+    EXPECT_TRUE(engine.freeBatch(batch_id).ok());
+    EXPECT_TRUE(engine.unregisterLocalMemory(failing_buf.data(), kBufLen).ok());
+    EXPECT_TRUE(engine.unregisterLocalMemory(pending_buf.data(), kBufLen).ok());
+}
+
+TEST(EngineFailoverE2E,
+     WaitTransferCompletionUsesProgressBatchWhenPollDisabled) {
+    auto cfg = makeMinimalP2PConfig();
+    cfg->set("enable_auto_failover_on_poll", false);
+    TransferEngineImpl engine(cfg);
+    ASSERT_TRUE(engine.available());
+
+    CorruptedRdmaBatch batch;
+    submitCorruptedRdmaBatch(engine, batch, 0xA9);
+
+    EXPECT_TRUE(engine.waitTransferCompletion(batch.batch_id).ok());
+    EXPECT_EQ(batch.fake_tcp->submit_calls.load(), 1);
+
+    EXPECT_TRUE(
+        engine.unregisterLocalMemory(batch.buf.data(), batch.buf.size()).ok());
+}
+
+TEST(EngineFailoverE2E, TransferSyncUsesProgressBatchWhenPollDisabled) {
+    auto cfg = makeMinimalP2PConfig();
+    cfg->set("enable_auto_failover_on_poll", false);
+    TransferEngineImpl engine(cfg);
+    ASSERT_TRUE(engine.available());
+
+    auto fake_rdma = std::make_shared<FakeTransport>(RDMA);
+    auto fake_tcp = std::make_shared<FakeTransport>(TCP);
+
+    FaultPolicy rdma_policy;
+    rdma_policy.status_corrupt_rate = 1.0;
+    auto proxied_rdma =
+        std::make_shared<FaultProxyTransport>(fake_rdma, rdma_policy);
+
+    std::string seg_name = engine.getSegmentName();
+    ASSERT_TRUE(proxied_rdma->install(seg_name, nullptr, nullptr).ok());
+    ASSERT_TRUE(fake_tcp->install(seg_name, nullptr, nullptr).ok());
+    engine.swapTransportForTest(RDMA, proxied_rdma);
+    engine.swapTransportForTest(TCP, fake_tcp);
+
+    constexpr size_t kBufLen = 4096;
+    std::vector<uint8_t> buf(kBufLen, 0xB0);
+    ASSERT_TRUE(engine.registerLocalMemory(buf.data(), kBufLen).ok());
+
+    Request req;
+    req.opcode = Request::WRITE;
+    req.source = buf.data();
+    req.target_id = LOCAL_SEGMENT_ID;
+    req.target_offset = reinterpret_cast<uint64_t>(buf.data());
+    req.length = kBufLen;
+
+    EXPECT_TRUE(engine.transferSync({req}).ok());
+    EXPECT_EQ(fake_rdma->submit_calls.load(), 1);
+    EXPECT_EQ(fake_tcp->submit_calls.load(), 1);
+
+    EXPECT_TRUE(engine.unregisterLocalMemory(buf.data(), kBufLen).ok());
+}
+
+TEST(EngineFailoverE2E, ProgressBatchAdvancesExactlyOneStepPerCall) {
+    auto cfg = makeMinimalP2PConfig();
+    cfg->set("enable_auto_failover_on_poll", false);
+    TransferEngineImpl engine(cfg);
+    ASSERT_TRUE(engine.available());
+
+    // Reach COMPLETED only on the third poll; earlier polls stay PENDING.
+    auto fake_rdma = std::make_shared<FakeTransport>(
+        RDMA, FakeTransport::StatusFactory{},
+        [](const Request& req, int poll_count) {
+            if (poll_count < 3) {
+                return TransferStatus{TransferStatusEnum::PENDING, 0};
+            }
+            return TransferStatus{TransferStatusEnum::COMPLETED, req.length};
+        });
+    auto fake_tcp = std::make_shared<FakeTransport>(TCP);
+
+    std::string seg_name = engine.getSegmentName();
+    ASSERT_TRUE(fake_rdma->install(seg_name, nullptr, nullptr).ok());
+    ASSERT_TRUE(fake_tcp->install(seg_name, nullptr, nullptr).ok());
+    engine.swapTransportForTest(RDMA, fake_rdma);
+    engine.swapTransportForTest(TCP, fake_tcp);
+
+    constexpr size_t kBufLen = 4096;
+    std::vector<uint8_t> buf(kBufLen, 0xC0);
+    ASSERT_TRUE(engine.registerLocalMemory(buf.data(), kBufLen).ok());
+
+    BatchID batch_id = engine.allocateBatch(1);
+    ASSERT_NE(batch_id, (BatchID)0);
+
+    Request req;
+    req.opcode = Request::WRITE;
+    req.source = buf.data();
+    req.target_id = LOCAL_SEGMENT_ID;
+    req.target_offset = reinterpret_cast<uint64_t>(buf.data());
+    req.length = kBufLen;
+    ASSERT_TRUE(engine.submitTransfer(batch_id, {req}).ok());
+
+    // Each progressBatch call must perform exactly one poll on the underlying
+    // transport — no internal loop until completion.
+    TransferStatus overall_status{};
+    ASSERT_TRUE(engine.progressBatch(batch_id, overall_status).ok());
+    EXPECT_EQ(overall_status.s, TransferStatusEnum::PENDING);
+    EXPECT_EQ(fake_rdma->status_calls.load(), 1);
+
+    ASSERT_TRUE(engine.progressBatch(batch_id, overall_status).ok());
+    EXPECT_EQ(overall_status.s, TransferStatusEnum::PENDING);
+    EXPECT_EQ(fake_rdma->status_calls.load(), 2);
+
+    ASSERT_TRUE(engine.progressBatch(batch_id, overall_status).ok());
+    EXPECT_EQ(overall_status.s, TransferStatusEnum::COMPLETED);
+    EXPECT_EQ(fake_rdma->status_calls.load(), 3);
+    EXPECT_EQ(fake_tcp->submit_calls.load(), 0);
+
+    EXPECT_TRUE(engine.freeBatch(batch_id).ok());
+    EXPECT_TRUE(engine.unregisterLocalMemory(buf.data(), kBufLen).ok());
+}
+
+// ---------------------------------------------------------------------------
+// P1b: Both transports keep failing at status stage -> failover limit reached.
+// ---------------------------------------------------------------------------
+
+TEST(EngineFailoverE2E, BothTransportsFailExhaustsFailoverBudget) {
+    auto cfg = makeMinimalP2PConfig();
+    cfg->set("max_failover_attempts", 2);  // tighten budget
+    TransferEngineImpl engine(cfg);
+    ASSERT_TRUE(engine.available());
+
+    auto fake_rdma = std::make_shared<FakeTransport>(RDMA);
+    auto fake_tcp = std::make_shared<FakeTransport>(TCP);
+
+    FaultPolicy always_corrupt;
+    always_corrupt.status_corrupt_rate = 1.0;
+
+    auto proxied_rdma =
+        std::make_shared<FaultProxyTransport>(fake_rdma, always_corrupt);
+    auto proxied_tcp =
+        std::make_shared<FaultProxyTransport>(fake_tcp, always_corrupt);
+
+    std::string seg_name = engine.getSegmentName();
+    ASSERT_TRUE(proxied_rdma->install(seg_name, nullptr, nullptr).ok());
+    ASSERT_TRUE(proxied_tcp->install(seg_name, nullptr, nullptr).ok());
+
+    engine.swapTransportForTest(RDMA, proxied_rdma);
+    engine.swapTransportForTest(TCP, proxied_tcp);
+
+    constexpr size_t kBufLen = 4096;
+    std::vector<uint8_t> buf(kBufLen, 0xEF);
+    ASSERT_TRUE(engine.registerLocalMemory(buf.data(), kBufLen).ok());
+
+    BatchID batch_id = engine.allocateBatch(8);
+    Request req;
+    req.opcode = Request::WRITE;
+    req.source = buf.data();
+    req.target_id = LOCAL_SEGMENT_ID;
+    req.target_offset = reinterpret_cast<uint64_t>(buf.data());
+    req.length = kBufLen;
+
+    ASSERT_TRUE(engine.submitTransfer(batch_id, {req}).ok());
+
+    auto final_status = pollUntilDone(engine, batch_id, 0);
+    EXPECT_EQ(final_status.s, TransferStatusEnum::FAILED)
+        << "after exhausting failover budget, task must be permanently FAILED";
+
+    EXPECT_TRUE(engine.freeBatch(batch_id).ok());
+    EXPECT_TRUE(engine.unregisterLocalMemory(buf.data(), kBufLen).ok());
+}
+
+// ---------------------------------------------------------------------------
+// PerRequestFaultProxy
+//
+// FaultProxyTransport corrupts completions uniformly based on a rate. For
+// tests that need per-request control (e.g. "fail task0 but not task1"),
+// this subclass inspects each Request at submit time, records the sub_task
+// indices of "poisoned" requests, and flips only those from COMPLETED to
+// FAILED in getTransferStatus(). Completions for non-poisoned requests
+// pass through unchanged. Submit is never rejected.
+//
+// Tests submit each request in its own one-request batch, so the engine's
+// failover logic routes each task individually.
+// ---------------------------------------------------------------------------
+
+class PerRequestFaultProxy : public FaultProxyTransport {
+   public:
+    using Predicate = std::function<bool(const Request&)>;
+
+    PerRequestFaultProxy(std::shared_ptr<Transport> real, Predicate pred)
+        : FaultProxyTransport(std::move(real), FaultPolicy{}),
+          should_fail_(std::move(pred)) {}
+
+    Status submitTransferTasks(
+        SubBatchRef batch, const std::vector<Request>& request_list) override {
+        // Remember (sub_batch, sub_task_id) pairs for "poisoned" requests
+        // so getTransferStatus() can flip only those to FAILED later. The
+        // sub_batch pointer must be part of the key: different engine
+        // batches (BatchID) have different SubBatchRefs per transport, so
+        // sub_task_id=0 in batch A is a different task than sub_task_id=0
+        // in batch B. Keying only on sub_task_id would cross-contaminate.
+        const int base = static_cast<int>(batch->size());
+        for (size_t i = 0; i < request_list.size(); ++i) {
+            if (should_fail_(request_list[i])) {
+                poisoned_.insert({batch, base + static_cast<int>(i)});
+            }
+        }
+        return FaultProxyTransport::submitTransferTasks(batch, request_list);
+    }
+
+    Status getTransferStatus(SubBatchRef batch, int task_id,
+                             TransferStatus& status) override {
+        auto s = FaultProxyTransport::getTransferStatus(batch, task_id, status);
+        if (!s.ok()) return s;
+        if (poisoned_.count({batch, task_id}) &&
+            status.s == TransferStatusEnum::COMPLETED) {
+            status.s = TransferStatusEnum::FAILED;
+        }
+        return s;
+    }
+
+   private:
+    Predicate should_fail_;
+    std::set<std::pair<SubBatchRef, int>> poisoned_;
+};
+
+// ---------------------------------------------------------------------------
+// A: Mixed faults across many independent single-request submissions.
+//
+// Submit 10 one-request batches. For each submission, RDMA's completion
+// reports FAILED with 30% probability. Assert every task ultimately
+// COMPLETES and that RDMA/TCP submit counts add up consistently: every
+// failed RDMA completion must be followed by a TCP success.
+// ---------------------------------------------------------------------------
+
+TEST(EngineFailoverE2E, MixedFaultsAcrossManySubmissions) {
+    auto cfg = makeMinimalP2PConfig();
+    TransferEngineImpl engine(cfg);
+    ASSERT_TRUE(engine.available());
+
+    auto fake_rdma = std::make_shared<FakeTransport>(RDMA);
+    auto fake_tcp = std::make_shared<FakeTransport>(TCP);
+
+    FaultPolicy rdma_policy;
+    rdma_policy.status_corrupt_rate = 0.3;  // 30% of completions flip to FAILED
+    auto proxied_rdma =
+        std::make_shared<FaultProxyTransport>(fake_rdma, rdma_policy);
+
+    std::string seg_name = engine.getSegmentName();
+    ASSERT_TRUE(proxied_rdma->install(seg_name, nullptr, nullptr).ok());
+    ASSERT_TRUE(fake_tcp->install(seg_name, nullptr, nullptr).ok());
+    engine.swapTransportForTest(RDMA, proxied_rdma);
+    engine.swapTransportForTest(TCP, fake_tcp);
+
+    constexpr size_t kBufLen = 4096;
+    constexpr int kNumTasks = 10;
+    std::vector<uint8_t> buf(kBufLen * kNumTasks, 0x77);
+    ASSERT_TRUE(engine.registerLocalMemory(buf.data(), buf.size()).ok());
+
+    std::vector<BatchID> batches;
+    batches.reserve(kNumTasks);
+    for (int i = 0; i < kNumTasks; ++i) {
+        BatchID b = engine.allocateBatch(1);
+        ASSERT_NE(b, (BatchID)0);
+
+        Request req;
+        req.opcode = Request::WRITE;
+        req.source = buf.data() + (size_t)i * kBufLen;
+        req.target_id = LOCAL_SEGMENT_ID;
+        req.target_offset = reinterpret_cast<uint64_t>(req.source);
+        req.length = kBufLen;
+
+        ASSERT_TRUE(engine.submitTransfer(b, {req}).ok());
+        batches.push_back(b);
+    }
+
+    int completed = 0;
+    for (int i = 0; i < kNumTasks; ++i) {
+        auto ts = pollUntilDone(engine, batches[i], 0);
+        EXPECT_EQ(ts.s, TransferStatusEnum::COMPLETED)
+            << "task " << i << " did not complete";
+        if (ts.s == TransferStatusEnum::COMPLETED) ++completed;
+    }
+    EXPECT_EQ(completed, kNumTasks);
+
+    // Sanity on submit counts under status-corruption injection:
+    //   - Every task hits RDMA at submit time (submit itself succeeds,
+    //     the proxy only corrupts getTransferStatus), so rdma_ok == kNumTasks.
+    //   - A corrupted completion triggers a TCP failover (+1 tcp).
+    //     A clean completion stays on RDMA (+0 tcp).
+    //   - Therefore tcp_ok equals the number of corrupted completions,
+    //     which is in [0, kNumTasks].
+    const int rdma_ok = fake_rdma->submit_calls.load();
+    const int tcp_ok = fake_tcp->submit_calls.load();
+    EXPECT_EQ(rdma_ok, kNumTasks)
+        << "every task must attempt RDMA first (submit is always accepted)";
+    EXPECT_GE(tcp_ok, 0);
+    EXPECT_LE(tcp_ok, kNumTasks);
+
+    // With 30% corruption rate over 10 tasks, both branches are exercised
+    // with overwhelming probability (0.7^10 ~= 2.8% no failover; 0.3^10 ~=
+    // 6e-6 all failover). We don't assert strict counts because this is
+    // rate-based.
+
+    for (auto b : batches) EXPECT_TRUE(engine.freeBatch(b).ok());
+    EXPECT_TRUE(engine.unregisterLocalMemory(buf.data(), buf.size()).ok());
+}
+
+// ---------------------------------------------------------------------------
+// C1: max_failover_attempts = 0 disables failover entirely.
+//
+// resubmitTransferTask bumps failover_count and compares it against the
+// budget before anything else. With budget=0, the very first attempted
+// failover is rejected (++count == 1 > 0), so a single RDMA fault must
+// result in a permanently FAILED task without touching TCP.
+// ---------------------------------------------------------------------------
+
+TEST(EngineFailoverE2E, MaxFailoverAttemptsZeroDisablesFailover) {
+    auto cfg = makeMinimalP2PConfig();
+    cfg->set("max_failover_attempts", 0);
+    TransferEngineImpl engine(cfg);
+    ASSERT_TRUE(engine.available());
+
+    auto fake_rdma = std::make_shared<FakeTransport>(RDMA);
+    auto fake_tcp = std::make_shared<FakeTransport>(TCP);
+
+    FaultPolicy always_fail;
+    always_fail.status_corrupt_rate = 1.0;  // every completion flips to FAILED
+    auto proxied_rdma =
+        std::make_shared<FaultProxyTransport>(fake_rdma, always_fail);
+
+    std::string seg_name = engine.getSegmentName();
+    ASSERT_TRUE(proxied_rdma->install(seg_name, nullptr, nullptr).ok());
+    ASSERT_TRUE(fake_tcp->install(seg_name, nullptr, nullptr).ok());
+    engine.swapTransportForTest(RDMA, proxied_rdma);
+    engine.swapTransportForTest(TCP, fake_tcp);
+
+    constexpr size_t kBufLen = 4096;
+    std::vector<uint8_t> buf(kBufLen, 0x11);
+    ASSERT_TRUE(engine.registerLocalMemory(buf.data(), kBufLen).ok());
+
+    BatchID batch_id = engine.allocateBatch(1);
+    Request req;
+    req.opcode = Request::WRITE;
+    req.source = buf.data();
+    req.target_id = LOCAL_SEGMENT_ID;
+    req.target_offset = reinterpret_cast<uint64_t>(buf.data());
+    req.length = kBufLen;
+    ASSERT_TRUE(engine.submitTransfer(batch_id, {req}).ok());
+
+    auto ts = pollUntilDone(engine, batch_id, 0);
+    EXPECT_EQ(ts.s, TransferStatusEnum::FAILED)
+        << "with budget=0 the first fault must be permanent";
+
+    // TCP must never have been touched: budget=0 means no failover attempt.
+    EXPECT_EQ(fake_tcp->submit_calls.load(), 0);
+
+    EXPECT_TRUE(engine.freeBatch(batch_id).ok());
+    EXPECT_TRUE(engine.unregisterLocalMemory(buf.data(), kBufLen).ok());
+}
+
+// C2: max_failover_attempts = 1 allows exactly one failover.
+//
+// RDMA submit fails -> engine spends its single budget to switch to TCP
+// -> TCP succeeds -> task COMPLETES. A symmetric run where *both* fake
+// transports always fail is covered by BothTransportsFailExhaustsFailoverBudget
+// at budget=2; here we just confirm the happy-path boundary.
+
+TEST(EngineFailoverE2E, MaxFailoverAttemptsOneAllowsSingleFailover) {
+    auto cfg = makeMinimalP2PConfig();
+    cfg->set("max_failover_attempts", 1);
+    TransferEngineImpl engine(cfg);
+    ASSERT_TRUE(engine.available());
+
+    auto fake_rdma = std::make_shared<FakeTransport>(RDMA);
+    auto fake_tcp = std::make_shared<FakeTransport>(TCP);
+
+    FaultPolicy rdma_fail;
+    rdma_fail.status_corrupt_rate = 1.0;  // every completion flips to FAILED
+    auto proxied_rdma =
+        std::make_shared<FaultProxyTransport>(fake_rdma, rdma_fail);
+
+    std::string seg_name = engine.getSegmentName();
+    ASSERT_TRUE(proxied_rdma->install(seg_name, nullptr, nullptr).ok());
+    ASSERT_TRUE(fake_tcp->install(seg_name, nullptr, nullptr).ok());
+    engine.swapTransportForTest(RDMA, proxied_rdma);
+    engine.swapTransportForTest(TCP, fake_tcp);
+
+    constexpr size_t kBufLen = 4096;
+    std::vector<uint8_t> buf(kBufLen, 0x22);
+    ASSERT_TRUE(engine.registerLocalMemory(buf.data(), kBufLen).ok());
+
+    BatchID batch_id = engine.allocateBatch(1);
+    Request req;
+    req.opcode = Request::WRITE;
+    req.source = buf.data();
+    req.target_id = LOCAL_SEGMENT_ID;
+    req.target_offset = reinterpret_cast<uint64_t>(buf.data());
+    req.length = kBufLen;
+    ASSERT_TRUE(engine.submitTransfer(batch_id, {req}).ok());
+
+    auto ts = pollUntilDone(engine, batch_id, 0);
+    EXPECT_EQ(ts.s, TransferStatusEnum::COMPLETED)
+        << "budget=1 must permit exactly one failover to TCP";
+    EXPECT_EQ(fake_tcp->submit_calls.load(), 1);
+
+    EXPECT_TRUE(engine.freeBatch(batch_id).ok());
+    EXPECT_TRUE(engine.unregisterLocalMemory(buf.data(), kBufLen).ok());
+}
+
+// ---------------------------------------------------------------------------
+// D: Tasks in the same logical workload maintain independent failover
+//    state.
+//
+// Workload is two submissions:
+//   - task0: RDMA *and* TCP always fail -> must end FAILED
+//   - task1: RDMA succeeds              -> must end COMPLETED, no TCP hit
+//
+// This pins down that one task's exhausted failover budget does not
+// "infect" another task: the engine must track failover_count per-task,
+// not per-batch, per-transport, or per-engine. A regression that made
+// the counter global or batch-scoped would flip task1 to FAILED or
+// trigger a spurious TCP submit.
+//
+// Uses PerRequestFaultProxy so we can make RDMA fail *only* for the
+// specific buffer address of task0.
+// ---------------------------------------------------------------------------
+
+TEST(EngineFailoverE2E, PerTaskFailoverCountsAreIndependent) {
+    auto cfg = makeMinimalP2PConfig();
+    TransferEngineImpl engine(cfg);
+    ASSERT_TRUE(engine.available());
+
+    auto fake_rdma = std::make_shared<FakeTransport>(RDMA);
+    auto fake_tcp = std::make_shared<FakeTransport>(TCP);
+
+    constexpr size_t kBufLen = 4096;
+    std::vector<uint8_t> buf0(kBufLen, 0xAA);
+    std::vector<uint8_t> buf1(kBufLen, 0xBB);
+
+    const uint64_t failing_addr = reinterpret_cast<uint64_t>(buf0.data());
+
+    auto proxied_rdma = std::make_shared<PerRequestFaultProxy>(
+        fake_rdma, [failing_addr](const Request& r) {
+            return r.target_offset == failing_addr;
+        });
+    auto proxied_tcp = std::make_shared<PerRequestFaultProxy>(
+        fake_tcp, [failing_addr](const Request& r) {
+            return r.target_offset == failing_addr;
+        });
+
+    std::string seg_name = engine.getSegmentName();
+    ASSERT_TRUE(proxied_rdma->install(seg_name, nullptr, nullptr).ok());
+    ASSERT_TRUE(proxied_tcp->install(seg_name, nullptr, nullptr).ok());
+    engine.swapTransportForTest(RDMA, proxied_rdma);
+    engine.swapTransportForTest(TCP, proxied_tcp);
+
+    ASSERT_TRUE(engine.registerLocalMemory(buf0.data(), kBufLen).ok());
+    ASSERT_TRUE(engine.registerLocalMemory(buf1.data(), kBufLen).ok());
+
+    auto submit = [&](uint8_t* source) -> BatchID {
+        BatchID b = engine.allocateBatch(1);
+        EXPECT_NE(b, (BatchID)0);
+        Request req;
+        req.opcode = Request::WRITE;
+        req.source = source;
+        req.target_id = LOCAL_SEGMENT_ID;
+        req.target_offset = reinterpret_cast<uint64_t>(source);
+        req.length = kBufLen;
+        EXPECT_TRUE(engine.submitTransfer(b, {req}).ok());
+        return b;
+    };
+
+    BatchID b0 = submit(buf0.data());  // must FAIL on both
+    BatchID b1 = submit(buf1.data());  // must SUCCEED on RDMA
+
+    auto ts0 = pollUntilDone(engine, b0, 0);
+    auto ts1 = pollUntilDone(engine, b1, 0);
+
+    EXPECT_EQ(ts0.s, TransferStatusEnum::FAILED)
+        << "task0 should exhaust all transports and end FAILED";
+    EXPECT_EQ(ts1.s, TransferStatusEnum::COMPLETED)
+        << "task1 must be unaffected by task0's failover exhaustion";
+
+    // Under status-corruption injection submits always succeed; the proxy
+    // corrupts only getTransferStatus. So both tasks hit RDMA at submit
+    // (rdma +2). task0's RDMA completion is corrupted -> failover to TCP
+    // -> TCP submit (+1), also corrupted -> no more transports -> FAILED.
+    // task1 completes cleanly on RDMA with no TCP touch.
+    EXPECT_EQ(fake_rdma->submit_calls.load(), 2);
+    EXPECT_EQ(fake_tcp->submit_calls.load(), 1);
+
+    EXPECT_TRUE(engine.freeBatch(b0).ok());
+    EXPECT_TRUE(engine.freeBatch(b1).ok());
+    EXPECT_TRUE(engine.unregisterLocalMemory(buf0.data(), kBufLen).ok());
+    EXPECT_TRUE(engine.unregisterLocalMemory(buf1.data(), kBufLen).ok());
+}
+
+}  // namespace
+}  // namespace tent
+}  // namespace mooncake
diff --git a/mooncake-transfer-engine/tent/tests/failover_test.cpp b/mooncake-transfer-engine/tent/tests/failover_test.cpp
index 8e58ef04..6c5b9441 100644
--- a/mooncake-transfer-engine/tent/tests/failover_test.cpp
+++ b/mooncake-transfer-engine/tent/tests/failover_test.cpp
@@ -15,7 +15,6 @@
 #include <gtest/gtest.h>
 
 #include <memory>
-#include <set>
 #include <string>
 
 #include "tent/common/config.h"
@@ -112,27 +111,31 @@ TEST(FailoverConfigTest, ZeroDisablesFailover) {
     EXPECT_GT(task.failover_count, max_attempts);
 }
 
+TEST(FailoverConfigTest, AutoFailoverOnPollDefaultEnabled) {
+    auto conf = std::make_shared<Config>();
+    EXPECT_TRUE(conf->get("enable_auto_failover_on_poll", true));
+}
+
+TEST(FailoverConfigTest, AutoFailoverOnPollCanBeDisabled) {
+    auto conf = std::make_shared<Config>();
+    conf->set("enable_auto_failover_on_poll", false);
+    EXPECT_FALSE(conf->get("enable_auto_failover_on_poll", true));
+}
+
 // ---------------------------------------------------------------------------
 // TransportType name coverage (tests the static helper indirectly via
 // the enum values — the function itself is file-local in the .cpp, so we
 // verify the enum values are well-defined and usable in switch)
 // ---------------------------------------------------------------------------
 
-TEST(TransportTypeTest, AllEnumValuesDistinct) {
-    // Verify no accidental collisions in the enum
-    std::set<int> seen;
-    TransportType types[] = {RDMA,    MNNVL, SHM,          NVLINK, GDS,
-                             IOURING, TCP,   AscendDirect, UNSPEC};
-    for (auto t : types) {
-        EXPECT_TRUE(seen.insert(static_cast<int>(t)).second)
-            << "Duplicate TransportType value: " << static_cast<int>(t);
-    }
-}
-
-TEST(TransportTypeTest, SupportedCount) {
-    // kSupportedTransportTypes should cover all types except UNSPEC
-    EXPECT_EQ(kSupportedTransportTypes, 8u);
-    EXPECT_EQ(static_cast<int>(UNSPEC), 8);
+TEST(TransportTypeTest, UnspecIsZeroSlot) {
+    // UNSPEC must be value 0 so that zero-initialized tent_request /
+    // tent_memory_options structs (the C ABI) default to UNSPEC
+    // rather than silently pinning to whichever transport happened
+    // to occupy slot 0.
+    EXPECT_EQ(static_cast<int>(UNSPEC), 0);
+    EXPECT_EQ(kSupportedTransportTypes, static_cast<int>(kNumTransportTypes));
+    EXPECT_GT(kSupportedTransportTypes, static_cast<int>(UNSPEC));
 }
 
 // ---------------------------------------------------------------------------
@@ -214,6 +217,92 @@ TEST(FailoverStateMachineTest, StagingBypassesPriorityIncrement) {
     EXPECT_EQ(task.xport_priority, 0);  // Should NOT have incremented
 }
 
+// ---------------------------------------------------------------------------
+// Batch status aggregation correctness
+// ---------------------------------------------------------------------------
+
+// Verifies the invariant: a batch with one permanently-FAILED task and
+// another still-PENDING task must report PENDING (not FAILED) overall,
+// because the PENDING task may still complete.  The old code latched
+// overall_status to FAILED as soon as any task was terminal.
+TEST(BatchStatusAggregationTest, PendingTaskPreventsEarlyBatchFailure) {
+    // Two tasks in a batch scenario (simulated):
+    // Task A: permanently FAILED (exhausted failover budget)
+    // Task B: still PENDING (retrying on secondary transport)
+    //
+    // Expected overall status: PENDING (not FAILED), because Task B is
+    // still in-flight.  Only once Task B reaches a terminal state should
+    // the batch become terminal.
+
+    struct MockTask {
+        TransferStatusEnum status{PENDING};
+        bool derived{false};
+        size_t length{1024};
+    };
+
+    auto aggregateStatus =
+        [](const std::vector<MockTask>& tasks) -> TransferStatusEnum {
+        size_t success_tasks = 0;
+        size_t failed_tasks = 0;
+        size_t total_tasks = 0;
+        for (auto& t : tasks) {
+            if (t.derived) continue;
+            total_tasks++;
+            if (t.status == COMPLETED)
+                success_tasks++;
+            else if (t.status != PENDING)
+                failed_tasks++;
+        }
+        if (success_tasks == total_tasks) return COMPLETED;
+        if (success_tasks + failed_tasks == total_tasks) return FAILED;
+        return PENDING;
+    };
+
+    // Case 1: one FAILED + one PENDING → overall PENDING
+    {
+        std::vector<MockTask> tasks = {{FAILED, false, 1024},
+                                       {PENDING, false, 1024}};
+        EXPECT_EQ(aggregateStatus(tasks), PENDING);
+    }
+
+    // Case 2: one FAILED + one COMPLETED → overall FAILED
+    {
+        std::vector<MockTask> tasks = {{FAILED, false, 1024},
+                                       {COMPLETED, false, 1024}};
+        EXPECT_EQ(aggregateStatus(tasks), FAILED);
+    }
+
+    // Case 3: all COMPLETED → overall COMPLETED
+    {
+        std::vector<MockTask> tasks = {{COMPLETED, false, 1024},
+                                       {COMPLETED, false, 1024}};
+        EXPECT_EQ(aggregateStatus(tasks), COMPLETED);
+    }
+
+    // Case 4: all PENDING → overall PENDING
+    {
+        std::vector<MockTask> tasks = {{PENDING, false, 1024},
+                                       {PENDING, false, 1024}};
+        EXPECT_EQ(aggregateStatus(tasks), PENDING);
+    }
+
+    // Case 5: derived tasks are skipped
+    {
+        std::vector<MockTask> tasks = {{FAILED, false, 1024},
+                                       {FAILED, true, 1024},  // derived: skip
+                                       {PENDING, false, 1024}};
+        EXPECT_EQ(aggregateStatus(tasks), PENDING);
+    }
+
+    // Case 6: all non-derived are terminal with at least one failure
+    {
+        std::vector<MockTask> tasks = {{COMPLETED, false, 1024},
+                                       {FAILED, true, 1024},  // derived: skip
+                                       {FAILED, false, 1024}};
+        EXPECT_EQ(aggregateStatus(tasks), FAILED);
+    }
+}
+
 }  // namespace
 }  // namespace tent
 }  // namespace mooncake
diff --git a/mooncake-transfer-engine/tent/tests/hip_bandwidth_bench.cpp b/mooncake-transfer-engine/tent/tests/hip_bandwidth_bench.cpp
new file mode 100644
index 00000000..a265a808
--- /dev/null
+++ b/mooncake-transfer-engine/tent/tests/hip_bandwidth_bench.cpp
@@ -0,0 +1,228 @@
+// Copyright 2025 KVCache.AI
+//
+// Standalone intra-node HIP bandwidth benchmark.
+// Measures H2D, D2H, D2D (same GPU), and GPU-to-GPU (P2P via XGMI)
+// bandwidth across a sweep of transfer sizes.
+//
+// Build:
+//   hipcc -O3 -o hip_bandwidth_bench hip_bandwidth_bench.cpp
+//
+// Run:
+//   ./hip_bandwidth_bench              # all GPUs, all directions
+//   ./hip_bandwidth_bench 0 1          # GPU 0 as src, GPU 1 as dst
+
+#include <hip/hip_runtime.h>
+#include <stdio.h>
+#include <stdlib.h>
+#include <string.h>
+#include <algorithm>
+#include <vector>
+
+#define CHECK_HIP(call)                                                  \
+    do {                                                                 \
+        hipError_t err = (call);                                         \
+        if (err != hipSuccess) {                                         \
+            fprintf(stderr, "HIP error %s:%d: %s\n", __FILE__, __LINE__, \
+                    hipGetErrorString(err));                             \
+            exit(EXIT_FAILURE);                                          \
+        }                                                                \
+    } while (0)
+
+static const size_t kWarmupIter = 5;
+static const size_t kBenchIter = 20;
+
+// Returns elapsed microseconds for `iters` copies of `bytes`.
+static double benchCopy(void* dst, const void* src, size_t bytes,
+                        hipMemcpyKind kind, hipStream_t stream, size_t iters) {
+    // warmup
+    for (size_t i = 0; i < kWarmupIter; i++)
+        CHECK_HIP(hipMemcpyAsync(dst, src, bytes, kind, stream));
+    CHECK_HIP(hipStreamSynchronize(stream));
+
+    hipEvent_t start, stop;
+    CHECK_HIP(hipEventCreate(&start));
+    CHECK_HIP(hipEventCreate(&stop));
+
+    CHECK_HIP(hipEventRecord(start, stream));
+    for (size_t i = 0; i < iters; i++)
+        CHECK_HIP(hipMemcpyAsync(dst, src, bytes, kind, stream));
+    CHECK_HIP(hipEventRecord(stop, stream));
+    CHECK_HIP(hipStreamSynchronize(stream));
+
+    float ms = 0;
+    CHECK_HIP(hipEventElapsedTime(&ms, start, stop));
+    CHECK_HIP(hipEventDestroy(start));
+    CHECK_HIP(hipEventDestroy(stop));
+    return (double)ms * 1000.0;  // → microseconds
+}
+
+static void printHeader() {
+    printf("\n%-18s  %10s  %12s  %12s\n", "Size", "Iters", "BW (GB/s)",
+           "Lat (us)");
+    printf("%s\n", std::string(58, '-').c_str());
+}
+
+static void runBench(const char* label, void* dst, const void* src,
+                     size_t max_bytes, hipMemcpyKind kind, hipStream_t stream) {
+    printf("\n[%s]\n", label);
+    printHeader();
+
+    for (size_t bytes = 4096; bytes <= max_bytes; bytes *= 4) {
+        size_t iters = std::max<size_t>(kBenchIter, 200 * 1024 * 1024 / bytes);
+        double us = benchCopy(dst, src, bytes, kind, stream, iters);
+        double bw = (double)bytes * iters / (us / 1e6) / 1e9;  // GB/s
+        double lat = us / iters;
+        printf("%-18zu  %10zu  %12.3f  %12.2f\n", bytes, iters, bw, lat);
+    }
+}
+
+int main(int argc, char** argv) {
+    int gpu_count = 0;
+    CHECK_HIP(hipGetDeviceCount(&gpu_count));
+    printf("Detected %d AMD GPU(s)\n", gpu_count);
+
+    for (int g = 0; g < gpu_count; g++) {
+        hipDeviceProp_t prop;
+        CHECK_HIP(hipGetDeviceProperties(&prop, g));
+        printf("  GPU %d: %s  VRAM %.0f GiB  PCIe %04x:%02x:%02x.0\n", g,
+               prop.name,
+               (double)prop.totalGlobalMem / (1024.0 * 1024.0 * 1024.0),
+               prop.pciDomainID, prop.pciBusID, prop.pciDeviceID);
+    }
+
+    int src_gpu = (argc >= 2) ? atoi(argv[1]) : 0;
+    int dst_gpu = (argc >= 3) ? atoi(argv[2]) : (gpu_count > 1 ? 1 : 0);
+
+    const size_t kMaxBytes = 4ULL * 1024 * 1024 * 1024;  // 4 GiB
+    const size_t kHostBytes = 512ULL * 1024 * 1024;      // 512 MiB host
+
+    // --- allocate host pinned ---
+    void* h_src = nullptr;
+    void* h_dst = nullptr;
+    CHECK_HIP(hipHostMalloc(&h_src, kHostBytes, hipHostMallocDefault));
+    CHECK_HIP(hipHostMalloc(&h_dst, kHostBytes, hipHostMallocDefault));
+    memset(h_src, 0xAA, kHostBytes);
+
+    // --- allocate device buffers on src_gpu ---
+    CHECK_HIP(hipSetDevice(src_gpu));
+    void* d_src = nullptr;
+    CHECK_HIP(hipMalloc(&d_src, kMaxBytes));
+    CHECK_HIP(hipMemset(d_src, 0xBB, kMaxBytes));
+
+    hipStream_t stream_src;
+    CHECK_HIP(hipStreamCreate(&stream_src));
+
+    // --- allocate device buffer on dst_gpu ---
+    void* d_dst = nullptr;
+    hipStream_t stream_dst;
+    if (dst_gpu != src_gpu) {
+        CHECK_HIP(hipSetDevice(dst_gpu));
+        CHECK_HIP(hipMalloc(&d_dst, kMaxBytes));
+        CHECK_HIP(hipMemset(d_dst, 0xCC, kMaxBytes));
+        CHECK_HIP(hipStreamCreate(&stream_dst));
+    } else {
+        d_dst = (char*)d_src + kMaxBytes / 2;
+        stream_dst = stream_src;
+    }
+
+    // ----------------------------------------------------------------
+    printf("\n=== H2D: Host → GPU %d ===\n", src_gpu);
+    CHECK_HIP(hipSetDevice(src_gpu));
+    runBench("H2D", d_src, h_src, kHostBytes, hipMemcpyHostToDevice,
+             stream_src);
+
+    // ----------------------------------------------------------------
+    printf("\n=== D2H: GPU %d → Host ===\n", src_gpu);
+    runBench("D2H", h_dst, d_src, kHostBytes, hipMemcpyDeviceToHost,
+             stream_src);
+
+    // ----------------------------------------------------------------
+    printf("\n=== D2D (intra-GPU %d): same-device copy ===\n", src_gpu);
+    runBench("D2D intra", (char*)d_src + kMaxBytes / 2, d_src, kMaxBytes / 2,
+             hipMemcpyDeviceToDevice, stream_src);
+
+    // ----------------------------------------------------------------
+    if (dst_gpu != src_gpu) {
+        // Check P2P access
+        int can_access = 0;
+        hipDeviceCanAccessPeer(&can_access, src_gpu, dst_gpu);
+        if (can_access) {
+            CHECK_HIP(hipSetDevice(src_gpu));
+            CHECK_HIP(hipDeviceEnablePeerAccess(dst_gpu, 0));
+        }
+
+        char label[64];
+        snprintf(label, sizeof(label), "GPU %d → GPU %d (P2P %s)", src_gpu,
+                 dst_gpu, can_access ? "enabled" : "disabled");
+        printf("\n=== %s ===\n", label);
+
+        CHECK_HIP(hipSetDevice(src_gpu));
+        runBench(label, d_dst, d_src, kMaxBytes, hipMemcpyDeviceToDevice,
+                 stream_src);
+
+        // Reverse direction
+        snprintf(label, sizeof(label), "GPU %d → GPU %d (reverse)", dst_gpu,
+                 src_gpu);
+        printf("\n=== %s ===\n", label);
+        CHECK_HIP(hipSetDevice(dst_gpu));
+        runBench(label, d_src, d_dst, kMaxBytes, hipMemcpyDeviceToDevice,
+                 stream_dst);
+    }
+
+    // ----------------------------------------------------------------
+    // All-pairs GPU-to-GPU if no specific args
+    if (argc < 3 && gpu_count > 1) {
+        printf("\n=== All-pairs GPU P2P bandwidth (large block = 1GiB) ===\n");
+        printf("%-8s  %-8s  %12s\n", "Src", "Dst", "BW (GB/s)");
+        printf("%s\n", std::string(32, '-').c_str());
+
+        std::vector<void*> dbufs(gpu_count, nullptr);
+        for (int g = 0; g < gpu_count; g++) {
+            CHECK_HIP(hipSetDevice(g));
+            CHECK_HIP(hipMalloc(&dbufs[g], 1ULL * 1024 * 1024 * 1024));
+            CHECK_HIP(hipMemset(dbufs[g], g, 1ULL * 1024 * 1024 * 1024));
+        }
+
+        const size_t kP2PBytes = 1ULL * 1024 * 1024 * 1024;
+        for (int s = 0; s < gpu_count; s++) {
+            for (int d = 0; d < gpu_count; d++) {
+                if (s == d) {
+                    printf("GPU%-5d  GPU%-5d  %12s\n", s, d, "-");
+                    continue;
+                }
+                CHECK_HIP(hipSetDevice(s));
+                int ca = 0;
+                hipDeviceCanAccessPeer(&ca, s, d);
+                if (ca) hipDeviceEnablePeerAccess(d, 0);
+
+                hipStream_t st;
+                CHECK_HIP(hipStreamCreate(&st));
+                double us = benchCopy(dbufs[d], dbufs[s], kP2PBytes,
+                                      hipMemcpyDeviceToDevice, st, kBenchIter);
+                double bw = (double)kP2PBytes * kBenchIter / (us / 1e6) / 1e9;
+                printf("GPU%-5d  GPU%-5d  %12.2f\n", s, d, bw);
+                CHECK_HIP(hipStreamDestroy(st));
+            }
+        }
+
+        for (int g = 0; g < gpu_count; g++) {
+            CHECK_HIP(hipSetDevice(g));
+            CHECK_HIP(hipFree(dbufs[g]));
+        }
+    }
+
+    // cleanup
+    CHECK_HIP(hipSetDevice(src_gpu));
+    CHECK_HIP(hipFree(d_src));
+    CHECK_HIP(hipStreamDestroy(stream_src));
+    if (dst_gpu != src_gpu) {
+        CHECK_HIP(hipSetDevice(dst_gpu));
+        CHECK_HIP(hipFree(d_dst));
+        CHECK_HIP(hipStreamDestroy(stream_dst));
+    }
+    CHECK_HIP(hipHostFree(h_src));
+    CHECK_HIP(hipHostFree(h_dst));
+
+    printf("\nDone.\n");
+    return 0;
+}
diff --git a/mooncake-transfer-engine/tent/tests/ip_utils_test.cpp b/mooncake-transfer-engine/tent/tests/ip_utils_test.cpp
new file mode 100644
index 00000000..b6cc6b45
--- /dev/null
+++ b/mooncake-transfer-engine/tent/tests/ip_utils_test.cpp
@@ -0,0 +1,125 @@
+// Copyright 2025 KVCache.AI
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include <gtest/gtest.h>
+
+#include "tent/common/utils/ip.h"
+
+namespace mooncake {
+namespace tent {
+namespace {
+
+TEST(TentIpUtilsTest, ParsesBracketedIpv6WithPort) {
+    auto [host, port] = parseHostNameWithPort("[2001:db8::1]:9000", 1234);
+
+    EXPECT_EQ(host, "2001:db8::1");
+    EXPECT_EQ(port, 9000);
+}
+
+TEST(TentIpUtilsTest, ParsesBracketedIpv6WithoutPort) {
+    auto [host, port] = parseHostNameWithPort("[2001:db8::1]", 1234);
+
+    EXPECT_EQ(host, "2001:db8::1");
+    EXPECT_EQ(port, 1234);
+}
+
+TEST(TentIpUtilsTest, ParsesRawIpv6WithoutPort) {
+    auto [host, port] = parseHostNameWithPort("2001:db8::1", 1234);
+
+    EXPECT_EQ(host, "2001:db8::1");
+    EXPECT_EQ(port, 1234);
+}
+
+TEST(TentIpUtilsTest, ParsesIpv4WithPort) {
+    auto [host, port] = parseHostNameWithPort("192.168.1.10:9000", 1234);
+
+    EXPECT_EQ(host, "192.168.1.10");
+    EXPECT_EQ(port, 9000);
+}
+
+TEST(TentIpUtilsTest, ParsesHostnameWithPort) {
+    auto [host, port] = parseHostNameWithPort("example.com:9000", 1234);
+
+    EXPECT_EQ(host, "example.com");
+    EXPECT_EQ(port, 9000);
+}
+
+TEST(TentIpUtilsTest, InvalidBracketedIpv6UsesDefaultPort) {
+    auto [host, port] = parseHostNameWithPort("[2001::db8::1]:9000", 1234);
+
+    EXPECT_TRUE(host.empty());
+    EXPECT_EQ(port, 1234);
+}
+
+TEST(TentIpUtilsTest, InvalidRawIpv6UsesDefaultPort) {
+    auto [host, port] = parseHostNameWithPort("2001::db8::1", 1234);
+
+    EXPECT_TRUE(host.empty());
+    EXPECT_EQ(port, 1234);
+}
+
+TEST(TentIpUtilsTest, BracketedIpv6WithInvalidPortUsesDefaultPort) {
+    auto [host, port] = parseHostNameWithPort("[2001:db8::1]:bad", 1234);
+
+    EXPECT_EQ(host, "2001:db8::1");
+    EXPECT_EQ(port, 1234);
+}
+
+TEST(TentIpUtilsTest, EmptyInputUsesDefaultPort) {
+    auto [host, port] = parseHostNameWithPort("", 1234);
+
+    EXPECT_TRUE(host.empty());
+    EXPECT_EQ(port, 1234);
+}
+
+TEST(TentIpUtilsTest, BuildsBracketedIpv6Endpoint) {
+    EXPECT_EQ(buildIpAddrWithPort("2001:db8::1", 9000, true),
+              "[2001:db8::1]:9000");
+}
+
+TEST(TentIpUtilsTest, ParsesScopeIdWithoutPort) {
+    auto [host, port] =
+        parseHostNameWithPort("fe80::a236:bcff:fecb:a1be%eno2", 1234);
+
+    EXPECT_EQ(host, "fe80::a236:bcff:fecb:a1be%eno2");
+    EXPECT_EQ(port, 1234);
+}
+
+TEST(TentIpUtilsTest, ParsesScopeIdWithPort) {
+    auto [host, port] =
+        parseHostNameWithPort("fe80::a236:bcff:fecb:a1be%eno2:15773", 1234);
+
+    EXPECT_EQ(host, "fe80::a236:bcff:fecb:a1be%eno2");
+    EXPECT_EQ(port, 15773);
+}
+
+TEST(TentIpUtilsTest, ParsesBracketedScopeIdWithPort) {
+    auto [host, port] =
+        parseHostNameWithPort("[fe80::a236:bcff:fecb:a1be%eno2]:15773", 1234);
+
+    EXPECT_EQ(host, "fe80::a236:bcff:fecb:a1be%eno2");
+    EXPECT_EQ(port, 15773);
+}
+
+TEST(TentIpUtilsTest, ScopeIdWithInvalidPortUsesDefault) {
+    auto [host, port] =
+        parseHostNameWithPort("fe80::a236:bcff:fecb:a1be%eno2:bad", 1234);
+
+    EXPECT_EQ(host, "fe80::a236:bcff:fecb:a1be%eno2");
+    EXPECT_EQ(port, 1234);
+}
+
+}  // namespace
+}  // namespace tent
+}  // namespace mooncake
diff --git a/mooncake-transfer-engine/tent/tests/progress_worker_test.cpp b/mooncake-transfer-engine/tent/tests/progress_worker_test.cpp
new file mode 100644
index 00000000..feae180d
--- /dev/null
+++ b/mooncake-transfer-engine/tent/tests/progress_worker_test.cpp
@@ -0,0 +1,527 @@
+// Copyright 2026 KVCache.AI
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+//
+// Tests the ProgressWorker skeleton (issue #2116, follow-up to PR #2160).
+// Goals:
+//   * default-off behavior is byte-identical to the pre-worker world;
+//   * with the worker enabled and enable_auto_failover_on_poll=false, a
+//     caller that only submits + observes status (never calls
+//     progressBatch / waitTransferCompletion) still sees its batch
+//     progress through failover;
+//   * one notify advances the engine by exactly one progress step;
+//   * freeBatch racing the worker, including cross-thread free, is safe;
+//   * worker shuts down cleanly on engine destruction.
+
+#include <gtest/gtest.h>
+
+#include <atomic>
+#include <chrono>
+#include <cstdint>
+#include <cstdlib>
+#include <functional>
+#include <memory>
+#include <string>
+#include <thread>
+#include <utility>
+#include <vector>
+
+#include "tent/common/config.h"
+#include "tent/common/types.h"
+#include "tent/runtime/segment.h"
+#include "tent/runtime/transfer_engine_impl.h"
+#include "tent/runtime/transport.h"
+#include "tent/transport/fault_proxy/fault_proxy_transport.h"
+
+namespace mooncake {
+namespace tent {
+namespace {
+
+// ---------------------------------------------------------------------------
+// FakeTransport — same minimal shape used by engine_failover_e2e_test.cpp.
+// Kept local to avoid cross-test linkage; sources of truth diverging is OK
+// because we only exercise the "completes / status-can-be-overridden" surface.
+// ---------------------------------------------------------------------------
+
+class FakeSubBatch : public Transport::SubBatch {
+   public:
+    size_t size() const override { return task_count; }
+    size_t task_count = 0;
+    std::vector<Request> requests;
+    std::vector<TransferStatus> statuses;
+    std::vector<int> poll_counts;
+};
+
+class FakeTransport : public Transport {
+   public:
+    using StatusFactory = std::function<TransferStatus(const Request&)>;
+    using PollStatusFactory =
+        std::function<TransferStatus(const Request&, int)>;
+
+    explicit FakeTransport(TransportType self_type,
+                           StatusFactory status_factory = {},
+                           PollStatusFactory poll_status_factory = {})
+        : self_type_(self_type),
+          status_factory_(std::move(status_factory)),
+          poll_status_factory_(std::move(poll_status_factory)) {
+        caps.dram_to_dram = true;
+    }
+
+    std::atomic<int> install_calls{0};
+    std::atomic<int> submit_calls{0};
+    std::atomic<int> status_calls{0};
+    std::atomic<int> add_mem_calls{0};
+
+    Status install(std::string& /*local_segment_name*/,
+                   std::shared_ptr<ControlService> /*metadata*/,
+                   std::shared_ptr<Topology> /*local_topology*/,
+                   std::shared_ptr<Config> /*conf*/ = nullptr) override {
+        ++install_calls;
+        return Status::OK();
+    }
+
+    Status allocateSubBatch(SubBatchRef& batch, size_t /*max_size*/) override {
+        batch = new FakeSubBatch();
+        return Status::OK();
+    }
+
+    Status freeSubBatch(SubBatchRef& batch) override {
+        delete batch;
+        batch = nullptr;
+        return Status::OK();
+    }
+
+    Status submitTransferTasks(
+        SubBatchRef batch, const std::vector<Request>& request_list) override {
+        ++submit_calls;
+        auto* fb = static_cast<FakeSubBatch*>(batch);
+        for (const auto& req : request_list) {
+            if (status_factory_) {
+                fb->statuses.push_back(status_factory_(req));
+            } else {
+                fb->statuses.push_back(
+                    {TransferStatusEnum::COMPLETED, req.length});
+            }
+            fb->requests.push_back(req);
+            fb->poll_counts.push_back(0);
+            fb->task_count++;
+        }
+        return Status::OK();
+    }
+
+    Status getTransferStatus(SubBatchRef batch, int task_id,
+                             TransferStatus& status) override {
+        ++status_calls;
+        auto* fb = static_cast<FakeSubBatch*>(batch);
+        if (task_id < 0 || task_id >= (int)fb->statuses.size()) {
+            return Status::InvalidArgument("bad task_id" LOC_MARK);
+        }
+        ++fb->poll_counts[task_id];
+        if (poll_status_factory_) {
+            status = poll_status_factory_(fb->requests[task_id],
+                                          fb->poll_counts[task_id]);
+        } else {
+            status = fb->statuses[task_id];
+        }
+        return Status::OK();
+    }
+
+    Status addMemoryBuffer(BufferDesc& desc,
+                           const MemoryOptions& /*options*/) override {
+        ++add_mem_calls;
+        desc.transports.push_back(self_type_);
+        return Status::OK();
+    }
+
+    Status addMemoryBuffer(std::vector<BufferDesc>& desc_list,
+                           const MemoryOptions& options) override {
+        for (auto& d : desc_list) {
+            auto s = addMemoryBuffer(d, options);
+            if (!s.ok()) return s;
+        }
+        return Status::OK();
+    }
+
+    Status removeMemoryBuffer(BufferDesc& /*desc*/) override {
+        return Status::OK();
+    }
+
+    Status allocateLocalMemory(void** addr, size_t size,
+                               MemoryOptions& /*options*/) override {
+        *addr = std::malloc(size);
+        if (!*addr) return Status::InternalError("malloc failed" LOC_MARK);
+        return Status::OK();
+    }
+
+    Status freeLocalMemory(void* addr, size_t /*size*/) override {
+        std::free(addr);
+        return Status::OK();
+    }
+
+    bool warmupMemory(void* /*addr*/, size_t /*length*/) override {
+        return false;
+    }
+
+    const char* getName() const override {
+        return self_type_ == RDMA ? "<fake-rdma>" : "<fake-tcp>";
+    }
+
+   private:
+    TransportType self_type_;
+    StatusFactory status_factory_;
+    PollStatusFactory poll_status_factory_;
+};
+
+std::shared_ptr<Config> makeMinimalP2PConfig() {
+    auto cfg = std::make_shared<Config>();
+    cfg->set("metadata_type", "p2p");
+    cfg->set("metadata_servers", "");
+    cfg->set("rpc_server_hostname", "127.0.0.1");
+    cfg->set("rpc_server_port", "0");
+    cfg->set("log_level", "warning");
+    cfg->set("merge_requests", false);
+
+    cfg->set("transports/tcp/enable", false);
+    cfg->set("transports/shm/enable", false);
+    cfg->set("transports/rdma/enable", false);
+    cfg->set("transports/io_uring/enable", false);
+    cfg->set("transports/nvlink/enable", false);
+    cfg->set("transports/mnnvl/enable", false);
+    cfg->set("transports/gds/enable", false);
+    cfg->set("transports/ascend_direct/enable", false);
+
+    cfg->set("max_failover_attempts", 3);
+    return cfg;
+}
+
+// ---------------------------------------------------------------------------
+// 1. Default config: worker is not constructed, notifyBatchMaybeReady is a
+// no-op, and behavior matches PR #2160 exactly.
+// ---------------------------------------------------------------------------
+
+TEST(ProgressWorker, DisabledByDefaultLeavesBehaviorUnchanged) {
+    auto cfg = makeMinimalP2PConfig();
+    TransferEngineImpl engine(cfg);
+    ASSERT_TRUE(engine.available());
+
+    auto fake_rdma = std::make_shared<FakeTransport>(RDMA);
+    auto fake_tcp = std::make_shared<FakeTransport>(TCP);
+    std::string seg = engine.getSegmentName();
+    ASSERT_TRUE(fake_rdma->install(seg, nullptr, nullptr).ok());
+    ASSERT_TRUE(fake_tcp->install(seg, nullptr, nullptr).ok());
+    engine.swapTransportForTest(RDMA, fake_rdma);
+    engine.swapTransportForTest(TCP, fake_tcp);
+
+    constexpr size_t kBufLen = 4096;
+    std::vector<uint8_t> buf(kBufLen, 0x10);
+    ASSERT_TRUE(engine.registerLocalMemory(buf.data(), kBufLen).ok());
+
+    BatchID batch_id = engine.allocateBatch(1);
+    ASSERT_NE(batch_id, (BatchID)0);
+
+    Request req;
+    req.opcode = Request::WRITE;
+    req.source = buf.data();
+    req.target_id = LOCAL_SEGMENT_ID;
+    req.target_offset = reinterpret_cast<uint64_t>(buf.data());
+    req.length = kBufLen;
+    ASSERT_TRUE(engine.submitTransfer(batch_id, {req}).ok());
+
+    // No-op when the worker isn't constructed.
+    engine.notifyBatchMaybeReady(batch_id);
+    engine.notifyBatchMaybeReady((BatchID)0);
+
+    TransferStatus status{};
+    ASSERT_TRUE(engine.getTransferStatus(batch_id, status).ok());
+    EXPECT_EQ(status.s, TransferStatusEnum::COMPLETED);
+    EXPECT_EQ(fake_tcp->submit_calls.load(), 0);
+
+    EXPECT_TRUE(engine.freeBatch(batch_id).ok());
+    EXPECT_TRUE(engine.unregisterLocalMemory(buf.data(), kBufLen).ok());
+}
+
+// ---------------------------------------------------------------------------
+// 2. Worker drives failover when the caller does not poll with
+// allow_failover. This is the integration shape mooncake-pg needs.
+// ---------------------------------------------------------------------------
+
+TEST(ProgressWorker, ProgressesWithoutPollAutoFailover) {
+    auto cfg = makeMinimalP2PConfig();
+    cfg->set("enable_auto_failover_on_poll", false);
+    cfg->set("enable_progress_worker", true);
+    TransferEngineImpl engine(cfg);
+    ASSERT_TRUE(engine.available());
+
+    auto fake_rdma = std::make_shared<FakeTransport>(RDMA);
+    auto fake_tcp = std::make_shared<FakeTransport>(TCP);
+
+    FaultPolicy rdma_policy;
+    rdma_policy.status_corrupt_rate = 1.0;
+    auto proxied_rdma =
+        std::make_shared<FaultProxyTransport>(fake_rdma, rdma_policy);
+
+    std::string seg = engine.getSegmentName();
+    ASSERT_TRUE(proxied_rdma->install(seg, nullptr, nullptr).ok());
+    ASSERT_TRUE(fake_tcp->install(seg, nullptr, nullptr).ok());
+    engine.swapTransportForTest(RDMA, proxied_rdma);
+    engine.swapTransportForTest(TCP, fake_tcp);
+
+    constexpr size_t kBufLen = 4096;
+    std::vector<uint8_t> buf(kBufLen, 0xC1);
+    ASSERT_TRUE(engine.registerLocalMemory(buf.data(), kBufLen).ok());
+
+    BatchID batch_id = engine.allocateBatch(1);
+    ASSERT_NE(batch_id, (BatchID)0);
+
+    Request req;
+    req.opcode = Request::WRITE;
+    req.source = buf.data();
+    req.target_id = LOCAL_SEGMENT_ID;
+    req.target_offset = reinterpret_cast<uint64_t>(buf.data());
+    req.length = kBufLen;
+    ASSERT_TRUE(engine.submitTransfer(batch_id, {req}).ok());
+
+    // Drive the worker until terminal. We deliberately never call
+    // progressBatch / waitTransferCompletion here — only
+    // notifyBatchMaybeReady + observation-only getTransferStatus
+    // (which, with auto-failover-on-poll disabled, will not advance failover
+    // by itself).
+    TransferStatus status{};
+    const auto deadline =
+        std::chrono::steady_clock::now() + std::chrono::milliseconds(2000);
+    while (std::chrono::steady_clock::now() < deadline) {
+        engine.notifyBatchMaybeReady(batch_id);
+        std::this_thread::sleep_for(std::chrono::milliseconds(5));
+        if (fake_tcp->submit_calls.load() == 0) continue;
+        status = {};
+        ASSERT_TRUE(engine.getTransferStatus(batch_id, status).ok());
+        if (status.s == TransferStatusEnum::COMPLETED) break;
+    }
+    EXPECT_EQ(status.s, TransferStatusEnum::COMPLETED)
+        << "progress worker must drive failover when caller never calls "
+           "progressBatch";
+    EXPECT_EQ(fake_rdma->submit_calls.load(), 1);
+    EXPECT_GE(fake_tcp->submit_calls.load(), 1);
+
+    EXPECT_TRUE(engine.freeBatch(batch_id).ok());
+    EXPECT_TRUE(engine.unregisterLocalMemory(buf.data(), kBufLen).ok());
+}
+
+// ---------------------------------------------------------------------------
+// 3. One notify == one progress step. The worker must not loop internally
+// until completion; the next step requires another notify.
+// ---------------------------------------------------------------------------
+
+TEST(ProgressWorker, SingleNotifyAdvancesOneStep) {
+    auto cfg = makeMinimalP2PConfig();
+    cfg->set("enable_auto_failover_on_poll", false);
+    cfg->set("enable_progress_worker", true);
+    TransferEngineImpl engine(cfg);
+    ASSERT_TRUE(engine.available());
+
+    auto fake_rdma = std::make_shared<FakeTransport>(
+        RDMA, FakeTransport::StatusFactory{},
+        [](const Request& req, int poll_count) {
+            if (poll_count < 3) {
+                return TransferStatus{TransferStatusEnum::PENDING, 0};
+            }
+            return TransferStatus{TransferStatusEnum::COMPLETED, req.length};
+        });
+    auto fake_tcp = std::make_shared<FakeTransport>(TCP);
+
+    std::string seg = engine.getSegmentName();
+    ASSERT_TRUE(fake_rdma->install(seg, nullptr, nullptr).ok());
+    ASSERT_TRUE(fake_tcp->install(seg, nullptr, nullptr).ok());
+    engine.swapTransportForTest(RDMA, fake_rdma);
+    engine.swapTransportForTest(TCP, fake_tcp);
+
+    constexpr size_t kBufLen = 4096;
+    std::vector<uint8_t> buf(kBufLen, 0xC2);
+    ASSERT_TRUE(engine.registerLocalMemory(buf.data(), kBufLen).ok());
+
+    BatchID batch_id = engine.allocateBatch(1);
+    ASSERT_NE(batch_id, (BatchID)0);
+
+    Request req;
+    req.opcode = Request::WRITE;
+    req.source = buf.data();
+    req.target_id = LOCAL_SEGMENT_ID;
+    req.target_offset = reinterpret_cast<uint64_t>(buf.data());
+    req.length = kBufLen;
+    ASSERT_TRUE(engine.submitTransfer(batch_id, {req}).ok());
+
+    // Initial state: nothing polled yet.
+    EXPECT_EQ(fake_rdma->status_calls.load(), 0);
+
+    // Drive exactly one progress step via the worker. We can't observe the
+    // step instantly, but we can wait until status_calls increments by 1
+    // and then assert it does NOT keep climbing to 3 on its own.
+    engine.notifyBatchMaybeReady(batch_id);
+    const auto step_deadline =
+        std::chrono::steady_clock::now() + std::chrono::milliseconds(500);
+    while (std::chrono::steady_clock::now() < step_deadline &&
+           fake_rdma->status_calls.load() == 0) {
+        std::this_thread::sleep_for(std::chrono::milliseconds(2));
+    }
+    ASSERT_EQ(fake_rdma->status_calls.load(), 1)
+        << "worker should issue exactly one poll for a single notify";
+
+    // Give the worker a generous window to misbehave. status_calls must
+    // stay at 1 because we did not notify again and the engine did not
+    // reach a terminal state.
+    std::this_thread::sleep_for(std::chrono::milliseconds(50));
+    EXPECT_EQ(fake_rdma->status_calls.load(), 1)
+        << "worker must not loop internally until completion";
+
+    EXPECT_TRUE(engine.freeBatch(batch_id).ok());
+    EXPECT_TRUE(engine.unregisterLocalMemory(buf.data(), kBufLen).ok());
+}
+
+// ---------------------------------------------------------------------------
+// 4. freeBatch races with worker notifications. With ASAN/UBSAN this
+// catches missing alive_batches_ / progress_mutex_ coverage.
+// ---------------------------------------------------------------------------
+
+TEST(ProgressWorker, FreeBatchRacesWithWorker) {
+    auto cfg = makeMinimalP2PConfig();
+    cfg->set("enable_auto_failover_on_poll", false);
+    cfg->set("enable_progress_worker", true);
+    TransferEngineImpl engine(cfg);
+    ASSERT_TRUE(engine.available());
+
+    auto fake_rdma = std::make_shared<FakeTransport>(RDMA);
+    auto fake_tcp = std::make_shared<FakeTransport>(TCP);
+    std::string seg = engine.getSegmentName();
+    ASSERT_TRUE(fake_rdma->install(seg, nullptr, nullptr).ok());
+    ASSERT_TRUE(fake_tcp->install(seg, nullptr, nullptr).ok());
+    engine.swapTransportForTest(RDMA, fake_rdma);
+    engine.swapTransportForTest(TCP, fake_tcp);
+
+    constexpr size_t kBufLen = 4096;
+    std::vector<uint8_t> buf(kBufLen, 0xC3);
+    ASSERT_TRUE(engine.registerLocalMemory(buf.data(), kBufLen).ok());
+
+    // Concurrently spam stale notifications from a second thread while the
+    // main thread submits, frees, and re-allocates batches.
+    std::atomic<bool> stop{false};
+    std::atomic<BatchID> latest{0};
+    std::thread spammer([&] {
+        while (!stop.load(std::memory_order_acquire)) {
+            BatchID bid = latest.load(std::memory_order_acquire);
+            engine.notifyBatchMaybeReady(bid);
+            std::this_thread::sleep_for(std::chrono::microseconds(50));
+        }
+    });
+
+    constexpr int kRounds = 100;
+    for (int i = 0; i < kRounds; ++i) {
+        BatchID batch_id = engine.allocateBatch(1);
+        ASSERT_NE(batch_id, (BatchID)0);
+
+        Request req;
+        req.opcode = Request::WRITE;
+        req.source = buf.data();
+        req.target_id = LOCAL_SEGMENT_ID;
+        req.target_offset = reinterpret_cast<uint64_t>(buf.data());
+        req.length = kBufLen;
+        ASSERT_TRUE(engine.submitTransfer(batch_id, {req}).ok());
+
+        latest.store(batch_id, std::memory_order_release);
+        engine.notifyBatchMaybeReady(batch_id);
+
+        // Free immediately; the worker may pick the notification up after
+        // free. The progress_mutex_ + alive_batches_ guard must keep this
+        // safe.
+        EXPECT_TRUE(engine.freeBatch(batch_id).ok());
+    }
+
+    stop.store(true, std::memory_order_release);
+    spammer.join();
+
+    EXPECT_TRUE(engine.unregisterLocalMemory(buf.data(), kBufLen).ok());
+}
+
+TEST(ProgressWorker, BatchMayBeFreedFromDifferentThread) {
+    auto cfg = makeMinimalP2PConfig();
+    TransferEngineImpl engine(cfg);
+    ASSERT_TRUE(engine.available());
+
+    auto fake_rdma = std::make_shared<FakeTransport>(RDMA);
+    std::string seg = engine.getSegmentName();
+    ASSERT_TRUE(fake_rdma->install(seg, nullptr, nullptr).ok());
+    engine.swapTransportForTest(RDMA, fake_rdma);
+
+    constexpr size_t kBufLen = 4096;
+    std::vector<uint8_t> buf(kBufLen, 0xC4);
+    ASSERT_TRUE(engine.registerLocalMemory(buf.data(), kBufLen).ok());
+
+    BatchID batch_id = 0;
+    std::thread allocator([&] { batch_id = engine.allocateBatch(1); });
+    allocator.join();
+    ASSERT_NE(batch_id, (BatchID)0);
+
+    Request req;
+    req.opcode = Request::WRITE;
+    req.source = buf.data();
+    req.target_id = LOCAL_SEGMENT_ID;
+    req.target_offset = reinterpret_cast<uint64_t>(buf.data());
+    req.length = kBufLen;
+    ASSERT_TRUE(engine.submitTransfer(batch_id, {req}).ok());
+
+    TransferStatus status{};
+    ASSERT_TRUE(engine.getTransferStatus(batch_id, status).ok());
+    ASSERT_EQ(status.s, TransferStatusEnum::COMPLETED);
+
+    std::atomic<bool> free_ok{false};
+    std::thread freer([&] { free_ok.store(engine.freeBatch(batch_id).ok()); });
+    freer.join();
+    EXPECT_TRUE(free_ok.load());
+    EXPECT_TRUE(engine.getTransferStatus(batch_id, status).IsInvalidArgument());
+
+    EXPECT_TRUE(engine.unregisterLocalMemory(buf.data(), kBufLen).ok());
+}
+
+// ---------------------------------------------------------------------------
+// 5. Engine teardown joins the worker cleanly even with pending notifies.
+// ---------------------------------------------------------------------------
+
+TEST(ProgressWorker, EngineDestructorJoinsWorker) {
+    auto cfg = makeMinimalP2PConfig();
+    cfg->set("enable_auto_failover_on_poll", false);
+    cfg->set("enable_progress_worker", true);
+    {
+        TransferEngineImpl engine(cfg);
+        ASSERT_TRUE(engine.available());
+
+        auto fake_rdma = std::make_shared<FakeTransport>(RDMA);
+        auto fake_tcp = std::make_shared<FakeTransport>(TCP);
+        std::string seg = engine.getSegmentName();
+        ASSERT_TRUE(fake_rdma->install(seg, nullptr, nullptr).ok());
+        ASSERT_TRUE(fake_tcp->install(seg, nullptr, nullptr).ok());
+        engine.swapTransportForTest(RDMA, fake_rdma);
+        engine.swapTransportForTest(TCP, fake_tcp);
+
+        // Push some notifies for non-existent batches; worker must reject
+        // them via alive_batches_ check and stay alive.
+        for (int i = 0; i < 8; ++i) {
+            engine.notifyBatchMaybeReady((BatchID)(uintptr_t)0xdeadbeef);
+        }
+    }
+    // If teardown hangs or crashes here, gtest fails this test on timeout
+    // / signal — no further assert needed.
+    SUCCEED();
+}
+
+}  // namespace
+}  // namespace tent
+}  // namespace mooncake
diff --git a/mooncake-transfer-engine/tent/tests/rail_monitor_test.cpp b/mooncake-transfer-engine/tent/tests/rail_monitor_test.cpp
new file mode 100644
index 00000000..b625f39b
--- /dev/null
+++ b/mooncake-transfer-engine/tent/tests/rail_monitor_test.cpp
@@ -0,0 +1,194 @@
+// Copyright 2026 KVCache.AI
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include <gtest/gtest.h>
+
+#include <chrono>
+#include <thread>
+
+#include "tent/common/config.h"
+#include "tent/transport/rdma/rail_monitor.h"
+#include "tent/runtime/topology.h"
+
+namespace mooncake {
+namespace tent {
+namespace {
+
+// ---------------------------------------------------------------------------
+// Helper: build a minimal single-NIC Topology via Topology::parse().
+//
+// JSON keys expected by parse(): "nics" / "mems"
+// NicType enum:  NIC_RDMA=0, NIC_TCP=1, NIC_UNKNOWN=2
+// MemType enum:  MEM_HOST=0, MEM_CUDA=1, ...
+// device_list keys: "rank0", "rank1", ...
+// ---------------------------------------------------------------------------
+
+static std::shared_ptr<Topology> makeSingleNicTopology(const std::string& nic,
+                                                       int numa_node = 0) {
+    // One RDMA NIC (type=0) and one CUDA mem region (type=1) referencing it.
+    auto json_str = R"({
+        "nics": [{"name": ")" +
+                    nic + R"(", "type": 0, "numa_node": )" +
+                    std::to_string(numa_node) + R"(}],
+        "mems": [{
+            "name": "cuda0",
+            "type": 1,
+            "numa_node": )" +
+                    std::to_string(numa_node) +
+                    R"(,
+            "device_list": {"rank0": [0]}
+        }]
+    })";
+    auto topo = std::make_shared<Topology>();
+    auto status = topo->parse(json_str);
+    if (!status.ok()) {
+        ADD_FAILURE() << "Topology::parse failed: " << status.ToString();
+    }
+    return topo;
+}
+
+// ---------------------------------------------------------------------------
+// markRecovered resets error_count so failures start accumulating fresh
+// ---------------------------------------------------------------------------
+
+TEST(RailMonitorRecoverTest, RecoverResetsErrorCount) {
+    auto local = makeSingleNicTopology("mlx5_0");
+    auto remote = makeSingleNicTopology("mlx5_1");
+    RailMonitor rail;
+    ASSERT_TRUE(rail.load(local.get(), remote.get()).ok());
+    ASSERT_TRUE(rail.ready());
+
+    // Initially available
+    EXPECT_TRUE(rail.available(0, 0));
+
+    // One failure — not yet past default threshold (3)
+    rail.markFailed(0, 0);
+    EXPECT_TRUE(rail.available(0, 0));  // error_count=1, not paused
+
+    // A successful transfer — reset error_count back to 0
+    rail.markRecovered(0, 0);
+    EXPECT_TRUE(rail.available(0, 0));
+
+    // Failure again — counter starts fresh from 0, one hit is not enough
+    rail.markFailed(0, 0);
+    EXPECT_TRUE(rail.available(0, 0));
+}
+
+// ---------------------------------------------------------------------------
+// markRecovered un-pauses a rail that reached the failure threshold
+// ---------------------------------------------------------------------------
+
+TEST(RailMonitorRecoverTest, RecoverUnpausesPausedRail) {
+    auto local = makeSingleNicTopology("mlx5_0");
+    auto remote = makeSingleNicTopology("mlx5_1");
+    RailMonitor rail;
+    ASSERT_TRUE(rail.load(local.get(), remote.get()).ok());
+
+    // Drive error_count to the default threshold (3) to trigger pause
+    for (int i = 0; i < 3; ++i) rail.markFailed(0, 0);
+    EXPECT_FALSE(rail.available(0, 0))
+        << "Rail should be paused after 3 failures";
+
+    // A successful transfer proves the path is live — should un-pause
+    // immediately
+    rail.markRecovered(0, 0);
+    EXPECT_TRUE(rail.available(0, 0))
+        << "Rail should be available after recovery";
+}
+
+// ---------------------------------------------------------------------------
+// markRecovered on an unknown NIC pair is a no-op (no crash / no assert)
+// ---------------------------------------------------------------------------
+
+TEST(RailMonitorRecoverTest, RecoverUnknownPairIsNoop) {
+    auto local = makeSingleNicTopology("mlx5_0");
+    auto remote = makeSingleNicTopology("mlx5_1");
+    RailMonitor rail;
+    ASSERT_TRUE(rail.load(local.get(), remote.get()).ok());
+
+    // NIC IDs 5 and 7 are not in the topology — must not crash
+    EXPECT_NO_FATAL_FAILURE(rail.markRecovered(5, 7));
+}
+
+// ---------------------------------------------------------------------------
+// After recovery, findBestRemoteDevice maps back to the (only) available rail
+// ---------------------------------------------------------------------------
+
+TEST(RailMonitorRecoverTest, FindBestAfterRecovery) {
+    auto local = makeSingleNicTopology("mlx5_0");
+    auto remote = makeSingleNicTopology("mlx5_1");
+    RailMonitor rail;
+    ASSERT_TRUE(rail.load(local.get(), remote.get()).ok());
+
+    // Pause the only available rail
+    for (int i = 0; i < 3; ++i) rail.markFailed(0, 0);
+    EXPECT_FALSE(rail.available(0, 0));
+
+    // Recovery must rebuild best_mapping so findBestRemoteDevice works again
+    rail.markRecovered(0, 0);
+    EXPECT_TRUE(rail.available(0, 0));
+    int best = rail.findBestRemoteDevice(/*local_nic=*/0, /*remote_numa=*/0);
+    EXPECT_EQ(best, 0) << "Recovered rail should be the best remote device";
+}
+
+// ---------------------------------------------------------------------------
+// After recovery, the cooldown on the next pause must start from the
+// configured initial value, not a doubled value left over from the
+// previous cycle.
+//
+// Uses error_threshold=1 and cooldown=1s so each single failure triggers
+// a pause. If cooldown is correctly reset on recovery, the second pause
+// expires in ~1s; if the cooldown had carried over (bug), the second
+// pause would expire in ~2s.
+// ---------------------------------------------------------------------------
+
+TEST(RailMonitorRecoverTest, CooldownDoesNotCarryOverAfterRecovery) {
+    auto local = makeSingleNicTopology("mlx5_0");
+    auto remote = makeSingleNicTopology("mlx5_1");
+
+    Config cfg;
+    cfg.set(RailMonitor::kCfgErrorThreshold, 1);    // pause on first failure
+    cfg.set(RailMonitor::kCfgErrorWindowSecs, 60);  // wide: no window resets
+    cfg.set(RailMonitor::kCfgCooldownSecs, 1);      // small initial cooldown
+
+    RailMonitor rail;
+    ASSERT_TRUE(rail.load(local.get(), remote.get(), "", &cfg).ok());
+
+    // First pause cycle: single failure arms resume_time at now+1s.
+    rail.markFailed(0, 0);
+    EXPECT_FALSE(rail.available(0, 0));
+
+    // Recover: must clear st.cooldown so the next pause uses 1s again,
+    // not the 1s left over from cycle 1 (which would double to 2s).
+    rail.markRecovered(0, 0);
+    EXPECT_TRUE(rail.available(0, 0));
+
+    // Second pause cycle: single failure must arm resume_time at now+1s.
+    rail.markFailed(0, 0);
+    EXPECT_FALSE(rail.available(0, 0));
+
+    // Wait 1.5s: longer than the initial 1s cooldown, shorter than the
+    // 2s value the bug would produce. If cooldown was correctly reset on
+    // recovery, available() returns true; if it carried over, available()
+    // stays false until ~2s elapses.
+    std::this_thread::sleep_for(std::chrono::milliseconds(1500));
+    EXPECT_TRUE(rail.available(0, 0))
+        << "After recovery, the next pause must use the initial cooldown "
+           "(1s); staying paused past 1.5s indicates cooldown carried over "
+           "from the previous cycle.";
+}
+
+}  // namespace
+}  // namespace tent
+}  // namespace mooncake
diff --git a/mooncake-transfer-engine/tent/tests/request_merge_test.cpp b/mooncake-transfer-engine/tent/tests/request_merge_test.cpp
index 591bd17f..2b5fdef4 100644
--- a/mooncake-transfer-engine/tent/tests/request_merge_test.cpp
+++ b/mooncake-transfer-engine/tent/tests/request_merge_test.cpp
@@ -113,6 +113,105 @@ TEST(RequestMergeTest, KeepsNonContiguousRequestsSplit) {
     EXPECT_EQ(merged.task_lookup.at(1), 1u);
 }
 
+TEST(RequestMergeTest, DoesNotMergeAcrossDifferentTransportHint) {
+    std::array<char, 2048> source{};
+    const auto source_addr =
+        static_cast<uint64_t>(reinterpret_cast<uintptr_t>(source.data()));
+
+    std::vector<Request> requests = {
+        makeWriteRequest(source.data(), 0, 4096, 1024),
+        makeWriteRequest(source.data(), 1024, 5120, 1024),
+    };
+    requests[0].transport_hint = TransportType::RDMA;
+    requests[1].transport_hint = TransportType::TCP;
+
+    std::vector<RequestBoundaryInfo> boundaries = {
+        {makeBufferKey(source_addr, 2048), makeBufferKey(4096, 2048)},
+        {makeBufferKey(source_addr, 2048), makeBufferKey(4096, 2048)},
+    };
+
+    auto merged = mergeRequests(requests, boundaries, true);
+
+    // Mixed hints across an otherwise-contiguous run must stay split,
+    // otherwise the merged request would carry only one hint and the
+    // other half's pinning would silently disappear.
+    ASSERT_EQ(merged.request_list.size(), 2u);
+    EXPECT_EQ(merged.request_list[0].length, 1024u);
+    EXPECT_EQ(merged.request_list[1].length, 1024u);
+    EXPECT_NE(merged.request_list[0].transport_hint,
+              merged.request_list[1].transport_hint);
+}
+
+TEST(RequestMergeTest, MergesAdjacentRequestsWithSameTransportHint) {
+    std::array<char, 2048> source{};
+    const auto source_addr =
+        static_cast<uint64_t>(reinterpret_cast<uintptr_t>(source.data()));
+
+    std::vector<Request> requests = {
+        makeWriteRequest(source.data(), 0, 4096, 1024),
+        makeWriteRequest(source.data(), 1024, 5120, 1024),
+    };
+    requests[0].transport_hint = TransportType::TCP;
+    requests[1].transport_hint = TransportType::TCP;
+
+    std::vector<RequestBoundaryInfo> boundaries = {
+        {makeBufferKey(source_addr, 2048), makeBufferKey(4096, 2048)},
+        {makeBufferKey(source_addr, 2048), makeBufferKey(4096, 2048)},
+    };
+
+    auto merged = mergeRequests(requests, boundaries, true);
+
+    // Regression guard: adding transport_hint must not break merging
+    // when the hint is the same on both sides.
+    ASSERT_EQ(merged.request_list.size(), 1u);
+    EXPECT_EQ(merged.request_list[0].length, 2048u);
+    EXPECT_EQ(merged.request_list[0].transport_hint, TransportType::TCP);
+    EXPECT_EQ(merged.task_lookup.at(0), 0u);
+    EXPECT_EQ(merged.task_lookup.at(1), 0u);
+}
+
+TEST(RequestMergeTest, MixedHintBatchStillMergesPerHint) {
+    std::array<char, 4096> source{};
+    const auto source_addr =
+        static_cast<uint64_t>(reinterpret_cast<uintptr_t>(source.data()));
+
+    // Four requests interleave RDMA / TCP in submission order, but within
+    // each hint group both source and target are contiguous. After the
+    // hint-aware sort, the two RDMA requests fold into one, and the two
+    // TCP requests fold into the other.
+    //
+    // RDMA source pair: bytes [0, 2048) -> target [4096, 6144)
+    // TCP  source pair: bytes [2048, 4096) -> target [8192, 10240)
+    std::vector<Request> requests = {
+        makeWriteRequest(source.data(), 0, 4096, 1024),     // RDMA
+        makeWriteRequest(source.data(), 2048, 8192, 1024),  // TCP
+        makeWriteRequest(source.data(), 1024, 5120, 1024),  // RDMA (contig)
+        makeWriteRequest(source.data(), 3072, 9216, 1024),  // TCP  (contig)
+    };
+    requests[0].transport_hint = TransportType::RDMA;
+    requests[1].transport_hint = TransportType::TCP;
+    requests[2].transport_hint = TransportType::RDMA;
+    requests[3].transport_hint = TransportType::TCP;
+
+    std::vector<RequestBoundaryInfo> boundaries = {
+        {makeBufferKey(source_addr, 4096), makeBufferKey(4096, 6144)},
+        {makeBufferKey(source_addr, 4096), makeBufferKey(8192, 2048)},
+        {makeBufferKey(source_addr, 4096), makeBufferKey(4096, 6144)},
+        {makeBufferKey(source_addr, 4096), makeBufferKey(8192, 2048)},
+    };
+
+    auto merged = mergeRequests(requests, boundaries, true);
+
+    ASSERT_EQ(merged.request_list.size(), 2u);
+    // 0,2 (RDMA, contiguous on src+target) → one merged request, len 2048.
+    // 1,3 (TCP, contiguous on src+target)  → one merged request, len 2048.
+    EXPECT_EQ(merged.request_list[0].length, 2048u);
+    EXPECT_EQ(merged.request_list[1].length, 2048u);
+    EXPECT_EQ(merged.task_lookup.at(0), merged.task_lookup.at(2));
+    EXPECT_EQ(merged.task_lookup.at(1), merged.task_lookup.at(3));
+    EXPECT_NE(merged.task_lookup.at(0), merged.task_lookup.at(1));
+}
+
 }  // namespace
 }  // namespace tent
 }  // namespace mooncake
diff --git a/mooncake-transfer-engine/tent/tests/rocm_platform_test.cpp b/mooncake-transfer-engine/tent/tests/rocm_platform_test.cpp
new file mode 100644
index 00000000..d0b85f95
--- /dev/null
+++ b/mooncake-transfer-engine/tent/tests/rocm_platform_test.cpp
@@ -0,0 +1,249 @@
+// Copyright 2025 KVCache.AI
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include <gtest/gtest.h>
+
+#include <hip/hip_runtime.h>
+#include <memory>
+#include <string>
+#include <vector>
+
+#include "tent/common/config.h"
+#include "tent/platform/rocm.h"
+#include "tent/runtime/platform.h"
+#include "tent/runtime/topology.h"
+
+namespace mooncake {
+namespace tent {
+namespace {
+
+// Helper: returns the number of visible AMD GPUs (0 if ROCm not available).
+static int getHipDeviceCount() {
+    int count = 0;
+    if (hipGetDeviceCount(&count) != hipSuccess) return 0;
+    return count;
+}
+
+// ---------------------------------------------------------------------------
+// HIPStreamPool unit tests
+// ---------------------------------------------------------------------------
+
+TEST(HIPStreamPoolTest, AcquireAndRelease) {
+    if (getHipDeviceCount() == 0) GTEST_SKIP() << "No AMD GPU available";
+
+    HIPStreamPool pool;
+    HIPStreamHandle handle;
+    Status s = pool.acquire(handle);
+    EXPECT_TRUE(s.ok()) << s;
+    EXPECT_NE(handle.get(), nullptr);
+    // Destructor releases back to pool automatically — no assertion needed.
+}
+
+TEST(HIPStreamPoolTest, AcquireMultipleStreams) {
+    if (getHipDeviceCount() == 0) GTEST_SKIP() << "No AMD GPU available";
+
+    HIPStreamPool pool;
+    constexpr int kStreams = 4;
+    std::vector<HIPStreamHandle> handles(kStreams);
+    for (int i = 0; i < kStreams; i++) {
+        Status s = pool.acquire(handles[i]);
+        EXPECT_TRUE(s.ok()) << s;
+        EXPECT_NE(handles[i].get(), nullptr);
+    }
+}
+
+TEST(HIPStreamPoolTest, AcquireByExplicitDeviceId) {
+    if (getHipDeviceCount() == 0) GTEST_SKIP() << "No AMD GPU available";
+
+    HIPStreamPool pool;
+    HIPStreamHandle handle;
+    Status s = pool.acquire(handle, 0);
+    EXPECT_TRUE(s.ok()) << s;
+    EXPECT_NE(handle.get(), nullptr);
+}
+
+TEST(HIPStreamPoolTest, InvalidDeviceIdReturnsError) {
+    HIPStreamPool pool;
+    HIPStreamHandle handle;
+    // Device ID -5 is always invalid.
+    Status s = pool.acquire(handle, -5);
+    EXPECT_FALSE(s.ok());
+}
+
+// ---------------------------------------------------------------------------
+// RocmPlatform unit tests
+// ---------------------------------------------------------------------------
+
+class RocmPlatformTest : public ::testing::Test {
+   protected:
+    void SetUp() override {
+        if (getHipDeviceCount() == 0) {
+            GTEST_SKIP() << "No AMD GPU available";
+        }
+        auto conf = std::make_shared<Config>();
+        platform_ = std::make_shared<RocmPlatform>(conf);
+    }
+
+    std::shared_ptr<RocmPlatform> platform_;
+};
+
+TEST_F(RocmPlatformTest, TypeString) { EXPECT_EQ(platform_->type(), "rocm"); }
+
+TEST_F(RocmPlatformTest, ProbeDiscoversMems) {
+    std::vector<Topology::NicEntry> nics;
+    std::vector<Topology::MemEntry> mems;
+    Status s = platform_->probe(nics, mems);
+    EXPECT_TRUE(s.ok()) << s;
+    // At least one ROCm memory entry should be found.
+    bool found_rocm = false;
+    for (const auto& m : mems) {
+        if (m.type == Topology::MEM_ROCM) {
+            found_rocm = true;
+            EXPECT_EQ(m.name.substr(0, 5), "rocm:");
+        }
+    }
+    EXPECT_TRUE(found_rocm) << "probe() found no MEM_ROCM entries";
+}
+
+TEST_F(RocmPlatformTest, AllocateAndFreeDeviceMemory) {
+    MemoryOptions opts;
+    opts.location = "rocm:0";
+    void* ptr = nullptr;
+    constexpr size_t kSize = 1024 * 1024;  // 1 MiB
+
+    Status s = platform_->allocate(&ptr, kSize, opts);
+    ASSERT_TRUE(s.ok()) << s;
+    ASSERT_NE(ptr, nullptr);
+
+    EXPECT_EQ(platform_->getMemoryType(ptr), MTYPE_ROCM);
+
+    Status fs = platform_->free(ptr, kSize);
+    EXPECT_TRUE(fs.ok()) << fs;
+}
+
+TEST_F(RocmPlatformTest, AllocateAndFreeCpuMemory) {
+    MemoryOptions opts;
+    opts.location = "cpu:0";
+    void* ptr = nullptr;
+    constexpr size_t kSize = 4096;
+
+    Status s = platform_->allocate(&ptr, kSize, opts);
+    ASSERT_TRUE(s.ok()) << s;
+    ASSERT_NE(ptr, nullptr);
+
+    EXPECT_EQ(platform_->getMemoryType(ptr), MTYPE_CPU);
+
+    Status fs = platform_->free(ptr, kSize);
+    EXPECT_TRUE(fs.ok()) << fs;
+}
+
+TEST_F(RocmPlatformTest, CopyDeviceToDevice) {
+    MemoryOptions opts;
+    opts.location = "rocm:0";
+    constexpr size_t kSize = 256;
+    void *src = nullptr, *dst = nullptr;
+
+    ASSERT_TRUE(platform_->allocate(&src, kSize, opts).ok());
+    ASSERT_TRUE(platform_->allocate(&dst, kSize, opts).ok());
+
+    // Fill src with a pattern via hipMemset.
+    ASSERT_EQ(hipMemset(src, 0xAB, kSize), hipSuccess);
+
+    Status cs = platform_->copy(dst, src, kSize);
+    EXPECT_TRUE(cs.ok()) << cs;
+
+    // Verify the copy by reading back to host.
+    std::vector<uint8_t> host_buf(kSize, 0);
+    ASSERT_EQ(hipMemcpy(host_buf.data(), dst, kSize, hipMemcpyDeviceToHost),
+              hipSuccess);
+    for (size_t i = 0; i < kSize; i++) {
+        EXPECT_EQ(host_buf[i], 0xAB) << "mismatch at byte " << i;
+    }
+
+    EXPECT_TRUE(platform_->free(src, kSize).ok());
+    EXPECT_TRUE(platform_->free(dst, kSize).ok());
+}
+
+TEST_F(RocmPlatformTest, CopyHostToDevice) {
+    MemoryOptions dev_opts;
+    dev_opts.location = "rocm:0";
+    MemoryOptions cpu_opts;
+    cpu_opts.location = "cpu:0";
+    constexpr size_t kSize = 512;
+    void *dev = nullptr, *host_mem = nullptr;
+
+    ASSERT_TRUE(platform_->allocate(&dev, kSize, dev_opts).ok());
+    ASSERT_TRUE(platform_->allocate(&host_mem, kSize, cpu_opts).ok());
+
+    memset(host_mem, 0xCD, kSize);
+
+    Status cs = platform_->copy(dev, host_mem, kSize);
+    EXPECT_TRUE(cs.ok()) << cs;
+
+    std::vector<uint8_t> verify(kSize, 0);
+    ASSERT_EQ(hipMemcpy(verify.data(), dev, kSize, hipMemcpyDeviceToHost),
+              hipSuccess);
+    for (size_t i = 0; i < kSize; i++) {
+        EXPECT_EQ(verify[i], 0xCD) << "mismatch at byte " << i;
+    }
+
+    EXPECT_TRUE(platform_->free(dev, kSize).ok());
+    EXPECT_TRUE(platform_->free(host_mem, kSize).ok());
+}
+
+TEST_F(RocmPlatformTest, GetLocationDeviceMemory) {
+    MemoryOptions opts;
+    opts.location = "rocm:0";
+    constexpr size_t kSize = 4096;
+    void* ptr = nullptr;
+
+    ASSERT_TRUE(platform_->allocate(&ptr, kSize, opts).ok());
+    auto locs = platform_->getLocation(ptr, kSize);
+    ASSERT_FALSE(locs.empty());
+    EXPECT_EQ(locs[0].location, "rocm:0");
+    EXPECT_TRUE(platform_->free(ptr, kSize).ok());
+}
+
+TEST_F(RocmPlatformTest, GetMemoryTypeUnknownPointer) {
+    // An arbitrary user-space stack pointer should not be MTYPE_ROCM.
+    int local_var = 42;
+    MemoryType mtype = platform_->getMemoryType(&local_var);
+    // CPU or UNKNOWN — never MTYPE_ROCM.
+    EXPECT_NE(mtype, MTYPE_ROCM);
+}
+
+// ---------------------------------------------------------------------------
+// Multi-GPU tests (skipped when only one GPU is present)
+// ---------------------------------------------------------------------------
+
+TEST_F(RocmPlatformTest, AllocateOnSecondGpu) {
+    if (getHipDeviceCount() < 2) GTEST_SKIP() << "Requires >= 2 AMD GPUs";
+
+    MemoryOptions opts;
+    opts.location = "rocm:1";
+    void* ptr = nullptr;
+    constexpr size_t kSize = 1024;
+
+    Status s = platform_->allocate(&ptr, kSize, opts);
+    ASSERT_TRUE(s.ok()) << s;
+    ASSERT_NE(ptr, nullptr);
+
+    EXPECT_EQ(platform_->getMemoryType(ptr), MTYPE_ROCM);
+
+    EXPECT_TRUE(platform_->free(ptr, kSize).ok());
+}
+
+}  // namespace
+}  // namespace tent
+}  // namespace mooncake
diff --git a/mooncake-transfer-engine/tent/tests/runtime_queue_dispatch_test.cpp b/mooncake-transfer-engine/tent/tests/runtime_queue_dispatch_test.cpp
new file mode 100644
index 00000000..2a94c9e5
--- /dev/null
+++ b/mooncake-transfer-engine/tent/tests/runtime_queue_dispatch_test.cpp
@@ -0,0 +1,533 @@
+// Copyright 2026 KVCache.AI
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include <gtest/gtest.h>
+
+#include <atomic>
+#include <chrono>
+#include <cstdlib>
+#include <functional>
+#include <memory>
+#include <string>
+#include <thread>
+#include <utility>
+#include <vector>
+
+#include "tent/common/config.h"
+#include "tent/common/types.h"
+#include "tent/runtime/segment.h"
+#include "tent/runtime/transfer_engine_impl.h"
+#include "tent/runtime/transport.h"
+
+namespace mooncake {
+namespace tent {
+namespace {
+
+class FakeSubBatch : public Transport::SubBatch {
+   public:
+    size_t size() const override { return task_count; }
+
+    size_t task_count = 0;
+    std::vector<Request> requests;
+    std::vector<TransferStatus> statuses;
+    std::vector<int> poll_counts;
+};
+
+class FakeTransport : public Transport {
+   public:
+    using PollStatusFactory =
+        std::function<TransferStatus(const Request&, int)>;
+
+    explicit FakeTransport(TransportType self_type,
+                           PollStatusFactory poll_status_factory = {},
+                           bool notify_on_submit = false)
+        : self_type_(self_type),
+          poll_status_factory_(std::move(poll_status_factory)),
+          notify_on_submit_(notify_on_submit) {
+        caps.dram_to_dram = true;
+    }
+
+    std::atomic<int> submit_calls{0};
+    std::atomic<int> status_calls{0};
+
+    Status install(std::string&, std::shared_ptr<ControlService>,
+                   std::shared_ptr<Topology>,
+                   std::shared_ptr<Config> = nullptr) override {
+        return Status::OK();
+    }
+
+    Status allocateSubBatch(SubBatchRef& batch, size_t) override {
+        batch = new FakeSubBatch();
+        return Status::OK();
+    }
+
+    Status freeSubBatch(SubBatchRef& batch) override {
+        delete static_cast<FakeSubBatch*>(batch);
+        batch = nullptr;
+        return Status::OK();
+    }
+
+    Status submitTransferTasks(SubBatchRef batch,
+                               const std::vector<Request>& requests) override {
+        ++submit_calls;
+        auto* fake = static_cast<FakeSubBatch*>(batch);
+        for (const auto& request : requests) {
+            fake->requests.push_back(request);
+            fake->statuses.push_back(
+                {TransferStatusEnum::COMPLETED, request.length});
+            fake->poll_counts.push_back(0);
+            ++fake->task_count;
+        }
+        if (notify_on_submit_) batch->notifyProgress();
+        return Status::OK();
+    }
+
+    Status getTransferStatus(SubBatchRef batch, int task_id,
+                             TransferStatus& status) override {
+        ++status_calls;
+        auto* fake = static_cast<FakeSubBatch*>(batch);
+        if (task_id < 0 || task_id >= (int)fake->statuses.size()) {
+            return Status::InvalidArgument("bad task_id" LOC_MARK);
+        }
+        ++fake->poll_counts[task_id];
+        if (poll_status_factory_) {
+            status = poll_status_factory_(fake->requests[task_id],
+                                          fake->poll_counts[task_id]);
+        } else {
+            status = fake->statuses[task_id];
+        }
+        return Status::OK();
+    }
+
+    Status addMemoryBuffer(BufferDesc& desc, const MemoryOptions&) override {
+        desc.transports.push_back(self_type_);
+        return Status::OK();
+    }
+
+    Status addMemoryBuffer(std::vector<BufferDesc>& desc_list,
+                           const MemoryOptions& options) override {
+        for (auto& desc : desc_list) {
+            auto status = addMemoryBuffer(desc, options);
+            if (!status.ok()) return status;
+        }
+        return Status::OK();
+    }
+
+    Status removeMemoryBuffer(BufferDesc&) override { return Status::OK(); }
+
+    Status allocateLocalMemory(void** addr, size_t size,
+                               MemoryOptions&) override {
+        *addr = std::malloc(size);
+        return *addr ? Status::OK()
+                     : Status::InternalError("malloc failed" LOC_MARK);
+    }
+
+    Status freeLocalMemory(void* addr, size_t) override {
+        std::free(addr);
+        return Status::OK();
+    }
+
+    bool warmupMemory(void*, size_t) override { return false; }
+
+    const char* getName() const override { return "<fake-rdma>"; }
+
+   private:
+    TransportType self_type_;
+    PollStatusFactory poll_status_factory_;
+    bool notify_on_submit_;
+};
+
+std::shared_ptr<Config> makeRuntimeQueueConfig(size_t max_dispatch_owners,
+                                               size_t max_dispatch_bytes,
+                                               bool merge_requests = false) {
+    auto cfg = std::make_shared<Config>();
+    cfg->set("metadata_type", "p2p");
+    cfg->set("metadata_servers", "");
+    cfg->set("rpc_server_hostname", "127.0.0.1");
+    cfg->set("rpc_server_port", "0");
+    cfg->set("log_level", "warning");
+    cfg->set("merge_requests", merge_requests);
+    cfg->set("enable_runtime_queue", true);
+    cfg->set("runtime_queue/max_outstanding_owners", 16UL);
+    cfg->set("runtime_queue/max_outstanding_bytes", 1UL << 20);
+    cfg->set("runtime_queue/max_dispatch_owners", max_dispatch_owners);
+    cfg->set("runtime_queue/max_dispatch_bytes", max_dispatch_bytes);
+    cfg->set("runtime_queue/staging_owner_reserve", 0UL);
+    cfg->set("runtime_queue/staging_byte_reserve", 0UL);
+    cfg->set("runtime_queue/progress_fallback_interval_us", 50000UL);
+
+    cfg->set("transports/tcp/enable", false);
+    cfg->set("transports/shm/enable", false);
+    cfg->set("transports/rdma/enable", false);
+    cfg->set("transports/io_uring/enable", false);
+    cfg->set("transports/nvlink/enable", false);
+    cfg->set("transports/mnnvl/enable", false);
+    cfg->set("transports/gds/enable", false);
+    cfg->set("transports/ascend_direct/enable", false);
+    return cfg;
+}
+
+void installFakeRdma(TransferEngineImpl& engine,
+                     const std::shared_ptr<FakeTransport>& fake_rdma) {
+    std::string seg_name = engine.getSegmentName();
+    ASSERT_TRUE(fake_rdma->install(seg_name, nullptr, nullptr).ok());
+    engine.swapTransportForTest(RDMA, fake_rdma);
+}
+
+Request makeLocalWrite(uint8_t* ptr, size_t length) {
+    Request request;
+    request.opcode = Request::WRITE;
+    request.source = ptr;
+    request.target_id = LOCAL_SEGMENT_ID;
+    request.target_offset = reinterpret_cast<uint64_t>(ptr);
+    request.length = length;
+    request.transport_hint = RDMA;
+    return request;
+}
+
+TEST(RuntimeQueueDispatch, RejectsOverfullBatchBeforePublishingTasks) {
+    auto cfg = makeRuntimeQueueConfig(1, 1UL << 20);
+    TransferEngineImpl engine(cfg);
+    ASSERT_TRUE(engine.available());
+
+    auto fake_rdma = std::make_shared<FakeTransport>(RDMA);
+    installFakeRdma(engine, fake_rdma);
+
+    constexpr size_t kReqLen = 4096;
+    std::vector<uint8_t> buffer(kReqLen * 2, 0x11);
+    ASSERT_TRUE(engine.registerLocalMemory(buffer.data(), buffer.size()).ok());
+
+    BatchID batch = engine.allocateBatch(1);
+    ASSERT_NE(batch, (BatchID)0);
+
+    auto status = engine.submitTransfer(
+        batch, {makeLocalWrite(buffer.data(), kReqLen),
+                makeLocalWrite(buffer.data() + kReqLen, kReqLen)});
+    EXPECT_TRUE(status.IsTooManyRequests()) << status.ToString();
+    EXPECT_EQ(fake_rdma->submit_calls.load(), 0);
+
+    TransferStatus task_status{};
+    EXPECT_TRUE(
+        engine.getTransferStatus(batch, 0, task_status).IsInvalidArgument());
+
+    EXPECT_TRUE(engine.freeBatch(batch).ok());
+    EXPECT_TRUE(
+        engine.unregisterLocalMemory(buffer.data(), buffer.size()).ok());
+}
+
+TEST(RuntimeQueueDispatch, RejectsOwnerLargerThanDispatchByteWindow) {
+    constexpr size_t kReqLen = 4096;
+    auto cfg = makeRuntimeQueueConfig(1, kReqLen - 1);
+    TransferEngineImpl engine(cfg);
+    ASSERT_TRUE(engine.available());
+
+    auto fake_rdma = std::make_shared<FakeTransport>(RDMA);
+    installFakeRdma(engine, fake_rdma);
+
+    std::vector<uint8_t> buffer(kReqLen, 0x77);
+    ASSERT_TRUE(engine.registerLocalMemory(buffer.data(), buffer.size()).ok());
+
+    BatchID batch = engine.allocateBatch(1);
+    ASSERT_NE(batch, (BatchID)0);
+
+    auto status =
+        engine.submitTransfer(batch, {makeLocalWrite(buffer.data(), kReqLen)});
+    EXPECT_TRUE(status.IsTooManyRequests()) << status.ToString();
+    EXPECT_EQ(fake_rdma->submit_calls.load(), 0);
+
+    TransferStatus task_status{};
+    EXPECT_TRUE(
+        engine.getTransferStatus(batch, 0, task_status).IsInvalidArgument());
+
+    EXPECT_TRUE(engine.freeBatch(batch).ok());
+    EXPECT_TRUE(
+        engine.unregisterLocalMemory(buffer.data(), buffer.size()).ok());
+}
+
+TEST(RuntimeQueueDispatch, RejectsEmptyDispatchWindowConfig) {
+    auto cfg = makeRuntimeQueueConfig(0, 1UL << 20);
+    TransferEngineImpl engine(cfg);
+
+    EXPECT_FALSE(engine.available());
+
+    cfg = makeRuntimeQueueConfig(1, 0);
+    TransferEngineImpl byte_window_engine(cfg);
+
+    EXPECT_FALSE(byte_window_engine.available());
+}
+
+TEST(RuntimeQueueDispatch, DispatchesOnlyOneWindowOnSubmit) {
+    auto cfg = makeRuntimeQueueConfig(1, 1UL << 20);
+    TransferEngineImpl engine(cfg);
+    ASSERT_TRUE(engine.available());
+
+    auto fake_rdma = std::make_shared<FakeTransport>(RDMA);
+    installFakeRdma(engine, fake_rdma);
+
+    constexpr size_t kReqLen = 4096;
+    std::vector<uint8_t> buffer(kReqLen * 2, 0x22);
+    ASSERT_TRUE(engine.registerLocalMemory(buffer.data(), buffer.size()).ok());
+
+    BatchID batch = engine.allocateBatch(2);
+    ASSERT_NE(batch, (BatchID)0);
+
+    ASSERT_TRUE(
+        engine
+            .submitTransfer(batch,
+                            {makeLocalWrite(buffer.data(), kReqLen),
+                             makeLocalWrite(buffer.data() + kReqLen, kReqLen)})
+            .ok());
+    EXPECT_EQ(fake_rdma->submit_calls.load(), 1);
+
+    TransferStatus first{};
+    ASSERT_TRUE(engine.getTransferStatus(batch, 0, first).ok());
+    EXPECT_EQ(first.s, TransferStatusEnum::COMPLETED);
+    EXPECT_EQ(fake_rdma->submit_calls.load(), 2);
+
+    TransferStatus second{};
+    ASSERT_TRUE(engine.getTransferStatus(batch, 1, second).ok());
+    EXPECT_EQ(second.s, TransferStatusEnum::COMPLETED);
+
+    EXPECT_TRUE(engine.freeBatch(batch).ok());
+    EXPECT_TRUE(
+        engine.unregisterLocalMemory(buffer.data(), buffer.size()).ok());
+}
+
+TEST(RuntimeQueueDispatch, KeepsDispatchWindowUntilOwnerIsTerminal) {
+    auto cfg = makeRuntimeQueueConfig(1, 1UL << 20);
+    TransferEngineImpl engine(cfg);
+    ASSERT_TRUE(engine.available());
+
+    std::atomic<bool> complete_first{false};
+    auto fake_rdma = std::make_shared<FakeTransport>(
+        RDMA, [&complete_first](const Request& request, int) {
+            if (!complete_first.load()) {
+                return TransferStatus{TransferStatusEnum::PENDING, 0};
+            }
+            return TransferStatus{TransferStatusEnum::COMPLETED,
+                                  request.length};
+        });
+    installFakeRdma(engine, fake_rdma);
+
+    constexpr size_t kReqLen = 4096;
+    std::vector<uint8_t> buffer(kReqLen * 2, 0x55);
+    ASSERT_TRUE(engine.registerLocalMemory(buffer.data(), buffer.size()).ok());
+
+    BatchID batch = engine.allocateBatch(2);
+    ASSERT_NE(batch, (BatchID)0);
+
+    ASSERT_TRUE(
+        engine
+            .submitTransfer(batch,
+                            {makeLocalWrite(buffer.data(), kReqLen),
+                             makeLocalWrite(buffer.data() + kReqLen, kReqLen)})
+            .ok());
+    EXPECT_EQ(fake_rdma->submit_calls.load(), 1);
+
+    TransferStatus first{};
+    ASSERT_TRUE(engine.getTransferStatus(batch, 0, first).ok());
+    EXPECT_EQ(first.s, TransferStatusEnum::PENDING);
+    EXPECT_EQ(fake_rdma->submit_calls.load(), 1);
+
+    TransferStatus second{};
+    ASSERT_TRUE(engine.getTransferStatus(batch, 1, second).ok());
+    EXPECT_EQ(second.s, TransferStatusEnum::PENDING);
+    EXPECT_EQ(fake_rdma->submit_calls.load(), 1);
+
+    complete_first.store(true);
+    ASSERT_TRUE(engine.getTransferStatus(batch, 0, first).ok());
+    EXPECT_EQ(first.s, TransferStatusEnum::COMPLETED);
+    EXPECT_EQ(fake_rdma->submit_calls.load(), 2);
+
+    ASSERT_TRUE(engine.getTransferStatus(batch, 1, second).ok());
+    EXPECT_EQ(second.s, TransferStatusEnum::COMPLETED);
+
+    EXPECT_TRUE(engine.freeBatch(batch).ok());
+    EXPECT_TRUE(
+        engine.unregisterLocalMemory(buffer.data(), buffer.size()).ok());
+}
+
+TEST(RuntimeQueueDispatch, ProgressWorkerRefillsWindowFromTransportNotify) {
+    auto cfg = makeRuntimeQueueConfig(1, 1UL << 20);
+    cfg->set("enable_progress_worker", true);
+    cfg->set("runtime_queue/progress_fallback_interval_us", 0UL);
+    TransferEngineImpl engine(cfg);
+    ASSERT_TRUE(engine.available());
+
+    auto fake_rdma = std::make_shared<FakeTransport>(
+        RDMA, FakeTransport::PollStatusFactory{}, true);
+    installFakeRdma(engine, fake_rdma);
+
+    constexpr size_t kReqLen = 4096;
+    std::vector<uint8_t> buffer(kReqLen * 2, 0x66);
+    ASSERT_TRUE(engine.registerLocalMemory(buffer.data(), buffer.size()).ok());
+
+    BatchID batch = engine.allocateBatch(2);
+    ASSERT_NE(batch, (BatchID)0);
+
+    ASSERT_TRUE(
+        engine
+            .submitTransfer(batch,
+                            {makeLocalWrite(buffer.data(), kReqLen),
+                             makeLocalWrite(buffer.data() + kReqLen, kReqLen)})
+            .ok());
+    EXPECT_EQ(fake_rdma->submit_calls.load(), 1);
+
+    const auto deadline =
+        std::chrono::steady_clock::now() + std::chrono::milliseconds(1000);
+    while (std::chrono::steady_clock::now() < deadline &&
+           fake_rdma->submit_calls.load() < 2) {
+        std::this_thread::sleep_for(std::chrono::milliseconds(2));
+    }
+    EXPECT_EQ(fake_rdma->submit_calls.load(), 2);
+
+    TransferStatus status{};
+    ASSERT_TRUE(engine.getTransferStatus(batch, status).ok());
+    EXPECT_EQ(status.s, TransferStatusEnum::COMPLETED);
+
+    EXPECT_TRUE(engine.freeBatch(batch).ok());
+    EXPECT_TRUE(
+        engine.unregisterLocalMemory(buffer.data(), buffer.size()).ok());
+}
+
+TEST(RuntimeQueueDispatch, RuntimeQueueDrainsWithoutUserPolling) {
+    auto cfg = makeRuntimeQueueConfig(1, 1UL << 20);
+    cfg->set("runtime_queue/progress_fallback_interval_us", 1000UL);
+    TransferEngineImpl engine(cfg);
+    ASSERT_TRUE(engine.available());
+
+    auto fake_rdma = std::make_shared<FakeTransport>(
+        RDMA, [](const Request& request, int poll_count) {
+            if (poll_count == 1) {
+                return TransferStatus{TransferStatusEnum::PENDING, 0};
+            }
+            return TransferStatus{TransferStatusEnum::COMPLETED,
+                                  request.length};
+        });
+    installFakeRdma(engine, fake_rdma);
+
+    constexpr size_t kReqLen = 4096;
+    std::vector<uint8_t> buffer(kReqLen * 2, 0x88);
+    ASSERT_TRUE(engine.registerLocalMemory(buffer.data(), buffer.size()).ok());
+
+    BatchID batch = engine.allocateBatch(2);
+    ASSERT_NE(batch, (BatchID)0);
+
+    ASSERT_TRUE(
+        engine
+            .submitTransfer(batch,
+                            {makeLocalWrite(buffer.data(), kReqLen),
+                             makeLocalWrite(buffer.data() + kReqLen, kReqLen)})
+            .ok());
+    EXPECT_EQ(fake_rdma->submit_calls.load(), 1);
+
+    const auto deadline =
+        std::chrono::steady_clock::now() + std::chrono::milliseconds(1000);
+    while (std::chrono::steady_clock::now() < deadline &&
+           (fake_rdma->submit_calls.load() < 2 ||
+            fake_rdma->status_calls.load() < 4)) {
+        std::this_thread::sleep_for(std::chrono::milliseconds(2));
+    }
+    EXPECT_EQ(fake_rdma->submit_calls.load(), 2);
+    EXPECT_GE(fake_rdma->status_calls.load(), 4);
+
+    TransferStatus status{};
+    ASSERT_TRUE(engine.getTransferStatus(batch, status).ok());
+    EXPECT_EQ(status.s, TransferStatusEnum::COMPLETED);
+
+    EXPECT_TRUE(engine.freeBatch(batch).ok());
+    EXPECT_TRUE(
+        engine.unregisterLocalMemory(buffer.data(), buffer.size()).ok());
+}
+
+TEST(RuntimeQueueDispatch, EarlyFreeReclaimsAfterQueuedCompletion) {
+    auto cfg = makeRuntimeQueueConfig(1, 1UL << 20);
+    TransferEngineImpl engine(cfg);
+    ASSERT_TRUE(engine.available());
+
+    auto fake_rdma = std::make_shared<FakeTransport>(
+        RDMA, [](const Request& request, int poll_count) {
+            if (poll_count == 1) {
+                return TransferStatus{TransferStatusEnum::PENDING, 0};
+            }
+            return TransferStatus{TransferStatusEnum::COMPLETED,
+                                  request.length};
+        });
+    installFakeRdma(engine, fake_rdma);
+
+    constexpr size_t kReqLen = 4096;
+    std::vector<uint8_t> buffer(kReqLen, 0x33);
+    ASSERT_TRUE(engine.registerLocalMemory(buffer.data(), buffer.size()).ok());
+
+    BatchID batch = engine.allocateBatch(1);
+    ASSERT_NE(batch, (BatchID)0);
+    ASSERT_TRUE(
+        engine.submitTransfer(batch, {makeLocalWrite(buffer.data(), kReqLen)})
+            .ok());
+
+    ASSERT_TRUE(engine.freeBatch(batch).ok());
+
+    const auto deadline =
+        std::chrono::steady_clock::now() + std::chrono::milliseconds(1000);
+    while (std::chrono::steady_clock::now() < deadline &&
+           fake_rdma->status_calls.load() < 2) {
+        std::this_thread::sleep_for(std::chrono::milliseconds(2));
+    }
+    EXPECT_GE(fake_rdma->status_calls.load(), 2);
+
+    TransferStatus status{};
+    EXPECT_TRUE(engine.getTransferStatus(batch, 0, status).IsInvalidArgument());
+    EXPECT_TRUE(
+        engine.unregisterLocalMemory(buffer.data(), buffer.size()).ok());
+}
+
+TEST(RuntimeQueueDispatch, PollingDerivedTaskCompletesMergedOwner) {
+    auto cfg = makeRuntimeQueueConfig(1, 1UL << 20, true);
+    TransferEngineImpl engine(cfg);
+    ASSERT_TRUE(engine.available());
+
+    auto fake_rdma = std::make_shared<FakeTransport>(RDMA);
+    installFakeRdma(engine, fake_rdma);
+
+    constexpr size_t kReqLen = 4096;
+    std::vector<uint8_t> buffer(kReqLen * 2, 0x44);
+    ASSERT_TRUE(engine.registerLocalMemory(buffer.data(), buffer.size()).ok());
+
+    BatchID batch = engine.allocateBatch(2);
+    ASSERT_NE(batch, (BatchID)0);
+
+    ASSERT_TRUE(
+        engine
+            .submitTransfer(batch,
+                            {makeLocalWrite(buffer.data(), kReqLen),
+                             makeLocalWrite(buffer.data() + kReqLen, kReqLen)})
+            .ok());
+    EXPECT_EQ(fake_rdma->submit_calls.load(), 1);
+
+    TransferStatus derived{};
+    ASSERT_TRUE(engine.getTransferStatus(batch, 1, derived).ok());
+    EXPECT_EQ(derived.s, TransferStatusEnum::COMPLETED);
+
+    TransferStatus owner{};
+    ASSERT_TRUE(engine.getTransferStatus(batch, 0, owner).ok());
+    EXPECT_EQ(owner.s, TransferStatusEnum::COMPLETED);
+
+    EXPECT_TRUE(engine.freeBatch(batch).ok());
+    EXPECT_TRUE(
+        engine.unregisterLocalMemory(buffer.data(), buffer.size()).ok());
+}
+
+}  // namespace
+}  // namespace tent
+}  // namespace mooncake
diff --git a/mooncake-transfer-engine/tent/tests/sunrise_link_transport_test.cpp b/mooncake-transfer-engine/tent/tests/sunrise_link_transport_test.cpp
new file mode 100644
index 00000000..a49993c9
--- /dev/null
+++ b/mooncake-transfer-engine/tent/tests/sunrise_link_transport_test.cpp
@@ -0,0 +1,209 @@
+#include <gtest/gtest.h>
+
+#include <chrono>
+#include <cstdint>
+#include <memory>
+#include <string>
+#include <sys/wait.h>
+#include <thread>
+#include <unistd.h>
+#include <vector>
+
+#include "gpu_vendor/sunrise.h"
+#include "tent/common/config.h"
+#include "tent/common/types.h"
+#include "tent/transfer_engine.h"
+
+namespace mooncake {
+namespace tent {
+namespace {
+
+void CheckTangError(tangError_t result, const char* message) {
+    ASSERT_EQ(result, tangSuccess)
+        << message << " (error=" << result
+        << ", detail=" << tangGetErrorString(result) << ")";
+}
+
+void WaitBatchDone(TransferEngine* engine, BatchID batch_id) {
+    TransferStatus status;
+    for (int i = 0; i < 5000; ++i) {
+        auto s = engine->getTransferStatus(batch_id, status);
+        ASSERT_TRUE(s.ok()) << "getTransferStatus failed: " << s.ToString();
+        if (status.s == TransferStatusEnum::COMPLETED) return;
+        ASSERT_NE(status.s, TransferStatusEnum::FAILED);
+        std::this_thread::sleep_for(std::chrono::milliseconds(1));
+    }
+    FAIL() << "timeout waiting batch completion";
+}
+
+void RunWriteReadCase(int server_gpu, int client_gpu) {
+    constexpr size_t kDataLength = 4 * 1024 * 1024;
+    const std::string server_local_name =
+        "sunrise_ut_server_" + std::to_string(getpid());
+
+    int ready_pipe[2];
+    int stop_pipe[2];
+    ASSERT_EQ(pipe(ready_pipe), 0);
+    ASSERT_EQ(pipe(stop_pipe), 0);
+
+    pid_t child = fork();
+    ASSERT_GE(child, 0);
+    if (child == 0) {
+        close(ready_pipe[0]);
+        close(stop_pipe[1]);
+        auto server_conf = std::make_shared<Config>();
+        server_conf->set("metadata_type", "p2p");
+        server_conf->set("metadata_servers", "P2PHANDSHAKE");
+        server_conf->set("local_segment_name", server_local_name);
+        server_conf->set("transports/sunrise_link/enable", true);
+        auto server_engine = std::make_unique<TransferEngine>(server_conf);
+        if (!server_engine->available()) _exit(2);
+        void* server_buf = nullptr;
+        if (tangSetDevice(server_gpu) != tangSuccess) _exit(3);
+        if (tangMalloc(&server_buf, kDataLength * 2) != tangSuccess) _exit(4);
+        MemoryOptions server_opts;
+        server_opts.location = "cuda:" + std::to_string(server_gpu);
+        auto s = server_engine->registerLocalMemory(server_buf, kDataLength * 2,
+                                                    server_opts);
+        if (!s.ok()) _exit(5);
+        const std::string seg = server_engine->getSegmentName();
+        uint32_t len = static_cast<uint32_t>(seg.size());
+        if (write(ready_pipe[1], &len, sizeof(len)) != sizeof(len)) _exit(6);
+        if (write(ready_pipe[1], seg.data(), len) != static_cast<ssize_t>(len))
+            _exit(7);
+        char stop = 0;
+        (void)read(stop_pipe[0], &stop, 1);
+        (void)server_engine->unregisterLocalMemory(server_buf);
+        tangSetDevice(server_gpu);
+        tangFree(server_buf);
+        _exit(0);
+    }
+
+    close(ready_pipe[1]);
+    close(stop_pipe[0]);
+    uint32_t seg_len = 0;
+    ssize_t got = read(ready_pipe[0], &seg_len, sizeof(seg_len));
+    if (got != static_cast<ssize_t>(sizeof(seg_len))) {
+        int wstatus = 0;
+        (void)waitpid(child, &wstatus, 0);
+        if (got == 0 && WIFEXITED(wstatus)) {
+            GTEST_SKIP() << "server init failed for case " << server_gpu << "->"
+                         << client_gpu
+                         << ", child exit=" << WEXITSTATUS(wstatus);
+        }
+        ASSERT_EQ(got, static_cast<ssize_t>(sizeof(seg_len)));
+    }
+    ASSERT_GT(seg_len, 0u);
+    std::string server_segment(seg_len, '\0');
+    ASSERT_EQ(read(ready_pipe[0], server_segment.data(), seg_len),
+              static_cast<ssize_t>(seg_len));
+
+    int gpu_count = 0;
+    CheckTangError(tangGetDeviceCount(&gpu_count), "tangGetDeviceCount");
+    if (gpu_count <= 0) {
+        const char stop = 'Q';
+        (void)write(stop_pipe[1], &stop, 1);
+        int wstatus = 0;
+        (void)waitpid(child, &wstatus, 0);
+        GTEST_SKIP() << "no sunrise GPU detected";
+    }
+    if (server_gpu < 0 || client_gpu < 0 || server_gpu >= gpu_count ||
+        client_gpu >= gpu_count) {
+        const char stop = 'Q';
+        (void)write(stop_pipe[1], &stop, 1);
+        int wstatus = 0;
+        (void)waitpid(child, &wstatus, 0);
+        GTEST_SKIP() << "insufficient GPUs for case " << server_gpu << "->"
+                     << client_gpu << ", detected: " << gpu_count;
+    }
+
+    auto client_conf = std::make_shared<Config>();
+    client_conf->set("metadata_type", "p2p");
+    client_conf->set("metadata_servers", "P2PHANDSHAKE");
+    client_conf->set("local_segment_name", "sunrise_ut_client");
+    client_conf->set("transports/sunrise_link/enable", true);
+    auto client_engine = std::make_unique<TransferEngine>(client_conf);
+    ASSERT_TRUE(client_engine->available());
+
+    void* client_buf = nullptr;
+    CheckTangError(tangSetDevice(client_gpu), "tangSetDevice(client)");
+    CheckTangError(tangMalloc(&client_buf, kDataLength * 2),
+                   "tangMalloc(client)");
+    MemoryOptions client_opts;
+    client_opts.location = "cuda:" + std::to_string(client_gpu);
+    auto s = client_engine->registerLocalMemory(client_buf, kDataLength * 2,
+                                                client_opts);
+    ASSERT_TRUE(s.ok()) << s.ToString();
+
+    SegmentID segment_id = 0;
+    for (int i = 0; i < 100; ++i) {
+        s = client_engine->openSegment(segment_id, server_segment);
+        if (s.ok()) break;
+        std::this_thread::sleep_for(std::chrono::milliseconds(20));
+    }
+    ASSERT_TRUE(s.ok()) << s.ToString();
+
+    SegmentInfo segment_info;
+    s = client_engine->getSegmentInfo(segment_id, segment_info);
+    ASSERT_TRUE(s.ok()) << s.ToString();
+    ASSERT_FALSE(segment_info.buffers.empty());
+    const uint64_t remote_base = segment_info.buffers[0].base;
+
+    std::vector<char> host_data(kDataLength, 'S');
+    CheckTangError(tangSetDevice(client_gpu), "tangSetDevice(client copy)");
+    CheckTangError(tangMemcpy(client_buf, host_data.data(), kDataLength,
+                              tangMemcpyHostToDevice),
+                   "tangMemcpy H2D");
+
+    BatchID batch = client_engine->allocateBatch(1);
+    Request req;
+    req.opcode = Request::WRITE;
+    req.length = kDataLength;
+    req.source = static_cast<uint8_t*>(client_buf);
+    req.target_id = segment_id;
+    req.target_offset = remote_base;
+    s = client_engine->submitTransfer(batch, {req});
+    ASSERT_TRUE(s.ok()) << s.ToString();
+    WaitBatchDone(client_engine.get(), batch);
+    s = client_engine->freeBatch(batch);
+    ASSERT_TRUE(s.ok()) << s.ToString();
+
+    batch = client_engine->allocateBatch(1);
+    req.opcode = Request::READ;
+    req.source = static_cast<uint8_t*>(client_buf) + kDataLength;
+    s = client_engine->submitTransfer(batch, {req});
+    ASSERT_TRUE(s.ok()) << s.ToString();
+    WaitBatchDone(client_engine.get(), batch);
+    s = client_engine->freeBatch(batch);
+    ASSERT_TRUE(s.ok()) << s.ToString();
+
+    std::vector<char> host_check(kDataLength);
+    CheckTangError(tangSetDevice(client_gpu), "tangSetDevice(client verify)");
+    CheckTangError(tangMemcpy(host_check.data(),
+                              static_cast<uint8_t*>(client_buf) + kDataLength,
+                              kDataLength, tangMemcpyDeviceToHost),
+                   "tangMemcpy D2H");
+    for (size_t i = 0; i < kDataLength; ++i) ASSERT_EQ(host_check[i], 'S');
+
+    s = client_engine->unregisterLocalMemory(client_buf);
+    ASSERT_TRUE(s.ok()) << s.ToString();
+    CheckTangError(tangSetDevice(client_gpu), "tangSetDevice(client free)");
+    CheckTangError(tangFree(client_buf), "tangFree(client)");
+
+    const char stop = 'Q';
+    ASSERT_EQ(write(stop_pipe[1], &stop, 1), 1);
+    int wstatus = 0;
+    ASSERT_EQ(waitpid(child, &wstatus, 0), child);
+    ASSERT_TRUE(WIFEXITED(wstatus));
+    ASSERT_EQ(WEXITSTATUS(wstatus), 0);
+}
+
+TEST(SunriseLinkTransportTest, WriteAndRead_0_to_1) { RunWriteReadCase(0, 1); }
+
+TEST(SunriseLinkTransportTest, WriteAndRead_1_to_0) { RunWriteReadCase(1, 0); }
+
+TEST(SunriseLinkTransportTest, WriteAndRead_2_to_4) { RunWriteReadCase(2, 4); }
+
+}  // namespace
+}  // namespace tent
+}  // namespace mooncake
diff --git a/mooncake-transfer-engine/tent/tests/transfer_engine_config_override_test.cpp b/mooncake-transfer-engine/tent/tests/transfer_engine_config_override_test.cpp
index 706f7a53..8704924b 100644
--- a/mooncake-transfer-engine/tent/tests/transfer_engine_config_override_test.cpp
+++ b/mooncake-transfer-engine/tent/tests/transfer_engine_config_override_test.cpp
@@ -260,6 +260,29 @@ class TestHttpMetadataServer {
 };
 #endif
 
+TEST(TransferEngineConfigOverrideTest,
+     LegacyRdmaBindAddressEnvLoadsIntoTentConfig) {
+    EnvVarGuard guard("MC_RDMA_BIND_ADDRESS", "10.0.0.2");
+
+    Config config;
+    ASSERT_TRUE(ConfigHelper().loadFromEnv(config).ok());
+
+    EXPECT_EQ(config.get("transports/rdma/bind_address", ""), "10.0.0.2");
+}
+
+TEST(TransferEngineConfigOverrideTest,
+     LegacyRdmaBindAddressEnvOverridesMcTentConf) {
+    EnvVarGuard conf_guard(
+        "MC_TENT_CONF",
+        R"({"transports":{"rdma":{"bind_address":"10.0.0.1"}}})");
+    EnvVarGuard bind_guard("MC_RDMA_BIND_ADDRESS", "10.0.0.2");
+
+    Config config;
+    ASSERT_TRUE(ConfigHelper().loadFromEnv(config).ok());
+
+    EXPECT_EQ(config.get("transports/rdma/bind_address", ""), "10.0.0.2");
+}
+
 TEST(TransferEngineConfigOverrideTest,
      ExplicitMetadataOverridesDriveSuccessfulHttpInitialization) {
 #ifdef _WIN32
diff --git a/mooncake-transfer-engine/tent/tests/transport_hint_test.cpp b/mooncake-transfer-engine/tent/tests/transport_hint_test.cpp
new file mode 100644
index 00000000..84cfdc5b
--- /dev/null
+++ b/mooncake-transfer-engine/tent/tests/transport_hint_test.cpp
@@ -0,0 +1,351 @@
+// Copyright 2026 KVCache.AI
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+//
+// End-to-end tests for Request::transport_hint added to
+// TransferEngineImpl::submitTransfer:
+//   * hint == UNSPEC keeps existing policy-driven behavior;
+//   * hint pinned to RDMA / TCP routes the request onto that transport;
+//   * hint to a disabled / capability-mismatched transport returns
+//     InvalidArgument and does not mutate batch state;
+//   * out-of-range hint is rejected up front.
+
+#include <gtest/gtest.h>
+
+#include <atomic>
+#include <cstdlib>
+#include <cstring>
+#include <memory>
+#include <string>
+#include <vector>
+
+#include "tent/common/config.h"
+#include "tent/common/types.h"
+#include "tent/runtime/segment.h"
+#include "tent/runtime/transfer_engine_impl.h"
+#include "tent/runtime/transport.h"
+
+namespace mooncake {
+namespace tent {
+namespace {
+
+// ---------------------------------------------------------------------------
+// FakeTransport: minimal Transport that records itself on BufferDesc, claims
+// dram_to_dram capability, and completes transfers synchronously. The
+// per-instance submit_calls counter is what each test inspects to learn
+// which transport the engine routed onto for a given request.
+// ---------------------------------------------------------------------------
+class FakeSubBatch : public Transport::SubBatch {
+   public:
+    size_t size() const override { return task_count; }
+    size_t task_count = 0;
+    std::vector<TransferStatus> statuses;
+};
+
+class FakeTransport : public Transport {
+   public:
+    explicit FakeTransport(TransportType self_type) : self_type_(self_type) {
+        caps.dram_to_dram = true;
+    }
+
+    std::atomic<int> submit_calls{0};
+
+    Status install(std::string&, std::shared_ptr<ControlService>,
+                   std::shared_ptr<Topology>,
+                   std::shared_ptr<Config> = nullptr) override {
+        return Status::OK();
+    }
+
+    Status allocateSubBatch(SubBatchRef& batch, size_t) override {
+        batch = new FakeSubBatch();
+        return Status::OK();
+    }
+    Status freeSubBatch(SubBatchRef& batch) override {
+        delete static_cast<FakeSubBatch*>(batch);
+        batch = nullptr;
+        return Status::OK();
+    }
+
+    Status submitTransferTasks(SubBatchRef batch,
+                               const std::vector<Request>& reqs) override {
+        ++submit_calls;
+        auto* fb = static_cast<FakeSubBatch*>(batch);
+        for (const auto& req : reqs) {
+            fb->statuses.push_back({TransferStatusEnum::COMPLETED, req.length});
+            fb->task_count++;
+        }
+        return Status::OK();
+    }
+
+    Status getTransferStatus(SubBatchRef batch, int task_id,
+                             TransferStatus& status) override {
+        auto* fb = static_cast<FakeSubBatch*>(batch);
+        if (task_id < 0 || task_id >= (int)fb->statuses.size())
+            return Status::InvalidArgument("bad task_id" LOC_MARK);
+        status = fb->statuses[task_id];
+        return Status::OK();
+    }
+
+    Status addMemoryBuffer(BufferDesc& desc, const MemoryOptions&) override {
+        desc.transports.push_back(self_type_);
+        return Status::OK();
+    }
+    Status addMemoryBuffer(std::vector<BufferDesc>& list,
+                           const MemoryOptions& options) override {
+        for (auto& d : list) {
+            auto s = addMemoryBuffer(d, options);
+            if (!s.ok()) return s;
+        }
+        return Status::OK();
+    }
+    Status removeMemoryBuffer(BufferDesc&) override { return Status::OK(); }
+
+    Status allocateLocalMemory(void** addr, size_t size,
+                               MemoryOptions&) override {
+        *addr = std::malloc(size);
+        return *addr ? Status::OK()
+                     : Status::InternalError("malloc failed" LOC_MARK);
+    }
+    Status freeLocalMemory(void* addr, size_t) override {
+        std::free(addr);
+        return Status::OK();
+    }
+    bool warmupMemory(void*, size_t) override { return false; }
+    const char* getName() const override {
+        return self_type_ == RDMA ? "<fake-rdma>" : "<fake-tcp>";
+    }
+
+   private:
+    TransportType self_type_;
+};
+
+// ---------------------------------------------------------------------------
+// Helpers
+// ---------------------------------------------------------------------------
+
+std::shared_ptr<Config> makeMinimalP2PConfig() {
+    auto cfg = std::make_shared<Config>();
+    cfg->set("metadata_type", "p2p");
+    cfg->set("metadata_servers", "");
+    cfg->set("rpc_server_hostname", "127.0.0.1");
+    cfg->set("rpc_server_port", "0");
+    cfg->set("log_level", "warning");
+    cfg->set("merge_requests", false);
+    cfg->set("transports/tcp/enable", false);
+    cfg->set("transports/shm/enable", false);
+    cfg->set("transports/rdma/enable", false);
+    cfg->set("transports/io_uring/enable", false);
+    cfg->set("transports/nvlink/enable", false);
+    cfg->set("transports/mnnvl/enable", false);
+    cfg->set("transports/gds/enable", false);
+    cfg->set("transports/ascend_direct/enable", false);
+    return cfg;
+}
+
+void installFakeRdmaAndTcp(TransferEngineImpl& engine,
+                           std::shared_ptr<FakeTransport>& rdma,
+                           std::shared_ptr<FakeTransport>& tcp) {
+    rdma = std::make_shared<FakeTransport>(RDMA);
+    tcp = std::make_shared<FakeTransport>(TCP);
+    std::string seg_name = engine.getSegmentName();
+    ASSERT_TRUE(rdma->install(seg_name, nullptr, nullptr).ok());
+    ASSERT_TRUE(tcp->install(seg_name, nullptr, nullptr).ok());
+    engine.swapTransportForTest(RDMA, rdma);
+    engine.swapTransportForTest(TCP, tcp);
+}
+
+Request makeLocalWriteRequest(uint8_t* buf, size_t len) {
+    Request r;
+    r.opcode = Request::WRITE;
+    r.source = buf;
+    r.target_id = LOCAL_SEGMENT_ID;
+    r.target_offset = reinterpret_cast<uint64_t>(buf);
+    r.length = len;
+    return r;
+}
+
+// ---------------------------------------------------------------------------
+// 1. hint=RDMA / hint=TCP routing
+// ---------------------------------------------------------------------------
+
+TEST(TransportHint, RoutesPinnedRequestsToHintedTransport) {
+    auto cfg = makeMinimalP2PConfig();
+    TransferEngineImpl engine(cfg);
+    ASSERT_TRUE(engine.available());
+
+    std::shared_ptr<FakeTransport> fake_rdma, fake_tcp;
+    installFakeRdmaAndTcp(engine, fake_rdma, fake_tcp);
+
+    constexpr size_t kBufLen = 4096;
+    std::vector<uint8_t> buf(kBufLen, 0x01);
+    ASSERT_TRUE(engine.registerLocalMemory(buf.data(), kBufLen).ok());
+
+    // Two batches, each with one hinted request. Use separate batches so
+    // submit_calls is unambiguous about which transport handled which.
+    BatchID b_rdma = engine.allocateBatch(2);
+    BatchID b_tcp = engine.allocateBatch(2);
+    ASSERT_NE(b_rdma, (BatchID)0);
+    ASSERT_NE(b_tcp, (BatchID)0);
+
+    Request r_rdma = makeLocalWriteRequest(buf.data(), kBufLen);
+    r_rdma.transport_hint = RDMA;
+    Request r_tcp = makeLocalWriteRequest(buf.data(), kBufLen);
+    r_tcp.transport_hint = TCP;
+
+    ASSERT_TRUE(engine.submitTransfer(b_rdma, {r_rdma}).ok());
+    ASSERT_TRUE(engine.submitTransfer(b_tcp, {r_tcp}).ok());
+
+    EXPECT_EQ(fake_rdma->submit_calls.load(), 1)
+        << "RDMA hint should route onto the RDMA fake exactly once";
+    EXPECT_EQ(fake_tcp->submit_calls.load(), 1)
+        << "TCP hint should route onto the TCP fake exactly once";
+
+    EXPECT_TRUE(engine.freeBatch(b_rdma).ok());
+    EXPECT_TRUE(engine.freeBatch(b_tcp).ok());
+    EXPECT_TRUE(engine.unregisterLocalMemory(buf.data(), kBufLen).ok());
+}
+
+// ---------------------------------------------------------------------------
+// 2. Disabled transport rejected before batch state mutates
+// ---------------------------------------------------------------------------
+
+TEST(TransportHint, RejectsHintForTransportNotInstalled) {
+    auto cfg = makeMinimalP2PConfig();
+    TransferEngineImpl engine(cfg);
+    ASSERT_TRUE(engine.available());
+
+    std::shared_ptr<FakeTransport> fake_rdma, fake_tcp;
+    installFakeRdmaAndTcp(engine, fake_rdma, fake_tcp);
+    // Intentionally leave GDS uninstalled.
+
+    constexpr size_t kBufLen = 1024;
+    std::vector<uint8_t> buf(kBufLen, 0x77);
+    ASSERT_TRUE(engine.registerLocalMemory(buf.data(), kBufLen).ok());
+
+    BatchID batch_id = engine.allocateBatch(2);
+    Request bad = makeLocalWriteRequest(buf.data(), kBufLen);
+    bad.transport_hint = GDS;
+
+    auto status = engine.submitTransfer(batch_id, {bad});
+    EXPECT_FALSE(status.ok());
+    EXPECT_TRUE(status.IsInvalidArgument());
+    EXPECT_EQ(fake_rdma->submit_calls.load(), 0);
+    EXPECT_EQ(fake_tcp->submit_calls.load(), 0);
+
+    EXPECT_TRUE(engine.freeBatch(batch_id).ok());
+    EXPECT_TRUE(engine.unregisterLocalMemory(buf.data(), kBufLen).ok());
+}
+
+// ---------------------------------------------------------------------------
+// 3. Out-of-range hint value rejected (defends against uninitialized C
+//    struct fields and stray casts)
+// ---------------------------------------------------------------------------
+
+TEST(TransportHint, RejectsHintWithOutOfRangeValue) {
+    auto cfg = makeMinimalP2PConfig();
+    TransferEngineImpl engine(cfg);
+    ASSERT_TRUE(engine.available());
+
+    std::shared_ptr<FakeTransport> fake_rdma, fake_tcp;
+    installFakeRdmaAndTcp(engine, fake_rdma, fake_tcp);
+
+    constexpr size_t kBufLen = 1024;
+    std::vector<uint8_t> buf(kBufLen);
+    ASSERT_TRUE(engine.registerLocalMemory(buf.data(), kBufLen).ok());
+
+    BatchID batch_id = engine.allocateBatch(2);
+
+    Request too_big = makeLocalWriteRequest(buf.data(), kBufLen);
+    too_big.transport_hint = static_cast<TransportType>(123);
+    EXPECT_TRUE(engine.submitTransfer(batch_id, {too_big}).IsInvalidArgument());
+
+    Request negative = makeLocalWriteRequest(buf.data(), kBufLen);
+    negative.transport_hint = static_cast<TransportType>(-1);
+    EXPECT_TRUE(
+        engine.submitTransfer(batch_id, {negative}).IsInvalidArgument());
+
+    // Neither submission should have reached any transport.
+    EXPECT_EQ(fake_rdma->submit_calls.load(), 0);
+    EXPECT_EQ(fake_tcp->submit_calls.load(), 0);
+
+    EXPECT_TRUE(engine.freeBatch(batch_id).ok());
+    EXPECT_TRUE(engine.unregisterLocalMemory(buf.data(), kBufLen).ok());
+}
+
+// ---------------------------------------------------------------------------
+// 4. UNSPEC keeps original policy-driven behavior intact
+// ---------------------------------------------------------------------------
+
+TEST(TransportHint, UnspecHintPreservesPolicyDrivenSelection) {
+    auto cfg = makeMinimalP2PConfig();
+    TransferEngineImpl engine(cfg);
+    ASSERT_TRUE(engine.available());
+
+    std::shared_ptr<FakeTransport> fake_rdma, fake_tcp;
+    installFakeRdmaAndTcp(engine, fake_rdma, fake_tcp);
+
+    constexpr size_t kBufLen = 4096;
+    std::vector<uint8_t> buf(kBufLen, 0xAA);
+    ASSERT_TRUE(engine.registerLocalMemory(buf.data(), kBufLen).ok());
+
+    BatchID batch_id = engine.allocateBatch(2);
+    Request r = makeLocalWriteRequest(buf.data(), kBufLen);
+    // r.transport_hint defaults to UNSPEC.
+
+    ASSERT_TRUE(engine.submitTransfer(batch_id, {r}).ok());
+
+    // The default policy prefers RDMA over TCP for memory segments, so
+    // an UNSPEC request should land on RDMA. This guards the early-out
+    // path that the change explicitly preserves.
+    EXPECT_EQ(fake_rdma->submit_calls.load(), 1);
+    EXPECT_EQ(fake_tcp->submit_calls.load(), 0);
+
+    EXPECT_TRUE(engine.freeBatch(batch_id).ok());
+    EXPECT_TRUE(engine.unregisterLocalMemory(buf.data(), kBufLen).ok());
+}
+
+// ---------------------------------------------------------------------------
+// 5. Mixed hints in one batch: each request lands on its hinted transport
+// ---------------------------------------------------------------------------
+
+TEST(TransportHint, MixedHintsInBatchSplitAcrossTransports) {
+    auto cfg = makeMinimalP2PConfig();
+    TransferEngineImpl engine(cfg);
+    ASSERT_TRUE(engine.available());
+
+    std::shared_ptr<FakeTransport> fake_rdma, fake_tcp;
+    installFakeRdmaAndTcp(engine, fake_rdma, fake_tcp);
+
+    constexpr size_t kBufLen = 8192;
+    std::vector<uint8_t> buf(kBufLen, 0xBB);
+    ASSERT_TRUE(engine.registerLocalMemory(buf.data(), kBufLen).ok());
+
+    BatchID batch_id = engine.allocateBatch(4);
+
+    Request r_rdma = makeLocalWriteRequest(buf.data(), 4096);
+    r_rdma.transport_hint = RDMA;
+    Request r_tcp = makeLocalWriteRequest(buf.data() + 4096, 4096);
+    r_tcp.transport_hint = TCP;
+
+    ASSERT_TRUE(engine.submitTransfer(batch_id, {r_rdma, r_tcp}).ok());
+
+    EXPECT_EQ(fake_rdma->submit_calls.load(), 1);
+    EXPECT_EQ(fake_tcp->submit_calls.load(), 1);
+
+    EXPECT_TRUE(engine.freeBatch(batch_id).ok());
+    EXPECT_TRUE(engine.unregisterLocalMemory(buf.data(), kBufLen).ok());
+}
+
+}  // namespace
+}  // namespace tent
+}  // namespace mooncake
diff --git a/mooncake-transfer-engine/tent/tests/transport_selector_test.cpp b/mooncake-transfer-engine/tent/tests/transport_selector_test.cpp
new file mode 100644
index 00000000..a479c8e4
--- /dev/null
+++ b/mooncake-transfer-engine/tent/tests/transport_selector_test.cpp
@@ -0,0 +1,646 @@
+// Copyright 2026 KVCache.AI
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include <gtest/gtest.h>
+
+#include <memory>
+#include <string>
+
+#include "tent/common/config.h"
+#include "tent/common/types.h"
+#include "tent/runtime/transport_selector.h"
+#include "tent/runtime/transport.h"
+
+namespace mooncake {
+namespace tent {
+namespace {
+
+// ---------------------------------------------------------------------------
+// Minimal FakeTransport for availability checking
+// ---------------------------------------------------------------------------
+
+class FakeTransport : public Transport {
+   public:
+    explicit FakeTransport(TransportType type) : type_(type) {
+        // Access protected member directly (inherited from Transport)
+    }
+
+    Status install(std::string&, std::shared_ptr<ControlService>,
+                   std::shared_ptr<Topology>,
+                   std::shared_ptr<Config> = nullptr) override {
+        return Status::OK();
+    }
+
+    Status allocateSubBatch(SubBatchRef&, size_t) override {
+        return Status::OK();
+    }
+
+    Status freeSubBatch(SubBatchRef&) override { return Status::OK(); }
+
+    Status submitTransferTasks(SubBatchRef,
+                               const std::vector<Request>&) override {
+        return Status::OK();
+    }
+
+    Status getTransferStatus(SubBatchRef, int, TransferStatus&) override {
+        return Status::OK();
+    }
+
+    const char* getName() const override { return "fake"; }
+
+    // Helper to set capabilities (accessing protected member)
+    void setDramToFile(bool val) { caps.dram_to_file = val; }
+    void setGpuToFile(bool val) { caps.gpu_to_file = val; }
+    void setDramToDram(bool val) { caps.dram_to_dram = val; }
+    void setGpuToGpu(bool val) { caps.gpu_to_gpu = val; }
+    void setDramToGpu(bool val) { caps.dram_to_gpu = val; }
+    void setGpuToDram(bool val) { caps.gpu_to_dram = val; }
+
+   private:
+    TransportType type_;
+};
+
+// ---------------------------------------------------------------------------
+// Test default policies match original behavior
+// ---------------------------------------------------------------------------
+
+TEST(TransportSelectorTest, DefaultPoliciesFileSegment) {
+    auto conf = std::make_shared<Config>();
+    TransportSelector selector(conf);
+
+    // Create fake transports
+    std::array<std::shared_ptr<Transport>, kSupportedTransportTypes>
+        transports{};
+    transports[GDS] = std::make_shared<FakeTransport>(GDS);
+    transports[IOURING] = std::make_shared<FakeTransport>(IOURING);
+    transports[RDMA] = std::make_shared<FakeTransport>(RDMA);
+
+    // Set capabilities
+    auto* gds = static_cast<FakeTransport*>(transports[GDS].get());
+    auto* iouring = static_cast<FakeTransport*>(transports[IOURING].get());
+    gds->setDramToFile(true);
+    iouring->setDramToFile(true);
+
+    // File segment with CPU memory should prefer GDS first
+    SelectionContext ctx;
+    ctx.segment_type = SegmentType::File;
+    ctx.same_machine = true;
+    ctx.local_memory_type = MTYPE_CPU;
+    ctx.remote_memory_type = MTYPE_CPU;
+    ctx.transfer_size = 4096;
+    ctx.priority_level = 0;
+    ctx.buffer_transports = nullptr;
+
+    auto result = selector.select(ctx, transports);
+    EXPECT_EQ(result.transport, GDS) << "File segment should prefer GDS first";
+}
+
+TEST(TransportSelectorTest, DefaultPoliciesFileSegmentGpuMemory) {
+    auto conf = std::make_shared<Config>();
+    TransportSelector selector(conf);
+
+    // Create fake transports
+    std::array<std::shared_ptr<Transport>, kSupportedTransportTypes>
+        transports{};
+    transports[GDS] = std::make_shared<FakeTransport>(GDS);
+    transports[IOURING] = std::make_shared<FakeTransport>(IOURING);
+
+    // Set capabilities for GPU
+    auto* gds = static_cast<FakeTransport*>(transports[GDS].get());
+    auto* iouring = static_cast<FakeTransport*>(transports[IOURING].get());
+    gds->setGpuToFile(true);
+    iouring->setGpuToFile(true);
+
+    // File segment with GPU memory
+    SelectionContext ctx;
+    ctx.segment_type = SegmentType::File;
+    ctx.same_machine = true;
+    ctx.local_memory_type = MTYPE_CUDA;
+    ctx.remote_memory_type = MTYPE_CPU;
+    ctx.transfer_size = 4096;
+    ctx.priority_level = 0;
+    ctx.buffer_transports = nullptr;
+
+    auto result = selector.select(ctx, transports);
+    EXPECT_EQ(result.transport, GDS)
+        << "File segment with GPU should prefer GDS first";
+}
+
+TEST(TransportSelectorTest, DefaultPoliciesMemorySegment) {
+    auto conf = std::make_shared<Config>();
+    TransportSelector selector(conf);
+
+    // Create fake transports
+    std::array<std::shared_ptr<Transport>, kSupportedTransportTypes>
+        transports{};
+    transports[RDMA] = std::make_shared<FakeTransport>(RDMA);
+    transports[TCP] = std::make_shared<FakeTransport>(TCP);
+
+    // Set capabilities
+    auto* rdma = static_cast<FakeTransport*>(transports[RDMA].get());
+    rdma->setDramToDram(true);
+
+    // Memory segment should use buffer_transports order
+    std::vector<TransportType> buffer_transports = {RDMA, TCP};
+
+    SelectionContext ctx;
+    ctx.segment_type = SegmentType::Memory;
+    ctx.same_machine = false;
+    ctx.local_memory_type = MTYPE_CPU;
+    ctx.remote_memory_type = MTYPE_CPU;
+    ctx.transfer_size = 4096;
+    ctx.priority_level = 0;
+    ctx.buffer_transports = &buffer_transports;
+
+    auto result = selector.select(ctx, transports);
+    EXPECT_EQ(result.transport, RDMA)
+        << "Should use first in buffer_transports";
+}
+
+// ---------------------------------------------------------------------------
+// Test transport type name parsing
+// ---------------------------------------------------------------------------
+
+TEST(TransportSelectorTest, TransportTypeNameMapping) {
+    EXPECT_EQ(TransportSelector::transportTypeName(UNSPEC), "unspec");
+    EXPECT_EQ(TransportSelector::transportTypeName(RDMA), "rdma");
+    EXPECT_EQ(TransportSelector::transportTypeName(MNNVL), "mnnvl");
+    EXPECT_EQ(TransportSelector::transportTypeName(SHM), "shm");
+    EXPECT_EQ(TransportSelector::transportTypeName(NVLINK), "nvlink");
+    EXPECT_EQ(TransportSelector::transportTypeName(GDS), "gds");
+    EXPECT_EQ(TransportSelector::transportTypeName(IOURING), "io_uring");
+    EXPECT_EQ(TransportSelector::transportTypeName(TCP), "tcp");
+    EXPECT_EQ(TransportSelector::transportTypeName(AscendDirect), "ascend");
+    EXPECT_EQ(TransportSelector::transportTypeName(SUNRISE_LINK),
+              "sunrise_link");
+}
+
+TEST(TransportSelectorTest, ParseTransportType) {
+    EXPECT_EQ(TransportSelector::parseTransportType("unspec"), UNSPEC);
+    EXPECT_EQ(TransportSelector::parseTransportType("rdma"), RDMA);
+    EXPECT_EQ(TransportSelector::parseTransportType("mnnvl"), MNNVL);
+    EXPECT_EQ(TransportSelector::parseTransportType("shm"), SHM);
+    EXPECT_EQ(TransportSelector::parseTransportType("nvlink"), NVLINK);
+    EXPECT_EQ(TransportSelector::parseTransportType("gds"), GDS);
+    EXPECT_EQ(TransportSelector::parseTransportType("io_uring"), IOURING);
+    EXPECT_EQ(TransportSelector::parseTransportType("tcp"), TCP);
+    EXPECT_EQ(TransportSelector::parseTransportType("ascend"), AscendDirect);
+    EXPECT_EQ(TransportSelector::parseTransportType("sunrise_link"),
+              SUNRISE_LINK);
+    EXPECT_EQ(TransportSelector::parseTransportType("unknown"), UNSPEC);
+}
+
+// ---------------------------------------------------------------------------
+// Test legacy mode
+// ---------------------------------------------------------------------------
+
+TEST(TransportSelectorTest, LegacyModeDefaultOff) {
+    auto conf = std::make_shared<Config>();
+    TransportSelector selector(conf);
+
+    EXPECT_FALSE(selector.isLegacyMode())
+        << "Legacy mode should be off by default";
+}
+
+TEST(TransportSelectorTest, LegacyModeCanBeEnabled) {
+    auto conf = std::make_shared<Config>();
+    TransportSelector selector(conf);
+
+    selector.setLegacyMode(true);
+    EXPECT_TRUE(selector.isLegacyMode());
+
+    selector.setLegacyMode(false);
+    EXPECT_FALSE(selector.isLegacyMode());
+}
+
+// ---------------------------------------------------------------------------
+// Test transport availability based on capabilities
+// ---------------------------------------------------------------------------
+
+TEST(TransportSelectorTest, TransportCapabilityDramToDram) {
+    auto conf = std::make_shared<Config>();
+    TransportSelector selector(conf);
+
+    std::array<std::shared_ptr<Transport>, kSupportedTransportTypes>
+        transports{};
+    transports[RDMA] = std::make_shared<FakeTransport>(RDMA);
+    auto* rdma = static_cast<FakeTransport*>(transports[RDMA].get());
+    rdma->setDramToDram(true);
+
+    std::vector<TransportType> buffer_transports = {RDMA};
+
+    SelectionContext ctx;
+    ctx.segment_type = SegmentType::Memory;
+    ctx.local_memory_type = MTYPE_CPU;
+    ctx.remote_memory_type = MTYPE_CPU;
+    ctx.buffer_transports = &buffer_transports;
+
+    auto result = selector.select(ctx, transports);
+    EXPECT_EQ(result.transport, RDMA)
+        << "RDMA should be available for CPU-to-CPU";
+}
+
+TEST(TransportSelectorTest, TransportCapabilityGpuToGpu) {
+    auto conf = std::make_shared<Config>();
+    TransportSelector selector(conf);
+
+    std::array<std::shared_ptr<Transport>, kSupportedTransportTypes>
+        transports{};
+    transports[RDMA] = std::make_shared<FakeTransport>(RDMA);
+    auto* rdma = static_cast<FakeTransport*>(transports[RDMA].get());
+    rdma->setGpuToGpu(true);
+
+    std::vector<TransportType> buffer_transports = {RDMA};
+
+    SelectionContext ctx;
+    ctx.segment_type = SegmentType::Memory;
+    ctx.local_memory_type = MTYPE_CUDA;
+    ctx.remote_memory_type = MTYPE_CUDA;
+    ctx.buffer_transports = &buffer_transports;
+
+    auto result = selector.select(ctx, transports);
+    EXPECT_EQ(result.transport, RDMA)
+        << "RDMA should be available for CUDA-to-CUDA";
+}
+
+TEST(TransportSelectorTest, TransportCapabilityDramToGpu) {
+    auto conf = std::make_shared<Config>();
+    TransportSelector selector(conf);
+
+    std::array<std::shared_ptr<Transport>, kSupportedTransportTypes>
+        transports{};
+    transports[RDMA] = std::make_shared<FakeTransport>(RDMA);
+    auto* rdma = static_cast<FakeTransport*>(transports[RDMA].get());
+    rdma->setDramToGpu(true);
+
+    std::vector<TransportType> buffer_transports = {RDMA};
+
+    SelectionContext ctx;
+    ctx.segment_type = SegmentType::Memory;
+    ctx.local_memory_type = MTYPE_CPU;
+    ctx.remote_memory_type = MTYPE_CUDA;
+    ctx.buffer_transports = &buffer_transports;
+
+    auto result = selector.select(ctx, transports);
+    EXPECT_EQ(result.transport, RDMA)
+        << "RDMA should be available for CPU-to-CUDA";
+}
+
+TEST(TransportSelectorTest, TransportCapabilityGpuToDram) {
+    auto conf = std::make_shared<Config>();
+    TransportSelector selector(conf);
+
+    std::array<std::shared_ptr<Transport>, kSupportedTransportTypes>
+        transports{};
+    transports[RDMA] = std::make_shared<FakeTransport>(RDMA);
+    auto* rdma = static_cast<FakeTransport*>(transports[RDMA].get());
+    rdma->setGpuToDram(true);
+
+    std::vector<TransportType> buffer_transports = {RDMA};
+
+    SelectionContext ctx;
+    ctx.segment_type = SegmentType::Memory;
+    ctx.local_memory_type = MTYPE_CUDA;
+    ctx.remote_memory_type = MTYPE_CPU;
+    ctx.buffer_transports = &buffer_transports;
+
+    auto result = selector.select(ctx, transports);
+    EXPECT_EQ(result.transport, RDMA)
+        << "RDMA should be available for CUDA-to-CPU";
+}
+
+TEST(TransportSelectorTest, FileSegmentDramToFile) {
+    auto conf = std::make_shared<Config>();
+    TransportSelector selector(conf);
+
+    std::array<std::shared_ptr<Transport>, kSupportedTransportTypes>
+        transports{};
+    transports[GDS] = std::make_shared<FakeTransport>(GDS);
+    auto* gds = static_cast<FakeTransport*>(transports[GDS].get());
+    gds->setDramToFile(true);
+
+    SelectionContext ctx;
+    ctx.segment_type = SegmentType::File;
+    ctx.local_memory_type = MTYPE_CPU;
+    ctx.buffer_transports = nullptr;
+
+    auto result = selector.select(ctx, transports);
+    EXPECT_EQ(result.transport, GDS)
+        << "GDS should be available for File segment with CPU memory";
+}
+
+TEST(TransportSelectorTest, FileSegmentGpuToFile) {
+    auto conf = std::make_shared<Config>();
+    TransportSelector selector(conf);
+
+    std::array<std::shared_ptr<Transport>, kSupportedTransportTypes>
+        transports{};
+    transports[GDS] = std::make_shared<FakeTransport>(GDS);
+    auto* gds = static_cast<FakeTransport*>(transports[GDS].get());
+    gds->setGpuToFile(true);
+
+    SelectionContext ctx;
+    ctx.segment_type = SegmentType::File;
+    ctx.local_memory_type = MTYPE_CUDA;
+    ctx.buffer_transports = nullptr;
+
+    auto result = selector.select(ctx, transports);
+    EXPECT_EQ(result.transport, GDS)
+        << "GDS should be available for File segment with CUDA memory";
+}
+
+// ---------------------------------------------------------------------------
+// Test priority_offset for fallback
+// ---------------------------------------------------------------------------
+
+TEST(TransportSelectorTest, PriorityOffsetFallback) {
+    auto conf = std::make_shared<Config>();
+    TransportSelector selector(conf);
+
+    std::array<std::shared_ptr<Transport>, kSupportedTransportTypes>
+        transports{};
+    transports[RDMA] = std::make_shared<FakeTransport>(RDMA);
+    transports[TCP] = std::make_shared<FakeTransport>(TCP);
+
+    auto* rdma = static_cast<FakeTransport*>(transports[RDMA].get());
+    rdma->setDramToDram(true);
+    auto* tcp = static_cast<FakeTransport*>(transports[TCP].get());
+    tcp->setDramToDram(true);
+
+    std::vector<TransportType> buffer_transports = {RDMA, TCP};
+
+    SelectionContext ctx;
+    ctx.segment_type = SegmentType::Memory;
+    ctx.same_machine = false;
+    ctx.local_memory_type = MTYPE_CPU;
+    ctx.remote_memory_type = MTYPE_CPU;
+    ctx.buffer_transports = &buffer_transports;
+
+    // priority_offset=0: first transport
+    auto result = selector.select(ctx, transports, 0);
+    EXPECT_EQ(result.transport, RDMA);
+
+    // priority_offset=1: second transport (fallback)
+    result = selector.select(ctx, transports, 1);
+    EXPECT_EQ(result.transport, TCP);
+
+    // priority_offset=2: no more transports
+    result = selector.select(ctx, transports, 2);
+    EXPECT_EQ(result.transport, UNSPEC);
+}
+
+// ---------------------------------------------------------------------------
+// Test device mask handling
+// ---------------------------------------------------------------------------
+
+TEST(TransportSelectorTest, DeviceMaskDefaultAll) {
+    auto conf = std::make_shared<Config>();
+    TransportSelector selector(conf);
+
+    SelectionContext ctx;
+    ctx.segment_type = SegmentType::Memory;
+    ctx.buffer_transports = nullptr;
+
+    std::array<std::shared_ptr<Transport>, kSupportedTransportTypes>
+        transports{};
+
+    auto result = selector.select(ctx, transports);
+    EXPECT_EQ(result.device_mask, ~0ULL)
+        << "Default device mask should be all devices";
+}
+
+// ---------------------------------------------------------------------------
+// Test NVLINK and SHM same-machine constraint (via select)
+// ---------------------------------------------------------------------------
+
+TEST(TransportSelectorTest, NvLinkRequiresSameMachine) {
+    auto conf = std::make_shared<Config>();
+    TransportSelector selector(conf);
+
+    std::array<std::shared_ptr<Transport>, kSupportedTransportTypes>
+        transports{};
+    transports[NVLINK] = std::make_shared<FakeTransport>(NVLINK);
+    auto* nvlink = static_cast<FakeTransport*>(transports[NVLINK].get());
+    nvlink->setGpuToGpu(true);
+
+    std::vector<TransportType> buffer_transports = {NVLINK};
+
+    SelectionContext ctx;
+    ctx.segment_type = SegmentType::Memory;
+    ctx.same_machine = false;  // Remote
+    ctx.local_memory_type = MTYPE_CUDA;
+    ctx.remote_memory_type = MTYPE_CUDA;
+    ctx.buffer_transports = &buffer_transports;
+
+    // NVLINK should not be available for remote transfers
+    auto result = selector.select(ctx, transports);
+    EXPECT_EQ(result.transport, UNSPEC)
+        << "NVLINK should not be available for remote transfers";
+}
+
+TEST(TransportSelectorTest, NvLinkAvailableSameMachine) {
+    auto conf = std::make_shared<Config>();
+    TransportSelector selector(conf);
+
+    std::array<std::shared_ptr<Transport>, kSupportedTransportTypes>
+        transports{};
+    transports[NVLINK] = std::make_shared<FakeTransport>(NVLINK);
+    auto* nvlink = static_cast<FakeTransport*>(transports[NVLINK].get());
+    nvlink->setGpuToGpu(true);
+
+    std::vector<TransportType> buffer_transports = {NVLINK};
+
+    SelectionContext ctx;
+    ctx.segment_type = SegmentType::Memory;
+    ctx.same_machine = true;
+    ctx.local_memory_type = MTYPE_CUDA;
+    ctx.remote_memory_type = MTYPE_CUDA;
+    ctx.buffer_transports = &buffer_transports;
+
+    auto result = selector.select(ctx, transports);
+    EXPECT_EQ(result.transport, NVLINK)
+        << "NVLINK should be available for same-machine transfers";
+}
+
+// ---------------------------------------------------------------------------
+// Test ROCm memory type
+// ---------------------------------------------------------------------------
+
+TEST(TransportSelectorTest, RocmMemoryTypeSupported) {
+    auto conf = std::make_shared<Config>();
+    TransportSelector selector(conf);
+
+    std::array<std::shared_ptr<Transport>, kSupportedTransportTypes>
+        transports{};
+    transports[RDMA] = std::make_shared<FakeTransport>(RDMA);
+    auto* rdma = static_cast<FakeTransport*>(transports[RDMA].get());
+    rdma->setGpuToGpu(true);
+
+    std::vector<TransportType> buffer_transports = {RDMA};
+
+    SelectionContext ctx;
+    ctx.segment_type = SegmentType::Memory;
+    ctx.same_machine = true;
+    ctx.local_memory_type = MTYPE_ROCM;
+    ctx.remote_memory_type = MTYPE_ROCM;
+    ctx.buffer_transports = &buffer_transports;
+
+    auto result = selector.select(ctx, transports);
+    EXPECT_EQ(result.transport, RDMA)
+        << "RDMA should be available for ROCm-to-ROCm";
+}
+
+// ---------------------------------------------------------------------------
+// Test with configuration (policy-based selection)
+// ---------------------------------------------------------------------------
+
+TEST(TransportSelectorTest, ConfigBasedPolicySelection) {
+    auto conf = std::make_shared<Config>();
+
+    // Set up a custom policy via JSON config
+    conf->set("policy", json::array());
+    auto policies = conf->getArray<json>("policy");
+
+    json policy;
+    policy["name"] = "test_memory_policy";
+    policy["segment_type"] = "memory";
+    policy["transports"] = {"tcp", "rdma"};  // Prefer TCP over RDMA
+
+    // We can't easily modify the config's internal JSON structure,
+    // so this test verifies the selector at least loads without error
+
+    TransportSelector selector(conf);
+
+    // Default behavior should still work
+    std::array<std::shared_ptr<Transport>, kSupportedTransportTypes>
+        transports{};
+    transports[RDMA] = std::make_shared<FakeTransport>(RDMA);
+    transports[TCP] = std::make_shared<FakeTransport>(TCP);
+
+    auto* rdma = static_cast<FakeTransport*>(transports[RDMA].get());
+    rdma->setDramToDram(true);
+    auto* tcp = static_cast<FakeTransport*>(transports[TCP].get());
+    tcp->setDramToDram(true);
+
+    std::vector<TransportType> buffer_transports = {RDMA, TCP};
+
+    SelectionContext ctx;
+    ctx.segment_type = SegmentType::Memory;
+    ctx.same_machine = false;
+    ctx.local_memory_type = MTYPE_CPU;
+    ctx.remote_memory_type = MTYPE_CPU;
+    ctx.buffer_transports = &buffer_transports;
+
+    auto result = selector.select(ctx, transports);
+    // With default policies, should use buffer_transports order (RDMA first)
+    EXPECT_EQ(result.transport, RDMA);
+}
+
+// ---------------------------------------------------------------------------
+// hint: per-request transport_hint hook on select()
+// ---------------------------------------------------------------------------
+
+TEST(TransportSelectorTest, FailoverFromHintAdvancesByIndex) {
+    // Policy order [TCP, RDMA], hint=RDMA. Hint is prepended; failover by
+    // bumping transport_index walks past the hint into the rest of raw
+    // (with RDMA de-duped), so index=1 returns TCP — never returns RDMA
+    // again.
+    auto conf = std::make_shared<Config>();
+    TransportSelector selector(conf);
+
+    std::array<std::shared_ptr<Transport>, kSupportedTransportTypes>
+        transports{};
+    transports[RDMA] = std::make_shared<FakeTransport>(RDMA);
+    transports[TCP] = std::make_shared<FakeTransport>(TCP);
+    static_cast<FakeTransport*>(transports[RDMA].get())->setDramToDram(true);
+    static_cast<FakeTransport*>(transports[TCP].get())->setDramToDram(true);
+
+    std::vector<TransportType> buffer_transports = {TCP, RDMA};
+    SelectionContext ctx;
+    ctx.segment_type = SegmentType::Memory;
+    ctx.same_machine = false;
+    ctx.local_memory_type = MTYPE_CPU;
+    ctx.remote_memory_type = MTYPE_CPU;
+    ctx.buffer_transports = &buffer_transports;
+
+    EXPECT_EQ(
+        selector.select(ctx, transports, /*index=*/0, /*hint=*/RDMA).transport,
+        RDMA);
+    EXPECT_EQ(
+        selector.select(ctx, transports, /*index=*/1, /*hint=*/RDMA).transport,
+        TCP);
+    // Past the end -> UNSPEC.
+    EXPECT_EQ(
+        selector.select(ctx, transports, /*index=*/2, /*hint=*/RDMA).transport,
+        UNSPEC);
+}
+
+TEST(TransportSelectorTest, HintIsPrependedToCandidateList) {
+    auto conf = std::make_shared<Config>();
+    TransportSelector selector(conf);
+
+    std::array<std::shared_ptr<Transport>, kSupportedTransportTypes>
+        transports{};
+    transports[RDMA] = std::make_shared<FakeTransport>(RDMA);
+    transports[TCP] = std::make_shared<FakeTransport>(TCP);
+    static_cast<FakeTransport*>(transports[RDMA].get())->setDramToDram(true);
+    static_cast<FakeTransport*>(transports[TCP].get())->setDramToDram(true);
+
+    // Default-policy memory path uses buffer_transports order = [RDMA, TCP].
+    // hint=TCP must put TCP first, RDMA second.
+    std::vector<TransportType> buffer_transports = {RDMA, TCP};
+    SelectionContext ctx;
+    ctx.segment_type = SegmentType::Memory;
+    ctx.same_machine = false;
+    ctx.local_memory_type = MTYPE_CPU;
+    ctx.remote_memory_type = MTYPE_CPU;
+    ctx.buffer_transports = &buffer_transports;
+
+    auto r0 = selector.select(ctx, transports, /*index=*/0, /*hint=*/TCP);
+    EXPECT_EQ(r0.transport, TCP);
+    auto r1 = selector.select(ctx, transports, /*index=*/1, /*hint=*/TCP);
+    EXPECT_EQ(r1.transport, RDMA);
+}
+
+TEST(TransportSelectorTest, HintNotInMatchingPolicyReturnsUnspec) {
+    // Selector mode: the matching policy's transports list is the
+    // authorization whitelist. A hint that is not in it produces UNSPEC
+    // (downstream the engine turns this into task = FAILED).
+    auto conf = std::make_shared<Config>();
+    TransportSelector selector(conf);
+
+    std::array<std::shared_ptr<Transport>, kSupportedTransportTypes>
+        transports{};
+    transports[RDMA] = std::make_shared<FakeTransport>(RDMA);
+    transports[TCP] = std::make_shared<FakeTransport>(TCP);
+    static_cast<FakeTransport*>(transports[RDMA].get())->setDramToDram(true);
+    static_cast<FakeTransport*>(transports[TCP].get())->setDramToDram(true);
+
+    // Buffer is only registered with TCP. hint=RDMA must fail even though
+    // RDMA is installed and capable.
+    std::vector<TransportType> buffer_transports = {TCP};
+    SelectionContext ctx;
+    ctx.segment_type = SegmentType::Memory;
+    ctx.same_machine = false;
+    ctx.local_memory_type = MTYPE_CPU;
+    ctx.remote_memory_type = MTYPE_CPU;
+    ctx.buffer_transports = &buffer_transports;
+
+    auto r = selector.select(ctx, transports, /*index=*/0, /*hint=*/RDMA);
+    EXPECT_EQ(r.transport, UNSPEC);
+}
+
+}  // namespace
+}  // namespace tent
+}  // namespace mooncake
diff --git a/mooncake-transfer-engine/tests/CMakeLists.txt b/mooncake-transfer-engine/tests/CMakeLists.txt
index 1c630b59..a36fc058 100644
--- a/mooncake-transfer-engine/tests/CMakeLists.txt
+++ b/mooncake-transfer-engine/tests/CMakeLists.txt
@@ -1,127 +1,261 @@
 set(WORKSPACE "${CMAKE_CURRENT_SOURCE_DIR}")
 
-if (USE_HIP)
+if(USE_HIP)
   file(GLOB TEST_SOURCES "*.cpp")
   hipify_files(TEST_SOURCES)
 
-  file(RELATIVE_PATH EXAMPLE_REL_PATH "${CMAKE_SOURCE_DIR}" "${CMAKE_CURRENT_SOURCE_DIR}")
+  file(RELATIVE_PATH EXAMPLE_REL_PATH "${CMAKE_SOURCE_DIR}"
+       "${CMAKE_CURRENT_SOURCE_DIR}")
   set(WORKSPACE "${CMAKE_BINARY_DIR}/${EXAMPLE_REL_PATH}")
 endif()
 
 add_executable(rdma_transport_test ${WORKSPACE}/rdma_transport_test.cpp)
-target_link_libraries(rdma_transport_test PUBLIC transfer_engine gflags::gflags glog::glog)
+target_link_libraries(rdma_transport_test PUBLIC transfer_engine gflags::gflags
+                                                 glog::glog)
 # add_test(NAME rdma_transport_test COMMAND rdma_transport_test)
 
 add_executable(transport_uint_test ${WORKSPACE}/transport_uint_test.cpp)
-target_link_libraries(transport_uint_test PUBLIC transfer_engine gtest gtest_main )
+target_link_libraries(transport_uint_test PUBLIC transfer_engine gtest
+                                                 gtest_main)
 add_test(NAME transport_uint_test COMMAND transport_uint_test)
 
+add_executable(endpoint_store_test ${WORKSPACE}/endpoint_store_test.cpp)
+target_link_libraries(endpoint_store_test PUBLIC transfer_engine gtest
+                                                 gtest_main)
+add_test(NAME endpoint_store_test COMMAND endpoint_store_test)
+
+# Integration test for the monitorWorker reclaim tick (issue #1845). Self-skips
+# when no RDMA device is present, so safe to register with ctest.
+add_executable(endpoint_store_integration_test
+               ${WORKSPACE}/endpoint_store_integration_test.cpp)
+target_link_libraries(endpoint_store_integration_test PUBLIC transfer_engine
+                                                             gtest gtest_main)
+add_test(NAME endpoint_store_integration_test
+         COMMAND endpoint_store_integration_test)
+set_tests_properties(endpoint_store_integration_test PROPERTIES LABELS "rdma")
+
 add_executable(rdma_transport_test2 ${WORKSPACE}/rdma_transport_test2.cpp)
-target_link_libraries(rdma_transport_test2 PUBLIC transfer_engine gtest gtest_main )
+target_link_libraries(rdma_transport_test2 PUBLIC transfer_engine gtest
+                                                  gtest_main)
 # add_test(NAME rdma_transport_test2 COMMAND rdma_transport_test2)
 
 add_executable(rdma_loopback_test ${WORKSPACE}/rdma_loopback_test.cpp)
-target_link_libraries(rdma_loopback_test PUBLIC transfer_engine gtest gtest_main )
+target_link_libraries(rdma_loopback_test PUBLIC transfer_engine gtest
+                                                gtest_main)
 # add_test(NAME rdma_loopback_test COMMAND rdma_loopback_test)
 
 # This test verifies endpoint re-establishment in RDMATransport.
-# Intended for manual testing only.
-add_executable(rdma_endpoint_reestablish_test ${WORKSPACE}/rdma_endpoint_reestablish_test.cpp)
-target_link_libraries(rdma_endpoint_reestablish_test PUBLIC transfer_engine gtest gtest_main )
+add_executable(rdma_endpoint_reestablish_test
+               ${WORKSPACE}/rdma_endpoint_reestablish_test.cpp)
+target_link_libraries(rdma_endpoint_reestablish_test PUBLIC transfer_engine
+                                                            gtest gtest_main)
+if(UNIX AND NOT APPLE)
+  target_link_options(
+    rdma_endpoint_reestablish_test PRIVATE "-Wl,--wrap=ibv_modify_qp"
+    "-Wl,--wrap=ibv_query_gid" "-Wl,--wrap=_ibv_query_gid_ex")
+endif()
+add_test(NAME rdma_endpoint_reestablish_test
+         COMMAND rdma_endpoint_reestablish_test)
+set_tests_properties(rdma_endpoint_reestablish_test PROPERTIES LABELS "rdma")
 
-if (USE_CXL)
-    add_executable(cxl_transport_test ${WORKSPACE}/cxl_transport_test.cpp)
-    target_link_libraries(cxl_transport_test PUBLIC transfer_engine gtest gtest_main )
-    add_test(NAME cxl_transport_test COMMAND cxl_transport_test)
+if(USE_CXL)
+  add_executable(cxl_transport_test ${WORKSPACE}/cxl_transport_test.cpp)
+  target_link_libraries(cxl_transport_test PUBLIC transfer_engine gtest
+                                                  gtest_main)
+  add_test(NAME cxl_transport_test COMMAND cxl_transport_test)
 endif()
 
-if (USE_NVMEOF)
-    add_executable(nvmeof_transport_test ${WORKSPACE}/nvmeof_transport_test.cpp)
-    target_link_libraries(nvmeof_transport_test PUBLIC transfer_engine gtest gtest_main )
-    # add_test(NAME nvmeof_transport_test COMMAND nvmeof_transport_test)
+if(USE_NVMEOF)
+  add_executable(nvmeof_transport_test ${WORKSPACE}/nvmeof_transport_test.cpp)
+  target_link_libraries(nvmeof_transport_test PUBLIC transfer_engine gtest
+                                                     gtest_main)
+  # add_test(NAME nvmeof_transport_test COMMAND nvmeof_transport_test)
 endif()
 
-if (USE_TCP)
-add_executable(tcp_transport_test ${WORKSPACE}/tcp_transport_test.cpp)
-target_link_libraries(tcp_transport_test PUBLIC transfer_engine gtest gtest_main )
-add_test(NAME tcp_transport_test COMMAND tcp_transport_test)
+if(USE_TCP)
+  add_executable(tcp_transport_test ${WORKSPACE}/tcp_transport_test.cpp)
+  target_link_libraries(tcp_transport_test PUBLIC transfer_engine gtest
+                                                  gtest_main)
+  add_test(NAME tcp_transport_test COMMAND tcp_transport_test)
 endif()
 
-if (USE_MNNVL)
-    add_executable(nvlink_transport_test ${WORKSPACE}/nvlink_transport_test.cpp)
-    target_link_libraries(nvlink_transport_test PUBLIC transfer_engine gtest gtest_main )
-    add_test(NAME nvlink_transport_test COMMAND nvlink_transport_test)
+add_executable(tcp_address_validation_test
+               ${WORKSPACE}/tcp_address_validation_test.cpp)
+target_link_libraries(tcp_address_validation_test PUBLIC gtest gtest_main)
+add_test(NAME tcp_address_validation_test COMMAND tcp_address_validation_test)
+
+if(USE_MNNVL)
+  add_executable(nvlink_transport_test ${WORKSPACE}/nvlink_transport_test.cpp)
+  target_link_libraries(nvlink_transport_test PUBLIC transfer_engine gtest
+                                                     gtest_main)
+  add_test(NAME nvlink_transport_test COMMAND nvlink_transport_test)
 endif()
 
-if (USE_UBSHMEM)
-    add_executable(ubshmem_transport_test ${WORKSPACE}/ubshmem_transport_test.cpp)
-    target_link_libraries(ubshmem_transport_test PUBLIC transfer_engine gtest gtest_main ascendcl)
-    add_test(NAME ubshmem_transport_test COMMAND ubshmem_transport_test)
+if(USE_HIP)
+  add_executable(hip_transport_test ${WORKSPACE}/hip_transport_test.cpp)
+  target_link_libraries(hip_transport_test PUBLIC transfer_engine gtest
+                                                  gtest_main)
+  add_test(NAME hip_transport_test COMMAND hip_transport_test)
 endif()
 
-if (USE_EFA)
-    add_executable(efa_transport_test ${WORKSPACE}/efa_transport_test.cpp)
-    target_link_libraries(efa_transport_test PUBLIC transfer_engine gtest gtest_main)
-    add_test(NAME efa_transport_test COMMAND efa_transport_test)
+if(USE_UBSHMEM)
+  add_executable(ubshmem_transport_test ${WORKSPACE}/ubshmem_transport_test.cpp)
+  target_link_libraries(ubshmem_transport_test PUBLIC transfer_engine gtest
+                                                      gtest_main ascendcl)
+  add_test(NAME ubshmem_transport_test COMMAND ubshmem_transport_test)
+endif()
 
-    add_executable(efa_c_api_test ${WORKSPACE}/efa_c_api_test.cpp)
-    target_link_libraries(efa_c_api_test PUBLIC transfer_engine gtest gtest_main)
-    add_test(NAME efa_c_api_test COMMAND efa_c_api_test)
+if(USE_EFA)
+  add_executable(efa_transport_test ${WORKSPACE}/efa_transport_test.cpp)
+  target_link_libraries(efa_transport_test PUBLIC transfer_engine gtest
+                                                  gtest_main)
+  add_test(NAME efa_transport_test COMMAND efa_transport_test)
 
-    add_executable(efa_single_nic_large_mr_test ${WORKSPACE}/efa_single_nic_large_mr_test.cpp)
-    target_link_libraries(efa_single_nic_large_mr_test PUBLIC transfer_engine gflags::gflags glog::glog)
+  add_executable(efa_c_api_test ${WORKSPACE}/efa_c_api_test.cpp)
+  target_link_libraries(efa_c_api_test PUBLIC transfer_engine gtest gtest_main)
+  add_test(NAME efa_c_api_test COMMAND efa_c_api_test)
 
-    add_executable(efa_transfer_test ${WORKSPACE}/efa_transfer_test.cpp)
-    target_link_libraries(efa_transfer_test PUBLIC transfer_engine gflags::gflags glog::glog)
+  add_executable(efa_single_nic_large_mr_test
+                 ${WORKSPACE}/efa_single_nic_large_mr_test.cpp)
+  target_link_libraries(efa_single_nic_large_mr_test
+                        PUBLIC transfer_engine gflags::gflags glog::glog)
+
+  add_executable(efa_transfer_test ${WORKSPACE}/efa_transfer_test.cpp)
+  target_link_libraries(efa_transfer_test PUBLIC transfer_engine gflags::gflags
+                                                 glog::glog)
+
+  # GPU (CUDA device memory) loopback test — reproduces the EFA SHM intra-node
+  # segfault on FI_HMEM_CUDA buffers (ofiwg/libfabric#12328) and validates
+  # EfaContext::tryLoopbackCopy.  Needs CUDA headers/libs.
+  if(USE_CUDA)
+    add_executable(efa_gpu_loopback_test ${WORKSPACE}/efa_gpu_loopback_test.cpp)
+    # Resolve CUDA include dirs / runtime via CUDAToolkit instead of a hardcoded
+    # /usr/local/cuda/include, so the test builds wherever CUDA lives (e.g. the
+    # DLAMI pip venv layout).  find_package is idempotent and may not have run
+    # yet in this scope (top-level only calls it under WITH_EP), so request it
+    # here; fall back to the legacy path if the module variant is unavailable.
+    find_package(CUDAToolkit QUIET)
+    if(CUDAToolkit_FOUND)
+      target_include_directories(efa_gpu_loopback_test
+                                 PRIVATE ${CUDAToolkit_INCLUDE_DIRS})
+      target_link_libraries(
+        efa_gpu_loopback_test PUBLIC transfer_engine gtest gtest_main
+                                     CUDA::cudart)
+    else()
+      target_include_directories(efa_gpu_loopback_test
+                                 PRIVATE /usr/local/cuda/include)
+      target_link_libraries(efa_gpu_loopback_test PUBLIC transfer_engine gtest
+                                                         gtest_main cudart)
+    endif()
+    add_test(NAME efa_gpu_loopback_test COMMAND efa_gpu_loopback_test)
+  endif()
+endif()
+
+if(USE_CXI)
+  add_executable(cxi_transport_test ${WORKSPACE}/cxi_transport_test.cpp)
+  add_executable(cxi_transfer_test ${WORKSPACE}/cxi_transfer_test.cpp)
+  add_executable(cxi_unit_tests ${WORKSPACE}/cxi_unit_tests.cpp)
+  target_link_libraries(cxi_transport_test PUBLIC transfer_engine gtest
+                                                  gtest_main)
+  target_link_libraries(cxi_transfer_test PUBLIC transfer_engine gtest
+                                                 gtest_main)
+  target_link_libraries(cxi_unit_tests PUBLIC transfer_engine gtest gtest_main)
+  add_test(NAME cxi_transport_test COMMAND cxi_transport_test)
+  add_test(NAME cxi_unit_tests COMMAND cxi_unit_test)
 endif()
 
 # UB transport test with URMA endpoint and mock support
-if (USE_UB)
-    add_executable(ub_transport_test ${WORKSPACE}/ub_transport_test.cpp)
-    target_link_libraries(ub_transport_test PUBLIC transfer_engine gtest gtest_main glog::glog pthread)
-    target_include_directories(ub_transport_test PRIVATE ${urma_INCLUDE_DIR})
-    add_test(NAME ub_transport_test COMMAND ub_transport_test)
+if(USE_UB)
+  add_executable(ub_transport_test ${WORKSPACE}/ub_transport_test.cpp)
+  target_link_libraries(ub_transport_test PUBLIC transfer_engine gtest
+                                                 gtest_main glog::glog pthread)
+  target_include_directories(ub_transport_test PRIVATE ${urma_INCLUDE_DIR})
+  # Built but not registered with ctest: this test may still have race
+  # conditions and other stability issues, so keep it out of CI for now. Run
+  # manually with ./ub_transport_test. add_test(NAME ub_transport_test COMMAND
+  # ub_transport_test)
+endif()
+
+if(USE_SUNRISE)
+  add_executable(sunrise_link_transport_test
+                 ${WORKSPACE}/sunrise_link_transport_test.cpp)
+  target_include_directories(sunrise_link_transport_test
+                             PRIVATE ${MC_TANGRT_ROOT}/include)
+  target_link_libraries(sunrise_link_transport_test PUBLIC transfer_engine
+                                                           gtest gtest_main)
+  add_test(NAME sunrise_link_transport_test COMMAND sunrise_link_transport_test)
+
+  add_executable(sunrise_link_transport_runtime_test
+                 ${WORKSPACE}/sunrise_link_transport_runtime_test.cpp)
+  target_include_directories(sunrise_link_transport_runtime_test
+                             PRIVATE ${MC_TANGRT_ROOT}/include)
+  target_link_libraries(sunrise_link_transport_runtime_test
+                        PUBLIC transfer_engine gtest gtest_main)
+  add_test(NAME sunrise_link_transport_runtime_test
+           COMMAND sunrise_link_transport_runtime_test)
+
+  add_executable(sunrise_link_transport_unit_test
+                 ${WORKSPACE}/sunrise_link_transport_unit_test.cpp)
+  target_include_directories(sunrise_link_transport_unit_test
+                             PRIVATE ${MC_TANGRT_ROOT}/include)
+  target_link_libraries(sunrise_link_transport_unit_test
+                        PUBLIC transfer_engine gtest gtest_main)
+  add_test(NAME sunrise_link_transport_unit_test
+           COMMAND sunrise_link_transport_unit_test)
 endif()
 
 add_executable(transfer_metadata_test ${WORKSPACE}/transfer_metadata_test.cpp)
-target_link_libraries(transfer_metadata_test PUBLIC transfer_engine gtest gtest_main)
+target_link_libraries(transfer_metadata_test PUBLIC transfer_engine gtest
+                                                    gtest_main)
 add_test(NAME transfer_metadata_test COMMAND transfer_metadata_test)
 
+add_executable(config_test ${WORKSPACE}/config_test.cpp)
+target_link_libraries(config_test PUBLIC transfer_engine gtest gtest_main)
+add_test(NAME config_test COMMAND config_test)
+
+add_executable(rdma_gid_probe_test ${WORKSPACE}/rdma_gid_probe_test.cpp)
+target_link_libraries(rdma_gid_probe_test PUBLIC transfer_engine gtest
+                                                 gtest_main)
+add_test(NAME rdma_gid_probe_test COMMAND rdma_gid_probe_test)
+
+add_executable(rdma_context_reprobe_test
+               ${WORKSPACE}/rdma_context_reprobe_test.cpp)
+target_link_libraries(rdma_context_reprobe_test PUBLIC transfer_engine gtest
+                                                       gtest_main)
+add_test(NAME rdma_context_reprobe_test COMMAND rdma_context_reprobe_test)
+
 add_executable(topology_test ${WORKSPACE}/topology_test.cpp)
 target_link_libraries(topology_test PUBLIC transfer_engine gtest gtest_main)
 add_test(NAME topology_test COMMAND topology_test)
 
 add_executable(memory_location_test ${WORKSPACE}/memory_location_test.cpp)
-target_link_libraries(memory_location_test PUBLIC transfer_engine gtest gtest_main)
+target_link_libraries(memory_location_test PUBLIC transfer_engine gtest
+                                                  gtest_main)
 add_test(NAME memory_location_test COMMAND memory_location_test)
 
 add_executable(common_test ${WORKSPACE}/common_test.cpp)
 target_link_libraries(common_test PUBLIC transfer_engine gtest gtest_main)
 add_test(NAME common_test COMMAND common_test)
 
-if (USE_ASCEND_DIRECT)
-    # AscendDirectTransport unit test with mock implementations
-    # Mock implementations are included in the test file via anonymous namespace
-    add_executable(ascend_direct_transport_test
-        ${WORKSPACE}/ascend_direct_transport_test.cpp
-    )
-    target_link_libraries(ascend_direct_transport_test PUBLIC
-        transfer_engine
-        gtest
-        gtest_main
-        metadef
-    )
-    # Allow shared library undefined symbols (Ascend libraries have circular deps)
-    set_target_properties(ascend_direct_transport_test PROPERTIES
-        LINK_FLAGS "-Wl,--allow-shlib-undefined"
-    )
+if(USE_ASCEND_DIRECT)
+  # AscendDirectTransport unit test with mock implementations Mock
+  # implementations are included in the test file via anonymous namespace
+  add_executable(ascend_direct_transport_test
+                 ${WORKSPACE}/ascend_direct_transport_test.cpp)
+  target_link_libraries(ascend_direct_transport_test
+                        PUBLIC transfer_engine gtest gtest_main metadef)
+  # Allow shared library undefined symbols (Ascend libraries have circular deps)
+  set_target_properties(ascend_direct_transport_test
+                        PROPERTIES LINK_FLAGS "-Wl,--allow-shlib-undefined")
 
-    add_test(NAME ascend_direct_transport_test COMMAND ascend_direct_transport_test)
+  add_test(NAME ascend_direct_transport_test
+           COMMAND ascend_direct_transport_test)
 endif()
 
 # Multi-protocol transport test (only when ENABLE_MULTI_PROTOCOL is ON)
-if (ENABLE_MULTI_PROTOCOL)
-    add_executable(mp_transport_test ${WORKSPACE}/mp_transport_test.cpp)
-    target_link_libraries(mp_transport_test PUBLIC transfer_engine gtest gtest_main)
-    add_test(NAME mp_transport_test COMMAND mp_transport_test)
+if(ENABLE_MULTI_PROTOCOL)
+  add_executable(mp_transport_test ${WORKSPACE}/mp_transport_test.cpp)
+  target_link_libraries(mp_transport_test PUBLIC transfer_engine gtest
+                                                 gtest_main)
+  add_test(NAME mp_transport_test COMMAND mp_transport_test)
 endif()
diff --git a/mooncake-transfer-engine/tests/ascend_direct_transport_test.cpp b/mooncake-transfer-engine/tests/ascend_direct_transport_test.cpp
index f2de2fe3..46dbd084 100644
--- a/mooncake-transfer-engine/tests/ascend_direct_transport_test.cpp
+++ b/mooncake-transfer-engine/tests/ascend_direct_transport_test.cpp
@@ -21,6 +21,7 @@
 #include <cstdlib>
 #include "transport/ascend_transport/ascend_direct_transport/ascend_direct_transport.h"
 #include "transport/ascend_transport/ascend_direct_transport/context_manager.h"
+#include "transport/ascend_transport/ascend_direct_transport/utils.h"
 #include "ascend_allocator.h"
 #include "transport/ascend_transport/ascend_direct_transport/adxl_compat.h"
 #include "transfer_metadata.h"
@@ -329,6 +330,7 @@ static int g_transfer_count = 0;
 static int g_transfer_async_count = 0;
 static int g_register_mem_count = 0;
 static int g_deregister_mem_count = 0;
+static std::string g_last_connect_target;
 
 namespace adxl_mock {
 void reset() {
@@ -353,6 +355,7 @@ void reset() {
     g_transfer_async_count = 0;
     g_register_mem_count = 0;
     g_deregister_mem_count = 0;
+    g_last_connect_target.clear();
 }
 
 void set_connect_result(adxl::Status status) {
@@ -446,6 +449,11 @@ std::vector<uintptr_t> get_deregistered_mem_handles() {
     std::lock_guard<std::mutex> lock(g_mutex);
     return g_deregistered_mem_handles;
 }
+
+std::string get_last_connect_target() {
+    std::lock_guard<std::mutex> lock(g_mutex);
+    return g_last_connect_target;
+}
 }  // namespace adxl_mock
 
 }  // namespace
@@ -474,6 +482,7 @@ Status AdxlEngine::Connect(const AscendString& remote_engine,
     (void)timeout_in_millis;
     std::lock_guard<std::mutex> lock(g_mutex);
     g_connected.insert(std::string(remote_engine.GetString()));
+    g_last_connect_target = remote_engine.GetString();
     g_connect_count++;
     return g_connect_result;
 }
@@ -583,6 +592,8 @@ Status AdxlEngine::DeregisterMem(MemHandle mem_handle) {
 class AscendDirectTransportTest : public ::testing::Test {
    protected:
     void SetUp() override {
+        // Unit tests mock explicit Connect(); disable auto-connect by default.
+        setenv("ASCEND_AUTO_CONNECT", "0", 1);
         mock_acl::reset();
         adxl_mock::reset();
 
@@ -604,6 +615,7 @@ class AscendDirectTransportTest : public ::testing::Test {
     }
 
     void TearDown() override {
+        unsetenv("ASCEND_AUTO_CONNECT");
         ContextManager::getInstance().finalize();
         google::ShutdownGoogleLogging();
     }
@@ -645,6 +657,18 @@ class AscendDirectTransportTest : public ::testing::Test {
         meta->addLocalSegment(s.segment_id, s.name, std::move(remote_desc));
     }
 
+    void addMultiEndpointRemoteSegment(
+        std::shared_ptr<TransferMetadata> meta, int segment_id,
+        const std::string& name, const std::string& host_ip,
+        const std::vector<std::string>& endpoints) {
+        auto remote_desc = std::make_shared<TransferMetadata::SegmentDesc>();
+        remote_desc->name = name;
+        remote_desc->protocol = "ascend";
+        remote_desc->rank_info.hostIp = host_ip;
+        remote_desc->rank_info.endpoints = endpoints;
+        meta->addLocalSegment(segment_id, name, std::move(remote_desc));
+    }
+
     std::shared_ptr<TransferMetadata> startRemoteMetadataServer(
         const RemoteSegmentSetup& s) {
         auto metadata = std::make_shared<TransferMetadata>("P2PHANDSHAKE");
@@ -1065,6 +1089,40 @@ TEST_F(AscendDirectTransportTest, DummyReal_Roce_LocalCopy_Success) {
     globalConfig().ascend_agent_mode = false;
 }
 
+TEST_F(AscendDirectTransportTest,
+       DummyReal_Roce_SameHostDifferentProcess_SelectsSameIndex) {
+    globalConfig().ascend_agent_mode = true;
+    setenv("HCCL_INTRA_ROCE_ENABLE", "1", 1);
+    constexpr int kDeviceCount = 4;
+    mock_acl::set_device_count(kDeviceCount);
+    g_device_id = 0;
+    ContextManager::getInstance().finalize();
+    ASSERT_TRUE(ContextManager::getInstance().initialize());
+    auto transport = createTransport();
+    ASSERT_NE(transport, nullptr);
+    ASSERT_EQ(transport->registerLocalMemory(test_buffer_src_, kRegisterMemSize,
+                                             "cpu:0", true, true),
+              0);
+
+    // Same host with multiple endpoints - dummy-real should use same-index
+    std::vector<std::string> endpoints = {"127.0.0.1:9000", "127.0.0.1:9100",
+                                          "127.0.0.1:9200", "127.0.0.1:9300"};
+    addMultiEndpointRemoteSegment(transport->meta(), 1, "dummy_real_same_host",
+                                  "127.0.0.1", endpoints);
+
+    initTestData(kTransferBufSize);
+    auto result = runRemoteTransfer(transport.get(), test_buffer_src_, 1,
+                                    0x10000, kTransferBufSize);
+    ASSERT_TRUE(result.finished);
+    // dummy-real uses same-index pairing: engine_id=0 -> endpoints[0]
+    // Different process so no local copy, ADXL connect happens
+    EXPECT_EQ(adxl_mock::get_last_connect_target(), "127.0.0.1:9000")
+        << "Dummy-real should use same-index pairing without offset";
+
+    unsetenv("HCCL_INTRA_ROCE_ENABLE");
+    globalConfig().ascend_agent_mode = false;
+}
+
 // -----------------------------------------------------------------------------
 // Memory registration tests
 // -----------------------------------------------------------------------------
@@ -1215,7 +1273,7 @@ TEST_F(AscendDirectTransportTest, Batch_ExceedsCapacity) {
     transport->freeBatchID(batch_id);
 }
 
-TEST_F(AscendDirectTransportTest, Batch_allocateBatchID_freeBatchID_Success) {
+TEST_F(AscendDirectTransportTest, Batch_AllocateFree_Success) {
     auto transport = createTransport();
     ASSERT_NE(transport, nullptr);
 
@@ -1313,7 +1371,7 @@ TEST_F(AscendDirectTransportTest, LocalCopy_Sync_aclrtMemcpyFailure) {
     transport->freeBatchID(batch_id);
 }
 
-TEST_F(AscendDirectTransportTest, LocalCopy_Async_aclrtMemcpyAsyncFailure) {
+TEST_F(AscendDirectTransportTest, LocalCopy_Async_MemcpyFailure) {
     auto transport = createTransport();
     ASSERT_NE(transport, nullptr);
 
@@ -1619,7 +1677,7 @@ TEST_F(AscendDirectTransportTest,
     EXPECT_EQ(adxl_mock::get_disconnect_count(), 1);
 }
 
-TEST_F(AscendDirectTransportTest, RemoteTransfer_Async_TransferAsyncFailure) {
+TEST_F(AscendDirectTransportTest, RemoteTransfer_Async_TransferFailure) {
     adxl_mock::set_transfer_async_result(adxl::FAILED);
 
     auto transport = createTransport(true);
@@ -1695,6 +1753,385 @@ TEST_F(AscendDirectTransportTest, RemoteTransfer_Async_TransferTimeout) {
                                   "(GetTransferStatus stays WAITING)";
 }
 
+TEST_F(AscendDirectTransportTest,
+       RemoteTransfer_Sync_FailureWithAutoConnect_SkipsDisconnect) {
+    setenv("ASCEND_AUTO_CONNECT", "1", 1);
+    adxl_mock::set_transfer_result(adxl::FAILED);
+
+    auto transport = createTransport();
+    ASSERT_NE(transport, nullptr);
+    ASSERT_EQ(transport->registerLocalMemory(test_buffer_src_, kRegisterMemSize,
+                                             "cpu:0", true, true),
+              0);
+    addRemoteSegment(transport->meta(),
+                     {1, "remote_server", "192.168.1.100", 30000});
+
+    initTestData(kTransferBufSize);
+    auto result = runRemoteTransfer(transport.get(), test_buffer_src_, 1,
+                                    0x10000, kTransferBufSize);
+    ASSERT_TRUE(result.finished);
+    EXPECT_TRUE(result.failed);
+    EXPECT_EQ(adxl_mock::get_disconnect_count(), 0);
+}
+
+TEST_F(AscendDirectTransportTest,
+       RemoteTransfer_Async_SubmitFailureWithAutoConnect_SkipsDisconnect) {
+    setenv("ASCEND_AUTO_CONNECT", "1", 1);
+    adxl_mock::set_transfer_async_result(adxl::FAILED);
+
+    auto transport = createTransport(true);
+    ASSERT_NE(transport, nullptr);
+    ASSERT_EQ(transport->registerLocalMemory(test_buffer_src_, kRegisterMemSize,
+                                             "cpu:0", true, true),
+              0);
+    addRemoteSegment(transport->meta(),
+                     {1, "remote_server", "192.168.1.100", 30000});
+
+    initTestData(kTransferBufSize);
+    auto result = runRemoteTransfer(transport.get(), test_buffer_src_, 1,
+                                    0x10000, kTransferBufSize);
+    ASSERT_TRUE(result.finished);
+    EXPECT_TRUE(result.failed);
+    EXPECT_EQ(adxl_mock::get_disconnect_count(), 0);
+}
+
+TEST_F(AscendDirectTransportTest,
+       RemoteTransfer_Async_GetStatusFailureWithAutoConnect_SkipsDisconnect) {
+    setenv("ASCEND_AUTO_CONNECT", "1", 1);
+    adxl_mock::set_transfer_async_result(adxl::SUCCESS);
+    adxl_mock::set_get_transfer_status_result(adxl::FAILED);
+
+    auto transport = createTransport(true);
+    ASSERT_NE(transport, nullptr);
+    ASSERT_EQ(transport->registerLocalMemory(test_buffer_src_, kRegisterMemSize,
+                                             "cpu:0", true, true),
+              0);
+    addRemoteSegment(transport->meta(),
+                     {1, "remote_server", "192.168.1.100", 30000});
+
+    initTestData(kTransferBufSize);
+    auto result = runRemoteTransfer(transport.get(), test_buffer_src_, 1,
+                                    0x10000, kTransferBufSize);
+    ASSERT_TRUE(result.finished);
+    EXPECT_TRUE(result.failed);
+    EXPECT_EQ(adxl_mock::get_disconnect_count(), 0);
+}
+
+TEST_F(AscendDirectTransportTest,
+       RemoteTransfer_Async_TimeoutWithAutoConnect_StillDisconnects) {
+    setenv("ASCEND_AUTO_CONNECT", "1", 1);
+    setenv("ASCEND_TRANSFER_TIMEOUT", "100", 1);
+    adxl_mock::set_transfer_async_result(adxl::SUCCESS);
+    adxl_mock::set_get_transfer_status_result(adxl::SUCCESS);
+    adxl_mock::set_transfer_status_enum(adxl::TransferStatus::WAITING);
+
+    auto transport = createTransport(true);
+    ASSERT_NE(transport, nullptr);
+    ASSERT_EQ(transport->registerLocalMemory(test_buffer_src_, kRegisterMemSize,
+                                             "cpu:0", true, true),
+              0);
+    addRemoteSegment(transport->meta(),
+                     {1, "remote_server", "192.168.1.100", 30000});
+
+    initTestData(kTransferBufSize);
+    auto result = runRemoteTransfer(transport.get(), test_buffer_src_, 1,
+                                    0x10000, kTransferBufSize);
+    ASSERT_TRUE(result.finished);
+    EXPECT_TRUE(result.failed);
+    EXPECT_EQ(adxl_mock::get_disconnect_count(), 1);
+}
+
+// -----------------------------------------------------------------------------
+// Standalone mode tests (non-dummy-real, non-fabric-mem)
+// -----------------------------------------------------------------------------
+
+TEST_F(AscendDirectTransportTest, Standalone_RemoteHost_SelectsPhyDevEndpoint) {
+    globalConfig().ascend_agent_mode = false;
+    unsetenv("HCCL_INTRA_ROCE_ENABLE");
+    constexpr int kDeviceCount = 4;
+    mock_acl::set_device_count(kDeviceCount);
+    g_device_id = 2;
+    auto transport = createTransport();
+    ASSERT_NE(transport, nullptr);
+    ASSERT_EQ(transport->registerLocalMemory(test_buffer_src_, kRegisterMemSize,
+                                             "cpu:0", true, true),
+              0);
+
+    // Remote host with multiple endpoints (different from local host)
+    std::vector<std::string> endpoints = {"10.0.0.1:5000", "10.0.0.1:5100",
+                                          "10.0.0.1:5200", "10.0.0.1:5300"};
+    addMultiEndpointRemoteSegment(transport->meta(), 1, "remote_host",
+                                  "10.0.0.1", endpoints);
+
+    initTestData(kTransferBufSize);
+    auto result = runRemoteTransfer(transport.get(), test_buffer_src_, 1,
+                                    0x10000, kTransferBufSize);
+    ASSERT_TRUE(result.finished);
+    EXPECT_FALSE(result.failed);
+    // phy_dev=2, different host -> selects endpoints[2]
+    EXPECT_EQ(adxl_mock::get_last_connect_target(), "10.0.0.1:5200");
+}
+
+TEST_F(AscendDirectTransportTest,
+       Standalone_SameHostDifferentProcess_OffsetsByOne) {
+    globalConfig().ascend_agent_mode = false;
+    unsetenv("HCCL_INTRA_ROCE_ENABLE");
+    constexpr int kDeviceCount = 4;
+    mock_acl::set_device_count(kDeviceCount);
+    g_device_id = 1;
+    auto transport = createTransport();
+    ASSERT_NE(transport, nullptr);
+    ASSERT_EQ(transport->registerLocalMemory(test_buffer_src_, kRegisterMemSize,
+                                             "cpu:0", true, true),
+              0);
+
+    // Same host (127.0.0.1) with multiple endpoints (different process)
+    std::vector<std::string> endpoints = {"127.0.0.1:6000", "127.0.0.1:6100",
+                                          "127.0.0.1:6200", "127.0.0.1:6300"};
+    addMultiEndpointRemoteSegment(transport->meta(), 1, "same_host_process",
+                                  "127.0.0.1", endpoints);
+
+    initTestData(kTransferBufSize);
+    auto result = runRemoteTransfer(transport.get(), test_buffer_src_, 1,
+                                    0x10000, kTransferBufSize);
+    ASSERT_TRUE(result.finished);
+    EXPECT_FALSE(result.failed);
+    // phy_dev=1, same host -> offset +1 -> selects endpoints[2]
+    EXPECT_EQ(adxl_mock::get_last_connect_target(), "127.0.0.1:6200");
+}
+
+TEST_F(AscendDirectTransportTest, Standalone_SameHostOffset_WrapsAround) {
+    globalConfig().ascend_agent_mode = false;
+    unsetenv("HCCL_INTRA_ROCE_ENABLE");
+    constexpr int kDeviceCount = 4;
+    mock_acl::set_device_count(kDeviceCount);
+    g_device_id = 3;  // last device: (3+1)%4 = 0
+    auto transport = createTransport();
+    ASSERT_NE(transport, nullptr);
+    ASSERT_EQ(transport->registerLocalMemory(test_buffer_src_, kRegisterMemSize,
+                                             "cpu:0", true, true),
+              0);
+
+    std::vector<std::string> endpoints = {"127.0.0.1:7000", "127.0.0.1:7100",
+                                          "127.0.0.1:7200", "127.0.0.1:7300"};
+    addMultiEndpointRemoteSegment(transport->meta(), 1, "same_host_wrap",
+                                  "127.0.0.1", endpoints);
+
+    initTestData(kTransferBufSize);
+    auto result = runRemoteTransfer(transport.get(), test_buffer_src_, 1,
+                                    0x10000, kTransferBufSize);
+    ASSERT_TRUE(result.finished);
+    EXPECT_FALSE(result.failed);
+    // phy_dev=3, same host -> (3+1)%4=0 -> selects endpoints[0]
+    EXPECT_EQ(adxl_mock::get_last_connect_target(), "127.0.0.1:7000");
+}
+
+TEST_F(AscendDirectTransportTest,
+       Standalone_FabricMem_SameHostUsesFrontEndpoint) {
+    // Fabric mem is a Store-TE feature, so the TE must be store-init for the
+    // transport to capture use_fabric_mem_=true.
+    globalConfig().ascend_agent_mode = false;
+    globalConfig().ascend_store_te_init = true;
+    globalConfig().ascend_use_fabric_mem = true;
+    unsetenv("HCCL_INTRA_ROCE_ENABLE");
+    constexpr int kDeviceCount = 4;
+    mock_acl::set_device_count(kDeviceCount);
+    g_device_id = 1;
+    auto transport = createTransport();
+    ASSERT_NE(transport, nullptr);
+    ASSERT_EQ(transport->registerLocalMemory(test_buffer_src_, kRegisterMemSize,
+                                             "cpu:0", true, true),
+              0);
+
+    std::vector<std::string> endpoints = {"127.0.0.1:8000", "127.0.0.1:8100",
+                                          "127.0.0.1:8200", "127.0.0.1:8300"};
+    addMultiEndpointRemoteSegment(transport->meta(), 1, "fabric_same_host",
+                                  "127.0.0.1", endpoints);
+
+    initTestData(kTransferBufSize);
+    auto result = runRemoteTransfer(transport.get(), test_buffer_src_, 1,
+                                    0x10000, kTransferBufSize);
+    ASSERT_TRUE(result.finished);
+    EXPECT_FALSE(result.failed);
+    EXPECT_EQ(adxl_mock::get_last_connect_target(), "127.0.0.1:8000")
+        << "Fabric-mem standalone should not apply same-host offset";
+
+    globalConfig().ascend_use_fabric_mem = false;
+    globalConfig().ascend_store_te_init = false;
+}
+
+// A non-Store (e.g. P2P/HCCS) TE must NOT inherit a Store TE's fabric flag that
+// leaked into the process-global config: with ascend_use_fabric_mem=true but
+// ascend_store_te_init=false, the transport must capture use_fabric_mem_=false
+// and behave like a normal standalone TE (same-host offset +1), not the fabric
+// path (front endpoint).
+TEST_F(AscendDirectTransportTest,
+       Standalone_FabricFlagWithoutStoreTe_DoesNotInheritFabric) {
+    globalConfig().ascend_agent_mode = false;
+    globalConfig().ascend_store_te_init = false;
+    globalConfig().ascend_use_fabric_mem = true;
+    unsetenv("HCCL_INTRA_ROCE_ENABLE");
+    constexpr int kDeviceCount = 4;
+    mock_acl::set_device_count(kDeviceCount);
+    g_device_id = 1;
+    auto transport = createTransport();
+    ASSERT_NE(transport, nullptr);
+    ASSERT_EQ(transport->registerLocalMemory(test_buffer_src_, kRegisterMemSize,
+                                             "cpu:0", true, true),
+              0);
+
+    std::vector<std::string> endpoints = {"127.0.0.1:9000", "127.0.0.1:9100",
+                                          "127.0.0.1:9200", "127.0.0.1:9300"};
+    addMultiEndpointRemoteSegment(transport->meta(), 1, "p2p_no_fabric",
+                                  "127.0.0.1", endpoints);
+
+    initTestData(kTransferBufSize);
+    auto result = runRemoteTransfer(transport.get(), test_buffer_src_, 1,
+                                    0x10000, kTransferBufSize);
+    ASSERT_TRUE(result.finished);
+    EXPECT_FALSE(result.failed);
+    // phy_dev=1, same host, fabric NOT inherited -> offset +1 -> endpoints[2]
+    EXPECT_EQ(adxl_mock::get_last_connect_target(), "127.0.0.1:9200")
+        << "Non-Store TE must not inherit fabric; should apply same-host "
+           "offset";
+
+    globalConfig().ascend_use_fabric_mem = false;
+}
+
+// -----------------------------------------------------------------------------
+// Roce mode detection (HCCL_INTRA_ROCE_ENABLE / ASCEND_GLOBAL_RESOURCE_CONFIG)
+// -----------------------------------------------------------------------------
+
+TEST(RoceModeDetectionTest, GlobalResourceConfig_StringRoceDesc) {
+    EXPECT_TRUE(HasRoceProtocolDescInGlobalResourceConfig(
+        R"({"comm_resource_config.protocol_desc":"roce:device"})"));
+}
+
+TEST(RoceModeDetectionTest, GlobalResourceConfig_ArrayRoceDesc) {
+    EXPECT_TRUE(HasRoceProtocolDescInGlobalResourceConfig(
+        R"({"comm_resource_config.protocol_desc":["hccs:device","roce:host"]})"));
+}
+
+TEST(RoceModeDetectionTest, GlobalResourceConfig_NonRoceDesc) {
+    EXPECT_FALSE(HasRoceProtocolDescInGlobalResourceConfig(
+        R"({"comm_resource_config.protocol_desc":"hccs:device"})"));
+    EXPECT_FALSE(HasRoceProtocolDescInGlobalResourceConfig(
+        R"({"comm_resource_config.listen_port":26666})"));
+}
+
+TEST(RoceModeDetectionTest, GlobalResourceConfig_NestedRoceDesc) {
+    EXPECT_TRUE(HasRoceProtocolDescInGlobalResourceConfig(
+        R"({"comm_resource_config":{"protocol_desc":"roce:device"}})"));
+}
+
+TEST(RoceModeDetectionTest, GlobalResourceConfig_InvalidJson) {
+    EXPECT_FALSE(HasRoceProtocolDescInGlobalResourceConfig("{not json"));
+    EXPECT_FALSE(HasRoceProtocolDescInGlobalResourceConfig(
+        R"("comm_resource_config.protocol_desc":"roce:device")"));
+}
+
+TEST(RoceModeDetectionTest, IsRoceModeEnabled_FromGlobalResourceConfig) {
+    unsetenv("HCCL_INTRA_ROCE_ENABLE");
+    setenv("ASCEND_GLOBAL_RESOURCE_CONFIG",
+           R"({"comm_resource_config.protocol_desc":"roce:device"})", 1);
+    EXPECT_TRUE(IsRoceModeEnabled());
+    unsetenv("ASCEND_GLOBAL_RESOURCE_CONFIG");
+}
+
+TEST(RoceModeDetectionTest, IsRoceModeEnabled_FromHcclEnv) {
+    unsetenv("ASCEND_GLOBAL_RESOURCE_CONFIG");
+    setenv("HCCL_INTRA_ROCE_ENABLE", "1", 1);
+    EXPECT_TRUE(IsRoceModeEnabled());
+    unsetenv("HCCL_INTRA_ROCE_ENABLE");
+}
+
+// -----------------------------------------------------------------------------
+// Store vs P2P protocol split via a single ASCEND_GLOBAL_RESOURCE_CONFIG:
+// top-level = default (P2P, e.g. HCCS), optional "store" sub-object overrides
+// it for a Store-init TE (e.g. RoCE). Resolution is gated on
+// ascend_store_te_init.
+// -----------------------------------------------------------------------------
+
+namespace {
+// Top-level default = HCCS (P2P), "store" override = RoCE (Store).
+constexpr const char* kDualProtocolConfig =
+    R"({"comm_resource_config":{"protocol_desc":"hccs:device"},)"
+    R"("store":{"comm_resource_config":{"protocol_desc":"roce:device"}}})";
+
+struct StoreTeInitScope {
+    explicit StoreTeInitScope(bool v) {
+        globalConfig().ascend_store_te_init = v;
+    }
+    ~StoreTeInitScope() { globalConfig().ascend_store_te_init = false; }
+};
+}  // namespace
+
+TEST(StoreResourceConfigSplitTest, NoStoreKey_PassthroughVerbatimBothRoles) {
+    const char* cfg =
+        R"({"comm_resource_config":{"protocol_desc":"hccs:device"}})";
+    {
+        StoreTeInitScope store(true);
+        EXPECT_EQ(ResolveAscendGlobalResourceConfig(cfg), std::string(cfg));
+    }
+    {
+        StoreTeInitScope store(false);
+        EXPECT_EQ(ResolveAscendGlobalResourceConfig(cfg), std::string(cfg));
+    }
+}
+
+TEST(StoreResourceConfigSplitTest, EmptyOrNull_ReturnsEmpty) {
+    StoreTeInitScope store(true);
+    EXPECT_TRUE(ResolveAscendGlobalResourceConfig(nullptr).empty());
+    EXPECT_TRUE(ResolveAscendGlobalResourceConfig("").empty());
+}
+
+TEST(StoreResourceConfigSplitTest, StoreRole_SelectsStoreSubtreeRoce) {
+    StoreTeInitScope store(true);
+    std::string resolved =
+        ResolveAscendGlobalResourceConfig(kDualProtocolConfig);
+    // Store TE gets the "store" subtree (RoCE), with the "store" key stripped.
+    EXPECT_TRUE(HasRoceProtocolDescInGlobalResourceConfig(resolved.c_str()));
+    EXPECT_EQ(resolved.find("hccs"), std::string::npos);
+    EXPECT_EQ(resolved.find("\"store\""), std::string::npos);
+}
+
+TEST(StoreResourceConfigSplitTest, DefaultRole_StripsStoreKeyHccs) {
+    StoreTeInitScope store(false);
+    std::string resolved =
+        ResolveAscendGlobalResourceConfig(kDualProtocolConfig);
+    // P2P/default TE gets the top-level (HCCS) with the "store" key removed.
+    EXPECT_FALSE(HasRoceProtocolDescInGlobalResourceConfig(resolved.c_str()));
+    EXPECT_NE(resolved.find("hccs"), std::string::npos);
+    EXPECT_EQ(resolved.find("roce"), std::string::npos);
+    EXPECT_EQ(resolved.find("\"store\""), std::string::npos);
+}
+
+TEST(StoreResourceConfigSplitTest,
+     StoreRoleMissingStoreKey_FallsBackToDefault) {
+    const char* cfg =
+        R"({"comm_resource_config":{"protocol_desc":"hccs:device"}})";
+    StoreTeInitScope store(true);
+    // No "store" key -> store TE falls back to the (verbatim) default config.
+    EXPECT_EQ(ResolveAscendGlobalResourceConfig(cfg), std::string(cfg));
+    EXPECT_FALSE(HasRoceProtocolDescInGlobalResourceConfig(
+        ResolveAscendGlobalResourceConfig(cfg).c_str()));
+}
+
+// The headline case: one env var, Store TE -> RoCE, P2P TE -> HCCS.
+TEST(StoreResourceConfigSplitTest, IsRoceModeEnabled_StoreRoceP2pHccs) {
+    unsetenv("HCCL_INTRA_ROCE_ENABLE");
+    setenv("ASCEND_GLOBAL_RESOURCE_CONFIG", kDualProtocolConfig, 1);
+    {
+        StoreTeInitScope store(true);
+        EXPECT_TRUE(IsRoceModeEnabled()) << "Store TE should resolve to RoCE";
+    }
+    {
+        StoreTeInitScope store(false);
+        EXPECT_FALSE(IsRoceModeEnabled()) << "P2P TE should resolve to HCCS";
+    }
+    unsetenv("ASCEND_GLOBAL_RESOURCE_CONFIG");
+}
+
 int main(int argc, char** argv) {
     ::testing::InitGoogleTest(&argc, argv);
     return RUN_ALL_TESTS();
diff --git a/mooncake-transfer-engine/tests/common_test.cpp b/mooncake-transfer-engine/tests/common_test.cpp
index 887ed426..b07971be 100644
--- a/mooncake-transfer-engine/tests/common_test.cpp
+++ b/mooncake-transfer-engine/tests/common_test.cpp
@@ -112,6 +112,27 @@ TEST(ParseHostNameWithPort, HostWithoutPort) {
     EXPECT_EQ(port, kDefaultPort);
 }
 
+TEST(HostPortHelpers, DetectsExplicitPort) {
+    EXPECT_TRUE(hasExplicitPort("8.8.8.8:4321"));
+    EXPECT_TRUE(hasExplicitPort("example.com:4321"));
+    EXPECT_TRUE(hasExplicitPort("[2001:db8::1]:4321"));
+    EXPECT_FALSE(hasExplicitPort("2001:db8::1"));
+    EXPECT_FALSE(hasExplicitPort("[2001:db8::1]"));
+}
+
+TEST(HostPortHelpers, ExtractsHostWithoutPort) {
+    EXPECT_EQ(getHostNameWithoutPort("8.8.8.8:4321"), "8.8.8.8");
+    EXPECT_EQ(getHostNameWithoutPort("example.com:4321"), "example.com");
+    EXPECT_EQ(getHostNameWithoutPort("[2001:db8::1]:4321"), "2001:db8::1");
+    EXPECT_EQ(getHostNameWithoutPort("[2001:db8::1]"), "2001:db8::1");
+    EXPECT_EQ(getHostNameWithoutPort("2001:db8::1"), "2001:db8::1");
+}
+
+TEST(HostPortHelpers, BuildsBracketedIpv6Endpoint) {
+    EXPECT_EQ(buildHostNameWithPort("2001:db8::1", 4321), "[2001:db8::1]:4321");
+    EXPECT_EQ(buildHostNameWithPort("8.8.8.8", 4321), "8.8.8.8:4321");
+}
+
 //------------------------------------------------------------------------------
 // parsePortAndDevice
 //------------------------------------------------------------------------------
diff --git a/mooncake-transfer-engine/tests/config_test.cpp b/mooncake-transfer-engine/tests/config_test.cpp
new file mode 100644
index 00000000..85c28652
--- /dev/null
+++ b/mooncake-transfer-engine/tests/config_test.cpp
@@ -0,0 +1,162 @@
+// Copyright 2024 KVCache.AI
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include <gtest/gtest.h>
+
+#include <cstdlib>
+
+#include "config.h"
+
+namespace mooncake {
+namespace {
+
+class PkeyIndexEnvTest : public ::testing::Test {
+   protected:
+    void TearDown() override {
+        ::unsetenv("MC_PKEY_INDEX");
+        ::unsetenv("MC_AUTO_GID_MAX_RETRIES");
+        ::unsetenv("MC_IB_SL");
+    }
+};
+
+TEST_F(PkeyIndexEnvTest, DefaultIsZeroWhenUnset) {
+    ::unsetenv("MC_PKEY_INDEX");
+    GlobalConfig config;
+    loadGlobalConfig(config);
+    EXPECT_EQ(config.pkey_index, 0);
+}
+
+TEST_F(PkeyIndexEnvTest, ValidOverrideIsApplied) {
+    ASSERT_EQ(::setenv("MC_PKEY_INDEX", "7", 1), 0);
+    GlobalConfig config;
+    loadGlobalConfig(config);
+    EXPECT_EQ(config.pkey_index, 7);
+}
+
+TEST_F(PkeyIndexEnvTest, MaxBoundaryIsApplied) {
+    ASSERT_EQ(::setenv("MC_PKEY_INDEX", "65535", 1), 0);
+    GlobalConfig config;
+    loadGlobalConfig(config);
+    EXPECT_EQ(config.pkey_index, 65535);
+}
+
+TEST_F(PkeyIndexEnvTest, OutOfRangeIsIgnored) {
+    ASSERT_EQ(::setenv("MC_PKEY_INDEX", "70000", 1), 0);
+    GlobalConfig config;
+    config.pkey_index = 3;  // sentinel preserved when env var is rejected
+    loadGlobalConfig(config);
+    EXPECT_EQ(config.pkey_index, 3);
+}
+
+TEST_F(PkeyIndexEnvTest, NegativeIsIgnored) {
+    ASSERT_EQ(::setenv("MC_PKEY_INDEX", "-1", 1), 0);
+    GlobalConfig config;
+    config.pkey_index = 5;
+    loadGlobalConfig(config);
+    EXPECT_EQ(config.pkey_index, 5);
+}
+
+TEST_F(PkeyIndexEnvTest, NonNumericKeepsDefault) {
+    ASSERT_EQ(::setenv("MC_PKEY_INDEX", "abc", 1), 0);
+    GlobalConfig config;
+    config.pkey_index = 9;
+    loadGlobalConfig(config);
+    EXPECT_EQ(config.pkey_index, 9);
+}
+
+TEST_F(PkeyIndexEnvTest, EmptyStringKeepsDefault) {
+    ASSERT_EQ(::setenv("MC_PKEY_INDEX", "", 1), 0);
+    GlobalConfig config;
+    config.pkey_index = 4;
+    loadGlobalConfig(config);
+    EXPECT_EQ(config.pkey_index, 4);
+}
+
+TEST_F(PkeyIndexEnvTest, AutoGidRetriesDefaultsToTwoWhenUnset) {
+    ::unsetenv("MC_AUTO_GID_MAX_RETRIES");
+    GlobalConfig config;
+    config.auto_gid_max_retries = 2;
+    loadGlobalConfig(config);
+    EXPECT_EQ(config.auto_gid_max_retries, 2);
+}
+
+TEST_F(PkeyIndexEnvTest, AutoGidRetriesAcceptsValidOverride) {
+    ASSERT_EQ(::setenv("MC_AUTO_GID_MAX_RETRIES", "0", 1), 0);
+    GlobalConfig config;
+    loadGlobalConfig(config);
+    EXPECT_EQ(config.auto_gid_max_retries, 0);
+}
+
+TEST_F(PkeyIndexEnvTest, AutoGidRetriesRejectsOutOfRangeOverride) {
+    ASSERT_EQ(::setenv("MC_AUTO_GID_MAX_RETRIES", "99", 1), 0);
+    GlobalConfig config;
+    config.auto_gid_max_retries = 5;
+    loadGlobalConfig(config);
+    EXPECT_EQ(config.auto_gid_max_retries, 5);
+}
+
+TEST_F(PkeyIndexEnvTest, IbSlDefaultsToMinusOneWhenUnset) {
+    ::unsetenv("MC_IB_SL");
+    GlobalConfig config;
+    loadGlobalConfig(config);
+    EXPECT_EQ(config.ib_service_level, -1);
+}
+
+TEST_F(PkeyIndexEnvTest, IbSlValidOverrideIsApplied) {
+    ASSERT_EQ(::setenv("MC_IB_SL", "3", 1), 0);
+    GlobalConfig config;
+    loadGlobalConfig(config);
+    EXPECT_EQ(config.ib_service_level, 3);
+}
+
+TEST_F(PkeyIndexEnvTest, IbSlMinBoundaryIsApplied) {
+    ASSERT_EQ(::setenv("MC_IB_SL", "0", 1), 0);
+    GlobalConfig config;
+    loadGlobalConfig(config);
+    EXPECT_EQ(config.ib_service_level, 0);
+}
+
+TEST_F(PkeyIndexEnvTest, IbSlMaxBoundaryIsApplied) {
+    ASSERT_EQ(::setenv("MC_IB_SL", "15", 1), 0);
+    GlobalConfig config;
+    loadGlobalConfig(config);
+    EXPECT_EQ(config.ib_service_level, 15);
+}
+
+TEST_F(PkeyIndexEnvTest, IbSlOutOfRangeIsIgnored) {
+    ASSERT_EQ(::setenv("MC_IB_SL", "16", 1), 0);
+    GlobalConfig config;
+    config.ib_service_level = 7;  // sentinel preserved when env var is rejected
+    loadGlobalConfig(config);
+    EXPECT_EQ(config.ib_service_level, 7);
+}
+
+TEST_F(PkeyIndexEnvTest, IbSlNegativeIsIgnored) {
+    ASSERT_EQ(::setenv("MC_IB_SL", "-1", 1), 0);
+    GlobalConfig config;
+    config.ib_service_level = 5;
+    loadGlobalConfig(config);
+    EXPECT_EQ(config.ib_service_level, 5);
+}
+
+TEST_F(PkeyIndexEnvTest, IbSlNonNumericKeepsDefault) {
+    ASSERT_EQ(::setenv("MC_IB_SL", "abc", 1), 0);
+    GlobalConfig config;
+    config.ib_service_level = 9;
+    loadGlobalConfig(config);
+    EXPECT_EQ(config.ib_service_level, 9);
+}
+
+}  // namespace
+}  // namespace mooncake
diff --git a/mooncake-transfer-engine/tests/cxi_transfer_test.cpp b/mooncake-transfer-engine/tests/cxi_transfer_test.cpp
new file mode 100644
index 00000000..763058f4
--- /dev/null
+++ b/mooncake-transfer-engine/tests/cxi_transfer_test.cpp
@@ -0,0 +1,605 @@
+// Copyright 2024 KVCache.AI
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+// CXI multi-NIC transfer test
+//
+// Target node: allocate buffers, register on all NICs, wait.
+// Initiator node: allocate receive buffer, register, pull from target.
+//
+// Usage:
+//   # Target (holds KV cache):
+//   ./cxi_transfer_test --mode target --server <target_ip>:12345 \
+//       --num_bufs 4 --buf_size_gb 1
+//
+//   # Initiator (pulls data):
+//   ./cxi_transfer_test --mode initiator --server <initiator_ip>:12346 \
+//       --target <target_ip>:12345 --num_bufs 4 --buf_size_gb 1 \
+//       --transfer_mb 368
+
+#include <gflags/gflags.h>
+#include <glog/logging.h>
+#include <signal.h>
+#include <sys/mman.h>
+
+#include <algorithm>
+#include <atomic>
+#include <chrono>
+#include <cstdint>
+#include <cstring>
+#include <numeric>
+#include <string>
+#include <thread>
+#include <vector>
+#include <numaif.h>
+
+#include "transfer_engine.h"
+#include "transport/cxi_transport/cxi_transport.h"
+
+#if defined(USE_CUDA) || defined(USE_HIP)
+#include <cuda_alike.h>
+#endif
+
+using namespace mooncake;
+
+DEFINE_string(mode, "target", "Running mode: target or initiator");
+DEFINE_string(server, "", "Local server name, e.g. 172.31.6.162:12345");
+DEFINE_string(target, "", "Target server name (initiator mode only)");
+DEFINE_string(metadata, "P2PHANDSHAKE", "Metadata server");
+DEFINE_int32(num_bufs, 200, "Number of buffers to allocate and register");
+DEFINE_double(buf_size_gb, 2.0, "Size of each buffer in GB");
+DEFINE_double(transfer_mb, 368.0,
+              "Transfer size per iteration in MB (initiator)");
+DEFINE_int32(iterations, 50, "Number of benchmark iterations");
+DEFINE_int32(warmup, 5, "Number of warmup iterations");
+DEFINE_int32(batch_size, 1, "Batch size for each transfer submission");
+DEFINE_int32(threads, 1, "Number of initiator worker threads");
+DEFINE_uint64(block_size, 65536, "Block size for transfer requests (64KB)");
+DEFINE_bool(use_device, false, "Use CUDA device memory for transfer");
+
+static std::atomic<bool> g_running(true);
+
+static void signalHandler(int) {
+    g_running.store(false, std::memory_order_relaxed);
+}
+
+static void setupSignalHandler() {
+    struct sigaction sa;
+    sa.sa_handler = signalHandler;
+    sigemptyset(&sa.sa_mask);
+    sa.sa_flags = 0;
+    sigaction(SIGINT, &sa, nullptr);
+    sigaction(SIGTERM, &sa, nullptr);
+}
+
+static void* allocateHugepage(size_t size, int node_id) {
+    void* buf = mmap(nullptr, size, PROT_READ | PROT_WRITE,
+                     MAP_PRIVATE | MAP_ANONYMOUS | MAP_HUGETLB, -1, 0);
+    if (buf == MAP_FAILED) {
+        LOG(WARNING) << "Hugepage mmap failed (" << strerror(errno)
+                     << "), falling back to regular pages";
+        buf = mmap(nullptr, size, PROT_READ | PROT_WRITE,
+                   MAP_PRIVATE | MAP_ANONYMOUS, -1, 0);
+        if (buf == MAP_FAILED) return nullptr;
+    }
+    unsigned long nodemask = 1UL << node_id;
+
+    mbind(buf, size, MPOL_BIND, &nodemask,
+          /* maxnode */ sizeof(nodemask) * 8, 0);
+    return buf;
+}
+
+#if defined(USE_CUDA) || defined(USE_HIP)
+static constexpr bool can_use_device = true;
+
+static void* allocateDevice(size_t size, int device) {
+    cudaSetDevice(device);
+    void* buf = nullptr;
+    cudaError_t err = cudaMalloc(&buf, size);
+    if (err != cudaSuccess) {
+        return nullptr;
+    }
+    return buf;
+}
+
+static int setDeviceMem(void* buf, size_t size, int device, int val) {
+    cudaSetDevice(device);
+    cudaError_t err = cudaMemset(buf, val, size);
+    if (err != cudaSuccess) {
+        return -1;
+    }
+    return 0;
+}
+
+static bool checkDeviceMem(void* buf, int device, uint8_t expected) {
+    cudaSetDevice(device);
+    uint8_t actual;
+    size_t offset = 42;
+    cudaError_t err =
+        cudaMemcpy(&actual, buf + offset, 1, cudaMemcpyDeviceToHost);
+    if (err != cudaSuccess) return false;
+    return expected == actual;
+}
+
+static int freeDeviceMem(void* buf, int device) {
+    cudaSetDevice(device);
+    cudaError_t err = cudaFree(buf);
+    if (err != cudaSuccess) return -1;
+    return 0;
+}
+
+#else
+static constexpr bool can_use_device = false;
+#endif
+
+static int runTarget(TransferEngine* engine) {
+    size_t buf_bytes =
+        static_cast<size_t>(FLAGS_buf_size_gb * 1024 * 1024 * 1024);
+    int num_bufs = FLAGS_num_bufs;
+
+    LOG(INFO) << "=== Target Node ===";
+    LOG(INFO) << "Registering " << num_bufs << " x " << FLAGS_buf_size_gb
+              << " GB = " << num_bufs * FLAGS_buf_size_gb << " GB";
+
+    // Allocate buffers
+    std::vector<void*> bufs;
+    bufs.reserve(num_bufs);
+    LOG(INFO) << "Allocating " << num_bufs << " buffers...";
+    for (int i = 0; i < num_bufs; ++i) {
+        void* buf;
+        if (!FLAGS_use_device) {
+            buf = allocateHugepage(buf_bytes, 0);
+            if (!buf) {
+                LOG(ERROR) << "Allocation failed at buffer " << i;
+                for (auto* p : bufs) munmap(p, buf_bytes);
+                return 1;
+            }
+            buf = memset(buf, (uint8_t)42, buf_bytes);
+            if (!buf) {
+                LOG(ERROR) << "Memset failed at buffer " << i;
+                for (auto* p : bufs) munmap(p, buf_bytes);
+                return 1;
+            }
+        } else {
+#if defined(USE_CUDA) || defined(USE_HIP)
+            int buf_device = i % 4;
+            buf = allocateDevice(buf_bytes, buf_device);
+            if (!buf) {
+                LOG(ERROR) << "allocation failed at buffer " << i;
+                for (int j = 0; j < i; j++) {
+                    freeDeviceMem(bufs[j], j % 4);
+                }
+                return 1;
+            }
+            if (setDeviceMem(buf, buf_bytes, buf_device, 42) != 0) {
+                LOG(ERROR) << "failed to set memory for buffer " << i;
+                for (int j = 0; j < i; j++) {
+                    freeDeviceMem(bufs[j], j % 4);
+                }
+                return 1;
+            }
+#endif
+        }
+
+        bufs.push_back(buf);
+        if ((i + 1) % 50 == 0 || i == num_bufs - 1)
+            LOG(INFO) << "  allocated " << (i + 1) << "/" << num_bufs;
+    }
+
+    // Register all buffers
+    LOG(INFO) << "Registering " << num_bufs << " buffers on all NICs...";
+    auto t0 = std::chrono::steady_clock::now();
+    for (int i = 0; i < num_bufs; ++i) {
+        int ret = engine->registerLocalMemory(bufs[i], buf_bytes, "*", true);
+        if (ret != 0) {
+            LOG(ERROR) << "registerLocalMemory failed at buffer " << i
+                       << ": ret=" << ret;
+            for (int j = 0; j < i; ++j) engine->unregisterLocalMemory(bufs[j]);
+            for (auto* p : bufs) munmap(p, buf_bytes);
+            return 1;
+        }
+        if ((i + 1) % 50 == 0 || i == num_bufs - 1) {
+            auto now = std::chrono::steady_clock::now();
+            double elapsed = std::chrono::duration<double>(now - t0).count();
+            LOG(INFO) << "  registered " << (i + 1) << "/" << num_bufs << " ("
+                      << elapsed << "s)";
+        }
+    }
+    auto t1 = std::chrono::steady_clock::now();
+    LOG(INFO) << "Registration complete: "
+              << std::chrono::duration<double>(t1 - t0).count() << "s";
+
+    LOG(INFO) << "Target ready. First buffer at " << bufs[0]
+              << ". Waiting for initiator (Ctrl+C to stop)...";
+
+    if (FLAGS_use_device) cudaSetDevice(0);
+    while (g_running) sleep(1);
+
+    LOG(INFO) << "Shutting down target...";
+    for (auto* p : bufs) {
+        engine->unregisterLocalMemory(p);
+        if (!FLAGS_use_device) {
+            munmap(p, buf_bytes);
+        } else {
+#if defined(USE_CUDA) || defined(USE_HIP)
+            freeDeviceMem(p, buf_bytes);
+#endif
+        }
+    }
+    return 0;
+}
+
+struct LatencyStats {
+    double avg_ms;
+    double p50_ms;
+    double p99_ms;
+    double throughput_gbs;
+};
+
+static LatencyStats computeStats(std::vector<double>& latencies_ms,
+                                 size_t transfer_bytes) {
+    std::sort(latencies_ms.begin(), latencies_ms.end());
+    double sum = std::accumulate(latencies_ms.begin(), latencies_ms.end(), 0.0);
+    size_t n = latencies_ms.size();
+    LatencyStats stats;
+    stats.avg_ms = sum / n;
+    stats.p50_ms = latencies_ms[n / 2];
+    stats.p99_ms = latencies_ms[std::min(n - 1, (size_t)(n * 0.99))];
+    stats.throughput_gbs = (transfer_bytes / 1e9) / (stats.p50_ms / 1000.0);
+    return stats;
+}
+
+static int runInitiator(TransferEngine* engine) {
+    size_t transfer_bytes =
+        static_cast<size_t>(FLAGS_transfer_mb * 1024 * 1024);
+
+    LOG(INFO) << "=== Initiator Node ===";
+    LOG(INFO) << "Target: " << FLAGS_target;
+    LOG(INFO) << "Transfer size: " << FLAGS_transfer_mb << " MB";
+    LOG(INFO) << "Threads: " << FLAGS_threads;
+
+    if (FLAGS_target.empty()) {
+        LOG(ERROR) << "--target required in initiator mode";
+        return 1;
+    }
+
+    // Allocate local receive buffer (one per thread)
+    size_t recv_bytes = transfer_bytes;
+    std::vector<void*> recv_bufs(FLAGS_threads);
+    for (int t = 0; t < FLAGS_threads; ++t) {
+        if (!FLAGS_use_device) {
+            recv_bufs[t] = allocateHugepage(recv_bytes, 0);
+            if (!recv_bufs[t]) {
+                LOG(ERROR) << "Failed to allocate receive buffer for thread "
+                           << t;
+                return 1;
+            }
+            recv_bufs[t] = memset(recv_bufs[t], 66, recv_bytes);
+            if (!recv_bufs[t]) {
+                LOG(ERROR) << "Failed to memset receive buffer for thread "
+                           << t;
+                return 1;
+            }
+        } else {
+#if defined(USE_CUDA) || defined(USE_HIP)
+            int tid_device = t % 4;
+            recv_bufs[t] = allocateDevice(recv_bytes, tid_device);
+            if (!recv_bufs[t]) {
+                LOG(ERROR) << "Failed to allocate receive buffer for thread "
+                           << t;
+                return 1;
+            }
+            if (setDeviceMem(recv_bufs[t], recv_bytes, tid_device, 66) != 0) {
+                LOG(ERROR) << "Failed to set data on receive buffer for thread "
+                           << t;
+            }
+#endif
+        }
+
+        int ret =
+            engine->registerLocalMemory(recv_bufs[t], recv_bytes, "*", true);
+        if (ret != 0) {
+            LOG(ERROR) << "Failed to register receive buffer: " << ret;
+            return 1;
+        }
+    }
+    LOG(INFO) << "Allocated and registered " << FLAGS_threads
+              << " receive buffers of " << recv_bytes / 1e6 << " MB each";
+
+    // Open remote segment
+    auto segment_id = engine->openSegment(FLAGS_target);
+    if (segment_id < 0) {
+        LOG(ERROR) << "openSegment failed for " << FLAGS_target;
+        return 1;
+    }
+
+    // Get remote buffer info
+    auto segment_desc = engine->getMetadata()->getSegmentDescByID(segment_id);
+    if (!segment_desc || segment_desc->buffers.empty()) {
+        LOG(ERROR) << "No remote buffers found";
+        return 1;
+    }
+    size_t num_remote_bufs = segment_desc->buffers.size();
+    LOG(INFO) << "Remote has " << num_remote_bufs << " buffers";
+    LOG(INFO) << "First buffer: addr=0x" << std::hex
+              << segment_desc->buffers[0].addr << std::dec
+              << " size=" << segment_desc->buffers[0].length;
+
+    for (auto buffer : segment_desc->buffers) {
+        for (int i = 0; i <= 4; i++) {
+            LOG(INFO) << "retry " << i << " selected device: "
+                      << segment_desc->topology.selectDevice(buffer.name, i);
+        }
+        for (int i = 0; i < buffer.rkey.size(); i++) {
+            LOG(INFO) << "remote device " << i << " has rkey " << buffer.rkey[i]
+                      << "\n";
+        }
+    }
+
+    // Connection warmup: small transfer to establish endpoints
+    LOG(INFO) << "Warming up connection...";
+    {
+        size_t warmup_size = std::min(transfer_bytes, (size_t)(64 * 1024));
+        auto batch_id = engine->allocateBatchID(1);
+        TransferRequest req;
+        req.opcode = TransferRequest::READ;
+        req.source = (uint8_t*)recv_bufs[0];
+        req.target_id = segment_id;
+        req.target_offset = segment_desc->buffers[0].addr;
+        req.length = warmup_size;
+        auto s = engine->submitTransfer(batch_id, {req});
+        if (!s.ok()) {
+            LOG(ERROR) << "Warmup transfer failed: " << s.ToString();
+            return 1;
+        }
+        while (true) {
+            TransferStatus status;
+            engine->getTransferStatus(batch_id, 0, status);
+            if (status.s == TransferStatusEnum::COMPLETED) break;
+            if (status.s == TransferStatusEnum::FAILED) {
+                LOG(ERROR) << "Warmup transfer FAILED";
+                return 1;
+            }
+        }
+        if (recv_bufs[0]) engine->freeBatchID(batch_id);
+    }
+    LOG(INFO) << "Connection ready.";
+
+    // Refresh segment desc after connection warmup (endpoints are now up)
+    engine->syncSegmentCache(FLAGS_target);
+    segment_desc = engine->getMetadata()->getSegmentDescByID(segment_id);
+
+    // Worker function: each thread runs its own transfer loop
+    struct ThreadResult {
+        std::vector<double> latencies;
+        int errors = 0;
+    };
+
+    auto workerFn = [&](int tid, int warmup_iters, int bench_iters,
+                        ThreadResult* result) {
+        void* my_recv = recv_bufs[tid];
+        size_t buf_idx =
+            (tid + 1) % num_remote_bufs;  // one thread <-> one buffer
+        for (int w = 0; w < warmup_iters; ++w) {
+            uint64_t raddr = segment_desc->buffers[buf_idx].addr;
+            size_t rlen = segment_desc->buffers[buf_idx].length;
+            size_t xfer = std::min(transfer_bytes, rlen);
+
+            auto bid = engine->allocateBatchID(1);
+            TransferRequest req;
+            req.opcode = TransferRequest::READ;
+            req.source = (uint8_t*)my_recv;
+            req.target_id = segment_id;
+            req.target_offset = raddr;
+            req.length = xfer;
+
+            engine->submitTransfer(bid, {req});
+            while (true) {
+                TransferStatus st;
+                engine->getTransferStatus(bid, 0, st);
+                if (st.s == TransferStatusEnum::COMPLETED ||
+                    st.s == TransferStatusEnum::FAILED)
+                    break;
+            }
+            engine->freeBatchID(bid);
+        }
+
+        for (int i = 0; i < bench_iters; ++i) {
+            uint64_t raddr = segment_desc->buffers[buf_idx].addr;
+            size_t rlen = segment_desc->buffers[buf_idx].length;
+            size_t xfer = std::min(transfer_bytes, rlen);
+
+            auto t0 = std::chrono::steady_clock::now();
+            auto bid = engine->allocateBatchID(1);
+            TransferRequest req;
+            req.opcode = TransferRequest::READ;
+            req.source = (uint8_t*)my_recv;
+            req.target_id = segment_id;
+            req.target_offset = raddr;
+            req.length = xfer;
+            auto s = engine->submitTransfer(bid, {req});
+            if (!s.ok()) {
+                result->errors++;
+                engine->freeBatchID(bid);
+                continue;
+            }
+            bool ok = false;
+            while (true) {
+                TransferStatus st;
+                engine->getTransferStatus(bid, 0, st);
+                if (st.s == TransferStatusEnum::COMPLETED) {
+                    ok = true;
+                    break;
+                }
+                if (st.s == TransferStatusEnum::FAILED) {
+                    LOG(ERROR)
+                        << "thread " << tid << " has failed the transfer\n";
+                    result->errors++;
+                    break;
+                }
+            }
+            engine->freeBatchID(bid);
+            if (ok) {
+                auto t1 = std::chrono::steady_clock::now();
+                result->latencies.push_back(
+                    std::chrono::duration<double, std::milli>(t1 - t0).count());
+            }
+        }
+    };
+
+    // Run warmup + benchmark with threads
+    int num_threads = FLAGS_threads;
+    LOG(INFO) << "Running with " << num_threads << " threads, " << FLAGS_warmup
+              << " warmup + " << FLAGS_iterations
+              << " bench iterations per thread...";
+
+    std::vector<ThreadResult> results(num_threads);
+    std::vector<std::thread> threads;
+
+    auto wall_t0 = std::chrono::steady_clock::now();
+    for (int t = 0; t < num_threads; ++t) {
+        threads.emplace_back(workerFn, t, FLAGS_warmup, FLAGS_iterations,
+                             &results[t]);
+    }
+    for (auto& th : threads) th.join();
+    auto wall_t1 = std::chrono::steady_clock::now();
+    double wall_ms =
+        std::chrono::duration<double, std::milli>(wall_t1 - wall_t0).count();
+
+    // Aggregate results
+    std::vector<double> all_latencies;
+    int total_errors = 0;
+    for (auto& r : results) {
+        all_latencies.insert(all_latencies.end(), r.latencies.begin(),
+                             r.latencies.end());
+        total_errors += r.errors;
+    }
+
+    if (all_latencies.empty()) {
+        LOG(ERROR) << "All transfers failed";
+        return 1;
+    }
+
+    auto stats = computeStats(all_latencies, transfer_bytes);
+    size_t total_xfers = all_latencies.size();
+    double total_bytes = (double)total_xfers * transfer_bytes;
+    double agg_throughput = total_bytes / 1e9 / (wall_ms / 1000.0);
+
+    LOG(INFO) << "=== Results (" << num_threads << " threads) ===";
+    LOG(INFO) << "Transfer: " << FLAGS_transfer_mb << " MB x " << total_xfers
+              << " transfers";
+    LOG(INFO) << "Wall time: " << wall_ms << " ms";
+    LOG(INFO) << "Per-transfer p50: " << stats.p50_ms << " ms"
+              << "  p99: " << stats.p99_ms << " ms";
+    LOG(INFO) << "Per-transfer throughput: " << stats.throughput_gbs << " GB/s";
+    LOG(INFO) << "Aggregate throughput: " << agg_throughput << " GB/s";
+    LOG(INFO) << "Errors: " << total_errors;
+
+    // Per-thread stats
+    for (int t = 0; t < num_threads; ++t) {
+        if (results[t].latencies.empty()) continue;
+        auto ts = computeStats(results[t].latencies, transfer_bytes);
+        LOG(INFO) << "  Thread " << t << ": p50=" << ts.p50_ms
+                  << "ms  tput=" << ts.throughput_gbs << " GB/s"
+                  << "  iters=" << results[t].latencies.size()
+                  << "  errors=" << results[t].errors;
+    }
+
+    // Cleanup
+    for (int t = 0; t < FLAGS_threads; ++t) {
+        if (!FLAGS_use_device) {
+            if (((uint8_t*)recv_bufs[t])[66] != 42) {
+                LOG(ERROR) << "transfer has corrupted data, expected to find "
+                              "42 but found "
+                           << ((int)((uint8_t*)recv_bufs[t])[66]);
+            } else {
+                LOG(INFO) << "transfer was successful";
+            }
+            engine->unregisterLocalMemory(recv_bufs[t]);
+            munmap(recv_bufs[t], recv_bytes);
+        } else {
+#if defined(USE_CUDA) || defined(USE_HIP)
+            bool result = checkDeviceMem(recv_bufs[t], t % 4, 42);
+            if (!result) {
+                LOG(ERROR)
+                    << "transfer has corrupted data, expected to find 42";
+            } else {
+                LOG(INFO) << "transfer was successful\n";
+            }
+            engine->unregisterLocalMemory(recv_bufs[t]);
+            freeDeviceMem(recv_bufs[t], t % 4);
+#endif
+        }
+    }
+    return total_errors > 0 ? 1 : 0;
+}
+
+int main(int argc, char** argv) {
+    google::InitGoogleLogging(argv[0]);
+    gflags::ParseCommandLineFlags(&argc, &argv, true);
+    FLAGS_logtostderr = 1;
+
+    setupSignalHandler();
+
+    if (FLAGS_use_device && !can_use_device) {
+        LOG(ERROR)
+            << "Cannot use device memory, recompile with cmake -DUSE_CUDA!";
+        return 1;
+    }
+
+    if (FLAGS_use_device) {
+        LOG(INFO) << "Using device memory";
+    }
+
+    if (FLAGS_server.empty()) {
+        LOG(ERROR) << "--server is required";
+        return 1;
+    }
+
+    // Parse host:port
+    auto colon = FLAGS_server.rfind(':');
+    std::string host = FLAGS_server.substr(0, colon);
+    uint64_t port = 12345;
+    if (colon != std::string::npos)
+        port = std::stoull(FLAGS_server.substr(colon + 1));
+
+    auto engine = std::make_unique<TransferEngine>(false);
+    int ret = engine->init(FLAGS_metadata, FLAGS_server, host, port);
+    if (ret != 0) {
+        LOG(ERROR) << "Engine init failed: " << ret;
+        return 1;
+    }
+
+    // Discover topology and install CXI transport (all NICs)
+    engine->getLocalTopology()->discover({});
+    auto* xport = engine->installTransport("cxi", nullptr);
+    if (!xport) {
+        LOG(ERROR) << "installTransport(cxi) failed";
+        return 1;
+    }
+
+    std::string actual_server = engine->getLocalIpAndPort();
+    LOG(INFO) << "Actual server name (use this for --target): "
+              << actual_server;
+
+    if (FLAGS_mode == "target") {
+        ret = runTarget(engine.get());
+    } else if (FLAGS_mode == "initiator") {
+        ret = runInitiator(engine.get());
+    } else {
+        LOG(ERROR) << "Unknown mode: " << FLAGS_mode;
+        ret = 1;
+    }
+
+    return ret;
+}
diff --git a/mooncake-transfer-engine/tests/cxi_transport_test.cpp b/mooncake-transfer-engine/tests/cxi_transport_test.cpp
new file mode 100644
index 00000000..a8cb4cc0
--- /dev/null
+++ b/mooncake-transfer-engine/tests/cxi_transport_test.cpp
@@ -0,0 +1,640 @@
+// Copyright 2024 KVCache.AI
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include <gflags/gflags.h>
+#include <glog/logging.h>
+#include <gtest/gtest.h>
+#include <sys/time.h>
+
+#include <cstdlib>
+#include <memory>
+
+#include "transfer_engine.h"
+#include "transport/cxi_transport/cxi_transport.h"
+#include "transport/transport.h"
+#include "cuda_alike.h"
+
+#if defined(USE_CUDA) || defined(USE_HIP)
+#define USE_GPU
+#endif
+
+using namespace mooncake;
+
+namespace mooncake {
+
+static void *allocateMemoryPool(size_t size, int socket_id) {
+    return numa_alloc_onnode(size, socket_id);
+}
+
+static void freeMemoryPool(void *addr, size_t size) { numa_free(addr, size); }
+
+#ifdef USE_GPU
+static void *allocateMemoryPoolDevice(size_t size, int device) {
+    void *devPtr = nullptr;
+    cudaError_t err = cudaMalloc(&devPtr, size);
+    cudaDeviceSynchronize();
+    if (err != cudaSuccess) {
+        LOG(ERROR) << "failed to alloc " << size << " bytes on cuda device "
+                   << device;
+        return nullptr;
+    }
+    return devPtr;
+}
+
+static void freeMemoryPoolDevice(void *addr) {
+    cudaError_t err = cudaFree(addr);
+    if (err != cudaSuccess) {
+        LOG(ERROR) << "failed to free cuda memory @ " << addr;
+    }
+}
+#endif
+
+// ---------------------------------------------------------------------------
+// CXI Transport Test Fixture
+//
+// This test uses the P2PHANDSHAKE metadata backend and performs loopback
+// transfers (local_server_name == segment_id), similar to the TCP transport
+// tests. It requires CXI hardware to be present (fi_info -p cxi must succeed).
+//
+// Environment variables:
+//   MC_METADATA_SERVER  - metadata backend (default: P2PHANDSHAKE)
+//   MC_LOCAL_SERVER_NAME - local server name (default: 127.0.0.1:12345)
+// ---------------------------------------------------------------------------
+class CXITransportTest : public ::testing::Test {
+   protected:
+    void SetUp() override {
+        google::InitGoogleLogging("CXITransportTest");
+        FLAGS_logtostderr = 1;
+
+        const char *env = std::getenv("MC_METADATA_SERVER");
+        metadata_server_ = env ? env : "P2PHANDSHAKE";
+        LOG(INFO) << "metadata_server: " << metadata_server_;
+
+        env = std::getenv("MC_LOCAL_SERVER_NAME");
+        local_server_name_ = env ? env : "127.0.0.1:12345";
+        LOG(INFO) << "local_server_name: " << local_server_name_;
+    }
+
+    void TearDown() override { google::ShutdownGoogleLogging(); }
+
+    // Helper: create engine, install CXI transport, register memory
+    struct EngineSetup {
+        std::unique_ptr<TransferEngine> engine;
+        Transport *xport;
+        void *addr;
+        size_t buffer_size;
+        SegmentID segment_id;
+    };
+
+#ifdef USE_GPU
+
+    EngineSetup createEngineDevice(size_t buffer_size = 1ull << 30) {
+        EngineSetup s;
+        s.buffer_size = buffer_size;
+
+        s.engine = std::make_unique<TransferEngine>(false);
+        // Manually discover topology to populate CXI device list
+        // (same pattern as the Python binding in transfer_engine_py.cpp)
+        s.engine->getLocalTopology()->discover({});
+        auto hp = parseHostNameWithPort(local_server_name_);
+        int rc = s.engine->init(metadata_server_, local_server_name_,
+                                hp.first.c_str(), hp.second);
+        EXPECT_EQ(rc, 0) << "engine->init failed";
+
+        s.xport = s.engine->installTransport("cxi", nullptr);
+        EXPECT_NE(s.xport, nullptr) << "installTransport(\"cxi\") failed";
+
+        s.addr = allocateMemoryPoolDevice(buffer_size,
+                                          0);  // allocate mempool on cuda:0
+        EXPECT_NE(s.addr, nullptr) << "allocateMemoryPool failed";
+
+        rc = s.engine->registerLocalMemory(s.addr, buffer_size,
+                                           GPU_PREFIX + "0");
+        EXPECT_EQ(rc, 0) << "registerLocalMemory failed";
+
+        // Use actual RPC address (P2PHANDSHAKE picks a random port)
+        auto actual_addr = s.engine->getLocalIpAndPort();
+        s.segment_id = s.engine->openSegment(actual_addr);
+        return s;
+    }
+
+    void destroyEngineDevice(EngineSetup &s) {
+        if (s.engine && s.addr) {
+            s.engine->unregisterLocalMemory(s.addr);
+        }
+        if (s.addr) {
+            freeMemoryPoolDevice(s.addr);
+            s.addr = nullptr;
+        }
+    }
+
+#endif
+
+    EngineSetup createEngine(size_t buffer_size = 1ull << 30) {
+        EngineSetup s;
+        s.buffer_size = buffer_size;
+
+        s.engine = std::make_unique<TransferEngine>(false);
+        // Manually discover topology to populate CXI device list
+        // (same pattern as the Python binding in transfer_engine_py.cpp)
+        s.engine->getLocalTopology()->discover({});
+        auto hp = parseHostNameWithPort(local_server_name_);
+        int rc = s.engine->init(metadata_server_, local_server_name_,
+                                hp.first.c_str(), hp.second);
+        EXPECT_EQ(rc, 0) << "engine->init failed";
+
+        s.xport = s.engine->installTransport("cxi", nullptr);
+        EXPECT_NE(s.xport, nullptr) << "installTransport(\"cxi\") failed";
+
+        s.addr = allocateMemoryPool(buffer_size, 0);
+        EXPECT_NE(s.addr, nullptr) << "allocateMemoryPool failed";
+
+        rc = s.engine->registerLocalMemory(s.addr, buffer_size, "cpu:0");
+        EXPECT_EQ(rc, 0) << "registerLocalMemory failed";
+
+        // Use actual RPC address (P2PHANDSHAKE picks a random port)
+        auto actual_addr = s.engine->getLocalIpAndPort();
+        s.segment_id = s.engine->openSegment(actual_addr);
+        return s;
+    }
+
+    void destroyEngine(EngineSetup &s) {
+        if (s.engine && s.addr) {
+            s.engine->unregisterLocalMemory(s.addr);
+        }
+        if (s.addr) {
+            freeMemoryPool(s.addr, s.buffer_size);
+            s.addr = nullptr;
+        }
+    }
+
+    // Helper: submit a single transfer and poll until completion
+    bool submitAndWait(TransferEngine *engine, SegmentID segment_id,
+                       void *source, uint64_t target_offset, size_t length,
+                       TransferRequest::OpCode opcode) {
+        auto batch_id = engine->allocateBatchID(1);
+
+        TransferRequest entry;
+        entry.opcode = opcode;
+        entry.length = length;
+        entry.source = (uint8_t *)source;
+        entry.target_id = segment_id;
+        entry.target_offset = target_offset;
+
+        Status s = engine->submitTransfer(batch_id, {entry});
+        if (!s.ok()) {
+            LOG(ERROR) << "submitTransfer failed: " << s.ToString();
+            engine->freeBatchID(batch_id);
+            return false;
+        }
+
+        // Poll for completion with timeout
+        const int kMaxPollIterations = 1000000;
+        TransferStatus status;
+        for (int i = 0; i < kMaxPollIterations; ++i) {
+            s = engine->getTransferStatus(batch_id, 0, status);
+            if (!s.ok()) {
+                LOG(ERROR) << "getTransferStatus failed: " << s.ToString();
+                engine->freeBatchID(batch_id);
+                return false;
+            }
+            if (status.s == TransferStatusEnum::COMPLETED) {
+                engine->freeBatchID(batch_id);
+                return true;
+            }
+            if (status.s == TransferStatusEnum::FAILED) {
+                LOG(ERROR) << "Transfer FAILED";
+                engine->freeBatchID(batch_id);
+                return false;
+            }
+        }
+        LOG(ERROR) << "Transfer timed out";
+        engine->freeBatchID(batch_id);
+        return false;
+    }
+
+    std::string metadata_server_;
+    std::string local_server_name_;
+};
+
+// Test 1: Verify CXI transport can be installed
+TEST_F(CXITransportTest, InstallTransport) {
+    auto engine = std::make_unique<TransferEngine>(false);
+    engine->getLocalTopology()->discover({});
+    auto hp = parseHostNameWithPort(local_server_name_);
+    int rc = engine->init(metadata_server_, local_server_name_,
+                          hp.first.c_str(), hp.second);
+    ASSERT_EQ(rc, 0);
+
+    Transport *xport = engine->installTransport("cxi", nullptr);
+    ASSERT_NE(xport, nullptr)
+        << "CXI transport should be installable on CXI hardware";
+}
+
+// Test 2: Basic loopback write
+TEST_F(CXITransportTest, LoopbackWrite) {
+    auto setup = createEngine();
+
+    auto segment_desc =
+        setup.engine->getMetadata()->getSegmentDescByID(setup.segment_id);
+    ASSERT_NE(segment_desc, nullptr);
+    uint64_t remote_base = (uint64_t)segment_desc->buffers[0].addr;
+
+    const size_t kDataLength = 4096;
+
+    // Fill source buffer with known data
+    memset(setup.addr, 0xAB, kDataLength);
+
+    bool ok = submitAndWait(setup.engine.get(), setup.segment_id, setup.addr,
+                            remote_base, kDataLength, TransferRequest::WRITE);
+    EXPECT_TRUE(ok) << "Loopback write should succeed";
+
+    destroyEngine(setup);
+}
+
+#ifdef USE_GPU
+TEST_F(CXITransportTest, LoopbackWriteDevice) {
+    cudaSetDevice(0);
+    auto setup = createEngineDevice();
+
+    auto segment_desc =
+        setup.engine->getMetadata()->getSegmentDescByID(setup.segment_id);
+    ASSERT_NE(segment_desc, nullptr);
+    uint64_t remote_base = (uint64_t)segment_desc->buffers[0].addr;
+
+    const size_t kDataLength = 4096;
+
+    // Fill source buffer with known data
+    cudaMemset(setup.addr, 0xAB, kDataLength);
+    cudaDeviceSynchronize();
+
+    bool ok = submitAndWait(setup.engine.get(), setup.segment_id, setup.addr,
+                            remote_base, kDataLength, TransferRequest::WRITE);
+    EXPECT_TRUE(ok) << "Loopback write should succeed";
+
+    destroyEngineDevice(setup);
+}
+#endif
+
+// Test 3: Write then read, verify data integrity
+TEST_F(CXITransportTest, WriteAndRead) {
+    auto setup = createEngine();
+
+    auto segment_desc =
+        setup.engine->getMetadata()->getSegmentDescByID(setup.segment_id);
+    ASSERT_NE(segment_desc, nullptr);
+    uint64_t remote_base = (uint64_t)segment_desc->buffers[0].addr;
+
+    const size_t kDataLength = 4096000;
+    uint8_t *buf = (uint8_t *)setup.addr;
+
+    // Fill first half with random data
+    for (size_t i = 0; i < kDataLength; ++i) buf[i] = 'a' + lrand48() % 26;
+
+    // Write local -> remote (loopback)
+    bool ok = submitAndWait(setup.engine.get(), setup.segment_id, buf,
+                            remote_base, kDataLength, TransferRequest::WRITE);
+    ASSERT_TRUE(ok) << "Write should succeed";
+
+    // Read remote -> local (into second half of buffer)
+    ok = submitAndWait(setup.engine.get(), setup.segment_id, buf + kDataLength,
+                       remote_base, kDataLength, TransferRequest::READ);
+    ASSERT_TRUE(ok) << "Read should succeed";
+
+    // Verify data integrity
+    EXPECT_EQ(0, memcmp(buf, buf + kDataLength, kDataLength))
+        << "Read-back data should match written data";
+
+    destroyEngine(setup);
+}
+
+#ifdef USE_GPU
+TEST_F(CXITransportTest, WriteAndReadDevice) {
+    cudaSetDevice(0);
+    auto setup = createEngineDevice();
+
+    auto segment_desc =
+        setup.engine->getMetadata()->getSegmentDescByID(setup.segment_id);
+    ASSERT_NE(segment_desc, nullptr);
+    uint64_t remote_base = (uint64_t)segment_desc->buffers[0].addr;
+
+    const size_t kDataLength = 4096000;
+    uint8_t *buf = (uint8_t *)setup.addr;
+
+    uint8_t random_data = 'a' + lrand48() % 26;
+    cudaMemset(buf, random_data, kDataLength);
+    cudaDeviceSynchronize();
+
+    // Write local -> remote (loopback)
+    bool ok = submitAndWait(setup.engine.get(), setup.segment_id, buf,
+                            remote_base, kDataLength, TransferRequest::WRITE);
+    ASSERT_TRUE(ok) << "Write should succeed";
+
+    // Read remote -> local (into second half of buffer)
+    ok = submitAndWait(setup.engine.get(), setup.segment_id, buf + kDataLength,
+                       remote_base, kDataLength, TransferRequest::READ);
+    ASSERT_TRUE(ok) << "Read should succeed";
+    // copy data back to host
+    uint8_t *dataHost = (uint8_t *)allocateMemoryPool(2 * kDataLength, 0);
+    cudaMemcpy(dataHost, buf, 2 * kDataLength, cudaMemcpyDeviceToHost);
+    cudaDeviceSynchronize();
+    // Verify data integrity
+    EXPECT_EQ(0, memcmp(dataHost, dataHost + kDataLength, kDataLength))
+        << "Read-back data should match written data";
+
+    destroyEngineDevice(setup);
+}
+#endif
+
+// Test 4: Multiple sequential writes in a batch
+TEST_F(CXITransportTest, MultiWrite) {
+    auto setup = createEngine();
+
+    auto segment_desc =
+        setup.engine->getMetadata()->getSegmentDescByID(setup.segment_id);
+    ASSERT_NE(segment_desc, nullptr);
+    uint64_t remote_base = (uint64_t)segment_desc->buffers[0].addr;
+
+    const size_t kDataLength = 65536;
+    const int kBatchSize = 16;
+
+    auto batch_id = setup.engine->allocateBatchID(kBatchSize);
+
+    std::vector<TransferRequest> requests;
+    for (int i = 0; i < kBatchSize; ++i) {
+        TransferRequest entry;
+        entry.opcode = TransferRequest::WRITE;
+        entry.length = kDataLength;
+        entry.source = (uint8_t *)setup.addr + i * kDataLength;
+        entry.target_id = setup.segment_id;
+        entry.target_offset = remote_base + i * kDataLength;
+        requests.push_back(entry);
+    }
+
+    Status s = setup.engine->submitTransfer(batch_id, requests);
+    ASSERT_TRUE(s.ok()) << "submitTransfer failed: " << s.ToString();
+
+    // Poll all tasks until completion
+    for (int task_id = 0; task_id < kBatchSize; ++task_id) {
+        TransferStatus status;
+        const int kMaxPollIterations = 1000000;
+        for (int i = 0; i < kMaxPollIterations; ++i) {
+            s = setup.engine->getTransferStatus(batch_id, task_id, status);
+            ASSERT_TRUE(s.ok());
+            if (status.s == TransferStatusEnum::COMPLETED) break;
+            ASSERT_NE(status.s, TransferStatusEnum::FAILED)
+                << "Task " << task_id << " failed";
+        }
+        ASSERT_EQ(status.s, TransferStatusEnum::COMPLETED)
+            << "Task " << task_id << " did not complete";
+    }
+
+    s = setup.engine->freeBatchID(batch_id);
+    ASSERT_TRUE(s.ok());
+
+    destroyEngine(setup);
+}
+
+// Test 5: Stress test - multiple batches to verify no CQ overflow
+TEST_F(CXITransportTest, StressMultipleBatches) {
+    auto setup = createEngine();
+
+    auto segment_desc =
+        setup.engine->getMetadata()->getSegmentDescByID(setup.segment_id);
+    ASSERT_NE(segment_desc, nullptr);
+    uint64_t remote_base = (uint64_t)segment_desc->buffers[0].addr;
+
+    const size_t kDataLength = 65536;
+    const int kBatchSize = 8;
+    const int kNumBatches = 20;
+
+    for (int batch = 0; batch < kNumBatches; ++batch) {
+        auto batch_id = setup.engine->allocateBatchID(kBatchSize);
+
+        std::vector<TransferRequest> requests;
+        for (int i = 0; i < kBatchSize; ++i) {
+            TransferRequest entry;
+            entry.opcode = TransferRequest::WRITE;
+            entry.length = kDataLength;
+            entry.source =
+                (uint8_t *)setup.addr + (i + batch * kBatchSize) * kDataLength;
+            entry.target_id = setup.segment_id;
+            entry.target_offset =
+                remote_base + (i + batch * kBatchSize) * kDataLength;
+            requests.push_back(entry);
+        }
+
+        Status s = setup.engine->submitTransfer(batch_id, requests);
+        ASSERT_TRUE(s.ok())
+            << "Batch " << batch << " submitTransfer failed: " << s.ToString();
+
+        // Wait for all tasks in batch
+        for (int task_id = 0; task_id < kBatchSize; ++task_id) {
+            TransferStatus status;
+            const int kMaxPollIterations = 1000000;
+            for (int i = 0; i < kMaxPollIterations; ++i) {
+                s = setup.engine->getTransferStatus(batch_id, task_id, status);
+                ASSERT_TRUE(s.ok());
+                if (status.s == TransferStatusEnum::COMPLETED) break;
+                ASSERT_NE(status.s, TransferStatusEnum::FAILED)
+                    << "Batch " << batch << " task " << task_id << " failed";
+            }
+            ASSERT_EQ(status.s, TransferStatusEnum::COMPLETED)
+                << "Batch " << batch << " task " << task_id
+                << " did not complete";
+        }
+
+        s = setup.engine->freeBatchID(batch_id);
+        ASSERT_TRUE(s.ok());
+    }
+
+    destroyEngine(setup);
+}
+
+// Test 6: warmupSegment on loopback peer
+//
+// Exercises CxiTransport::warmupSegment() which is the C++ entry point
+// behind the warmup_efa_segment() Python binding / warmupSegment() C API.
+// Loopback is enough to cover the handshake + fi_av_insert path AND the
+// idempotent short-circuit on the second call.
+TEST_F(CXITransportTest, WarmupSegmentLoopback) {
+    auto setup = createEngine();
+
+    auto *cxi = dynamic_cast<CxiTransport *>(setup.xport);
+    ASSERT_NE(cxi, nullptr) << "installTransport did not return a CxiTransport";
+
+    // First call: should connect every (local NIC x peer NIC) pair.
+    int rc = cxi->warmupSegment(setup.engine->getLocalIpAndPort());
+    EXPECT_EQ(rc, 0) << "warmupSegment should succeed on loopback";
+
+    // Second call: should short-circuit (all endpoints already connected).
+    rc = cxi->warmupSegment(setup.engine->getLocalIpAndPort());
+    EXPECT_EQ(rc, 0) << "warmupSegment should be idempotent";
+
+    // Empty / self-name: short-circuit path returning 0 without touching AV.
+    rc = cxi->warmupSegment("");
+    EXPECT_EQ(rc, 0) << "warmupSegment(\"\") should be a no-op";
+
+    destroyEngine(setup);
+}
+
+// Test 7: warmupSegment on a non-existent segment name should fail cleanly
+// (no crash, no hang) rather than blocking for the poll timeout.
+TEST_F(CXITransportTest, WarmupSegmentNotFound) {
+    auto setup = createEngine();
+
+    auto *cxi = dynamic_cast<CxiTransport *>(setup.xport);
+    ASSERT_NE(cxi, nullptr);
+
+    int rc = cxi->warmupSegment("127.0.0.1:1");  // not openSegment'd
+    EXPECT_NE(rc, 0) << "warmupSegment should fail for unknown segment";
+
+    destroyEngine(setup);
+}
+
+// Test 8: registerLocalMemoryBatch / unregisterLocalMemoryBatch round-trip.
+// Covers the batched MR path which the single-buffer tests above never hit.
+TEST_F(CXITransportTest, RegisterMemoryBatch) {
+    auto engine = std::make_unique<TransferEngine>(false);
+    engine->getLocalTopology()->discover({});
+    auto hp = parseHostNameWithPort(local_server_name_);
+    int rc = engine->init(metadata_server_, local_server_name_,
+                          hp.first.c_str(), hp.second);
+    ASSERT_EQ(rc, 0);
+
+    Transport *xport = engine->installTransport("cxi", nullptr);
+    ASSERT_NE(xport, nullptr);
+
+    const size_t kBufSize = 4ull << 20;  // 4 MB each
+    const int kNumBufs = 4;
+    std::vector<void *> addrs;
+    std::vector<BufferEntry> entries;
+    for (int i = 0; i < kNumBufs; ++i) {
+        void *a = allocateMemoryPool(kBufSize, 0);
+        ASSERT_NE(a, nullptr);
+        addrs.push_back(a);
+        entries.push_back({a, kBufSize});
+    }
+
+    rc = engine->registerLocalMemoryBatch(entries, "cpu:0");
+    EXPECT_EQ(rc, 0) << "registerLocalMemoryBatch should succeed";
+
+    rc = engine->unregisterLocalMemoryBatch(addrs);
+    EXPECT_EQ(rc, 0) << "unregisterLocalMemoryBatch should succeed";
+
+    for (void *a : addrs) freeMemoryPool(a, kBufSize);
+}
+
+// Test 9: Larger transfer (64 MB total split into 1 MB slices) to exercise
+// the WR / CQ pacing logic in CxiContext::submitSlicesOnPeer beyond what the
+// 16 x 64 KB MultiWrite test reaches.
+TEST_F(CXITransportTest, LargeTransfer) {
+    const size_t kBufSize = 128ull << 20;  // 128 MB
+    auto setup = createEngine(kBufSize);
+
+    auto segment_desc =
+        setup.engine->getMetadata()->getSegmentDescByID(setup.segment_id);
+    ASSERT_NE(segment_desc, nullptr);
+    uint64_t remote_base = (uint64_t)segment_desc->buffers[0].addr;
+
+    const size_t kSliceLen = 1ull << 20;  // 1 MB per slice
+    const int kNumSlices = 64;            // 64 MB total
+    ASSERT_LE(static_cast<size_t>(kNumSlices) * kSliceLen, kBufSize / 2);
+
+    // Fill first half with known data
+    uint8_t *buf = (uint8_t *)setup.addr;
+    for (size_t i = 0; i < static_cast<size_t>(kNumSlices) * kSliceLen; ++i)
+        buf[i] = (uint8_t)(i & 0xFF);
+
+    auto batch_id = setup.engine->allocateBatchID(kNumSlices);
+    std::vector<TransferRequest> requests;
+    requests.reserve(kNumSlices);
+    for (int i = 0; i < kNumSlices; ++i) {
+        TransferRequest entry;
+        entry.opcode = TransferRequest::WRITE;
+        entry.length = kSliceLen;
+        entry.source = buf + i * kSliceLen;
+        entry.target_id = setup.segment_id;
+        entry.target_offset = remote_base + (kBufSize / 2) + i * kSliceLen;
+        requests.push_back(entry);
+    }
+
+    Status s = setup.engine->submitTransfer(batch_id, requests);
+    ASSERT_TRUE(s.ok()) << "submitTransfer failed: " << s.ToString();
+
+    for (int task_id = 0; task_id < kNumSlices; ++task_id) {
+        TransferStatus status;
+        const int kMaxPollIterations = 2000000;
+        int i = 0;
+        for (; i < kMaxPollIterations; ++i) {
+            s = setup.engine->getTransferStatus(batch_id, task_id, status);
+            ASSERT_TRUE(s.ok());
+            if (status.s == TransferStatusEnum::COMPLETED) break;
+            ASSERT_NE(status.s, TransferStatusEnum::FAILED);
+        }
+        ASSERT_EQ(status.s, TransferStatusEnum::COMPLETED)
+            << "task " << task_id << " did not complete";
+    }
+
+    s = setup.engine->freeBatchID(batch_id);
+    ASSERT_TRUE(s.ok());
+
+    // Verify byte-level integrity of the last slice (spot check).
+    EXPECT_EQ(0, memcmp(buf + (kNumSlices - 1) * kSliceLen,
+                        buf + (kBufSize / 2) + (kNumSlices - 1) * kSliceLen,
+                        kSliceLen));
+
+    destroyEngine(setup);
+}
+
+// Test 10: Repeated open/close of the same remote segment must not leak AV
+// slots or break loopback transfers — this is the setPeerNicPath-detach path
+// that target restarts depend on under the SRD shared-endpoint model.
+TEST_F(CXITransportTest, RepeatedOpenSegment) {
+    auto setup = createEngine();
+
+    auto actual_addr = setup.engine->getLocalIpAndPort();
+
+    // First write via setup.segment_id (from createEngine()).
+    auto segment_desc =
+        setup.engine->getMetadata()->getSegmentDescByID(setup.segment_id);
+    ASSERT_NE(segment_desc, nullptr);
+    uint64_t remote_base = (uint64_t)segment_desc->buffers[0].addr;
+    memset(setup.addr, 0xCD, 4096);
+    EXPECT_TRUE(submitAndWait(setup.engine.get(), setup.segment_id, setup.addr,
+                              remote_base, 4096, TransferRequest::WRITE));
+
+    // Re-open same segment several times; each should return a working handle
+    // and subsequent writes should still succeed.
+    for (int i = 0; i < 5; ++i) {
+        SegmentID sid = setup.engine->openSegment(actual_addr);
+        ASSERT_NE(sid, (SegmentID)-1);
+        auto desc = setup.engine->getMetadata()->getSegmentDescByID(sid);
+        ASSERT_NE(desc, nullptr);
+        uint64_t base = (uint64_t)desc->buffers[0].addr;
+        ASSERT_TRUE(submitAndWait(setup.engine.get(), sid, setup.addr, base,
+                                  4096, TransferRequest::WRITE))
+            << "write #" << i << " after re-open failed";
+    }
+
+    destroyEngine(setup);
+}
+
+}  // namespace mooncake
+
+int main(int argc, char **argv) {
+    gflags::ParseCommandLineFlags(&argc, &argv, false);
+    ::testing::InitGoogleTest(&argc, argv);
+    return RUN_ALL_TESTS();
+}
diff --git a/mooncake-transfer-engine/tests/cxi_unit_tests.cpp b/mooncake-transfer-engine/tests/cxi_unit_tests.cpp
new file mode 100644
index 00000000..d5e954c1
--- /dev/null
+++ b/mooncake-transfer-engine/tests/cxi_unit_tests.cpp
@@ -0,0 +1,124 @@
+#include <gflags/gflags.h>
+#include <glog/logging.h>
+#include <gtest/gtest.h>
+#include <sys/time.h>
+
+#include <cstdlib>
+#include <fstream>
+#include <iomanip>
+#include <memory>
+#include <unistd.h>
+#include <fcntl.h>
+#include <sys/mman.h>
+
+#include "transfer_engine.h"
+#include "transport/transport.h"
+#include "transport/cxi_transport/cxi_transport.h"
+#include "common.h"
+
+using namespace mooncake;
+
+namespace mooncake {
+
+class CxiTransportUnitTest : public ::testing::Test {
+   protected:
+    void SetUp() override {
+        metadata_server_ = "P2PHANDSHAKE";
+        LOG(INFO) << "metadata_server: " << metadata_server_;
+
+        local_server_name_ = "127.0.0.1:12345";
+        LOG(INFO) << "local_server_name: " << local_server_name_;
+    }
+
+    void TearDown() override {}
+
+    std::string metadata_server_;
+    std::string local_server_name_;
+};
+
+TEST_F(CxiTransportUnitTest, TestInit) {
+    std::unique_ptr<TransferEngine> engine;
+    std::string device_filter = "cxi0";
+    engine = std::make_unique<TransferEngine>(
+        false, std::vector<std::string>{device_filter});
+    EXPECT_NE(engine, nullptr) << "failed to allocate TransferEngine";
+
+    engine->getLocalTopology()->discover({});
+    auto hp = parseHostNameWithPort(local_server_name_);
+    int rc = engine->init(metadata_server_, local_server_name_,
+                          hp.first.c_str(), hp.second);
+    EXPECT_EQ(rc, 0) << "engine->init failed";
+
+    Transport* xport = engine->installTransport("cxi", nullptr);
+    EXPECT_NE(xport, nullptr) << "installTransport(\"cxi\") failed";
+    CxiTransport* cxi_xport = dynamic_cast<CxiTransport*>(xport);
+    EXPECT_NE(cxi_xport, nullptr) << "initialization for cxi failed";
+
+    engine.reset();
+
+    LOG(INFO) << "default init ok";
+
+    device_filter = "invalid_device";  // dummy device name
+    engine = std::make_unique<TransferEngine>(
+        false, std::vector<std::string>{device_filter});
+    EXPECT_NE(engine, nullptr) << "failed to allocate TransferEngine";
+    engine->getLocalTopology()->discover({device_filter});
+    hp = parseHostNameWithPort(local_server_name_);
+    rc = engine->init(metadata_server_, local_server_name_, hp.first.c_str(),
+                      hp.second);
+    EXPECT_EQ(rc, 0) << "engine->init failed";
+    xport = engine->installTransport("cxi", nullptr);
+    EXPECT_EQ(xport, nullptr)
+        << "installTransport(\"cxi\") should fail for invalid_device";
+}
+
+TEST_F(CxiTransportUnitTest, TestAllocate) {
+    std::unique_ptr<TransferEngine> engine;
+    engine = std::make_unique<TransferEngine>(false);
+    EXPECT_NE(engine, nullptr) << "failed to allocate TransferEngine";
+
+    engine->getLocalTopology()->discover({});
+    auto hp = parseHostNameWithPort(local_server_name_);
+    int rc = engine->init(metadata_server_, local_server_name_,
+                          hp.first.c_str(), hp.second);
+    EXPECT_EQ(rc, 0) << "engine->init failed";
+
+    Transport* xport = engine->installTransport("cxi", nullptr);
+    EXPECT_NE(xport, nullptr) << "installTransport(\"cxi\") failed";
+    CxiTransport* cxi_xport = dynamic_cast<CxiTransport*>(xport);
+    EXPECT_NE(cxi_xport, nullptr) << "initialization for cxi failed";
+
+    size_t buffer_size = 128 * 1024;
+    auto buffer = std::vector<uint8_t>(buffer_size);
+    size_t dummy_buffer_size = 128;
+    auto dummy_buffer = std::vector<uint8_t>(dummy_buffer_size);
+
+    // check invalid device
+    int retcode = engine->registerLocalMemory(
+        (void*)dummy_buffer.data(), dummy_buffer_size, "invalid_device:0");
+    EXPECT_EQ(retcode, ERR_DEVICE_NOT_FOUND)
+        << "invalid device memory should not be registered!";
+
+    retcode =
+        engine->registerLocalMemory((void*)buffer.data(), buffer_size, "cpu:0");
+    EXPECT_EQ(retcode, 0) << "unable to register memory";
+
+    retcode = engine->registerLocalMemory(
+        (void*)(buffer.data() + (buffer_size >> 1)), buffer_size >> 1, "cpu:1");
+    EXPECT_EQ(retcode, ERR_ADDRESS_OVERLAPPED)
+        << "should not be able to register overlapping MRs!";
+
+    engine->unregisterLocalMemory(buffer.data());
+}
+
+};  // namespace mooncake
+
+int main(int argc, char** argv) {
+    gflags::ParseCommandLineFlags(&argc, &argv, false);
+    ::testing::InitGoogleTest(&argc, argv);
+    google::InitGoogleLogging("CxiTransportUnitTest");
+    FLAGS_logtostderr = 1;
+    int rc = RUN_ALL_TESTS();
+    google::ShutdownGoogleLogging();
+    return rc;
+}
diff --git a/mooncake-transfer-engine/tests/efa_gpu_loopback_test.cpp b/mooncake-transfer-engine/tests/efa_gpu_loopback_test.cpp
new file mode 100644
index 00000000..a7735d77
--- /dev/null
+++ b/mooncake-transfer-engine/tests/efa_gpu_loopback_test.cpp
@@ -0,0 +1,324 @@
+// Copyright 2024 KVCache.AI
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+// EFA GPU (CUDA device memory) loopback test.
+//
+// This is the GPU-memory counterpart of efa_transport_test.cpp, which only
+// exercises loopback on host (numa) buffers.  The distinction matters: the
+// libfabric EFA provider's SHM intra-node path does a host memcpy into
+// FI_HMEM_CUDA device buffers and segfaults on the first same-host transfer
+// (ofiwg/libfabric#12328).  Host-memory loopback never hits that bug, so it
+// went unnoticed by the existing tests.
+//
+// Mooncake's fix (EfaContext::tryLoopbackCopy) detects same-process
+// self-loopback and satisfies it with a GPU-aware cudaMemcpy instead of
+// routing it over EFA.  This test reproduces the crash WITHOUT the fix and
+// verifies correct, byte-accurate data movement WITH it — for both WRITE and
+// READ opcodes (the two copy in opposite directions).
+//
+// Requires EFA hardware (fi_info -p efa) AND at least one CUDA GPU.  Self-
+// skips cleanly when either is absent.
+
+#include <cuda_runtime.h>
+#include <gflags/gflags.h>
+#include <glog/logging.h>
+#include <gtest/gtest.h>
+
+#include <cstdlib>
+#include <memory>
+#include <vector>
+
+#include "transfer_engine.h"
+#include "transport/efa_transport/efa_transport.h"
+#include "transport/transport.h"
+
+using namespace mooncake;
+
+namespace mooncake {
+
+namespace {
+bool cudaAvailable() {
+    int n = 0;
+    return cudaGetDeviceCount(&n) == cudaSuccess && n > 0;
+}
+}  // namespace
+
+class EFAGpuLoopbackTest : public ::testing::Test {
+   protected:
+    void SetUp() override {
+        google::InitGoogleLogging("EFAGpuLoopbackTest");
+        FLAGS_logtostderr = 1;
+
+        const char *env = std::getenv("MC_METADATA_SERVER");
+        metadata_server_ = env ? env : "P2PHANDSHAKE";
+
+        env = std::getenv("MC_LOCAL_SERVER_NAME");
+        local_server_name_ = env ? env : "127.0.0.1:12345";
+    }
+
+    void TearDown() override { google::ShutdownGoogleLogging(); }
+
+    struct EngineSetup {
+        std::unique_ptr<TransferEngine> engine;
+        Transport *xport = nullptr;
+        void *dev_buf = nullptr;  // CUDA device buffer
+        size_t buffer_size = 0;
+        SegmentID segment_id = 0;
+        bool ok = false;
+    };
+
+    // Create engine, install EFA transport, allocate + register a CUDA
+    // device buffer, and open our own segment for loopback.
+    EngineSetup createEngine(size_t buffer_size = 1ull << 26 /* 64 MB */) {
+        EngineSetup s;
+        s.buffer_size = buffer_size;
+
+        if (cudaSetDevice(0) != cudaSuccess) {
+            LOG(WARNING) << "cudaSetDevice(0) failed";
+            return s;
+        }
+
+        s.engine = std::make_unique<TransferEngine>(false);
+        s.engine->getLocalTopology()->discover({});
+        auto hp = parseHostNameWithPort(local_server_name_);
+        int rc = s.engine->init(metadata_server_, local_server_name_,
+                                hp.first.c_str(), hp.second);
+        EXPECT_EQ(rc, 0) << "engine->init failed";
+        if (rc != 0) return s;
+
+        s.xport = s.engine->installTransport("efa", nullptr);
+        EXPECT_NE(s.xport, nullptr) << "installTransport(\"efa\") failed";
+        if (!s.xport) return s;
+
+        cudaError_t cerr = cudaMalloc(&s.dev_buf, buffer_size);
+        EXPECT_EQ(cerr, cudaSuccess)
+            << "cudaMalloc failed: " << cudaGetErrorString(cerr);
+        if (cerr != cudaSuccess) return s;
+
+        // Register as GPU memory so the EFA transport tags the MR with
+        // FI_HMEM_CUDA (the registration path under test).
+        rc = s.engine->registerLocalMemory(s.dev_buf, buffer_size, "cuda:0");
+        EXPECT_EQ(rc, 0) << "registerLocalMemory(cuda:0) failed";
+        if (rc != 0) return s;
+
+        auto actual_addr = s.engine->getLocalIpAndPort();
+        s.segment_id = s.engine->openSegment(actual_addr);
+        s.ok = true;
+        return s;
+    }
+
+    void destroyEngine(EngineSetup &s) {
+        if (s.engine && s.dev_buf) {
+            s.engine->unregisterLocalMemory(s.dev_buf);
+        }
+        if (s.dev_buf) {
+            cudaFree(s.dev_buf);
+            s.dev_buf = nullptr;
+        }
+    }
+
+    bool submitAndWait(TransferEngine *engine, SegmentID segment_id,
+                       void *source, uint64_t target_offset, size_t length,
+                       TransferRequest::OpCode opcode) {
+        auto batch_id = engine->allocateBatchID(1);
+        TransferRequest entry;
+        entry.opcode = opcode;
+        entry.length = length;
+        entry.source = (uint8_t *)source;
+        entry.target_id = segment_id;
+        entry.target_offset = target_offset;
+
+        Status s = engine->submitTransfer(batch_id, {entry});
+        if (!s.ok()) {
+            LOG(ERROR) << "submitTransfer failed: " << s.ToString();
+            engine->freeBatchID(batch_id);
+            return false;
+        }
+
+        TransferStatus status;
+        const int kMaxPollIterations = 2000000;
+        for (int i = 0; i < kMaxPollIterations; ++i) {
+            s = engine->getTransferStatus(batch_id, 0, status);
+            if (!s.ok()) {
+                engine->freeBatchID(batch_id);
+                return false;
+            }
+            if (status.s == TransferStatusEnum::COMPLETED) {
+                engine->freeBatchID(batch_id);
+                return true;
+            }
+            if (status.s == TransferStatusEnum::FAILED) {
+                LOG(ERROR) << "Transfer FAILED";
+                engine->freeBatchID(batch_id);
+                return false;
+            }
+        }
+        LOG(ERROR) << "Transfer timed out";
+        engine->freeBatchID(batch_id);
+        return false;
+    }
+
+    std::string metadata_server_;
+    std::string local_server_name_;
+};
+
+// Test 1: GPU loopback WRITE must not crash.
+//
+// Without the fix this segfaults inside the EFA provider's SHM path
+// (host memcpy into a device pointer).  With the fix it completes via a
+// local cudaMemcpy.
+TEST_F(EFAGpuLoopbackTest, GpuLoopbackWrite) {
+    if (!cudaAvailable()) GTEST_SKIP() << "No CUDA GPU present";
+    auto setup = createEngine();
+    if (!setup.ok) GTEST_SKIP() << "EFA/CUDA setup unavailable";
+
+    auto segment_desc =
+        setup.engine->getMetadata()->getSegmentDescByID(setup.segment_id);
+    ASSERT_NE(segment_desc, nullptr);
+    uint64_t remote_base = (uint64_t)segment_desc->buffers[0].addr;
+
+    const size_t kDataLength = 4096;
+    ASSERT_EQ(cudaMemset(setup.dev_buf, 0xAB, kDataLength), cudaSuccess);
+    ASSERT_EQ(cudaDeviceSynchronize(), cudaSuccess);
+
+    // Write src(=dev_buf) -> remote(=dev_buf + offset), both on the GPU.
+    bool ok = submitAndWait(setup.engine.get(), setup.segment_id, setup.dev_buf,
+                            remote_base + (setup.buffer_size / 2), kDataLength,
+                            TransferRequest::WRITE);
+    EXPECT_TRUE(ok) << "GPU loopback write should succeed";
+
+    destroyEngine(setup);
+}
+
+// Test 2: GPU loopback WRITE then READ, with byte-accurate verification.
+//
+// The first half holds the source pattern; we WRITE it into the second
+// half, scribble the first half, then READ the second half back into the
+// first half and compare.  This exercises BOTH copy directions of
+// tryLoopbackCopy (WRITE: src->dst, READ: dst->src) and proves the data is
+// actually moved correctly, not merely "did not crash".
+TEST_F(EFAGpuLoopbackTest, GpuLoopbackWriteThenRead) {
+    if (!cudaAvailable()) GTEST_SKIP() << "No CUDA GPU present";
+    auto setup = createEngine();
+    if (!setup.ok) GTEST_SKIP() << "EFA/CUDA setup unavailable";
+
+    auto segment_desc =
+        setup.engine->getMetadata()->getSegmentDescByID(setup.segment_id);
+    ASSERT_NE(segment_desc, nullptr);
+    uint64_t remote_base = (uint64_t)segment_desc->buffers[0].addr;
+
+    const size_t kDataLength = 4ull << 20;  // 4 MB
+    ASSERT_LE(2 * kDataLength, setup.buffer_size);
+
+    uint8_t *dev = (uint8_t *)setup.dev_buf;
+    const uint64_t second_half = setup.buffer_size / 2;
+
+    // Build a known pattern on the host and copy it into the first half.
+    std::vector<uint8_t> pattern(kDataLength);
+    for (size_t i = 0; i < kDataLength; ++i)
+        pattern[i] = (uint8_t)(i * 131 + 7);
+    ASSERT_EQ(
+        cudaMemcpy(dev, pattern.data(), kDataLength, cudaMemcpyHostToDevice),
+        cudaSuccess);
+
+    // WRITE: first half -> second half (loopback).
+    ASSERT_TRUE(submitAndWait(setup.engine.get(), setup.segment_id, dev,
+                              remote_base + second_half, kDataLength,
+                              TransferRequest::WRITE))
+        << "WRITE should succeed";
+
+    // Corrupt the first half so the READ-back has to actually move bytes.
+    ASSERT_EQ(cudaMemset(dev, 0x00, kDataLength), cudaSuccess);
+    ASSERT_EQ(cudaDeviceSynchronize(), cudaSuccess);
+
+    // READ: second half -> first half (loopback, opposite direction).
+    ASSERT_TRUE(submitAndWait(setup.engine.get(), setup.segment_id, dev,
+                              remote_base + second_half, kDataLength,
+                              TransferRequest::READ))
+        << "READ should succeed";
+
+    // Verify the first half now matches the original pattern.
+    std::vector<uint8_t> readback(kDataLength, 0xFF);
+    ASSERT_EQ(
+        cudaMemcpy(readback.data(), dev, kDataLength, cudaMemcpyDeviceToHost),
+        cudaSuccess);
+    EXPECT_EQ(0, memcmp(readback.data(), pattern.data(), kDataLength))
+        << "READ-back GPU data should match the written pattern";
+
+    destroyEngine(setup);
+}
+
+// Test 3: batch of GPU loopback writes (multiple slices), to exercise the
+// per-slice loopback short-circuit inside the grouping loop.
+TEST_F(EFAGpuLoopbackTest, GpuLoopbackMultiWrite) {
+    if (!cudaAvailable()) GTEST_SKIP() << "No CUDA GPU present";
+    auto setup = createEngine();
+    if (!setup.ok) GTEST_SKIP() << "EFA/CUDA setup unavailable";
+
+    auto segment_desc =
+        setup.engine->getMetadata()->getSegmentDescByID(setup.segment_id);
+    ASSERT_NE(segment_desc, nullptr);
+    uint64_t remote_base = (uint64_t)segment_desc->buffers[0].addr;
+
+    const size_t kSliceLen = 65536;
+    const int kBatchSize = 16;
+    const uint64_t dst_region = setup.buffer_size / 2;
+    ASSERT_LE(dst_region + (uint64_t)kBatchSize * kSliceLen, setup.buffer_size);
+
+    ASSERT_EQ(cudaMemset(setup.dev_buf, 0x5A, kBatchSize * kSliceLen),
+              cudaSuccess);
+    ASSERT_EQ(cudaDeviceSynchronize(), cudaSuccess);
+
+    auto batch_id = setup.engine->allocateBatchID(kBatchSize);
+    std::vector<TransferRequest> requests;
+    for (int i = 0; i < kBatchSize; ++i) {
+        TransferRequest entry;
+        entry.opcode = TransferRequest::WRITE;
+        entry.length = kSliceLen;
+        entry.source = (uint8_t *)setup.dev_buf + i * kSliceLen;
+        entry.target_id = setup.segment_id;
+        entry.target_offset = remote_base + dst_region + i * kSliceLen;
+        requests.push_back(entry);
+    }
+
+    Status s = setup.engine->submitTransfer(batch_id, requests);
+    ASSERT_TRUE(s.ok()) << "submitTransfer failed: " << s.ToString();
+
+    for (int task_id = 0; task_id < kBatchSize; ++task_id) {
+        TransferStatus status;
+        const int kMaxPollIterations = 2000000;
+        for (int i = 0; i < kMaxPollIterations; ++i) {
+            s = setup.engine->getTransferStatus(batch_id, task_id, status);
+            ASSERT_TRUE(s.ok());
+            if (status.s == TransferStatusEnum::COMPLETED) break;
+            ASSERT_NE(status.s, TransferStatusEnum::FAILED)
+                << "task " << task_id << " failed";
+        }
+        ASSERT_EQ(status.s, TransferStatusEnum::COMPLETED)
+            << "task " << task_id << " did not complete";
+    }
+
+    s = setup.engine->freeBatchID(batch_id);
+    ASSERT_TRUE(s.ok());
+
+    destroyEngine(setup);
+}
+
+}  // namespace mooncake
+
+int main(int argc, char **argv) {
+    gflags::ParseCommandLineFlags(&argc, &argv, false);
+    ::testing::InitGoogleTest(&argc, argv);
+    return RUN_ALL_TESTS();
+}
diff --git a/mooncake-transfer-engine/tests/efa_transport_test.cpp b/mooncake-transfer-engine/tests/efa_transport_test.cpp
index 066421ac..704a3384 100644
--- a/mooncake-transfer-engine/tests/efa_transport_test.cpp
+++ b/mooncake-transfer-engine/tests/efa_transport_test.cpp
@@ -21,6 +21,7 @@
 #include <memory>
 
 #include "transfer_engine.h"
+#include "transport/efa_transport/efa_transport.h"
 #include "transport/transport.h"
 
 using namespace mooncake;
@@ -329,6 +330,176 @@ TEST_F(EFATransportTest, StressMultipleBatches) {
     destroyEngine(setup);
 }
 
+// Test 6: warmupSegment on loopback peer
+//
+// Exercises EfaTransport::warmupSegment() which is the C++ entry point
+// behind the warmup_efa_segment() Python binding / warmupEfaSegment() C API.
+// Loopback is enough to cover the handshake + fi_av_insert path AND the
+// idempotent short-circuit on the second call.
+TEST_F(EFATransportTest, WarmupSegmentLoopback) {
+    auto setup = createEngine();
+
+    auto *efa = dynamic_cast<EfaTransport *>(setup.xport);
+    ASSERT_NE(efa, nullptr)
+        << "installTransport did not return an EfaTransport";
+
+    // First call: should connect every (local NIC x peer NIC) pair.
+    int rc = efa->warmupSegment(setup.engine->getLocalIpAndPort());
+    EXPECT_EQ(rc, 0) << "warmupSegment should succeed on loopback";
+
+    // Second call: should short-circuit (all endpoints already connected).
+    rc = efa->warmupSegment(setup.engine->getLocalIpAndPort());
+    EXPECT_EQ(rc, 0) << "warmupSegment should be idempotent";
+
+    // Empty / self-name: short-circuit path returning 0 without touching AV.
+    rc = efa->warmupSegment("");
+    EXPECT_EQ(rc, 0) << "warmupSegment(\"\") should be a no-op";
+
+    destroyEngine(setup);
+}
+
+// Test 7: warmupSegment on a non-existent segment name should fail cleanly
+// (no crash, no hang) rather than blocking for the poll timeout.
+TEST_F(EFATransportTest, WarmupSegmentNotFound) {
+    auto setup = createEngine();
+
+    auto *efa = dynamic_cast<EfaTransport *>(setup.xport);
+    ASSERT_NE(efa, nullptr);
+
+    int rc = efa->warmupSegment("127.0.0.1:1");  // not openSegment'd
+    EXPECT_NE(rc, 0) << "warmupSegment should fail for unknown segment";
+
+    destroyEngine(setup);
+}
+
+// Test 8: registerLocalMemoryBatch / unregisterLocalMemoryBatch round-trip.
+// Covers the batched MR path which the single-buffer tests above never hit.
+TEST_F(EFATransportTest, RegisterMemoryBatch) {
+    auto engine = std::make_unique<TransferEngine>(false);
+    engine->getLocalTopology()->discover({});
+    auto hp = parseHostNameWithPort(local_server_name_);
+    int rc = engine->init(metadata_server_, local_server_name_,
+                          hp.first.c_str(), hp.second);
+    ASSERT_EQ(rc, 0);
+
+    Transport *xport = engine->installTransport("efa", nullptr);
+    ASSERT_NE(xport, nullptr);
+
+    const size_t kBufSize = 4ull << 20;  // 4 MB each
+    const int kNumBufs = 4;
+    std::vector<void *> addrs;
+    std::vector<BufferEntry> entries;
+    for (int i = 0; i < kNumBufs; ++i) {
+        void *a = allocateMemoryPool(kBufSize, 0);
+        ASSERT_NE(a, nullptr);
+        addrs.push_back(a);
+        entries.push_back({a, kBufSize});
+    }
+
+    rc = engine->registerLocalMemoryBatch(entries, "cpu:0");
+    EXPECT_EQ(rc, 0) << "registerLocalMemoryBatch should succeed";
+
+    rc = engine->unregisterLocalMemoryBatch(addrs);
+    EXPECT_EQ(rc, 0) << "unregisterLocalMemoryBatch should succeed";
+
+    for (void *a : addrs) freeMemoryPool(a, kBufSize);
+}
+
+// Test 9: Larger transfer (64 MB total split into 1 MB slices) to exercise
+// the WR / CQ pacing logic in EfaContext::submitSlicesOnPeer beyond what the
+// 16 x 64 KB MultiWrite test reaches.
+TEST_F(EFATransportTest, LargeTransfer) {
+    const size_t kBufSize = 128ull << 20;  // 128 MB
+    auto setup = createEngine(kBufSize);
+
+    auto segment_desc =
+        setup.engine->getMetadata()->getSegmentDescByID(setup.segment_id);
+    ASSERT_NE(segment_desc, nullptr);
+    uint64_t remote_base = (uint64_t)segment_desc->buffers[0].addr;
+
+    const size_t kSliceLen = 1ull << 20;  // 1 MB per slice
+    const int kNumSlices = 64;            // 64 MB total
+    ASSERT_LE(static_cast<size_t>(kNumSlices) * kSliceLen, kBufSize / 2);
+
+    // Fill first half with known data
+    uint8_t *buf = (uint8_t *)setup.addr;
+    for (size_t i = 0; i < static_cast<size_t>(kNumSlices) * kSliceLen; ++i)
+        buf[i] = (uint8_t)(i & 0xFF);
+
+    auto batch_id = setup.engine->allocateBatchID(kNumSlices);
+    std::vector<TransferRequest> requests;
+    requests.reserve(kNumSlices);
+    for (int i = 0; i < kNumSlices; ++i) {
+        TransferRequest entry;
+        entry.opcode = TransferRequest::WRITE;
+        entry.length = kSliceLen;
+        entry.source = buf + i * kSliceLen;
+        entry.target_id = setup.segment_id;
+        entry.target_offset = remote_base + (kBufSize / 2) + i * kSliceLen;
+        requests.push_back(entry);
+    }
+
+    Status s = setup.engine->submitTransfer(batch_id, requests);
+    ASSERT_TRUE(s.ok()) << "submitTransfer failed: " << s.ToString();
+
+    for (int task_id = 0; task_id < kNumSlices; ++task_id) {
+        TransferStatus status;
+        const int kMaxPollIterations = 2000000;
+        int i = 0;
+        for (; i < kMaxPollIterations; ++i) {
+            s = setup.engine->getTransferStatus(batch_id, task_id, status);
+            ASSERT_TRUE(s.ok());
+            if (status.s == TransferStatusEnum::COMPLETED) break;
+            ASSERT_NE(status.s, TransferStatusEnum::FAILED);
+        }
+        ASSERT_EQ(status.s, TransferStatusEnum::COMPLETED)
+            << "task " << task_id << " did not complete";
+    }
+
+    s = setup.engine->freeBatchID(batch_id);
+    ASSERT_TRUE(s.ok());
+
+    // Verify byte-level integrity of the last slice (spot check).
+    EXPECT_EQ(0, memcmp(buf + (kNumSlices - 1) * kSliceLen,
+                        buf + (kBufSize / 2) + (kNumSlices - 1) * kSliceLen,
+                        kSliceLen));
+
+    destroyEngine(setup);
+}
+
+// Test 10: Repeated open/close of the same remote segment must not leak AV
+// slots or break loopback transfers — this is the setPeerNicPath-detach path
+// that target restarts depend on under the SRD shared-endpoint model.
+TEST_F(EFATransportTest, RepeatedOpenSegment) {
+    auto setup = createEngine();
+
+    auto actual_addr = setup.engine->getLocalIpAndPort();
+
+    // First write via setup.segment_id (from createEngine()).
+    auto segment_desc =
+        setup.engine->getMetadata()->getSegmentDescByID(setup.segment_id);
+    ASSERT_NE(segment_desc, nullptr);
+    uint64_t remote_base = (uint64_t)segment_desc->buffers[0].addr;
+    memset(setup.addr, 0xCD, 4096);
+    EXPECT_TRUE(submitAndWait(setup.engine.get(), setup.segment_id, setup.addr,
+                              remote_base, 4096, TransferRequest::WRITE));
+
+    // Re-open same segment several times; each should return a working handle
+    // and subsequent writes should still succeed.
+    for (int i = 0; i < 5; ++i) {
+        SegmentID sid = setup.engine->openSegment(actual_addr);
+        ASSERT_NE(sid, (SegmentID)-1);
+        auto desc = setup.engine->getMetadata()->getSegmentDescByID(sid);
+        ASSERT_NE(desc, nullptr);
+        uint64_t base = (uint64_t)desc->buffers[0].addr;
+        ASSERT_TRUE(submitAndWait(setup.engine.get(), sid, setup.addr, base,
+                                  4096, TransferRequest::WRITE))
+            << "write #" << i << " after re-open failed";
+    }
+
+    destroyEngine(setup);
+}
+
 }  // namespace mooncake
 
 int main(int argc, char **argv) {
diff --git a/mooncake-transfer-engine/tests/endpoint_store_integration_test.cpp b/mooncake-transfer-engine/tests/endpoint_store_integration_test.cpp
new file mode 100644
index 00000000..2972fa12
--- /dev/null
+++ b/mooncake-transfer-engine/tests/endpoint_store_integration_test.cpp
@@ -0,0 +1,128 @@
+// Copyright 2026 KVCache.AI
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+//
+// Integration test for #1845. Verifies the end-to-end wiring of the fix:
+// WorkerPool::monitorWorker actually calls RdmaContext::reclaimEndpoints at
+// ~1 Hz, causing quiescent entries in the endpoint store's waiting_list_ to
+// drain without any further insertion traffic. The unit tests in
+// endpoint_store_test.cpp verify the reclaim method itself; this file
+// verifies that the scheduler invokes it.
+//
+// Requires an RDMA device. Passes on soft-RoCE (`rdma_rxe`) as well as real
+// NICs. Self-skips (GTEST_SKIP) when no device is present, so it is safe to
+// register with ctest on CI runners without RDMA.
+//
+// Environment override: set MC_TEST_DEVICE_NAME to force a specific device;
+// otherwise the first device returned by ibv_get_device_list is used.
+
+#include <gtest/gtest.h>
+#include <infiniband/verbs.h>
+
+#include <chrono>
+#include <cstdlib>
+#include <memory>
+#include <string>
+#include <thread>
+
+#include "config.h"
+#include "transport/rdma_transport/endpoint_store.h"
+#include "transport/rdma_transport/rdma_context.h"
+#include "transport/rdma_transport/rdma_endpoint.h"
+#include "transport/rdma_transport/rdma_transport.h"
+
+#if defined(__has_feature)
+#define MC_HAS_FEATURE(x) __has_feature(x)
+#else
+#define MC_HAS_FEATURE(x) 0
+#endif
+#if defined(__SANITIZE_ADDRESS__) || MC_HAS_FEATURE(address_sanitizer)
+#include <sanitizer/lsan_interface.h>
+#define MC_LSAN_IGNORE_OBJECT(p) __lsan_ignore_object(p)
+#else
+#define MC_LSAN_IGNORE_OBJECT(p) ((void)(p))
+#endif
+
+using namespace mooncake;
+
+namespace {
+
+std::string pickRdmaDevice() {
+    const char *override_name = std::getenv("MC_TEST_DEVICE_NAME");
+    if (override_name && *override_name) return override_name;
+    int num_devices = 0;
+    ibv_device **list = ibv_get_device_list(&num_devices);
+    if (!list) return "";
+    if (num_devices == 0) {
+        ibv_free_device_list(list);
+        return "";
+    }
+    std::string name = ibv_get_device_name(list[0]);
+    ibv_free_device_list(list);
+    return name;
+}
+
+// Build an RdmaEndPoint with no QPs and active_=false. The store's reclaim
+// path only inspects hasOutstandingSlice(), which for an endpoint with empty
+// qp_list_ reduces to !active_. Safe to destruct because qp_list_ is empty.
+std::shared_ptr<RdmaEndPoint> makeQuiescentEndpoint(RdmaContext &ctx) {
+    auto ep = std::make_shared<RdmaEndPoint>(ctx);
+    ep->set_active(false);
+    return ep;
+}
+
+// Verifies the full fix wiring: after construct() spawns monitorWorker, a
+// quiescent entry injected into the store's waiting_list_ is drained by the
+// scheduler within ~1.5 s with no further insertion traffic.
+TEST(EndpointStoreIntegration, MonitorWorkerTickDrainsWaitingList) {
+    const std::string device = pickRdmaDevice();
+    if (device.empty()) {
+        GTEST_SKIP() << "no RDMA device available — integration test requires "
+                        "rxe0, mlx5, or similar. Set MC_TEST_DEVICE_NAME to "
+                        "override.";
+    }
+
+    // RdmaTransport's destructor dereferences metadata_ which is null until
+    // init(); leak the engine to avoid touching that path. Marked ignored so
+    // LSAN under ASAN builds doesn't flag this intentional leak.
+    auto *transport = new RdmaTransport();
+    MC_LSAN_IGNORE_OBJECT(transport);
+    auto context = std::make_shared<RdmaContext>(*transport, device);
+    auto &config = globalConfig();
+    int rc = context->construct(config.num_cq_per_ctx,
+                                config.num_comp_channels_per_ctx, config.port,
+                                config.gid_index, config.max_cqe,
+                                /*max_endpoints=*/4);
+    if (rc != 0) {
+        GTEST_SKIP() << "RdmaContext::construct failed on device " << device
+                     << " (rc=" << rc << "); no usable RDMA device on this "
+                     << "host (e.g., CI runners may enumerate a phantom "
+                     << "mlx5_0 without a working port).";
+    }
+
+    context->testOnlyInsertWaiting(makeQuiescentEndpoint(*context));
+    context->testOnlyInsertWaiting(makeQuiescentEndpoint(*context));
+    context->testOnlyInsertWaiting(makeQuiescentEndpoint(*context));
+    ASSERT_EQ(context->waitingListSize(), 3u);
+
+    // monitorWorker's reclaim tick fires every ~1 s. Give it enough margin
+    // for scheduling jitter but keep the test fast.
+    std::this_thread::sleep_for(std::chrono::milliseconds(1500));
+
+    EXPECT_EQ(context->waitingListSize(), 0u)
+        << "monitorWorker must call reclaimEndpoints within ~1 s. If this "
+           "fails, either the periodic tick in worker_pool.cpp was removed or "
+           "reclaim is failing on quiescent entries.";
+}
+
+}  // namespace
diff --git a/mooncake-transfer-engine/tests/endpoint_store_test.cpp b/mooncake-transfer-engine/tests/endpoint_store_test.cpp
new file mode 100644
index 00000000..5f3b1216
--- /dev/null
+++ b/mooncake-transfer-engine/tests/endpoint_store_test.cpp
@@ -0,0 +1,178 @@
+// Copyright 2026 KVCache.AI
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+// Regression coverage for #1845. Asserts that
+// SIEVEEndpointStore::reclaimEndpoint drains quiescent entries from
+// waiting_list_ without requiring a subsequent insertEndpoint call. This is the
+// invariant the periodic-reclaim tick in monitorWorker depends on.
+
+#include <gtest/gtest.h>
+
+#include <memory>
+
+#include "transport/rdma_transport/endpoint_store.h"
+#include "transport/rdma_transport/rdma_context.h"
+#include "transport/rdma_transport/rdma_endpoint.h"
+#include "transport/rdma_transport/rdma_transport.h"
+
+#if defined(__has_feature)
+#define MC_HAS_FEATURE(x) __has_feature(x)
+#else
+#define MC_HAS_FEATURE(x) 0
+#endif
+#if defined(__SANITIZE_ADDRESS__) || MC_HAS_FEATURE(address_sanitizer)
+#include <sanitizer/lsan_interface.h>
+#define MC_LSAN_IGNORE_OBJECT(p) __lsan_ignore_object(p)
+#else
+#define MC_LSAN_IGNORE_OBJECT(p) ((void)(p))
+#endif
+
+using namespace mooncake;
+
+namespace {
+
+// Build an RdmaEndPoint that owns zero QPs and has active_=false. construct()
+// is deliberately not called — the store's reclaim logic only inspects
+// hasOutstandingSlice(), which for an endpoint with empty qp_list_ returns
+// whatever active_ is.
+std::shared_ptr<RdmaEndPoint> makeQuiescentEndpoint(RdmaContext& ctx) {
+    auto ep = std::make_shared<RdmaEndPoint>(ctx);
+    ep->set_active(false);
+    return ep;
+}
+
+std::shared_ptr<RdmaEndPoint> makeActiveEndpoint(RdmaContext& ctx) {
+    // Default ctor leaves active_=true.
+    return std::make_shared<RdmaEndPoint>(ctx);
+}
+
+class EndpointStoreTest : public ::testing::Test {
+   protected:
+    // Leaked on purpose: RdmaTransport's destructor dereferences metadata_,
+    // which is null when the engine was never init()ed. We only need a live
+    // reference for RdmaContext's constructor; the engine object is otherwise
+    // unused by the reclaim logic under test.
+    RdmaTransport* transport_ = nullptr;
+    std::unique_ptr<RdmaContext> ctx_;
+
+    void SetUp() override {
+        transport_ = new RdmaTransport();
+        // Intentional leak: ~RdmaTransport dereferences metadata_, which is
+        // null until install(). Marking it ignored keeps LSAN under ASAN
+        // builds from flagging this one allocation while still catching
+        // real leaks elsewhere.
+        MC_LSAN_IGNORE_OBJECT(transport_);
+        ctx_ = std::make_unique<RdmaContext>(*transport_, "unused");
+    }
+};
+
+// The core invariant behind #1845's fix: reclaimEndpoint must drain quiescent
+// entries on its own, without needing a subsequent insertEndpoint to trigger
+// it. Before the fix, reclaim ran only on insertion, so if insertions stopped
+// (e.g., all peers died), waiting_list_ grew unboundedly. The periodic tick
+// from monitorWorker calls this method every second; this test asserts its
+// contract in isolation.
+TEST_F(EndpointStoreTest, ReclaimDrainsQuiescentEntries) {
+    SIEVEEndpointStore store(/*max_size=*/4);
+
+    constexpr size_t kN = 10;
+    for (size_t i = 0; i < kN; ++i) {
+        store.testOnlyInsertWaiting(makeQuiescentEndpoint(*ctx_));
+    }
+    EXPECT_EQ(store.waitingListSize(), kN);
+
+    store.reclaimEndpoint();
+    EXPECT_EQ(store.waitingListSize(), 0u)
+        << "reclaimEndpoint must drain quiescent entries with no insertion "
+           "prerequisite";
+}
+
+// Negative control: reclaim must leave entries in place if they still report
+// outstanding slices. Ensures we didn't break the hasOutstandingSlice gate.
+TEST_F(EndpointStoreTest, ReclaimLeavesActiveEntries) {
+    SIEVEEndpointStore store(4);
+
+    store.testOnlyInsertWaiting(makeActiveEndpoint(*ctx_));
+    store.testOnlyInsertWaiting(makeActiveEndpoint(*ctx_));
+    store.testOnlyInsertWaiting(makeQuiescentEndpoint(*ctx_));
+    EXPECT_EQ(store.waitingListSize(), 3u);
+
+    store.reclaimEndpoint();
+    EXPECT_EQ(store.waitingListSize(), 2u)
+        << "reclaim should drop only the quiescent endpoint, keep the two "
+           "active ones";
+}
+
+TEST_F(EndpointStoreTest, ReclaimIsIdempotentWhenEmpty) {
+    SIEVEEndpointStore store(4);
+
+    store.reclaimEndpoint();
+    EXPECT_EQ(store.waitingListSize(), 0u);
+
+    store.testOnlyInsertWaiting(makeQuiescentEndpoint(*ctx_));
+    store.reclaimEndpoint();
+    EXPECT_EQ(store.waitingListSize(), 0u);
+
+    store.reclaimEndpoint();  // second call is a no-op
+    EXPECT_EQ(store.waitingListSize(), 0u);
+}
+
+// Demonstrates the #1845 failure mode: once insertions stop but evictions
+// keep landing in the waiting list, nothing drains them without an explicit
+// reclaim call. Before this fix, reclaimEndpoint ran only from insertEndpoint,
+// so "many evictions, no new peers to connect to" meant waiting_list_ grew
+// without bound. This test simulates that workload without any RDMA or
+// scheduler; the assertion is a strict "zero reclaim calls leaves the leak
+// at its peak."
+TEST_F(EndpointStoreTest, LeakManifestsWithoutReclaimCall) {
+    SIEVEEndpointStore store(/*max_size=*/4);
+
+    constexpr size_t kEvictions = 1118;  // match reporter's eviction count
+    for (size_t i = 0; i < kEvictions; ++i) {
+        store.testOnlyInsertWaiting(makeQuiescentEndpoint(*ctx_));
+    }
+
+    // Without a reclaim call the leak is at its peak.
+    EXPECT_EQ(store.waitingListSize(), kEvictions)
+        << "baseline confirmation: waiting_list_ accumulates as expected";
+
+    // The fix is a 1 Hz invocation of this single method from monitorWorker.
+    // One call is enough to drain the entire backlog (because the entries are
+    // quiescent by the time the peer-death path finishes). This is the
+    // invariant the PR relies on.
+    store.reclaimEndpoint();
+    EXPECT_EQ(store.waitingListSize(), 0u)
+        << "a single reclaim call drains the full backlog once insertions "
+           "stop; this is what the periodic tick in monitorWorker provides";
+}
+
+// Guards against a future regression that re-breaks the reclaim contract —
+// e.g., someone changing reclaimEndpoint to no-op when endpoint_map_ is
+// empty, on the incorrect assumption that reclaim only runs from
+// insertEndpoint. Walking 1000 quiescent entries should still drain them.
+TEST_F(EndpointStoreTest, ReclaimDoesNotRequireActiveMap) {
+    SIEVEEndpointStore store(4);
+    EXPECT_EQ(store.getSize(), 0u);  // endpoint_map_ empty
+
+    for (size_t i = 0; i < 1000; ++i) {
+        store.testOnlyInsertWaiting(makeQuiescentEndpoint(*ctx_));
+    }
+    EXPECT_EQ(store.getSize(), 0u);  // still empty
+    EXPECT_EQ(store.waitingListSize(), 1000u);
+
+    store.reclaimEndpoint();
+    EXPECT_EQ(store.waitingListSize(), 0u);
+}
+
+}  // namespace
diff --git a/mooncake-transfer-engine/tests/hip_transport_test.cpp b/mooncake-transfer-engine/tests/hip_transport_test.cpp
new file mode 100644
index 00000000..255b195a
--- /dev/null
+++ b/mooncake-transfer-engine/tests/hip_transport_test.cpp
@@ -0,0 +1,133 @@
+// Copyright 2025 Mooncake Authors
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include <gflags/gflags.h>
+#include <glog/logging.h>
+#include <gtest/gtest.h>
+
+#include <memory>
+#include <string>
+
+#include "cuda_alike.h"
+#include "transfer_engine.h"
+#include "transfer_metadata.h"  // P2PHANDSHAKE
+#include "transport/transport.h"
+
+using namespace mooncake;
+
+// startAsyncTransfer() switches the active device to the source GPU; if it does
+// not restore it, the calling thread (which also launches the engine's compute
+// kernels) is left on the wrong GPU and the next kernel fails with
+// hipErrorInvalidDevice. This pins the caller to one GPU, transfers from a
+// buffer on a different GPU, and asserts the active device is unchanged.
+
+namespace {
+constexpr size_t kLen = 64 * 1024;
+
+void* allocOnDevice(size_t size, int device) {
+    EXPECT_EQ(cudaSetDevice(device), cudaSuccess);
+    void* ptr = nullptr;
+    EXPECT_EQ(cudaMalloc(&ptr, size), cudaSuccess);
+    return ptr;
+}
+}  // namespace
+
+TEST(HipTransportTest, RestoresActiveDeviceAfterTransfer) {
+    int device_count = 0;
+    ASSERT_EQ(cudaGetDeviceCount(&device_count), cudaSuccess);
+    if (device_count < 2) {
+        GTEST_SKIP() << "Needs >= 2 GPUs: the transfer source must live on a "
+                        "different device than the calling thread's device.";
+    }
+
+    const int kCallerDevice = 0;  // device the engine thread runs on
+    const int kSourceDevice = 1;  // KV source lives on a different GPU
+
+    // P2PHANDSHAKE: self-contained loopback, no external metadata server.
+    auto engine = std::make_unique<TransferEngine>(false);
+    const std::string server_name = "127.0.0.1:17813";
+    if (engine->init(P2PHANDSHAKE, server_name, "127.0.0.1", 17813) != 0) {
+        GTEST_SKIP() << "TransferEngine init failed in this environment.";
+    }
+
+    Transport* transport = engine->installTransport("hip", nullptr);
+    if (transport == nullptr) {
+        GTEST_SKIP()
+            << "HIP transport unavailable (built without -DUSE_HIP=ON?).";
+    }
+
+    // Both buffers on kSourceDevice so the source GPU differs from the caller.
+    void* src = allocOnDevice(kLen, kSourceDevice);
+    void* dst = allocOnDevice(kLen, kSourceDevice);
+    ASSERT_EQ(engine->registerLocalMemory(
+                  src, kLen, "cuda:" + std::to_string(kSourceDevice)),
+              0);
+    ASSERT_EQ(engine->registerLocalMemory(
+                  dst, kLen, "cuda:" + std::to_string(kSourceDevice)),
+              0);
+
+    auto segment_id = engine->openSegment(server_name);
+    ASSERT_GE(segment_id, 0);
+
+    ASSERT_EQ(cudaSetDevice(kSourceDevice), cudaSuccess);
+    ASSERT_EQ(cudaMemset(src, 0xAB, kLen), cudaSuccess);
+    ASSERT_EQ(cudaDeviceSynchronize(), cudaSuccess);
+
+    // Pin the caller to a different device than the source, as the engine does.
+    ASSERT_EQ(cudaSetDevice(kCallerDevice), cudaSuccess);
+
+    auto batch_id = engine->allocateBatchID(1);
+    TransferRequest entry;
+    entry.opcode = TransferRequest::WRITE;
+    entry.length = kLen;
+    entry.source = src;
+    entry.target_id = segment_id;
+    entry.target_offset = reinterpret_cast<uint64_t>(dst);
+    Status s = engine->submitTransfer(batch_id, {entry});
+    ASSERT_TRUE(s.ok());
+
+    // Invariant: the caller's device is unchanged (== kSourceDevice before
+    // fix).
+    int active_after_submit = -1;
+    ASSERT_EQ(cudaGetDevice(&active_after_submit), cudaSuccess);
+    EXPECT_EQ(active_after_submit, kCallerDevice)
+        << "startAsyncTransfer() must restore the caller's active device. "
+           "Leaving it on the source GPU corrupts the engine thread's HIP "
+           "context, so its next kernel launch fails with "
+           "hipErrorInvalidDevice.";
+
+    // Drain the transfer (also a basic functional check).
+    TransferStatus status;
+    do {
+        ASSERT_TRUE(engine->getTransferStatus(batch_id, 0, status).ok());
+    } while (status.s == TransferStatusEnum::WAITING);
+    EXPECT_EQ(status.s, TransferStatusEnum::COMPLETED);
+
+    int active_after_wait = -1;
+    ASSERT_EQ(cudaGetDevice(&active_after_wait), cudaSuccess);
+    EXPECT_EQ(active_after_wait, kCallerDevice);
+
+    engine->freeBatchID(batch_id);
+    engine->unregisterLocalMemory(src);
+    engine->unregisterLocalMemory(dst);
+    (void)cudaSetDevice(kSourceDevice);
+    (void)cudaFree(src);
+    (void)cudaFree(dst);
+}
+
+int main(int argc, char** argv) {
+    gflags::ParseCommandLineFlags(&argc, &argv, false);
+    ::testing::InitGoogleTest(&argc, argv);
+    return RUN_ALL_TESTS();
+}
diff --git a/mooncake-transfer-engine/tests/rdma_context_reprobe_test.cpp b/mooncake-transfer-engine/tests/rdma_context_reprobe_test.cpp
new file mode 100644
index 00000000..319d57b2
--- /dev/null
+++ b/mooncake-transfer-engine/tests/rdma_context_reprobe_test.cpp
@@ -0,0 +1,144 @@
+// Copyright 2026 KVCache.AI
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include <gtest/gtest.h>
+
+#include <array>
+#include <cstring>
+#include <memory>
+#include <string>
+
+#include "common.h"
+#include "error.h"
+#include "transfer_metadata.h"
+#include "transport/rdma_transport/rdma_context.h"
+#include "transport/rdma_transport/rdma_transport.h"
+
+#if defined(__has_feature)
+#define MC_HAS_FEATURE(x) __has_feature(x)
+#else
+#define MC_HAS_FEATURE(x) 0
+#endif
+#if defined(__SANITIZE_ADDRESS__) || MC_HAS_FEATURE(address_sanitizer)
+#include <sanitizer/lsan_interface.h>
+#define MC_LSAN_IGNORE_OBJECT(p) __lsan_ignore_object(p)
+#else
+#define MC_LSAN_IGNORE_OBJECT(p) ((void)(p))
+#endif
+
+using namespace mooncake;
+
+namespace mooncake {
+
+class RdmaTransportTestPeer {
+   public:
+    static void bindMetadata(RdmaTransport &transport,
+                             std::shared_ptr<TransferMetadata> metadata,
+                             std::string local_server_name) {
+        transport.metadata_ = std::move(metadata);
+        transport.local_server_name_ = std::move(local_server_name);
+    }
+};
+
+class RdmaContextTestPeer {
+   public:
+    static void seedAutoGidState(RdmaContext &context, ibv_context *verbs_ctx,
+                                 uint8_t port, uint16_t lid, const ibv_gid &gid,
+                                 int gid_index) {
+        context.context_ = verbs_ctx;
+        context.port_ = port;
+        context.lid_ = lid;
+        context.gid_ = gid;
+        context.gid_index_ = gid_index;
+        context.auto_gid_selection_enabled_ = true;
+    }
+
+    static void disableContextForTeardown(RdmaContext &context) {
+        context.context_ = nullptr;
+    }
+};
+
+}  // namespace mooncake
+
+namespace {
+
+ibv_gid makeGid(const std::array<uint8_t, 16> &bytes) {
+    ibv_gid gid = {};
+    std::memcpy(gid.raw, bytes.data(), bytes.size());
+    return gid;
+}
+
+std::string formatGid(const std::array<uint8_t, 16> &bytes) {
+    std::string gid;
+    char buf[4] = {0};
+    for (size_t i = 0; i < bytes.size(); ++i) {
+        std::snprintf(buf, sizeof(buf), "%02x", bytes[i]);
+        gid += i == 0 ? buf : std::string(":") + buf;
+    }
+    return gid;
+}
+
+constexpr std::array<uint8_t, 16> kCurrentGid = {
+    0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
+    0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x11};
+class RdmaContextReprobeTest : public ::testing::Test {
+   protected:
+    void SetUp() override {
+        transport_ = new RdmaTransport();
+        MC_LSAN_IGNORE_OBJECT(transport_);
+        metadata_ = std::make_shared<TransferMetadata>(P2PHANDSHAKE);
+        RdmaTransportTestPeer::bindMetadata(*transport_, metadata_,
+                                            "local-rdma-segment");
+
+        auto local_desc = std::make_shared<TransferMetadata::SegmentDesc>();
+        local_desc->name = "local-rdma-segment";
+        local_desc->protocol = "rdma";
+        local_desc->devices.push_back(
+            {"synthetic0", 23, formatGid(kCurrentGid), ""});
+        ASSERT_EQ(
+            metadata_->addLocalSegment(LOCAL_SEGMENT_ID, "local-rdma-segment",
+                                       std::move(local_desc)),
+            0);
+
+        context_ = new RdmaContext(*transport_, "synthetic0");
+        MC_LSAN_IGNORE_OBJECT(context_);
+        RdmaContextTestPeer::seedAutoGidState(
+            *context_, reinterpret_cast<ibv_context *>(0x1), /*port=*/1,
+            /*lid=*/23, makeGid(kCurrentGid), /*gid_index=*/0);
+    }
+
+    std::shared_ptr<TransferMetadata::SegmentDesc> localDesc() const {
+        return metadata_->getSegmentDescByID(LOCAL_SEGMENT_ID);
+    }
+
+    RdmaTransport *transport_ = nullptr;
+    std::shared_ptr<TransferMetadata> metadata_;
+    RdmaContext *context_ = nullptr;
+};
+
+TEST_F(RdmaContextReprobeTest,
+       ReprobeStopsWhenExpectedSelectionDoesNotMatchCurrentState) {
+    auto before_desc = localDesc();
+    ASSERT_TRUE(before_desc);
+
+    bool changed = context_->reprobeAutoGid({formatGid(kCurrentGid), 9}, {});
+
+    EXPECT_FALSE(changed);
+    EXPECT_EQ(context_->gidIndex(), 0);
+    EXPECT_EQ(context_->gid(), formatGid(kCurrentGid));
+    auto after_desc = localDesc();
+    EXPECT_EQ(after_desc.get(), before_desc.get());
+}
+
+}  // namespace
diff --git a/mooncake-transfer-engine/tests/rdma_endpoint_reestablish_test.cpp b/mooncake-transfer-engine/tests/rdma_endpoint_reestablish_test.cpp
index 3498b9bf..77f58192 100644
--- a/mooncake-transfer-engine/tests/rdma_endpoint_reestablish_test.cpp
+++ b/mooncake-transfer-engine/tests/rdma_endpoint_reestablish_test.cpp
@@ -16,49 +16,169 @@
  * RDMA Endpoint Re-establishment Test
  *
  * Purpose:
- * This test verifies that TE correctly handles endpoint re-establish
- * during simulated Initiator restarts.
+ * This test verifies that TE correctly handles endpoint re-establish during
+ * simulated initiator restarts, and that classic RDMA can recover from a first
+ * RTR/EINVAL by reprobeing the next auto-selected local GID.
  *
  * How to run:
- * 1. Start etcd:
- *    etcd --listen-client-urls http://127.0.0.1:18222 \
- *      --advertise-client-urls http://127.0.0.1:18222
- *
- * 2. Run test
- *    sudo env MC_METADATA_SERVER=127.0.0.1:18222 \
- *      MC_TARGET_SERVER_NAME=127.0.0.1:12345 \
- *      MC_INITIATOR_SERVER_NAME=127.0.0.1:12346 \
- *      MC_TARGET_DEVICE_NAME=erdma_0 MC_INITIATOR_DEVICE_NAME=erdma_1 \
- *      ./build/mooncake-transfer-engine/tests/rdma_endpoint_reestablish_test
+ *   sudo env MC_METADATA_SERVER=P2PHANDSHAKE \
+ *     MC_TARGET_SERVER_NAME=127.0.0.1:12345 \
+ *     MC_INITIATOR_SERVER_NAME=127.0.0.1:12346 \
+ *     MC_TARGET_DEVICE_NAME=erdma_0 MC_INITIATOR_DEVICE_NAME=erdma_1 \
+ *     ./build/mooncake-transfer-engine/tests/rdma_endpoint_reestablish_test
  */
 
-#include <gflags/gflags.h>
-#include <glog/logging.h>
-#include <gtest/gtest.h>
-#include <numa.h>
-#include <sys/time.h>
-#include <unistd.h>
-
+#include <algorithm>
+#include <cerrno>
 #include <cstdlib>
 #include <fstream>
 #include <iomanip>
 #include <memory>
+#include <mutex>
 #include <sstream>
+#include <string>
+#include <unordered_map>
 #include <vector>
 
+#include <gflags/gflags.h>
+#include <glog/logging.h>
+#include <gtest/gtest.h>
+#include <infiniband/verbs.h>
+#include <numa.h>
+#include <sys/time.h>
+#include <unistd.h>
+
+#include "common.h"
 #include "transfer_engine.h"
 #include "transport/transport.h"
-#include "common.h"
 
 using namespace mooncake;
 
-// Size of the pre-registered memory.
-constexpr size_t kRAMBufSize = 256ull << 24;  // 256 MB
+namespace {
 
-// Actual data payload size for RDMA Read/Write.
-constexpr size_t kDataLength = 16ull << 24;  // 16MB
+constexpr size_t kRAMBufSize = 256ull << 24;
+constexpr size_t kDataLength = 16ull << 24;
 
-std::string formatDeviceNames(const std::string &device_names) {
+bool usesP2PHandshake(const std::string& metadata_server) {
+    return metadata_server == P2PHANDSHAKE;
+}
+
+std::vector<std::string> getAvailableRdmaDevices() {
+    int num_devices = 0;
+    ibv_device** device_list = ibv_get_device_list(&num_devices);
+    std::vector<std::string> devices;
+    if (device_list == nullptr) {
+        return devices;
+    }
+    devices.reserve(num_devices);
+    for (int i = 0; i < num_devices; ++i) {
+        devices.emplace_back(ibv_get_device_name(device_list[i]));
+    }
+    ibv_free_device_list(device_list);
+    return devices;
+}
+
+struct RtrFaultInjectionState {
+    std::mutex mu;
+    bool synthetic_gid_swap_enabled = false;
+    std::string synthetic_gid_device;
+    bool fail_first_rtr_einval = false;
+    std::string fail_rtr_device;
+    int injected_failures = 0;
+    std::unordered_map<std::string, std::vector<int>> rtr_sgid_history;
+    std::unordered_map<std::string, std::vector<std::string>> rtr_gid_history;
+} g_rtr_fault_injection_state;
+
+std::string formatGidBytes(const uint8_t* raw) {
+    std::ostringstream oss;
+    oss << std::hex << std::setfill('0');
+    for (size_t i = 0; i < 16; ++i) {
+        if (i != 0) {
+            oss << ":";
+        }
+        oss << std::setw(2) << static_cast<int>(raw[i]);
+    }
+    return oss.str();
+}
+
+void resetRtrFaultInjectionState() {
+    std::lock_guard<std::mutex> guard(g_rtr_fault_injection_state.mu);
+    g_rtr_fault_injection_state.synthetic_gid_swap_enabled = false;
+    g_rtr_fault_injection_state.synthetic_gid_device.clear();
+    g_rtr_fault_injection_state.fail_first_rtr_einval = false;
+    g_rtr_fault_injection_state.fail_rtr_device.clear();
+    g_rtr_fault_injection_state.injected_failures = 0;
+    g_rtr_fault_injection_state.rtr_sgid_history.clear();
+    g_rtr_fault_injection_state.rtr_gid_history.clear();
+}
+
+void configureRtrFaultInjection(const std::string& device_name) {
+    std::lock_guard<std::mutex> guard(g_rtr_fault_injection_state.mu);
+    g_rtr_fault_injection_state.synthetic_gid_swap_enabled = true;
+    g_rtr_fault_injection_state.synthetic_gid_device = device_name;
+    g_rtr_fault_injection_state.fail_first_rtr_einval = true;
+    g_rtr_fault_injection_state.fail_rtr_device = device_name;
+    g_rtr_fault_injection_state.injected_failures = 0;
+    g_rtr_fault_injection_state.rtr_sgid_history.clear();
+    g_rtr_fault_injection_state.rtr_gid_history.clear();
+}
+
+std::vector<int> getRtrSgidHistory(const std::string& device_name) {
+    std::lock_guard<std::mutex> guard(g_rtr_fault_injection_state.mu);
+    auto iter = g_rtr_fault_injection_state.rtr_sgid_history.find(device_name);
+    if (iter == g_rtr_fault_injection_state.rtr_sgid_history.end()) {
+        return {};
+    }
+    return iter->second;
+}
+
+int getInjectedFailureCount() {
+    std::lock_guard<std::mutex> guard(g_rtr_fault_injection_state.mu);
+    return g_rtr_fault_injection_state.injected_failures;
+}
+
+std::vector<std::string> getRtrGidHistory(const std::string& device_name) {
+    std::lock_guard<std::mutex> guard(g_rtr_fault_injection_state.mu);
+    auto iter = g_rtr_fault_injection_state.rtr_gid_history.find(device_name);
+    if (iter == g_rtr_fault_injection_state.rtr_gid_history.end()) {
+        return {};
+    }
+    return iter->second;
+}
+
+void recordRtrAttempt(const std::string& device_name, int sgid_index,
+                      const std::string& gid) {
+    std::lock_guard<std::mutex> guard(g_rtr_fault_injection_state.mu);
+    g_rtr_fault_injection_state.rtr_sgid_history[device_name].push_back(
+        sgid_index);
+    g_rtr_fault_injection_state.rtr_gid_history[device_name].push_back(gid);
+}
+
+int maybeSwapSyntheticGidIndex(const std::string& device_name, int gid_index) {
+    std::lock_guard<std::mutex> guard(g_rtr_fault_injection_state.mu);
+    if (!g_rtr_fault_injection_state.synthetic_gid_swap_enabled ||
+        g_rtr_fault_injection_state.synthetic_gid_device != device_name) {
+        return gid_index;
+    }
+    if (gid_index == 0) return 1;
+    if (gid_index == 1) return 0;
+    return gid_index;
+}
+
+bool shouldInjectRtrEinval(const std::string& device_name, int sgid_index) {
+    std::lock_guard<std::mutex> guard(g_rtr_fault_injection_state.mu);
+    if (!g_rtr_fault_injection_state.fail_first_rtr_einval ||
+        g_rtr_fault_injection_state.fail_rtr_device != device_name ||
+        sgid_index != 0) {
+        return false;
+    }
+    g_rtr_fault_injection_state.fail_first_rtr_einval = false;
+    g_rtr_fault_injection_state.synthetic_gid_swap_enabled = false;
+    ++g_rtr_fault_injection_state.injected_failures;
+    return true;
+}
+
+std::string formatDeviceNames(const std::string& device_names) {
     std::stringstream ss(device_names);
     std::string item;
     std::vector<std::string> tokens;
@@ -76,7 +196,7 @@ std::string formatDeviceNames(const std::string &device_names) {
     return formatted;
 }
 
-std::string makeNicPriorityMatrix(const std::string &device_name) {
+std::string makeNicPriorityMatrix(const std::string& device_name) {
     auto formatted_devices = formatDeviceNames(device_name);
     return "{\"cpu:0\": [[" + formatted_devices +
            "],[]], "
@@ -84,8 +204,8 @@ std::string makeNicPriorityMatrix(const std::string &device_name) {
            formatted_devices + "],[]]}";
 }
 
-void wait_for_transfer(TransferEngine *engine, BatchID batch_id,
-                       const std::string &op_name) {
+void waitForTransfer(TransferEngine* engine, BatchID batch_id,
+                     const std::string& op_name) {
     bool completed = false;
     TransferStatus status;
     while (!completed) {
@@ -103,27 +223,26 @@ void wait_for_transfer(TransferEngine *engine, BatchID batch_id,
 
 struct TEContext {
     std::unique_ptr<TransferEngine> engine_{};
-    uint8_t *local_addr_{};
+    uint8_t* local_addr_{};
     bool segment_opened_{false};
     SegmentHandle segment_handle_{};
     uint64_t remote_base_{};
 
-    TEContext(const std::string &local_server_name,
-              const std::string &metadata_server, const std::string &segment_id,
-              const std::string &device_name) {
+    TEContext(const std::string& local_server_name,
+              const std::string& metadata_server, const std::string& segment_id,
+              const std::string& device_name) {
         engine_ = std::make_unique<TransferEngine>(false);
         auto hostname_port = parseHostNameWithPort(local_server_name);
         engine_->init(metadata_server, local_server_name, hostname_port.first,
                       hostname_port.second);
 
         auto nic_priority_matrix = makeNicPriorityMatrix(device_name);
-        void *args[2] = {const_cast<char *>(nic_priority_matrix.c_str()),
+        void* args[2] = {const_cast<char*>(nic_priority_matrix.c_str()),
                          nullptr};
-
-        Transport *xport = engine_->installTransport("rdma", args);
+        Transport* xport = engine_->installTransport("rdma", args);
         LOG_ASSERT(xport);
 
-        local_addr_ = static_cast<uint8_t *>(numa_alloc_onnode(kRAMBufSize, 0));
+        local_addr_ = static_cast<uint8_t*>(numa_alloc_onnode(kRAMBufSize, 0));
         memset(local_addr_, 0, kDataLength);
 
         int rc =
@@ -145,16 +264,21 @@ struct TEContext {
         numa_free(local_addr_, kRAMBufSize);
         if (segment_opened_) engine_->closeSegment(segment_handle_);
     }
+
+    std::string localSegmentName() const {
+        return engine_->getLocalIpAndPort();
+    }
 };
 
 class RDMAEndpointReestablishTest : public ::testing::Test {
    protected:
     void SetUp() override {
+        resetRtrFaultInjectionState();
         google::InitGoogleLogging("RDMAEndpointReestablishTest");
         FLAGS_logtostderr = true;
 
-        const char *env = std::getenv("MC_METADATA_SERVER");
-        metadata_server = env ? env : "127.0.0.1:18222";
+        const char* env = std::getenv("MC_METADATA_SERVER");
+        metadata_server = env ? env : P2PHANDSHAKE;
         LOG(INFO) << "metadata_server: " << metadata_server;
 
         env = std::getenv("MC_TARGET_SERVER_NAME");
@@ -165,16 +289,99 @@ class RDMAEndpointReestablishTest : public ::testing::Test {
         initiator_server_name = env ? env : "127.0.0.1:12346";
         LOG(INFO) << "initiator_server_name: " << initiator_server_name;
 
+        auto devices = getAvailableRdmaDevices();
+        if (devices.size() < 2) {
+            GTEST_SKIP() << "Need at least two RDMA devices, found "
+                         << devices.size();
+        }
+
         env = std::getenv("MC_TARGET_DEVICE_NAME");
-        target_device_name = env ? env : "erdma_0";
+        target_device_name = env ? env : devices[0];
         LOG(INFO) << "target_device_name: " << target_device_name;
 
         env = std::getenv("MC_INITIATOR_DEVICE_NAME");
-        initiator_device_name = env ? env : "erdma_1";
+        initiator_device_name = env ? env : devices[1];
         LOG(INFO) << "initiator_device_name: " << initiator_device_name;
     }
 
-    void TearDown() override { google::ShutdownGoogleLogging(); }
+    void TearDown() override {
+        google::ShutdownGoogleLogging();
+        resetRtrFaultInjectionState();
+    }
+
+    void runEndpointReestablishScenario(const std::string& target_device,
+                                        const std::string& initiator_device) {
+        LOG(INFO) << "========== Setting up Target ==========";
+        TEContext target_ctx(target_server_name, metadata_server, "",
+                             target_device);
+        const std::string target_segment_name =
+            usesP2PHandshake(metadata_server) ? target_ctx.localSegmentName()
+                                              : target_server_name;
+        LOG(INFO) << "Resolved target segment name: " << target_segment_name;
+        LOG(INFO)
+            << "Target is up. Waiting for RDMA connections and operations...";
+
+        LOG(INFO) << "========== Phase 1: Start, Connect & Write ==========";
+        {
+            TEContext init_ctx(initiator_server_name, metadata_server,
+                               target_segment_name, initiator_device);
+            for (size_t i = 0; i < kDataLength; ++i) {
+                init_ctx.local_addr_[i] = static_cast<uint8_t>(i % 256);
+            }
+
+            LOG(INFO) << "Writing " << kDataLength << " bytes to Target...";
+            auto batch_id = init_ctx.engine_->allocateBatchID(1);
+            TransferRequest entry;
+            entry.opcode = TransferRequest::WRITE;
+            entry.length = kDataLength;
+            entry.source = init_ctx.local_addr_;
+            entry.target_id = init_ctx.segment_handle_;
+            entry.target_offset = init_ctx.remote_base_;
+
+            Status s = init_ctx.engine_->submitTransfer(batch_id, {entry});
+            ASSERT_EQ(s, Status::OK());
+            waitForTransfer(init_ctx.engine_.get(), batch_id, "WRITE");
+            LOG(INFO) << "Phase 1: Write Completed. Tearing down connection...";
+        }
+
+        LOG(INFO) << "Simulating Initiator Crash/Restart... Waiting 2 seconds.";
+        sleep(2);
+
+        LOG(INFO)
+            << "========== Phase 2: Restart, Re-establish Endpoint & Read "
+               "==========";
+        {
+            TEContext init_ctx(initiator_server_name, metadata_server,
+                               target_segment_name, initiator_device);
+            LOG(INFO) << "Reading data back over new Endpoint...";
+            auto batch_id = init_ctx.engine_->allocateBatchID(1);
+            TransferRequest entry;
+            entry.opcode = TransferRequest::READ;
+            entry.length = kDataLength;
+            entry.source = init_ctx.local_addr_;
+            entry.target_id = init_ctx.segment_handle_;
+            entry.target_offset = init_ctx.remote_base_;
+
+            Status s = init_ctx.engine_->submitTransfer(batch_id, {entry});
+            ASSERT_EQ(s, Status::OK());
+            waitForTransfer(init_ctx.engine_.get(), batch_id, "READ");
+
+            bool ok = true;
+            for (size_t i = 0; i < kDataLength; ++i) {
+                if (init_ctx.local_addr_[i] != static_cast<uint8_t>(i % 256)) {
+                    ok = false;
+                    LOG(ERROR) << "Data mismatch at offset " << i
+                               << ", expected " << (i % 256) << ", got "
+                               << (int)init_ctx.local_addr_[i];
+                    break;
+                }
+            }
+
+            ASSERT_TRUE(ok) << "Endpoint Reconstruction Verification Failed!";
+            LOG(INFO) << ">>> ENDPOINT RECONSTRUCTION VERIFICATION: "
+                         "\033[32mSUCCESS\033";
+        }
+    }
 
     std::string metadata_server;
     std::string target_server_name;
@@ -184,77 +391,94 @@ class RDMAEndpointReestablishTest : public ::testing::Test {
 };
 
 TEST_F(RDMAEndpointReestablishTest, EndpointReestablish) {
-    // 1. Setup Target (will stay alive until test function exits)
-    LOG(INFO) << "========== Setting up Target ==========";
-    TEContext target_ctx(target_server_name, metadata_server, "",
-                         target_device_name);
-    LOG(INFO) << "Target is up. Waiting for RDMA connections and operations...";
-
-    // 2. Phase 1: Initiator Start, Connect & Write
-    LOG(INFO) << "========== Phase 1: Start, Connect & Write ==========";
-    {
-        TEContext init_ctx(initiator_server_name, metadata_server,
-                           target_server_name, initiator_device_name);
-
-        // Fill buffer with test pattern
-        for (size_t i = 0; i < kDataLength; ++i) {
-            init_ctx.local_addr_[i] = static_cast<uint8_t>(i % 256);
-        }
-
-        LOG(INFO) << "Writing " << kDataLength << " bytes to Target...";
-        auto batch_id = init_ctx.engine_->allocateBatchID(1);
-        TransferRequest entry;
-        entry.opcode = TransferRequest::WRITE;
-        entry.length = kDataLength;
-        entry.source = init_ctx.local_addr_;
-        entry.target_id = init_ctx.segment_handle_;
-        entry.target_offset = init_ctx.remote_base_;
-
-        Status s = init_ctx.engine_->submitTransfer(batch_id, {entry});
-        ASSERT_EQ(s, Status::OK());
-
-        wait_for_transfer(init_ctx.engine_.get(), batch_id, "WRITE");
-
-        LOG(INFO) << "Phase 1: Write Completed. Tearing down connection...";
-    }
-
-    LOG(INFO) << "Simulating Initiator Crash/Restart... Waiting 2 seconds.";
-    sleep(2);
-
-    // 3. Phase 2: Restart, Re-establish Endpoint & Read
-    LOG(INFO) << "========== Phase 2: Restart, Re-establish Endpoint & Read "
-                 "==========";
-    {
-        TEContext init_ctx(initiator_server_name, metadata_server,
-                           target_server_name, initiator_device_name);
-
-        LOG(INFO) << "Reading data back over new Endpoint...";
-        auto batch_id = init_ctx.engine_->allocateBatchID(1);
-        TransferRequest entry;
-        entry.opcode = TransferRequest::READ;
-        entry.length = kDataLength;
-        entry.source = init_ctx.local_addr_;
-        entry.target_id = init_ctx.segment_handle_;
-        entry.target_offset = init_ctx.remote_base_;
-
-        Status s = init_ctx.engine_->submitTransfer(batch_id, {entry});
-        ASSERT_EQ(s, Status::OK());
-
-        wait_for_transfer(init_ctx.engine_.get(), batch_id, "READ");
-
-        bool ok = true;
-        for (size_t i = 0; i < kDataLength; ++i) {
-            if (init_ctx.local_addr_[i] != static_cast<uint8_t>(i % 256)) {
-                ok = false;
-                LOG(ERROR) << "Data mismatch at offset " << i << ", expected "
-                           << (i % 256) << ", got "
-                           << (int)init_ctx.local_addr_[i];
-                break;
-            }
-        }
-
-        ASSERT_TRUE(ok) << "Endpoint Reconstruction Verification Failed!";
-        LOG(INFO)
-            << ">>> ENDPOINT RECONSTRUCTION VERIFICATION: \033[32mSUCCESS\033";
-    }
+    runEndpointReestablishScenario(target_device_name, initiator_device_name);
+}
+
+TEST_F(RDMAEndpointReestablishTest, EndpointReestablishReverseDevices) {
+    runEndpointReestablishScenario(initiator_device_name, target_device_name);
+}
+
+TEST_F(RDMAEndpointReestablishTest, ActiveHandshakeRetriesAfterAutoGidReprobe) {
+    configureRtrFaultInjection(initiator_device_name);
+    runEndpointReestablishScenario(target_device_name, initiator_device_name);
+
+    EXPECT_EQ(getInjectedFailureCount(), 1);
+    auto sgid_history = getRtrSgidHistory(initiator_device_name);
+    auto gid_history = getRtrGidHistory(initiator_device_name);
+    ASSERT_EQ(sgid_history.size(), gid_history.size());
+    ASSERT_GE(sgid_history.size(), 2u);
+    EXPECT_EQ(sgid_history.front(), 0);
+    EXPECT_FALSE(gid_history.front().empty());
+    EXPECT_TRUE(std::any_of(
+        gid_history.begin() + 1, gid_history.end(),
+        [&](const std::string& gid) { return gid != gid_history.front(); }));
+}
+
+TEST_F(RDMAEndpointReestablishTest,
+       PassiveHandshakeRetriesAfterAutoGidReprobe) {
+    configureRtrFaultInjection(target_device_name);
+    runEndpointReestablishScenario(target_device_name, initiator_device_name);
+
+    EXPECT_EQ(getInjectedFailureCount(), 1);
+    auto sgid_history = getRtrSgidHistory(target_device_name);
+    auto gid_history = getRtrGidHistory(target_device_name);
+    ASSERT_EQ(sgid_history.size(), gid_history.size());
+    ASSERT_GE(sgid_history.size(), 2u);
+    EXPECT_EQ(sgid_history.front(), 0);
+    EXPECT_FALSE(gid_history.front().empty());
+    EXPECT_TRUE(std::any_of(
+        gid_history.begin() + 1, gid_history.end(),
+        [&](const std::string& gid) { return gid != gid_history.front(); }));
+}
+
+}  // namespace
+
+extern "C" int __real__ibv_query_gid_ex(ibv_context* context, uint8_t port_num,
+                                        int gid_index,
+                                        struct ibv_gid_entry* entry,
+                                        uint32_t flags, size_t entry_size);
+
+extern "C" int __wrap__ibv_query_gid_ex(ibv_context* context, uint8_t port_num,
+                                        int gid_index,
+                                        struct ibv_gid_entry* entry,
+                                        uint32_t flags, size_t entry_size) {
+    const std::string device_name = ibv_get_device_name(context->device);
+    int wrapped_gid_index = maybeSwapSyntheticGidIndex(device_name, gid_index);
+    return __real__ibv_query_gid_ex(context, port_num, wrapped_gid_index, entry,
+                                    flags, entry_size);
+}
+
+extern "C" int __real_ibv_query_gid(ibv_context* context, uint8_t port_num,
+                                    int gid_index, union ibv_gid* gid);
+
+extern "C" int __wrap_ibv_query_gid(ibv_context* context, uint8_t port_num,
+                                    int gid_index, union ibv_gid* gid) {
+    const std::string device_name = ibv_get_device_name(context->device);
+    int wrapped_gid_index = maybeSwapSyntheticGidIndex(device_name, gid_index);
+    return __real_ibv_query_gid(context, port_num, wrapped_gid_index, gid);
+}
+
+extern "C" int __real_ibv_modify_qp(ibv_qp* qp, ibv_qp_attr* attr,
+                                    int attr_mask);
+
+extern "C" int __wrap_ibv_modify_qp(ibv_qp* qp, ibv_qp_attr* attr,
+                                    int attr_mask) {
+    if (qp != nullptr && attr != nullptr && attr->qp_state == IBV_QPS_RTR &&
+        (attr_mask & IBV_QP_AV)) {
+        const std::string device_name =
+            ibv_get_device_name(qp->context->device);
+        int sgid_index = attr->ah_attr.grh.sgid_index;
+        union ibv_gid actual_gid = {};
+        std::string gid_string;
+        if (__real_ibv_query_gid(qp->context, attr->ah_attr.port_num,
+                                 sgid_index, &actual_gid) == 0) {
+            gid_string = formatGidBytes(actual_gid.raw);
+        }
+        recordRtrAttempt(device_name, sgid_index, gid_string);
+        if (shouldInjectRtrEinval(device_name, sgid_index)) {
+            errno = EINVAL;
+            return -1;
+        }
+    }
+    return __real_ibv_modify_qp(qp, attr, attr_mask);
 }
diff --git a/mooncake-transfer-engine/tests/rdma_gid_probe_test.cpp b/mooncake-transfer-engine/tests/rdma_gid_probe_test.cpp
new file mode 100644
index 00000000..740222bc
--- /dev/null
+++ b/mooncake-transfer-engine/tests/rdma_gid_probe_test.cpp
@@ -0,0 +1,402 @@
+// Copyright 2026 KVCache.AI
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include <gtest/gtest.h>
+
+#include <vector>
+
+#include "transport/rdma_transport/rdma_gid_probe.h"
+
+using namespace mooncake;
+
+namespace {
+
+AutoGidCandidate makeCandidate(int gid_index, uint32_t gid_type,
+                               bool has_network_device, bool is_ipv4_mapped,
+                               bool is_link_local_ipv6,
+                               bool is_overlay_network = false,
+                               bool is_overlay_ipv4 = false,
+                               bool is_null_gid = false,
+                               bool query_succeeded = true,
+                               std::string gid = "") {
+    AutoGidCandidate candidate;
+    candidate.gid_index = gid_index;
+    candidate.gid =
+        gid.empty() ? "gid-" + std::to_string(gid_index) : std::move(gid);
+    candidate.gid_type = gid_type;
+    candidate.has_network_device = has_network_device;
+    candidate.is_ipv4_mapped = is_ipv4_mapped;
+    candidate.is_link_local_ipv6 = is_link_local_ipv6;
+    candidate.is_overlay_network = is_overlay_network;
+    candidate.is_overlay_ipv4 = is_overlay_ipv4;
+    candidate.is_null_gid = is_null_gid;
+    candidate.query_succeeded = query_succeeded;
+    return candidate;
+}
+
+TEST(RdmaGidProbeTest, PrefersNetworkBackedRoutableCandidate) {
+    std::vector<AutoGidCandidate> candidates = {
+        makeCandidate(/*gid_index=*/0, IBV_GID_TYPE_ROCE_V2,
+                      /*has_network_device=*/false,
+                      /*is_ipv4_mapped=*/true,
+                      /*is_link_local_ipv6=*/false),
+        makeCandidate(/*gid_index=*/1, IBV_GID_TYPE_ROCE_V2,
+                      /*has_network_device=*/true,
+                      /*is_ipv4_mapped=*/true,
+                      /*is_link_local_ipv6=*/false),
+    };
+
+    auto selection = selectBestAutoGidCandidate(candidates);
+    ASSERT_TRUE(selection.has_value());
+    EXPECT_EQ(selection->gid_index, 1);
+    EXPECT_EQ(selection->candidate_class,
+              AutoGidCandidateClass::kNetworkRoutable);
+}
+
+TEST(RdmaGidProbeTest, DemotesLinkLocalBehindRoutableNetworkCandidate) {
+    std::vector<AutoGidCandidate> candidates = {
+        makeCandidate(/*gid_index=*/0, IBV_GID_TYPE_ROCE_V2,
+                      /*has_network_device=*/true,
+                      /*is_ipv4_mapped=*/false,
+                      /*is_link_local_ipv6=*/true),
+        makeCandidate(/*gid_index=*/1, IBV_GID_TYPE_ROCE_V2,
+                      /*has_network_device=*/true,
+                      /*is_ipv4_mapped=*/true,
+                      /*is_link_local_ipv6=*/false),
+    };
+
+    auto selection = selectBestAutoGidCandidate(candidates);
+    ASSERT_TRUE(selection.has_value());
+    EXPECT_EQ(selection->gid_index, 1);
+    EXPECT_EQ(selection->candidate_class,
+              AutoGidCandidateClass::kNetworkRoutable);
+}
+
+TEST(RdmaGidProbeTest, DemotesOverlayCandidateBehindNormalNetworkCandidate) {
+    std::vector<AutoGidCandidate> candidates = {
+        makeCandidate(/*gid_index=*/0, IBV_GID_TYPE_ROCE_V2,
+                      /*has_network_device=*/true,
+                      /*is_ipv4_mapped=*/true,
+                      /*is_link_local_ipv6=*/false,
+                      /*is_overlay_network=*/true),
+        makeCandidate(/*gid_index=*/1, IBV_GID_TYPE_ROCE_V2,
+                      /*has_network_device=*/true,
+                      /*is_ipv4_mapped=*/true,
+                      /*is_link_local_ipv6=*/false),
+    };
+
+    auto selection = selectBestAutoGidCandidate(candidates);
+    ASSERT_TRUE(selection.has_value());
+    EXPECT_EQ(selection->gid_index, 1);
+}
+
+TEST(RdmaGidProbeTest,
+     PrefersNoNetworkRoutableOverDegradedNetworkBackedCandidate) {
+    std::vector<AutoGidCandidate> candidates = {
+        makeCandidate(/*gid_index=*/0, IBV_GID_TYPE_ROCE_V2,
+                      /*has_network_device=*/true,
+                      /*is_ipv4_mapped=*/false,
+                      /*is_link_local_ipv6=*/true),
+        makeCandidate(/*gid_index=*/1, IBV_GID_TYPE_ROCE_V2,
+                      /*has_network_device=*/false,
+                      /*is_ipv4_mapped=*/true,
+                      /*is_link_local_ipv6=*/false),
+    };
+
+    auto selection = selectBestAutoGidCandidate(candidates);
+    ASSERT_TRUE(selection.has_value());
+    EXPECT_EQ(selection->gid_index, 1);
+    EXPECT_EQ(selection->candidate_class,
+              AutoGidCandidateClass::kNoNetworkRoutable);
+}
+
+TEST(RdmaGidProbeTest, KeepsNoNetworkFallbackAsLastResort) {
+    std::vector<AutoGidCandidate> candidates = {
+        makeCandidate(/*gid_index=*/3, IBV_GID_TYPE_ROCE_V2,
+                      /*has_network_device=*/false,
+                      /*is_ipv4_mapped=*/true,
+                      /*is_link_local_ipv6=*/false),
+    };
+
+    auto selection = selectBestAutoGidCandidate(candidates);
+    ASSERT_TRUE(selection.has_value());
+    EXPECT_EQ(selection->gid_index, 3);
+    EXPECT_EQ(selection->candidate_class,
+              AutoGidCandidateClass::kNoNetworkRoutable);
+}
+
+TEST(RdmaGidProbeTest, FallsBackToFirstNonzeroCandidateWhenNeeded) {
+    std::vector<AutoGidCandidate> candidates = {
+        makeCandidate(/*gid_index=*/0, IBV_GID_TYPE_ROCE_V1,
+                      /*has_network_device=*/false,
+                      /*is_ipv4_mapped=*/false,
+                      /*is_link_local_ipv6=*/false,
+                      /*is_overlay_network=*/false,
+                      /*is_overlay_ipv4=*/false,
+                      /*is_null_gid=*/false),
+        makeCandidate(/*gid_index=*/2, IBV_GID_TYPE_ROCE_V1,
+                      /*has_network_device=*/false,
+                      /*is_ipv4_mapped=*/false,
+                      /*is_link_local_ipv6=*/false,
+                      /*is_overlay_network=*/false,
+                      /*is_overlay_ipv4=*/false,
+                      /*is_null_gid=*/false),
+    };
+
+    auto selection = selectBestAutoGidCandidate(candidates);
+    ASSERT_TRUE(selection.has_value());
+    EXPECT_EQ(selection->gid_index, 0);
+    EXPECT_EQ(selection->candidate_class,
+              AutoGidCandidateClass::kFallbackNonzero);
+}
+
+TEST(RdmaGidProbeTest, DoesNotTreatIbCandidateAsLinkLocalIpv6Penalty) {
+    std::vector<AutoGidCandidate> candidates = {
+        makeCandidate(/*gid_index=*/0, IBV_GID_TYPE_IB,
+                      /*has_network_device=*/true,
+                      /*is_ipv4_mapped=*/false,
+                      /*is_link_local_ipv6=*/true),
+        makeCandidate(/*gid_index=*/1, IBV_GID_TYPE_ROCE_V2,
+                      /*has_network_device=*/false,
+                      /*is_ipv4_mapped=*/true,
+                      /*is_link_local_ipv6=*/false),
+    };
+
+    auto selection = selectBestAutoGidCandidate(candidates);
+    ASSERT_TRUE(selection.has_value());
+    EXPECT_EQ(selection->gid_index, 0);
+    EXPECT_EQ(selection->candidate_class,
+              AutoGidCandidateClass::kNetworkRoutable);
+}
+
+TEST(RdmaGidProbeTest, SkipsInvalidAndNullCandidates) {
+    std::vector<AutoGidCandidate> candidates = {
+        makeCandidate(/*gid_index=*/0, IBV_GID_TYPE_ROCE_V2,
+                      /*has_network_device=*/true,
+                      /*is_ipv4_mapped=*/true,
+                      /*is_link_local_ipv6=*/false,
+                      /*is_overlay_network=*/false,
+                      /*is_overlay_ipv4=*/false,
+                      /*is_null_gid=*/false,
+                      /*query_succeeded=*/false),
+        makeCandidate(/*gid_index=*/1, IBV_GID_TYPE_ROCE_V2,
+                      /*has_network_device=*/true,
+                      /*is_ipv4_mapped=*/true,
+                      /*is_link_local_ipv6=*/false,
+                      /*is_overlay_network=*/false,
+                      /*is_overlay_ipv4=*/false,
+                      /*is_null_gid=*/true),
+        makeCandidate(/*gid_index=*/2, IBV_GID_TYPE_ROCE_V2,
+                      /*has_network_device=*/true,
+                      /*is_ipv4_mapped=*/true,
+                      /*is_link_local_ipv6=*/false),
+    };
+
+    auto selection = selectBestAutoGidCandidate(candidates);
+    ASSERT_TRUE(selection.has_value());
+    EXPECT_EQ(selection->gid_index, 2);
+}
+
+TEST(RdmaGidProbeTest, KeepsStableOrderingWithinSameCandidateClass) {
+    std::vector<AutoGidCandidate> candidates = {
+        makeCandidate(/*gid_index=*/3, IBV_GID_TYPE_ROCE_V2,
+                      /*has_network_device=*/true,
+                      /*is_ipv4_mapped=*/true,
+                      /*is_link_local_ipv6=*/false),
+        makeCandidate(/*gid_index=*/1, IBV_GID_TYPE_ROCE_V2,
+                      /*has_network_device=*/true,
+                      /*is_ipv4_mapped=*/true,
+                      /*is_link_local_ipv6=*/false),
+    };
+
+    auto ranked = rankAutoGidCandidates(candidates);
+    ASSERT_EQ(ranked.size(), 2u);
+    EXPECT_EQ(ranked[0].gid_index, 1);
+    EXPECT_EQ(ranked[1].gid_index, 3);
+}
+
+TEST(RdmaGidProbeTest, ReprobeStillPicksBestCandidateFromFreshSnapshot) {
+    std::vector<AutoGidCandidate> candidates = {
+        makeCandidate(/*gid_index=*/1, IBV_GID_TYPE_ROCE_V2,
+                      /*has_network_device=*/true,
+                      /*is_ipv4_mapped=*/true,
+                      /*is_link_local_ipv6=*/false),
+        makeCandidate(/*gid_index=*/3, IBV_GID_TYPE_ROCE_V2,
+                      /*has_network_device=*/false,
+                      /*is_ipv4_mapped=*/true,
+                      /*is_link_local_ipv6=*/false),
+    };
+
+    auto selection = selectBestAutoGidCandidate(candidates);
+    ASSERT_TRUE(selection.has_value());
+    EXPECT_EQ(selection->gid_index, 1);
+    EXPECT_EQ(selection->candidate_class,
+              AutoGidCandidateClass::kNetworkRoutable);
+}
+
+TEST(RdmaGidProbeTest, ReprobeDetectsSameIndexGidRefresh) {
+    std::vector<AutoGidCandidate> candidates = {
+        makeCandidate(/*gid_index=*/1, IBV_GID_TYPE_ROCE_V2,
+                      /*has_network_device=*/true,
+                      /*is_ipv4_mapped=*/true,
+                      /*is_link_local_ipv6=*/false,
+                      /*is_overlay_network=*/false,
+                      /*is_overlay_ipv4=*/false,
+                      /*is_null_gid=*/false,
+                      /*query_succeeded=*/true,
+                      /*gid=*/"00:11:22"),
+        makeCandidate(/*gid_index=*/3, IBV_GID_TYPE_ROCE_V2,
+                      /*has_network_device=*/false,
+                      /*is_ipv4_mapped=*/true,
+                      /*is_link_local_ipv6=*/false),
+    };
+
+    auto selection = reselectAutoGidCandidate(
+        candidates, /*current_gid_index=*/1, /*current_gid=*/"00:11:21");
+    ASSERT_TRUE(selection.has_value());
+    EXPECT_EQ(selection->gid_index, 1);
+    EXPECT_EQ(selection->gid, "00:11:22");
+    EXPECT_EQ(selection->candidate_class,
+              AutoGidCandidateClass::kNetworkRoutable);
+}
+
+TEST(RdmaGidProbeTest, ReprobeSkipsRetryWhenBestSelectionDidNotChange) {
+    std::vector<AutoGidCandidate> candidates = {
+        makeCandidate(/*gid_index=*/1, IBV_GID_TYPE_ROCE_V2,
+                      /*has_network_device=*/true,
+                      /*is_ipv4_mapped=*/true,
+                      /*is_link_local_ipv6=*/false,
+                      /*is_overlay_network=*/false,
+                      /*is_overlay_ipv4=*/false,
+                      /*is_null_gid=*/false,
+                      /*query_succeeded=*/true,
+                      /*gid=*/"00:11:22"),
+    };
+
+    auto selection = reselectAutoGidCandidate(
+        candidates, /*current_gid_index=*/1, /*current_gid=*/"00:11:22");
+    EXPECT_FALSE(selection.has_value());
+}
+
+TEST(RdmaGidProbeTest, ReprobeSkipsAlreadyTriedCandidates) {
+    std::vector<AutoGidCandidate> candidates = {
+        makeCandidate(/*gid_index=*/1, IBV_GID_TYPE_ROCE_V2,
+                      /*has_network_device=*/true,
+                      /*is_ipv4_mapped=*/true,
+                      /*is_link_local_ipv6=*/false,
+                      /*is_overlay_network=*/false,
+                      /*is_overlay_ipv4=*/false,
+                      /*is_null_gid=*/false,
+                      /*query_succeeded=*/true,
+                      /*gid=*/"00:11:22"),
+        makeCandidate(/*gid_index=*/3, IBV_GID_TYPE_ROCE_V2,
+                      /*has_network_device=*/false,
+                      /*is_ipv4_mapped=*/true,
+                      /*is_link_local_ipv6=*/false,
+                      /*is_overlay_network=*/false,
+                      /*is_overlay_ipv4=*/false,
+                      /*is_null_gid=*/false,
+                      /*query_succeeded=*/true,
+                      /*gid=*/"00:11:33"),
+    };
+
+    std::vector<AutoGidSelectionIdentity> tried = {
+        {1, "00:11:22"},
+    };
+    auto selection = reselectAutoGidCandidate(
+        candidates, /*current_gid_index=*/1, /*current_gid=*/"00:11:21", tried);
+    ASSERT_TRUE(selection.has_value());
+    EXPECT_EQ(selection->gid_index, 3);
+    EXPECT_EQ(selection->gid, "00:11:33");
+}
+
+TEST(RdmaGidProbeTest, DetectsSameIndexGidByteChangesAsSelectionChanges) {
+    EXPECT_TRUE(didAutoGidSelectionChange(/*previous_gid_index=*/1,
+                                          /*previous_gid=*/"00:11:22",
+                                          /*current_gid_index=*/1,
+                                          /*current_gid=*/"00:11:23"));
+
+    EXPECT_FALSE(didAutoGidSelectionChange(/*previous_gid_index=*/1,
+                                           /*previous_gid=*/"00:11:22",
+                                           /*current_gid_index=*/1,
+                                           /*current_gid=*/"00:11:22"));
+}
+
+TEST(RdmaGidProbeTest, HandshakeRetryRespectsConfiguredRetryBudget) {
+    EXPECT_TRUE(shouldAttemptAutoGidHandshakeRetry(
+        /*auto_gid_selection_enabled=*/true,
+        /*retry_count=*/0,
+        /*max_retries=*/2,
+        /*failure_happened_at_rtr=*/true, EINVAL));
+
+    EXPECT_TRUE(shouldAttemptAutoGidHandshakeRetry(
+        /*auto_gid_selection_enabled=*/true,
+        /*retry_count=*/1,
+        /*max_retries=*/2,
+        /*failure_happened_at_rtr=*/true, EINVAL));
+
+    EXPECT_FALSE(shouldAttemptAutoGidHandshakeRetry(
+        /*auto_gid_selection_enabled=*/false,
+        /*retry_count=*/0,
+        /*max_retries=*/2,
+        /*failure_happened_at_rtr=*/true, EINVAL));
+
+    EXPECT_FALSE(shouldAttemptAutoGidHandshakeRetry(
+        /*auto_gid_selection_enabled=*/true,
+        /*retry_count=*/2,
+        /*max_retries=*/2,
+        /*failure_happened_at_rtr=*/true, EINVAL));
+
+    EXPECT_FALSE(shouldAttemptAutoGidHandshakeRetry(
+        /*auto_gid_selection_enabled=*/true,
+        /*retry_count=*/0,
+        /*max_retries=*/0,
+        /*failure_happened_at_rtr=*/true, EINVAL));
+}
+
+TEST(RdmaGidProbeTest, HandshakeRetryOnlyTriggersForRtrEinval) {
+    EXPECT_FALSE(shouldAttemptAutoGidHandshakeRetry(
+        /*auto_gid_selection_enabled=*/true,
+        /*retry_count=*/0,
+        /*max_retries=*/2,
+        /*failure_happened_at_rtr=*/false, EINVAL));
+
+    EXPECT_FALSE(shouldAttemptAutoGidHandshakeRetry(
+        /*auto_gid_selection_enabled=*/true,
+        /*retry_count=*/0,
+        /*max_retries=*/2,
+        /*failure_happened_at_rtr=*/true, ENOENT));
+}
+
+TEST(RdmaGidProbeTest, RetryActionRequiresObservedOrReprobedChange) {
+    EXPECT_EQ(decideAutoGidRetryAction(
+                  /*reprobe_changed=*/false, /*previous_gid_index=*/1,
+                  /*previous_gid=*/"00:11:22", /*current_gid_index=*/1,
+                  /*current_gid=*/"00:11:22"),
+              AutoGidRetryAction::kDoNotRetry);
+
+    EXPECT_EQ(decideAutoGidRetryAction(
+                  /*reprobe_changed=*/true, /*previous_gid_index=*/1,
+                  /*previous_gid=*/"00:11:22", /*current_gid_index=*/1,
+                  /*current_gid=*/"00:11:23"),
+              AutoGidRetryAction::kRetryWithReprobedGid);
+
+    EXPECT_EQ(decideAutoGidRetryAction(
+                  /*reprobe_changed=*/false, /*previous_gid_index=*/1,
+                  /*previous_gid=*/"00:11:22", /*current_gid_index=*/1,
+                  /*current_gid=*/"00:11:23"),
+              AutoGidRetryAction::kRetryWithObservedChange);
+}
+
+}  // namespace
diff --git a/mooncake-transfer-engine/tests/rdma_transport_test.cpp b/mooncake-transfer-engine/tests/rdma_transport_test.cpp
index e05596e6..6c6e747b 100644
--- a/mooncake-transfer-engine/tests/rdma_transport_test.cpp
+++ b/mooncake-transfer-engine/tests/rdma_transport_test.cpp
@@ -48,8 +48,9 @@
 #include <cufile.h>
 #endif
 
-#if defined(USE_CUDA) || defined(USE_MUSA) || defined(USE_HIP) || \
-    defined(USE_MLU) || defined(USE_MACA)
+#if defined(USE_CUDA) || defined(USE_MUSA) || defined(USE_HIP) ||  \
+    defined(USE_MLU) || defined(USE_MACA) || defined(USE_HYGON) || \
+    defined(USE_COREX)
 
 #include <cassert>
 
@@ -100,7 +101,7 @@ std::string pickBackend() {
 #if defined(USE_MLU)
     return "mlu";
 #elif defined(USE_CUDA) || defined(USE_MUSA) || defined(USE_HIP) || \
-    defined(USE_MACA)
+    defined(USE_MACA) || defined(USE_HYGON) || defined(USE_COREX)
     return FLAGS_use_vram ? "gpu" : "cpu";
 #else
     return "cpu";
@@ -112,7 +113,7 @@ int pickDevId(const std::string &backend) {
         return FLAGS_device_id;
     }
 #if defined(USE_CUDA) || defined(USE_MUSA) || defined(USE_HIP) || \
-    defined(USE_MACA)
+    defined(USE_MACA) || defined(USE_HYGON) || defined(USE_COREX)
     if (backend == "gpu") {
         return FLAGS_gpu_id;
     }
@@ -127,7 +128,7 @@ void validateBackend(const std::string &backend) {
         return;
     }
 #if defined(USE_CUDA) || defined(USE_MUSA) || defined(USE_HIP) || \
-    defined(USE_MACA)
+    defined(USE_MACA) || defined(USE_HYGON) || defined(USE_COREX)
     if (backend == "gpu") {
         return;
     }
@@ -170,12 +171,13 @@ bool validateTransferSizes() {
     return true;
 }
 
-void setBackendDevice(const std::string &backend) {
+[[maybe_unused]] void setBackendDevice(const std::string &backend) {
     if (!usesDeviceMemory(backend)) {
         return;
     }
-#if defined(USE_CUDA) || defined(USE_MUSA) || defined(USE_HIP) || \
-    defined(USE_MLU) || defined(USE_MACA)
+#if defined(USE_CUDA) || defined(USE_MUSA) || defined(USE_HIP) ||  \
+    defined(USE_MLU) || defined(USE_MACA) || defined(USE_HYGON) || \
+    defined(USE_COREX)
     checkCudaError(cudaSetDevice(pickDevId(backend)), "Failed to set device");
 #else
     LOG(FATAL) << "Device memory backend is not available in this build";
@@ -185,8 +187,9 @@ void setBackendDevice(const std::string &backend) {
 void *allocateMemoryPool(size_t size, int socket_id,
                          const std::string &backend) {
     if (usesDeviceMemory(backend)) {
-#if defined(USE_CUDA) || defined(USE_MUSA) || defined(USE_HIP) || \
-    defined(USE_MLU) || defined(USE_MACA)
+#if defined(USE_CUDA) || defined(USE_MUSA) || defined(USE_HIP) ||  \
+    defined(USE_MLU) || defined(USE_MACA) || defined(USE_HYGON) || \
+    defined(USE_COREX)
         setBackendDevice(backend);
         void *d_buf = nullptr;
         checkCudaError(cudaMalloc(&d_buf, size),
@@ -202,8 +205,9 @@ void *allocateMemoryPool(size_t size, int socket_id,
 
 void freeMemoryPool(void *addr, size_t size, const std::string &backend) {
     if (usesDeviceMemory(backend)) {
-#if defined(USE_CUDA) || defined(USE_MUSA) || defined(USE_HIP) || \
-    defined(USE_MLU) || defined(USE_MACA)
+#if defined(USE_CUDA) || defined(USE_MUSA) || defined(USE_HIP) ||  \
+    defined(USE_MLU) || defined(USE_MACA) || defined(USE_HYGON) || \
+    defined(USE_COREX)
         cudaFree(addr);
         return;
 #else
@@ -216,8 +220,9 @@ void freeMemoryPool(void *addr, size_t size, const std::string &backend) {
 void copyFromHost(void *dst, const void *src, size_t size,
                   const std::string &backend) {
     if (usesDeviceMemory(backend)) {
-#if defined(USE_CUDA) || defined(USE_MUSA) || defined(USE_HIP) || \
-    defined(USE_MLU) || defined(USE_MACA)
+#if defined(USE_CUDA) || defined(USE_MUSA) || defined(USE_HIP) ||  \
+    defined(USE_MLU) || defined(USE_MACA) || defined(USE_HYGON) || \
+    defined(USE_COREX)
         checkCudaError(cudaMemcpy(dst, src, size, cudaMemcpyHostToDevice),
                        "Failed to copy host data to device");
         return;
@@ -231,8 +236,9 @@ void copyFromHost(void *dst, const void *src, size_t size,
 void copyToHost(void *dst, const void *src, size_t size,
                 const std::string &backend) {
     if (usesDeviceMemory(backend)) {
-#if defined(USE_CUDA) || defined(USE_MUSA) || defined(USE_HIP) || \
-    defined(USE_MLU) || defined(USE_MACA)
+#if defined(USE_CUDA) || defined(USE_MUSA) || defined(USE_HIP) ||  \
+    defined(USE_MLU) || defined(USE_MACA) || defined(USE_HYGON) || \
+    defined(USE_COREX)
         checkCudaError(cudaMemcpy(dst, src, size, cudaMemcpyDeviceToHost),
                        "Failed to copy device data to host");
         return;
diff --git a/mooncake-transfer-engine/tests/sunrise_link_transport_runtime_test.cpp b/mooncake-transfer-engine/tests/sunrise_link_transport_runtime_test.cpp
new file mode 100644
index 00000000..04c2abbd
--- /dev/null
+++ b/mooncake-transfer-engine/tests/sunrise_link_transport_runtime_test.cpp
@@ -0,0 +1,131 @@
+#include <gtest/gtest.h>
+
+#include <memory>
+
+#include "gpu_vendor/sunrise.h"
+#include "transfer_engine.h"
+#include "transport/transport.h"
+
+using namespace mooncake;
+
+namespace {
+
+static tangError_t allocOnDevice(size_t size, void** buffer) {
+    *buffer = nullptr;
+    if (tangSetDevice(0) != tangSuccess) return tangErrorInvalidDevice;
+    return tangMalloc(buffer, size);
+}
+
+}  // namespace
+
+TEST(SunriseLinkTransportRuntimeTest,
+     InstallDoesNotBreakFirstDeviceAllocation) {
+    int gpu_count = 0;
+    if (tangGetDeviceCount(&gpu_count) != tangSuccess || gpu_count <= 0) {
+        GTEST_SKIP() << "Sunrise device is unavailable";
+    }
+
+    auto engine = std::make_unique<TransferEngine>(false);
+    engine->init("P2PHANDSHAKE", "sunrise_runtime_test:12345");
+
+    Transport* transport = engine->installTransport("sunrise_link", nullptr);
+    ASSERT_NE(transport, nullptr);
+
+    ASSERT_EQ(tangSetDevice(0), tangSuccess);
+    void* buffer = nullptr;
+    tangError_t ret = tangMalloc(&buffer, 4096);
+    ASSERT_EQ(ret, tangSuccess) << tangGetErrorString(ret);
+    ASSERT_NE(buffer, nullptr);
+
+    ASSERT_EQ(tangFree(buffer), tangSuccess);
+}
+
+TEST(SunriseLinkTransportRuntimeTest,
+     InstallDoesNotChangeLargeAllocationBehavior) {
+    int gpu_count = 0;
+    if (tangGetDeviceCount(&gpu_count) != tangSuccess || gpu_count <= 0) {
+        GTEST_SKIP() << "Sunrise device is unavailable";
+    }
+
+    void* before = nullptr;
+    tangError_t before_ret = allocOnDevice(1 << 20, &before);
+    if (before_ret == tangSuccess && before != nullptr) {
+        ASSERT_EQ(tangFree(before), tangSuccess);
+    }
+
+    auto engine = std::make_unique<TransferEngine>(false);
+    engine->init("P2PHANDSHAKE", "sunrise_runtime_test:12346");
+    Transport* transport = engine->installTransport("sunrise_link", nullptr);
+    ASSERT_NE(transport, nullptr);
+
+    void* after = nullptr;
+    tangError_t after_ret = allocOnDevice(1 << 20, &after);
+    EXPECT_EQ(after_ret, before_ret)
+        << "before=" << tangGetErrorString(before_ret)
+        << " after=" << tangGetErrorString(after_ret);
+    if (after_ret == tangSuccess && after != nullptr) {
+        ASSERT_EQ(tangFree(after), tangSuccess);
+    }
+}
+
+TEST(SunriseLinkTransportRuntimeTest, EngineInitDoesNotBreakDeviceState) {
+    int gpu_count = 0;
+    if (tangGetDeviceCount(&gpu_count) != tangSuccess || gpu_count <= 0) {
+        GTEST_SKIP() << "Sunrise device is unavailable";
+    }
+
+    auto engine = std::make_unique<TransferEngine>(false);
+    engine->init("P2PHANDSHAKE", "sunrise_runtime_test:12347");
+
+    int current_dev = -1;
+    tangError_t set_ret = tangSetDevice(0);
+    tangError_t get_ret = tangGetDevice(&current_dev);
+    EXPECT_EQ(set_ret, tangSuccess) << tangGetErrorString(set_ret);
+    EXPECT_EQ(get_ret, tangSuccess) << tangGetErrorString(get_ret);
+    EXPECT_EQ(current_dev, 0);
+}
+
+TEST(SunriseLinkTransportRuntimeTest,
+     MetadataConstructionDoesNotBreakDeviceState) {
+    int gpu_count = 0;
+    if (tangGetDeviceCount(&gpu_count) != tangSuccess || gpu_count <= 0) {
+        GTEST_SKIP() << "Sunrise device is unavailable";
+    }
+
+    auto metadata = std::make_shared<TransferMetadata>("P2PHANDSHAKE");
+    ASSERT_NE(metadata, nullptr);
+
+    int current_dev = -1;
+    tangError_t set_ret = tangSetDevice(0);
+    tangError_t get_ret = tangGetDevice(&current_dev);
+    EXPECT_EQ(set_ret, tangSuccess) << tangGetErrorString(set_ret);
+    EXPECT_EQ(get_ret, tangSuccess) << tangGetErrorString(get_ret);
+    EXPECT_EQ(current_dev, 0);
+}
+
+TEST(SunriseLinkTransportRuntimeTest,
+     DestroyingOneTransportDoesNotBreakAnother) {
+    int gpu_count = 0;
+    if (tangGetDeviceCount(&gpu_count) != tangSuccess || gpu_count <= 0) {
+        GTEST_SKIP() << "Sunrise device is unavailable";
+    }
+
+    auto engine1 = std::make_unique<TransferEngine>(false);
+    engine1->init("P2PHANDSHAKE", "sunrise_runtime_test:12350");
+    Transport* transport1 = engine1->installTransport("sunrise_link", nullptr);
+    ASSERT_NE(transport1, nullptr);
+
+    {
+        auto engine2 = std::make_unique<TransferEngine>(false);
+        engine2->init("P2PHANDSHAKE", "sunrise_runtime_test:12351");
+        Transport* transport2 =
+            engine2->installTransport("sunrise_link", nullptr);
+        ASSERT_NE(transport2, nullptr);
+    }
+
+    ASSERT_EQ(tangSetDevice(0), tangSuccess);
+    void* buffer = nullptr;
+    tangError_t ret = tangMalloc(&buffer, 4096);
+    EXPECT_EQ(ret, tangSuccess) << tangGetErrorString(ret);
+    if (buffer) tangFree(buffer);
+}
diff --git a/mooncake-transfer-engine/tests/sunrise_link_transport_test.cpp b/mooncake-transfer-engine/tests/sunrise_link_transport_test.cpp
new file mode 100644
index 00000000..57541402
--- /dev/null
+++ b/mooncake-transfer-engine/tests/sunrise_link_transport_test.cpp
@@ -0,0 +1,136 @@
+#include <gflags/gflags.h>
+#include <glog/logging.h>
+#include <gtest/gtest.h>
+
+#include <memory>
+#include <vector>
+
+#include "gpu_vendor/sunrise.h"
+#include "transfer_engine.h"
+#include "transport/transport.h"
+
+using namespace mooncake;
+
+DEFINE_string(metadata_server, "127.0.0.1:2379", "etcd server host address");
+DEFINE_string(local_server_name, "sunrise_server:12345", "Local server name");
+DEFINE_string(segment_id, "sunrise_server:12345", "Segment ID to access data");
+DEFINE_int32(gpu_id, 0, "Sunrise device ID to use");
+
+static void checkSunriseError(tangError_t result, const char* message) {
+    if (result != tangSuccess) {
+        LOG(ERROR) << message << " (Error code: " << result << " - "
+                   << tangGetErrorString(result) << ")";
+        std::exit(EXIT_FAILURE);
+    }
+}
+
+static void* allocateSunriseBuffer(size_t size, int gpu_id) {
+    checkSunriseError(tangSetDevice(gpu_id), "Failed to set device");
+    void* buffer = nullptr;
+    checkSunriseError(tangMalloc(&buffer, size),
+                      "Failed to allocate Sunrise device memory");
+    return buffer;
+}
+
+static void freeSunriseBuffer(void* addr) {
+    checkSunriseError(tangFree(addr), "Failed to free Sunrise device memory");
+}
+
+TEST(SunriseLinkTransportTest, WriteAndRead) {
+    const size_t kDataLength = 4096000;
+    int gpu_id = FLAGS_gpu_id;
+
+    auto server_engine = std::make_unique<TransferEngine>(false);
+    server_engine->init(FLAGS_metadata_server, FLAGS_local_server_name);
+    Transport* server_transport =
+        server_engine->installTransport("sunrise_link", nullptr);
+    ASSERT_NE(server_transport, nullptr);
+
+    void* server_buffer = allocateSunriseBuffer(kDataLength * 2, gpu_id);
+    int rc = server_engine->registerLocalMemory(
+        server_buffer, kDataLength * 2, "cuda:" + std::to_string(gpu_id));
+    ASSERT_EQ(rc, 0);
+
+    auto segment_id = server_engine->openSegment(FLAGS_segment_id);
+
+    auto client_engine = std::make_unique<TransferEngine>(false);
+    client_engine->init(FLAGS_metadata_server, "sunrise_client:12346");
+    Transport* client_transport =
+        client_engine->installTransport("sunrise_link", nullptr);
+    ASSERT_NE(client_transport, nullptr);
+
+    void* client_buffer = allocateSunriseBuffer(kDataLength * 2, gpu_id);
+    rc = client_engine->registerLocalMemory(client_buffer, kDataLength * 2,
+                                            "cuda:" + std::to_string(gpu_id));
+    ASSERT_EQ(rc, 0);
+
+    {
+        std::vector<char> host_data(kDataLength, 'A');
+        checkSunriseError(tangMemcpy(client_buffer, host_data.data(),
+                                     kDataLength, tangMemcpyHostToDevice),
+                          "Memcpy to client_buffer");
+
+        auto batch_id = client_engine->allocateBatchID(1);
+        TransferRequest entry;
+        entry.opcode = TransferRequest::WRITE;
+        entry.length = kDataLength;
+        entry.source = client_buffer;
+        entry.target_id = segment_id;
+        entry.target_offset = reinterpret_cast<uint64_t>(server_buffer);
+        Status s = client_engine->submitTransfer(batch_id, {entry});
+        ASSERT_TRUE(s.ok());
+
+        TransferStatus status;
+        do {
+            s = client_engine->getTransferStatus(batch_id, 0, status);
+            ASSERT_TRUE(s.ok());
+        } while (status.s == TransferStatusEnum::WAITING);
+
+        ASSERT_EQ(status.s, TransferStatusEnum::COMPLETED);
+        s = client_engine->freeBatchID(batch_id);
+        ASSERT_TRUE(s.ok());
+    }
+
+    {
+        auto batch_id = client_engine->allocateBatchID(1);
+        TransferRequest entry;
+        entry.opcode = TransferRequest::READ;
+        entry.length = kDataLength;
+        entry.source = static_cast<char*>(client_buffer) + kDataLength;
+        entry.target_id = segment_id;
+        entry.target_offset = reinterpret_cast<uint64_t>(server_buffer);
+        Status s = client_engine->submitTransfer(batch_id, {entry});
+        ASSERT_TRUE(s.ok());
+
+        TransferStatus status;
+        do {
+            s = client_engine->getTransferStatus(batch_id, 0, status);
+            ASSERT_TRUE(s.ok());
+        } while (status.s == TransferStatusEnum::WAITING);
+
+        ASSERT_EQ(status.s, TransferStatusEnum::COMPLETED);
+        s = client_engine->freeBatchID(batch_id);
+        ASSERT_TRUE(s.ok());
+    }
+
+    std::vector<char> host_check(kDataLength);
+    checkSunriseError(
+        tangMemcpy(host_check.data(),
+                   static_cast<char*>(client_buffer) + kDataLength, kDataLength,
+                   tangMemcpyDeviceToHost),
+        "Memcpy from client_buffer");
+    for (size_t i = 0; i < kDataLength; ++i) {
+        ASSERT_EQ(host_check[i], 'A');
+    }
+
+    client_engine->unregisterLocalMemory(client_buffer);
+    freeSunriseBuffer(client_buffer);
+    server_engine->unregisterLocalMemory(server_buffer);
+    freeSunriseBuffer(server_buffer);
+}
+
+int main(int argc, char** argv) {
+    gflags::ParseCommandLineFlags(&argc, &argv, false);
+    ::testing::InitGoogleTest(&argc, argv);
+    return RUN_ALL_TESTS();
+}
diff --git a/mooncake-transfer-engine/tests/sunrise_link_transport_unit_test.cpp b/mooncake-transfer-engine/tests/sunrise_link_transport_unit_test.cpp
new file mode 100644
index 00000000..6b70184d
--- /dev/null
+++ b/mooncake-transfer-engine/tests/sunrise_link_transport_unit_test.cpp
@@ -0,0 +1,21 @@
+#include "transport/sunrise_link_transport/sunrise_link_transport.h"
+
+#include <gtest/gtest.h>
+
+#include <limits>
+
+using namespace mooncake;
+
+TEST(SunriseLinkTransportUnitTest, RejectsWrappedRemoteRanges) {
+    EXPECT_FALSE(SunriseLinkTransport::requestRangeFitsInBufferForTest(
+        std::numeric_limits<uint64_t>::max() - 50, 80,
+        std::numeric_limits<uint64_t>::max() - 100, 64));
+}
+
+TEST(SunriseLinkTransportUnitTest, ParsesDeviceIdStrictly) {
+    EXPECT_EQ(SunriseLinkTransport::parseDeviceIdForTest("cuda:7"), 7);
+    EXPECT_EQ(SunriseLinkTransport::parseDeviceIdForTest("cuda:"), -1);
+    EXPECT_EQ(SunriseLinkTransport::parseDeviceIdForTest("cuda:not-a-number"),
+              -1);
+    EXPECT_EQ(SunriseLinkTransport::parseDeviceIdForTest("cuda:7x"), -1);
+}
diff --git a/mooncake-transfer-engine/tests/tcp_address_validation_test.cpp b/mooncake-transfer-engine/tests/tcp_address_validation_test.cpp
new file mode 100644
index 00000000..d10966e3
--- /dev/null
+++ b/mooncake-transfer-engine/tests/tcp_address_validation_test.cpp
@@ -0,0 +1,126 @@
+// Copyright 2024 KVCache.AI
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include <gtest/gtest.h>
+
+#include <cstdint>
+#include <vector>
+
+namespace {
+
+struct BufferRange {
+    uint64_t addr;
+    uint64_t length;
+};
+
+// Mirrors TcpTransport::validateAddress logic exactly.
+bool validateAddress(uint64_t addr, uint64_t size,
+                     const std::vector<BufferRange>& buffers) {
+    if (size == 0) return false;
+    if (addr + size < addr) return false;
+
+    for (const auto& buffer : buffers) {
+        if (buffer.addr + buffer.length < buffer.addr) continue;
+        if (buffer.addr <= addr && addr + size <= buffer.addr + buffer.length)
+            return true;
+    }
+    return false;
+}
+
+class TcpAddressValidationTest : public ::testing::Test {
+   protected:
+    std::vector<BufferRange> buffers_;
+
+    void SetUp() override {
+        buffers_ = {
+            {0x1000, 0x2000},     // [0x1000, 0x3000)
+            {0x10000, 0x100000},  // [0x10000, 0x110000)
+        };
+    }
+};
+
+TEST_F(TcpAddressValidationTest, ExactMatch) {
+    EXPECT_TRUE(validateAddress(0x1000, 0x2000, buffers_));
+    EXPECT_TRUE(validateAddress(0x10000, 0x100000, buffers_));
+}
+
+TEST_F(TcpAddressValidationTest, WithinBounds) {
+    EXPECT_TRUE(validateAddress(0x1000, 0x100, buffers_));
+    EXPECT_TRUE(validateAddress(0x1500, 0x500, buffers_));
+    EXPECT_TRUE(validateAddress(0x2FFF, 1, buffers_));
+    EXPECT_TRUE(validateAddress(0x50000, 0x1000, buffers_));
+}
+
+TEST_F(TcpAddressValidationTest, OutOfBounds) {
+    EXPECT_FALSE(validateAddress(0x500, 0x100, buffers_));
+    EXPECT_FALSE(validateAddress(0x3000, 0x100, buffers_));
+    EXPECT_FALSE(validateAddress(0x5000, 0x1000, buffers_));
+    EXPECT_FALSE(validateAddress(0x200000, 0x100, buffers_));
+}
+
+TEST_F(TcpAddressValidationTest, PartialOverlap) {
+    EXPECT_FALSE(validateAddress(0x2F00, 0x200, buffers_));
+    EXPECT_FALSE(validateAddress(0x0F00, 0x200, buffers_));
+    EXPECT_FALSE(validateAddress(0x10F000, 0x2000, buffers_));
+}
+
+TEST_F(TcpAddressValidationTest, ZeroSize) {
+    EXPECT_FALSE(validateAddress(0x1000, 0, buffers_));
+    EXPECT_FALSE(validateAddress(0x0, 0, buffers_));
+}
+
+TEST_F(TcpAddressValidationTest, IntegerOverflow) {
+    EXPECT_FALSE(validateAddress(UINT64_MAX, 1, buffers_));
+    EXPECT_FALSE(validateAddress(UINT64_MAX - 10, 100, buffers_));
+    EXPECT_FALSE(validateAddress(1, UINT64_MAX, buffers_));
+}
+
+TEST_F(TcpAddressValidationTest, EmptyBufferList) {
+    std::vector<BufferRange> empty;
+    EXPECT_FALSE(validateAddress(0x1000, 0x100, empty));
+}
+
+TEST_F(TcpAddressValidationTest, SingleByteAccess) {
+    EXPECT_TRUE(validateAddress(0x1000, 1, buffers_));
+    EXPECT_TRUE(validateAddress(0x2FFF, 1, buffers_));
+    EXPECT_FALSE(validateAddress(0x3000, 1, buffers_));
+    EXPECT_FALSE(validateAddress(0x0FFF, 1, buffers_));
+}
+
+TEST_F(TcpAddressValidationTest, AdjacentBuffers) {
+    std::vector<BufferRange> adjacent = {
+        {0x1000, 0x1000},  // [0x1000, 0x2000)
+        {0x2000, 0x1000},  // [0x2000, 0x3000)
+    };
+    EXPECT_TRUE(validateAddress(0x1000, 0x1000, adjacent));
+    EXPECT_TRUE(validateAddress(0x2000, 0x1000, adjacent));
+    // Spanning two buffers should fail
+    EXPECT_FALSE(validateAddress(0x1800, 0x1000, adjacent));
+}
+
+TEST_F(TcpAddressValidationTest, LargeBuffer) {
+    std::vector<BufferRange> large = {
+        {0, UINT64_MAX},
+    };
+    EXPECT_TRUE(validateAddress(0, 1, large));
+    EXPECT_TRUE(validateAddress(0, UINT64_MAX, large));
+    EXPECT_TRUE(validateAddress(UINT64_MAX - 1, 1, large));
+}
+
+}  // namespace
+
+int main(int argc, char** argv) {
+    ::testing::InitGoogleTest(&argc, argv);
+    return RUN_ALL_TESTS();
+}
diff --git a/mooncake-transfer-engine/tests/tcp_transport_test.cpp b/mooncake-transfer-engine/tests/tcp_transport_test.cpp
index 0e34a4f7..6fde3ac7 100644
--- a/mooncake-transfer-engine/tests/tcp_transport_test.cpp
+++ b/mooncake-transfer-engine/tests/tcp_transport_test.cpp
@@ -28,7 +28,7 @@
 #endif
 
 #if defined(USE_CUDA) || defined(USE_MUSA) || defined(USE_HIP) || \
-    defined(USE_MACA)
+    defined(USE_MACA) || defined(USE_HYGON) || defined(USE_COREX)
 #include <cassert>
 #include "common/base/status.h"
 
@@ -45,7 +45,7 @@ static void checkCudaError(cudaError_t result, const char *message) {
 #include "transport/transport.h"
 
 #if defined(USE_CUDA) || defined(USE_MUSA) || defined(USE_HIP) || \
-    defined(USE_MACA)
+    defined(USE_MACA) || defined(USE_HYGON) || defined(USE_COREX)
 DEFINE_int32(gpu_id, 0, "GPU ID to use");
 #endif
 
diff --git a/mooncake-transfer-engine/tests/topology_test.cpp b/mooncake-transfer-engine/tests/topology_test.cpp
index 5b39d5d5..f8e38b34 100644
--- a/mooncake-transfer-engine/tests/topology_test.cpp
+++ b/mooncake-transfer-engine/tests/topology_test.cpp
@@ -2,6 +2,7 @@
 
 #include <glog/logging.h>
 #include <gtest/gtest.h>
+#include <cstdlib>
 
 #include "transfer_metadata.h"
 #include "memory_location.h"
@@ -16,6 +17,13 @@ TEST(ToplogyTest, GetTopologyMatrix) {
     ASSERT_EQ(topology.toString(), json_str);
 }
 
+TEST(ToplogyTest, DiscoverWithMcTeFilters) {
+    ASSERT_EQ(::setenv("MC_TE_FILTERS", "erdma_0", 1), 0);
+    mooncake::Topology topology;
+    EXPECT_EQ(topology.discover(), 0);
+    ::unsetenv("MC_TE_FILTERS");
+}
+
 TEST(ToplogyTest, TestEmpty) {
     mooncake::Topology topology;
     std::string json_str =
diff --git a/mooncake-transfer-engine/ubshmem-allocator/CMakeLists.txt b/mooncake-transfer-engine/ubshmem-allocator/CMakeLists.txt
index 1ee8af11..f6389f74 100644
--- a/mooncake-transfer-engine/ubshmem-allocator/CMakeLists.txt
+++ b/mooncake-transfer-engine/ubshmem-allocator/CMakeLists.txt
@@ -1,18 +1,16 @@
-# Build ubshmem allocator and output to build directory
+include(${CMAKE_CURRENT_SOURCE_DIR}/../fabric_allocator.cmake)
 
-add_custom_target(build_ubshmem_allocator DEPENDS transfer_engine)
-
-get_target_property(INCLUDE_DIRS build_ubshmem_allocator INCLUDE_DIRECTORIES)
-string(REPLACE ";" " " INCLUDE_DIRS_STR "${INCLUDE_DIRS}")
-
-if (USE_UBSHMEM)
-    add_custom_command(
-        TARGET build_ubshmem_allocator
-        COMMAND ${CMAKE_COMMAND} -E make_directory ${CMAKE_CURRENT_BINARY_DIR}
-        COMMAND bash build.sh ${CMAKE_CURRENT_BINARY_DIR} "${INCLUDE_DIRS_STR}"
-        WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}
-        COMMENT "Building ubshmem allocator to ${CMAKE_CURRENT_BINARY_DIR}"
-    )
+set(_enable_ubshmem_allocator_build FALSE)
+if(USE_UBSHMEM)
+  set(_enable_ubshmem_allocator_build TRUE)
 endif()
 
-set_property(TARGET build_ubshmem_allocator PROPERTY EXCLUDE_FROM_ALL FALSE)
\ No newline at end of file
+add_fabric_allocator_build_target(
+  TARGET_NAME
+  build_ubshmem_allocator
+  BUILD_SCRIPT
+  ${CMAKE_CURRENT_SOURCE_DIR}/build.sh
+  COMMENT
+  "Building ubshmem allocator to ${CMAKE_CURRENT_BINARY_DIR}"
+  ENABLE_BUILD
+  ${_enable_ubshmem_allocator_build})
diff --git a/mooncake-transfer-engine/ubshmem-allocator/build.sh b/mooncake-transfer-engine/ubshmem-allocator/build.sh
index 99ac821d..4b13d114 100644
--- a/mooncake-transfer-engine/ubshmem-allocator/build.sh
+++ b/mooncake-transfer-engine/ubshmem-allocator/build.sh
@@ -2,30 +2,16 @@
 
 set -e
 
+source "$(dirname "$(readlink -f "$0")")/../scripts/allocator_build_common.sh"
+
 # Get output directory from command line argument, default to current directory
 OUTPUT_DIR=${1:-.}
 
-# Get include directories from second argument (if provided)
-INCLUDE_LIST=""
-if [ $# -ge 2 ]; then
-    INCLUDE_LIST=${2}
-fi
-
-# Add include directory for cuda (relative to build.sh location)
-SCRIPT_DIR=$(cd "$(dirname "${BASH_SOURCE[0]}")" &>/dev/null && pwd)
-INCLUDE_LIST="${INCLUDE_LIST:+${INCLUDE_LIST} }${SCRIPT_DIR}/../include"
-
-# Process include directories into flags
-INCLUDE_FLAGS=""
-if [ -n "$INCLUDE_LIST" ]; then
-    INCLUDE_FLAGS=$(echo "$INCLUDE_LIST" | tr ' ' '\n' | sed 's/^/-I/' | paste -sd' ' -)
-fi
+prepare_allocator_build_env "$OUTPUT_DIR" "${2:-}"
 
 echo "Building ubshmem fabric allocator to: $OUTPUT_DIR"
-# Create output directory if it doesn't exist
-mkdir -p "$OUTPUT_DIR"
 
-CPP_FILE=$(dirname $(readlink -f $0))/ubshmem_fabric_allocator.cpp
+CPP_FILE="${SCRIPT_DIR}/ubshmem_fabric_allocator.cpp"
 
 # Detect CPU architecture
 CURRENT_CPU=$(uname -m)
diff --git a/mooncake-transfer-engine/ubshmem-allocator/ubshmem_fabric_allocator.cpp b/mooncake-transfer-engine/ubshmem-allocator/ubshmem_fabric_allocator.cpp
index 8540dc47..b76e67b3 100644
--- a/mooncake-transfer-engine/ubshmem-allocator/ubshmem_fabric_allocator.cpp
+++ b/mooncake-transfer-engine/ubshmem-allocator/ubshmem_fabric_allocator.cpp
@@ -3,9 +3,9 @@
 
 #include <iostream>
 
-extern "C" {
+namespace {
 
-bool mc_probe_ub_fabric_support(int device_id) {
+bool ProbeAllocatorBackend(int device_id) {
     aclError res = aclrtSetDevice(device_id);
     if (res != ACL_ERROR_NONE) {
         std::cerr << "Set device failed: " << device_id << ", result" << res;
@@ -19,25 +19,22 @@ bool mc_probe_ub_fabric_support(int device_id) {
     prop.memAttr = ACL_HBM_MEM_HUGE;
 
     aclrtDrvMemHandle handle;
-    // try to allocate 2M fabric mem
     size_t size = 2 * 1024 * 1024;
 
     res = aclrtMallocPhysical(&handle, size, &prop, 0);
     if (res == ACL_ERROR_NONE) {
-        aclrtFreePhysical(handle);  // success → clean up
+        aclrtFreePhysical(handle);
         return true;
-    } else {
-        return false;
     }
+    return false;
 }
 
-void *mc_ub_fabric_malloc(ssize_t size, int device) {
+void *AllocateFabricMemory(ssize_t size, int device) {
     size_t granularity = 0;
     aclrtPhysicalMemProp prop = {};
     aclrtDrvMemHandle handle;
     void *ptr = nullptr;
 
-    // Align size to 2M
     const size_t alignment = 2 * 1024 * 1024;
     size = (size + alignment - 1) & ~(alignment - 1);
 
@@ -54,7 +51,6 @@ void *mc_ub_fabric_malloc(ssize_t size, int device) {
         return nullptr;
     }
     uint64_t page_type = 1;
-    // now granularity is reserved to 0
     result =
         aclrtReserveMemAddress(&ptr, size, granularity, nullptr, page_type);
     if (result != ACL_ERROR_NONE) {
@@ -73,7 +69,8 @@ void *mc_ub_fabric_malloc(ssize_t size, int device) {
     return ptr;
 }
 
-void mc_ub_fabric_free(void *ptr, int device) {
+void FreeFabricMemory(void *ptr, int device) {
+    (void)device;
     aclrtDrvMemHandle handle;
     if (!ptr) {
         return;
@@ -83,10 +80,35 @@ void mc_ub_fabric_free(void *ptr, int device) {
         std::cerr << "aclrtMemRetainAllocationHandle failed: " << result
                   << "\n";
         return;
-    } else {
-        (void)aclrtUnmapMem(ptr);
-        (void)aclrtReleaseMemAddress(ptr);
     }
+    (void)aclrtUnmapMem(ptr);
+    (void)aclrtReleaseMemAddress(ptr);
     (void)aclrtFreePhysical(handle);
 }
+
+}  // namespace
+
+extern "C" {
+
+bool mc_probe_ub_fabric_support(int device_id) {
+    return ProbeAllocatorBackend(device_id);
+}
+
+int mc_allocator_probe(int device_id) {
+    return ProbeAllocatorBackend(device_id) ? 1 : 0;
+}
+
+void *mc_allocator_malloc(ssize_t size, int device) {
+    return AllocateFabricMemory(size, device);
+}
+
+void *mc_ub_fabric_malloc(ssize_t size, int device) {
+    return mc_allocator_malloc(size, device);
+}
+
+void mc_allocator_free(void *ptr, int device) { FreeFabricMemory(ptr, device); }
+
+void mc_ub_fabric_free(void *ptr, int device) {
+    mc_allocator_free(ptr, device);
+}
 }
diff --git a/mooncake-wheel/mooncake/__init__.py b/mooncake-wheel/mooncake/__init__.py
index ebc08cad..f1bb787f 100644
--- a/mooncake-wheel/mooncake/__init__.py
+++ b/mooncake-wheel/mooncake/__init__.py
@@ -1,2 +1,6 @@
-# Import for backward compatibility
+"""Mooncake public Python package."""
+
+from mooncake.buffer_pool import BufferPool, RegisteredBufferPool
+
+__all__ = ["BufferPool", "RegisteredBufferPool"]
 
diff --git a/mooncake-wheel/mooncake/buffer_pool.py b/mooncake-wheel/mooncake/buffer_pool.py
new file mode 100644
index 00000000..bb4d39ed
--- /dev/null
+++ b/mooncake-wheel/mooncake/buffer_pool.py
@@ -0,0 +1,12 @@
+from __future__ import annotations
+
+"""Public Python entrypoint for Mooncake local-buffer pools."""
+
+try:
+    from mooncake.store import BufferPool
+except (ImportError, AttributeError):  # pragma: no cover - depends on built extension
+    BufferPool = None  # type: ignore[assignment]
+
+RegisteredBufferPool = BufferPool
+
+__all__ = ["BufferPool", "RegisteredBufferPool"]
diff --git a/mooncake-wheel/mooncake/mooncake_config.py b/mooncake-wheel/mooncake/mooncake_config.py
index b4920d0f..e3112968 100644
--- a/mooncake-wheel/mooncake/mooncake_config.py
+++ b/mooncake-wheel/mooncake/mooncake_config.py
@@ -53,21 +53,63 @@ from typing import Optional
 DEFAULT_GLOBAL_SEGMENT_SIZE = 3355443200  # 3.125 GiB
 DEFAULT_LOCAL_BUFFER_SIZE = 1073741824  # 1.0 GiB
 
+_SIZE_SUFFIXES = [
+    ("kb", 1024),
+    ("mb", 1024 ** 2),
+    ("gb", 1024 ** 3),
+    ("tb", 1024 ** 4),
+    ("k", 1024),
+    ("m", 1024 ** 2),
+    ("g", 1024 ** 3),
+    ("t", 1024 ** 4),
+    ("b", 1),
+]
+
+
 def _parse_segment_size(value) -> int:
     if isinstance(value, int):
         return value
     if isinstance(value, str):
         s = value.strip().lower()
-        if s.endswith("gb"):
-            num = s[:-2].strip()
-            if not num:
-                raise ValueError(
-                    "Invalid segment size: missing number before 'gb'"
-                )
-            return int(num) * 1024 * 1024 * 1024
-        return int(s)
+        if not s:
+            raise ValueError("Invalid segment size: empty string")
+        for suffix, multiplier in _SIZE_SUFFIXES:
+            if s.endswith(suffix):
+                num = s[: -len(suffix)].strip()
+                if not num:
+                    raise ValueError(
+                        f"Invalid segment size: missing number before '{suffix}'"
+                    )
+                return int(float(num) * multiplier)
+        return int(float(s))
     return int(value)
 
+
+def _parse_bool(value) -> bool:
+    """Interpret a config boolean that may arrive as a real bool, number, or string.
+
+    Config files and environment variables sometimes carry booleans as strings
+    ("true", "false", "yes", "no", "on", "off", "0", "1"). Using bool() directly
+    would treat any non-empty string (including "false") as True, so parse the
+    common textual forms and reject anything unrecognized instead of silently
+    defaulting to False, mirroring how _parse_segment_size rejects bad input.
+    """
+    if isinstance(value, bool):
+        return value
+    if isinstance(value, (int, float)):
+        return bool(value)
+    if value is None:
+        return False
+    s = str(value).strip().lower()
+    if not s:
+        return False
+    if s in ("true", "1", "yes", "on"):
+        return True
+    if s in ("false", "0", "no", "off"):
+        return False
+    raise ValueError(f"Invalid boolean value: {value!r}")
+
+
 @dataclass
 class MooncakeConfig:
     """The configuration class for Mooncake.
@@ -85,6 +127,8 @@ class MooncakeConfig:
             (e.g., "mlx5_0", "erdma_0", or "auto-discovery").
             Required when protocol is "rdma", optional for other protocols.
         master_server_address (str): The address of the master server.
+        enable_ssd_offload (bool): Enable SSD offload. Default is False.
+        ssd_offload_path (str): The path to the SSD directory for offloading.
 
     Example of configuration file:
         {
@@ -94,7 +138,9 @@ class MooncakeConfig:
             "local_buffer_size": 1073741824,
             "protocol": "tcp",
             "device_name": "",
-            "master_server_address": "localhost:8081"
+            "master_server_address": "localhost:8081",
+            "enable_ssd_offload": true,
+            "ssd_offload_path": "/nvme/mooncake_offload"
         }
         
         For RDMA:
@@ -105,7 +151,9 @@ class MooncakeConfig:
             "local_buffer_size": 1073741824,
             "protocol": "rdma",
             "device_name": "mlx5_0",
-            "master_server_address": "master:8081"
+            "master_server_address": "master:8081",
+            "enable_ssd_offload": true,
+            "ssd_offload_path": "/nvme/mooncake_offload"
         }
     """
     local_hostname: str
@@ -115,6 +163,8 @@ class MooncakeConfig:
     protocol: str
     device_name: Optional[str]
     master_server_address: str
+    enable_ssd_offload: bool = False
+    ssd_offload_path: str = ""
 
     @staticmethod
     def from_file(file_path: str) -> 'MooncakeConfig':
@@ -141,6 +191,8 @@ class MooncakeConfig:
             protocol=config.get("protocol", "tcp"),
             device_name=config.get("device_name", ""),
             master_server_address=config.get("master_server_address"),
+            enable_ssd_offload=_parse_bool(config.get("enable_ssd_offload", False)),
+            ssd_offload_path=str(config.get("ssd_offload_path", "")),
         )
 
     @staticmethod
@@ -167,5 +219,7 @@ class MooncakeConfig:
                 protocol=os.getenv("MOONCAKE_PROTOCOL", "tcp"),
                 device_name=os.getenv("MOONCAKE_DEVICE", ""),
                 master_server_address=os.getenv("MOONCAKE_MASTER"),
+                enable_ssd_offload=_parse_bool(os.getenv("MOONCAKE_OFFLOAD_ENABLED", "false")),
+                ssd_offload_path=os.getenv("MOONCAKE_OFFLOAD_FILE_STORAGE_PATH", ""),
             )
         return MooncakeConfig.from_file(config_file_path)
\ No newline at end of file
diff --git a/mooncake-wheel/mooncake/mooncake_connector_v1.py b/mooncake-wheel/mooncake/mooncake_connector_v1.py
index ee30914c..18873b38 100644
--- a/mooncake-wheel/mooncake/mooncake_connector_v1.py
+++ b/mooncake-wheel/mooncake/mooncake_connector_v1.py
@@ -88,6 +88,7 @@ class RecvReqMeta:
     local_block_ids: list[int]
     remote_host: str
     remote_port: int
+    remote_request_id: Optional[ReqId] = None
 
 
 @dataclass
@@ -112,7 +113,8 @@ class MooncakeConnectorMetadata(KVConnectorMetadata):
             self.reqs_to_recv[request_id] = RecvReqMeta(
                 local_block_ids=local_block_ids,
                 remote_host=kv_transfer_params["remote_host"],
-                remote_port=kv_transfer_params["remote_port"])
+                remote_port=kv_transfer_params["remote_port"],
+                remote_request_id=kv_transfer_params.get("remote_request_id"))
         else:
             self.reqs_to_send[request_id] = local_block_ids
 
@@ -382,10 +384,12 @@ class MooncakeConnectorScheduler:
         if delay_free_blocks:
             self._reqs_need_send[request.request_id] = block_ids
 
-        return delay_free_blocks, dict(do_remote_prefill=True,
-                                       do_remote_decode=False,
-                                       remote_host=self.side_channel_host,
-                                       remote_port=self.side_channel_port)
+        return delay_free_blocks, dict(
+            do_remote_prefill=True,
+            do_remote_decode=False,
+            remote_host=self.side_channel_host,
+            remote_port=self.side_channel_port,
+            remote_request_id=request.request_id)
 
 
 class MooncakeConnectorWorker:
@@ -812,12 +816,14 @@ class MooncakeConnectorWorker:
 
         return finished_sending_reqs or None, finished_recving_reqs or None
 
-    async def receive_kv(self, path: str, req_blocks: list[tuple[str, list[int]]]):
-        req_ids, block_ids = map(list, zip(*req_blocks))
+    async def receive_kv(
+        self, path: str, req_blocks: list[tuple[str, str, list[int]]]
+    ):
+        local_req_ids, remote_req_ids, block_ids = map(list, zip(*req_blocks))
         metadata = MooncakeAgentMetadata(
             remote_hostname=self.hostname,
             remote_port=self.rpc_port,
-            request_ids=req_ids,
+            request_ids=remote_req_ids,
             kv_caches_base_addr=self.kv_caches_base_addr,
             block_ids=block_ids,
         )
@@ -826,7 +832,9 @@ class MooncakeConnectorWorker:
         logger.debug(
             "Size of encoded MooncakeAgentMetadata: %d bytes", len(encoded_data)
         )
-        logger.debug("Sending kv transfer request for %s on path: %s", req_ids, path)
+        logger.debug(
+            "Sending kv transfer request for %s on path: %s "
+            "(local requests: %s)", remote_req_ids, path, local_req_ids)
 
         # Send query for the request.
         sock: zmq.asyncio.Socket = make_zmq_socket(
@@ -839,20 +847,24 @@ class MooncakeConnectorWorker:
             if ret_msg != TRANS_DONE:
                 logger.error(
                     "Error happens during transferring kvcache for %s, see logs in prefiller.",  # noqa: E501
-                    req_ids,
+                    remote_req_ids,
                 )
                 return
         except zmq.ContextTerminated:
             logger.debug("ZMQ context terminated, exiting Mooncake receiver thread.")
         except Exception as e:
-            logger.error("MooncakeAgentMetadata transfer failed for %s: %s", req_ids, e)
+            logger.error(
+                "MooncakeAgentMetadata transfer failed for %s: %s",
+                remote_req_ids, e)
             return
         finally:
             sock.close()
 
-        self.finished_recving_reqs.update(req_ids)
+        self.finished_recving_reqs.update(local_req_ids)
 
-        logger.debug("pulling kv_caches for %s finished", req_ids)
+        logger.debug(
+            "pulling kv_caches for %s finished (local requests: %s)",
+            remote_req_ids, local_req_ids)
 
     def group_kv_pull(self, metadata: MooncakeConnectorMetadata):
         kv_pulls = defaultdict(list)
@@ -862,7 +874,12 @@ class MooncakeConnectorWorker:
                 "Num local_block_ids: %s.", req_id, len(meta.local_block_ids))
             path = make_zmq_path("tcp", meta.remote_host,
                                  meta.remote_port + self.tp_rank)
-            kv_pulls[path].append((req_id, meta.local_block_ids))
+            remote_req_id = meta.remote_request_id or req_id
+            if remote_req_id != req_id:
+                logger.debug(
+                    "request %s will pull remote kv for producer request %s",
+                    req_id, remote_req_id)
+            kv_pulls[path].append((req_id, remote_req_id, meta.local_block_ids))
 
         return kv_pulls
 
diff --git a/mooncake-wheel/mooncake/mooncake_ep_buffer.py b/mooncake-wheel/mooncake/mooncake_ep_buffer.py
index c37502b5..b3deb7e7 100644
--- a/mooncake-wheel/mooncake/mooncake_ep_buffer.py
+++ b/mooncake-wheel/mooncake/mooncake_ep_buffer.py
@@ -1,3 +1,4 @@
+import os
 import torch
 import torch.distributed as dist
 from typing import Any, Callable, List, Tuple, Optional, Union
@@ -63,20 +64,17 @@ class EventOverlap:
 
 class Buffer:
     def __init__(self, group: dist.ProcessGroup, num_ep_buffer_bytes: int = 0):
-        from mooncake import ep, pg
+        from mooncake import ep
 
         # Initialize the CPP runtime
         self.rank = group.rank()
         self.group_size = group.size()
         self.group = group
         self.num_ep_buffer_bytes = num_ep_buffer_bytes
-        # Get the index of the closest NIC
-        self.backend = self.group._get_backend(torch.device("cuda"))
-        preferred_hca = pg.get_preferred_hca(
-            self.backend, f"cuda:{torch.cuda.current_device()}"
-        )
+        self.backend = self.group
+        # NIC auto-detection happens inside ep.Buffer via Topology::discover().
         self.runtime = ep.Buffer(
-            self.rank, self.group_size, num_ep_buffer_bytes, preferred_hca
+            self.rank, self.group_size, num_ep_buffer_bytes
         )
         # Fallback flag and buffers.
         # Note: `sync_nvlink_ipc_handles()` can mutate C++ `ibgda_disabled_` (True->False when
@@ -125,57 +123,46 @@ class Buffer:
                 for _ in range(self.group_size)
             ]
             dist.all_to_all(remote_qpns, local_qpns, self.group)
-            remote_qpns = torch.cat(remote_qpns).tolist()
+            peer_qpns = [remote_qpns[r].tolist() for r in range(self.group_size)]
 
-            if self.runtime.is_roce():
-                (subnet_prefix, interface_id) = self.runtime.get_gid()
-
-                subnet_prefix = torch.tensor(
-                    [subnet_prefix], dtype=torch.int64, device="cuda"
-                )
-                subnet_prefixes = [
-                    torch.empty(1, dtype=torch.int64, device="cuda")
-                    for _ in range(self.group_size)
-                ]
-                dist.all_gather(subnet_prefixes, subnet_prefix, self.group)
-                subnet_prefixes = torch.cat(subnet_prefixes).tolist()
-
-                interface_id = torch.tensor(
-                    [interface_id], dtype=torch.int64, device="cuda"
-                )
-                interface_ids = [
-                    torch.empty(1, dtype=torch.int64, device="cuda")
-                    for _ in range(self.group_size)
-                ]
-                dist.all_gather(interface_ids, interface_id, self.group)
-                interface_ids = torch.cat(interface_ids).tolist()
-
-                from mooncake.ep import get_active_ranks
-                active_ranks_mask = get_active_ranks(self.backend).tolist()
-                self.runtime.sync_roce(
-                    raddrs, rkeys, remote_qpns, subnet_prefixes, interface_ids,
-                    active_ranks_mask
-                )
-            else:
-                local_lids = self.runtime.get_local_lids()
-                local_lids = list(
-                    torch.unbind(
-                        torch.tensor(local_lids, dtype=torch.int32, device="cuda").view(
-                            -1, all_to_all_size
-                        )
+            local_lids = self.runtime.get_local_lids()
+            local_lids = list(
+                torch.unbind(
+                    torch.tensor(local_lids, dtype=torch.int32, device="cuda").view(
+                        -1, all_to_all_size
                     )
                 )
-                remote_lids = [
-                    torch.empty(all_to_all_size, dtype=torch.int32, device="cuda")
-                    for _ in range(self.group_size)
-                ]
-                dist.all_to_all(remote_lids, local_lids, self.group)
-                remote_lids = torch.cat(remote_lids).tolist()
+            )
+            remote_lids = [
+                torch.empty(all_to_all_size, dtype=torch.int32, device="cuda")
+                for _ in range(self.group_size)
+            ]
+            dist.all_to_all(remote_lids, local_lids, self.group)
+            peer_lids = [remote_lids[r].tolist() for r in range(self.group_size)]
 
-                from mooncake.ep import get_active_ranks
-                active_ranks_mask = get_active_ranks(self.backend).tolist()
-                self.runtime.sync_ib(raddrs, rkeys, remote_qpns, remote_lids,
-                                     active_ranks_mask)
+            (subnet_prefix, interface_id) = self.runtime.get_gid()
+            subnet_prefix_t = torch.tensor([subnet_prefix], dtype=torch.int64, device="cuda")
+            subnet_prefixes_list = [
+                torch.empty(1, dtype=torch.int64, device="cuda")
+                for _ in range(self.group_size)
+            ]
+            dist.all_gather(subnet_prefixes_list, subnet_prefix_t, self.group)
+            subnet_prefixes = torch.cat(subnet_prefixes_list).tolist()
+
+            interface_id_t = torch.tensor([interface_id], dtype=torch.int64, device="cuda")
+            interface_ids_list = [
+                torch.empty(1, dtype=torch.int64, device="cuda")
+                for _ in range(self.group_size)
+            ]
+            dist.all_gather(interface_ids_list, interface_id_t, self.group)
+            interface_ids = torch.cat(interface_ids_list).tolist()
+
+            from mooncake.ep import get_active_ranks
+            active_ranks_mask = get_active_ranks(self.backend).tolist()
+            self.runtime.sync_ibgda_peers(
+                raddrs, rkeys, peer_qpns, peer_lids,
+                subnet_prefixes, interface_ids, active_ranks_mask
+            )
 
         try:
             local_handle_ints = self.runtime.get_ipc_handle()
@@ -247,6 +234,20 @@ class Buffer:
         EventOverlap,
         Callable,
     ]:
+        # MUSA does not support cooperative grid sync, so the C++ runtime
+        # splits no-hook calls into SEND -> phase-ack -> RECV instead of using
+        # a single cooperative kernel.  async_finish still returns a stream
+        # event, but it is not the CUDA single-kernel cooperative path.
+        if os.getenv("MOONCAKE_EP_USE_MUSA") and async_finish:
+            import warnings
+
+            warnings.warn(
+                "MUSA async_finish uses split SEND/RECV kernels plus a stream "
+                "event, not CUDA cooperative single-kernel async semantics.",
+                RuntimeWarning,
+                stacklevel=2,
+            )
+
         if self._use_fallback:
             from mooncake.ep import get_active_ranks
 
@@ -329,6 +330,17 @@ class Buffer:
         return_recv_hook: bool = False,
         out: Optional[torch.Tensor] = None,
     ) -> Tuple[torch.Tensor, EventOverlap, Callable]:
+        # Same MUSA split-kernel behavior as dispatch().
+        if os.getenv("MOONCAKE_EP_USE_MUSA") and async_finish:
+            import warnings
+
+            warnings.warn(
+                "MUSA async_finish uses split SEND/RECV kernels plus a stream "
+                "event, not CUDA cooperative single-kernel async semantics.",
+                RuntimeWarning,
+                stacklevel=2,
+            )
+
         (
             src_info,
             layout_range,
diff --git a/mooncake-wheel/mooncake/mooncake_ssd_register.py b/mooncake-wheel/mooncake/mooncake_ssd_register.py
new file mode 100644
index 00000000..3a5e6e18
--- /dev/null
+++ b/mooncake-wheel/mooncake/mooncake_ssd_register.py
@@ -0,0 +1,328 @@
+#!/usr/bin/env python3
+# python3 -m mooncake.mooncake_ssd_register --master_server_address=192.168.65.81:50051 --spdk_target_info="ip:192.168.65.56 path:/home/spdk" --spdk_target_info="ip:192.168.65.57 path:/root/spdk"
+
+import argparse
+import json
+import logging
+import time
+import re
+import shlex
+from typing import List, Dict, Any
+
+import paramiko
+
+from mooncake.store import MooncakeDistributedNoFRegister
+from mooncake.mooncake_config import MooncakeConfig
+
+
+class MooncakeNoFRegister:
+    """
+    Registers SSDs from remote SPDK targets to Mooncake master server.
+    """
+
+    def __init__(self, cli_config: dict = None, spdk_targets: List[str] = None):
+        self.register = None
+        self.config_list: List[Dict[str, Any]] = []
+        self.cli_config = cli_config or {}
+        self.spdk_targets = spdk_targets or []
+        self._setup_logging()
+
+        try:
+            # Only support getting SSD info from remote SPDK targets
+            if self.spdk_targets:
+                # Get SSD info from remote SPDK targets
+                master_server_address = self.cli_config.get('master_server_address')
+                if not master_server_address:
+                    raise ValueError("master_server_address is required when using spdk_target_info")
+                self.config_list = self._get_remote_ssd_info(master_server_address)
+            else:
+                raise ValueError("spdk_target_info is required")
+
+            # Apply CLI overrides to every config (if key exists)
+            for config in self.config_list:
+                for key, value in self.cli_config.items():
+                    if key in config:
+                        # Convert trsvcid/nsid to int if needed
+                        if key in ("trsvcid", "nsid"):
+                            try:
+                                config[key] = int(value)
+                            except ValueError:
+                                logging.warning(f"Invalid integer for {key}: {value}")
+                        else:
+                            config[key] = value
+
+            # Remove duplicate SSDs based on their unique identifiers
+            self._remove_duplicate_ssds()
+
+            logging.info("Loaded %d SSD configuration(s)", len(self.config_list))
+        except Exception as e:
+            logging.error("Configuration load failed: %s", e)
+            raise
+
+
+
+    def _remove_duplicate_ssds(self):
+        """
+        Remove duplicate SSD configurations from the config list
+        """
+        unique_configs = []
+        seen_keys = set()
+
+        for config in self.config_list:
+            # Generate unique key directly without calling _get_ssd_unique_key
+            key = (config['nqn'], config['nsid'], config['traddr'], config['trsvcid'])
+            if key not in seen_keys:
+                seen_keys.add(key)
+                unique_configs.append(config)
+
+        if len(unique_configs) < len(self.config_list):
+            logging.info(f"Removed {len(self.config_list) - len(unique_configs)} duplicate SSD configuration(s)")
+
+        self.config_list = unique_configs
+
+
+
+    def _parse_spdk_target_info(self, target_info: str) -> Dict[str, str]:
+        """
+        Parse spdk target info string like "ip:192.168.65.56 path:/home"
+        """
+        result = {}
+        parts = re.findall(r'(\w+):([^\s]+)', target_info)
+        for key, value in parts:
+            result[key] = value
+        return result
+
+    def _execute_ssh_command(self, ip: str, command: str, path: str) -> str:
+        """
+        Execute command on remote server via SSH
+        """
+        ssh = paramiko.SSHClient()
+        ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
+        try:
+            ssh.connect(
+                ip,
+                port=int(self.cli_config.get('port', 22)),
+                username=self.cli_config.get('username', 'root'),
+                password=self.cli_config.get('password'),
+                key_filename=self.cli_config.get('key_file'),
+                timeout=10
+            )
+
+            # Try multiple possible paths to find the RPC script
+            possible_paths = [
+                path,  # Direct path provided by user
+                f"{path}/spdk"  # Common case: spdk is a subdirectory
+            ]
+
+            for test_path in possible_paths:
+                full_command = f"cd {shlex.quote(test_path)} && test -f scripts/rpc.py"
+                stdin, stdout, stderr = ssh.exec_command(full_command, timeout=5)
+                if stdout.channel.recv_exit_status() == 0:
+                    # Found the script, execute the actual command
+                    full_command = f"cd {shlex.quote(test_path)} && {command}"
+                    stdin, stdout, stderr = ssh.exec_command(full_command, timeout=30)
+                    exit_status = stdout.channel.recv_exit_status()
+                    output = stdout.read().decode('utf-8')
+                    error = stderr.read().decode('utf-8')
+                    if exit_status != 0:
+                        logging.error(f"SSH command error on {ip}: {error}")
+                        raise RuntimeError(f"SSH command failed: {error or output}")
+                    return output
+
+            # If we get here, none of the paths worked
+            raise RuntimeError(f"Could not find scripts/rpc.py in any of the possible paths: {possible_paths}")
+        finally:
+            ssh.close()
+
+    def _get_remote_ssd_info(self, master_server_address: str) -> List[Dict[str, Any]]:
+        """
+        Get SSD info from remote SPDK targets
+        """
+        ssd_configs = []
+
+        for target_info in self.spdk_targets:
+            target = self._parse_spdk_target_info(target_info)
+            ip = target.get('ip')
+            path = target.get('path')
+
+            if not ip or not path:
+                logging.error(f"Invalid target info: {target_info}")
+                continue
+
+            logging.info(f"Getting SSD info from target: {ip} (path: {path})")
+
+            try:
+                # Get subsystems info
+                subsystems_cmd = "./scripts/rpc.py nvmf_get_subsystems"
+                subsystems_output = self._execute_ssh_command(ip, subsystems_cmd, path)
+                subsystems = json.loads(subsystems_output)
+
+                # Process each subsystem
+                for subsystem in subsystems:
+                    if subsystem.get('subtype') != 'NVMe':
+                        continue
+
+                    nqn = subsystem.get('nqn')
+                    listen_addresses = subsystem.get('listen_addresses', [])
+
+                    if not nqn or not listen_addresses:
+                        continue
+
+                    # Get transport info from first listen address
+                    traddr = listen_addresses[0].get('traddr')
+                    trsvcid = listen_addresses[0].get('trsvcid')
+
+                    if not traddr or not trsvcid:
+                        continue
+
+                    # Process each namespace
+                    namespaces = subsystem.get('namespaces', [])
+                    for namespace in namespaces:
+                        nsid = namespace.get('nsid')
+                        bdev_name = namespace.get('bdev_name')
+
+                        if not nsid or not bdev_name:
+                            continue
+
+                        # Get bdev info to calculate size
+                        bdev_cmd = f"./scripts/rpc.py bdev_get_bdevs -b {shlex.quote(bdev_name)}"
+                        bdev_output = self._execute_ssh_command(ip, bdev_cmd, path)
+                        bdevs = json.loads(bdev_output)
+
+                        if not bdevs:
+                            continue
+
+                        bdev = bdevs[0]
+                        block_size = bdev.get('block_size', 512)
+                        num_blocks = bdev.get('num_blocks', 0)
+                        size = block_size * num_blocks
+
+                        # Create SSD config
+                        ssd_config = {
+                            'nqn': nqn,
+                            'nsid': nsid,
+                            'traddr': traddr,
+                            'trsvcid': int(trsvcid),  # Ensure trsvcid is integer
+                            'base': 0,
+                            'size': size,
+                            'master_server_address': master_server_address,
+                            'metadata_server': ''
+                        }
+
+                        ssd_configs.append(ssd_config)
+                        logging.info(f"Found SSD: nqn={nqn}, nsid={nsid}, traddr={traddr}, size={size}")
+
+            except Exception as e:
+                logging.error(f"Failed to get SSD info from {ip}: {e}")
+                continue
+
+        if not ssd_configs:
+            raise RuntimeError("No SSD information found from remote targets")
+
+        return ssd_configs
+
+    def _setup_logging(self):
+        logging.basicConfig(
+            level=logging.INFO,
+            format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
+        )
+
+    def start_ssd_service(self):
+        success_count = 0
+        skipped_count = 0
+        failed_count = 0
+        total = len(self.config_list)
+
+        for i, cfg in enumerate(self.config_list):
+            try:
+                logging.info("Registering SSD %d/%d: nqn=%s, traddr=%s", i + 1, total, cfg.get("nqn"), cfg.get("traddr"))
+
+                # Create register instance and register SSD
+                self.register = MooncakeDistributedNoFRegister()
+                ret = self.register.real_register(
+                    cfg["nqn"],
+                    cfg["nsid"],
+                    cfg["traddr"],
+                    cfg["trsvcid"],
+                    cfg["base"],
+                    cfg["size"],
+                    cfg["master_server_address"]
+                )
+
+                if ret != 0:
+                    raise RuntimeError(f"Registration failed with code {ret}")
+
+                logging.info("Register SSD %d/%d succeeded", i + 1, total)
+                success_count += 1
+
+            except Exception as e:
+                # Check if the error is due to the segment already existing on the server
+                if "SEGMENT_ALREADY_EXISTS" in str(e) or "segment already exists" in str(e):
+                    logging.info("SSD %d/%d (nqn=%s, traddr=%s) already registered on server, skipping",
+                                i + 1, total, cfg.get("nqn"), cfg.get("traddr"))
+                    skipped_count += 1
+                else:
+                    logging.error("Failed to register SSD %d/%d: %s", i + 1, total, e)
+                    failed_count += 1
+
+        # Summary
+        logging.info("SSD registration summary:")
+        logging.info("- Total SSDs: %d", total)
+        logging.info("- Successfully registered: %d", success_count)
+        logging.info("- Already registered (skipped): %d", skipped_count)
+        logging.info("- Failed: %d", failed_count)
+
+        # Return success if all SSDs were either registered or already existed
+        return failed_count == 0
+
+
+def parse_arguments():
+    parser = argparse.ArgumentParser(description='Mooncake SSD Register with REST API')
+    parser.add_argument('--master_server_address', type=str,
+                        help='Master server address (e.g., 192.168.65.81:50051)',
+                        required=True)
+    parser.add_argument('--spdk_target_info', action='append',
+                        help='SPDK target information (e.g., "ip:192.168.65.56 path:/home")',
+                        required=True)
+    parser.add_argument('--username', type=str, default='root',
+                        help='SSH username for target nodes (default: root)')
+    parser.add_argument('--port', type=int, default=22,
+                        help='SSH port for target nodes (default: 22)')
+    parser.add_argument('--password', type=str,
+                        help='SSH password for target nodes')
+    parser.add_argument('--key-file', type=str,
+                        help='SSH private key file path')
+    parser.add_argument('-D', '--define', action='append',
+                        help='Override configuration fields globally (e.g., -Dtrsvcid=4420)',
+                        default=[])
+    return parser.parse_args()
+
+
+def main():
+    args = parse_arguments()
+
+    cli_config = {}
+    for item in args.define:
+        if '=' in item:
+            key, value = item.split('=', 1)
+            cli_config[key] = value
+        else:
+            logging.warning(f"Ignoring invalid CLI config: {item}")
+
+    # Add master_server_address to cli_config if provided
+    if args.master_server_address:
+        cli_config['master_server_address'] = args.master_server_address
+    cli_config['username'] = args.username
+    cli_config['port'] = args.port
+    if args.password:
+        cli_config['password'] = args.password
+    if args.key_file:
+        cli_config['key_file'] = args.key_file
+
+    register = MooncakeNoFRegister(cli_config, args.spdk_target_info)
+    success = register.start_ssd_service()
+    if not success:
+        exit(1)
+
+if __name__ == "__main__":
+    main()
diff --git a/mooncake-wheel/mooncake/mooncake_ssd_unregister.py b/mooncake-wheel/mooncake/mooncake_ssd_unregister.py
new file mode 100644
index 00000000..ec84b21b
--- /dev/null
+++ b/mooncake-wheel/mooncake/mooncake_ssd_unregister.py
@@ -0,0 +1,361 @@
+#!/usr/bin/env python3
+# python3 -m mooncake.mooncake_ssd_unregister --master_server_address=192.168.65.81:50051 --spdk_target_info="ip:192.168.65.56 ns:1 nqn:nqn.2016-06.io.spdk:cnode1"
+
+import argparse
+import logging
+import json
+import re
+import shlex
+from typing import List, Dict, Any
+
+import paramiko
+
+from mooncake.store import MooncakeDistributedNoFRegister
+from mooncake.mooncake_config import MooncakeConfig
+
+
+class MooncakeNoFUnregister:
+    """
+    Unregisters SSDs from Mooncake master server by endpoint information.
+    """
+
+    def __init__(self, cli_config: dict = None, spdk_targets: List[str] = None):
+        self.register = None
+        self.config_list: List[Dict[str, Any]] = []
+        self.cli_config = cli_config or {}
+        self.spdk_targets = spdk_targets or []
+        self._setup_logging()
+
+        try:
+            # Only support getting SSD info from command line
+            if self.spdk_targets:
+                # Get SSD info from command line parameters
+                master_server_address = self.cli_config.get('master_server_address')
+                if not master_server_address:
+                    raise ValueError("master_server_address is required when using spdk_target_info")
+                self.config_list = self._parse_spdk_targets(master_server_address)
+            else:
+                raise ValueError("spdk_target_info is required")
+
+            # Apply CLI overrides to every config (if key exists)
+            for config in self.config_list:
+                for key, value in self.cli_config.items():
+                    if key in config:
+                        # Convert trsvcid/nsid to int if needed
+                        if key in ("trsvcid", "nsid"):
+                            try:
+                                config[key] = int(value)
+                            except ValueError:
+                                logging.warning(f"Invalid integer for {key}: {value}")
+                        else:
+                            config[key] = value
+
+            # Remove duplicate SSDs based on their unique identifiers
+            self._remove_duplicate_ssds()
+
+            logging.info("Loaded %d SSD configuration(s) to unregister", len(self.config_list))
+        except Exception as e:
+            logging.error("Configuration load failed: %s", e)
+            raise
+
+    def _remove_duplicate_ssds(self):
+        """
+        Remove duplicate SSD configurations from the config list
+        """
+        unique_configs = []
+        seen_keys = set()
+
+        for config in self.config_list:
+            # Generate unique key directly
+            key = (config['nqn'], config['nsid'], config['traddr'], config['trsvcid'])
+            if key not in seen_keys:
+                seen_keys.add(key)
+                unique_configs.append(config)
+
+        if len(unique_configs) < len(self.config_list):
+            logging.info(f"Removed {len(self.config_list) - len(unique_configs)} duplicate SSD configuration(s)")
+
+        self.config_list = unique_configs
+
+    def _execute_ssh_command(self, ip: str, command: str, path: str) -> str:
+        """
+        Execute command on remote server via SSH
+        """
+        ssh = paramiko.SSHClient()
+        ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
+        try:
+            ssh.connect(
+                ip,
+                port=int(self.cli_config.get('port', 22)),
+                username=self.cli_config.get('username', 'root'),
+                password=self.cli_config.get('password'),
+                key_filename=self.cli_config.get('key_file'),
+                timeout=10
+            )
+
+            # Try multiple possible paths to find the RPC script
+            possible_paths = [
+                path,  # Direct path provided by user
+                f"{path}/spdk"  # Common case: spdk is a subdirectory
+            ]
+
+            for test_path in possible_paths:
+                full_command = f"cd {shlex.quote(test_path)} && test -f scripts/rpc.py"
+                stdin, stdout, stderr = ssh.exec_command(full_command, timeout=5)
+                if stdout.channel.recv_exit_status() == 0:
+                    # Found the script, execute the actual command
+                    full_command = f"cd {shlex.quote(test_path)} && {command}"
+                    stdin, stdout, stderr = ssh.exec_command(full_command, timeout=30)
+                    exit_status = stdout.channel.recv_exit_status()
+                    output = stdout.read().decode('utf-8')
+                    error = stderr.read().decode('utf-8')
+                    if exit_status != 0:
+                        logging.error(f"SSH command error on {ip}: {error}")
+                        raise RuntimeError(f"SSH command failed: {error or output}")
+                    return output
+
+            # If we get here, none of the paths worked
+            raise RuntimeError(f"Could not find scripts/rpc.py in any of the possible paths: {possible_paths}")
+        finally:
+            ssh.close()
+
+    def _parse_spdk_targets(self, master_server_address: str) -> List[Dict[str, Any]]:
+        """
+        Parse spdk target info strings like "ip:192.168.65.56" or "ip:192.168.65.56 ns:2"
+        or "ip:192.168.65.56 path:/home/spdk" to get actual namespaces from SPDK target
+        """
+        ssd_configs = []
+
+        for target_info in self.spdk_targets:
+            # Parse target info
+            target = {}
+            parts = target_info.split()
+            for part in parts:
+                if ':' in part:
+                    key, value = part.split(':', 1)
+                    target[key.strip()] = value.strip()
+
+            # Validate required fields
+            if 'ip' not in target:
+                raise ValueError("spdk_target_info must contain 'ip' field")
+
+            ip = target['ip']
+            specified_ns = int(target['ns']) if 'ns' in target else None
+            path = target.get('path')
+
+            # We need to know the NQN to build the te_endpoint
+            # For now, we'll use a default NQN pattern (this should be improved)
+            default_nqn = "nqn.2016-06.io.spdk:cnode1"
+            nqn = target.get('nqn', default_nqn)
+
+            # Default transport parameters
+            trsvcid = int(target.get('port', '4420'))
+            trtype = target.get('trtype', 'RDMA')
+
+            # Create SSD config for each namespace (or specified ns only)
+            if specified_ns is not None:
+                # Unregister specific namespace
+                ssd_config = {
+                    'nqn': nqn,
+                    'nsid': specified_ns,
+                    'traddr': ip,
+                    'trsvcid': trsvcid,
+                    'base': 0,
+                    'size': 0,  # Size is not needed for unregister
+                    'master_server_address': master_server_address,
+                    'metadata_server': ''
+                }
+                ssd_configs.append(ssd_config)
+                logging.info(f"Will unregister SSD: nqn={nqn}, nsid={specified_ns}, traddr={ip}")
+            elif path is not None:
+                # Get actual namespaces from SPDK target
+                logging.info(f"Getting namespace info from target: {ip} (path: {path})")
+                try:
+                    # Get subsystems info
+                    subsystems_cmd = "./scripts/rpc.py nvmf_get_subsystems"
+                    subsystems_output = self._execute_ssh_command(ip, subsystems_cmd, path)
+                    subsystems = json.loads(subsystems_output)
+
+                    # Process each subsystem
+                    for subsystem in subsystems:
+                        if subsystem.get('subtype') != 'NVMe':
+                            continue
+
+                        subsystem_nqn = subsystem.get('nqn')
+                        listen_addresses = subsystem.get('listen_addresses', [])
+
+                        if not subsystem_nqn or not listen_addresses:
+                            continue
+
+                        # Get transport info from first listen address
+                        traddr = listen_addresses[0].get('traddr')
+                        target_trsvcid = listen_addresses[0].get('trsvcid')
+
+                        if not traddr or not target_trsvcid:
+                            continue
+
+                        # Use the nqn from the subsystem if not specified
+                        current_nqn = nqn if nqn != default_nqn else subsystem_nqn
+                        current_trsvcid = int(target_trsvcid) if target_trsvcid else trsvcid
+
+                        # Process each namespace
+                        namespaces = subsystem.get('namespaces', [])
+                        for namespace in namespaces:
+                            nsid = namespace.get('nsid')
+
+                            if not nsid:
+                                continue
+
+                            # Create SSD config
+                            ssd_config = {
+                                'nqn': current_nqn,
+                                'nsid': nsid,
+                                'traddr': traddr,
+                                'trsvcid': current_trsvcid,
+                                'base': 0,
+                                'size': 0,  # Size is not needed for unregister
+                                'master_server_address': master_server_address,
+                                'metadata_server': ''
+                            }
+                            ssd_configs.append(ssd_config)
+                            logging.info(f"Will unregister SSD: nqn={current_nqn}, nsid={nsid}, traddr={traddr}")
+
+                except Exception as e:
+                    logging.error(f"Failed to get namespace info from {ip}: {e}")
+                    logging.warning("Falling back to unregistering default namespace (nsid=1)")
+                    # Fall back to unregistering default namespace
+                    ssd_config = {
+                        'nqn': nqn,
+                        'nsid': 1,
+                        'traddr': ip,
+                        'trsvcid': trsvcid,
+                        'base': 0,
+                        'size': 0,
+                        'master_server_address': master_server_address,
+                        'metadata_server': ''
+                    }
+                    ssd_configs.append(ssd_config)
+                    logging.info(f"Will unregister SSD (fallback): nqn={nqn}, nsid=1, traddr={ip}")
+            else:
+                # No path provided, unregister default namespace only
+                logging.warning("No 'path' provided in spdk_target_info, cannot query actual namespaces from SPDK target")
+                logging.warning("Will unregister default namespace (nsid=1) only")
+                ssd_config = {
+                    'nqn': nqn,
+                    'nsid': 1,
+                    'traddr': ip,
+                    'trsvcid': trsvcid,
+                    'base': 0,
+                    'size': 0,
+                    'master_server_address': master_server_address,
+                    'metadata_server': ''
+                }
+                ssd_configs.append(ssd_config)
+                logging.info(f"Will unregister SSD: nqn={nqn}, nsid=1, traddr={ip}")
+
+        if not ssd_configs:
+            raise RuntimeError("No SSD configuration generated from target info")
+
+        return ssd_configs
+
+    def _setup_logging(self):
+        logging.basicConfig(
+            level=logging.INFO,
+            format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
+        )
+
+    def start_ssd_unregister_service(self):
+        success_count = 0
+        failed_count = 0
+        total = len(self.config_list)
+
+        for i, cfg in enumerate(self.config_list):
+            try:
+                logging.info("Unregistering SSD %d/%d: nqn=%s, traddr=%s, nsid=%d", i + 1, total, cfg.get("nqn"), cfg.get("traddr"), cfg.get("nsid"))
+
+                # Create register instance and unregister SSD
+                self.register = MooncakeDistributedNoFRegister()
+                ret = self.register.real_unregister_by_endpoint(
+                    cfg["nqn"],
+                    cfg["nsid"],
+                    cfg["traddr"],
+                    cfg["trsvcid"],
+                    cfg["master_server_address"]
+                )
+
+                if ret != 0:
+                    raise RuntimeError(f"Unregistration failed with code {ret}")
+
+                logging.info("Unregister SSD %d/%d succeeded", i + 1, total)
+                success_count += 1
+
+            except Exception as e:
+                logging.error("Failed to unregister SSD %d/%d: %s", i + 1, total, e)
+                failed_count += 1
+
+        # Summary
+        logging.info("SSD unregistration summary:")
+        logging.info("- Total SSDs: %d", total)
+        logging.info("- Successfully unregistered: %d", success_count)
+        logging.info("- Failed: %d", failed_count)
+
+        # Return success if all SSDs were unregistered successfully
+        return failed_count == 0
+
+
+def parse_arguments():
+    parser = argparse.ArgumentParser(description='Mooncake SSD Unregister with REST API')
+    parser.add_argument('--master_server_address', type=str,
+                        help='Master server address (e.g., 192.168.65.81:50051)',
+                        default='192.168.65.81:50051')
+    parser.add_argument('--spdk_target_info', action='append',
+                        help='SPDK target information (e.g., "ip:192.168.65.56 ns:2 nqn:nqn.2016-06.io.spdk:cnode1" or "ip:192.168.65.56 path:/home/spdk")',
+                        default=None)
+    parser.add_argument('--username', type=str, default='root',
+                        help='SSH username for target nodes (default: root)')
+    parser.add_argument('--port', type=int, default=22,
+                        help='SSH port for target nodes (default: 22)')
+    parser.add_argument('--password', type=str,
+                        help='SSH password for target nodes')
+    parser.add_argument('--key-file', type=str,
+                        help='SSH private key file path')
+    parser.add_argument('-D', '--define', action='append',
+                        help='Override configuration fields globally (e.g., -Dtrsvcid=4420)',
+                        default=[])
+    args = parser.parse_args()
+
+    # If no spdk_target_info is provided, use default
+    if args.spdk_target_info is None:
+        args.spdk_target_info = ['ip:192.168.65.56 ns:1 nqn:nqn.2016-06.io.spdk:cnode1']
+
+    return args
+
+
+def main():
+    args = parse_arguments()
+
+    cli_config = {}
+    for item in args.define:
+        if '=' in item:
+            key, value = item.split('=', 1)
+            cli_config[key] = value
+        else:
+            logging.warning(f"Ignoring invalid CLI config: {item}")
+
+    # Add master_server_address to cli_config if provided
+    if args.master_server_address:
+        cli_config['master_server_address'] = args.master_server_address
+    cli_config['username'] = args.username
+    cli_config['port'] = args.port
+    if args.password:
+        cli_config['password'] = args.password
+    if args.key_file:
+        cli_config['key_file'] = args.key_file
+
+    unregister = MooncakeNoFUnregister(cli_config, args.spdk_target_info)
+    success = unregister.start_ssd_unregister_service()
+    if not success:
+        exit(1)
+
+if __name__ == "__main__":
+    main()
diff --git a/mooncake-wheel/mooncake/mooncake_store_service.py b/mooncake-wheel/mooncake/mooncake_store_service.py
index 6dae3416..a9158407 100644
--- a/mooncake-wheel/mooncake/mooncake_store_service.py
+++ b/mooncake-wheel/mooncake/mooncake_store_service.py
@@ -23,6 +23,15 @@ def _timed_handler(operation_name, handler):
     return wrapper
 
 
+def _shm_name_to_path(name):
+    if not isinstance(name, str) or not name:
+        return None
+    normalized = name[1:] if name.startswith("/") else name
+    if not normalized or "/" in normalized or normalized in {".", ".."}:
+        return None
+    return f"/dev/shm/{normalized}"
+
+
 class MooncakeStoreService:
     """
     Mooncake Store Service with REST API.
@@ -53,6 +62,12 @@ class MooncakeStoreService:
         self.config = None
         self._setup_logging()
 
+        # State for /api/reconfigure (Prefill/Decode mode switch)
+        self.current_mode = "prefill"          # "prefill" or "decode"
+        self.mounted_segment_ids = []          # persisted segment_ids from last decode mount
+        self.last_mount_info = {}              # last mount parameters for debugging
+        self._state_lock = asyncio.Lock()      # serialize reconfigure/mount/unmount state changes
+
         try:
             if config_path:
                 self.config = MooncakeConfig.from_file(config_path)
@@ -113,7 +128,10 @@ class MooncakeStoreService:
                     self.config.local_buffer_size,
                     self.config.protocol,
                     self.config.device_name,
-                    self.config.master_server_address
+                    self.config.master_server_address,
+                    None,
+                    self.config.enable_ssd_offload,
+                    self.config.ssd_offload_path
                 )
 
                 if ret != 0:
@@ -143,6 +161,11 @@ class MooncakeStoreService:
     async def start_http_service(self, port: int = 8080):
         app = web.Application(client_max_size=1024 * 1024 * 100)  # 100MB limit
         app.add_routes([
+            web.post('/api/reconfigure', _timed_handler("RECONFIGURE", self.handle_reconfigure)),
+            web.post('/api/mount_shm', _timed_handler("MOUNT_SHM", self.handle_mount_shm)),
+            web.post('/api/unmount_shm', _timed_handler("UNMOUNT_SHM", self.handle_unmount_shm)),
+            web.post('/api/mount', _timed_handler("MOUNT", self.handle_mount)),
+            web.post('/api/unmount', _timed_handler("UNMOUNT", self.handle_unmount)),
             web.put('/api/put', _timed_handler("PUT", self.handle_put)),
             web.get('/api/get/{key}', _timed_handler("GET", self.handle_get)),
             web.get('/api/exist/{key}', _timed_handler("EXIST", self.handle_exist)),
@@ -157,19 +180,297 @@ class MooncakeStoreService:
         return True
 
     # REST API handlers
+    async def handle_reconfigure(self, request):
+        try:
+            data = await request.json()
+            mode = data.get("mode", "").lower()
+
+            if mode == "decode":
+                path = data.get("path")
+                size = data.get("size")
+                offset = data.get("offset", 0)
+                protocol = data.get("protocol", self.config.protocol)
+                location = data.get("location", "")
+
+                if not path or size is None:
+                    return web.Response(
+                        status=400,
+                        text=json.dumps({"error": "Missing path or size for decode mode"}),
+                        content_type="application/json"
+                    )
+
+                async with self._state_lock:
+                    # If already in decode mode with mounted segments, unmount them first
+                    if self.mounted_segment_ids:
+                        logging.info("Reconfigure decode: unmounting previous segments before remount")
+                        ret = self.store.unmount_segment(self.mounted_segment_ids)
+                        if ret != 0:
+                            return web.Response(
+                                status=500,
+                                text=json.dumps({"error": f"Unmount of previous segments failed, ret={ret}"}),
+                                content_type="application/json"
+                            )
+                        self.mounted_segment_ids.clear()
+
+                    result = self.store.mount_segment(path, size, offset, protocol, location)
+                    if result["ret"] != 0:
+                        self.current_mode = "prefill"
+                        self.mounted_segment_ids.clear()
+                        self.last_mount_info.clear()
+                        return web.Response(
+                            status=500,
+                            text=json.dumps(
+                                {
+                                    "error": (
+                                        f"Mount failed, ret={result['ret']}; "
+                                        "rolled back to prefill"
+                                    ),
+                                    "mode": self.current_mode,
+                                }
+                            ),
+                            content_type="application/json"
+                        )
+
+                    self.mounted_segment_ids = list(result["segment_ids"])
+                    self.current_mode = "decode"
+                    self.last_mount_info = {
+                        "path": path, "offset": offset, "size": size,
+                        "protocol": protocol, "location": location
+                    }
+
+                return web.Response(
+                    status=200,
+                    text=json.dumps({
+                        "status": "success",
+                        "mode": self.current_mode,
+                        "segment_ids": self.mounted_segment_ids,
+                    }),
+                    content_type="application/json"
+                )
+
+            elif mode == "prefill":
+                async with self._state_lock:
+                    if self.mounted_segment_ids:
+                        ret = self.store.unmount_segment(self.mounted_segment_ids)
+                        if ret != 0:
+                            return web.Response(
+                                status=500,
+                                text=json.dumps({"error": f"Unmount failed, ret={ret}"}),
+                                content_type="application/json"
+                            )
+                        self.mounted_segment_ids.clear()
+
+                    self.current_mode = "prefill"
+                    self.last_mount_info.clear()
+
+                return web.Response(
+                    status=200,
+                    text=json.dumps({"status": "success", "mode": self.current_mode}),
+                    content_type="application/json"
+                )
+
+            else:
+                return web.Response(
+                    status=400,
+                    text=json.dumps({"error": "Invalid mode. Use 'decode' or 'prefill'"}),
+                    content_type="application/json"
+                )
+        except Exception as e:
+            logging.error("RECONFIGURE error: %s", e)
+            return web.Response(
+                status=500,
+                text=json.dumps({"error": str(e)}),
+                content_type="application/json"
+            )
+
+    async def handle_mount_shm(self, request):
+        try:
+            data = await request.json()
+            name = data.get("name")
+            path = _shm_name_to_path(name)
+            offset = data.get("offset", 0)
+            size = data.get("size")
+            protocol = data.get("protocol", self.config.protocol)
+            location = data.get("location", "")
+
+            if not path or size is None:
+                return web.Response(
+                    status=400,
+                    text=json.dumps({"error": "Missing or invalid name or size"}),
+                    content_type="application/json"
+                )
+
+            result = self.store.mount_segment(path, size, offset, protocol, location)
+            if result["ret"] != 0:
+                return web.Response(
+                    status=500,
+                    text=json.dumps({"error": f"Mount failed, ret={result['ret']}"}),
+                    content_type="application/json"
+                )
+
+            return web.Response(
+                status=200,
+                text=json.dumps(
+                    {
+                        "status": "success",
+                        "segment_ids": list(result["segment_ids"]),
+                    }
+                ),
+                content_type="application/json",
+            )
+        except Exception as e:
+            logging.error("MOUNT_SHM error: %s", e)
+            return web.Response(
+                status=500,
+                text=json.dumps({"error": str(e)}),
+                content_type="application/json"
+            )
+
+    async def handle_unmount_shm(self, request):
+        try:
+            data = await request.json()
+            segment_ids = data.get("segment_ids", [])
+            if isinstance(segment_ids, str):
+                segment_ids = [segment_ids]
+            if not segment_ids:
+                return web.Response(
+                    status=400,
+                    text=json.dumps({"error": "Missing segment_ids"}),
+                    content_type="application/json",
+                )
+
+            grace_period_seconds = data.get("grace_period_seconds", 0)
+            failed_segment_ids = []
+            async with self._state_lock:
+                for sid in segment_ids:
+                    ret = self.store.unmount_segment([sid], grace_period_seconds)
+                    if ret != 0:
+                        failed_segment_ids.append(sid)
+                        continue
+                    if sid in self.mounted_segment_ids:
+                        self.mounted_segment_ids.remove(sid)
+                if not self.mounted_segment_ids:
+                    self.current_mode = "prefill"
+
+            if failed_segment_ids:
+                return web.Response(
+                    status=500,
+                    text=json.dumps(
+                        {
+                            "error": "Unmount failed for one or more segments",
+                            "failed_segment_ids": failed_segment_ids,
+                        }
+                    ),
+                    content_type="application/json",
+                )
+
+            return web.Response(
+                status=200,
+                text=json.dumps({"status": "success"}),
+                content_type="application/json",
+            )
+        except Exception as e:
+            logging.error("UNMOUNT_SHM error: %s", e)
+            return web.Response(
+                status=500,
+                text=json.dumps({"error": str(e)}),
+                content_type="application/json"
+            )
+
+    async def handle_mount(self, request):
+        try:
+            data = await request.json()
+            size = data.get("size")
+            protocol = data.get("protocol", self.config.protocol)
+            location = data.get("location", "")
+
+            if type(size) is not int or size <= 0:
+                return web.Response(
+                    status=400,
+                    text=json.dumps({"error": "Invalid size, must be a positive integer"}),
+                    content_type="application/json"
+                )
+
+            result = self.store.allocate_and_mount_segment(size, protocol, location)
+            if result["ret"] != 0:
+                return web.Response(
+                    status=500,
+                    text=json.dumps({"error": f"Allocate and mount failed, ret={result['ret']}"}),
+                    content_type="application/json"
+                )
+
+            return web.Response(
+                status=200,
+                text=json.dumps(
+                    {
+                        "status": "success",
+                        "segment_ids": list(result["segment_ids"]),
+                        "allocated_size": result["allocated_size"],
+                    }
+                ),
+                content_type="application/json",
+            )
+        except Exception as e:
+            logging.error("MOUNT error: %s", e)
+            return web.Response(
+                status=500,
+                text=json.dumps({"error": str(e)}),
+                content_type="application/json"
+            )
+
+    async def handle_unmount(self, request):
+        try:
+            data = await request.json()
+            segment_ids = data.get("segment_ids", [])
+            if isinstance(segment_ids, str):
+                segment_ids = [segment_ids]
+            if not segment_ids:
+                return web.Response(
+                    status=400,
+                    text=json.dumps({"error": "Missing segment_ids"}),
+                    content_type="application/json",
+                )
+
+            grace_period_seconds = data.get("grace_period_seconds", 0)
+            ret = self.store.unmount_and_free_segment(
+                segment_ids, grace_period_seconds
+            )
+            if ret != 0:
+                return web.Response(
+                    status=500,
+                    text=json.dumps(
+                        {"error": f"Unmount and free failed, ret={ret}"}
+                    ),
+                    content_type="application/json",
+                )
+
+            return web.Response(
+                status=200,
+                text=json.dumps({"status": "success"}),
+                content_type="application/json",
+            )
+        except Exception as e:
+            logging.error("UNMOUNT error: %s", e)
+            return web.Response(
+                status=500,
+                text=json.dumps({"error": str(e)}),
+                content_type="application/json"
+            )
+
     async def handle_put(self, request):
         try:
             data = await request.json()
             key = data.get('key')
-            value = data.get('value').encode()
+            raw_value = data.get('value')
 
-            if not key or not value:
+            if not key or raw_value is None:
                 return web.Response(
                     status=400,
                     text=json.dumps({'error': 'Missing key or value'}),
                     content_type='application/json'
                 )
 
+            value = raw_value.encode()
             ret = self.store.put(key, value)
             if ret != 0:
                 return web.Response(
@@ -194,14 +495,42 @@ class MooncakeStoreService:
     async def handle_get(self, request):
         try:
             key = request.match_info['key']
-            value = self.store.get(key)
+            exists = self.store.is_exist(key)
 
-            if not value:
+            if exists == 0:
                 return web.Response(
                     status=404,
                     text=json.dumps({'error': 'Key not found'}),
                     content_type='application/json'
                 )
+            if exists < 0:
+                return web.Response(
+                    status=500,
+                    text=json.dumps({'error': 'Exist check failed'}),
+                    content_type='application/json'
+                )
+
+            value = self.store.get(key)
+            if value is None:
+                return web.Response(
+                    status=500,
+                    text=json.dumps({'error': 'GET operation failed'}),
+                    content_type='application/json'
+                )
+            if value == b"":
+                exists = self.store.is_exist(key)
+                if exists == 0:
+                    return web.Response(
+                        status=404,
+                        text=json.dumps({'error': 'Key not found'}),
+                        content_type='application/json'
+                    )
+                if exists < 0:
+                    return web.Response(
+                        status=500,
+                        text=json.dumps({'error': 'Exist check failed'}),
+                        content_type='application/json'
+                    )
 
             return web.Response(
                 status=200,
@@ -221,9 +550,16 @@ class MooncakeStoreService:
             key = request.match_info['key']
             exists = self.store.is_exist(key)
 
+            if exists < 0:
+                return web.Response(
+                    status=500,
+                    text=json.dumps({'error': 'Exist check failed'}),
+                    content_type='application/json'
+                )
+
             return web.Response(
                 status=200,
-                text=json.dumps({'exists': bool(exists)}),
+                text=json.dumps({'exists': exists > 0}),
                 content_type='application/json'
             )
         except Exception as e:
@@ -339,5 +675,11 @@ async def main():
         await service.stop()
         raise
 
+
+def sync_main():
+    """Synchronous entry point for ``mc_store_rest_server`` CLI."""
+    return asyncio.run(main())
+
+
 if __name__ == "__main__":
-    asyncio.run(main())
+    sync_main()
diff --git a/mooncake-wheel/mooncake/spdk_tgt_create.py b/mooncake-wheel/mooncake/spdk_tgt_create.py
new file mode 100644
index 00000000..f5d85543
--- /dev/null
+++ b/mooncake-wheel/mooncake/spdk_tgt_create.py
@@ -0,0 +1,520 @@
+#!/usr/bin/env python3
+# Tool to remotely create SPDK targets on multiple nodes
+# Usage: python3 -m mooncake.spdk_tgt_create --spdk_target_info="ip:192.168.65.56 path:/home/spdk pci:0000:01:00.0,0000:02:00.0" --spdk_target_info="ip:192.168.65.57 path:/home/spdk"
+
+import argparse
+import logging
+import paramiko
+import re
+import shlex
+import time
+from typing import List, Dict, Any, Optional
+
+
+class SPDKTgtCreator:
+    """
+    Remotely creates SPDK targets on multiple nodes via SSH.
+    """
+
+    DEFAULT_TRANSPORT_OPTIONS = {
+        'trtype': 'RDMA',
+        'max_queue_depth': 128,
+        'max_io_qpairs_per_ctrlr': 127,
+        'max_io_size': 4096,
+        'in_capsule_data_size': 131072,
+        'io_unit_size': 131072,
+        'max_aq_depth': 128,
+        'num_shared_buffers': 4096,
+        'buf_cache_size': 32,
+    }
+
+    TRANSPORT_RPC_FLAGS = {
+        'trtype': '-t',
+        'max_queue_depth': '-q',
+        'max_io_qpairs_per_ctrlr': '-m',
+        'max_io_size': '-c',
+        'in_capsule_data_size': '-i',
+        'io_unit_size': '-u',
+        'max_aq_depth': '-a',
+        'num_shared_buffers': '-n',
+        'buf_cache_size': '-b',
+    }
+
+    def __init__(self, spdk_targets: List[str], transport_options: Dict[str, Any] = None, core_mask: str = '0xff'):
+        self.spdk_targets = spdk_targets
+        self.core_mask = core_mask
+        self.transport_options = dict(self.DEFAULT_TRANSPORT_OPTIONS)
+        if transport_options:
+            self.transport_options.update(transport_options)
+        self._setup_logging()
+        self.target_configs = self._parse_spdk_targets()
+
+    def _setup_logging(self):
+        logging.basicConfig(
+            level=logging.INFO,
+            format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
+        )
+        self.logger = logging.getLogger(self.__class__.__name__)
+
+    def _parse_spdk_targets(self) -> List[Dict[str, Any]]:
+        """
+        Parse SPDK target information from command line arguments.
+        Format: "ip:<ip> path:<spdk_path> [pci:<pci1>,<pci2> ...]"
+        """
+        target_configs = []
+
+        for target_info in self.spdk_targets:
+            target = {
+                'ip': None,
+                'path': None,
+                'pci_devices': []
+            }
+
+            # Split the target info by spaces
+            parts = target_info.split()
+
+            # Simple state machine to parse the target info
+            state = None  # Can be 'ip', 'path', or 'pci'
+
+            for part in parts:
+                if ':' in part:
+                    # This is a key-value pair
+                    key, value = part.split(':', 1)
+                    key = key.strip()
+                    value = value.strip()
+
+                    if key == 'ip':
+                        target['ip'] = value
+                        state = 'ip'
+                    elif key == 'path':
+                        target['path'] = value
+                        state = 'path'
+                    elif key == 'pci':
+                        # Parse PCI devices separated by commas
+                        if value:
+                            # Split by commas and strip whitespace
+                            pci_list = [dev.strip() for dev in value.split(',') if dev.strip()]
+                            target['pci_devices'].extend(pci_list)
+                        state = 'pci'
+                    elif state == 'pci':
+                        pci_list = [dev.strip() for dev in part.split(',') if dev.strip()]
+                        target['pci_devices'].extend(pci_list)
+                else:
+                    # This is a continuation of the current state
+                    if state == 'path':
+                        # Path might contain spaces (unlikely but possible)
+                        target['path'] += ' ' + part
+                    elif state == 'pci':
+                        pci_list = [dev.strip() for dev in part.split(',') if dev.strip()]
+                        target['pci_devices'].extend(pci_list)
+
+            # Validate required fields
+            if not target['ip']:
+                raise ValueError("Each spdk_target_info must contain 'ip' field")
+            if not target['path']:
+                raise ValueError("Each spdk_target_info must contain 'path' field")
+
+            target_configs.append(target)
+            pci_info = target['pci_devices'] if target['pci_devices'] else 'auto-discover'
+            self.logger.info(f"Parsed target: IP={target['ip']}, Path={target['path']}, PCI devices={pci_info}")
+
+        return target_configs
+
+    def _ssh_connect(self, ip: str, username: str = 'root', password: str = None, key_file: str = None) -> paramiko.SSHClient:
+        """
+        Establish an SSH connection to the target host.
+        """
+        ssh = paramiko.SSHClient()
+        ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
+
+        try:
+            if key_file:
+                self.logger.info(f"Connecting to {ip} using key file {key_file}")
+                ssh.connect(ip, username=username, key_filename=key_file)
+            else:
+                self.logger.info(f"Connecting to {ip} using password authentication")
+                ssh.connect(ip, username=username, password=password)
+            return ssh
+        except Exception as e:
+            self.logger.error(f"Failed to connect to {ip}: {e}")
+            raise
+
+    def _execute_command(self, ssh: paramiko.SSHClient, command: str, working_dir: str = None, sudo: bool = False, log_errors: bool = True, timeout: Optional[int] = 30) -> tuple:
+        """
+        Execute a command on the remote host via SSH.
+        """
+        if working_dir:
+            command = f"cd {working_dir} && {command}"
+
+        if sudo:
+            command = f"sudo {command}"
+
+        self.logger.debug(f"Executing command: {command}")
+
+        stdin, stdout, stderr = ssh.exec_command(command, timeout=timeout)
+        exit_status = stdout.channel.recv_exit_status()
+        output = stdout.read().decode('utf-8')
+        error = stderr.read().decode('utf-8')
+
+        if output:
+            self.logger.debug(f"Command output: {output}")
+        if error:
+            self.logger.debug(f"Command error: {error}")
+        if exit_status != 0:
+            if log_errors:
+                self.logger.error(f"Command failed with exit code {exit_status}: {command}")
+                if output:
+                    self.logger.error(f"Command output: {output}")
+                self.logger.error(f"Error output: {error}")
+            raise RuntimeError(f"Command execution failed: {error or output}")
+
+        return output, error
+
+    def _discover_nvme_pci_devices(self, ssh: paramiko.SSHClient) -> List[str]:
+        """
+        Discover SPDK-ready or unmounted NVMe controller PCI addresses on the target host.
+        """
+        self.logger.info("No PCI devices specified, discovering SPDK-ready or unmounted NVMe PCI devices")
+        output, _ = self._execute_command(
+            ssh,
+            r"""for dev in /sys/bus/pci/devices/*; do
+  class=$(cat "$dev/class" 2>/dev/null || true)
+  case "$class" in
+    0x0108*) ;;
+    *) continue ;;
+  esac
+  pci=$(basename "$dev")
+  driver=$(basename "$(readlink "$dev/driver" 2>/dev/null)" 2>/dev/null || true)
+  case "$driver" in
+    vfio-pci|uio_pci_generic|igb_uio)
+      echo "USE $pci $driver"
+      continue
+      ;;
+  esac
+  has_block=0
+  mounted=0
+  for block in /sys/block/nvme*n*; do
+    [ -e "$block" ] || continue
+    real_device=$(readlink -f "$block/device" 2>/dev/null || true)
+    case "$real_device" in
+      *"/$pci"/*|*"/$pci") ;;
+      *) continue ;;
+    esac
+    has_block=1
+    disk=$(basename "$block")
+    if lsblk -nr -o MOUNTPOINT "/dev/$disk" 2>/dev/null | grep -q '[^[:space:]]'; then
+        mounted=1
+    fi
+  done
+  if [ "$has_block" -eq 0 ]; then
+    echo "SKIP $pci no_block_device"
+  elif [ "$mounted" -eq 0 ]; then
+    echo "USE $pci"
+  else
+    echo "SKIP $pci mounted"
+  fi
+done"""
+        )
+
+        pci_devices = []
+        pci_pattern = re.compile(
+            r'^(?:[0-9a-fA-F]{4}:)?[0-9a-fA-F]{2}:[0-9a-fA-F]{2}\.[0-7]$'
+        )
+        for line in output.splitlines():
+            fields = line.split()
+            if len(fields) < 2:
+                continue
+            action, pci = fields[0], fields[1]
+            if action == 'USE' and pci_pattern.match(pci):
+                pci_devices.append(pci)
+            elif action == 'SKIP' and pci_pattern.match(pci):
+                reason = ' '.join(fields[2:]) or 'not eligible'
+                self.logger.warning(f"Skipping NVMe PCI device {pci}: {reason}")
+
+        if not pci_devices:
+            raise RuntimeError("No SPDK-ready or unmounted NVMe PCI devices found on the target host")
+
+        self.logger.info(f"Selected NVMe PCI devices for SPDK setup: {', '.join(pci_devices)}")
+        return pci_devices
+
+    def _filter_spdk_ready_pci_devices(self, ssh: paramiko.SSHClient, pci_devices: List[str], strict: bool) -> List[str]:
+        """
+        Keep PCI devices that are bound to an SPDK-compatible userspace driver.
+        """
+        if not pci_devices:
+            return []
+
+        pci_args = ' '.join(shlex.quote(pci) for pci in pci_devices)
+        output, _ = self._execute_command(
+            ssh,
+            f"""for pci in {pci_args}; do
+  driver=$(basename "$(readlink "/sys/bus/pci/devices/$pci/driver" 2>/dev/null)" 2>/dev/null || true)
+  case "$driver" in
+    vfio-pci|uio_pci_generic|igb_uio) echo "READY $pci $driver" ;;
+    "") echo "NOT_READY $pci no_driver" ;;
+    *) echo "NOT_READY $pci $driver" ;;
+  esac
+done"""
+        )
+
+        ready_devices = []
+        not_ready_devices = []
+        for line in output.splitlines():
+            fields = line.split()
+            if len(fields) < 3:
+                continue
+            state, pci, driver = fields[0], fields[1], fields[2]
+            if state == 'READY':
+                ready_devices.append(pci)
+            elif state == 'NOT_READY':
+                not_ready_devices.append((pci, driver))
+
+        if not_ready_devices:
+            details = ', '.join(f"{pci} ({driver})" for pci, driver in not_ready_devices)
+            if strict:
+                raise RuntimeError(
+                    "Some PCI devices are not available to SPDK after setup.sh: "
+                    f"{details}. Check whether they are mounted or still bound to the kernel driver."
+                )
+            self.logger.warning(f"Skipping PCI devices not available to SPDK after setup.sh: {details}")
+
+        if not ready_devices:
+            raise RuntimeError("No PCI devices are available to SPDK after setup.sh")
+
+        self.logger.info(f"PCI devices available to SPDK: {', '.join(ready_devices)}")
+        return ready_devices
+
+    def _start_spdk_tgt(self, ssh: paramiko.SSHClient, spdk_path: str) -> None:
+        """
+        Start the SPDK NVMF target service in the background.
+        """
+        # Check if tgt is already running
+        try:
+            self._execute_command(ssh, "pgrep -x nvmf_tgt", log_errors=False)
+            self.logger.info("SPDK tgt service is already running, stopping it first")
+            self._execute_command(ssh, "pkill -x nvmf_tgt")
+            time.sleep(2)  # Give it time to stop
+        except RuntimeError:
+            self.logger.debug("SPDK tgt service is not running, will start it")
+
+        # Start tgt in the background using absolute path
+        self.logger.info(f"Starting SPDK tgt service with core mask {self.core_mask}")
+        tgt_binary = f"{spdk_path}/build/bin/nvmf_tgt"
+        log_file = f"{spdk_path}/tgt.log"
+        self._execute_command(
+            ssh,
+            f"nohup {tgt_binary} -m {shlex.quote(self.core_mask)} > {log_file} 2>&1 &",
+            timeout=None
+        )
+        time.sleep(3)  # Give it time to start
+
+    def _setup_spdk(self, ssh: paramiko.SSHClient, spdk_path: str, pci_devices: List[str]) -> None:
+        """
+        Setup SPDK with the specified PCI devices.
+        """
+        self.logger.info(f"Setting up SPDK with PCI devices: {', '.join(pci_devices)}")
+        pci_allowed = ' '.join(pci_devices)
+        setup_script = f"{spdk_path}/scripts/setup.sh"
+        self._execute_command(ssh, f"PCI_ALLOWED='{pci_allowed}' {setup_script}", sudo=True)
+
+    def _format_transport_options(self) -> str:
+        """
+        Format transport options for the SPDK nvmf_create_transport RPC.
+        """
+        formatted_options = []
+        for option, flag in self.TRANSPORT_RPC_FLAGS.items():
+            value = self.transport_options[option]
+            formatted_options.extend([flag, shlex.quote(str(value))])
+        return ' '.join(formatted_options)
+
+    def _create_transport(self, ssh: paramiko.SSHClient, spdk_path: str) -> None:
+        """
+        Create NVMe-oF transport for SPDK.
+        """
+        self.logger.info(f"Creating {self.transport_options['trtype']} transport")
+        rpc_script = f"{spdk_path}/scripts/rpc.py"
+        transport_options = self._format_transport_options()
+        self._execute_command(ssh, f"{rpc_script} nvmf_create_transport {transport_options}")
+
+    def _create_bdevs(self, ssh: paramiko.SSHClient, spdk_path: str, pci_devices: List[str]) -> List[str]:
+        """
+        Create block devices for the specified PCI devices.
+        """
+        bdevs = []
+        rpc_script = f"{spdk_path}/scripts/rpc.py"
+        for i, pci in enumerate(pci_devices):
+            bdev_name = f"Nvme{i}"
+            self.logger.info(f"Creating bdev {bdev_name} for PCI {pci}")
+            self._execute_command(ssh, f"{rpc_script} bdev_nvme_attach_controller -b {bdev_name} -t PCIe -a {pci}")
+            bdevs.append(f"{bdev_name}n1")
+            self.logger.info(f"Attached PCI {pci} as bdev {bdev_name}n1")
+        return bdevs
+
+    def _create_subsystem(self, ssh: paramiko.SSHClient, spdk_path: str) -> str:
+        """
+        Create an NVMF subsystem.
+        """
+        subsystem_nqn = "nqn.2016-06.io.spdk:cnode1"
+        self.logger.info(f"Creating subsystem {subsystem_nqn}")
+        rpc_script = f"{spdk_path}/scripts/rpc.py"
+        self._execute_command(ssh, f"{rpc_script} nvmf_create_subsystem {subsystem_nqn} -a -s SPDK00000000000001 -m 12")
+        return subsystem_nqn
+
+    def _add_namespaces(self, ssh: paramiko.SSHClient, spdk_path: str, subsystem_nqn: str, bdevs: List[str]) -> None:
+        """
+        Add namespaces to the subsystem.
+        """
+        rpc_script = f"{spdk_path}/scripts/rpc.py"
+        for bdev in bdevs:
+            self.logger.info(f"Adding namespace {bdev} to subsystem {subsystem_nqn}")
+            self._execute_command(ssh, f"{rpc_script} nvmf_subsystem_add_ns {subsystem_nqn} {bdev}")
+
+    def _add_listener(self, ssh: paramiko.SSHClient, spdk_path: str, subsystem_nqn: str, ip: str) -> None:
+        """
+        Add a listener to the subsystem.
+        """
+        trtype = self.transport_options['trtype']
+        self.logger.info(f"Adding {trtype} listener on {ip}:4420")
+        rpc_script = f"{spdk_path}/scripts/rpc.py"
+        self._execute_command(ssh, f"{rpc_script} nvmf_subsystem_add_listener {subsystem_nqn} -t {shlex.quote(str(trtype))} -a {ip} -s 4420")
+
+    def deploy_target(self, target_config: Dict[str, Any]) -> bool:
+        """
+        Deploy SPDK target on a single node.
+        """
+        ip = target_config['ip']
+        spdk_path = target_config['path']
+        pci_devices = target_config['pci_devices']
+
+        self.logger.info(f"Deploying SPDK target on {ip}")
+
+        try:
+            # Establish SSH connection
+            ssh = self._ssh_connect(ip)
+
+            try:
+                auto_discovered = not pci_devices
+                if not pci_devices:
+                    pci_devices = self._discover_nvme_pci_devices(ssh)
+                else:
+                    self.logger.info(f"Using explicitly specified PCI devices for {ip}: {', '.join(pci_devices)}")
+
+                # Setup SPDK
+                self._setup_spdk(ssh, spdk_path, pci_devices)
+                pci_devices = self._filter_spdk_ready_pci_devices(
+                    ssh,
+                    pci_devices,
+                    strict=not auto_discovered
+                )
+                self.logger.info(f"Target {ip} will expose PCI devices: {', '.join(pci_devices)}")
+
+                # Start tgt service
+                self._start_spdk_tgt(ssh, spdk_path)
+
+                # Create transport
+                self._create_transport(ssh, spdk_path)
+
+                # Create bdevs
+                bdevs = self._create_bdevs(ssh, spdk_path, pci_devices)
+
+                # Create subsystem
+                subsystem_nqn = self._create_subsystem(ssh, spdk_path)
+
+                # Add namespaces
+                self._add_namespaces(ssh, spdk_path, subsystem_nqn, bdevs)
+
+                # Add listener
+                self._add_listener(ssh, spdk_path, subsystem_nqn, ip)
+
+                self.logger.info(f"Successfully deployed SPDK target on {ip}")
+                return True
+            finally:
+                ssh.close()
+        except Exception as e:
+            self.logger.error(f"Failed to deploy SPDK target on {ip}: {e}")
+            return False
+
+    def deploy_all_targets(self) -> bool:
+        """
+        Deploy SPDK targets on all specified nodes.
+        """
+        success_count = 0
+        total_count = len(self.target_configs)
+
+        for i, target_config in enumerate(self.target_configs):
+            self.logger.info(f"=== Deploying target {i+1}/{total_count} ===")
+            if self.deploy_target(target_config):
+                success_count += 1
+
+        self.logger.info("=== Deployment Summary ===")
+        self.logger.info(f"Total targets: {total_count}")
+        self.logger.info(f"Successfully deployed: {success_count}")
+        self.logger.info(f"Failed: {total_count - success_count}")
+
+        return success_count == total_count
+
+
+def parse_arguments():
+    parser = argparse.ArgumentParser(description='SPDK Target Creator Tool')
+    parser.add_argument(
+        '--spdk_target_info',
+        action='append',
+        help=('SPDK target information (e.g., "ip:192.168.65.56 '
+              'path:/home/spdk pci:0000:01:00.0,0000:02:00.0"). If pci is '
+              'omitted, SPDK-ready or unmounted NVMe devices on the target are used.'),
+        required=True
+    )
+    parser.add_argument('--core-mask', type=str, default='0xff',
+                        help='CPU core mask used to start nvmf_tgt with -m (default: 0xff)')
+    parser.add_argument('--transport-type', type=str, default='RDMA',
+                        help='NVMe-oF transport type for nvmf_create_transport (default: RDMA)')
+    parser.add_argument('--max-queue-depth', type=int, default=128,
+                        help='Max number of outstanding I/O per queue (default: 128)')
+    parser.add_argument('--max-io-qpairs-per-ctrlr', type=int, default=127,
+                        help='Max number of I/O qpairs per controller (default: 127)')
+    parser.add_argument('--max-io-size', type=int, default=4096,
+                        help='Max I/O size in bytes (default: 4096)')
+    parser.add_argument('--in-capsule-data-size', type=int, default=131072,
+                        help='Max in-capsule data size in bytes (default: 131072)')
+    parser.add_argument('--io-unit-size', type=int, default=131072,
+                        help='I/O unit size in bytes (default: 131072)')
+    parser.add_argument('--max-aq-depth', type=int, default=128,
+                        help='Max number of admin commands per admin queue (default: 128)')
+    parser.add_argument('--num-shared-buffers', type=int, default=4096,
+                        help='Number of pooled data buffers available to the transport (default: 4096)')
+    parser.add_argument('--buf-cache-size', type=int, default=32,
+                        help='Number of shared buffers reserved for each poll group (default: 32)')
+    parser.add_argument('--username', type=str, default='root',
+                        help='SSH username for target nodes (default: root)')
+    parser.add_argument('--password', type=str,
+                        help='SSH password for target nodes')
+    parser.add_argument('--key-file', type=str,
+                        help='SSH private key file path')
+    return parser.parse_args()
+
+
+def main():
+    args = parse_arguments()
+    transport_options = {
+        'trtype': args.transport_type,
+        'max_queue_depth': args.max_queue_depth,
+        'max_io_qpairs_per_ctrlr': args.max_io_qpairs_per_ctrlr,
+        'max_io_size': args.max_io_size,
+        'in_capsule_data_size': args.in_capsule_data_size,
+        'io_unit_size': args.io_unit_size,
+        'max_aq_depth': args.max_aq_depth,
+        'num_shared_buffers': args.num_shared_buffers,
+        'buf_cache_size': args.buf_cache_size,
+    }
+
+    try:
+        creator = SPDKTgtCreator(args.spdk_target_info, transport_options, args.core_mask)
+        success = creator.deploy_all_targets()
+        exit(0 if success else 1)
+    except Exception as e:
+        logging.error(f"Error: {e}")
+        exit(1)
+
+
+if __name__ == '__main__':
+    main()
diff --git a/mooncake-wheel/mooncake/structured_object_store.py b/mooncake-wheel/mooncake/structured_object_store.py
new file mode 100644
index 00000000..b13dbde9
--- /dev/null
+++ b/mooncake-wheel/mooncake/structured_object_store.py
@@ -0,0 +1,4548 @@
+from __future__ import annotations
+
+import ctypes
+import io
+import json
+import uuid
+from concurrent.futures import Future, ThreadPoolExecutor, as_completed
+from contextlib import contextmanager
+from dataclasses import dataclass, field
+from typing import Any, Iterator, Literal, Mapping, Optional, Protocol, Sequence
+
+import numpy as np
+
+try:
+    import mooncake.store as _mooncake_store
+except Exception:  # pragma: no cover - depends on built extension
+    _mooncake_store = None  # type: ignore[assignment]
+
+import msgpack as _msgpack
+
+DEFAULT_BUNDLE_CHUNK_BYTES = 512 * 1024**2
+AUTO_PARALLEL_MIN_BYTES = 4 * 1024**3
+AUTO_PARALLEL_MIN_CHUNKS = 8
+MISSING_OBJECT_ERROR = (
+    -704
+)  # Mooncake remove returns -704 for an already-missing object.
+STRUCTURED_FIELD_SPECS_KEY = "__mooncake_structured_fields__"
+
+
+class BundleStore(Protocol):
+    def put(self, key: str, value: Any) -> int: ...
+
+    def get(self, key: str) -> bytes: ...
+
+    def remove(self, key: str, force: bool = False) -> int: ...
+
+
+@dataclass(frozen=True)
+class BundleTransferPolicy:
+    """Controls generic bundle transfer parallelism."""
+
+    max_inflight_put: int = 1
+    put_mode: Literal["auto", "batch", "parallel"] = "auto"
+    copy_mode: Literal["auto", "zero_copy", "copy"] = "auto"
+
+
+@dataclass
+class RemoteBundleRef:
+    """Reference to a generic named-buffer bundle stored in Mooncake."""
+
+    manifest_key: str
+    manifest: dict[str, Any]
+
+
+@dataclass(frozen=True, init=False)
+class StructuredObjectPayload:
+    """Structured object encoded as metadata plus named buffers."""
+
+    metadata: Mapping[str, Any]
+    buffers: Mapping[str, Any]
+
+    def __init__(
+        self,
+        metadata: Optional[Mapping[str, Any]] = None,
+        buffers: Optional[Mapping[str, Any]] = None,
+    ) -> None:
+        if buffers is None:
+            raise TypeError("StructuredObjectPayload requires buffers")
+        object.__setattr__(self, "metadata", {} if metadata is None else metadata)
+        object.__setattr__(self, "buffers", buffers)
+
+
+@dataclass
+class StructuredObjectResult:
+    """Materialized structured object data fetched from Mooncake."""
+
+    metadata: dict[str, Any]
+    objects: dict[str, Any]
+
+
+@dataclass(frozen=True)
+class StructuredFieldLocation:
+    """Location of a DataProto field inside stage-level structured objects."""
+
+    stage: str
+    member: str
+    section: Literal["batch", "non_tensor_batch"]
+
+
+@dataclass
+class MooncakeDataProtoRef:
+    """Lightweight DataProto handle backed by structured object refs."""
+
+    batch_size: int
+    stage_refs: dict[str, RemoteBundleRef]
+    field_index: dict[str, StructuredFieldLocation]
+    meta_info: dict[str, Any]
+    namespace: str = "default"
+    partition: str = "default"
+    global_indexes: list[int] | None = None
+    encoded_non_tensor: dict[str, Any] = field(default_factory=dict)
+
+
+@dataclass(frozen=True)
+class _TensorPayload:
+    tensor: Any
+
+
+@dataclass(frozen=True)
+class _TensorObjectBufferPayload:
+    ptr: int
+    size: int
+    owner: Any
+    batch_size: int | None = None
+
+    def __len__(self) -> int:
+        if self.batch_size is None:
+            raise TypeError(
+                "tensor object buffer requires batch_size for DataProto use"
+            )
+        return self.batch_size
+
+
+@dataclass(frozen=True)
+class _RawDestinationBuffer:
+    ptr: int
+    size: int
+    owner: Any
+    pre_registered: bool = False
+
+
+@dataclass(frozen=True)
+class StructuredMemberSlice:
+    """Slice description for one structured member."""
+
+    axis: int
+    start: int
+    end: int
+    step: int = 1
+
+
+@dataclass(frozen=True)
+class StructuredObjectReadSpec:
+    """Lazy read plan for structured object materialization."""
+
+    ref: RemoteBundleRef | Mapping[str, Any]
+    member_names: tuple[str, ...] | None = None
+    member_slices: tuple[tuple[str, StructuredMemberSlice], ...] = ()
+
+    def select_members(
+        self, names: Optional[Sequence[str]]
+    ) -> "StructuredObjectReadSpec":
+        """Return a new spec that materializes only the selected members."""
+        return StructuredObjectReadSpec(
+            ref=self.ref,
+            member_names=None if names is None else tuple(names),
+            member_slices=self.member_slices,
+        )
+
+    def slice_member(
+        self,
+        name: str,
+        axis: int,
+        start: int,
+        end: int,
+        step: int = 1,
+    ) -> "StructuredObjectReadSpec":
+        """Return a new spec with a slice applied to one member."""
+        slice_spec = StructuredMemberSlice(axis=axis, start=start, end=end, step=step)
+        member_slices = dict(self.member_slices)
+        member_slices[name] = slice_spec
+        return StructuredObjectReadSpec(
+            ref=self.ref,
+            member_names=self.member_names,
+            member_slices=tuple(member_slices.items()),
+        )
+
+    def member_slice(self, name: str) -> StructuredMemberSlice | None:
+        """Return the configured slice for a member, if any."""
+        return dict(self.member_slices).get(name)
+
+
+@dataclass(frozen=True)
+class _NdarrayReadPlan:
+    dtype: np.dtype[Any]
+    full_shape: tuple[int, ...]
+    output_shape: tuple[int, ...]
+    byte_offset: int
+    byte_length: int
+    step: int
+    cover_row_count: int
+
+
+@dataclass(frozen=True)
+class _DataProtoRowSelection:
+    count: int
+    member_slice: StructuredMemberSlice | None = None
+    indices: tuple[int, ...] | None = None
+
+
+DATAPROTO_REF_HANDLE_TYPE = "mooncake_dataproto_ref"
+DATAPROTO_REF_HANDLE_VERSION = 1
+DataProtoRefLike = MooncakeDataProtoRef | Mapping[str, Any]
+
+
+def is_dataproto_ref_handle(value: Any) -> bool:
+    return isinstance(value, Mapping) and value.get("type") == DATAPROTO_REF_HANDLE_TYPE
+
+
+def export_dataproto_ref(ref: MooncakeDataProtoRef) -> dict[str, Any]:
+    """Export a DataProto ref as a JSON-safe transport handle."""
+    if not isinstance(ref, MooncakeDataProtoRef):
+        raise TypeError(f"expected MooncakeDataProtoRef, got {type(ref).__name__}")
+    handle = {
+        "type": DATAPROTO_REF_HANDLE_TYPE,
+        "version": DATAPROTO_REF_HANDLE_VERSION,
+        "kind": "bundle_stages",
+        "batch_size": int(ref.batch_size),
+        "namespace": ref.namespace,
+        "partition": ref.partition,
+        "global_indexes": _json_safe_value(ref.global_indexes),
+        "stage_refs": {
+            stage: {"manifest_key": stage_ref.manifest_key}
+            for stage, stage_ref in ref.stage_refs.items()
+        },
+        "field_index": {
+            name: {
+                "stage": location.stage,
+                "member": location.member,
+                "section": location.section,
+            }
+            for name, location in ref.field_index.items()
+        },
+        "meta_info": _json_safe_value(ref.meta_info),
+        "encoded_non_tensor": _json_safe_value(ref.encoded_non_tensor),
+    }
+    json.dumps(handle, ensure_ascii=False)
+    return handle
+
+
+def import_dataproto_ref(handle: Mapping[str, Any]) -> MooncakeDataProtoRef:
+    """Import a JSON-safe DataProto transport handle into a lazy ref."""
+    if not is_dataproto_ref_handle(handle):
+        raise ValueError("not a Mooncake DataProto ref handle")
+    if int(handle.get("version", -1)) != DATAPROTO_REF_HANDLE_VERSION:
+        raise ValueError(
+            f"unsupported DataProto ref handle version: {handle.get('version')!r}"
+        )
+    if handle.get("kind") != "bundle_stages":
+        raise ValueError(
+            f"unsupported DataProto ref handle kind: {handle.get('kind')!r}"
+        )
+    stage_refs: dict[str, RemoteBundleRef] = {}
+    for stage, stage_ref in _require_mapping(
+        handle.get("stage_refs"), "stage_refs"
+    ).items():
+        manifest_key = _require_mapping(stage_ref, f"stage_refs[{stage!r}]").get(
+            "manifest_key"
+        )
+        if not isinstance(stage, str) or not isinstance(manifest_key, str):
+            raise ValueError(
+                "DataProto ref handle stage refs must contain string manifest_key values"
+            )
+        stage_refs[stage] = RemoteBundleRef(manifest_key=manifest_key, manifest={})
+    field_index: dict[str, StructuredFieldLocation] = {}
+    for name, raw_location in _require_mapping(
+        handle.get("field_index"), "field_index"
+    ).items():
+        location = _require_mapping(raw_location, f"field_index[{name!r}]")
+        section = location.get("section")
+        if section not in {"batch", "non_tensor_batch"}:
+            raise ValueError(f"invalid DataProto field section: {section!r}")
+        stage = location.get("stage")
+        member = location.get("member")
+        if (
+            not isinstance(name, str)
+            or not isinstance(stage, str)
+            or not isinstance(member, str)
+        ):
+            raise ValueError("DataProto ref handle field locations must be strings")
+        if stage not in stage_refs:
+            raise ValueError(
+                f"DataProto ref handle field {name!r} references unknown stage {stage!r}"
+            )
+        field_index[name] = StructuredFieldLocation(
+            stage=stage, member=member, section=section
+        )
+    return MooncakeDataProtoRef(
+        batch_size=int(handle["batch_size"]),
+        stage_refs=stage_refs,
+        field_index=field_index,
+        meta_info=dict(_require_mapping(handle.get("meta_info", {}), "meta_info")),
+        namespace=str(handle.get("namespace", "default")),
+        partition=str(handle.get("partition", "default")),
+        global_indexes=_import_global_indexes(handle.get("global_indexes")),
+        encoded_non_tensor=dict(
+            _require_mapping(handle.get("encoded_non_tensor", {}), "encoded_non_tensor")
+        ),
+    )
+
+
+def _resolve_dataproto_ref(ref: DataProtoRefLike) -> MooncakeDataProtoRef:
+    if isinstance(ref, MooncakeDataProtoRef):
+        return ref
+    if is_dataproto_ref_handle(ref):
+        return import_dataproto_ref(ref)
+    raise TypeError(
+        f"expected MooncakeDataProtoRef or DataProto handle, got {type(ref).__name__}"
+    )
+
+
+def _require_mapping(value: Any, name: str) -> Mapping[str, Any]:
+    if not isinstance(value, Mapping):
+        raise ValueError(f"DataProto ref handle {name} must be a mapping")
+    return value
+
+
+def _import_global_indexes(value: Any) -> list[int] | None:
+    if value is None:
+        return None
+    if not isinstance(value, list):
+        raise ValueError("DataProto ref handle global_indexes must be a list or null")
+    indexes: list[int] = []
+    for index in value:
+        if isinstance(index, bool) or not isinstance(index, int):
+            raise ValueError(
+                "DataProto ref handle global_indexes must contain integers"
+            )
+        indexes.append(index)
+    return indexes
+
+
+def _json_safe_value(value: Any) -> Any:
+    if value is None or isinstance(value, (str, int, float, bool)):
+        return value
+    if isinstance(value, np.generic):
+        return value.item()
+    if isinstance(value, np.ndarray):
+        return _json_safe_value(value.tolist())
+    if isinstance(value, Mapping):
+        return {str(key): _json_safe_value(val) for key, val in value.items()}
+    if isinstance(value, (list, tuple)):
+        return [_json_safe_value(item) for item in value]
+    raise TypeError(
+        f"DataProto ref handles support only JSON-safe values, got {type(value).__name__}"
+    )
+
+
+class MooncakeBundleTransfer:
+    """Transfer structured objects through Mooncake, with a low-level bundle fallback."""
+
+    def __init__(
+        self,
+        store: BundleStore,
+        key_prefix: str = "bundle",
+        default_chunk_bytes: int = DEFAULT_BUNDLE_CHUNK_BYTES,
+        buffer_pool: Any = None,
+    ) -> None:
+        """Initialize a bundle transfer helper with a configurable default chunk size."""
+        self.store = store
+        self.key_prefix = _normalize_key_prefix(key_prefix)
+        self.default_chunk_bytes = _validate_chunk_bytes(default_chunk_bytes)
+        self._transport = _MooncakePayloadTransport(store, buffer_pool=buffer_pool)
+        self._bundle_store = _BundleManifestStore(
+            store=store,
+            transport=self._transport,
+            key_prefix=self.key_prefix,
+            default_chunk_bytes=self.default_chunk_bytes,
+        )
+        self._structured_store = _StructuredObjectLayer(self._bundle_store)
+
+    def put_bundle(
+        self,
+        meta: bytes | bytearray | memoryview,
+        buffers: Mapping[str, Any],
+        partition: str = "default",
+        chunk_bytes: Optional[int] = None,
+        policy: Optional[BundleTransferPolicy] = None,
+        max_inflight_put: Optional[int] = None,
+        pre_registered_buffers: Optional[Mapping[str, bool]] = None,
+    ) -> RemoteBundleRef:
+        """Store raw metadata bytes plus named buffers as a low-level bundle."""
+        return self._bundle_store.put_bundle(
+            meta=meta,
+            buffers=buffers,
+            partition=partition,
+            chunk_bytes=chunk_bytes,
+            policy=policy,
+            max_inflight_put=max_inflight_put,
+            pre_registered_buffers=pre_registered_buffers,
+        )
+
+    def remove_bundle(self, ref: RemoteBundleRef | Mapping[str, Any]) -> None:
+        """Remove all Mooncake objects that belong to a stored bundle."""
+        self._bundle_store.remove_bundle(ref)
+
+    def put_structured_object(
+        self,
+        payload: StructuredObjectPayload,
+        partition: str = "default",
+        chunk_bytes: Optional[int] = None,
+        policy: Optional[BundleTransferPolicy] = None,
+        max_inflight_put: Optional[int] = None,
+        pre_registered_buffers: Optional[Mapping[str, bool]] = None,
+    ) -> RemoteBundleRef:
+        """Store a structured object described by JSON metadata plus named members."""
+        return self._structured_store.put_structured_object(
+            payload=payload,
+            partition=partition,
+            chunk_bytes=chunk_bytes,
+            policy=policy,
+            max_inflight_put=max_inflight_put,
+            pre_registered_buffers=pre_registered_buffers,
+        )
+
+    def put_object(
+        self,
+        obj: Any,
+        partition: str = "default",
+        chunk_bytes: Optional[int] = None,
+        policy: Optional[BundleTransferPolicy] = None,
+        max_inflight_put: Optional[int] = None,
+    ) -> RemoteBundleRef:
+        """Store a mapping object or a single tensor/array value."""
+        if isinstance(obj, Mapping):
+            payload = StructuredObjectPayload(metadata={}, buffers=obj)
+        else:
+            payload = StructuredObjectPayload(
+                metadata={"__mooncake_wrapped_object__": True},
+                buffers={"value": obj},
+            )
+        return self.put_structured_object(
+            payload,
+            partition=partition,
+            chunk_bytes=chunk_bytes,
+            policy=policy,
+            max_inflight_put=max_inflight_put,
+        )
+
+    def get_object(self, ref: RemoteBundleRef | Mapping[str, Any]) -> Any:
+        """Materialize an object stored by put_object()."""
+        result = self.materialize(self.read_spec(ref))
+        if result.metadata.get("__mooncake_wrapped_object__"):
+            return result.objects["value"]
+        return result.objects
+
+    def read_spec(
+        self, ref: RemoteBundleRef | Mapping[str, Any]
+    ) -> StructuredObjectReadSpec:
+        """Create a lazy read spec for a structured object reference."""
+        return self._structured_store.read_spec(ref)
+
+    def materialize(self, spec: StructuredObjectReadSpec) -> StructuredObjectResult:
+        """Materialize a structured object read spec."""
+        return self._structured_store.materialize(spec)
+
+    def materialize_into(
+        self,
+        spec: StructuredObjectReadSpec,
+        destinations: Optional[Mapping[str, Any]],
+    ) -> StructuredObjectResult:
+        """Materialize a structured object read spec into caller-provided destinations when possible."""
+        return self._structured_store.materialize_into(spec, destinations)
+
+    def put_dataproto(
+        self,
+        data: Any,
+        *,
+        namespace: str = "default",
+        partition: str = "default",
+        stage: str = "default",
+        chunk_bytes: Optional[int] = None,
+        policy: Optional[BundleTransferPolicy] = None,
+    ) -> MooncakeDataProtoRef:
+        """Store a DataProto-like object as a stage-level structured object."""
+        return self._put_dataproto_stage(
+            None,
+            data,
+            namespace=namespace,
+            partition=partition,
+            stage=stage,
+            chunk_bytes=chunk_bytes,
+            policy=policy,
+            overwrite=False,
+        )
+
+    def append_dataproto_fields(
+        self,
+        ref: DataProtoRefLike,
+        data: Any,
+        *,
+        stage: str,
+        overwrite: bool = False,
+        chunk_bytes: Optional[int] = None,
+        policy: Optional[BundleTransferPolicy] = None,
+    ) -> MooncakeDataProtoRef:
+        """Append fields from a later DataProto stage without rewriting previous stages."""
+        ref = _resolve_dataproto_ref(ref)
+        return self._put_dataproto_stage(
+            ref,
+            data,
+            namespace=ref.namespace,
+            partition=ref.partition,
+            stage=stage,
+            chunk_bytes=chunk_bytes,
+            policy=policy,
+            overwrite=overwrite,
+        )
+
+    def dataproto_manifest_view(self, ref: DataProtoRefLike) -> dict[str, Any]:
+        """Return a DataProto field view derived from stage structured-object manifests."""
+        return _dataproto_manifest_view(self, _resolve_dataproto_ref(ref))
+
+    def get_dataproto(
+        self,
+        ref: DataProtoRefLike,
+        *,
+        fields: Optional[Sequence[str]] = None,
+        batch_fields: Optional[Sequence[str]] = None,
+        non_tensor_fields: Optional[Sequence[str]] = None,
+        meta_info_keys: Optional[Sequence[str]] = None,
+        data_cls: Optional[Any] = None,
+        destinations: Optional[Mapping[str, Any]] = None,
+        rows: slice | StructuredMemberSlice | Sequence[int] | None = None,
+    ) -> Any:
+        """Materialize selected DataProto fields from structured object refs."""
+        ref = _resolve_dataproto_ref(ref)
+        row_selection = _coerce_dataproto_row_selection(rows, ref.batch_size)
+        row_slice = None if row_selection is None else row_selection.member_slice
+        row_indices = None if row_selection is None else row_selection.indices
+        output_rows = ref.batch_size if row_selection is None else row_selection.count
+        batch_names, non_tensor_names = _resolve_dataproto_field_selection(
+            ref, fields, batch_fields, non_tensor_fields
+        )
+        batch: dict[str, Any] = {}
+        non_tensor_batch: dict[str, Any] = {}
+        destination_map = destinations or {}
+        requested = [
+            *[("batch", name) for name in batch_names],
+            *[("non_tensor_batch", name) for name in non_tensor_names],
+        ]
+        by_stage: dict[str, list[tuple[str, StructuredFieldLocation]]] = {}
+        encoded_requests: list[tuple[str, StructuredFieldLocation]] = []
+        for _section, name in requested:
+            location = ref.field_index[name]
+            if (
+                location.section == "non_tensor_batch"
+                and name in ref.encoded_non_tensor
+            ):
+                encoded_requests.append((name, location))
+                continue
+            by_stage.setdefault(location.stage, []).append((name, location))
+        for stage, entries in by_stage.items():
+            stage_ref = ref.stage_refs[stage]
+            members = [location.member for _name, location in entries]
+            stage_destinations = {
+                location.member: destination_map[name]
+                for name, location in entries
+                if name in destination_map
+            }
+            if row_indices is None:
+                spec = self.read_spec(stage_ref).select_members(members)
+                if row_slice is not None:
+                    for member in members:
+                        spec = spec.slice_member(
+                            member,
+                            axis=row_slice.axis,
+                            start=row_slice.start,
+                            end=row_slice.end,
+                            step=row_slice.step,
+                        )
+                result = self.materialize_into(spec, stage_destinations)
+                for name, location in entries:
+                    value = result.objects[location.member]
+                    if location.section == "batch":
+                        batch[name] = value
+                    else:
+                        non_tensor_batch[name] = value
+                continue
+            for name, location in entries:
+                value = self._read_dataproto_member_indices(
+                    stage_ref,
+                    location.member,
+                    row_indices,
+                    destination_map.get(name),
+                )
+                if location.section == "batch":
+                    batch[name] = value
+                else:
+                    non_tensor_batch[name] = value
+        for name, location in encoded_requests:
+            stage_ref = ref.stage_refs[location.stage]
+            encoded = ref.encoded_non_tensor[name]
+            if row_selection is None:
+                members = list(encoded["payload_members"].values())
+                result = self.materialize(
+                    self.read_spec(stage_ref).select_members(members)
+                )
+                payload = {
+                    payload_name: result.objects[member]
+                    for payload_name, member in encoded["payload_members"].items()
+                }
+                values = _decode_structured_non_tensor_encoded(
+                    encoded, payload, ref.batch_size, encoded.get("metadata")
+                )
+            elif row_indices is not None:
+                payload, metadata = self._read_structured_non_tensor_payload_indices(
+                    stage_ref,
+                    encoded,
+                    row_indices,
+                )
+                values = _decode_structured_non_tensor_encoded(
+                    encoded, payload, output_rows, metadata
+                )
+            else:
+                payload, metadata = self._read_structured_non_tensor_payload_slice(
+                    stage_ref,
+                    encoded,
+                    row_slice,
+                    ref.batch_size,
+                )
+                values = _decode_structured_non_tensor_encoded(
+                    encoded, payload, output_rows, metadata
+                )
+            array = np.empty(output_rows, dtype=object)
+            array[:] = values
+            non_tensor_batch[name] = array
+        meta_info = _select_mapping(ref.meta_info, meta_info_keys)
+        return _build_dataproto_like_result(
+            batch, non_tensor_batch, meta_info, data_cls
+        )
+
+    def _read_dataproto_member_indices(
+        self,
+        stage_ref: RemoteBundleRef,
+        member: str,
+        indices: Sequence[int],
+        destination: Any,
+    ) -> Any:
+        manifest = self._bundle_store.resolve_manifest(stage_ref)
+        metadata = _decode_structured_metadata(
+            self._bundle_store.read_payload(manifest["meta"])
+        )
+        field_spec = _structured_field_specs(metadata).get(member, {"encoding": "bytes"})
+        payload_spec = manifest["buffers"][member]
+        encoding = field_spec.get("encoding", "bytes")
+        if encoding == "ndarray":
+            return self._read_ndarray_member_indices(
+                member, payload_spec, field_spec, indices, destination
+            )
+        if encoding == "torch_tensor":
+            return self._read_torch_tensor_member_indices(
+                member, payload_spec, field_spec, indices, destination
+            )
+        raise ValueError(f"structured member {member} does not support indexed rows")
+
+    def _read_ndarray_member_indices(
+        self,
+        name: str,
+        payload_spec: Mapping[str, Any],
+        field_spec: Mapping[str, Any],
+        indices: Sequence[int],
+        destination: Any,
+    ) -> np.ndarray:
+        dtype = field_spec.get("dtype")
+        shape = field_spec.get("shape")
+        if not isinstance(dtype, str) or not isinstance(shape, list):
+            raise ValueError(
+                f"structured ndarray field {name} is missing dtype or shape"
+            )
+        dtype_obj = np.dtype(dtype)
+        full_shape = tuple(int(dim) for dim in shape)
+        if not full_shape:
+            raise ValueError("structured ndarray indexed read requires at least one dimension")
+        output_shape = (len(indices), *full_shape[1:])
+        if isinstance(destination, _RawDestinationBuffer):
+            nbytes = int(np.prod(output_shape, dtype=np.int64)) * dtype_obj.itemsize
+            if destination.size < nbytes:
+                raise ValueError(
+                    f"raw destination has {destination.size} bytes, expected at least {nbytes}"
+                )
+            target = np.ctypeslib.as_array(
+                (ctypes.c_uint8 * nbytes).from_address(destination.ptr)
+            ).view(dtype_obj).reshape(output_shape)
+            if not indices:
+                return target
+            row_width = dtype_obj.itemsize * int(np.prod(full_shape[1:], dtype=np.int64))
+            ranges = [
+                (row * row_width, out * row_width, row_width)
+                for out, row in enumerate(indices)
+            ]
+            if destination.pre_registered:
+                self._bundle_store.read_payload_ranges_into_raw_destination(
+                    payload_spec, destination.ptr, ranges
+                )
+            else:
+                data = bytearray(nbytes)
+                self._bundle_store.read_payload_ranges_into_bytearray(
+                    payload_spec, data, ranges
+                )
+                ctypes.memmove(destination.ptr, bytes(data), nbytes)
+            _ = destination.owner
+            return target
+        target = _resolve_ndarray_destination(name, destination, dtype_obj, output_shape)
+        if not indices:
+            return target
+        row_width = dtype_obj.itemsize * int(np.prod(full_shape[1:], dtype=np.int64))
+        self._bundle_store.read_payload_ranges_into_array(
+            payload_spec,
+            target.view(np.uint8).reshape(-1),
+            [
+                (row * row_width, out * row_width, row_width)
+                for out, row in enumerate(indices)
+            ],
+        )
+        return target
+
+    def _read_torch_tensor_member_indices(
+        self,
+        name: str,
+        payload_spec: Mapping[str, Any],
+        field_spec: Mapping[str, Any],
+        indices: Sequence[int],
+        destination: Any,
+    ) -> Any:
+        metadata_bytes = int(payload_spec.get("metadata_bytes", -1))
+        shape = field_spec.get("shape")
+        element_size = int(field_spec.get("element_size", 0))
+        if metadata_bytes < 0 or not isinstance(shape, list) or element_size <= 0:
+            raise ValueError(
+                f"structured tensor field {name} is missing slice metadata"
+            )
+        if not shape:
+            raise ValueError("structured tensor indexed read requires at least one dimension")
+        row_width = element_size * int(np.prod(shape[1:], dtype=np.int64))
+        data_length = len(indices) * row_width
+        metadata = self._bundle_store.read_payload_range(payload_spec, 0, metadata_bytes)
+        sliced_metadata = _slice_tensor_metadata(
+            metadata, (len(indices), *shape[1:]), data_length
+        )
+        if destination is not None:
+            if not isinstance(destination, (_TensorObjectBufferPayload, _RawDestinationBuffer)):
+                raise ValueError(
+                    f"structured tensor member {name} only supports tensor_object_buffer or raw_destination destinations"
+                )
+            expected_bytes = metadata_bytes + data_length
+            if destination.size < expected_bytes:
+                raise ValueError(
+                    f"tensor destination has {destination.size} bytes, expected at least {expected_bytes}"
+                )
+            ctypes.memmove(destination.ptr, sliced_metadata, metadata_bytes)
+            if data_length:
+                ranges = [
+                    (
+                        metadata_bytes + row * row_width,
+                        metadata_bytes + out * row_width,
+                        row_width,
+                    )
+                    for out, row in enumerate(indices)
+                ]
+                if isinstance(destination, _RawDestinationBuffer) and destination.pre_registered:
+                    self._bundle_store.read_payload_ranges_into_raw_destination(
+                        payload_spec, destination.ptr, ranges
+                    )
+                else:
+                    data = bytearray(data_length)
+                    self._bundle_store.read_payload_ranges_into_bytearray(
+                        payload_spec,
+                        data,
+                        [
+                            (source_offset, destination_offset - metadata_bytes, size)
+                            for source_offset, destination_offset, size in ranges
+                        ],
+                    )
+                    ctypes.memmove(
+                        destination.ptr + metadata_bytes, bytes(data), data_length
+                    )
+            _ = destination.owner
+            return destination
+        data = bytearray(data_length)
+        if data_length:
+            self._bundle_store.read_payload_ranges_into_bytearray(
+                payload_spec,
+                data,
+                [(metadata_bytes + row * row_width, out * row_width, row_width) for out, row in enumerate(indices)],
+            )
+        return _deserialize_tensor_payload(sliced_metadata + data)
+
+    def _read_structured_non_tensor_payload_indices(
+        self,
+        stage_ref: RemoteBundleRef,
+        encoded: Mapping[str, Any],
+        indices: Sequence[int],
+    ) -> tuple[dict[str, Any], Mapping[str, Any]]:
+        codec = encoded["codec"]
+        payload_members = encoded["payload_members"]
+        manifest = self._bundle_store.resolve_manifest(stage_ref)
+        stage_metadata = _decode_structured_metadata(
+            self._bundle_store.read_payload(manifest["meta"])
+        )
+        field_specs = _structured_field_specs(stage_metadata)
+        metadata = dict(encoded.get("metadata") or {})
+
+        def member(payload_name: str) -> str:
+            return payload_members[payload_name]
+
+        def member_dtype(payload_name: str) -> np.dtype[Any]:
+            spec = field_specs.get(member(payload_name), {})
+            dtype = spec.get("dtype")
+            if not isinstance(dtype, str):
+                raise ValueError(
+                    f"structured non-tensor payload {member(payload_name)} is missing dtype"
+                )
+            return np.dtype(dtype)
+
+        def read_member_indices(payload_name: str, member_indices: Sequence[int]) -> Any:
+            return self._read_dataproto_member_indices(
+                stage_ref, member(payload_name), member_indices, None
+            )
+
+        def read_data_ranges(
+            payload_name: str,
+            ranges: Sequence[tuple[int, int, int]],
+            dtype: np.dtype[Any] | None = None,
+        ) -> Any:
+            total_bytes = sum(byte_length for _src, _dst, byte_length in ranges)
+            payload_spec = manifest["buffers"][member(payload_name)]
+            if dtype is None:
+                data = bytearray(total_bytes)
+                self._bundle_store.read_payload_ranges_into_bytearray(
+                    payload_spec, data, ranges
+                )
+                return bytes(data)
+            array = np.empty(total_bytes // dtype.itemsize, dtype=dtype)
+            self._bundle_store.read_payload_ranges_into_array(
+                payload_spec, array.view(np.uint8).reshape(-1), ranges
+            )
+            return array
+
+        if _is_recursive_encoded_non_tensor(encoded):
+            metadata = _copy_recursive_metadata_for_leaf_updates(metadata)
+            recursive_payload: dict[str, Any] = {}
+            for node in metadata.get("nodes", []):
+                for key in ("missing_payload", "row_mask_payload", "lengths_payload"):
+                    payload_name = node.get(key)
+                    if payload_name is not None:
+                        recursive_payload[payload_name] = read_member_indices(
+                            payload_name, indices
+                        )
+            for leaf in metadata.get("leaves", []):
+                leaf_payload_members = leaf["payload_members"]
+                flat_payload_members = {
+                    name: payload_members[global_name]
+                    for name, global_name in leaf_payload_members.items()
+                    if name != "missing"
+                }
+                leaf_payload, _leaf_metadata = self._read_structured_non_tensor_payload_indices(
+                    stage_ref,
+                    {
+                        "codec": leaf["codec"],
+                        "metadata": leaf.get("metadata") or {},
+                        "payload_members": flat_payload_members,
+                    },
+                    indices,
+                )
+                leaf["metadata"] = _leaf_metadata
+                for name, value in leaf_payload.items():
+                    recursive_payload[leaf_payload_members[name]] = value
+                recursive_payload[leaf_payload_members["missing"]] = read_member_indices(
+                    leaf_payload_members["missing"], indices
+                )
+            return recursive_payload, metadata
+
+        if codec == "ndarray":
+            return {
+                "data": read_member_indices("data", indices),
+                "nulls": read_member_indices("nulls", indices),
+            }, metadata
+
+        if codec in {"ragged_tensor", "typed_ragged"}:
+            offsets = read_member_indices(
+                "offsets", [index for row in indices for index in (row, row + 1)]
+            )
+            begins = offsets[0::2]
+            ends = offsets[1::2]
+            item_counts = [int(end) - int(begin) for begin, end in zip(begins, ends)]
+            gathered_offsets = np.empty(len(indices) + 1, dtype=offsets.dtype)
+            gathered_offsets[0] = 0
+            for index, count in enumerate(item_counts):
+                gathered_offsets[index + 1] = int(gathered_offsets[index]) + count
+            dtype = member_dtype("data")
+            ranges = []
+            destination_item = 0
+            for begin, count in zip(begins, item_counts):
+                if count:
+                    ranges.append(
+                        (
+                            int(begin) * dtype.itemsize,
+                            destination_item * dtype.itemsize,
+                            count * dtype.itemsize,
+                        )
+                    )
+                destination_item += count
+            return {
+                "data": read_data_ranges("data", ranges, dtype),
+                "offsets": gathered_offsets,
+                "shapes": read_member_indices("shapes", indices),
+                "ndims": read_member_indices("ndims", indices),
+                "nulls": read_member_indices("nulls", indices),
+            }, metadata
+
+        if codec in {"media_bytes", "bytes_ragged", "utf8_ragged", "msgpack_ragged", "json_ragged"}:
+            offsets = read_member_indices(
+                "offsets", [index for row in indices for index in (row, row + 1)]
+            )
+            begins = offsets[0::2]
+            ends = offsets[1::2]
+            byte_counts = [int(end) - int(begin) for begin, end in zip(begins, ends)]
+            gathered_offsets = np.empty(len(indices) + 1, dtype=offsets.dtype)
+            gathered_offsets[0] = 0
+            ranges = []
+            destination_offset = 0
+            for index, (begin, count) in enumerate(zip(begins, byte_counts)):
+                gathered_offsets[index + 1] = destination_offset + count
+                if count:
+                    ranges.append((int(begin), destination_offset, count))
+                destination_offset += count
+            if "media_encodings" in metadata:
+                metadata["media_encodings"] = [
+                    metadata["media_encodings"][index] for index in indices
+                ]
+            return {
+                "data": read_data_ranges("data", ranges),
+                "offsets": gathered_offsets,
+                "nulls": read_member_indices("nulls", indices),
+            }, metadata
+
+        if codec == "media_list_ragged":
+            row_offsets = read_member_indices(
+                "row_offsets", [index for row in indices for index in (row, row + 1)]
+            )
+            item_begins = row_offsets[0::2]
+            item_ends = row_offsets[1::2]
+            item_counts = [
+                int(end) - int(begin) for begin, end in zip(item_begins, item_ends)
+            ]
+            gathered_row_offsets = np.empty(len(indices) + 1, dtype=row_offsets.dtype)
+            gathered_row_offsets[0] = 0
+            for index, count in enumerate(item_counts):
+                gathered_row_offsets[index + 1] = int(gathered_row_offsets[index]) + count
+            boundary_indices = [
+                boundary
+                for begin, end in zip(item_begins, item_ends)
+                for boundary in range(int(begin), int(end) + 1)
+            ]
+            byte_offsets = read_member_indices("byte_offsets", boundary_indices)
+            ranges = []
+            gathered_byte_offsets = np.empty(
+                int(gathered_row_offsets[-1]) + 1, dtype=byte_offsets.dtype
+            )
+            gathered_byte_offsets[0] = 0
+            destination_offset = 0
+            source_index = 0
+            destination_boundary = 1
+            media_encodings = []
+            for begin, count in zip(item_begins, item_counts):
+                for item in range(count):
+                    byte_begin = int(byte_offsets[source_index + item])
+                    byte_end = int(byte_offsets[source_index + item + 1])
+                    byte_count = byte_end - byte_begin
+                    if byte_count:
+                        ranges.append((byte_begin, destination_offset, byte_count))
+                    destination_offset += byte_count
+                    gathered_byte_offsets[destination_boundary] = destination_offset
+                    destination_boundary += 1
+                if "media_encodings" in metadata:
+                    media_encodings.extend(
+                        metadata["media_encodings"][int(begin) : int(begin) + count]
+                    )
+                source_index += count + 1
+            if "media_encodings" in metadata:
+                metadata["media_encodings"] = media_encodings
+            return {
+                "data": read_data_ranges("data", ranges),
+                "row_offsets": gathered_row_offsets,
+                "byte_offsets": gathered_byte_offsets,
+                "nulls": read_member_indices("nulls", indices),
+            }, metadata
+
+        raise ValueError(f"unknown structured non-tensor codec: {codec}")
+
+    def _read_structured_non_tensor_payload_slice(
+        self,
+        stage_ref: RemoteBundleRef,
+        encoded: Mapping[str, Any],
+        row_slice: StructuredMemberSlice,
+        total_rows: int,
+    ) -> tuple[dict[str, Any], Mapping[str, Any]]:
+        start, end, step = _normalized_member_slice(row_slice, total_rows)
+        if step != 1:
+            raise ValueError(
+                "DataProto structured non-tensor slicing currently requires step=1"
+            )
+        codec = encoded["codec"]
+        payload_members = encoded["payload_members"]
+        manifest = self._bundle_store.resolve_manifest(stage_ref)
+        metadata = dict(encoded.get("metadata") or {})
+
+        def read_member(payload_name: str, row_start: int, row_end: int) -> Any:
+            member = payload_members[payload_name]
+            spec = (
+                self.read_spec(stage_ref)
+                .select_members([member])
+                .slice_member(member, axis=0, start=row_start, end=row_end)
+            )
+            return self.materialize(spec).objects[member]
+
+        def read_bytes(payload_name: str, byte_start: int, byte_end: int) -> bytes:
+            member = payload_members[payload_name]
+            payload_spec = manifest["buffers"][member]
+            return self._bundle_store.read_payload_range(
+                payload_spec, byte_start, byte_end - byte_start
+            )
+
+        if _is_recursive_encoded_non_tensor(encoded):
+            metadata = _copy_recursive_metadata_for_leaf_updates(metadata)
+            recursive_payload: dict[str, Any] = {}
+            for node in metadata.get("nodes", []):
+                for key in ("missing_payload", "row_mask_payload", "lengths_payload"):
+                    payload_name = node.get(key)
+                    if payload_name is not None:
+                        recursive_payload[payload_name] = read_member(payload_name, start, end)
+            for leaf in metadata.get("leaves", []):
+                leaf_payload_members = leaf["payload_members"]
+                flat_payload_members = {
+                    name: payload_members[global_name]
+                    for name, global_name in leaf_payload_members.items()
+                    if name != "missing"
+                }
+                leaf_payload, _leaf_metadata = self._read_structured_non_tensor_payload_slice(
+                    stage_ref,
+                    {
+                        "codec": leaf["codec"],
+                        "metadata": leaf.get("metadata") or {},
+                        "payload_members": flat_payload_members,
+                    },
+                    row_slice,
+                    total_rows,
+                )
+                leaf["metadata"] = _leaf_metadata
+                for name, value in leaf_payload.items():
+                    recursive_payload[leaf_payload_members[name]] = value
+                recursive_payload[leaf_payload_members["missing"]] = read_member(
+                    leaf_payload_members["missing"], start, end
+                )
+            return recursive_payload, metadata
+
+        if codec == "ndarray":
+            return {
+                "data": read_member("data", start, end),
+                "nulls": read_member("nulls", start, end),
+            }, metadata
+
+        if codec in {"ragged_tensor", "typed_ragged"}:
+            offsets = read_member("offsets", start, end + 1)
+            base = int(offsets[0])
+            limit = int(offsets[-1])
+            offsets = offsets - base
+            return {
+                "data": read_member("data", base, limit),
+                "offsets": offsets,
+                "shapes": read_member("shapes", start, end),
+                "ndims": read_member("ndims", start, end),
+                "nulls": read_member("nulls", start, end),
+            }, metadata
+
+        if codec in {"media_bytes", "bytes_ragged", "utf8_ragged", "msgpack_ragged", "json_ragged"}:
+            offsets = read_member("offsets", start, end + 1)
+            base = int(offsets[0])
+            limit = int(offsets[-1])
+            offsets = offsets - base
+            if "media_encodings" in metadata:
+                metadata["media_encodings"] = metadata["media_encodings"][start:end]
+            return {
+                "data": read_bytes("data", base, limit),
+                "offsets": offsets,
+                "nulls": read_member("nulls", start, end),
+            }, metadata
+
+        if codec == "media_list_ragged":
+            row_offsets = read_member("row_offsets", start, end + 1)
+            item_start = int(row_offsets[0])
+            item_end = int(row_offsets[-1])
+            row_offsets = row_offsets - item_start
+            byte_offsets = read_member("byte_offsets", item_start, item_end + 1)
+            byte_start = int(byte_offsets[0])
+            byte_end = int(byte_offsets[-1])
+            byte_offsets = byte_offsets - byte_start
+            if "media_encodings" in metadata:
+                metadata["media_encodings"] = metadata["media_encodings"][
+                    item_start:item_end
+                ]
+            return {
+                "data": read_bytes("data", byte_start, byte_end),
+                "row_offsets": row_offsets,
+                "byte_offsets": byte_offsets,
+                "nulls": read_member("nulls", start, end),
+            }, metadata
+
+        raise ValueError(f"unknown structured non-tensor codec: {codec}")
+
+    def cleanup_dataproto(self, ref: DataProtoRefLike) -> None:
+        """Remove all structured object stages referenced by a DataProto handle."""
+        ref = _resolve_dataproto_ref(ref)
+        seen: set[str] = set()
+        for stage_ref in ref.stage_refs.values():
+            if stage_ref.manifest_key in seen:
+                continue
+            seen.add(stage_ref.manifest_key)
+            self.remove_bundle(stage_ref)
+
+    def _append_dataproto_stage_manifest(
+        self,
+        old_stage_ref: RemoteBundleRef,
+        payload: StructuredObjectPayload,
+        *,
+        partition: str,
+        chunk_bytes: Optional[int],
+        policy: Optional[BundleTransferPolicy],
+    ) -> RemoteBundleRef:
+        new_stage_ref = self.put_structured_object(
+            payload,
+            partition=partition,
+            chunk_bytes=chunk_bytes,
+            policy=policy,
+        )
+        try:
+            old_manifest = self._bundle_store.resolve_manifest(old_stage_ref)
+            new_manifest = self._bundle_store.resolve_manifest(new_stage_ref)
+            old_metadata = _decode_structured_metadata(
+                self._bundle_store.read_payload(old_manifest["meta"])
+            )
+            new_metadata = _decode_structured_metadata(
+                self._bundle_store.read_payload(new_manifest["meta"])
+            )
+            merged_metadata = _merge_structured_stage_metadata(
+                old_metadata, new_metadata
+            )
+            merged_buffers = dict(old_manifest["buffers"])
+            collisions = sorted(set(merged_buffers) & set(new_manifest["buffers"]))
+            if collisions:
+                raise ValueError(f"structured members already exist: {collisions}")
+            merged_buffers.update(new_manifest["buffers"])
+            merged_ref = self._bundle_store.put_bundle_manifest(
+                _encode_structured_metadata(merged_metadata),
+                merged_buffers,
+                partition=partition,
+                chunk_bytes=chunk_bytes,
+                policy=policy,
+                cleanup_keys=[
+                    self._bundle_store.manifest_key(old_stage_ref),
+                    *self._bundle_store.payload_keys(old_manifest["meta"]),
+                ],
+            )
+        except Exception:
+            self.remove_bundle(new_stage_ref)
+            raise
+        obsolete_keys = [
+            self._bundle_store.manifest_key(new_stage_ref),
+            *self._bundle_store.payload_keys(new_manifest["meta"]),
+        ]
+        self._bundle_store.remove_keys(obsolete_keys, strict=False)
+        return merged_ref
+
+    def _put_dataproto_stage(
+        self,
+        ref: MooncakeDataProtoRef | None,
+        data: Any,
+        *,
+        namespace: str,
+        partition: str,
+        stage: str,
+        chunk_bytes: Optional[int],
+        policy: Optional[BundleTransferPolicy],
+        overwrite: bool,
+    ) -> MooncakeDataProtoRef:
+        batch, non_tensor_batch, meta_info = _split_dataproto_like(data)
+        batch_size = _dataproto_batch_size(batch, non_tensor_batch)
+        if ref is not None and batch_size != ref.batch_size:
+            raise ValueError(
+                f"DataProto append batch size {batch_size} does not match ref batch size {ref.batch_size}"
+            )
+        cross_section_fields = sorted(set(batch) & set(non_tensor_batch))
+        if cross_section_fields:
+            raise ValueError(
+                f"DataProto batch and non_tensor_batch fields overlap: {cross_section_fields}"
+            )
+        buffers: dict[str, Any] = {}
+        field_updates: dict[str, StructuredFieldLocation] = {}
+        for name, value in batch.items():
+            member = f"batch.{name}"
+            buffers[member] = value
+            field_updates[name] = StructuredFieldLocation(stage, member, "batch")
+        encoded_updates: dict[str, Any] = {}
+        for name, value in non_tensor_batch.items():
+            if _should_encode_non_tensor_field(value):
+                encoded = _encode_structured_non_tensor_field(
+                    f"non_tensor_batch.{name}", value
+                )
+                payload_members: dict[str, str] = {}
+                for payload_name, payload_value in encoded.payload.items():
+                    member = f"non_tensor_batch.{name}.{payload_name}"
+                    buffers[member] = payload_value
+                    payload_members[payload_name] = member
+                field_updates[name] = StructuredFieldLocation(
+                    stage, f"non_tensor_batch.{name}", "non_tensor_batch"
+                )
+                encoded_updates[name] = {
+                    "codec": encoded.codec,
+                    "rows": encoded.rows,
+                    "metadata": encoded.metadata,
+                    "payload_members": payload_members,
+                }
+                continue
+            member = f"non_tensor_batch.{name}"
+            buffers[member] = value
+            field_updates[name] = StructuredFieldLocation(
+                stage, member, "non_tensor_batch"
+            )
+        if not buffers:
+            if ref is not None:
+                merged_meta_info = dict(ref.meta_info)
+                merged_meta_info.update(meta_info)
+                return MooncakeDataProtoRef(
+                    batch_size=ref.batch_size,
+                    stage_refs=dict(ref.stage_refs),
+                    field_index=dict(ref.field_index),
+                    meta_info=merged_meta_info,
+                    namespace=ref.namespace,
+                    partition=ref.partition,
+                    global_indexes=ref.global_indexes,
+                    encoded_non_tensor=dict(ref.encoded_non_tensor),
+                )
+            return MooncakeDataProtoRef(
+                batch_size=batch_size,
+                stage_refs={},
+                field_index={},
+                meta_info=dict(meta_info),
+                namespace=namespace,
+                partition=partition,
+                encoded_non_tensor={},
+            )
+        duplicates = (
+            sorted(set(ref.field_index) & set(field_updates)) if ref is not None else []
+        )
+        if duplicates and not overwrite:
+            raise ValueError(f"DataProto fields already exist: {duplicates}")
+        if ref is not None and overwrite:
+            dangling = [
+                name
+                for name, location in ref.field_index.items()
+                if location.stage == stage and name not in field_updates
+            ]
+            if dangling:
+                raise ValueError(
+                    f"DataProto overwrite for stage {stage!r} must include existing fields: {sorted(dangling)}"
+                )
+        payload = StructuredObjectPayload(
+            metadata={
+                "layout": "dataproto_stage",
+                "dataproto": {
+                    "version": 1,
+                    "namespace": namespace,
+                    "partition": partition,
+                    "stage": stage,
+                    "batch_size": batch_size,
+                },
+            },
+            buffers=buffers,
+        )
+        existing_stage_ref = ref.stage_refs.get(stage) if ref is not None else None
+        if existing_stage_ref is not None and not overwrite:
+            stage_ref = self._append_dataproto_stage_manifest(
+                existing_stage_ref,
+                payload,
+                partition=partition,
+                chunk_bytes=chunk_bytes,
+                policy=policy,
+            )
+        else:
+            stage_ref = self.put_structured_object(
+                payload,
+                partition=partition,
+                chunk_bytes=chunk_bytes,
+                policy=policy,
+            )
+            if (
+                existing_stage_ref is not None
+                and overwrite
+                and existing_stage_ref.manifest_key != stage_ref.manifest_key
+            ):
+                self.remove_bundle(existing_stage_ref)
+        if ref is None:
+            stage_refs = {stage: stage_ref}
+            field_index = dict(field_updates)
+            merged_meta_info = dict(meta_info)
+            encoded_non_tensor = dict(encoded_updates)
+        else:
+            stage_refs = dict(ref.stage_refs)
+            stage_refs[stage] = stage_ref
+            field_index = dict(ref.field_index)
+            field_index.update(field_updates)
+            merged_meta_info = dict(ref.meta_info)
+            merged_meta_info.update(meta_info)
+            encoded_non_tensor = dict(ref.encoded_non_tensor)
+            for name in field_updates:
+                encoded_non_tensor.pop(name, None)
+            encoded_non_tensor.update(encoded_updates)
+        return MooncakeDataProtoRef(
+            batch_size=batch_size,
+            stage_refs=stage_refs,
+            field_index=field_index,
+            meta_info=merged_meta_info,
+            namespace=namespace,
+            partition=partition,
+            global_indexes=None if ref is None else ref.global_indexes,
+            encoded_non_tensor=encoded_non_tensor,
+        )
+
+
+def _copy_transfer_policy(policy: BundleTransferPolicy) -> BundleTransferPolicy:
+    if policy.copy_mode == "copy":
+        return policy
+    return BundleTransferPolicy(
+        max_inflight_put=policy.max_inflight_put,
+        put_mode=policy.put_mode,
+        copy_mode="copy",
+    )
+
+
+def _dataproto_manifest_view(
+    transfer: MooncakeBundleTransfer, ref: MooncakeDataProtoRef
+) -> dict[str, Any]:
+    stage_manifests = {
+        stage: transfer._bundle_store.resolve_manifest(stage_ref)
+        for stage, stage_ref in ref.stage_refs.items()
+    }
+    stage_metadata = {
+        stage: _decode_structured_metadata(
+            transfer._bundle_store.read_payload(stage_manifest["meta"])
+        )
+        for stage, stage_manifest in stage_manifests.items()
+    }
+    fields: dict[str, dict[str, Any]] = {}
+    for name, location in ref.field_index.items():
+        metadata = stage_metadata[location.stage]
+        field_specs = _structured_field_specs(metadata)
+        member_spec = dict(field_specs.get(location.member, {"encoding": "bytes"}))
+        encoded = ref.encoded_non_tensor.get(name)
+        if encoded is not None:
+            member_spec = {
+                "encoding": "structured_non_tensor",
+                "codec": encoded["codec"],
+                "metadata": encoded.get("metadata"),
+                "payload_members": dict(encoded["payload_members"]),
+                "payload_specs": {
+                    payload_name: dict(field_specs.get(member, {"encoding": "bytes"}))
+                    for payload_name, member in encoded["payload_members"].items()
+                },
+            }
+        fields[name] = {
+            "section": location.section,
+            "stage": location.stage,
+            "member": location.member,
+            "spec": member_spec,
+        }
+    return {
+        "namespace": ref.namespace,
+        "partition": ref.partition,
+        "batch_size": ref.batch_size,
+        "batch_fields": {
+            name: info for name, info in fields.items() if info["section"] == "batch"
+        },
+        "non_tensor_fields": {
+            name: info
+            for name, info in fields.items()
+            if info["section"] == "non_tensor_batch"
+        },
+        "meta_info_keys": list(ref.meta_info),
+        "stages": {
+            stage: {
+                "manifest_key": stage_ref.manifest_key,
+                "dataproto": stage_metadata[stage].get("dataproto", {}),
+            }
+            for stage, stage_ref in ref.stage_refs.items()
+        },
+    }
+
+
+def _build_dataproto_like_result(
+    batch: dict[str, Any],
+    non_tensor_batch: dict[str, Any],
+    meta_info: dict[str, Any],
+    data_cls: Optional[Any],
+) -> Any:
+    payload = {
+        "batch": batch,
+        "non_tensor_batch": non_tensor_batch,
+        "meta_info": meta_info,
+    }
+    if data_cls is None or data_cls is dict:
+        return payload
+    from_dict = getattr(data_cls, "from_dict", None)
+    try:
+        if callable(from_dict):
+            return from_dict(batch, non_tensor_batch, meta_info=meta_info)
+        return data_cls(
+            batch=batch, non_tensor_batch=non_tensor_batch, meta_info=meta_info
+        )
+    except TypeError as error:
+        raise TypeError(
+            f"{getattr(data_cls, '__name__', data_cls)!r} cannot be constructed from "
+            "batch, non_tensor_batch, and meta_info"
+        ) from error
+
+
+def _split_dataproto_like(
+    data: Any,
+) -> tuple[dict[str, Any], dict[str, Any], dict[str, Any]]:
+    if isinstance(data, Mapping):
+        if _is_dataproto_envelope_mapping(data):
+            return (
+                _mapping_to_dict(data.get("batch")),
+                _mapping_to_dict(data.get("non_tensor_batch")),
+                _mapping_to_dict(data.get("meta_info")),
+            )
+        return dict(data), {}, {}
+    batch = _mapping_to_dict(getattr(data, "batch", None))
+    non_tensor_batch = _mapping_to_dict(getattr(data, "non_tensor_batch", None))
+    meta_info = _mapping_to_dict(getattr(data, "meta_info", None))
+    return batch, non_tensor_batch, meta_info
+
+
+def _is_dataproto_envelope_mapping(data: Mapping[str, Any]) -> bool:
+    envelope_keys = {"batch", "non_tensor_batch", "meta_info"}
+    if not data or not set(data).issubset(envelope_keys):
+        return False
+    return all(_is_mapping_like_or_none(value) for value in data.values())
+
+
+def _is_mapping_like_or_none(value: Any) -> bool:
+    return (
+        value is None
+        or isinstance(value, Mapping)
+        or callable(getattr(value, "items", None))
+    )
+
+
+def _mapping_to_dict(value: Any) -> dict[str, Any]:
+    if value is None:
+        return {}
+    if isinstance(value, Mapping):
+        return dict(value)
+    items = getattr(value, "items", None)
+    if callable(items):
+        return dict(items())
+    raise TypeError(f"expected mapping-like value, got {type(value).__name__}")
+
+
+def _dataproto_batch_size(
+    batch: Mapping[str, Any], non_tensor_batch: Mapping[str, Any]
+) -> int:
+    sizes = {
+        len(value) for values in (batch, non_tensor_batch) for value in values.values()
+    }
+    if not sizes:
+        return 0
+    if len(sizes) != 1:
+        raise ValueError(
+            f"DataProto fields have inconsistent batch sizes: {sorted(sizes)}"
+        )
+    return sizes.pop()
+
+
+def _resolve_dataproto_field_selection(
+    ref: MooncakeDataProtoRef,
+    fields: Optional[Sequence[str]],
+    batch_fields: Optional[Sequence[str]],
+    non_tensor_fields: Optional[Sequence[str]],
+) -> tuple[list[str], list[str]]:
+    if fields is not None and (
+        batch_fields is not None or non_tensor_fields is not None
+    ):
+        raise ValueError(
+            "fields cannot be combined with batch_fields or non_tensor_fields"
+        )
+    if fields is not None:
+        requested = list(fields)
+        _validate_dataproto_fields_exist(ref, requested)
+        return (
+            [name for name in requested if ref.field_index[name].section == "batch"],
+            [
+                name
+                for name in requested
+                if ref.field_index[name].section == "non_tensor_batch"
+            ],
+        )
+    if batch_fields is None and non_tensor_fields is None:
+        batch_names = [
+            name
+            for name, location in ref.field_index.items()
+            if location.section == "batch"
+        ]
+        non_tensor_names = [
+            name
+            for name, location in ref.field_index.items()
+            if location.section == "non_tensor_batch"
+        ]
+    else:
+        batch_names = [] if batch_fields is None else list(batch_fields)
+        non_tensor_names = (
+            [] if non_tensor_fields is None else list(non_tensor_fields)
+        )
+    _validate_dataproto_fields_exist(ref, [*batch_names, *non_tensor_names])
+    return batch_names, non_tensor_names
+
+
+def _validate_dataproto_fields_exist(
+    ref: MooncakeDataProtoRef, names: Sequence[str]
+) -> None:
+    missing = [name for name in names if name not in ref.field_index]
+    if missing:
+        raise KeyError(f"unknown DataProto fields: {missing}")
+
+
+def _coerce_dataproto_row_selection(
+    rows: slice | StructuredMemberSlice | Sequence[int] | None, total_rows: int
+) -> _DataProtoRowSelection | None:
+    if rows is None:
+        return None
+    if isinstance(rows, StructuredMemberSlice):
+        if rows.axis != 0:
+            raise ValueError("DataProto row slicing currently supports axis=0 only")
+        _normalized_member_slice(rows, total_rows)
+        return _DataProtoRowSelection(
+            count=_slice_length(rows, total_rows), member_slice=rows
+        )
+    if isinstance(rows, slice):
+        start, end, step = rows.indices(total_rows)
+        if step <= 0:
+            raise ValueError("DataProto row slicing step must be positive")
+        member_slice = StructuredMemberSlice(axis=0, start=start, end=end, step=step)
+        return _DataProtoRowSelection(
+            count=_slice_length(member_slice, total_rows), member_slice=member_slice
+        )
+    if isinstance(rows, Sequence) and not isinstance(rows, (str, bytes, bytearray)):
+        indices = tuple(_normalize_dataproto_row_index(index, total_rows) for index in rows)
+        return _DataProtoRowSelection(count=len(indices), indices=indices)
+    raise TypeError("DataProto rows must be a slice, StructuredMemberSlice, or row index sequence")
+
+
+def _normalize_dataproto_row_index(index: Any, total_rows: int) -> int:
+    if not isinstance(index, (int, np.integer)):
+        raise TypeError("DataProto row indices must be integers")
+    normalized = int(index)
+    if normalized < 0:
+        normalized += total_rows
+    if normalized < 0 or normalized >= total_rows:
+        raise IndexError(f"DataProto row index {index} out of range for {total_rows} rows")
+    return normalized
+
+
+def _slice_length(member_slice: StructuredMemberSlice, total_rows: int) -> int:
+    start, end, step = _normalized_member_slice(member_slice, total_rows)
+    if start >= end:
+        return 0
+    return 1 + (end - 1 - start) // step
+
+
+def _select_mapping(
+    value: Mapping[str, Any], keys: Optional[Sequence[str]]
+) -> dict[str, Any]:
+    if keys is None:
+        return dict(value)
+    return {key: value[key] for key in keys if key in value}
+
+
+def _should_encode_non_tensor_field(value: Any) -> bool:
+    return isinstance(value, np.ndarray) and value.dtype == object
+
+
+def _encode_structured_non_tensor_field(
+    path: str, value: np.ndarray
+) -> _EncodedStructuredLeaf:
+    values = list(value)
+    leaves: list[_InferredLeaf] = []
+    nodes: list[_InferredNode] = []
+    infer_structure(path, values, leaves, nodes)
+    if nodes and _should_encode_recursive_structure(leaves):
+        return _encode_recursive_structured_non_tensor_field(path, values, leaves, nodes)
+    decision = _choose_leaf_codec(values)
+    return _encode_structured_leaf(values, decision)
+
+
+def _should_encode_recursive_structure(leaves: Sequence[_InferredLeaf]) -> bool:
+    return any(_should_encode_recursive_leaf(leaf) for leaf in leaves)
+
+
+def _should_encode_recursive_leaf(leaf: _InferredLeaf) -> bool:
+    codec = leaf.decision.codec
+    if codec in {"ragged_tensor", "bytes_ragged", "media_bytes", "media_list_ragged"}:
+        return True
+    if codec == "typed_ragged":
+        return any(isinstance(value, np.ndarray) for value in _non_null(leaf.values))
+    return False
+
+
+def _recursive_leaf_decision(values: list[Any], decision: _CodecDecision) -> _CodecDecision:
+    if decision.accepted:
+        return decision
+    if all(value is None or isinstance(value, _Missing) for value in values):
+        return _CodecDecision(
+            True,
+            "json_ragged",
+            "all rows are null or missing",
+            "json",
+        )
+    return decision
+
+
+def _encode_recursive_structured_non_tensor_field(
+    path: str,
+    values: list[Any],
+    leaves: Sequence[_InferredLeaf],
+    nodes: Sequence[_InferredNode],
+) -> _EncodedStructuredLeaf:
+    payload: dict[str, Any] = {}
+    node_specs: list[dict[str, Any]] = []
+    for node_id, node in enumerate(nodes):
+        spec: dict[str, Any] = {
+            "id": node_id,
+            "path": node.path,
+            "node_type": node.node_type,
+            "children": list(node.children),
+        }
+        missing_payload_name = f"node.{node_id}.missing"
+        payload[missing_payload_name] = np.asarray(
+            [_lookup_structured_path(value, path, node.path) is MISSING for value in values],
+            dtype=np.bool_,
+        )
+        spec["missing_payload"] = missing_payload_name
+        if node.row_mask is not None:
+            payload_name = f"node.{node_id}.row_mask"
+            payload[payload_name] = np.asarray(node.row_mask, dtype=np.bool_)
+            spec["row_mask_payload"] = payload_name
+        if node.lengths is not None:
+            payload_name = f"node.{node_id}.lengths"
+            payload[payload_name] = np.asarray(node.lengths, dtype=np.int64)
+            spec["lengths_payload"] = payload_name
+        node_specs.append(spec)
+
+    leaf_specs: list[dict[str, Any]] = []
+    for leaf_id, leaf in enumerate(leaves):
+        missing = np.asarray(
+            [isinstance(value, _Missing) for value in leaf.values], dtype=np.bool_
+        )
+        codec_values = [None if is_missing else value for is_missing, value in zip(missing, leaf.values)]
+        decision = _recursive_leaf_decision(leaf.values, leaf.decision)
+        encoded = _encode_structured_leaf(codec_values, decision)
+        leaf_payload_members: dict[str, str] = {}
+        for payload_name, payload_value in encoded.payload.items():
+            recursive_payload_name = f"leaf.{leaf_id}.{payload_name}"
+            payload[recursive_payload_name] = payload_value
+            leaf_payload_members[payload_name] = recursive_payload_name
+        missing_payload_name = f"leaf.{leaf_id}.missing"
+        payload[missing_payload_name] = missing
+        leaf_payload_members["missing"] = missing_payload_name
+        leaf_specs.append(
+            {
+                "id": leaf_id,
+                "path": leaf.path,
+                "codec": encoded.codec,
+                "rows": encoded.rows,
+                "metadata": encoded.metadata,
+                "payload_members": leaf_payload_members,
+            }
+        )
+
+    return _EncodedStructuredLeaf(
+        codec="structured_recursive",
+        rows=len(values),
+        payload=payload,
+        metadata={
+            "schema_source": "inferred_from_runtime_values",
+            "structure_version": 1,
+            "root_path": path,
+            "nodes": node_specs,
+            "leaves": leaf_specs,
+        },
+    )
+
+
+def _is_recursive_encoded_non_tensor(encoded: Mapping[str, Any]) -> bool:
+    return encoded.get("codec") == "structured_recursive"
+
+
+def _structured_path_tokens(path: str) -> list[Any]:
+    tokens: list[Any] = []
+    index = 0
+    current = []
+    while index < len(path):
+        char = path[index]
+        if char == "\\":
+            index += 1
+            if index < len(path):
+                current.append(path[index])
+        elif char == ".":
+            tokens.append("".join(current))
+            current = []
+        elif char == "[":
+            if current:
+                tokens.append("".join(current))
+                current = []
+            end = path.index("]", index)
+            tokens.append(int(path[index + 1 : end]))
+            index = end
+        else:
+            current.append(char)
+        index += 1
+    if current:
+        tokens.append("".join(current))
+    return tokens
+
+
+def _lookup_structured_path(value: Any, root_path: str, path: str) -> Any:
+    root_tokens = _structured_path_tokens(root_path)
+    path_tokens = _structured_path_tokens(path)
+    current = value
+    for token in path_tokens[len(root_tokens) :]:
+        if isinstance(token, str):
+            if current is None:
+                return None
+            if isinstance(current, _Missing) or not isinstance(current, dict):
+                return MISSING
+            current = current.get(token, MISSING)
+            continue
+        if current is None:
+            return None
+        if isinstance(current, _Missing) or not isinstance(current, (list, tuple)):
+            return MISSING
+        current = current[token] if token < len(current) else MISSING
+    return current
+
+
+def _encode_structured_leaf(
+    values: list[Any], decision: _CodecDecision
+) -> _EncodedStructuredLeaf:
+    if not decision.accepted:
+        raise ValueError(f"unsupported structured non-tensor field: {decision.reason}")
+    codec = decision.codec
+    if codec == "ragged_tensor":
+        payload, metadata = _encode_ragged_tensor_values(values)
+    elif codec == "typed_ragged":
+        payload, metadata = _encode_typed_ragged_values(values)
+    elif codec == "media_list_ragged":
+        payload, metadata = _encode_media_list_values(values)
+    elif codec == "ndarray":
+        payload, metadata = _encode_numeric_scalar_values(values, decision)
+    elif codec == "bytes_ragged":
+        payload, metadata = _encode_bytes_like_values(values)
+    elif codec == "media_bytes":
+        payload, metadata = _encode_bytes_like_values(values)
+    elif codec == "utf8_ragged":
+        payload, metadata = _encode_bytes_like_values(
+            [None if value is None else value.encode("utf-8") for value in values]
+        )
+    elif codec == "msgpack_ragged":
+        payload, metadata = _encode_bytes_like_values(
+            [
+                None
+                if value is None
+                else _msgpack.packb(value, use_bin_type=True, strict_types=True)
+                for value in values
+            ]
+        )
+    elif codec == "json_ragged":
+        payload, metadata = _encode_bytes_like_values(
+            [
+                None
+                if value is None
+                else json.dumps(value, ensure_ascii=False, separators=(",", ":")).encode(
+                    "utf-8"
+                )
+                for value in values
+            ]
+        )
+    else:
+        raise ValueError(f"unsupported structured non-tensor codec: {codec}")
+    metadata.update(
+        {
+            "schema_source": "inferred_from_runtime_values"
+            if decision.accepted
+            else "fallback",
+            "decision_reason": decision.reason,
+            "normalized_type": decision.normalized_type,
+            **decision.metadata,
+        }
+    )
+    return _EncodedStructuredLeaf(
+        codec=codec, rows=len(values), payload=payload, metadata=metadata
+    )
+
+
+
+def _decode_structured_non_tensor_encoded(
+    encoded: Mapping[str, Any],
+    payload: dict[str, Any],
+    rows: int,
+    metadata: Optional[Mapping[str, Any]] = None,
+) -> list[Any]:
+    if _is_recursive_encoded_non_tensor(encoded):
+        return _decode_structured_recursive_field(encoded, payload, rows)
+    return _decode_structured_leaf(encoded["codec"], payload, rows, metadata)
+
+
+def _copy_recursive_metadata_for_leaf_updates(
+    metadata: Mapping[str, Any]
+) -> dict[str, Any]:
+    copied = dict(metadata)
+    copied["leaves"] = [dict(leaf) for leaf in metadata.get("leaves", [])]
+    return copied
+
+
+def _decode_structured_recursive_field(
+    encoded: Mapping[str, Any], payload: dict[str, Any], rows: int
+) -> list[Any]:
+    metadata = encoded.get("metadata") or {}
+    leaf_values: dict[str, list[Any]] = {}
+    for leaf in metadata.get("leaves", []):
+        leaf_payload_members = leaf["payload_members"]
+        leaf_payload = {
+            name: payload[global_name]
+            for name, global_name in leaf_payload_members.items()
+            if name != "missing"
+        }
+        values = _decode_structured_leaf(
+            leaf["codec"], leaf_payload, rows, leaf.get("metadata")
+        )
+        missing = payload[leaf_payload_members["missing"]]
+        leaf_values[leaf["path"]] = [
+            MISSING if bool(missing[row]) else values[row] for row in range(rows)
+        ]
+    return _reconstruct_structured_rows(metadata, payload, leaf_values, rows)
+
+
+def _child_path(node: Mapping[str, Any], child: Any) -> str:
+    if node["node_type"] == "dict":
+        return f"{node['path']}.{_escape_key(str(child))}"
+    return f"{node['path']}[{int(child)}]"
+
+
+def _path_depth(path: str) -> int:
+    return path.count(".") + path.count("[")
+
+
+def _reconstruct_structured_rows(
+    metadata: Mapping[str, Any],
+    payload: Mapping[str, Any],
+    leaf_values: Mapping[str, list[Any]],
+    rows: int,
+) -> list[Any]:
+    values_by_path: dict[str, list[Any]] = dict(leaf_values)
+    nodes = sorted(
+        metadata.get("nodes", []), key=lambda node: _path_depth(node["path"]), reverse=True
+    )
+    for node in nodes:
+        missing_payload = node.get("missing_payload")
+        missing = payload[missing_payload] if missing_payload is not None else None
+        row_mask_payload = node.get("row_mask_payload")
+        row_mask = payload[row_mask_payload] if row_mask_payload is not None else None
+        lengths_payload = node.get("lengths_payload")
+        lengths = payload[lengths_payload] if lengths_payload is not None else None
+        node_values = []
+        for row in range(rows):
+            if missing is not None and bool(missing[row]):
+                node_values.append(MISSING)
+                continue
+            if row_mask is not None and not bool(row_mask[row]):
+                node_values.append(None)
+                continue
+            if node["node_type"] == "dict":
+                item = {}
+                for child in node["children"]:
+                    child_values = values_by_path[_child_path(node, child)]
+                    child_value = child_values[row]
+                    if not isinstance(child_value, _Missing):
+                        item[child] = child_value
+                node_values.append(item)
+                continue
+            length = int(lengths[row]) if lengths is not None else len(node["children"])
+            item = [None] * length
+            for child in node["children"]:
+                index = int(child)
+                if index >= length:
+                    continue
+                child_values = values_by_path[_child_path(node, child)]
+                child_value = child_values[row]
+                if not isinstance(child_value, _Missing):
+                    item[index] = child_value
+            node_values.append(item)
+        values_by_path[node["path"]] = node_values
+    return values_by_path[metadata["root_path"]]
+
+
+def _decode_structured_leaf(
+    codec: str,
+    payload: dict[str, Any],
+    rows: int,
+    metadata: Optional[Mapping[str, Any]] = None,
+) -> list[Any]:
+    if codec == "ragged_tensor":
+        return _decode_ragged_tensor_values(payload, rows)
+    if codec == "typed_ragged":
+        return _decode_typed_ragged_values(payload, rows)
+    if codec == "media_list_ragged":
+        return _decode_media_list_values(payload, rows, metadata)
+    if codec == "ndarray":
+        return _decode_numeric_scalar_values(payload)
+    if codec == "media_bytes":
+        return _decode_bytes_like_values(payload, rows, metadata)
+    if codec == "bytes_ragged":
+        return _decode_bytes_like_values(payload, rows)
+    if codec == "utf8_ragged":
+        return [
+            None if value is None else value.decode("utf-8")
+            for value in _decode_bytes_like_values(payload, rows)
+        ]
+    if codec == "msgpack_ragged":
+        return [
+            None if value is None else _msgpack.unpackb(value, raw=False)
+            for value in _decode_bytes_like_values(payload, rows)
+        ]
+    if codec == "json_ragged":
+        return [
+            None if value is None else json.loads(value.decode("utf-8"))
+            for value in _decode_bytes_like_values(payload, rows)
+        ]
+    raise ValueError(f"unknown structured non-tensor codec: {codec}")
+
+
+def _encode_ragged_tensor_values(
+    values: list[Any],
+) -> tuple[dict[str, Any], dict[str, Any]]:
+    if _torch is None:
+        raise RuntimeError("torch is required to encode ragged tensor fields")
+    tensors: list[Any] = []
+    dtype = None
+    max_ndim = 0
+    for value in values:
+        if value is None:
+            tensors.append(None)
+            continue
+        tensor = value.detach()
+        if tensor.device.type != "cpu" or not tensor.is_contiguous():
+            tensor = tensor.cpu().contiguous()
+        dtype = tensor.dtype if dtype is None else dtype
+        if tensor.dtype != dtype:
+            raise ValueError(f"mixed tensor dtype: {dtype} vs {tensor.dtype}")
+        max_ndim = max(max_ndim, tensor.dim())
+        tensors.append(tensor)
+    offsets = _torch.zeros(len(tensors) + 1, dtype=_torch.int64)
+    ndims = _torch.zeros(len(tensors), dtype=_torch.int16)
+    shapes = _torch.zeros((len(tensors), max(max_ndim, 1)), dtype=_torch.int64)
+    nulls = np.asarray([tensor is None for tensor in tensors], dtype=np.bool_)
+    flat_parts = []
+    offset = 0
+    for row, tensor in enumerate(tensors):
+        if tensor is None:
+            offsets[row + 1] = offset
+            continue
+        flat = tensor.reshape(-1)
+        flat_parts.append(flat)
+        offset += flat.numel()
+        offsets[row + 1] = offset
+        ndims[row] = tensor.dim()
+        if tensor.dim() > 0:
+            shapes[row, : tensor.dim()] = _torch.tensor(
+                list(tensor.shape), dtype=_torch.int64
+            )
+    data_dtype = dtype or _torch.float32
+    data = (
+        _torch.cat(flat_parts) if flat_parts else _torch.empty((0,), dtype=data_dtype)
+    )
+    payload = {
+        "data": data.numpy(),
+        "offsets": offsets.numpy(),
+        "shapes": shapes.numpy(),
+        "ndims": ndims.numpy(),
+        "nulls": nulls,
+    }
+    return payload, {
+        "dtype": str(data_dtype),
+        "max_ndim": int(max_ndim),
+        "shape_policy": "ragged",
+    }
+
+
+def _decode_ragged_tensor_values(payload: dict[str, Any], rows: int) -> list[Any]:
+    if _torch is None:
+        raise RuntimeError("torch is required to decode ragged tensor fields")
+    data = _torch.from_numpy(payload["data"])
+    offsets = payload["offsets"]
+    shapes = payload["shapes"]
+    ndims = payload["ndims"]
+    nulls = payload["nulls"]
+    values = []
+    for row in range(rows):
+        if bool(nulls[row]):
+            values.append(None)
+            continue
+        begin = int(offsets[row])
+        end = int(offsets[row + 1])
+        ndim = int(ndims[row])
+        shape = tuple(int(v) for v in shapes[row, :ndim].tolist())
+        values.append(data[begin:end].reshape(shape))
+    return values
+
+
+def _encode_typed_ragged_values(
+    values: list[Any],
+) -> tuple[dict[str, Any], dict[str, Any]]:
+    source_arrays = [np.asarray(value) for value in values if value is not None]
+    dtype = np.result_type(*source_arrays) if source_arrays else np.dtype(np.int64)
+    arrays = [
+        np.asarray([], dtype=dtype)
+        if value is None
+        else np.ascontiguousarray(np.asarray(value, dtype=dtype))
+        for value in values
+    ]
+    max_ndim = max((array.ndim for array in arrays), default=0)
+    offsets = np.zeros(len(arrays) + 1, dtype=np.int64)
+    ndims = np.zeros(len(arrays), dtype=np.int16)
+    shapes = np.zeros((len(arrays), max(max_ndim, 1)), dtype=np.int64)
+    nulls = np.asarray([value is None for value in values], dtype=np.bool_)
+    flat_arrays = []
+    offset = 0
+    for row, array in enumerate(arrays):
+        flat = array.reshape(-1)
+        flat_arrays.append(flat)
+        offset += flat.size
+        offsets[row + 1] = offset
+        ndims[row] = array.ndim
+        if array.ndim > 0:
+            shapes[row, : array.ndim] = array.shape
+    data = np.concatenate(flat_arrays) if flat_arrays else np.asarray([], dtype=dtype)
+    return (
+        {
+            "data": data,
+            "offsets": offsets,
+            "shapes": shapes,
+            "ndims": ndims,
+            "nulls": nulls,
+        },
+        {"dtype": str(data.dtype), "max_ndim": int(max_ndim), "shape_policy": "ragged"},
+    )
+
+
+def _decode_typed_ragged_values(payload: dict[str, Any], rows: int) -> list[Any]:
+    data = payload["data"]
+    offsets = payload["offsets"]
+    shapes = payload["shapes"]
+    ndims = payload["ndims"]
+    nulls = payload["nulls"]
+    values = []
+    for row in range(rows):
+        if bool(nulls[row]):
+            values.append(None)
+            continue
+        begin = int(offsets[row])
+        end = int(offsets[row + 1])
+        ndim = int(ndims[row])
+        shape = tuple(int(v) for v in shapes[row, :ndim].tolist())
+        values.append(data[begin:end].reshape(shape).tolist())
+    return values
+
+
+def _encode_numeric_scalar_values(
+    values: list[Any], decision: _CodecDecision
+) -> tuple[dict[str, Any], dict[str, Any]]:
+    dtype = np.dtype(decision.metadata["dtype"])
+    nulls = np.asarray([value is None for value in values], dtype=np.bool_)
+    fill_value = False if dtype == np.dtype(np.bool_) else 0
+    data = np.asarray(
+        [fill_value if value is None else value for value in values], dtype=dtype
+    )
+    return {"data": data, "nulls": nulls}, {
+        "dtype": str(data.dtype),
+        "shape": list(data.shape),
+    }
+
+
+def _decode_numeric_scalar_values(payload: dict[str, Any]) -> list[Any]:
+    data = payload["data"]
+    nulls = payload["nulls"]
+    if not bool(nulls.any()):
+        return data.tolist()
+    values = data.astype(object)
+    values[nulls] = None
+    return values.tolist()
+
+
+def _value_to_media_bytes(value: Any) -> tuple[bytes, str | None, dict[str, Any]]:
+    if value is None:
+        return b"", None, {"kind": "null"}
+    if _is_bytes_like(value):
+        return bytes(value), None, {"kind": "bytes"}
+    if _is_pil_image(value):
+        image = (
+            value.convert(value.mode) if getattr(value, "readonly", False) else value
+        )
+        return (
+            image.tobytes(),
+            "image/raw",
+            {
+                "kind": "pil_raw",
+                "mode": image.mode,
+                "size": list(image.size),
+                "format": getattr(value, "format", None),
+            },
+        )
+    return bytes(value), None, {"kind": "bytes"}
+
+
+def _encode_bytes_like_values(
+    values: list[Any],
+) -> tuple[dict[str, Any], dict[str, Any]]:
+    offsets = [0]
+    parts = []
+    media_types = []
+    encodings = []
+    for value in values:
+        data, media_type, encoding = _value_to_media_bytes(value)
+        parts.append(data)
+        media_types.append(media_type)
+        encodings.append(encoding)
+        offsets.append(offsets[-1] + len(data))
+    payload_bytes = b"".join(parts)
+    metadata: dict[str, Any] = {}
+    if any(encoding.get("kind") == "pil_raw" for encoding in encodings):
+        metadata["media_encodings"] = encodings
+    non_null_media_types = sorted(
+        {media_type for media_type in media_types if media_type}
+    )
+    if non_null_media_types:
+        metadata["media_types"] = non_null_media_types
+        metadata["encode_source"] = "raw_pixels_from_object"
+    return (
+        {
+            "data": payload_bytes,
+            "offsets": np.asarray(offsets, dtype=np.int64),
+            "nulls": np.asarray([value is None for value in values], dtype=np.bool_),
+        },
+        metadata,
+    )
+
+
+def _decode_media_bytes(data: Any, encoding: Optional[Mapping[str, Any]]) -> Any:
+    if not encoding or encoding.get("kind") != "pil_raw":
+        return bytes(data)
+    try:
+        from PIL import Image
+    except ImportError:
+        return bytes(data)
+    image = Image.frombuffer(
+        encoding["mode"], tuple(encoding["size"]), data, "raw", encoding["mode"], 0, 1
+    )
+    image.format = encoding.get("format")
+    return image
+
+
+def _decode_bytes_like_values(
+    payload: dict[str, Any], rows: int, metadata: Optional[Mapping[str, Any]] = None
+) -> list[Any]:
+    data = payload["data"]
+    offsets = payload["offsets"]
+    nulls = payload["nulls"]
+    encodings = (metadata or {}).get("media_encodings", [])
+    values = []
+    for row in range(rows):
+        if bool(nulls[row]):
+            values.append(None)
+            continue
+        item = memoryview(data)[int(offsets[row]) : int(offsets[row + 1])]
+        encoding = encodings[row] if row < len(encodings) else None
+        values.append(_decode_media_bytes(item, encoding))
+    return values
+
+
+def _encode_media_list_values(
+    values: list[Any],
+) -> tuple[dict[str, Any], dict[str, Any]]:
+    row_offsets = [0]
+    byte_offsets = [0]
+    parts = []
+    media_types = []
+    encodings = []
+    for value in values:
+        items = [] if value is None else list(value)
+        for item in items:
+            data, media_type, encoding = _value_to_media_bytes(item)
+            parts.append(data)
+            media_types.append(media_type)
+            encodings.append(encoding)
+            byte_offsets.append(byte_offsets[-1] + len(data))
+        row_offsets.append(row_offsets[-1] + len(items))
+    metadata: dict[str, Any] = {
+        "encode_source": "raw_pixels_from_object",
+        "media_encodings": encodings,
+    }
+    non_null_media_types = sorted(
+        {media_type for media_type in media_types if media_type}
+    )
+    if non_null_media_types:
+        metadata["media_types"] = non_null_media_types
+    return (
+        {
+            "data": b"".join(parts),
+            "row_offsets": np.asarray(row_offsets, dtype=np.int64),
+            "byte_offsets": np.asarray(byte_offsets, dtype=np.int64),
+            "nulls": np.asarray([value is None for value in values], dtype=np.bool_),
+        },
+        metadata,
+    )
+
+
+def _decode_media_list_values(
+    payload: dict[str, Any], rows: int, metadata: Optional[Mapping[str, Any]] = None
+) -> list[Any]:
+    data = payload["data"]
+    row_offsets = payload["row_offsets"]
+    byte_offsets = payload["byte_offsets"]
+    nulls = payload["nulls"]
+    encodings = (metadata or {}).get("media_encodings", [])
+    values = []
+    for row in range(rows):
+        if bool(nulls[row]):
+            values.append(None)
+            continue
+        items = []
+        for item_index in range(int(row_offsets[row]), int(row_offsets[row + 1])):
+            item = memoryview(data)[
+                int(byte_offsets[item_index]) : int(byte_offsets[item_index + 1])
+            ]
+            encoding = encodings[item_index] if item_index < len(encodings) else None
+            items.append(_decode_media_bytes(item, encoding))
+        values.append(items)
+    return values
+
+
+class _StructuredObjectLayer:
+    def __init__(self, bundle_store: "_BundleManifestStore") -> None:
+        self._bundle_store = bundle_store
+
+    def put_structured_object(
+        self,
+        payload: StructuredObjectPayload,
+        partition: str,
+        chunk_bytes: Optional[int],
+        policy: Optional[BundleTransferPolicy],
+        max_inflight_put: Optional[int],
+        pre_registered_buffers: Optional[Mapping[str, bool]],
+    ) -> RemoteBundleRef:
+        metadata, buffers = _encode_structured_fields(payload.metadata, payload.buffers)
+        transfer_policy = self._bundle_store._policy(
+            policy, max_inflight_put=max_inflight_put
+        )
+        if transfer_policy.copy_mode != "copy":
+            _validate_pre_registered_structured_buffers(
+                payload.buffers, buffers, pre_registered_buffers
+            )
+        return self._bundle_store.put_bundle(
+            meta=_encode_structured_metadata(metadata),
+            buffers=buffers,
+            partition=partition,
+            chunk_bytes=chunk_bytes,
+            policy=transfer_policy,
+            max_inflight_put=None,
+            pre_registered_buffers=pre_registered_buffers,
+        )
+
+    def read_spec(
+        self, ref: RemoteBundleRef | Mapping[str, Any]
+    ) -> StructuredObjectReadSpec:
+        return StructuredObjectReadSpec(ref=ref)
+
+    def materialize(self, spec: StructuredObjectReadSpec) -> StructuredObjectResult:
+        return self.materialize_into(spec, destinations=None)
+
+    def materialize_into(
+        self,
+        spec: StructuredObjectReadSpec,
+        destinations: Optional[Mapping[str, Any]],
+    ) -> StructuredObjectResult:
+        metadata, buffers, selected = self._resolve_structured_read(spec)
+        field_specs = _structured_field_specs(metadata)
+        destination_map = destinations or {}
+        member_slices = dict(spec.member_slices)
+        objects = {
+            name: self._read_structured_member(
+                name=name,
+                payload_spec=buffers[name],
+                field_spec=field_specs.get(name, {"encoding": "bytes"}),
+                member_slice=member_slices.get(name),
+                destination=destination_map.get(name),
+            )
+            for name in selected
+        }
+        return StructuredObjectResult(metadata=dict(metadata), objects=objects)
+
+    def _resolve_structured_read(
+        self,
+        spec: StructuredObjectReadSpec,
+    ) -> tuple[dict[str, Any], Mapping[str, Any], list[str]]:
+        if spec.member_names is not None and not spec.member_names:
+            raise ValueError("structured read spec selected no members")
+        manifest = self._bundle_store.resolve_manifest(spec.ref)
+        metadata = _decode_structured_metadata(
+            self._bundle_store.read_payload(manifest["meta"])
+        )
+        buffers = manifest["buffers"]
+        selected = (
+            list(buffers) if spec.member_names is None else list(spec.member_names)
+        )
+        missing = [name for name in selected if name not in buffers]
+        if missing:
+            raise KeyError(f"unknown bundle buffers: {missing}")
+        return metadata, buffers, selected
+
+    def _read_structured_member(
+        self,
+        name: str,
+        payload_spec: Mapping[str, Any],
+        field_spec: Mapping[str, Any],
+        member_slice: StructuredMemberSlice | None,
+        destination: Any,
+    ) -> Any:
+        encoding = field_spec.get("encoding", "bytes")
+        if encoding == "bytes":
+            return self._read_bytes_member(
+                name, payload_spec, member_slice, destination
+            )
+        if encoding == "torch_tensor":
+            return self._read_torch_tensor_member(
+                name, payload_spec, field_spec, member_slice, destination
+            )
+        if encoding != "ndarray":
+            raise ValueError(f"unsupported structured field encoding: {encoding}")
+        return self._read_ndarray_member(
+            name, payload_spec, field_spec, member_slice, destination
+        )
+
+    def _read_torch_tensor_member(
+        self,
+        name: str,
+        payload_spec: Mapping[str, Any],
+        field_spec: Mapping[str, Any],
+        member_slice: StructuredMemberSlice | None,
+        destination: Any,
+    ) -> Any:
+        if member_slice is not None:
+            return self._read_sliced_torch_tensor_member(
+                name, payload_spec, field_spec, member_slice, destination
+            )
+        if destination is not None:
+            if not isinstance(destination, (_TensorObjectBufferPayload, _RawDestinationBuffer)):
+                raise ValueError(
+                    f"structured tensor member {name} only supports tensor_object_buffer or raw_destination destinations"
+                )
+            materialized = self._bundle_store.read_tensor_payload_into(
+                payload_spec, destination
+            )
+            return destination if materialized is None else materialized
+        if payload_spec.get("kind") == "tensor":
+            return self._bundle_store.read_tensor_payload(payload_spec)
+        if payload_spec.get("format") == "torch_save":
+            return _deserialize_torch_save_payload(
+                self._bundle_store.read_payload(payload_spec)
+            )
+        return _deserialize_tensor_payload(
+            self._bundle_store.read_payload(payload_spec)
+        )
+
+    def _read_sliced_torch_tensor_member(
+        self,
+        name: str,
+        payload_spec: Mapping[str, Any],
+        field_spec: Mapping[str, Any],
+        member_slice: StructuredMemberSlice,
+        destination: Any,
+    ) -> Any:
+        metadata_bytes = int(payload_spec.get("metadata_bytes", -1))
+        shape = field_spec.get("shape")
+        element_size = int(field_spec.get("element_size", 0))
+        if metadata_bytes < 0 or not isinstance(shape, list) or element_size <= 0:
+            raise ValueError(
+                f"structured tensor field {name} is missing slice metadata"
+            )
+        if member_slice.axis != 0:
+            raise ValueError("structured tensor slicing currently supports axis=0 only")
+        if not shape:
+            raise ValueError(
+                "structured tensor slicing requires at least one dimension"
+            )
+        start, end, step = _normalized_member_slice(member_slice, int(shape[0]))
+        if step != 1:
+            raise ValueError("structured tensor slicing currently requires step=1")
+        row_width = element_size * int(np.prod(shape[1:], dtype=np.int64))
+        data_offset = metadata_bytes + start * row_width
+        data_length = (end - start) * row_width
+        metadata = self._bundle_store.read_payload_range(
+            payload_spec, 0, metadata_bytes
+        )
+        sliced_metadata = _slice_tensor_metadata(
+            metadata, (end - start, *shape[1:]), data_length
+        )
+        if destination is not None:
+            if not isinstance(destination, (_TensorObjectBufferPayload, _RawDestinationBuffer)):
+                raise ValueError(
+                    f"structured tensor member {name} only supports tensor_object_buffer or raw_destination destinations"
+                )
+            expected_bytes = metadata_bytes + data_length
+            if destination.size < expected_bytes:
+                raise ValueError(
+                    f"tensor destination has {destination.size} bytes, expected at least {expected_bytes}"
+                )
+            ctypes.memmove(destination.ptr, sliced_metadata, metadata_bytes)
+            if isinstance(destination, _RawDestinationBuffer) and destination.pre_registered:
+                self._bundle_store.read_payload_range_into_raw_destination(
+                    payload_spec,
+                    destination.ptr,
+                    metadata_bytes,
+                    data_offset,
+                    data_length,
+                )
+            else:
+                data = self._bundle_store.read_payload_range(
+                    payload_spec, data_offset, data_length
+                )
+                ctypes.memmove(destination.ptr + metadata_bytes, data, data_length)
+            _ = destination.owner
+            return destination
+        data = self._bundle_store.read_payload_range(
+            payload_spec, data_offset, data_length
+        )
+        return _deserialize_tensor_payload(sliced_metadata + data)
+
+    def _read_bytes_member(
+        self,
+        name: str,
+        payload_spec: Mapping[str, Any],
+        member_slice: StructuredMemberSlice | None,
+        destination: Any,
+    ) -> bytes:
+        if member_slice is not None:
+            raise ValueError(f"structured bytes member {name} does not support slicing")
+        if destination is not None:
+            raise ValueError(
+                f"structured bytes member {name} does not support materialize_into"
+            )
+        return self._bundle_store.read_payload(payload_spec)
+
+    def _read_ndarray_member(
+        self,
+        name: str,
+        payload_spec: Mapping[str, Any],
+        field_spec: Mapping[str, Any],
+        member_slice: StructuredMemberSlice | None,
+        destination: Any,
+    ) -> np.ndarray:
+        dtype = field_spec.get("dtype")
+        shape = field_spec.get("shape")
+        if not isinstance(dtype, str) or not isinstance(shape, list):
+            raise ValueError(
+                f"structured ndarray field {name} is missing dtype or shape"
+            )
+        read_plan = _resolve_ndarray_read_plan(
+            tuple(int(dim) for dim in shape), np.dtype(dtype), member_slice
+        )
+        target = _resolve_ndarray_destination(
+            name, destination, read_plan.dtype, read_plan.output_shape
+        )
+        destination_view = target.view(np.uint8).reshape(-1)
+        if read_plan.byte_length == 0:
+            return target
+        if read_plan.step == 1:
+            self._bundle_store.read_payload_range_into_destination(
+                payload_spec,
+                destination_view,
+                read_plan.byte_offset,
+            )
+            return target
+        temp_shape = (read_plan.cover_row_count, *read_plan.full_shape[1:])
+        temp = np.empty(temp_shape, dtype=read_plan.dtype)
+        temp_view = temp.view(np.uint8).reshape(-1)
+        self._bundle_store.read_payload_range_into_destination(
+            payload_spec, temp_view, read_plan.byte_offset
+        )
+        target[...] = temp[:: read_plan.step]
+        return target
+
+
+class _BundleManifestStore:
+    def __init__(
+        self,
+        store: BundleStore,
+        transport: "_MooncakePayloadTransport",
+        key_prefix: str,
+        default_chunk_bytes: int,
+    ) -> None:
+        self._store = store
+        self._transport = transport
+        self._key_prefix = key_prefix
+        self._default_chunk_bytes = default_chunk_bytes
+
+    def put_bundle(
+        self,
+        meta: bytes | bytearray | memoryview,
+        buffers: Mapping[str, Any],
+        partition: str,
+        chunk_bytes: Optional[int],
+        policy: Optional[BundleTransferPolicy],
+        max_inflight_put: Optional[int],
+        pre_registered_buffers: Optional[Mapping[str, bool]] = None,
+    ) -> RemoteBundleRef:
+        _validate_key_segment(partition, "partition")
+        meta_view = _bytes_view(meta, "meta")
+        target_chunk_bytes = _validate_chunk_bytes(
+            self._default_chunk_bytes if chunk_bytes is None else chunk_bytes
+        )
+        transfer_policy = self._policy(policy, max_inflight_put=max_inflight_put)
+        object_id = f"{partition}/{uuid.uuid4().hex}"
+        base_key = f"{self._key_prefix}/{object_id}"
+        manifest_key = f"{base_key}/manifest"
+        written_keys: list[str] = []
+        buffer_specs: dict[str, Any] = {}
+        pre_registered_map = dict(pre_registered_buffers or {})
+        try:
+            meta_spec, meta_keys = self._put_payload(
+                f"{base_key}/meta",
+                meta_view,
+                target_chunk_bytes,
+                _copy_transfer_policy(transfer_policy),
+                pre_registered=False,
+            )
+            written_keys.extend(meta_keys)
+            for name, value in buffers.items():
+                _validate_key_segment(name, "buffer name")
+                payload_key = f"{base_key}/buffer/{name}"
+                if isinstance(value, (_TensorPayload, _TensorObjectBufferPayload)):
+                    payload_spec, payload_keys = self._put_tensor_payload(
+                        payload_key,
+                        value,
+                        transfer_policy,
+                    )
+                else:
+                    payload_spec, payload_keys = self._put_payload(
+                        payload_key,
+                        _bytes_view(value, name),
+                        target_chunk_bytes,
+                        transfer_policy,
+                        pre_registered=bool(pre_registered_map.get(name, False)),
+                    )
+                buffer_specs[name] = payload_spec
+                written_keys.extend(payload_keys)
+            manifest = {
+                "version": 1,
+                "layout": "bundle",
+                "object_id": object_id,
+                "meta": meta_spec,
+                "buffers": buffer_specs,
+            }
+            manifest_blob = _encode_manifest(manifest)
+            _check_status(
+                self._store.put(manifest_key, manifest_blob), "put", manifest_key
+            )
+            written_keys.append(manifest_key)
+        except Exception:
+            _cleanup_keys(self._store, written_keys, strict=False)
+            raise
+        return RemoteBundleRef(manifest_key=manifest_key, manifest=manifest)
+
+    def put_bundle_manifest(
+        self,
+        meta: bytes | bytearray | memoryview,
+        buffers: Mapping[str, Any],
+        *,
+        partition: str,
+        chunk_bytes: Optional[int],
+        policy: Optional[BundleTransferPolicy],
+        cleanup_keys: Optional[Sequence[str]] = None,
+    ) -> RemoteBundleRef:
+        _validate_key_segment(partition, "partition")
+        meta_view = _bytes_view(meta, "meta")
+        target_chunk_bytes = _validate_chunk_bytes(
+            self._default_chunk_bytes if chunk_bytes is None else chunk_bytes
+        )
+        transfer_policy = self._policy(policy)
+        object_id = f"{partition}/{uuid.uuid4().hex}"
+        base_key = f"{self._key_prefix}/{object_id}"
+        manifest_key = f"{base_key}/manifest"
+        written_keys: list[str] = []
+        try:
+            meta_spec, meta_keys = self._put_payload(
+                f"{base_key}/meta",
+                meta_view,
+                target_chunk_bytes,
+                _copy_transfer_policy(transfer_policy),
+                pre_registered=False,
+            )
+            written_keys.extend(meta_keys)
+            manifest = {
+                "version": 1,
+                "layout": "bundle",
+                "object_id": object_id,
+                "meta": meta_spec,
+                "buffers": dict(buffers),
+                "buffer_object_ids": sorted(
+                    {
+                        str(payload_spec["key"])
+                        .removeprefix(f"{self._key_prefix}/")
+                        .split("/buffer/", 1)[0]
+                        for payload_spec in buffers.values()
+                    }
+                ),
+            }
+            if cleanup_keys:
+                manifest["cleanup_keys"] = list(dict.fromkeys(cleanup_keys))
+            self._validate_manifest(manifest)
+            manifest_blob = _encode_manifest(manifest)
+            _check_status(
+                self._store.put(manifest_key, manifest_blob), "put", manifest_key
+            )
+            written_keys.append(manifest_key)
+        except Exception:
+            _cleanup_keys(self._store, written_keys, strict=False)
+            raise
+        return RemoteBundleRef(manifest_key=manifest_key, manifest=manifest)
+
+    def remove_bundle(self, ref: RemoteBundleRef | Mapping[str, Any]) -> None:
+        manifest = self.resolve_manifest(ref)
+        keys = self._payload_keys(manifest)
+        keys.extend(manifest.get("cleanup_keys", []))
+        keys.append(self._manifest_key(ref, manifest))
+        _cleanup_keys(self._store, keys, strict=True)
+
+    def manifest_key(self, ref: RemoteBundleRef | Mapping[str, Any]) -> str:
+        return self._manifest_key(ref, self.resolve_manifest(ref))
+
+    def payload_keys(self, payload_spec: Mapping[str, Any]) -> list[str]:
+        return [chunk["key"] for chunk in payload_spec["chunks"]]
+
+    def remove_keys(self, keys: Sequence[str], *, strict: bool) -> None:
+        _cleanup_keys(self._store, keys, strict=strict)
+
+    def resolve_manifest(
+        self, ref: RemoteBundleRef | Mapping[str, Any]
+    ) -> dict[str, Any]:
+        if isinstance(ref, RemoteBundleRef):
+            manifest = ref.manifest
+            if not manifest:
+                manifest = _decode_manifest(self._store.get(ref.manifest_key))
+        else:
+            manifest = ref.get("manifest")
+            if manifest is None:
+                manifest_key = ref.get("manifest_key")
+                if not isinstance(manifest_key, str):
+                    raise ValueError("bundle ref must include manifest_key")
+                manifest = _decode_manifest(self._store.get(manifest_key))
+        self._validate_manifest(manifest)
+        return manifest
+
+    def read_payload(self, payload_spec: Mapping[str, Any]) -> bytes:
+        return self._transport.read_payload(payload_spec)
+
+    def read_payload_range(
+        self, payload_spec: Mapping[str, Any], byte_offset: int, byte_length: int
+    ) -> bytes:
+        return self._transport.read_payload_range(
+            payload_spec, byte_offset, byte_length
+        )
+
+    def read_tensor_payload(self, payload_spec: Mapping[str, Any]) -> Any:
+        return self._transport.read_tensor_payload(payload_spec)
+
+    def read_tensor_payload_into(
+        self, payload_spec: Mapping[str, Any], destination: _TensorObjectBufferPayload
+    ) -> Any:
+        return self._transport.read_tensor_payload_into(payload_spec, destination)
+
+    def read_payload_range_into_destination(
+        self,
+        payload_spec: Mapping[str, Any],
+        destination: np.ndarray,
+        byte_offset: int,
+        destination_pre_registered: bool = False,
+    ) -> None:
+        self._transport.read_payload_range_into_destination(
+            payload_spec,
+            destination,
+            byte_offset,
+            destination_pre_registered=destination_pre_registered,
+        )
+
+    def read_payload_range_into_raw_destination(
+        self,
+        payload_spec: Mapping[str, Any],
+        destination_ptr: int,
+        destination_offset: int,
+        byte_offset: int,
+        byte_length: int,
+    ) -> None:
+        if not self._transport.read_payload_range_into_raw_destination(
+            payload_spec, destination_ptr, destination_offset, byte_offset, byte_length
+        ):
+            data = self.read_payload_range(payload_spec, byte_offset, byte_length)
+            ctypes.memmove(destination_ptr + destination_offset, data, byte_length)
+
+    def read_payload_ranges_into_array(
+        self,
+        payload_spec: Mapping[str, Any],
+        destination: np.ndarray,
+        ranges: Sequence[tuple[int, int, int]],
+    ) -> None:
+        if not ranges:
+            return
+        self._transport.read_payload_ranges_into_array(payload_spec, destination, ranges)
+
+    def read_payload_ranges_into_bytearray(
+        self,
+        payload_spec: Mapping[str, Any],
+        destination: bytearray,
+        ranges: Sequence[tuple[int, int, int]],
+    ) -> None:
+        if not ranges:
+            return
+        self._transport.read_payload_ranges_into_bytearray(
+            payload_spec, destination, ranges
+        )
+
+    def read_payload_ranges_into_raw_destination(
+        self,
+        payload_spec: Mapping[str, Any],
+        destination_ptr: int,
+        ranges: Sequence[tuple[int, int, int]],
+    ) -> None:
+        if not ranges:
+            return
+        if self._transport.read_payload_ranges_into_raw_destination(
+            payload_spec, destination_ptr, ranges
+        ):
+            return
+        for byte_offset, destination_offset, byte_length in ranges:
+            data = self.read_payload_range(payload_spec, byte_offset, byte_length)
+            ctypes.memmove(destination_ptr + destination_offset, data, byte_length)
+
+    def _put_tensor_payload(
+        self,
+        key: str,
+        value: _TensorPayload | _TensorObjectBufferPayload,
+        transfer_policy: BundleTransferPolicy,
+    ) -> tuple[dict[str, Any], list[str]]:
+        if isinstance(value, _TensorObjectBufferPayload):
+            total_bytes = self._transport.put_tensor_object_buffer(key, value)
+            return {
+                "key": key,
+                "bytes": total_bytes,
+                "chunks": [{"key": key, "bytes": total_bytes}],
+            }, [key]
+        if transfer_policy.copy_mode != "copy":
+            if transfer_policy.copy_mode == "zero_copy":
+                raise ValueError(
+                    "zero-copy structured tensor fields require a BufferPool tensor-object buffer"
+                )
+            tensor_spec = self._transport.put_tensor_payload_direct(key, value)
+            if tensor_spec is not None:
+                return tensor_spec, [key]
+            try:
+                total_bytes = self._transport.put_tensor_payload_from_pool(key, value)
+                return {
+                    "key": key,
+                    "bytes": total_bytes,
+                    "chunks": [{"key": key, "bytes": total_bytes}],
+                }, [key]
+            except RuntimeError:
+                pass
+        if not _has_tensor_codec_helpers() or value.tensor.dim() == 0:
+            payload = _torch_save_payload_bytes(value.tensor)
+            payload_spec, payload_keys = self._put_payload(
+                key,
+                memoryview(payload),
+                len(payload) or 1,
+                transfer_policy,
+                pre_registered=False,
+            )
+            payload_spec["format"] = "torch_save"
+            return payload_spec, payload_keys
+        payload, metadata_bytes = _tensor_payload_bytes(value)
+        payload_spec, payload_keys = self._put_payload(
+            key,
+            memoryview(payload),
+            len(payload) or 1,
+            transfer_policy,
+            pre_registered=False,
+        )
+        payload_spec["metadata_bytes"] = metadata_bytes
+        return payload_spec, payload_keys
+
+    def _put_payload(
+        self,
+        key: str,
+        value: memoryview,
+        chunk_bytes: int,
+        transfer_policy: BundleTransferPolicy,
+        pre_registered: bool,
+    ) -> tuple[dict[str, Any], list[str]]:
+        if len(value) == 0:
+            return {"key": key, "bytes": 0, "chunks": []}, []
+        chunks = _split_view(value, chunk_bytes)
+        chunk_keys = [
+            key if len(chunks) == 1 else f"{key}/chunk/{index}"
+            for index in range(len(chunks))
+        ]
+        written_keys = self._transport.put_payload_chunks(
+            chunk_keys,
+            chunks,
+            transfer_policy,
+            pre_registered=pre_registered,
+        )
+        payload_spec = {
+            "key": key,
+            "bytes": len(value),
+            "chunks": [
+                {"key": chunk_key, "bytes": len(chunk)}
+                for chunk_key, chunk in zip(chunk_keys, chunks)
+            ],
+        }
+        return payload_spec, written_keys
+
+    def _policy(
+        self,
+        policy: Optional[BundleTransferPolicy],
+        max_inflight_put: Optional[int] = None,
+    ) -> BundleTransferPolicy:
+        result = policy or BundleTransferPolicy()
+        if max_inflight_put is not None:
+            result = BundleTransferPolicy(
+                max_inflight_put=max_inflight_put,
+                put_mode=result.put_mode,
+                copy_mode=result.copy_mode,
+            )
+        if result.max_inflight_put < 1:
+            raise ValueError("max_inflight_put must be positive")
+        if result.put_mode not in {"auto", "batch", "parallel"}:
+            raise ValueError(f"unsupported put_mode: {result.put_mode}")
+        if result.copy_mode not in {"auto", "zero_copy", "copy"}:
+            raise ValueError(f"unsupported copy_mode: {result.copy_mode}")
+        return result
+
+    def _validate_manifest(self, manifest: Mapping[str, Any]) -> None:
+        if manifest.get("version") != 1 or manifest.get("layout") != "bundle":
+            raise ValueError("invalid bundle manifest")
+        object_id = manifest.get("object_id")
+        if not isinstance(object_id, str):
+            raise ValueError("bundle manifest object_id must be a string")
+        base_key = f"{self._key_prefix}/{object_id}"
+        buffer_object_ids = manifest.get("buffer_object_ids", [object_id])
+        if not isinstance(buffer_object_ids, list) or not all(
+            isinstance(item, str) for item in buffer_object_ids
+        ):
+            raise ValueError("bundle manifest buffer_object_ids must be a list of strings")
+        allowed_buffer_base_keys = [
+            f"{self._key_prefix}/{buffer_object_id}"
+            for buffer_object_id in buffer_object_ids
+        ]
+        self._validate_payload_spec(manifest.get("meta"), base_keys=[base_key])
+        cleanup_keys = manifest.get("cleanup_keys", [])
+        if not isinstance(cleanup_keys, list) or not all(
+            self._is_allowed_cleanup_key(item, allowed_buffer_base_keys)
+            for item in cleanup_keys
+        ):
+            raise ValueError("bundle manifest cleanup_keys are outside the bundle namespace")
+        buffers = manifest.get("buffers")
+        if not isinstance(buffers, dict):
+            raise ValueError("bundle manifest buffers must be a dict")
+        for name, payload_spec in buffers.items():
+            _validate_key_segment(name, "buffer name")
+            self._validate_payload_spec(
+                payload_spec, base_keys=allowed_buffer_base_keys
+            )
+
+    def _is_allowed_cleanup_key(self, key: Any, base_keys: Sequence[str]) -> bool:
+        if not isinstance(key, str):
+            return False
+        return any(
+            key == f"{base_key}/manifest" or key.startswith(f"{base_key}/")
+            for base_key in base_keys
+        )
+
+    def _validate_payload_spec(
+        self, payload_spec: Any, base_keys: Sequence[str] | None = None
+    ) -> None:
+        if not isinstance(payload_spec, dict):
+            raise ValueError("bundle payload spec must be a dict")
+        payload_key = payload_spec.get("key")
+        if not isinstance(payload_key, str) or (
+            base_keys is not None
+            and not any(payload_key.startswith(f"{base_key}/") for base_key in base_keys)
+        ):
+            raise ValueError("bundle payload key is outside the bundle namespace")
+        expected_bytes = int(payload_spec.get("bytes", -1))
+        if expected_bytes < 0:
+            raise ValueError("bundle payload bytes must be non-negative")
+        chunks = payload_spec.get("chunks")
+        if not isinstance(chunks, list) or (expected_bytes and not chunks):
+            raise ValueError("bundle payload chunks are invalid")
+        seen_keys = set()
+        total_bytes = 0
+        for chunk in chunks:
+            if not isinstance(chunk, dict):
+                raise ValueError("bundle chunk must be a dict")
+            key = chunk.get("key")
+            chunk_bytes = int(chunk.get("bytes", -1))
+            if not isinstance(key, str) or (
+                key != payload_key and not key.startswith(f"{payload_key}/")
+            ):
+                raise ValueError("bundle chunk key is outside the payload namespace")
+            if key in seen_keys:
+                raise ValueError("bundle chunk keys must be unique")
+            if chunk_bytes < 0:
+                raise ValueError("bundle chunk bytes must be non-negative")
+            seen_keys.add(key)
+            total_bytes += chunk_bytes
+        if total_bytes != expected_bytes:
+            raise ValueError(
+                f"bundle payload chunks total {total_bytes} bytes, expected {expected_bytes}"
+            )
+
+    def _manifest_key(
+        self, ref: RemoteBundleRef | Mapping[str, Any], manifest: Mapping[str, Any]
+    ) -> str:
+        expected = f"{self._key_prefix}/{manifest['object_id']}/manifest"
+        manifest_key = (
+            ref.manifest_key
+            if isinstance(ref, RemoteBundleRef)
+            else ref.get("manifest_key")
+        )
+        if manifest_key is None:
+            return expected
+        if manifest_key != expected:
+            raise ValueError("bundle manifest_key does not match manifest object_id")
+        return manifest_key
+
+    def _payload_keys(self, manifest: Mapping[str, Any]) -> list[str]:
+        return [
+            chunk["key"]
+            for payload_spec in [manifest["meta"], *manifest["buffers"].values()]
+            for chunk in payload_spec["chunks"]
+        ]
+
+
+class _MooncakePayloadTransport:
+    """Move payload bytes through Mooncake, preferring fast-path APIs and falling back to generic store calls."""
+
+    def __init__(self, store: BundleStore, buffer_pool: Any = None) -> None:
+        self._store = store
+        self._buffer_pool = buffer_pool
+        self._batch_put_from = getattr(store, "batch_put_from", None)
+        self._put_tensor_from = getattr(store, "put_tensor_from", None)
+        self._batch_get_into = getattr(store, "batch_get_into", None)
+        self._get_into = getattr(store, "get_into", None)
+        self._get_into_ranges = getattr(store, "get_into_ranges", None)
+        self._register_buffer = getattr(store, "register_buffer", None)
+        self._unregister_buffer = getattr(store, "unregister_buffer", None)
+
+    def put_payload_chunks(
+        self,
+        chunk_keys: Sequence[str],
+        chunks: Sequence[memoryview],
+        transfer_policy: BundleTransferPolicy,
+        pre_registered: bool,
+    ) -> list[str]:
+        if transfer_policy.copy_mode == "copy":
+            return self._put_chunks_direct(chunk_keys, chunks)
+        if not self._has_batch_put_support():
+            if transfer_policy.copy_mode == "zero_copy":
+                raise RuntimeError(
+                    "zero-copy put requested but batch_put_from is unavailable"
+                )
+            return self._put_chunks_direct(chunk_keys, chunks)
+        put_mode = self._resolve_put_mode(chunks, transfer_policy)
+        if put_mode == "batch":
+            self.batch_put_chunks_from(
+                chunk_keys, chunks, pre_registered=pre_registered
+            )
+            return list(chunk_keys)
+        return self._put_chunks_parallel(
+            list(chunk_keys),
+            list(chunks),
+            transfer_policy.max_inflight_put,
+            pre_registered=pre_registered,
+        )
+
+    def read_payload(self, payload_spec: Mapping[str, Any]) -> bytes:
+        expected_bytes = int(payload_spec["bytes"])
+        if expected_bytes == 0:
+            return b""
+        data = bytearray(expected_bytes)
+        self.read_payload_into(payload_spec, data)
+        return bytes(data)
+
+    def read_tensor_payload(self, payload_spec: Mapping[str, Any]) -> Any:
+        get_tensor = getattr(self._store, "get_tensor", None)
+        key = payload_spec["key"]
+        if not callable(get_tensor):
+            raise RuntimeError("structured tensor payload does not support get_tensor")
+        return get_tensor(key)
+
+    def read_payload_into(
+        self, payload_spec: Mapping[str, Any], destination: bytearray | np.ndarray
+    ) -> None:
+        chunks = payload_spec["chunks"]
+        offsets = _chunk_offsets(chunks)
+        if self._read_chunks_with_batch_get_into(chunks, offsets, destination):
+            return
+        for offset, chunk in zip(offsets, chunks):
+            self._read_chunk_with_get(chunk, destination, offset)
+
+    def read_tensor_payload_into(
+        self,
+        payload_spec: Mapping[str, Any],
+        destination: _TensorObjectBufferPayload,
+    ) -> Any:
+        if payload_spec.get("kind") == "tensor":
+            get_tensor_into = getattr(self._store, "get_tensor_into", None)
+            if not callable(get_tensor_into):
+                raise RuntimeError(
+                    "structured tensor payload does not support get_tensor_into"
+                )
+            result = get_tensor_into(
+                payload_spec["key"], destination.ptr, destination.size
+            )
+            _ = destination.owner
+            return result
+        expected_bytes = int(payload_spec["bytes"])
+        if destination.size < expected_bytes:
+            raise ValueError(
+                f"tensor destination has {destination.size} bytes, expected at least {expected_bytes}"
+            )
+        chunks = payload_spec["chunks"]
+        if not self._read_payload_range_into_raw_destination(
+            chunks, destination.ptr, 0, expected_bytes, allow_get_into=False
+        ):
+            offset = 0
+            for chunk in chunks:
+                data = self._store.get(chunk["key"])
+                size = int(chunk["bytes"])
+                if len(data) != size:
+                    raise RuntimeError(
+                        f"get failed for {chunk['key']}: expected {size} bytes, got {len(data)}"
+                    )
+                ctypes.memmove(destination.ptr + offset, data, size)
+                offset += size
+        _ = destination.owner
+        return None
+
+    def read_payload_range(
+        self,
+        payload_spec: Mapping[str, Any],
+        byte_offset: int,
+        byte_length: int,
+    ) -> bytes:
+        if byte_length == 0:
+            return b""
+        data = bytearray(byte_length)
+        if not self._read_payload_range_into_registered_destination(
+            payload_spec["chunks"], data, byte_offset, byte_length, False
+        ):
+            self._copy_payload_range_into_bytearray(
+                payload_spec["chunks"], data, byte_offset, byte_length
+            )
+        return bytes(data)
+
+    def read_payload_range_into_destination(
+        self,
+        payload_spec: Mapping[str, Any],
+        destination: np.ndarray,
+        byte_offset: int,
+        destination_pre_registered: bool,
+    ) -> None:
+        chunks = payload_spec["chunks"]
+        byte_length = _buffer_nbytes(destination)
+        if not self._read_payload_range_into_registered_destination(
+            chunks,
+            destination,
+            byte_offset,
+            byte_length,
+            destination_pre_registered=destination_pre_registered,
+        ):
+            self._copy_payload_range_into_destination(
+                chunks, destination, byte_offset, byte_length
+            )
+
+    def read_payload_range_into_raw_destination(
+        self,
+        payload_spec: Mapping[str, Any],
+        destination_ptr: int,
+        destination_offset: int,
+        byte_offset: int,
+        byte_length: int,
+    ) -> bool:
+        if byte_length == 0:
+            return True
+        return self._read_payload_range_into_raw_destination(
+            payload_spec["chunks"],
+            destination_ptr,
+            byte_offset,
+            byte_length,
+            allow_get_into=False,
+            destination_offset=destination_offset,
+        )
+
+    def read_payload_ranges_into_array(
+        self,
+        payload_spec: Mapping[str, Any],
+        destination: np.ndarray,
+        ranges: Sequence[tuple[int, int, int]],
+    ) -> None:
+        if self._read_payload_ranges_into_registered_destination(
+            payload_spec["chunks"], destination, ranges
+        ):
+            return
+        self._copy_payload_ranges_into_destination(
+            payload_spec["chunks"], destination.view(np.uint8).reshape(-1), ranges
+        )
+
+    def read_payload_ranges_into_bytearray(
+        self,
+        payload_spec: Mapping[str, Any],
+        destination: bytearray,
+        ranges: Sequence[tuple[int, int, int]],
+    ) -> None:
+        if self._read_payload_ranges_into_registered_destination(
+            payload_spec["chunks"], destination, ranges
+        ):
+            return
+        self._copy_payload_ranges_into_destination(
+            payload_spec["chunks"], destination, ranges
+        )
+
+    def read_payload_ranges_into_raw_destination(
+        self,
+        payload_spec: Mapping[str, Any],
+        destination_ptr: int,
+        ranges: Sequence[tuple[int, int, int]],
+    ) -> bool:
+        return self._read_payload_ranges_into_raw_destination(
+            payload_spec["chunks"], destination_ptr, ranges
+        )
+
+    def put_tensor_object_buffer(
+        self,
+        key: str,
+        value: _TensorObjectBufferPayload,
+    ) -> int:
+        put_tensor_from = self._put_tensor_from
+        if not callable(put_tensor_from):
+            raise RuntimeError("put_tensor_from is unavailable")
+        _check_status(
+            put_tensor_from(key, value.ptr, value.size), "put_tensor_from", key
+        )
+        _ = value.owner
+        return value.size
+
+    def put_tensor_payload_direct(
+        self,
+        key: str,
+        value: _TensorPayload,
+    ) -> dict[str, Any] | None:
+        put_tensor = getattr(self._store, "put_tensor", None)
+        if not callable(put_tensor):
+            return None
+        tensor = value.tensor
+        nbytes = int(getattr(tensor, "nbytes", 0))
+        metadata_bytes = _tensor_metadata_size()
+        total_bytes = metadata_bytes + nbytes
+        _check_status(put_tensor(key, tensor), "put_tensor", key)
+        return {
+            "key": key,
+            "kind": "tensor",
+            "bytes": total_bytes,
+            "dtype": str(getattr(tensor, "dtype", "")),
+            "shape": list(getattr(tensor, "shape", ())),
+            "chunks": [{"key": key, "bytes": total_bytes}],
+            "metadata_bytes": metadata_bytes,
+        }
+
+    def put_tensor_payload_from_pool(
+        self,
+        key: str,
+        value: _TensorPayload,
+    ) -> int:
+        if self._buffer_pool is None:
+            raise RuntimeError("structured tensor zero-copy requires a BufferPool")
+        put_tensor_from = self._put_tensor_from
+        if not callable(put_tensor_from):
+            raise RuntimeError("put_tensor_from is unavailable")
+        metadata, data_ptr, tensor_nbytes, owner = _tensor_payload_parts(value)
+        total_bytes = len(metadata) + tensor_nbytes
+        lease = self._buffer_pool.acquire(total_bytes)
+        view = None
+        try:
+            view = lease.buffer
+            view[: len(metadata)] = metadata
+            if tensor_nbytes:
+                ctypes.memmove(lease.ptr + len(metadata), data_ptr, tensor_nbytes)
+            view.release()
+            view = None
+            _check_status(
+                put_tensor_from(key, lease.ptr, total_bytes), "put_tensor_from", key
+            )
+            _ = owner
+            return total_bytes
+        finally:
+            if view is not None:
+                view.release()
+            lease.release()
+
+    def batch_put_chunks_from(
+        self,
+        chunk_keys: Sequence[str],
+        chunks: Sequence[memoryview],
+        pre_registered: bool,
+    ) -> None:
+        batch_put_from = self._batch_put_from
+        if not callable(batch_put_from):
+            raise RuntimeError("batch_put_from is unavailable")
+        if not chunk_keys:
+            return
+        prepared_chunks = [_prepare_chunk_source_buffer(chunk) for chunk in chunks]
+        buffer_ptrs = [ptr for _owner, ptr, _size in prepared_chunks]
+        sizes = [size for _owner, _ptr, size in prepared_chunks]
+        registered_ptrs = self._register_buffers(
+            buffer_ptrs, sizes, pre_registered, "bundle source payload"
+        )
+        try:
+            results = batch_put_from(list(chunk_keys), buffer_ptrs, sizes)
+            if len(results) != len(chunk_keys):
+                raise RuntimeError(
+                    f"batch_put_from returned {len(results)} results for {len(chunk_keys)} chunks"
+                )
+            for chunk_key, status in zip(chunk_keys, results):
+                _check_status(status, "batch_put_from", chunk_key)
+        except Exception:
+            _cleanup_keys(self._store, chunk_keys, strict=False)
+            raise
+        finally:
+            self._unregister_buffers(registered_ptrs, "bundle source payload")
+
+    def _put_chunks_direct(
+        self,
+        chunk_keys: Sequence[str],
+        chunks: Sequence[memoryview],
+    ) -> list[str]:
+        written_keys: list[str] = []
+        try:
+            for chunk_key, chunk in zip(chunk_keys, chunks):
+                _check_status(self._store.put(chunk_key, chunk), "put", chunk_key)
+                written_keys.append(chunk_key)
+        except Exception:
+            _cleanup_keys(self._store, written_keys, strict=False)
+            raise
+        return list(chunk_keys)
+
+    def _put_chunks_parallel(
+        self,
+        chunk_keys: list[str],
+        chunks: list[memoryview],
+        max_inflight_put: int,
+        pre_registered: bool,
+    ) -> list[str]:
+        groups = self._group_chunk_ranges(chunk_keys, chunks, max_inflight_put)
+        futures: list[Future[None]] = []
+        try:
+            with ThreadPoolExecutor(
+                max_workers=min(max_inflight_put, len(groups))
+            ) as executor:
+                futures = [
+                    executor.submit(
+                        self.batch_put_chunks_from,
+                        group_keys,
+                        group_chunks,
+                        pre_registered,
+                    )
+                    for group_keys, group_chunks in groups
+                ]
+                for future in as_completed(futures):
+                    future.result()
+        except Exception:
+            for future in futures:
+                future.cancel()
+            for future in futures:
+                if future.done() and not future.cancelled():
+                    try:
+                        future.result()
+                    except Exception:
+                        pass
+            _cleanup_keys(self._store, chunk_keys, strict=False)
+            raise
+        return chunk_keys
+
+    def _group_chunk_ranges(
+        self,
+        chunk_keys: Sequence[str],
+        chunks: Sequence[memoryview],
+        max_inflight_put: int,
+    ) -> list[tuple[list[str], list[memoryview]]]:
+        if not chunk_keys:
+            return []
+        group_count = max(1, min(max_inflight_put, len(chunks)))
+        group_size = (len(chunks) + group_count - 1) // group_count
+        return [
+            (
+                list(chunk_keys[start : start + group_size]),
+                list(chunks[start : start + group_size]),
+            )
+            for start in range(0, len(chunks), group_size)
+        ]
+
+    def _read_chunks_with_batch_get_into(
+        self,
+        chunks: Sequence[Mapping[str, Any]],
+        offsets: Sequence[int],
+        destination: bytearray | np.ndarray,
+    ) -> bool:
+        batch_get_into = self._batch_get_into
+        if not callable(batch_get_into) or not self._has_buffer_registration_support():
+            return False
+        with self._registered_buffer(destination, "bundle payload") as base_ptr:
+            keys = [chunk["key"] for chunk in chunks]
+            ptrs = [base_ptr + offset for offset in offsets]
+            sizes = [int(chunk["bytes"]) for chunk in chunks]
+            read_sizes = batch_get_into(keys, ptrs, sizes)
+            if len(read_sizes) != len(keys):
+                raise RuntimeError(
+                    f"batch_get_into returned {len(read_sizes)} results for {len(keys)} chunks"
+                )
+            for key, expected_size, actual_size in zip(keys, sizes, read_sizes):
+                if actual_size != expected_size:
+                    raise RuntimeError(
+                        f"batch_get_into failed for {key}: expected {expected_size}, got {actual_size}"
+                    )
+        return True
+
+    def _read_chunk_with_get(
+        self,
+        chunk: Mapping[str, Any],
+        destination: bytearray | np.ndarray,
+        offset: int,
+    ) -> None:
+        chunk_bytes = int(chunk["bytes"])
+        if chunk_bytes == 0:
+            return
+        data = self._store.get(chunk["key"])
+        if len(data) != chunk_bytes:
+            raise RuntimeError(
+                f"get failed for {chunk['key']}: expected {chunk_bytes} bytes, got {len(data)}"
+            )
+        destination[offset : offset + chunk_bytes] = data
+
+    def _read_payload_range_into_registered_destination(
+        self,
+        chunks: Sequence[Mapping[str, Any]],
+        destination: np.ndarray,
+        byte_offset: int,
+        byte_length: int,
+        destination_pre_registered: bool,
+    ) -> bool:
+        if not self._has_buffer_registration_support():
+            return False
+        with self._registered_buffer(
+            destination,
+            "structured ndarray payload",
+            pre_registered=destination_pre_registered,
+        ) as base_ptr:
+            return self._read_payload_range_into_raw_destination(
+                chunks, base_ptr, byte_offset, byte_length
+            )
+
+    def _read_payload_ranges_into_registered_destination(
+        self,
+        chunks: Sequence[Mapping[str, Any]],
+        destination: bytearray | np.ndarray,
+        ranges: Sequence[tuple[int, int, int]],
+    ) -> bool:
+        if not self._has_buffer_registration_support():
+            return False
+        with self._registered_buffer(destination, "structured ranged payload") as base_ptr:
+            return self._read_payload_ranges_into_raw_destination(
+                chunks, base_ptr, ranges
+            )
+
+    def _read_payload_ranges_into_raw_destination(
+        self,
+        chunks: Sequence[Mapping[str, Any]],
+        base_ptr: int,
+        ranges: Sequence[tuple[int, int, int]],
+    ) -> bool:
+        get_into_ranges = self._get_into_ranges
+        if not callable(get_into_ranges):
+            return False
+        fragments = []
+        for source_offset, destination_offset, byte_length in ranges:
+            fragments.extend(
+                _payload_range_fragments(chunks, source_offset, byte_length, destination_offset)
+            )
+        if not fragments:
+            return True
+        keys = [
+            key
+            for key, _chunk_size, _destination_offset, _source_offset, _size in fragments
+        ]
+        dst_offsets = [
+            [destination_offset]
+            for _key, _chunk_size, destination_offset, _source_offset, _size in fragments
+        ]
+        src_offsets = [
+            [source_offset]
+            for _key, _chunk_size, _destination_offset, source_offset, _size in fragments
+        ]
+        sizes = [
+            [size]
+            for _key, _chunk_size, _destination_offset, _source_offset, size in fragments
+        ]
+        results = get_into_ranges(
+            [base_ptr], [keys], [dst_offsets], [src_offsets], [sizes]
+        )
+        if len(results) != 1 or len(results[0]) != len(keys):
+            raise RuntimeError(
+                f"get_into_ranges returned invalid ranged result shape for {len(keys)} chunks"
+            )
+        for key, expected_sizes, actual_sizes in zip(keys, sizes, results[0]):
+            if len(actual_sizes) != len(expected_sizes):
+                raise RuntimeError(
+                    f"get_into_ranges returned invalid ranged fragment count for {key}"
+                )
+            for expected_size, actual_size in zip(expected_sizes, actual_sizes):
+                if actual_size != expected_size:
+                    raise RuntimeError(
+                        f"get_into_ranges failed for {key}: expected {expected_size}, got {actual_size}"
+                    )
+        return True
+
+    def _read_payload_range_into_raw_destination(
+        self,
+        chunks: Sequence[Mapping[str, Any]],
+        base_ptr: int,
+        byte_offset: int,
+        byte_length: int,
+        allow_get_into: bool = True,
+        destination_offset: int = 0,
+    ) -> bool:
+        get_into = self._get_into
+        if (
+            allow_get_into
+            and len(chunks) == 1
+            and byte_offset == 0
+            and callable(get_into)
+        ):
+            expected_size = int(chunks[0]["bytes"])
+            if expected_size == byte_length:
+                read_size = get_into(
+                    chunks[0]["key"], base_ptr + destination_offset, expected_size
+                )
+                if read_size != expected_size:
+                    raise RuntimeError(
+                        f"get_into failed for {chunks[0]['key']}: expected {expected_size}, got {read_size}"
+                    )
+                return True
+        get_into_ranges = self._get_into_ranges
+        if not callable(get_into_ranges):
+            return False
+        fragments = _payload_range_fragments(
+            chunks, byte_offset, byte_length, destination_offset
+        )
+        if not fragments:
+            return True
+        keys = [
+            key
+            for key, _chunk_size, _destination_offset, _source_offset, _size in fragments
+        ]
+        dst_offsets = [
+            [fragment_destination_offset]
+            for _key, _chunk_size, fragment_destination_offset, _source_offset, _size in fragments
+        ]
+        src_offsets = [
+            [source_offset]
+            for _key, _chunk_size, _destination_offset, source_offset, _size in fragments
+        ]
+        sizes = [
+            [size]
+            for _key, _chunk_size, _destination_offset, _source_offset, size in fragments
+        ]
+        results = get_into_ranges(
+            [base_ptr], [keys], [dst_offsets], [src_offsets], [sizes]
+        )
+        if len(results) != 1 or len(results[0]) != len(keys):
+            raise RuntimeError(
+                f"get_into_ranges returned invalid ranged result shape for {len(keys)} chunks"
+            )
+        for key, expected_sizes, actual_sizes in zip(keys, sizes, results[0]):
+            if len(actual_sizes) != len(expected_sizes):
+                raise RuntimeError(
+                    f"get_into_ranges returned invalid ranged fragment count for {key}"
+                )
+            for expected_size, actual_size in zip(expected_sizes, actual_sizes):
+                if actual_size != expected_size:
+                    raise RuntimeError(
+                        f"get_into_ranges failed for {key}: expected {expected_size}, got {actual_size}"
+                    )
+        return True
+
+    def _copy_payload_range_into_destination(
+        self,
+        chunks: Sequence[Mapping[str, Any]],
+        destination: np.ndarray,
+        byte_offset: int,
+        byte_length: int,
+    ) -> None:
+        data = bytearray(byte_length)
+        self._copy_payload_range_into_bytearray(chunks, data, byte_offset, byte_length)
+        destination[:byte_length] = np.frombuffer(data, dtype=np.uint8)
+
+    def _copy_payload_ranges_into_destination(
+        self,
+        chunks: Sequence[Mapping[str, Any]],
+        destination: bytearray | np.ndarray,
+        ranges: Sequence[tuple[int, int, int]],
+    ) -> None:
+        for source_offset, destination_offset, byte_length in ranges:
+            for (
+                key,
+                chunk_size,
+                fragment_destination_offset,
+                fragment_source_offset,
+                size,
+            ) in _payload_range_fragments(
+                chunks, source_offset, byte_length, destination_offset
+            ):
+                data = self._store.get(key)
+                if len(data) != chunk_size:
+                    raise RuntimeError(
+                        f"get failed for {key}: expected {chunk_size} bytes, got {len(data)}"
+                    )
+                destination[
+                    fragment_destination_offset : fragment_destination_offset + size
+                ] = data[fragment_source_offset : fragment_source_offset + size]
+
+    def _copy_payload_range_into_bytearray(
+        self,
+        chunks: Sequence[Mapping[str, Any]],
+        destination: bytearray,
+        byte_offset: int,
+        byte_length: int,
+    ) -> None:
+        for (
+            key,
+            chunk_size,
+            destination_offset,
+            source_offset,
+            size,
+        ) in _payload_range_fragments(chunks, byte_offset, byte_length):
+            data = self._store.get(key)
+            if len(data) != chunk_size:
+                raise RuntimeError(
+                    f"get failed for {key}: expected {chunk_size} bytes, got {len(data)}"
+                )
+            destination[destination_offset : destination_offset + size] = data[
+                source_offset : source_offset + size
+            ]
+
+    def _register_buffers(
+        self,
+        buffer_ptrs: Sequence[int],
+        sizes: Sequence[int],
+        pre_registered: bool,
+        label: str,
+    ) -> list[int]:
+        register_buffer = self._register_buffer
+        unregister_buffer = self._unregister_buffer
+        if pre_registered:
+            return []
+        if not (callable(register_buffer) and callable(unregister_buffer)):
+            raise RuntimeError(f"register_buffer APIs are unavailable for {label}")
+        registered_ptrs: list[int] = []
+        try:
+            for ptr, size in zip(buffer_ptrs, sizes):
+                if size == 0:
+                    continue
+                register_status = register_buffer(ptr, size)
+                if register_status == 0:
+                    registered_ptrs.append(ptr)
+                    continue
+                if not _is_duplicate_buffer_registration(register_status):
+                    _check_status(register_status, "register_buffer", label)
+        except Exception:
+            self._unregister_buffers(registered_ptrs, label)
+            raise
+        return registered_ptrs
+
+    def _unregister_buffers(self, registered_ptrs: Sequence[int], label: str) -> None:
+        unregister_buffer = self._unregister_buffer
+        if not callable(unregister_buffer):
+            raise RuntimeError(f"unregister_buffer APIs are unavailable for {label}")
+        for ptr in reversed(registered_ptrs):
+            _check_status(unregister_buffer(ptr), "unregister_buffer", label)
+
+    @contextmanager
+    def _registered_buffer(
+        self,
+        destination: bytearray | np.ndarray,
+        label: str,
+        pre_registered: bool = False,
+    ) -> Iterator[int]:
+        base_ptr = _buffer_ptr(destination)
+        if pre_registered:
+            yield base_ptr
+            return
+        registered_ptrs = self._register_buffers(
+            [base_ptr], [_buffer_nbytes(destination)], False, label
+        )
+        if not registered_ptrs:
+            yield base_ptr
+            return
+        succeeded = False
+        try:
+            yield base_ptr
+            succeeded = True
+        finally:
+            try:
+                self._unregister_buffers(registered_ptrs, label)
+            except Exception:
+                if succeeded:
+                    raise
+
+    def _resolve_put_mode(
+        self,
+        chunks: Sequence[memoryview],
+        transfer_policy: BundleTransferPolicy,
+    ) -> Literal["batch", "parallel"]:
+        if transfer_policy.put_mode == "parallel":
+            return "parallel"
+        if transfer_policy.put_mode == "batch":
+            return "batch"
+        if transfer_policy.max_inflight_put <= 1:
+            return "batch"
+        if len(chunks) < AUTO_PARALLEL_MIN_CHUNKS:
+            return "batch"
+        if sum(len(chunk) for chunk in chunks) < AUTO_PARALLEL_MIN_BYTES:
+            return "batch"
+        if min(transfer_policy.max_inflight_put, len(chunks)) < 2:
+            return "batch"
+        return "parallel"
+
+    def _has_batch_put_support(self) -> bool:
+        return (
+            callable(self._batch_put_from) and self._has_buffer_registration_support()
+        )
+
+    def _has_buffer_registration_support(self) -> bool:
+        return callable(self._register_buffer) and callable(self._unregister_buffer)
+
+
+def _resolve_ndarray_destination(
+    name: str,
+    destination: Any,
+    dtype: np.dtype[Any],
+    shape: tuple[int, ...],
+) -> np.ndarray:
+    if destination is None:
+        return np.empty(shape, dtype=dtype)
+    if isinstance(destination, _RawDestinationBuffer):
+        nbytes = int(np.prod(shape, dtype=np.int64)) * dtype.itemsize
+        if destination.size < nbytes:
+            raise ValueError(
+                f"raw destination has {destination.size} bytes, expected at least {nbytes}"
+            )
+        _ = destination.owner
+        return np.ctypeslib.as_array(
+            (ctypes.c_uint8 * nbytes).from_address(destination.ptr)
+        ).view(dtype).reshape(shape)
+    if not isinstance(destination, np.ndarray):
+        raise TypeError(
+            f"structured ndarray field {name} destination must be a numpy.ndarray or raw_destination"
+        )
+    if destination.dtype != dtype:
+        raise ValueError(
+            f"structured ndarray field {name} destination dtype mismatch: expected {dtype.str}, got {destination.dtype.str}"
+        )
+    if tuple(destination.shape) != shape:
+        raise ValueError(
+            f"structured ndarray field {name} destination shape mismatch: expected {shape}, got {tuple(destination.shape)}"
+        )
+    if not destination.flags["C_CONTIGUOUS"]:
+        raise ValueError(
+            f"structured ndarray field {name} destination must be C-contiguous"
+        )
+    if not destination.flags["WRITEABLE"]:
+        raise ValueError(
+            f"structured ndarray field {name} destination must be writeable"
+        )
+    return destination
+
+
+def _resolve_ndarray_read_plan(
+    shape: tuple[int, ...],
+    dtype: np.dtype[Any],
+    member_slice: StructuredMemberSlice | None,
+) -> _NdarrayReadPlan:
+    if member_slice is None:
+        return _NdarrayReadPlan(
+            dtype=dtype,
+            full_shape=shape,
+            output_shape=shape,
+            byte_offset=0,
+            byte_length=int(np.prod(shape, dtype=np.int64)) * dtype.itemsize
+            if shape
+            else dtype.itemsize,
+            step=1,
+            cover_row_count=shape[0] if shape else 1,
+        )
+    if member_slice.axis != 0:
+        raise ValueError("structured ndarray slicing currently supports axis=0 only")
+    if not shape:
+        raise ValueError("structured ndarray slicing requires at least one dimension")
+    total_rows = shape[0]
+    start, end, step = _normalized_member_slice(member_slice, total_rows)
+    row_width = dtype.itemsize * int(np.prod(shape[1:], dtype=np.int64))
+    if start >= end:
+        return _NdarrayReadPlan(
+            dtype=dtype,
+            full_shape=shape,
+            output_shape=(0, *shape[1:]),
+            byte_offset=start * row_width,
+            byte_length=0,
+            step=step,
+            cover_row_count=0,
+        )
+    selected_count = 1 + (end - 1 - start) // step
+    cover_start = start
+    cover_end = start + (selected_count - 1) * step + 1
+    return _NdarrayReadPlan(
+        dtype=dtype,
+        full_shape=shape,
+        output_shape=(selected_count, *shape[1:]),
+        byte_offset=cover_start * row_width,
+        byte_length=(cover_end - cover_start) * row_width,
+        step=step,
+        cover_row_count=cover_end - cover_start,
+    )
+
+
+def _normalized_member_slice(
+    member_slice: StructuredMemberSlice, total_rows: int
+) -> tuple[int, int, int]:
+    if member_slice.step <= 0:
+        raise ValueError("structured ndarray slicing step must be positive")
+    return slice(member_slice.start, member_slice.end, member_slice.step).indices(
+        total_rows
+    )
+
+
+def _bytes_view(value: Any, name: str) -> memoryview:
+    try:
+        view = memoryview(value)
+    except TypeError as error:
+        raise TypeError(
+            f"{name} must be bytes-like, got {type(value).__name__}"
+        ) from error
+    if not view.contiguous:
+        view = memoryview(bytes(view))
+    return view.cast("B")
+
+
+def _prepare_chunk_source_buffer(chunk: memoryview) -> tuple[Any, int, int]:
+    if len(chunk) == 0:
+        copied = ctypes.create_string_buffer(0)
+        return copied, ctypes.addressof(copied), 0
+    if chunk.c_contiguous and not chunk.readonly:
+        return chunk, ctypes.addressof(ctypes.c_char.from_buffer(chunk)), len(chunk)
+    copied = ctypes.create_string_buffer(bytes(chunk))
+    return copied, ctypes.addressof(copied), len(chunk)
+
+
+def _buffer_ptr(value: bytearray | np.ndarray) -> int:
+    return ctypes.addressof(ctypes.c_char.from_buffer(value))
+
+
+def _buffer_nbytes(value: bytearray | np.ndarray) -> int:
+    if isinstance(value, np.ndarray):
+        return int(value.nbytes)
+    return len(value)
+
+
+def _split_view(view: memoryview, chunk_bytes: int) -> list[memoryview]:
+    if len(view) == 0:
+        return [view]
+    return [
+        view[start : start + chunk_bytes] for start in range(0, len(view), chunk_bytes)
+    ]
+
+
+def _chunk_offsets(chunks: Sequence[Mapping[str, Any]]) -> list[int]:
+    offsets = [0]
+    for chunk in chunks[:-1]:
+        offsets.append(offsets[-1] + int(chunk["bytes"]))
+    return offsets
+
+
+def _payload_range_fragments(
+    chunks: Sequence[Mapping[str, Any]],
+    byte_offset: int,
+    byte_length: int,
+    destination_offset: int = 0,
+) -> list[tuple[str, int, int, int, int]]:
+    fragments: list[tuple[str, int, int, int, int]] = []
+    read_end = byte_offset + byte_length
+    chunk_offset = 0
+    for chunk in chunks:
+        chunk_size = int(chunk["bytes"])
+        chunk_end = chunk_offset + chunk_size
+        overlap_start = max(byte_offset, chunk_offset)
+        overlap_end = min(read_end, chunk_end)
+        if overlap_start < overlap_end:
+            fragments.append(
+                (
+                    chunk["key"],
+                    chunk_size,
+                    destination_offset + overlap_start - byte_offset,
+                    overlap_start - chunk_offset,
+                    overlap_end - overlap_start,
+                )
+            )
+        chunk_offset = chunk_end
+    return fragments
+
+
+def _validate_chunk_bytes(chunk_bytes: int) -> int:
+    if chunk_bytes <= 0:
+        raise ValueError("chunk_bytes must be positive")
+    return chunk_bytes
+
+
+def _normalize_key_prefix(key_prefix: str) -> str:
+    prefix = key_prefix.strip("/")
+    if not prefix or any(ord(char) < 32 for char in prefix):
+        raise ValueError("key_prefix must be a non-empty key prefix")
+    return prefix
+
+
+def _validate_key_segment(value: str, name: str) -> None:
+    if (
+        not isinstance(value, str)
+        or not value
+        or "/" in value
+        or any(ord(char) < 32 for char in value)
+    ):
+        raise ValueError(f"invalid bundle {name}: {value!r}")
+
+
+def _encode_manifest(manifest: Mapping[str, Any]) -> bytes:
+    return _encode_json_dict(manifest, "bundle manifest")
+
+
+def _decode_manifest(payload: bytes) -> dict[str, Any]:
+    return _decode_json_dict(payload, "bundle manifest")
+
+
+def _normalize_structured_metadata(metadata: Mapping[str, Any]) -> dict[str, Any]:
+    if not isinstance(metadata, Mapping):
+        raise TypeError("structured metadata must be a mapping")
+    normalized = dict(metadata)
+    if STRUCTURED_FIELD_SPECS_KEY in normalized:
+        raise ValueError(
+            f"structured metadata key {STRUCTURED_FIELD_SPECS_KEY!r} is reserved"
+        )
+    if not isinstance(normalized.get("layout", "structured"), str):
+        raise ValueError("structured metadata layout must be a string")
+    normalized.setdefault("layout", "structured")
+    return normalized
+
+
+def _encode_structured_fields(
+    metadata: Mapping[str, Any],
+    fields: Mapping[str, Any],
+) -> tuple[dict[str, Any], dict[str, Any]]:
+    normalized_metadata = _normalize_structured_metadata(metadata)
+    field_specs: dict[str, Any] = {}
+    encoded_fields: dict[str, Any] = {}
+    for name, value in fields.items():
+        _validate_key_segment(name, "buffer name")
+        spec, encoded_value = _encode_structured_field(value)
+        field_specs[name] = spec
+        encoded_fields[name] = encoded_value
+    if field_specs:
+        normalized_metadata[STRUCTURED_FIELD_SPECS_KEY] = field_specs
+    return normalized_metadata, encoded_fields
+
+
+def tensor_object_buffer(
+    ptr: int, size: int, owner: Any = None, batch_size: int | None = None
+) -> _TensorObjectBufferPayload:
+    return _TensorObjectBufferPayload(
+        ptr=int(ptr), size=int(size), owner=owner, batch_size=batch_size
+    )
+
+
+def raw_destination(
+    ptr: int, size: int, owner: Any = None, *, pre_registered: bool = False
+) -> _RawDestinationBuffer:
+    return _RawDestinationBuffer(
+        ptr=int(ptr), size=int(size), owner=owner, pre_registered=pre_registered
+    )
+
+
+def _encode_structured_field(value: Any) -> tuple[dict[str, Any], Any]:
+    if isinstance(value, _TensorObjectBufferPayload):
+        return {"encoding": "torch_tensor"}, value
+    if _torch is not None and isinstance(value, _torch.Tensor):
+        return _encode_torch_tensor_field(value)
+    if isinstance(value, np.ndarray):
+        array = np.ascontiguousarray(value)
+        return {
+            "encoding": "ndarray",
+            "dtype": array.dtype.str,
+            "shape": list(array.shape),
+        }, array.view(np.uint8).reshape(-1)
+    return {"encoding": "bytes"}, value
+
+
+def _encode_torch_tensor_field(value: Any) -> tuple[dict[str, Any], Any]:
+    return {
+        "encoding": "torch_tensor",
+        "dtype": str(value.dtype),
+        "shape": list(value.shape),
+        "element_size": int(value.element_size()),
+    }, _TensorPayload(tensor=value)
+
+
+def _tensor_payload_parts(value: _TensorPayload) -> tuple[bytes, int, int, Any]:
+    metadata, data_ptr, tensor_nbytes, owner = _tensor_codec_helper(
+        "_serialize_tensor"
+    )(value.tensor)
+    return bytes(metadata), int(data_ptr), int(tensor_nbytes), owner
+
+
+def _has_tensor_codec_helpers() -> bool:
+    return (
+        _mooncake_store is not None
+        and callable(getattr(_mooncake_store, "_serialize_tensor", None))
+        and callable(getattr(_mooncake_store, "_deserialize_tensor", None))
+    )
+
+
+def _tensor_metadata_size() -> int:
+    helper = (
+        None
+        if _mooncake_store is None
+        else getattr(_mooncake_store, "_tensor_metadata_size", None)
+    )
+    if callable(helper):
+        return int(helper())
+    return 0
+
+
+def _torch_save_payload_bytes(value: Any) -> bytes:
+    buffer = io.BytesIO()
+    _torch.save(value, buffer)
+    return buffer.getvalue()
+
+
+def _deserialize_torch_save_payload(payload: bytes) -> Any:
+    return _torch.load(io.BytesIO(payload), weights_only=True)
+
+
+def _slice_tensor_metadata(
+    metadata: bytes, shape: Sequence[int], data_bytes: int
+) -> bytes:
+    patched = bytearray(metadata)
+    ctypes.c_uint64.from_buffer(patched, 32).value = int(data_bytes)
+    global_shape_offset = 40
+    local_shape_offset = global_shape_offset + 8 * 8
+    for index, dim in enumerate(shape):
+        ctypes.c_int64.from_buffer(
+            patched, global_shape_offset + index * 8
+        ).value = int(dim)
+        ctypes.c_int64.from_buffer(patched, local_shape_offset + index * 8).value = int(
+            dim
+        )
+    return bytes(patched)
+
+
+def _tensor_payload_bytes(value: _TensorPayload) -> tuple[bytes, int]:
+    metadata, data_ptr, tensor_nbytes, _owner = _tensor_payload_parts(value)
+    if tensor_nbytes == 0:
+        return metadata, len(metadata)
+    return metadata + ctypes.string_at(data_ptr, tensor_nbytes), len(metadata)
+
+
+def _deserialize_tensor_payload(payload: bytes) -> Any:
+    return _tensor_codec_helper("_deserialize_tensor")(payload)
+
+
+def _tensor_codec_helper(name: str) -> Any:
+    helper = None if _mooncake_store is None else getattr(_mooncake_store, name, None)
+    if not callable(helper):
+        raise RuntimeError(
+            "mooncake.store tensor serialization helpers are required for structured tensor fields"
+        )
+    return helper
+
+
+def _validate_pre_registered_structured_buffers(
+    original_buffers: Mapping[str, Any],
+    encoded_buffers: Mapping[str, Any],
+    pre_registered_buffers: Optional[Mapping[str, bool]],
+) -> None:
+    if not pre_registered_buffers:
+        return
+    for name, pre_registered in pre_registered_buffers.items():
+        if not pre_registered:
+            continue
+        if name not in original_buffers or name not in encoded_buffers:
+            raise ValueError(f"unknown pre-registered structured buffer: {name}")
+        if not _is_same_writable_buffer(original_buffers[name], encoded_buffers[name]):
+            raise ValueError(
+                f"pre-registered structured buffer {name} must be the same writable contiguous buffer used for transfer"
+            )
+
+
+def _is_same_writable_buffer(original: Any, encoded: Any) -> bool:
+    try:
+        original_view = memoryview(original)
+        encoded_view = memoryview(encoded)
+    except TypeError:
+        return False
+    if not original_view.c_contiguous or not encoded_view.c_contiguous:
+        return False
+    if original_view.readonly or encoded_view.readonly:
+        return False
+    if original_view.nbytes != encoded_view.nbytes:
+        return False
+    if original_view.nbytes == 0:
+        return True
+    try:
+        original_bytes = original_view.cast("B")
+        encoded_bytes = encoded_view.cast("B")
+        original_ptr = ctypes.addressof(ctypes.c_char.from_buffer(original_bytes))
+        encoded_ptr = ctypes.addressof(ctypes.c_char.from_buffer(encoded_bytes))
+    except (BufferError, TypeError, ValueError):
+        return False
+    return original_ptr == encoded_ptr
+
+
+def _structured_field_specs(metadata: Mapping[str, Any]) -> dict[str, Any]:
+    field_specs = metadata.get(STRUCTURED_FIELD_SPECS_KEY, {})
+    if not isinstance(field_specs, dict):
+        raise ValueError("structured field specs must be a dict")
+    return field_specs
+
+
+def _merge_structured_stage_metadata(
+    old_metadata: Mapping[str, Any], new_metadata: Mapping[str, Any]
+) -> dict[str, Any]:
+    old_dataproto = old_metadata.get("dataproto")
+    new_dataproto = new_metadata.get("dataproto")
+    if old_dataproto != new_dataproto:
+        raise ValueError("DataProto stage metadata mismatch during manifest merge")
+    merged = dict(old_metadata)
+    old_specs = dict(_structured_field_specs(old_metadata))
+    new_specs = dict(_structured_field_specs(new_metadata))
+    collisions = sorted(set(old_specs) & set(new_specs))
+    if collisions:
+        raise ValueError(f"structured members already exist: {collisions}")
+    old_specs.update(new_specs)
+    if old_specs:
+        merged[STRUCTURED_FIELD_SPECS_KEY] = old_specs
+    return merged
+
+
+def _encode_structured_metadata(metadata: Mapping[str, Any]) -> bytes:
+    return _encode_json_dict(metadata, "structured metadata")
+
+
+def _decode_structured_metadata(payload: bytes) -> dict[str, Any]:
+    return _decode_json_dict(payload, "structured metadata")
+
+
+def _encode_json_dict(value: Mapping[str, Any], label: str) -> bytes:
+    try:
+        return json.dumps(value, separators=(",", ":")).encode("utf-8")
+    except TypeError as error:
+        raise TypeError(f"{label} must be JSON-serializable") from error
+
+
+def _decode_json_dict(payload: bytes, label: str) -> dict[str, Any]:
+    value = json.loads(payload.decode("utf-8"))
+    if not isinstance(value, dict):
+        raise ValueError(f"{label} must decode to a dict")
+    return value
+
+
+def _check_status(status: Any, operation: str, key: str) -> None:
+    if status not in (None, 0):
+        raise RuntimeError(f"{operation} failed for {key}: {status}")
+
+
+def _is_duplicate_buffer_registration(status: Any) -> bool:
+    # Mooncake returns -600 when a caller-owned buffer is already registered.
+    return status == -600
+
+
+def _cleanup_keys(store: BundleStore, keys: Sequence[str], strict: bool) -> None:
+    pending_keys = list(dict.fromkeys(keys))
+    batch_remove = getattr(store, "batch_remove", None)
+    if callable(batch_remove) and pending_keys:
+        try:
+            try:
+                results = batch_remove(pending_keys, True)
+            except TypeError:
+                results = batch_remove(pending_keys)
+            if len(results) != len(pending_keys):
+                raise RuntimeError(
+                    f"batch_remove returned {len(results)} results for {len(pending_keys)} keys"
+                )
+            failed_results = [
+                (key, status)
+                for key, status in zip(pending_keys, results)
+                if status not in (None, 0, MISSING_OBJECT_ERROR)
+            ]
+            if not failed_results:
+                return
+            pending_keys = [key for key, _status in failed_results]
+        except Exception:
+            if strict:
+                raise
+
+    retry_errors = []
+    for key in pending_keys:
+        try:
+            status = store.remove(key, True)
+        except KeyError:
+            continue
+        except Exception:
+            if strict:
+                raise
+            continue
+        if status not in (None, 0, MISSING_OBJECT_ERROR):
+            retry_errors.append((key, status))
+    if retry_errors and strict:
+        raise RuntimeError(
+            f"failed to remove {len(retry_errors)} Mooncake keys: {retry_errors[:3]}"
+        )
+
+
+# ---------------------------------------------------------------------------
+# Codec inference and recursive structure expansion
+# ---------------------------------------------------------------------------
+
+_INFER_MAX_STRUCT_KEYS = 64
+_INFER_MAX_LIST_LEN = 8
+_INFER_MAX_SAMPLE_ROWS = 128
+_INFER_MAX_JSON_BYTES = 1 << 20
+_INFER_MAX_DEPTH = 32
+
+try:
+    import torch as _torch
+except Exception:  # pragma: no cover
+    _torch = None  # type: ignore[assignment]
+
+
+@dataclass
+class _CodecDecision:
+    accepted: bool
+    codec: str
+    reason: str
+    normalized_type: str
+    metadata: dict[str, Any] = field(default_factory=dict)
+
+
+@dataclass
+class _InferredLeaf:
+    path: str
+    values: list[Any]
+    decision: _CodecDecision
+
+
+@dataclass
+class _EncodedStructuredLeaf:
+    codec: str
+    rows: int
+    payload: dict[str, Any]
+    metadata: dict[str, Any]
+
+
+class _Missing:
+    """Sentinel for dict keys absent from a row (distinct from value-is-None)."""
+
+    __slots__ = ()
+
+    def __repr__(self) -> str:
+        return "<MISSING>"
+
+
+MISSING = _Missing()
+
+
+@dataclass
+class _InferredNode:
+    path: str
+    node_type: str
+    children: list[Any]
+    lengths: list[int] | None = None
+    row_mask: list[bool] | None = None
+
+
+def _non_null(values: list[Any]) -> list[Any]:
+    return [v for v in values if v is not None and not isinstance(v, _Missing)]
+
+
+def _is_pil_image(value: Any) -> bool:
+    module = getattr(value.__class__, "__module__", None) or ""
+    return module.startswith("PIL.") and hasattr(value, "save")
+
+
+def _is_bytes_like(value: Any) -> bool:
+    return isinstance(value, (bytes, bytearray, memoryview))
+
+
+def _is_media_list(value: Any) -> bool:
+    return isinstance(value, (list, tuple)) and all(
+        _is_pil_image(item) or _is_bytes_like(item) for item in value
+    )
+
+
+def _can_media_list(values: list[Any]) -> _CodecDecision:
+    nn = _non_null(values)
+    if not nn:
+        return _CodecDecision(
+            False, "media_list_ragged", "all rows are null", "media list"
+        )
+    if not all(_is_media_list(v) for v in nn):
+        return _CodecDecision(
+            False, "media_list_ragged", "not all rows are media list", "media list"
+        )
+    if not any(len(v) > 0 for v in nn):
+        return _CodecDecision(
+            False, "media_list_ragged", "all media lists are empty", "media list"
+        )
+    return _CodecDecision(
+        True, "media_list_ragged", "all non-null rows are media list", "media list"
+    )
+
+
+def _check_all(
+    values: list[Any],
+    predicate: Any,
+    codec: str,
+    normalized_type: str,
+) -> _CodecDecision:
+    nn = _non_null(values)
+    if not nn:
+        return _CodecDecision(False, codec, "all rows are null", normalized_type)
+    if not all(predicate(v) for v in nn):
+        return _CodecDecision(
+            False, codec, f"not all rows are {normalized_type}", normalized_type
+        )
+    return _CodecDecision(
+        True, codec, f"all non-null rows are {normalized_type}", normalized_type
+    )
+
+
+def _can_tensor(values: list[Any]) -> _CodecDecision:
+    if _torch is None:
+        return _CodecDecision(
+            False, "ragged_tensor", "torch is not available", "torch.Tensor"
+        )
+    nn = _non_null(values)
+    if not nn:
+        return _CodecDecision(
+            False, "ragged_tensor", "all rows are null", "torch.Tensor"
+        )
+    if not all(isinstance(v, _torch.Tensor) for v in nn):
+        return _CodecDecision(
+            False, "ragged_tensor", "not all rows are Tensor", "torch.Tensor"
+        )
+    dtypes = sorted({str(v.dtype) for v in nn})
+    if len(dtypes) != 1:
+        return _CodecDecision(
+            False, "ragged_tensor", f"mixed dtypes: {dtypes}", "torch.Tensor"
+        )
+    ndims = {v.ndim for v in nn}
+    if len(ndims) != 1:
+        return _CodecDecision(
+            False, "ragged_tensor", f"mixed dimensions: {ndims}", "torch.Tensor"
+        )
+    return _CodecDecision(
+        True,
+        "ragged_tensor",
+        "all non-null rows are Tensor",
+        "torch.Tensor",
+        {"dtype": dtypes[0]},
+    )
+
+
+def _can_numeric_sequence(values: list[Any]) -> _CodecDecision:
+    nn = _non_null(values)
+    if not nn:
+        return _CodecDecision(
+            False, "typed_ragged", "all rows are null", "numeric sequence"
+        )
+    dtypes = []
+    for v in nn:
+        if isinstance(v, np.ndarray):
+            arr = v
+        elif isinstance(v, (list, tuple)):
+            try:
+                arr = np.asarray(v)
+            except (ValueError, TypeError):
+                return _CodecDecision(
+                    False,
+                    "typed_ragged",
+                    "row cannot be converted to ndarray",
+                    "numeric sequence",
+                )
+        else:
+            return _CodecDecision(
+                False, "typed_ragged", "row is not array-like", "numeric sequence"
+            )
+        if arr.dtype == object or not np.issubdtype(arr.dtype, np.number):
+            return _CodecDecision(
+                False,
+                "typed_ragged",
+                f"non-numeric dtype: {arr.dtype}",
+                "numeric sequence",
+            )
+        dtypes.append(arr.dtype)
+    try:
+        dtype = np.result_type(*dtypes)
+    except (TypeError, ValueError, OverflowError):
+        return _CodecDecision(
+            False, "typed_ragged", "cannot determine common dtype", "numeric sequence"
+        )
+    return _CodecDecision(
+        True,
+        "typed_ragged",
+        "all rows promote to common numeric dtype",
+        "numeric sequence",
+        {"dtype": str(dtype)},
+    )
+
+
+def _can_numeric_scalar(values: list[Any]) -> _CodecDecision:
+    nn = _non_null(values)
+    if not nn:
+        return _CodecDecision(False, "ndarray", "all rows are null", "numeric scalar")
+    if not all(isinstance(v, (bool, int, float, np.number)) for v in nn):
+        return _CodecDecision(
+            False, "ndarray", "not all rows are numeric scalar", "numeric scalar"
+        )
+    try:
+        dtype = np.result_type(*nn)
+    except (TypeError, ValueError, OverflowError):
+        return _CodecDecision(
+            False, "ndarray", "cannot determine common dtype", "numeric scalar"
+        )
+    if not np.issubdtype(dtype, np.number) and not np.issubdtype(dtype, np.bool_):
+        return _CodecDecision(
+            False, "ndarray", f"non-numeric dtype: {dtype}", "numeric scalar"
+        )
+    return _CodecDecision(
+        True,
+        "ndarray",
+        "all rows are numeric scalar",
+        "numeric scalar",
+        {"dtype": str(dtype)},
+    )
+
+
+def _can_msgpack(values: list[Any]) -> _CodecDecision:
+    nn = _non_null(values)
+    if not nn:
+        return _CodecDecision(False, "msgpack_ragged", "all rows are null", "msgpack")
+    if not all(isinstance(v, (dict, list, tuple)) for v in nn):
+        return _CodecDecision(
+            False, "msgpack_ragged", "not all rows are structured objects", "msgpack"
+        )
+    sampled_bytes = 0
+    for i, v in enumerate(nn):
+        try:
+            encoded = _msgpack.packb(v, use_bin_type=True, strict_types=True)
+        except (TypeError, ValueError, OverflowError):
+            return _CodecDecision(
+                False, "msgpack_ragged", "serialization failed", "msgpack"
+            )
+        if i < _INFER_MAX_SAMPLE_ROWS:
+            sampled_bytes += len(encoded)
+    if sampled_bytes > _INFER_MAX_JSON_BYTES:
+        return _CodecDecision(
+            False, "msgpack_ragged", "sampled payload too large", "msgpack"
+        )
+    return _CodecDecision(
+        True, "msgpack_ragged", "all rows pass msgpack serialization", "msgpack"
+    )
+
+
+def _can_json(values: list[Any]) -> _CodecDecision:
+    nn = _non_null(values)
+    if not nn:
+        return _CodecDecision(False, "json_ragged", "all rows are null", "json")
+    sampled_bytes = 0
+    for i, v in enumerate(nn):
+        try:
+            encoded = json.dumps(v, ensure_ascii=False, separators=(",", ":")).encode(
+                "utf-8"
+            )
+        except (TypeError, ValueError, OverflowError, RecursionError):
+            return _CodecDecision(False, "json_ragged", "serialization failed", "json")
+        if i < _INFER_MAX_SAMPLE_ROWS:
+            sampled_bytes += len(encoded)
+    if sampled_bytes > _INFER_MAX_JSON_BYTES:
+        return _CodecDecision(False, "json_ragged", "sampled payload too large", "json")
+    return _CodecDecision(
+        True, "json_ragged", "all rows pass JSON serialization", "json"
+    )
+
+
+_CODEC_PREDICATES: tuple[Any, ...] = (
+    _can_tensor,
+    _can_media_list,
+    _can_numeric_sequence,
+    _can_numeric_scalar,
+    lambda v: _check_all(v, _is_bytes_like, "bytes_ragged", "bytes-like"),
+    lambda v: _check_all(v, _is_pil_image, "media_bytes", "media"),
+    lambda v: _check_all(v, lambda x: isinstance(x, str), "utf8_ragged", "str"),
+    _can_msgpack,
+    _can_json,
+)
+
+
+def _choose_leaf_codec(values: list[Any]) -> _CodecDecision:
+    for predicate in _CODEC_PREDICATES:
+        decision = predicate(values)
+        if decision.accepted:
+            return decision
+    return _CodecDecision(
+        False, "pickle_ragged_fallback", "no optimized codec matched", "python object"
+    )
+
+
+def _try_expand_dict(values: list[Any]) -> list[str] | None:
+    nn = _non_null(values)
+    if not nn or not all(isinstance(v, dict) for v in nn):
+        return None
+    keys = {k for v in nn for k in v.keys()}
+    if not all(isinstance(k, str) for k in keys) or len(keys) > _INFER_MAX_STRUCT_KEYS:
+        return None
+    return sorted(keys)
+
+
+def _try_expand_list(values: list[Any]) -> tuple[int, list[int]] | None:
+    nn = _non_null(values)
+    if not nn or not all(isinstance(v, (list, tuple)) for v in nn):
+        return None
+    max_len = max(len(v) for v in nn)
+    if max_len > _INFER_MAX_LIST_LEN:
+        return None
+    if not all(
+        item is None or isinstance(item, (dict, list, tuple)) for v in nn for item in v
+    ):
+        return None
+    lengths = [len(v) if isinstance(v, (list, tuple)) else 0 for v in values]
+    return max_len, lengths
+
+
+def _escape_key(key: str) -> str:
+    return key.replace("\\", "\\\\").replace(".", "\\.").replace("[", "\\[")
+
+
+def infer_structure(
+    path: str,
+    values: list[Any],
+    leaves: list[_InferredLeaf],
+    nodes: list[_InferredNode],
+    *,
+    _depth: int = 0,
+) -> None:
+    """Recursively expand *values* into leaves and interior nodes.
+
+    *leaves* and *nodes* are output accumulators populated during recursion.
+    Absent dict keys are represented as ``MISSING`` in child columns
+    (distinct from ``None`` values).  Each interior node carries a
+    ``row_mask`` that records which rows had a real parent container.
+    """
+    if _depth > _INFER_MAX_DEPTH:
+        raise ValueError(
+            f"infer_structure exceeded max depth {_INFER_MAX_DEPTH} at {path!r}"
+        )
+    dict_keys = _try_expand_dict(values)
+    if dict_keys is not None:
+        row_mask = [isinstance(v, dict) for v in values]
+        nodes.append(_InferredNode(path, "dict", dict_keys, row_mask=row_mask))
+        for key in dict_keys:
+            child = [
+                v.get(key, MISSING) if isinstance(v, dict) else None for v in values
+            ]
+            infer_structure(
+                f"{path}.{_escape_key(key)}", child, leaves, nodes, _depth=_depth + 1
+            )
+        return
+    list_result = _try_expand_list(values)
+    if list_result is not None:
+        max_len, lengths = list_result
+        row_mask = [isinstance(v, (list, tuple)) for v in values]
+        nodes.append(
+            _InferredNode(
+                path, "list", list(range(max_len)), lengths, row_mask=row_mask
+            )
+        )
+        for index in range(max_len):
+            child = [
+                v[index]
+                if isinstance(v, (list, tuple)) and index < len(v)
+                else (MISSING if isinstance(v, (list, tuple)) else None)
+                for v in values
+            ]
+            infer_structure(f"{path}[{index}]", child, leaves, nodes, _depth=_depth + 1)
+        return
+    leaves.append(_InferredLeaf(path, values, _choose_leaf_codec(values)))
diff --git a/mooncake-wheel/pyproject.toml b/mooncake-wheel/pyproject.toml
index 90dcf32b..19253f4a 100644
--- a/mooncake-wheel/pyproject.toml
+++ b/mooncake-wheel/pyproject.toml
@@ -4,13 +4,13 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "mooncake-transfer-engine"
-version = "0.3.10.post2"
+version = "0.3.11.post1"
 description = "Python binding of a Mooncake library using pybind11"
 authors = [
     { name = "Mooncake Authors" }
 ]
 requires-python = ">=3.8"
-dependencies = ["aiohttp", "requests"]
+dependencies = ["aiohttp", "requests", "msgpack"]
 keywords = ["mooncake", "data transfer", "kv cache", "llm inference"]
 classifiers = [
     "Programming Language :: Python :: 3",
@@ -21,7 +21,7 @@ readme = ""
 
 [project.urls]
 Homepage = "https://github.com/kvcache-ai/Mooncake"
-Documentation = "https://github.com/kvcache-ai/Mooncake/tree/main/doc"
+Documentation = "https://kvcache-ai.github.io/Mooncake"
 Source = "https://github.com/kvcache-ai/Mooncake"
 Issues = "https://github.com/kvcache-ai/Mooncake/issues"
 
@@ -30,7 +30,7 @@ mooncake_master = "mooncake.cli:main"
 mooncake_client = "mooncake.cli_client:main"
 transfer_engine_bench = "mooncake.cli_bench:main"
 mooncake_http_metadata_server = "mooncake.http_metadata_server:main"
-mc_store_rest_server = "mooncake.mooncake_store_service:main"
+mc_store_rest_server = "mooncake.mooncake_store_service:sync_main"
 transfer_engine_topology_dump = "mooncake.transfer_engine_topology_dump:main"
 
 [tool.setuptools]
diff --git a/mooncake-wheel/tests/conftest.py b/mooncake-wheel/tests/conftest.py
new file mode 100644
index 00000000..30728561
--- /dev/null
+++ b/mooncake-wheel/tests/conftest.py
@@ -0,0 +1,19 @@
+from __future__ import annotations
+
+import importlib.util
+import sys
+from pathlib import Path
+
+ROOT = Path(__file__).resolve().parents[1]
+SOURCE_PACKAGE = ROOT / "mooncake"
+
+try:
+    import mooncake
+except ModuleNotFoundError:
+    if str(ROOT) not in sys.path:
+        sys.path.insert(0, str(ROOT))
+else:
+    if importlib.util.find_spec("mooncake.structured_object_store") is None:
+        mooncake_path = getattr(mooncake, "__path__", None)
+        if mooncake_path is not None and str(SOURCE_PACKAGE) not in mooncake_path:
+            mooncake_path.append(str(SOURCE_PACKAGE))
diff --git a/mooncake-wheel/tests/test_buffer_pool_native.py b/mooncake-wheel/tests/test_buffer_pool_native.py
new file mode 100644
index 00000000..14067aed
--- /dev/null
+++ b/mooncake-wheel/tests/test_buffer_pool_native.py
@@ -0,0 +1,231 @@
+from __future__ import annotations
+
+import gc
+import threading
+import time
+
+import pytest
+
+from mooncake.buffer_pool import BufferPool
+from mooncake.store import MooncakeDistributedStore
+
+
+def create_store() -> MooncakeDistributedStore:
+    store = MooncakeDistributedStore()
+    rc = store.setup(
+        "localhost",
+        "P2PHANDSHAKE",
+        16 * 1024 * 1024,
+        4 * 1024 * 1024,
+        "tcp",
+        "",
+        "127.0.0.1:50051",
+    )
+    assert rc == 0
+    return store
+
+
+def test_buffer_pool_acquires_released_buffer() -> None:
+    store = create_store()
+    pool = BufferPool(store, min_size_class=4096, alignment=4096)
+
+    lease = pool.acquire(1234)
+    lease.buffer[:4] = b"abcd"
+    assert bytes(lease.buffer[:4]) == b"abcd"
+    lease.release()
+
+    next_lease = pool.acquire(2048)
+    next_lease.release()
+    pool.close()
+
+
+def test_buffer_pool_uses_store_local_buffer_by_default() -> None:
+    store = create_store()
+    pool = BufferPool(store)
+
+    lease = pool.acquire(1024)
+    assert lease.size == 1024
+    lease.release()
+    pool.close()
+
+
+def test_buffer_pool_overflows_when_local_buffer_is_full() -> None:
+    store = create_store()
+    pool = BufferPool(store)
+
+    local_lease = pool.acquire(4 * 1024 * 1024)
+    overflow_lease = pool.acquire(1024, block=False)
+    assert overflow_lease.size == 1024
+
+    overflow_lease.release()
+    local_lease.release()
+    pool.close()
+
+
+def test_buffer_pool_prewarm_and_close() -> None:
+    store = create_store()
+    pool = BufferPool(
+        store,
+        1024 * 1024,
+        min_size_class=4096,
+        alignment=4096,
+        prewarm_size=1024,
+        prewarm_count=2,
+    )
+    lease1 = pool.acquire(1024)
+    lease2 = pool.acquire(1024)
+    lease1.release()
+    lease2.release()
+    pool.close()
+
+
+def test_buffer_pool_uses_local_buffer_alignment() -> None:
+    store = create_store()
+    pool = BufferPool(store, min_size_class=4096, alignment=65536)
+
+    lease = pool.acquire(1024)
+    assert lease.ptr % 64 == 0
+    lease.release()
+    pool.close()
+
+
+def test_buffer_pool_rejects_invalid_alignment() -> None:
+    store = create_store()
+    with pytest.raises(RuntimeError, match="alignment"):
+        BufferPool(store, 1024 * 1024, alignment=12345)
+
+
+@pytest.mark.parametrize("size", [0, 1, 128 * 1024 + 1])
+def test_buffer_pool_supports_arbitrary_sizes(size: int) -> None:
+    store = create_store()
+    pool = BufferPool(
+        store, 1024 * 1024, max_size_class=128 * 1024, alignment=4096
+    )
+
+    lease = pool.acquire(size)
+    assert lease.size == size
+    assert len(lease.buffer) == size
+    lease.release()
+    pool.close()
+
+
+def test_buffer_pool_nonblocking_exhaustion() -> None:
+    store = create_store()
+    pool = BufferPool(
+        store,
+        min_size_class=4096,
+        max_size_class=4096,
+        alignment=4096,
+        max_regions=1,
+        block_on_exhaustion=False,
+    )
+
+    lease = pool.acquire(1)
+    with pytest.raises(RuntimeError, match="exhausted"):
+        pool.acquire(1)
+    lease.release()
+    pool.close()
+
+
+def test_buffer_pool_timeout_allows_late_release() -> None:
+    store = create_store()
+    pool = BufferPool(
+        store,
+        min_size_class=4096,
+        max_size_class=4096,
+        alignment=4096,
+        max_regions=1,
+    )
+    lease = pool.acquire(1)
+    acquired = []
+
+    def acquire_after_release() -> None:
+        acquired.append(pool.acquire(1, timeout=1.0))
+
+    thread = threading.Thread(target=acquire_after_release)
+    thread.start()
+    time.sleep(0.1)
+    lease.release()
+    thread.join(timeout=5)
+
+    assert not thread.is_alive()
+    acquired[0].release()
+    pool.close()
+
+
+def test_buffer_pool_blocking_acquire_releases_gil() -> None:
+    store = create_store()
+    pool = BufferPool(
+        store,
+        min_size_class=4096,
+        max_size_class=4096,
+        alignment=4096,
+        max_regions=1,
+    )
+    lease = pool.acquire(1)
+    acquired = []
+
+    def acquire_after_release() -> None:
+        acquired.append(pool.acquire(1))
+
+    thread = threading.Thread(target=acquire_after_release)
+    thread.start()
+    time.sleep(0.1)
+    lease.release()
+    thread.join(timeout=5)
+
+    assert not thread.is_alive()
+    acquired[0].release()
+    pool.close()
+
+
+def test_buffer_pool_memoryview_keeps_lease_alive() -> None:
+    store = create_store()
+    pool = BufferPool(store, 4096, min_size_class=4096, alignment=4096)
+
+    lease = pool.acquire(4)
+    view = lease.buffer
+    view[:] = b"test"
+
+    assert bytes(view) == b"test"
+    with pytest.raises(RuntimeError, match="exported views"):
+        lease.release()
+    del view
+    gc.collect()
+    lease.release()
+    pool.close()
+
+
+def test_buffer_pool_releases_from_destructor() -> None:
+    store = create_store()
+    pool = BufferPool(store, 4096, min_size_class=4096, alignment=4096)
+
+    lease = pool.acquire(1)
+    del lease
+    gc.collect()
+
+    pool.close()
+
+
+def test_buffer_pool_rejects_close_with_active_lease() -> None:
+    store = create_store()
+    pool = BufferPool(store, 1024 * 1024, min_size_class=4096, alignment=4096)
+
+    lease = pool.acquire(1)
+    with pytest.raises(RuntimeError, match="active leases"):
+        pool.close()
+    lease.release()
+    pool.close()
+
+
+def test_buffer_pool_rejects_huge_size_overflow() -> None:
+    store = create_store()
+    pool = BufferPool(store, 1024 * 1024, min_size_class=4096, alignment=4096)
+
+    with pytest.raises(RuntimeError, match="overflow|capacity"):
+        pool.acquire((1 << 64) - 1)
+    pool.close()
+
+
+if __name__ == "__main__":
+    raise SystemExit(pytest.main([__file__]))
diff --git a/mooncake-wheel/tests/test_distributed_object_store.py b/mooncake-wheel/tests/test_distributed_object_store.py
index 52530d3b..9e5d9e96 100644
--- a/mooncake-wheel/tests/test_distributed_object_store.py
+++ b/mooncake-wheel/tests/test_distributed_object_store.py
@@ -38,6 +38,46 @@ def get_client(store, local_buffer_size_param=None):
     if retcode:
         raise RuntimeError(f"Failed to setup store client. Return code: {retcode}")
 
+
+def get_config_dict(global_segment_size, local_buffer_size):
+    """Build a config dictionary for the MooncakeDistributedStore setup wrapper."""
+    return {
+        "local_hostname": os.getenv("LOCAL_HOSTNAME", "localhost"),
+        "metadata_server": os.getenv(
+            "MC_METADATA_SERVER", "http://127.0.0.1:8080/metadata"
+        ),
+        "global_segment_size": global_segment_size,
+        "local_buffer_size": local_buffer_size,
+        "protocol": os.getenv("PROTOCOL", "tcp"),
+        "rdma_devices": os.getenv("DEVICE_NAME", "ibp6s0"),
+        "master_server_addr": os.getenv("MASTER_SERVER", "127.0.0.1:50051"),
+    }
+
+
+class TestConfigDictSetup(unittest.TestCase):
+    """Test configuration-dictionary setup through the Python store wrapper."""
+
+    def test_human_readable_sizes(self):
+        store = MooncakeDistributedStore()
+        self.addCleanup(store.close)
+
+        retcode = store.setup(get_config_dict("16MB", "16 MB"))
+        self.assertEqual(retcode, 0)
+
+        test_data = b"test_config_dict_human_readable_value"
+        key = f"test_config_dict_human_readable_key_{os.getpid()}"
+
+        self.assertEqual(store.put(key, test_data), 0)
+        self.assertEqual(store.get(key), test_data)
+
+    def test_unsupported_percentage_size(self):
+        store = MooncakeDistributedStore()
+        self.addCleanup(store.close)
+
+        retcode = store.setup(get_config_dict("50%", "16MB"))
+        self.assertNotEqual(retcode, 0)
+
+
 class TestZeroLocalBufferSize(unittest.TestCase):
     """Test class for zero local buffer size scenarios."""
     
diff --git a/mooncake-wheel/tests/test_import_structure.py b/mooncake-wheel/tests/test_import_structure.py
index 2dcd0c5b..e05c980f 100644
--- a/mooncake-wheel/tests/test_import_structure.py
+++ b/mooncake-wheel/tests/test_import_structure.py
@@ -16,12 +16,18 @@ class TestImportStructure(unittest.TestCase):
         # Verify direct access to TransferOpcode
         self.assertIsNotNone(mooncake.engine.TransferOpcode)
 
-        from mooncake.store import MooncakeDistributedStore
+        from mooncake.store import MooncakeDistributedStore, ReplicateConfig
 
         # Just verify we can create instances
         store = MooncakeDistributedStore()
+        config = ReplicateConfig()
+        config.group_ids = ["group-a", ""]
 
         self.assertIsNotNone(store)
+        self.assertEqual(config.group_ids, ["group-a", ""])
+
+        config.group_ids = None
+        self.assertIsNone(config.group_ids)
 
     def test_direct_import(self):
         """Test direct import of specific components."""
@@ -31,5 +37,11 @@ class TestImportStructure(unittest.TestCase):
         self.assertIsNotNone(TransferEngine)
         self.assertIsNotNone(TransferOpcode)
 
+    def test_top_level_buffer_pool_alias(self):
+        """Test top-level buffer-pool exports."""
+        from mooncake import BufferPool, RegisteredBufferPool
+
+        self.assertIs(BufferPool, RegisteredBufferPool)
+
 if __name__ == '__main__':
     unittest.main()
diff --git a/mooncake-wheel/tests/test_mooncake_backend.py b/mooncake-wheel/tests/test_mooncake_backend.py
deleted file mode 100644
index d058a591..00000000
--- a/mooncake-wheel/tests/test_mooncake_backend.py
+++ /dev/null
@@ -1,176 +0,0 @@
-import os
-import time
-import unittest
-import torch
-import torch.distributed as dist
-import torch.multiprocessing as mp
-from mooncake import pg
-
-
-def worker(rank, world_size, results, collective):
-    torch.cuda.set_device(rank)
-    dist.init_process_group(
-        backend="mooncake",
-        rank=rank,
-        world_size=world_size,
-        pg_options=pg.MooncakeBackendOptions(torch.zeros((world_size,), dtype=torch.int32, device="cuda")),
-    )
-
-    if collective == "all_reduce_sum":
-        tensor = torch.tensor([rank + 1], dtype=torch.int32, device="cuda")
-        dist.all_reduce(tensor, op=dist.ReduceOp.SUM)
-        results[rank] = tensor.item()
-    
-    elif collective == "all_reduce_product":
-        tensor = torch.tensor([2], dtype=torch.int32, device="cuda")
-        dist.all_reduce(tensor, op=dist.ReduceOp.PRODUCT)
-        results[rank] = tensor.item()
-
-    elif collective == "all_reduce_min":
-        tensor = torch.tensor([rank + 10], dtype=torch.int32, device="cuda")
-        dist.all_reduce(tensor, op=dist.ReduceOp.MIN)
-        results[rank] = tensor.item()
-
-    elif collective == "all_reduce_max":
-        tensor = torch.tensor([rank + 10], dtype=torch.int32, device="cuda")
-        dist.all_reduce(tensor, op=dist.ReduceOp.MAX)
-        results[rank] = tensor.item()
-
-    elif collective == "all_gather":
-        tensor = torch.tensor([rank], device="cuda")
-        gathered = [torch.zeros_like(tensor) for _ in range(world_size)]
-        dist.all_gather(gathered, tensor)
-        results[rank] = [t.item() for t in gathered]
-
-    elif collective == "barrier":
-        dist.barrier()
-        results[rank] = "ok"
-
-    elif collective == "gather":
-        tensor = torch.tensor([rank], dtype=torch.int32, device="cuda")
-        if rank == 0:
-            gather_list = [torch.zeros_like(tensor) for _ in range(world_size)]
-            dist.gather(tensor, gather_list, dst=0)
-            results[rank] = [t.item() for t in gather_list]
-        else:
-            dist.gather(tensor, dst=0)
-            results[rank] = None
-
-    elif collective == "scatter":
-        tensor = torch.zeros(1, dtype=torch.int32, device="cuda")
-        if rank == 0:
-            scatter_list = [torch.tensor([i], dtype=torch.int32, device="cuda") for i in range(world_size)]
-            dist.scatter(tensor, scatter_list, src=0)
-        else:
-            dist.scatter(tensor, src=0)
-        results[rank] = tensor.item()
-
-    elif collective == "reduce":
-        tensor = torch.tensor([1], dtype=torch.int32, device="cuda")
-        dist.reduce(tensor, dst=0, op=dist.ReduceOp.SUM)
-        if rank == 0:
-            results[rank] = tensor.item()
-        else:
-            results[rank] = None
-
-    else:
-        raise ValueError(f"Unsupported collective: {collective}")
-
-    while len(results) < world_size:
-        time.sleep(1)
-
-    dist.destroy_process_group()
-
-
-class TestMooncakeBackend(unittest.TestCase):
-    def setUp(self):
-        self.world_size = torch.cuda.device_count()
-        os.environ["MASTER_ADDR"] = "127.0.0.1"
-        os.environ["MASTER_PORT"] = "29500"
-
-    def tearDown(self):
-        pass
-
-    def _spawn_and_check(self, collective, expected_fn):
-        mp_manager = mp.Manager()
-        results = mp_manager.dict()
-        mp.spawn(
-            worker,
-            args=(self.world_size, results, collective),
-            nprocs=self.world_size,
-            join=True,
-        )
-
-        expected = expected_fn(self.world_size)
-        for r in range(self.world_size):
-            self.assertEqual(results[r], expected)
-
-    def test_allreduce_sum(self):
-        # Expected sum = 1 + 2 + 3 + 4 = 10
-        self._spawn_and_check("all_reduce_sum", lambda size: sum(range(1, size + 1)))
-    
-    def test_allreduce_product(self):
-        # Expected product(2, 2, 2, ……, 2) = 2 ^ N
-        self._spawn_and_check("all_reduce_product", lambda size: 2 ** size)
-
-    def test_allreduce_min(self):
-        # Expected Min(10, 11, ……) = 10
-        self._spawn_and_check("all_reduce_min", lambda size: 10)
-    
-    def test_allreduce_max(self):
-        # Expected Max(10, 11, ..., 10+N-1) = 10 + N - 1
-        self._spawn_and_check("all_reduce_max", lambda size: 10 + size - 1)
-
-    def test_allgather(self):
-        # Expected gather = [0, 1, 2, 3]
-        self._spawn_and_check("all_gather", lambda size: list(range(size)))
-
-    def test_barrier(self):
-        self._spawn_and_check("barrier", lambda size: "ok")
-
-    def test_gather(self):
-        # Expected gather (Root) = [0, 1, 2, ..., size-1]
-        # Expected gather (Others) = None
-        mp_manager = mp.Manager()
-        results = mp_manager.dict()
-        mp.spawn(
-            worker,
-            args=(self.world_size, results, "gather"),
-            nprocs=self.world_size,
-            join=True,
-        )
-        self.assertEqual(results[0], list(range(self.world_size)))
-        for r in range(1, self.world_size):
-            self.assertIsNone(results[r])
-
-    def test_scatter(self):
-        # Expected scatter (Rank i) = i
-        mp_manager = mp.Manager()
-        results = mp_manager.dict()
-        mp.spawn(
-            worker,
-            args=(self.world_size, results, "scatter"),
-            nprocs=self.world_size,
-            join=True,
-        )
-        for r in range(self.world_size):
-            self.assertEqual(results[r], r)
-
-    def test_reduce(self):
-        # Expected reduce (Root) = sum([1, 1, ..., 1]) = size
-        # Expected reduce (Others) = None
-        mp_manager = mp.Manager()
-        results = mp_manager.dict()
-        mp.spawn(
-            worker,
-            args=(self.world_size, results, "reduce"),
-            nprocs=self.world_size,
-            join=True,
-        )
-        self.assertEqual(results[0], self.world_size)
-        for r in range(1, self.world_size):
-            self.assertIsNone(results[r])
-
-
-if __name__ == "__main__":
-    unittest.main()
diff --git a/mooncake-wheel/tests/test_mooncake_backend_chunk.py b/mooncake-wheel/tests/test_mooncake_backend_chunk.py
deleted file mode 100644
index 01f9169a..00000000
--- a/mooncake-wheel/tests/test_mooncake_backend_chunk.py
+++ /dev/null
@@ -1,64 +0,0 @@
-import os
-import time
-import unittest
-import torch
-import torch.distributed as dist
-import torch.multiprocessing as mp
-from mooncake import pg
-
-N = 2 ** 24
-
-def worker(rank, world_size, results, collective):
-    torch.cuda.set_device(rank)
-    dist.init_process_group(
-        backend="mooncake-cpu",
-        rank=rank,
-        world_size=world_size,
-        pg_options=pg.MooncakeBackendOptions(torch.zeros((world_size,), dtype=torch.int32, device="cpu")),
-    )
-
-    if collective == "all_reduce":
-        tensor = torch.tensor([rank + 1] * N, dtype=torch.int32, device="cpu")
-        dist.all_reduce(tensor, op=dist.ReduceOp.SUM)
-        results[rank] = tensor[0].item()
-        assert torch.all(tensor == tensor[0].item()) 
-
-    else:
-        raise ValueError(f"Unsupported collective: {collective}")
-
-    while len(results) < world_size:
-        time.sleep(1)
-
-    dist.destroy_process_group()
-
-
-class TestMooncakeBackend(unittest.TestCase):
-    def setUp(self):
-        self.world_size = torch.cuda.device_count()
-        os.environ["MASTER_ADDR"] = "127.0.0.1"
-        os.environ["MASTER_PORT"] = "29500"
-
-    def tearDown(self):
-        pass
-
-    def _spawn_and_check(self, collective):
-        mp_manager = mp.Manager()
-        results = mp_manager.dict()
-        mp.spawn(
-            worker,
-            args=(self.world_size, results, collective),
-            nprocs=self.world_size,
-            join=True,
-        )
-
-        expected = sum(range(1, self.world_size + 1))
-        for r in range(self.world_size):
-            self.assertEqual(results[r], expected)
-
-    def test_allreduce_sum(self):
-        self._spawn_and_check("all_reduce")
-
-
-
-if __name__ == "__main__":
-    unittest.main()
diff --git a/mooncake-wheel/tests/test_mooncake_backend_cpu.py b/mooncake-wheel/tests/test_mooncake_backend_cpu.py
deleted file mode 100644
index 94c301f3..00000000
--- a/mooncake-wheel/tests/test_mooncake_backend_cpu.py
+++ /dev/null
@@ -1,148 +0,0 @@
-import os
-import time
-import unittest
-import torch
-import torch.distributed as dist
-import torch.multiprocessing as mp
-from mooncake import pg
-
-
-def worker(rank, world_size, results, collective):
-    dist.init_process_group(
-        backend="mooncake-cpu",
-        rank=rank,
-        world_size=world_size,
-        pg_options=pg.MooncakeBackendOptions(torch.zeros((world_size,), dtype=torch.int32, device="cpu")),
-    )
-
-    if collective == "all_reduce":
-        tensor = torch.tensor([rank + 1], dtype=torch.int32, device="cpu")
-        dist.all_reduce(tensor, op=dist.ReduceOp.SUM)
-        results[rank] = tensor.item()
-
-    elif collective == "all_gather":
-        tensor = torch.tensor([rank], device="cpu")
-        gathered = [torch.zeros_like(tensor) for _ in range(world_size)]
-        dist.all_gather(gathered, tensor)
-        results[rank] = [t.item() for t in gathered]
-
-    elif collective == "barrier":
-        dist.barrier()
-        results[rank] = "ok"
-
-    elif collective == "gather":
-        tensor = torch.tensor([rank], dtype=torch.int32, device="cpu")
-        if rank == 0:
-            gather_list = [torch.zeros_like(tensor) for _ in range(world_size)]
-            dist.gather(tensor, gather_list, dst=0)
-            results[rank] = [t.item() for t in gather_list]
-        else:
-            dist.gather(tensor, dst=0)
-            results[rank] = None
-
-    elif collective == "scatter":
-        tensor = torch.zeros(1, dtype=torch.int32, device="cpu")
-        if rank == 0:
-            scatter_list = [torch.tensor([i], dtype=torch.int32, device="cpu") for i in range(world_size)]
-            dist.scatter(tensor, scatter_list, src=0)
-        else:
-            dist.scatter(tensor, src=0)
-        results[rank] = tensor.item()
-
-    elif collective == "reduce":
-        tensor = torch.tensor([1], dtype=torch.int32, device="cpu")
-        dist.reduce(tensor, dst=0, op=dist.ReduceOp.SUM)
-        if rank == 0:
-            results[rank] = tensor.item()
-        else:
-            results[rank] = None
-
-    else:
-        raise ValueError(f"Unsupported collective: {collective}")
-
-    while len(results) < world_size:
-        time.sleep(1)
-
-    dist.destroy_process_group()
-
-
-class TestMooncakeBackend(unittest.TestCase):
-    def setUp(self):
-        self.world_size = 4
-        os.environ["MASTER_ADDR"] = "127.0.0.1"
-        os.environ["MASTER_PORT"] = "29500"
-
-    def tearDown(self):
-        pass
-
-    def _spawn_and_check(self, collective, expected_fn):
-        mp_manager = mp.Manager()
-        results = mp_manager.dict()
-        mp.spawn(
-            worker,
-            args=(self.world_size, results, collective),
-            nprocs=self.world_size,
-            join=True,
-        )
-
-        expected = expected_fn(self.world_size)
-        for r in range(self.world_size):
-            self.assertEqual(results[r], expected)
-
-    def test_allreduce(self):
-        # Expected sum = 1 + 2 + 3 + 4 = 10
-        self._spawn_and_check("all_reduce", lambda size: sum(range(1, size + 1)))
-
-    def test_allgather(self):
-        # Expected gather = [0, 1, 2, 3]
-        self._spawn_and_check("all_gather", lambda size: list(range(size)))
-
-    def test_barrier(self):
-        self._spawn_and_check("barrier", lambda size: "ok")
-
-    def test_gather(self):
-        # Expected gather (Root) = [0, 1, 2, ..., size-1]
-        # Expected gather (Others) = None
-        mp_manager = mp.Manager()
-        results = mp_manager.dict()
-        mp.spawn(
-            worker, 
-            args=(self.world_size, results, "gather"), 
-            nprocs=self.world_size, 
-            join=True
-        )
-        self.assertEqual(results[0], list(range(self.world_size)))
-        for r in range(1, self.world_size):
-            self.assertIsNone(results[r])
-
-    def test_scatter(self):
-        # Expected scatter (Rank i) = i
-        mp_manager = mp.Manager()
-        results = mp_manager.dict()
-        mp.spawn(
-            worker, 
-            args=(self.world_size, results, "scatter"), 
-            nprocs=self.world_size, 
-            join=True
-        )
-        for r in range(self.world_size):
-            self.assertEqual(results[r], r)
-
-    def test_reduce(self):
-        # Expected reduce (Root) = sum([1, 1, ..., 1]) = size
-        # Expected reduce (Others) = None
-        mp_manager = mp.Manager()
-        results = mp_manager.dict()
-        mp.spawn(
-            worker, 
-            args=(self.world_size, results, "reduce"), 
-            nprocs=self.world_size, 
-            join=True
-        )
-        self.assertEqual(results[0], self.world_size)
-        for r in range(1, self.world_size):
-            self.assertIsNone(results[r])
-
-
-if __name__ == "__main__":
-    unittest.main()
diff --git a/mooncake-wheel/tests/test_mooncake_backend_elastic.py b/mooncake-wheel/tests/test_mooncake_backend_elastic.py
deleted file mode 100644
index 249e7e98..00000000
--- a/mooncake-wheel/tests/test_mooncake_backend_elastic.py
+++ /dev/null
@@ -1,164 +0,0 @@
-import os
-import time
-import unittest
-
-import torch
-import torch.distributed as dist
-import torch.multiprocessing as mp
-
-from mooncake import pg
-
-USE_CUDA = os.getenv("MOONCAKE_TEST_USE_CUDA") == "1"
-TEST_BACKEND = "mooncake" if USE_CUDA else "mooncake-cpu"
-TEST_DEVICE = torch.device("cuda" if USE_CUDA else "cpu")
-
-os.environ["MASTER_ADDR"] = "127.0.0.1"
-os.environ["MASTER_PORT"] = "19000"
-
-broken_rank = 1
-
-
-def _set_device(rank):
-    if USE_CUDA:
-        torch.cuda.set_device(rank)
-
-
-def _elastic_worker(rank, num_processes, signals):
-    """Worker for testing elastic world size extension."""
-    _set_device(rank)
-    assert num_processes % 2 == 0
-    if rank < num_processes // 2:
-        # Ensure correct operation before extension
-        world_size = num_processes // 2
-        dist.init_process_group(
-            backend=TEST_BACKEND,
-            rank=rank,
-            world_size=world_size,
-        )
-        tensor = torch.tensor([rank + 1], dtype=torch.int32, device=TEST_DEVICE)
-        dist.all_reduce(tensor, op=dist.ReduceOp.SUM)
-        assert tensor.item() == sum(range(1, world_size + 1))
-        if rank == 0:
-            signals["extend"] = 1
-
-        backend = dist.group.WORLD._get_backend(TEST_DEVICE)
-        # Extend world
-        # Note: `extend_group_size_to` is non-blocking. Blocking will only
-        # occur at the first communication if some peers have not yet connected.
-        # This allows overlapping other operations between the group expansion
-        # and the first communication call.
-        pg.extend_group_size_to(backend, num_processes)
-    else:
-        while "extend" not in signals:
-            time.sleep(1)
-        dist.init_process_group(
-            backend=TEST_BACKEND,
-            rank=rank,
-            world_size=num_processes,
-        )
-
-    # Ensure correct operation after extension
-    tensor = torch.tensor([rank + 1], dtype=torch.int32, device=TEST_DEVICE)
-    dist.all_reduce(tensor, op=dist.ReduceOp.SUM)
-    assert tensor.item() == sum(range(1, num_processes + 1)), (
-        f"Rank {rank} expected {sum(range(1, num_processes + 1))}, "
-        f"get {tensor.item()}"
-    )
-
-
-def _deferred_recovery_worker(rank, num_processes, signals):
-    """Worker for testing deferred rank recovery join."""
-    if rank < num_processes:
-        _set_device(rank)
-        dist.init_process_group(
-            backend=TEST_BACKEND,
-            rank=rank,
-            world_size=num_processes,
-        )
-        backend = dist.group.WORLD._get_backend(TEST_DEVICE)
-        while pg.get_num_synced_ranks(backend) < num_processes:
-            time.sleep(0.1)
-        if rank == broken_rank:
-            return  # Simulate broken rank
-
-        expected_without_broken = sum(range(0, num_processes)) - broken_rank
-        tensor = torch.tensor([rank], dtype=torch.int32, device=TEST_DEVICE)
-        dist.all_reduce(tensor, op=dist.ReduceOp.SUM)
-        assert tensor.item() == expected_without_broken
-
-        time.sleep(5)
-        signals["recover"] = 1
-        while True:
-            (peer_state,) = pg.get_peer_state(backend, [broken_rank])
-            if peer_state:
-                break
-
-        # Healthy ranks keep making progress before the recovered rank joins.
-        tensor = torch.tensor([rank], dtype=torch.int32, device=TEST_DEVICE)
-        dist.all_reduce(tensor, op=dist.ReduceOp.SUM)
-        assert tensor.item() == expected_without_broken
-
-        while "join_ready" not in signals:
-            time.sleep(0.1)
-
-        pg.recover_ranks(backend, [broken_rank])
-
-        tensor = torch.tensor([rank], dtype=torch.int32, device=TEST_DEVICE)
-        dist.all_reduce(tensor, op=dist.ReduceOp.SUM)
-        assert tensor.item() == sum(range(0, num_processes)), (
-            f"Rank {rank} expected {sum(range(0, num_processes))}, "
-            f"get {tensor.item()}"
-        )
-    else:
-        _set_device(broken_rank)
-        while "recover" not in signals:
-            time.sleep(0.1)
-
-        dist.init_process_group(
-            backend=TEST_BACKEND,
-            rank=broken_rank,
-            world_size=num_processes,
-            pg_options=pg.MooncakeBackendOptions(
-                torch.ones((num_processes,), dtype=torch.int32, device=TEST_DEVICE),
-                True,
-            ),
-        )
-
-        backend = dist.group.WORLD._get_backend(TEST_DEVICE)
-
-        # Deferred join starts in a local-only mode so collectives stay self-contained.
-        tensor = torch.tensor([broken_rank], dtype=torch.int32, device=TEST_DEVICE)
-        dist.all_reduce(tensor, op=dist.ReduceOp.SUM)
-        assert tensor.item() == broken_rank
-
-        signals["join_ready"] = 1
-        pg.join_group(backend)
-
-        tensor = torch.tensor([broken_rank], dtype=torch.int32, device=TEST_DEVICE)
-        dist.all_reduce(tensor, op=dist.ReduceOp.SUM)
-        assert tensor.item() == sum(range(0, num_processes)), (
-            f"Rank {rank} expected {sum(range(0, num_processes))}, "
-            f"get {tensor.item()}"
-        )
-
-
-class TestMooncakeBackend(unittest.TestCase):
-    def test_elastic_extension(self):
-        num_processes = 4
-        mp_manager = mp.Manager()
-        signals = mp_manager.dict()
-        mp.spawn(_elastic_worker, args=(num_processes, signals), nprocs=num_processes)
-
-    def test_rank_recovery_deferred_join(self):
-        num_processes = 4
-        mp_manager = mp.Manager()
-        signals = mp_manager.dict()
-        mp.spawn(
-            _deferred_recovery_worker,
-            args=(num_processes, signals),
-            nprocs=num_processes + 1,
-        )
-
-
-if __name__ == "__main__":
-    unittest.main()
diff --git a/mooncake-wheel/tests/test_mooncake_backend_p2p_cpu.py b/mooncake-wheel/tests/test_mooncake_backend_p2p_cpu.py
deleted file mode 100644
index 0e8efbdb..00000000
--- a/mooncake-wheel/tests/test_mooncake_backend_p2p_cpu.py
+++ /dev/null
@@ -1,212 +0,0 @@
-import os
-import time
-import unittest
-
-import torch
-import torch.distributed as dist
-import torch.multiprocessing as mp
-
-from mooncake import pg
-
-
-def _worker_ring(rank: int, world_size: int, results):
-    os.environ.setdefault("MASTER_ADDR", "127.0.0.1")
-    os.environ.setdefault("MASTER_PORT", "29500")
-
-    dist.init_process_group(
-        backend="mooncake-cpu",
-        rank=rank,
-        world_size=world_size,
-        pg_options=pg.MooncakeBackendOptions(
-            torch.zeros((world_size,), dtype=torch.int32, device="cpu")
-        ),
-    )
-
-    send_tensor = torch.tensor([rank], dtype=torch.int64, device="cpu")
-    recv_tensor = torch.empty_like(send_tensor)
-
-    dst = (rank + 1) % world_size
-    src = (rank - 1 + world_size) % world_size
-
-    p2p_ops = [
-        dist.P2POp(op=dist.isend, tensor=send_tensor, peer=dst),
-        dist.P2POp(op=dist.irecv, tensor=recv_tensor, peer=src),
-    ]
-
-    works = dist.batch_isend_irecv(p2p_ops)
-    for w in works:
-        w.wait()
-
-    results[rank] = recv_tensor.item()
-
-    while len(results) < world_size:
-        time.sleep(0.1)
-
-    dist.destroy_process_group()
-
-
-def _worker_ordering(rank: int, results):
-    # Only two ranks participate in this test.
-    world_size = 2
-    os.environ.setdefault("MASTER_ADDR", "127.0.0.1")
-    os.environ.setdefault("MASTER_PORT", "29501")
-
-    dist.init_process_group(
-        backend="mooncake-cpu",
-        rank=rank,
-        world_size=world_size,
-        pg_options=pg.MooncakeBackendOptions(
-            torch.zeros((world_size,), dtype=torch.int32, device="cpu")
-        ),
-    )
-
-    num_msgs = 4
-
-    if rank == 0:
-        send_tensors = [
-            torch.tensor([i], dtype=torch.int64, device="cpu")
-            for i in range(num_msgs)
-        ]
-        ops = [
-            dist.P2POp(op=dist.isend, tensor=t, peer=1) for t in send_tensors
-        ]
-        works = dist.batch_isend_irecv(ops)
-        for w in works:
-            w.wait()
-        results[rank] = "ok"
-    else:
-        recv_tensors = [
-            torch.empty(1, dtype=torch.int64, device="cpu")
-            for _ in range(num_msgs)
-        ]
-        ops = [
-            dist.P2POp(op=dist.irecv, tensor=t, peer=0) for t in recv_tensors
-        ]
-        works = dist.batch_isend_irecv(ops)
-        for w in works:
-            w.wait()
-        results[rank] = [t.item() for t in recv_tensors]
-
-    while len(results) < world_size:
-        time.sleep(0.1)
-
-    dist.destroy_process_group()
-
-
-def _worker_multiple_senders(rank: int, world_size: int, results):
-    os.environ.setdefault("MASTER_ADDR", "127.0.0.1")
-    os.environ.setdefault("MASTER_PORT", "29502")
-
-    dist.init_process_group(
-        backend="mooncake-cpu",
-        rank=rank,
-        world_size=world_size,
-        pg_options=pg.MooncakeBackendOptions(
-            torch.zeros((world_size,), dtype=torch.int32, device="cpu")
-        ),
-    )
-
-    if rank == 0:
-        # Rank 0 sends to rank 1
-        send_tensor = torch.tensor([100], dtype=torch.int64, device="cpu")
-        recv_tensor = torch.empty_like(send_tensor)
-        p2p_ops = [
-            dist.P2POp(op=dist.isend, tensor=send_tensor, peer=1),
-            dist.P2POp(op=dist.irecv, tensor=recv_tensor, peer=1),
-        ]
-        works = dist.batch_isend_irecv(p2p_ops)
-        for w in works:
-            w.wait()
-        results[rank] = recv_tensor.item()
-    elif rank == 1:
-        # Rank 1 receives from both rank 0 and rank 2
-        recv_tensor_0 = torch.empty(1, dtype=torch.int64, device="cpu")
-        recv_tensor_2 = torch.empty(1, dtype=torch.int64, device="cpu")
-        send_tensor = torch.tensor([101], dtype=torch.int64, device="cpu")
-        p2p_ops = [
-            dist.P2POp(op=dist.irecv, tensor=recv_tensor_0, peer=0),
-            dist.P2POp(op=dist.irecv, tensor=recv_tensor_2, peer=2),
-            dist.P2POp(op=dist.isend, tensor=send_tensor, peer=0),
-        ]
-        works = dist.batch_isend_irecv(p2p_ops)
-        for w in works:
-            w.wait()
-        results[rank] = [recv_tensor_0.item(), recv_tensor_2.item()]
-    elif rank == 2:
-        # Rank 2 sends to rank 1
-        send_tensor = torch.tensor([200], dtype=torch.int64, device="cpu")
-        p2p_ops = [
-            dist.P2POp(op=dist.isend, tensor=send_tensor, peer=1),
-        ]
-        works = dist.batch_isend_irecv(p2p_ops)
-        for w in works:
-            w.wait()
-        results[rank] = "ok"
-    else:
-        # Other ranks don't participate
-        results[rank] = "skip"
-
-    while len(results) < world_size:
-        time.sleep(0.1)
-
-    dist.destroy_process_group()
-
-
-class TestMooncakeBackendP2PCPU(unittest.TestCase):
-    def test_ring_send_recv(self):
-        world_size = 4
-        mp_manager = mp.Manager()
-        results = mp_manager.dict()
-        mp.spawn(
-            _worker_ring,
-            args=(world_size, results),
-            nprocs=world_size,
-            join=True,
-        )
-
-        for rank in range(world_size):
-            expected = (rank - 1 + world_size) % world_size
-            self.assertEqual(results[rank], expected)
-
-    def test_ordering_between_two_ranks(self):
-        world_size = 2
-        mp_manager = mp.Manager()
-        results = mp_manager.dict()
-        mp.spawn(
-            _worker_ordering,
-            args=(results,),
-            nprocs=world_size,
-            join=True,
-        )
-
-        # Rank 1 should receive messages [0, 1, 2, 3] in order.
-        self.assertEqual(results[1], list(range(4)))
-
-    def test_multiple_senders_to_same_receiver(self):
-        """Test that multiple senders (rank 0 and rank 2) can send to the same receiver (rank 1)
-        without slot conflicts. This tests the receiver-based slot allocation mechanism."""
-        world_size = 4
-        mp_manager = mp.Manager()
-        results = mp_manager.dict()
-        mp.spawn(
-            _worker_multiple_senders,
-            args=(world_size, results),
-            nprocs=world_size,
-            join=True,
-        )
-
-        # Rank 0 should receive 101 from rank 1
-        self.assertEqual(results[0], 101)
-        # Rank 1 should receive 100 from rank 0 and 200 from rank 2
-        # Order may vary, so check that both values are present
-        self.assertIn(100, results[1])
-        self.assertIn(200, results[1])
-        self.assertEqual(len(results[1]), 2)
-        # Rank 2 should complete successfully
-        self.assertEqual(results[2], "ok")
-
-
-if __name__ == "__main__":
-    unittest.main()
-
-
diff --git a/mooncake-wheel/tests/test_mooncake_backend_perf.py b/mooncake-wheel/tests/test_mooncake_backend_perf.py
deleted file mode 100644
index 4f741ed3..00000000
--- a/mooncake-wheel/tests/test_mooncake_backend_perf.py
+++ /dev/null
@@ -1,107 +0,0 @@
-import os
-import torch
-import torch.distributed as dist
-import mooncake.pg
-import time
-import unittest
-import torch.multiprocessing as mp
-
-
-def run_latency_test(rank, world_size, backend, device, collective, data_size, results, num_iterations=100):
-    # Initialize the process group in each spawned process
-    torch.cuda.set_device(rank)
-    dist.init_process_group(backend=backend, rank=rank, world_size=world_size)
-
-    # Create a tensor for the collective operation
-    tensor = torch.rand(data_size, device=device)
-
-    gathered = [torch.zeros_like(tensor) for _ in range(world_size)]
-
-    # Warm up
-    for _ in range(num_iterations):
-        if collective == 'broadcast':
-            dist.broadcast(tensor, src=0)
-        elif collective == 'allreduce':
-            dist.all_reduce(tensor)
-        elif collective == 'allgather':
-            dist.all_gather(gathered, tensor)
-
-    # Synchronize before starting the test
-    torch.cuda.synchronize()
-
-    start = time.perf_counter()
-    for _ in range(num_iterations):
-        if collective == 'broadcast':
-            dist.broadcast(tensor, src=0)
-        elif collective == 'allreduce':
-            dist.all_reduce(tensor)
-        elif collective == 'allgather':
-            dist.all_gather(gathered, tensor)
-
-    torch.cuda.synchronize()
-
-    end = time.perf_counter()
-
-    # Calculate average time
-    avg_time = (end - start) / num_iterations
-
-    # Store the result
-    results[rank] = avg_time
-
-    while len(results) < world_size:
-        time.sleep(1)
-
-    dist.destroy_process_group()  # Destroy the process group after testing
-
-class TestMooncakeBackendPerf(unittest.TestCase):
-    def setUp(self):
-        self.world_size = torch.cuda.device_count()
-        os.environ["MASTER_ADDR"] = "127.0.0.1"
-        os.environ["MASTER_PORT"] = "29500"
-
-    def tearDown(self):
-        pass
-
-    def do_test(self, device, collective, data_size):
-        # Use mp.spawn to call the latency test
-        mp_manager = mp.Manager()
-        # Test mooncake
-        mooncake_results = mp_manager.dict()
-        mp.spawn(run_latency_test, args=(self.world_size, 'mooncake-cpu' if device == 'cpu' else 'mooncake', device, collective, data_size, mooncake_results), nprocs=self.world_size, join=True)
-        # Test baseline
-        baseline_results = mp_manager.dict()
-        mp.spawn(run_latency_test, args=(self.world_size, 'gloo' if device == 'cpu' else 'nccl', device, collective, data_size, baseline_results), nprocs=self.world_size, join=True)
-
-        # After all processes have completed, check the results
-        mooncake_latency = max(mooncake_results[r] for r in mooncake_results)
-        baseline_latency = max(baseline_results[r] for r in baseline_results)
-        self.assertLessEqual(mooncake_latency, 10 * baseline_latency,
-                             f"Latency of mooncake({device}) for {collective} with size {data_size} exceeded 10 times the baseline.")
-
-    # cpu + allgather
-    def test_cpu_allgather_1024(self):
-        self.do_test("cpu", "allgather", 1024)
-
-    # cpu + allreduce
-    def test_cpu_allreduce_1024(self):
-        self.do_test("cpu", "allreduce", 1024)
-
-    # cpu + broadcast
-    def test_cpu_broadcast_1024(self):
-        self.do_test("cpu", "broadcast", 1024)
-
-    # cuda + allgather
-    def test_cuda_allgather_1024(self):
-        self.do_test("cuda", "allgather", 1024)
-
-    # cuda + allreduce
-    def test_cuda_allreduce_1024(self):
-        self.do_test("cuda", "allreduce", 1024)
-
-    # cuda + broadcast
-    def test_cuda_broadcast_1024(self):
-        self.do_test("cuda", "broadcast", 1024)
-
-
-if __name__ == "__main__":
-    unittest.main()
diff --git a/mooncake-wheel/tests/test_mooncake_config.py b/mooncake-wheel/tests/test_mooncake_config.py
index 497da0c8..26ead113 100644
--- a/mooncake-wheel/tests/test_mooncake_config.py
+++ b/mooncake-wheel/tests/test_mooncake_config.py
@@ -3,7 +3,13 @@ import os
 import tempfile
 import unittest
 
-from mooncake.mooncake_config import MooncakeConfig, DEFAULT_GLOBAL_SEGMENT_SIZE, DEFAULT_LOCAL_BUFFER_SIZE
+from mooncake.mooncake_config import (
+    MooncakeConfig,
+    DEFAULT_GLOBAL_SEGMENT_SIZE,
+    DEFAULT_LOCAL_BUFFER_SIZE,
+    _parse_segment_size,
+)
+
 
 class TestMooncakeConfig(unittest.TestCase):
     def setUp(self):
@@ -19,7 +25,9 @@ class TestMooncakeConfig(unittest.TestCase):
             "global_segment_size": 3355443200,
             "local_buffer_size": 1073741824,
             "protocol": "tcp",
-            "device_name": "eth0"
+            "device_name": "eth0",
+            "enable_ssd_offload": True,
+            "ssd_offload_path": "/nvme/mooncake_offload"
         }
 
     def tearDown(self):
@@ -42,6 +50,8 @@ class TestMooncakeConfig(unittest.TestCase):
         self.assertEqual(config.local_buffer_size, 1073741824)
         self.assertEqual(config.protocol, "tcp")
         self.assertEqual(config.device_name, "eth0")
+        self.assertEqual(config.enable_ssd_offload, True)
+        self.assertEqual(config.ssd_offload_path, "/nvme/mooncake_offload")
 
     def test_load_with_default_values(self):
         """Test loading configuration with default values"""
@@ -57,6 +67,45 @@ class TestMooncakeConfig(unittest.TestCase):
         self.assertEqual(config.local_buffer_size, DEFAULT_LOCAL_BUFFER_SIZE)
         self.assertEqual(config.protocol, "tcp")
         self.assertEqual(config.device_name, "")
+        self.assertEqual(config.enable_ssd_offload, False)
+        self.assertEqual(config.ssd_offload_path, "")
+
+    def test_enable_ssd_offload_string_values(self):
+        """from_file must parse string booleans like load_from_env, and reject typos.
+
+        from_file used bool(value), so any non-empty string (including "false")
+        turned SSD offload on, disagreeing with load_from_env which parses the
+        string. _parse_bool now understands the common textual forms and raises
+        on anything unrecognized instead of silently defaulting to False.
+        """
+        required = {
+            "local_hostname": "localhost",
+            "metadata_server": "localhost:8080",
+            "master_server_address": "localhost:8081",
+        }
+        cases = [
+            ("false", False),
+            ("False", False),
+            ("0", False),
+            ("no", False),
+            ("off", False),
+            ("true", True),
+            ("1", True),
+            ("yes", True),
+            ("on", True),
+            (True, True),
+            (False, False),
+        ]
+        for raw, expected in cases:
+            with self.subTest(raw=raw):
+                self.write_config({**required, "enable_ssd_offload": raw})
+                config = MooncakeConfig.from_file(self.config_file)
+                self.assertEqual(config.enable_ssd_offload, expected)
+
+        # An unrecognized string is a config error, not a silent disable.
+        self.write_config({**required, "enable_ssd_offload": "notabool"})
+        with self.assertRaises(ValueError):
+            MooncakeConfig.from_file(self.config_file)
 
     def test_missing_required_field(self):
         """Test missing required field"""
@@ -93,6 +142,8 @@ class TestMooncakeConfig(unittest.TestCase):
         os.environ['MOONCAKE_GLOBAL_SEGMENT_SIZE'] = str(self.valid_config["global_segment_size"])
         os.environ['MOONCAKE_PROTOCOL'] = self.valid_config["protocol"]
         os.environ['MOONCAKE_DEVICE'] = self.valid_config["device_name"]
+        os.environ['MOONCAKE_OFFLOAD_ENABLED'] = str(self.valid_config["enable_ssd_offload"])
+        os.environ['MOONCAKE_OFFLOAD_FILE_STORAGE_PATH'] = self.valid_config["ssd_offload_path"]
 
         try:
             config = MooncakeConfig.load_from_env()
@@ -102,6 +153,8 @@ class TestMooncakeConfig(unittest.TestCase):
             self.assertEqual(config.global_segment_size, self.valid_config["global_segment_size"])
             self.assertEqual(config.protocol, self.valid_config["protocol"])
             self.assertEqual(config.device_name, self.valid_config["device_name"])
+            self.assertEqual(config.enable_ssd_offload, self.valid_config["enable_ssd_offload"])
+            self.assertEqual(config.ssd_offload_path, self.valid_config["ssd_offload_path"])
 
         finally:
             # Clean up environment variable
@@ -111,6 +164,8 @@ class TestMooncakeConfig(unittest.TestCase):
             del os.environ['MOONCAKE_GLOBAL_SEGMENT_SIZE']
             del os.environ['MOONCAKE_PROTOCOL']
             del os.environ['MOONCAKE_DEVICE']
+            del os.environ['MOONCAKE_OFFLOAD_ENABLED']
+            del os.environ['MOONCAKE_OFFLOAD_FILE_STORAGE_PATH']
 
     def test_load_from_env_missing(self):
         """Test loading configuration from environment variable when not set"""
@@ -118,5 +173,68 @@ class TestMooncakeConfig(unittest.TestCase):
             MooncakeConfig.load_from_env()
         self.assertIn("Neither the environment variable 'MOONCAKE_CONFIG_PATH' nor 'MOONCAKE_MASTER' is set.", str(cm.exception))
 
+
+class TestParseSegmentSize(unittest.TestCase):
+    def test_integer_passthrough(self):
+        self.assertEqual(_parse_segment_size(1024), 1024)
+        self.assertEqual(_parse_segment_size(0), 0)
+
+    def test_float_passthrough(self):
+        self.assertEqual(_parse_segment_size(1.5), 1)
+
+    def test_bytes_string(self):
+        self.assertEqual(_parse_segment_size("1024"), 1024)
+        self.assertEqual(_parse_segment_size("  2048  "), 2048)
+
+    def test_kb_suffix(self):
+        self.assertEqual(_parse_segment_size("1kb"), 1024)
+        self.assertEqual(_parse_segment_size("1KB"), 1024)
+        self.assertEqual(_parse_segment_size("512k"), 512 * 1024)
+        self.assertEqual(_parse_segment_size("1.5kb"), int(1.5 * 1024))
+
+    def test_mb_suffix(self):
+        self.assertEqual(_parse_segment_size("1mb"), 1024 ** 2)
+        self.assertEqual(_parse_segment_size("512MB"), 512 * 1024 ** 2)
+        self.assertEqual(_parse_segment_size("1m"), 1024 ** 2)
+
+    def test_gb_suffix(self):
+        self.assertEqual(_parse_segment_size("1gb"), 1024 ** 3)
+        self.assertEqual(_parse_segment_size("3GB"), 3 * 1024 ** 3)
+        self.assertEqual(_parse_segment_size("1g"), 1024 ** 3)
+        self.assertEqual(_parse_segment_size("1.5gb"), int(1.5 * 1024 ** 3))
+
+    def test_tb_suffix(self):
+        self.assertEqual(_parse_segment_size("1tb"), 1024 ** 4)
+        self.assertEqual(_parse_segment_size("1TB"), 1024 ** 4)
+        self.assertEqual(_parse_segment_size("1t"), 1024 ** 4)
+
+    def test_b_suffix(self):
+        self.assertEqual(_parse_segment_size("4096b"), 4096)
+        self.assertEqual(_parse_segment_size("4096B"), 4096)
+
+    def test_empty_string_raises(self):
+        with self.assertRaises(ValueError):
+            _parse_segment_size("")
+        with self.assertRaises(ValueError):
+            _parse_segment_size("   ")
+
+    def test_missing_number_raises(self):
+        with self.assertRaises(ValueError):
+            _parse_segment_size("gb")
+        with self.assertRaises(ValueError):
+            _parse_segment_size("mb")
+
+    def test_bare_float_string(self):
+        self.assertEqual(_parse_segment_size("1.5"), 1)
+        self.assertEqual(_parse_segment_size("1e9"), 1000000000)
+
+    def test_invalid_string_raises(self):
+        with self.assertRaises(ValueError):
+            _parse_segment_size("abc")
+
+    def test_whitespace_handling(self):
+        self.assertEqual(_parse_segment_size("  3 gb  "), 3 * 1024 ** 3)
+
+
 if __name__ == '__main__':
     unittest.main()
diff --git a/mooncake-wheel/tests/test_mooncake_store_service_api.py b/mooncake-wheel/tests/test_mooncake_store_service_api.py
new file mode 100644
index 00000000..62d28e0e
--- /dev/null
+++ b/mooncake-wheel/tests/test_mooncake_store_service_api.py
@@ -0,0 +1,727 @@
+#!/usr/bin/env python3
+import asyncio
+import json
+import sys
+import types
+import unittest
+from pathlib import Path
+from types import SimpleNamespace
+
+sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
+
+try:
+    from aiohttp import web as _unused_web  # noqa: F401
+except ModuleNotFoundError:
+    aiohttp_module = types.ModuleType("aiohttp")
+    web_module = types.ModuleType("aiohttp.web")
+
+    class Response:
+        def __init__(self, status=200, text="", body=None, content_type=None):
+            self.status = status
+            self.text = text
+            self.body = body
+            self.content_type = content_type
+
+    web_module.Response = Response
+    aiohttp_module.web = web_module
+    sys.modules["aiohttp"] = aiohttp_module
+    sys.modules["aiohttp.web"] = web_module
+
+try:
+    from mooncake.store import MooncakeDistributedStore as _unused_store  # noqa: F401
+except ModuleNotFoundError:
+    store_module = types.ModuleType("mooncake.store")
+
+    class MooncakeDistributedStore:
+        pass
+
+    store_module.MooncakeDistributedStore = MooncakeDistributedStore
+    sys.modules["mooncake.store"] = store_module
+
+from mooncake.mooncake_store_service import MooncakeStoreService, _shm_name_to_path
+
+
+class FakeStore:
+    def __init__(self):
+        self.mounted = {}
+        self.mount_calls = []
+        self.unmount_calls = []
+        self.fail_mount = False
+        self.unmount_failures = set()
+        self.allocated_mount_calls = []
+        self.free_unmount_calls = []
+
+    def mount_segment(self, path, size, offset, protocol, location):
+        self.mount_calls.append((path, size, offset, protocol, location))
+        if self.fail_mount:
+            return {"ret": -1, "segment_ids": []}
+        segment_id = "00000000-0000-0000-0000-000000000001"
+        self.mounted[segment_id] = {
+            "path": path,
+            "size": size,
+            "offset": offset,
+            "protocol": protocol,
+            "location": location,
+        }
+        return {"ret": 0, "segment_ids": [segment_id]}
+
+    def unmount_segment(self, segment_ids, grace_period_seconds=0):
+        self.unmount_calls.append((list(segment_ids), grace_period_seconds))
+        for segment_id in segment_ids:
+            if segment_id in self.unmount_failures:
+                return -1
+        for segment_id in segment_ids:
+            self.mounted.pop(segment_id, None)
+        return 0
+
+    def allocate_and_mount_segment(self, size, protocol, location):
+        self.allocated_mount_calls.append((size, protocol, location))
+        segment_id = "00000000-0000-0000-0000-000000000002"
+        return {
+            "ret": 0,
+            "segment_ids": [segment_id],
+            "allocated_size": 4096,
+        }
+
+    def unmount_and_free_segment(self, segment_ids, grace_period_seconds=0):
+        self.free_unmount_calls.append((list(segment_ids), grace_period_seconds))
+        return 0
+
+
+class FakeRequest:
+    def __init__(self, body):
+        self.body = body
+
+    async def json(self):
+        return self.body
+
+
+class StoreServiceApiTest(unittest.IsolatedAsyncioTestCase):
+    async def asyncSetUp(self):
+        self.fake_store = FakeStore()
+        self.service = MooncakeStoreService.__new__(MooncakeStoreService)
+        self.service.store = self.fake_store
+        self.service.config = SimpleNamespace(protocol="tcp")
+        self.service.current_mode = "prefill"
+        self.service.mounted_segment_ids = []
+        self.service.last_mount_info = {}
+        self.service._state_lock = asyncio.Lock()
+
+    async def test_mount_shm_then_unmount_shm_api(self):
+        mount_resp = await self.service.handle_mount_shm(
+            FakeRequest(
+                {
+                    "name": "mooncake-segment",
+                    "size": 4096,
+                    "offset": 128,
+                    "protocol": "tcp",
+                    "location": "cpu:0",
+                }
+            )
+        )
+        self.assertEqual(mount_resp.status, 200)
+        mount_body = json.loads(mount_resp.text)
+        self.assertEqual(mount_body["status"], "success")
+        self.assertEqual(
+            mount_body["segment_ids"],
+            ["00000000-0000-0000-0000-000000000001"],
+        )
+        self.assertEqual(
+            self.fake_store.mount_calls,
+            [("/dev/shm/mooncake-segment", 4096, 128, "tcp", "cpu:0")],
+        )
+
+        unmount_resp = await self.service.handle_unmount_shm(
+            FakeRequest({"segment_ids": mount_body["segment_ids"]})
+        )
+        self.assertEqual(unmount_resp.status, 200)
+        unmount_body = json.loads(unmount_resp.text)
+        self.assertEqual(unmount_body["status"], "success")
+        self.assertEqual(
+            self.fake_store.unmount_calls,
+            [(["00000000-0000-0000-0000-000000000001"], 0)],
+        )
+        self.assertEqual(self.service.current_mode, "prefill")
+
+    async def test_unmount_shm_passes_grace_period(self):
+        segment_id = "00000000-0000-0000-0000-000000000001"
+        self.service.current_mode = "decode"
+        self.service.mounted_segment_ids = [segment_id]
+
+        resp = await self.service.handle_unmount_shm(
+            FakeRequest({"segment_ids": [segment_id], "grace_period_seconds": 3})
+        )
+
+        self.assertEqual(resp.status, 200)
+        self.assertEqual(self.fake_store.unmount_calls, [([segment_id], 3)])
+        self.assertEqual(self.service.mounted_segment_ids, [])
+        self.assertEqual(self.service.current_mode, "prefill")
+
+    async def test_unmount_shm_updates_state_for_partial_success(self):
+        succeeded_id = "00000000-0000-0000-0000-000000000001"
+        failed_id = "00000000-0000-0000-0000-000000000002"
+        self.service.current_mode = "decode"
+        self.service.mounted_segment_ids = [succeeded_id, failed_id]
+        self.fake_store.unmount_failures = {failed_id}
+
+        resp = await self.service.handle_unmount_shm(
+            FakeRequest({"segment_ids": [succeeded_id, failed_id]})
+        )
+
+        self.assertEqual(resp.status, 500)
+        body = json.loads(resp.text)
+        self.assertEqual(body["failed_segment_ids"], [failed_id])
+        self.assertEqual(
+            self.fake_store.unmount_calls,
+            [([succeeded_id], 0), ([failed_id], 0)],
+        )
+        self.assertEqual(self.service.mounted_segment_ids, [failed_id])
+        self.assertEqual(self.service.current_mode, "decode")
+
+    async def test_reconfigure_decode_mount_failure_rolls_back_to_prefill(self):
+        old_id = "00000000-0000-0000-0000-000000000001"
+        self.service.current_mode = "decode"
+        self.service.mounted_segment_ids = [old_id]
+        self.service.last_mount_info = {
+            "path": "/dev/shm/old",
+            "offset": 0,
+            "size": 4096,
+            "protocol": "tcp",
+            "location": "",
+        }
+        self.fake_store.fail_mount = True
+
+        resp = await self.service.handle_reconfigure(
+            FakeRequest({"mode": "decode", "path": "/dev/shm/new", "size": 4096})
+        )
+
+        self.assertEqual(resp.status, 500)
+        body = json.loads(resp.text)
+        self.assertEqual(body["mode"], "prefill")
+        self.assertIn("rolled back to prefill", body["error"])
+        self.assertEqual(self.fake_store.unmount_calls, [([old_id], 0)])
+        self.assertEqual(self.service.mounted_segment_ids, [])
+        self.assertEqual(self.service.current_mode, "prefill")
+        self.assertEqual(self.service.last_mount_info, {})
+
+    async def test_mount_allocates_and_frees_on_unmount(self):
+        mount_resp = await self.service.handle_mount(
+            FakeRequest({"size": 1, "protocol": "tcp", "location": "cpu:0"})
+        )
+        self.assertEqual(mount_resp.status, 200)
+        mount_body = json.loads(mount_resp.text)
+        self.assertEqual(mount_body["status"], "success")
+        self.assertEqual(
+            mount_body["segment_ids"],
+            ["00000000-0000-0000-0000-000000000002"],
+        )
+        self.assertEqual(mount_body["allocated_size"], 4096)
+        self.assertEqual(self.fake_store.allocated_mount_calls, [(1, "tcp", "cpu:0")])
+
+        unmount_resp = await self.service.handle_unmount(
+            FakeRequest({"segment_ids": mount_body["segment_ids"]})
+        )
+        self.assertEqual(unmount_resp.status, 200)
+        unmount_body = json.loads(unmount_resp.text)
+        self.assertEqual(unmount_body["status"], "success")
+        self.assertEqual(
+            self.fake_store.free_unmount_calls,
+            [(["00000000-0000-0000-0000-000000000002"], 0)],
+        )
+
+    async def test_unmount_allocated_passes_grace_period(self):
+        segment_id = "00000000-0000-0000-0000-000000000002"
+
+        resp = await self.service.handle_unmount(
+            FakeRequest({"segment_ids": [segment_id], "grace_period_seconds": 4})
+        )
+
+        self.assertEqual(resp.status, 200)
+        self.assertEqual(self.fake_store.free_unmount_calls, [([segment_id], 4)])
+
+    async def test_mount_shm_requires_name_and_size(self):
+        resp = await self.service.handle_mount_shm(
+            FakeRequest({"name": "mooncake-segment"})
+        )
+        self.assertEqual(resp.status, 400)
+        body = json.loads(resp.text)
+        self.assertIn("name or size", body["error"])
+
+    async def test_mount_shm_rejects_path_like_name(self):
+        resp = await self.service.handle_mount_shm(
+            FakeRequest({"name": "../foo", "size": 4096})
+        )
+        self.assertEqual(resp.status, 400)
+        body = json.loads(resp.text)
+        self.assertIn("name or size", body["error"])
+
+    async def test_mount_requires_positive_size(self):
+        resp = await self.service.handle_mount(FakeRequest({"size": 0}))
+        self.assertEqual(resp.status, 400)
+        body = json.loads(resp.text)
+        self.assertIn("Invalid size", body["error"])
+
+    async def test_unmount_shm_requires_segment_ids(self):
+        resp = await self.service.handle_unmount_shm(FakeRequest({}))
+        self.assertEqual(resp.status, 400)
+        body = json.loads(resp.text)
+        self.assertIn("Missing segment_ids", body["error"])
+
+    # ==================== /api/put tests ====================
+
+    async def test_handle_put_success(self):
+        self.fake_store.put = lambda key, value: 0
+        resp = await self.service.handle_put(
+            FakeRequest({"key": "test_key", "value": "test_value"})
+        )
+        self.assertEqual(resp.status, 200)
+        body = json.loads(resp.text)
+        self.assertEqual(body["status"], "success")
+
+    async def test_handle_put_missing_key(self):
+        self.fake_store.put = lambda key, value: 0
+        resp = await self.service.handle_put(FakeRequest({"value": "test_value"}))
+        self.assertEqual(resp.status, 400)
+        body = json.loads(resp.text)
+        self.assertIn("Missing key or value", body["error"])
+
+    async def test_handle_put_missing_value(self):
+        self.fake_store.put = lambda key, value: 0
+        resp = await self.service.handle_put(FakeRequest({"key": "k"}))
+        self.assertEqual(resp.status, 500)
+
+    async def test_handle_put_store_failure(self):
+        self.fake_store.put = lambda key, value: -1
+        resp = await self.service.handle_put(FakeRequest({"key": "k", "value": "v"}))
+        self.assertEqual(resp.status, 500)
+        body = json.loads(resp.text)
+        self.assertIn("PUT operation failed", body["error"])
+
+    async def test_handle_put_empty_key(self):
+        self.fake_store.put = lambda key, value: 0
+        resp = await self.service.handle_put(FakeRequest({"key": "", "value": "v"}))
+        self.assertEqual(resp.status, 400)
+        body = json.loads(resp.text)
+        self.assertIn("Missing key or value", body["error"])
+
+    # ==================== /api/get/{key} tests ====================
+
+    async def test_handle_get_success(self):
+        self.fake_store.is_exist = lambda key: True
+        self.fake_store.get = lambda key: b"payload_bytes"
+        request = FakeRequest({})
+        request.match_info = {"key": "my_key"}
+        resp = await self.service.handle_get(request)
+        self.assertEqual(resp.status, 200)
+        self.assertEqual(resp.body, b"payload_bytes")
+
+    async def test_handle_get_not_found(self):
+        self.fake_store.is_exist = lambda key: 0
+        self.fake_store.get = lambda key: b""
+        request = FakeRequest({})
+        request.match_info = {"key": "missing_key"}
+        resp = await self.service.handle_get(request)
+        self.assertEqual(resp.status, 404)
+        body = json.loads(resp.text)
+        self.assertIn("Key not found", body["error"])
+
+    async def test_handle_get_exist_check_failure(self):
+        self.fake_store.is_exist = lambda key: -1
+        self.fake_store.get = lambda key: b""
+        request = FakeRequest({})
+        request.match_info = {"key": "error_key"}
+        resp = await self.service.handle_get(request)
+        self.assertEqual(resp.status, 500)
+        body = json.loads(resp.text)
+        self.assertIn("Exist check failed", body["error"])
+
+    async def test_handle_get_empty_bytes(self):
+        self.fake_store.is_exist = lambda key: True
+        self.fake_store.get = lambda key: b""
+        request = FakeRequest({})
+        request.match_info = {"key": "empty_value_key"}
+        resp = await self.service.handle_get(request)
+        self.assertEqual(resp.status, 200)
+        self.assertEqual(resp.body, b"")
+
+    async def test_handle_get_empty_bytes_rechecks_existence(self):
+        existence_results = iter([1, 0])
+        self.fake_store.is_exist = lambda key: next(existence_results)
+        self.fake_store.get = lambda key: b""
+        request = FakeRequest({})
+        request.match_info = {"key": "empty_value_key"}
+        resp = await self.service.handle_get(request)
+        self.assertEqual(resp.status, 404)
+        body = json.loads(resp.text)
+        self.assertIn("Key not found", body["error"])
+
+    async def test_handle_get_none_value_is_store_failure(self):
+        self.fake_store.is_exist = lambda key: True
+        self.fake_store.get = lambda key: None
+        request = FakeRequest({})
+        request.match_info = {"key": "empty_value_key"}
+        resp = await self.service.handle_get(request)
+        self.assertEqual(resp.status, 500)
+        body = json.loads(resp.text)
+        self.assertIn("GET operation failed", body["error"])
+
+    async def test_handle_get_store_exception(self):
+        def raise_error(key):
+            raise RuntimeError("store crashed")
+
+        self.fake_store.is_exist = lambda key: True
+        self.fake_store.get = raise_error
+        request = FakeRequest({})
+        request.match_info = {"key": "crash_key"}
+        resp = await self.service.handle_get(request)
+        self.assertEqual(resp.status, 500)
+        body = json.loads(resp.text)
+        self.assertIn("store crashed", body["error"])
+
+    # ==================== /api/exist/{key} tests ====================
+
+    async def test_handle_exist_true(self):
+        self.fake_store.is_exist = lambda key: True
+        request = FakeRequest({})
+        request.match_info = {"key": "existing_key"}
+        resp = await self.service.handle_exist(request)
+        self.assertEqual(resp.status, 200)
+        body = json.loads(resp.text)
+        self.assertTrue(body["exists"])
+
+    async def test_handle_exist_false(self):
+        self.fake_store.is_exist = lambda key: False
+        request = FakeRequest({})
+        request.match_info = {"key": "missing_key"}
+        resp = await self.service.handle_exist(request)
+        self.assertEqual(resp.status, 200)
+        body = json.loads(resp.text)
+        self.assertFalse(body["exists"])
+
+    async def test_handle_exist_store_exception(self):
+        def raise_error(key):
+            raise RuntimeError("exist check crashed")
+
+        self.fake_store.is_exist = raise_error
+        request = FakeRequest({})
+        request.match_info = {"key": "crash_key"}
+        resp = await self.service.handle_exist(request)
+        self.assertEqual(resp.status, 500)
+        body = json.loads(resp.text)
+        self.assertIn("exist check crashed", body["error"])
+
+    async def test_handle_exist_store_error(self):
+        # is_exist returns -1 when the store is unhealthy; this should surface
+        # as HTTP 500, not as HTTP 200 {"exists": true} (bool(-1) == True).
+        self.fake_store.is_exist = lambda key: -1
+        request = FakeRequest({})
+        request.match_info = {"key": "some_key"}
+        resp = await self.service.handle_exist(request)
+        self.assertEqual(resp.status, 500)
+        body = json.loads(resp.text)
+        # Assert the specific message so this pins the exists < 0 branch rather
+        # than any 500 (the except path returns {"error": str(e)} too).
+        self.assertEqual(body["error"], "Exist check failed")
+
+    # ==================== /api/remove/{key} tests ====================
+
+    async def test_handle_remove_success(self):
+        self.fake_store.remove = lambda key: 0
+        request = FakeRequest({})
+        request.match_info = {"key": "removable_key"}
+        resp = await self.service.handle_remove(request)
+        self.assertEqual(resp.status, 200)
+        body = json.loads(resp.text)
+        self.assertEqual(body["status"], "success")
+
+    async def test_handle_remove_failure(self):
+        self.fake_store.remove = lambda key: -1
+        request = FakeRequest({})
+        request.match_info = {"key": "stuck_key"}
+        resp = await self.service.handle_remove(request)
+        self.assertEqual(resp.status, 500)
+        body = json.loads(resp.text)
+        self.assertIn("Remove operation failed", body["error"])
+
+    async def test_handle_remove_store_exception(self):
+        def raise_error(key):
+            raise RuntimeError("remove crashed")
+
+        self.fake_store.remove = raise_error
+        request = FakeRequest({})
+        request.match_info = {"key": "crash_key"}
+        resp = await self.service.handle_remove(request)
+        self.assertEqual(resp.status, 500)
+        body = json.loads(resp.text)
+        self.assertIn("remove crashed", body["error"])
+
+    # ==================== /api/remove_all tests ====================
+
+    async def test_handle_remove_all_success(self):
+        self.fake_store.remove_all = lambda: 5
+        resp = await self.service.handle_remove_all(FakeRequest({}))
+        self.assertEqual(resp.status, 200)
+        body = json.loads(resp.text)
+        self.assertIn("5", body["status"])
+
+    async def test_handle_remove_all_zero_keys(self):
+        self.fake_store.remove_all = lambda: 0
+        resp = await self.service.handle_remove_all(FakeRequest({}))
+        self.assertEqual(resp.status, 200)
+        body = json.loads(resp.text)
+        self.assertIn("0", body["status"])
+
+    async def test_handle_remove_all_failure(self):
+        self.fake_store.remove_all = lambda: -1
+        resp = await self.service.handle_remove_all(FakeRequest({}))
+        self.assertEqual(resp.status, 500)
+        body = json.loads(resp.text)
+        self.assertIn("RemoveAll operation failed", body["error"])
+
+    async def test_handle_remove_all_store_exception(self):
+        def raise_error():
+            raise RuntimeError("remove_all crashed")
+
+        self.fake_store.remove_all = raise_error
+        resp = await self.service.handle_remove_all(FakeRequest({}))
+        self.assertEqual(resp.status, 500)
+        body = json.loads(resp.text)
+        self.assertIn("remove_all crashed", body["error"])
+
+    # ==================== /api/reconfigure tests ====================
+
+    async def test_reconfigure_decode_success(self):
+        self.fake_store.fail_mount = False
+        resp = await self.service.handle_reconfigure(
+            FakeRequest(
+                {
+                    "mode": "decode",
+                    "path": "/dev/shm/test",
+                    "size": 4096,
+                }
+            )
+        )
+        self.assertEqual(resp.status, 200)
+        body = json.loads(resp.text)
+        self.assertEqual(body["mode"], "decode")
+        self.assertEqual(body["status"], "success")
+        self.assertIn("segment_ids", body)
+        self.assertEqual(self.service.current_mode, "decode")
+
+    async def test_reconfigure_decode_missing_path(self):
+        resp = await self.service.handle_reconfigure(
+            FakeRequest({"mode": "decode", "size": 4096})
+        )
+        self.assertEqual(resp.status, 400)
+        body = json.loads(resp.text)
+        self.assertIn("Missing path or size", body["error"])
+
+    async def test_reconfigure_decode_missing_size(self):
+        resp = await self.service.handle_reconfigure(
+            FakeRequest({"mode": "decode", "path": "/dev/shm/test"})
+        )
+        self.assertEqual(resp.status, 400)
+        body = json.loads(resp.text)
+        self.assertIn("Missing path or size", body["error"])
+
+    async def test_reconfigure_prefill_success(self):
+        self.service.current_mode = "decode"
+        self.service.mounted_segment_ids = []
+        resp = await self.service.handle_reconfigure(FakeRequest({"mode": "prefill"}))
+        self.assertEqual(resp.status, 200)
+        body = json.loads(resp.text)
+        self.assertEqual(body["mode"], "prefill")
+        self.assertEqual(self.service.current_mode, "prefill")
+
+    async def test_reconfigure_prefill_unmounts_segments(self):
+        sid = "00000000-0000-0000-0000-000000000001"
+        self.service.current_mode = "decode"
+        self.service.mounted_segment_ids = [sid]
+        resp = await self.service.handle_reconfigure(FakeRequest({"mode": "prefill"}))
+        self.assertEqual(resp.status, 200)
+        self.assertEqual(self.fake_store.unmount_calls, [[sid]])
+        self.assertEqual(self.service.mounted_segment_ids, [])
+        self.assertEqual(self.service.current_mode, "prefill")
+
+    async def test_reconfigure_prefill_unmount_failure(self):
+        sid = "00000000-0000-0000-0000-000000000002"
+        self.service.current_mode = "decode"
+        self.service.mounted_segment_ids = [sid]
+        self.fake_store.unmount_failures = {sid}
+        resp = await self.service.handle_reconfigure(FakeRequest({"mode": "prefill"}))
+        self.assertEqual(resp.status, 500)
+        body = json.loads(resp.text)
+        self.assertIn("Unmount failed", body["error"])
+
+    async def test_reconfigure_invalid_mode(self):
+        resp = await self.service.handle_reconfigure(FakeRequest({"mode": "invalid"}))
+        self.assertEqual(resp.status, 400)
+        body = json.loads(resp.text)
+        self.assertIn("Invalid mode", body["error"])
+
+    async def test_reconfigure_empty_mode(self):
+        resp = await self.service.handle_reconfigure(FakeRequest({}))
+        self.assertEqual(resp.status, 400)
+        body = json.loads(resp.text)
+        self.assertIn("Invalid mode", body["error"])
+
+    async def test_reconfigure_decode_remount_unmounts_previous(self):
+        old_id = "00000000-0000-0000-0000-000000000001"
+        self.service.current_mode = "decode"
+        self.service.mounted_segment_ids = [old_id]
+        self.fake_store.fail_mount = False
+        resp = await self.service.handle_reconfigure(
+            FakeRequest(
+                {
+                    "mode": "decode",
+                    "path": "/dev/shm/new",
+                    "size": 8192,
+                }
+            )
+        )
+        self.assertEqual(resp.status, 200)
+        self.assertEqual(self.fake_store.unmount_calls, [[old_id]])
+        self.assertEqual(self.service.current_mode, "decode")
+
+    # ==================== /api/mount edge cases ====================
+
+    async def test_mount_negative_size(self):
+        resp = await self.service.handle_mount(FakeRequest({"size": -1}))
+        self.assertEqual(resp.status, 400)
+        body = json.loads(resp.text)
+        self.assertIn("Invalid size", body["error"])
+
+    async def test_mount_float_size(self):
+        resp = await self.service.handle_mount(FakeRequest({"size": 1.5}))
+        self.assertEqual(resp.status, 400)
+        body = json.loads(resp.text)
+        self.assertIn("Invalid size", body["error"])
+
+    async def test_mount_string_size(self):
+        resp = await self.service.handle_mount(FakeRequest({"size": "1024"}))
+        self.assertEqual(resp.status, 400)
+        body = json.loads(resp.text)
+        self.assertIn("Invalid size", body["error"])
+
+    async def test_mount_missing_size(self):
+        resp = await self.service.handle_mount(FakeRequest({}))
+        self.assertEqual(resp.status, 400)
+
+    # ==================== /api/unmount edge cases ====================
+
+    async def test_unmount_requires_segment_ids(self):
+        resp = await self.service.handle_unmount(FakeRequest({}))
+        self.assertEqual(resp.status, 400)
+        body = json.loads(resp.text)
+        self.assertIn("Missing segment_ids", body["error"])
+
+    async def test_unmount_empty_list(self):
+        resp = await self.service.handle_unmount(FakeRequest({"segment_ids": []}))
+        self.assertEqual(resp.status, 400)
+        body = json.loads(resp.text)
+        self.assertIn("Missing segment_ids", body["error"])
+
+    async def test_unmount_success(self):
+        resp = await self.service.handle_unmount(
+            FakeRequest({"segment_ids": ["00000000-0000-0000-0000-000000000002"]})
+        )
+        self.assertEqual(resp.status, 200)
+        body = json.loads(resp.text)
+        self.assertEqual(body["status"], "success")
+
+    # ==================== /api/mount_shm edge cases ====================
+
+    async def test_mount_shm_store_failure(self):
+        self.fake_store.fail_mount = True
+        resp = await self.service.handle_mount_shm(
+            FakeRequest({"name": "test-seg", "size": 4096})
+        )
+        self.assertEqual(resp.status, 500)
+        body = json.loads(resp.text)
+        self.assertIn("Mount failed", body["error"])
+
+    async def test_mount_shm_with_defaults(self):
+        resp = await self.service.handle_mount_shm(
+            FakeRequest({"name": "minimal-seg", "size": 2048})
+        )
+        self.assertEqual(resp.status, 200)
+        self.assertEqual(
+            self.fake_store.mount_calls,
+            [("/dev/shm/minimal-seg", 2048, 0, "tcp", "")],
+        )
+
+    async def test_mount_shm_empty_name(self):
+        resp = await self.service.handle_mount_shm(
+            FakeRequest({"name": "", "size": 4096})
+        )
+        self.assertEqual(resp.status, 400)
+
+    async def test_mount_shm_dot_name(self):
+        resp = await self.service.handle_mount_shm(
+            FakeRequest({"name": ".", "size": 4096})
+        )
+        self.assertEqual(resp.status, 400)
+
+    async def test_mount_shm_dotdot_name(self):
+        resp = await self.service.handle_mount_shm(
+            FakeRequest({"name": "..", "size": 4096})
+        )
+        self.assertEqual(resp.status, 400)
+
+    # ==================== /api/unmount_shm edge cases ====================
+
+    async def test_unmount_shm_string_segment_id_coercion(self):
+        resp = await self.service.handle_unmount_shm(
+            FakeRequest({"segment_ids": "00000000-0000-0000-0000-000000000001"})
+        )
+        self.assertEqual(resp.status, 200)
+        self.assertEqual(
+            self.fake_store.unmount_calls,
+            [["00000000-0000-0000-0000-000000000001"]],
+        )
+
+    async def test_unmount_shm_empty_list(self):
+        resp = await self.service.handle_unmount_shm(FakeRequest({"segment_ids": []}))
+        self.assertEqual(resp.status, 400)
+
+
+class ShmNameToPathTest(unittest.TestCase):
+    def test_valid_simple_name(self):
+        self.assertEqual(_shm_name_to_path("my-segment"), "/dev/shm/my-segment")
+
+    def test_valid_leading_slash(self):
+        self.assertEqual(_shm_name_to_path("/my-segment"), "/dev/shm/my-segment")
+
+    def test_rejects_empty(self):
+        self.assertIsNone(_shm_name_to_path(""))
+
+    def test_rejects_none(self):
+        self.assertIsNone(_shm_name_to_path(None))
+
+    def test_rejects_non_string(self):
+        self.assertIsNone(_shm_name_to_path(123))
+
+    def test_rejects_dot(self):
+        self.assertIsNone(_shm_name_to_path("."))
+
+    def test_rejects_dotdot(self):
+        self.assertIsNone(_shm_name_to_path(".."))
+
+    def test_rejects_path_traversal(self):
+        self.assertIsNone(_shm_name_to_path("../etc/passwd"))
+
+    def test_rejects_nested_path(self):
+        self.assertIsNone(_shm_name_to_path("subdir/file"))
+
+    def test_rejects_slash_only(self):
+        self.assertIsNone(_shm_name_to_path("/"))
+
+    def test_rejects_slash_dot(self):
+        self.assertIsNone(_shm_name_to_path("/."))
+
+    def test_rejects_slash_dotdot(self):
+        self.assertIsNone(_shm_name_to_path("/.."))
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/mooncake-wheel/tests/test_promotion_on_hit.py b/mooncake-wheel/tests/test_promotion_on_hit.py
new file mode 100644
index 00000000..19efc177
--- /dev/null
+++ b/mooncake-wheel/tests/test_promotion_on_hit.py
@@ -0,0 +1,492 @@
+"""Python-binding test for the L2->L1 promotion-on-hit feature.
+
+Mirror of ``test_offload_on_eviction.py``: that test asserts data flows
+DRAM -> SSD on eviction; this test asserts the reverse — once an object
+exists only on LOCAL_DISK and a client reads it enough times to clear the
+admission threshold, the master enqueues a promotion task and the client's
+next heartbeat tick stages a fresh MEMORY replica.
+
+Test scenario:
+  1. Push enough data to overflow DRAM, so eviction + offload turns warm
+     keys into LOCAL_DISK-only objects.
+  2. Identify one such key from the replica descriptors.
+  3. Read it repeatedly to clear ``promotion_admission_threshold``.
+  4. Wait long enough for one master heartbeat (which queues the task) plus
+     one FileStorage heartbeat (which executes it).
+  5. Assert the key now also has a MEMORY replica.
+
+Prerequisites:
+  - ``mooncake_master`` running with master config containing:
+      ``--enable_offload=true``
+      ``--offload_on_evict=true``
+      ``--promotion_on_hit=true``
+      ``--promotion_admission_threshold=1``  (any positive int; we send
+          enough reads to clear up to 4)
+      ``--root_fs_dir=<dir>``  (so the master tells the client to init
+          its FileStorage; without this, no offload heartbeat runs and
+          no LOCAL_DISK replicas are ever created)
+  - ``MOONCAKE_OFFLOAD_FILE_STORAGE_PATH=<dir>`` set on the client side.
+  - ``MOONCAKE_OFFLOAD_BUCKET_KEYS_LIMIT=10`` and
+    ``MOONCAKE_OFFLOAD_BUCKET_SIZE_LIMIT_BYTES=10485760`` on the client.
+    The default bucket-flush thresholds (500 keys / 256 MB) are too high
+    for a single-process test workload; the client's BucketStorageBackend
+    would buffer everything in the ungrouped pool and never write to disk.
+"""
+
+import os
+import statistics
+import time
+import unittest
+from mooncake.store import MooncakeDistributedStore
+
+DEFAULT_DEFAULT_KV_LEASE_TTL = 5000  # ms
+default_kv_lease_ttl = int(
+    os.getenv("DEFAULT_KV_LEASE_TTL", DEFAULT_DEFAULT_KV_LEASE_TTL)
+)
+
+# Keep the segment small so we can overflow it cheaply in CI.
+# 32 MB — small enough that the 96 x 1 MB workload reliably overflows past
+# the eviction high watermark even with OffsetAllocator block padding.
+# Larger values (e.g. 64 MB) sometimes fit the full workload because the
+# allocator's effective capacity exceeds the nominal segment size.
+SEGMENT_SIZE = int(os.getenv("SEGMENT_SIZE_BYTES", str(32 * 1024 * 1024)))
+LOCAL_BUFFER_SIZE = int(os.getenv("LOCAL_BUFFER_SIZE_BYTES", str(64 * 1024 * 1024)))
+
+# Wait window for a full master+client heartbeat round-trip. The default
+# heartbeat interval is 10s on both sides, so 25s is comfortably > 2 ticks.
+PROMOTION_WAIT_SECONDS = int(os.getenv("PROMOTION_WAIT_SECONDS", "25"))
+
+
+def get_client(store):
+    protocol = os.getenv("PROTOCOL", "tcp")
+    device_name = os.getenv("DEVICE_NAME", "")
+    local_hostname = os.getenv("LOCAL_HOSTNAME", "localhost")
+    metadata_server = os.getenv("MC_METADATA_SERVER", "P2PHANDSHAKE")
+    master_server_address = os.getenv("MASTER_SERVER", "127.0.0.1:50051")
+
+    retcode = store.setup(
+        local_hostname,
+        metadata_server,
+        SEGMENT_SIZE,
+        LOCAL_BUFFER_SIZE,
+        protocol,
+        device_name,
+        master_server_address,
+        None,  # engine
+        True,  # enable_ssd_offload — required for FileStorage / LOCAL_DISK
+    )
+    if retcode:
+        raise RuntimeError(f"Failed to setup store client. Return code: {retcode}")
+
+
+def _replica_types(descs, key):
+    """Return a list of replica type tags ('MEMORY', 'LOCAL_DISK', 'DISK')
+    for a key from a batch_get_replica_desc result. Tolerant to either
+    list-of-info or single-info shape."""
+    infos = descs.get(key) if isinstance(descs, dict) else None
+    if infos is None:
+        return []
+    if not isinstance(infos, (list, tuple)):
+        infos = [infos]
+    tags = []
+    for info in infos:
+        if hasattr(info, "is_memory_replica") and info.is_memory_replica():
+            tags.append("MEMORY")
+        elif hasattr(info, "is_local_disk_replica") and info.is_local_disk_replica():
+            tags.append("LOCAL_DISK")
+        elif hasattr(info, "is_disk_replica") and info.is_disk_replica():
+            tags.append("DISK")
+        else:
+            tags.append("UNKNOWN")
+    return tags
+
+
+class TestPromotionOnHit(unittest.TestCase):
+    """Python-binding test for the L2->L1 promotion-on-hit behavioral
+    contract."""
+
+    @classmethod
+    def setUpClass(cls):
+        cls.store = MooncakeDistributedStore()
+        get_client(cls.store)
+
+    def test_promotion_after_repeated_hits(self):
+        """A LOCAL_DISK-only key must regain a MEMORY replica after
+        clearing the admission threshold via repeated reads.
+
+        A failure here either means the master is not running with
+        ``promotion_on_hit: true``, or the FileStorage promotion executor
+        did not run / did not complete byte movement before the timeout."""
+
+        VALUE_SIZE = 1024 * 1024  # 1 MB
+        # Enough keys to comfortably overflow a 64 MB segment and trigger
+        # offload-on-evict for at least the older portion of the workload.
+        NUM_KEYS = 96
+
+        timestamp = int(time.time())
+        keys = [f"poh_{i}_{timestamp}" for i in range(NUM_KEYS)]
+        reference = {}
+
+        try:
+            # Phase 1: overflow DRAM, force offload-on-evict.
+            for key in keys:
+                value = os.urandom(VALUE_SIZE)
+                retcode = self.store.put(key, value)
+                # NO_AVAILABLE_HANDLE-style failures (-200) under pressure
+                # are expected and not fatal — we just need *some* keys to
+                # successfully land on LOCAL_DISK.
+                if retcode == 0:
+                    reference[key] = value
+
+            self.assertGreater(
+                len(reference), 0, "No PUTs succeeded — cannot run promotion test"
+            )
+
+            # Phase 2: wait long enough for offload heartbeat to flush the
+            # evicted keys to LOCAL_DISK.
+            time.sleep(PROMOTION_WAIT_SECONDS)
+
+            descs = self.store.batch_get_replica_desc(list(reference.keys()))
+            type_hist = {}
+            cold_key = None
+            for key in reference.keys():
+                types = _replica_types(descs, key)
+                hist_key = ",".join(sorted(set(types)))
+                type_hist[hist_key] = type_hist.get(hist_key, 0) + 1
+                if (
+                    cold_key is None
+                    and types
+                    and all("MEMORY" not in t for t in types)
+                    and any("LOCAL_DISK" in t for t in types)
+                ):
+                    cold_key = key
+            print(f"replica-type histogram after phase 2: {type_hist}")
+
+            self.assertIsNotNone(
+                cold_key,
+                "No LOCAL_DISK-only key found after eviction — is "
+                "offload_on_evict=true and the segment small enough to "
+                "overflow? master config / SEGMENT_SIZE_BYTES env may need "
+                "tuning. Histogram above shows the actual replica state.",
+            )
+
+            # Phase 3: clear the admission threshold via per-key reads.
+            # ``store.get`` goes through ``Client::Query`` → master's
+            # ``GetReplicaList``, so each call fires the promotion gate
+            # once; 4 calls comfortably clear any reasonable admission
+            # threshold. We also assert bit-exact bytes back, exercising
+            # the LOCAL_DISK read path end-to-end via the offload-RPC
+            # route (this transitively guards the read-side fix the
+            # promotion feature relies on; reads only return the
+            # offloaded bytes if the peer's offload-RPC + segment-name
+            # resolution work).
+            expected_bytes = reference[cold_key]
+            for _ in range(4):
+                got = self.store.get(cold_key)
+                self.assertEqual(
+                    got,
+                    expected_bytes,
+                    f"store.get on LOCAL_DISK-only key {cold_key} returned "
+                    f"wrong/empty bytes (got len={len(got) if got else 0}, "
+                    f"expected len={len(expected_bytes)}). The LOCAL_DISK "
+                    f"read path is broken — promotion cannot be tested.",
+                )
+
+            # Phase 4: wait for the master to enqueue the promotion task and
+            # for the client's next FileStorage heartbeat to execute it.
+            time.sleep(PROMOTION_WAIT_SECONDS)
+
+            descs_after = self.store.batch_get_replica_desc([cold_key])
+            types_after = _replica_types(descs_after, cold_key)
+            print(f"replica types for {cold_key} after promotion: {types_after}")
+            self.assertTrue(
+                "MEMORY" in types_after,
+                f"Expected a MEMORY replica for {cold_key} after promotion, "
+                f"got types={types_after}. Either promotion_on_hit is not "
+                f"enabled in the master, the admission threshold was not "
+                f"cleared, or the heartbeat tick did not fire within "
+                f"{PROMOTION_WAIT_SECONDS}s.",
+            )
+
+            # Phase 5: post-promotion bytes-back AND prove MEMORY was the
+            # replica actually served. SelectBestReplica should prefer the
+            # newly-COMPLETE MEMORY replica over the older LOCAL_DISK; if
+            # it does, no offload-RPC call is issued for this read.
+            #
+            # offload_rpc_read_count counts every invocation of
+            # batch_get_into_offload_object_internal — the single
+            # chokepoint for LOCAL_DISK reads served via peer offload-RPC.
+            # We snapshot the counter, do a read, then assert the counter
+            # didn't move. Bytes-back alone wouldn't distinguish MEMORY
+            # from LOCAL_DISK (the offload-RPC path returns correct bytes
+            # too); the counter delta is the hard correctness signal that
+            # the MEMORY path was used. The BenchPromotionLatency class
+            # below offers a heavier latency-based view of the same
+            # invariant when richer numbers are wanted.
+            offload_count_before = self.store.get_offload_rpc_read_count()
+            self.assertEqual(
+                self.store.get(cold_key),
+                expected_bytes,
+                f"post-promotion store.get on {cold_key} returned wrong "
+                f"bytes; the new MEMORY replica is metadata-visible but "
+                f"not byte-correct. NotifyPromotionSuccess flipped "
+                f"PROCESSING -> COMPLETE before the TE write finished, "
+                f"or the TE write landed at the wrong address.",
+            )
+            offload_count_after = self.store.get_offload_rpc_read_count()
+            self.assertEqual(
+                offload_count_after,
+                offload_count_before,
+                f"post-promotion store.get on {cold_key} bumped the "
+                f"offload-RPC counter from {offload_count_before} to "
+                f"{offload_count_after}: it served from LOCAL_DISK SSD, "
+                f"not the freshly-promoted MEMORY replica. The RFC's "
+                f"'subsequent reads avoid SSD' invariant is broken; "
+                f"check SelectBestReplica preference and the post-"
+                f"promotion replica order in metadata.",
+            )
+        finally:
+            for key in keys:
+                try:
+                    self.store.remove(key)
+                except Exception:
+                    pass
+            time.sleep(default_kv_lease_ttl / 1000 + 0.5)
+
+
+class TestPromotionDoesNotFire(unittest.TestCase):
+    """Negative-case e2e. With workloads that don't produce LOCAL_DISK
+    replicas, promotion has nothing to do and must not spuriously stage
+    MEMORY replicas. This is the inverse invariant of the positive test
+    and the only Python-level guard against:
+      - regressions that auto-create LOCAL_DISK at PutEnd (would silently
+        break offload semantics and trigger spurious promotions),
+      - regressions that fire promotion without a LOCAL_DISK source.
+    """
+
+    @classmethod
+    def setUpClass(cls):
+        cls.store = MooncakeDistributedStore()
+        get_client(cls.store)
+
+    def test_below_watermark_workload_stays_memory_only(self):
+        VALUE_SIZE = 1024  # 1 KB
+        NUM_KEYS = 16  # 16 KB total, well below the 32 MB segment
+
+        timestamp = int(time.time())
+        keys = [f"neg_poh_small_{i}_{timestamp}" for i in range(NUM_KEYS)]
+
+        try:
+            for key in keys:
+                value = os.urandom(VALUE_SIZE)
+                self.assertEqual(self.store.put(key, value), 0)
+
+            # One heartbeat cycle is plenty for any spurious offload to
+            # have surfaced.
+            time.sleep(PROMOTION_WAIT_SECONDS)
+
+            descs = self.store.batch_get_replica_desc(keys)
+            unexpected = []
+            for key in keys:
+                types = _replica_types(descs, key)
+                if not any(t == "MEMORY" for t in types):
+                    unexpected.append((key, types))
+                if any(t == "LOCAL_DISK" for t in types):
+                    unexpected.append((key, types))
+
+            self.assertEqual(
+                unexpected,
+                [],
+                f"Below-watermark workload produced unexpected non-MEMORY "
+                f"or LOCAL_DISK replicas: {unexpected[:5]}. Either "
+                f"eviction is firing prematurely or PutEnd is auto-"
+                f"creating LOCAL_DISK.",
+            )
+        finally:
+            for key in keys:
+                try:
+                    self.store.remove(key)
+                except Exception:
+                    pass
+            time.sleep(default_kv_lease_ttl / 1000 + 0.5)
+
+
+@unittest.skipUnless(
+    os.getenv("MC_BENCH_PROMOTION_LATENCY"),
+    "opt-in benchmark — set MC_BENCH_PROMOTION_LATENCY=1 to run. "
+    "Not part of CI; gives p50/p95/p99 latency comparison of LOCAL_DISK "
+    "reads (pre-promotion) vs MEMORY reads (post-promotion).",
+)
+class BenchPromotionLatency(unittest.TestCase):
+    """Ad-hoc latency benchmark for the L2->L1 promotion feature.
+
+    Runs the same overflow → eviction → promote workflow as
+    ``TestPromotionOnHit`` but:
+      - times every read with ``time.perf_counter``;
+      - aborts the pre-promotion sample loop the moment a read serves
+        from MEMORY (detected via the offload-RPC counter) to keep the
+        LOCAL_DISK distribution uncontaminated;
+      - prints p50/p95/p99 + min/max for both phases and a per-percentile
+        speedup table;
+      - asserts a conservative 1.3x p50 speedup so a regression where
+        MEMORY isn't being preferred breaks the bench.
+
+    Sample size is controlled by ``LATENCY_SAMPLES`` (default 1000).
+    For larger samples (e.g. N=10000) bump
+    ``MOONCAKE_OFFLOAD_HEARTBEAT_INTERVAL_SECONDS`` so the worker
+    doesn't race the pre-promotion read loop.
+
+    Expected speedup on docker tmpfs with 1 MiB objects is ~2x at every
+    percentile (most wall-time is Python/C++ marshalling, not the
+    SSD-vs-DRAM cost gap). On real SSDs + RDMA NICs expect 10-50x.
+    """
+
+    @classmethod
+    def setUpClass(cls):
+        cls.store = MooncakeDistributedStore()
+        get_client(cls.store)
+
+    def test_latency_comparison(self):
+        VALUE_SIZE = 1024 * 1024  # 1 MB
+        NUM_KEYS = 96
+
+        timestamp = int(time.time())
+        keys = [f"bench_poh_{i}_{timestamp}" for i in range(NUM_KEYS)]
+        reference = {}
+
+        try:
+            # Overflow DRAM, force offload-on-evict.
+            for key in keys:
+                value = os.urandom(VALUE_SIZE)
+                if self.store.put(key, value) == 0:
+                    reference[key] = value
+            self.assertGreater(len(reference), 0)
+
+            time.sleep(PROMOTION_WAIT_SECONDS)
+
+            descs = self.store.batch_get_replica_desc(list(reference.keys()))
+            cold_key = None
+            for key in reference.keys():
+                types = _replica_types(descs, key)
+                if (
+                    cold_key is None
+                    and types
+                    and all("MEMORY" not in t for t in types)
+                    and any("LOCAL_DISK" in t for t in types)
+                ):
+                    cold_key = key
+            self.assertIsNotNone(cold_key, "no LOCAL_DISK-only key after eviction")
+
+            # Pre-promotion: timed LOCAL_DISK reads, bail when a read
+            # serves from MEMORY (= worker raced the loop).
+            N = int(os.getenv("LATENCY_SAMPLES", "1000"))
+            expected_bytes = reference[cold_key]
+            pre_latencies_ms = []
+            for i in range(N):
+                count_before = self.store.get_offload_rpc_read_count()
+                t0 = time.perf_counter()
+                got = self.store.get(cold_key)
+                t1 = time.perf_counter()
+                count_after = self.store.get_offload_rpc_read_count()
+                self.assertEqual(got, expected_bytes)
+                if count_after == count_before:
+                    print(
+                        f"  pre-promotion loop: promotion fired at "
+                        f"sample {i}/{N}; stopping to keep the "
+                        f"LOCAL_DISK sample uncontaminated"
+                    )
+                    break
+                pre_latencies_ms.append((t1 - t0) * 1000.0)
+            self.assertGreater(
+                len(pre_latencies_ms),
+                1,
+                "collected <2 pre-promotion samples — promotion fired "
+                "before measurement could capture LOCAL_DISK latency. "
+                "Bump MOONCAKE_OFFLOAD_HEARTBEAT_INTERVAL_SECONDS or "
+                "lower LATENCY_SAMPLES.",
+            )
+
+            # Wait for promotion to complete.
+            time.sleep(PROMOTION_WAIT_SECONDS)
+            descs_after = self.store.batch_get_replica_desc([cold_key])
+            types_after = _replica_types(descs_after, cold_key)
+            self.assertIn("MEMORY", types_after)
+
+            # Post-promotion: timed MEMORY reads; the counter must stay
+            # put (otherwise at least one read still hit LOCAL_DISK).
+            count_before = self.store.get_offload_rpc_read_count()
+            post_latencies_ms = []
+            for _ in range(N):
+                t0 = time.perf_counter()
+                got = self.store.get(cold_key)
+                t1 = time.perf_counter()
+                self.assertEqual(got, expected_bytes)
+                post_latencies_ms.append((t1 - t0) * 1000.0)
+            count_after = self.store.get_offload_rpc_read_count()
+            self.assertEqual(
+                count_after,
+                count_before,
+                f"post-promotion reads bumped the offload-RPC counter "
+                f"from {count_before} to {count_after}: at least one "
+                f"read served from LOCAL_DISK SSD instead of MEMORY.",
+            )
+
+            def _pcts(samples):
+                """{p: value_ms} for p in [50, 95, 99]."""
+                if len(samples) >= 2:
+                    cuts = statistics.quantiles(samples, n=100)
+                    return {50: cuts[49], 95: cuts[94], 99: cuts[98]}
+                v = samples[0] if samples else 0.0
+                return {50: v, 95: v, 99: v}
+
+            def _speedup(a, b):
+                return a / b if b > 0 else float("inf")
+
+            pre = _pcts(pre_latencies_ms)
+            post = _pcts(post_latencies_ms)
+
+            print()
+            print(f"=== promotion latency comparison "
+                  f"(cold_key={cold_key}) ===")
+            print(
+                f"pre-promotion  (LOCAL_DISK via offload-RPC): "
+                f"n={len(pre_latencies_ms)}, "
+                f"p50={pre[50]:.2f} ms, p95={pre[95]:.2f} ms, "
+                f"p99={pre[99]:.2f} ms, "
+                f"min={min(pre_latencies_ms):.2f} ms, "
+                f"max={max(pre_latencies_ms):.2f} ms"
+            )
+            print(
+                f"post-promotion (MEMORY direct):              "
+                f"n={len(post_latencies_ms)}, "
+                f"p50={post[50]:.2f} ms, p95={post[95]:.2f} ms, "
+                f"p99={post[99]:.2f} ms, "
+                f"min={min(post_latencies_ms):.2f} ms, "
+                f"max={max(post_latencies_ms):.2f} ms"
+            )
+            print(
+                f"speedup: p50={_speedup(pre[50], post[50]):.1f}x  "
+                f"|  p95={_speedup(pre[95], post[95]):.1f}x  "
+                f"|  p99={_speedup(pre[99], post[99]):.1f}x"
+            )
+
+            self.assertLess(
+                post[50],
+                pre[50] / 1.3,
+                f"post-promotion p50 latency ({post[50]:.2f} ms) is not "
+                f"meaningfully lower than pre-promotion p50 "
+                f"({pre[50]:.2f} ms). Real-hardware speedup is typically "
+                f"10–50x; failing this 1.3x check means MEMORY likely "
+                f"isn't being preferred.",
+            )
+        finally:
+            for key in keys:
+                try:
+                    self.store.remove(key)
+                except Exception:
+                    pass
+            time.sleep(default_kv_lease_ttl / 1000 + 0.5)
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/mooncake-wheel/tests/test_put_get_tensor.py b/mooncake-wheel/tests/test_put_get_tensor.py
index 38733b9b..0c84f65a 100644
--- a/mooncake-wheel/tests/test_put_get_tensor.py
+++ b/mooncake-wheel/tests/test_put_get_tensor.py
@@ -153,5 +153,276 @@ class TestDistributedObjectStore(unittest.TestCase):
         self.store.remove(key_3d)
 
              
+    def _assert_empty_tensor_equal(self, expected, actual):
+        import torch
+
+        self.assertIsNotNone(actual)
+        self.assertEqual(tuple(actual.shape), tuple(expected.shape))
+        self.assertEqual(actual.dtype, expected.dtype)
+        self.assertEqual(actual.numel(), 0)
+        self.assertTrue(torch.equal(actual, expected))
+
+    def test_put_get_zero_tensor(self):
+        """Test storing and retrieving zero-sized PyTorch tensors."""
+        import torch
+
+        tensors = [
+            torch.empty((0,), dtype=torch.float32),
+            torch.empty((0, 4096), dtype=torch.float16),
+            torch.empty((2, 0, 3), dtype=torch.bfloat16),
+            torch.empty((1, 0), dtype=torch.int64),
+        ]
+        keys = [f"test_zero_tensor_{os.getpid()}_{i}" for i in range(len(tensors))]
+
+        for key, tensor in zip(keys, tensors):
+            self.assertEqual(self.store.put_tensor(key, tensor), 0)
+            self._assert_empty_tensor_equal(tensor, self.store.get_tensor(key))
+
+        batch_results = self.store.batch_get_tensor(keys)
+        self.assertEqual(len(batch_results), len(tensors))
+        for tensor, retrieved in zip(tensors, batch_results):
+            self._assert_empty_tensor_equal(tensor, retrieved)
+
+        for key in keys:
+            self.store.remove(key)
+
+    def test_zero_tensor_into_and_from(self):
+        """Test metadata-only zero tensor get_into and put_tensor_from."""
+        import torch
+
+        tensor = torch.empty((2, 0, 3), dtype=torch.float32)
+        key = f"test_zero_tensor_into_{os.getpid()}"
+        key_from = f"test_zero_tensor_from_{os.getpid()}"
+        self.assertEqual(self.store.put_tensor(key, tensor), 0)
+
+        buffer_size = 4096
+        buffer = (ctypes.c_ubyte * buffer_size)()
+        buffer_ptr = ctypes.addressof(buffer)
+        self.assertEqual(self.store.register_buffer(buffer_ptr, buffer_size), 0)
+        try:
+            total_length = self.store.get_size(key)
+            self.assertGreater(total_length, 0)
+            self.assertLess(total_length, buffer_size)
+            retrieved = self.store.get_tensor_into(key, buffer_ptr, buffer_size)
+            self._assert_empty_tensor_equal(tensor, retrieved)
+
+            self.assertEqual(self.store.put_tensor_from(key_from, buffer_ptr, total_length), 0)
+            self._assert_empty_tensor_equal(tensor, self.store.get_tensor(key_from))
+        finally:
+            self.store.unregister_buffer(buffer_ptr)
+            self.store.remove(key)
+            self.store.remove(key_from)
+
+    def test_zero_tensor_with_tp(self):
+        """Test zero-sized tensors through legacy tensor-parallel APIs."""
+        import torch
+
+        tensor = torch.empty((2, 0, 3), dtype=torch.float32)
+        key = f"test_zero_tensor_tp_{os.getpid()}"
+        tp_size = 2
+        split_dim = 0
+
+        self.assertEqual(
+            self.store.put_tensor_with_tp(key, tensor, tp_size=tp_size, split_dim=split_dim),
+            0,
+        )
+        for rank in range(tp_size):
+            shard = self.store.get_tensor_with_tp(key, tp_rank=rank, tp_size=tp_size)
+            expected = tensor.narrow(split_dim, rank, 1).contiguous()
+            self._assert_empty_tensor_equal(expected, shard)
+
+        for rank in range(tp_size):
+            self.store.remove(f"{key}_tp_{rank}")
+
+
+from mooncake.structured_object_store import (
+    MISSING,
+    _choose_leaf_codec,
+    _escape_key,
+    infer_structure,
+)
+
+
+class TestCodecInference(unittest.TestCase):
+
+    def test_tensor(self):
+        import torch
+        d = _choose_leaf_codec([torch.tensor([1, 2]), torch.tensor([3])])
+        self.assertTrue(d.accepted)
+        self.assertEqual(d.codec, "ragged_tensor")
+
+    def test_tensor_mixed_dtype_rejected(self):
+        import torch
+        d = _choose_leaf_codec([torch.tensor([1], dtype=torch.float32), torch.tensor([1], dtype=torch.int64)])
+        self.assertFalse(d.accepted)
+
+    def test_tensor_mixed_ndim_rejected(self):
+        import torch
+        d = _choose_leaf_codec([torch.tensor([1]), torch.tensor([[1, 2]])])
+        self.assertFalse(d.accepted)
+
+    def test_numeric_sequence(self):
+        d = _choose_leaf_codec([[1, 2, 3], [4, 5]])
+        self.assertTrue(d.accepted)
+        self.assertEqual(d.codec, "typed_ragged")
+
+    def test_bytes(self):
+        d = _choose_leaf_codec([b"hello", b"world"])
+        self.assertTrue(d.accepted)
+        self.assertEqual(d.codec, "bytes_ragged")
+
+    def test_text(self):
+        d = _choose_leaf_codec(["hello", "world"])
+        self.assertTrue(d.accepted)
+        self.assertEqual(d.codec, "utf8_ragged")
+
+    def test_json(self):
+        d = _choose_leaf_codec([{"a": 1}, {"b": 2}])
+        self.assertTrue(d.accepted)
+        self.assertEqual(d.codec, "msgpack_ragged")
+
+    def test_json_rejects_late_non_serializable(self):
+        values = [{"ok": i} for i in range(200)] + [object()]
+        d = _choose_leaf_codec(values)
+        self.assertFalse(d.accepted)
+
+    def test_scalar(self):
+        d = _choose_leaf_codec([1, 2.0, 3])
+        self.assertTrue(d.accepted)
+        self.assertEqual(d.codec, "ndarray")
+
+    def test_fallback(self):
+        d = _choose_leaf_codec([object(), object()])
+        self.assertFalse(d.accepted)
+        self.assertEqual(d.codec, "pickle_ragged_fallback")
+
+    def test_with_nulls(self):
+        d = _choose_leaf_codec(["hello", None, "world"])
+        self.assertTrue(d.accepted)
+        self.assertEqual(d.codec, "utf8_ragged")
+
+    def test_empty_values(self):
+        d = _choose_leaf_codec([])
+        self.assertFalse(d.accepted)
+
+    def test_all_none(self):
+        d = _choose_leaf_codec([None, None, None])
+        self.assertFalse(d.accepted)
+
+    def test_infer_flat(self):
+        leaves, nodes = [], []
+        infer_structure("root", ["a", "b", "c"], leaves, nodes)
+        self.assertEqual(len(leaves), 1)
+        self.assertEqual(len(nodes), 0)
+        self.assertEqual(leaves[0].decision.codec, "utf8_ragged")
+
+    def test_infer_dict(self):
+        leaves, nodes = [], []
+        infer_structure("root", [{"x": 1, "y": "a"}, {"x": 2, "y": "b"}], leaves, nodes)
+        self.assertEqual(len(nodes), 1)
+        self.assertEqual(nodes[0].node_type, "dict")
+        self.assertEqual(sorted(nodes[0].children), ["x", "y"])
+        self.assertEqual(sorted(leaf.path for leaf in leaves), ["root.x", "root.y"])
+
+    def test_infer_dict_with_none_rows(self):
+        leaves, nodes = [], []
+        infer_structure("r", [{"x": 1}, None, {"x": 3}], leaves, nodes)
+        self.assertEqual(len(nodes), 1)
+        self.assertEqual(len(leaves), 1)
+        self.assertEqual(leaves[0].path, "r.x")
+
+    def test_infer_nested(self):
+        leaves, nodes = [], []
+        infer_structure("r", [{"a": {"b": 1}}, {"a": {"b": 2}}], leaves, nodes)
+        self.assertEqual(len(nodes), 2)
+        self.assertEqual(leaves[0].path, "r.a.b")
+        self.assertEqual(leaves[0].decision.codec, "ndarray")
+
+    def test_infer_list(self):
+        leaves, nodes = [], []
+        infer_structure("r", [[{"k": 1}], [{"k": 2}]], leaves, nodes)
+        list_nodes = [n for n in nodes if n.node_type == "list"]
+        self.assertEqual(len(list_nodes), 1)
+        self.assertEqual(list_nodes[0].lengths, [1, 1])
+        dict_nodes = [n for n in nodes if n.node_type == "dict"]
+        self.assertEqual(len(dict_nodes), 1)
+        self.assertEqual(len(leaves), 1)
+        self.assertEqual(leaves[0].path, "r[0].k")
+
+    def test_infer_list_with_none_items(self):
+        leaves, nodes = [], []
+        infer_structure("r", [[{"k": 1}, None], [{"k": 2}, None]], leaves, nodes)
+        list_nodes = [n for n in nodes if n.node_type == "list"]
+        self.assertEqual(len(list_nodes), 1)
+
+    def test_dict_missing_key_vs_none_value(self):
+        leaves, nodes = [], []
+        infer_structure("r", [{"x": None}, {"y": 2}, None], leaves, nodes)
+        self.assertIsNotNone(nodes[0].row_mask)
+        self.assertEqual(nodes[0].row_mask, [True, True, False])
+        x_leaf = next(leaf for leaf in leaves if leaf.path == "r.x")
+        self.assertIsNone(x_leaf.values[0])
+        self.assertIsInstance(x_leaf.values[1], type(MISSING))
+        self.assertIsNone(x_leaf.values[2])
+
+    def test_infer_dict_of_tensors(self):
+        import torch
+
+        leaves, nodes = [], []
+        infer_structure(
+            "r",
+            [
+                {"tokens": torch.arange(2, dtype=torch.int64), "score": 1.0},
+                {"tokens": torch.arange(3, dtype=torch.int64), "score": 2.0},
+                {"tokens": torch.arange(1, dtype=torch.int64), "score": None},
+            ],
+            leaves,
+            nodes,
+        )
+        self.assertEqual(len(nodes), 1)
+        self.assertEqual(nodes[0].node_type, "dict")
+        by_path = {leaf.path: leaf for leaf in leaves}
+        self.assertEqual(by_path["r.tokens"].decision.codec, "ragged_tensor")
+        self.assertEqual(by_path["r.score"].decision.codec, "ndarray")
+
+    def test_infer_dict_of_tensors_missing_keys_and_null_rows(self):
+        import torch
+
+        leaves, nodes = [], []
+        rows = [
+            {"tokens": torch.arange(2, dtype=torch.float32), "label": None},
+            None,
+            {"label": 3},
+            {"tokens": None, "label": 4},
+        ]
+        infer_structure("r", rows, leaves, nodes)
+        self.assertEqual(nodes[0].row_mask, [True, False, True, True])
+        by_path = {leaf.path: leaf for leaf in leaves}
+        tokens = by_path["r.tokens"]
+        self.assertTrue(torch.equal(tokens.values[0], rows[0]["tokens"]))
+        self.assertIsNone(tokens.values[1])
+        self.assertIsInstance(tokens.values[2], type(MISSING))
+        self.assertIsNone(tokens.values[3])
+        self.assertEqual(tokens.decision.codec, "ragged_tensor")
+
+    def test_escape_key_in_path(self):
+        self.assertEqual(_escape_key("simple"), "simple")
+        self.assertEqual(_escape_key("a.b"), "a\\.b")
+        self.assertEqual(_escape_key("a[0]"), "a\\[0]")
+        self.assertEqual(_escape_key("a\\b"), "a\\\\b")
+
+    def test_dict_with_dot_key(self):
+        leaves, nodes = [], []
+        infer_structure("r", [{"a.b": 1}, {"a.b": 2}], leaves, nodes)
+        self.assertEqual(leaves[0].path, "r.a\\.b")
+
+    def test_depth_limit(self):
+        deep = 1
+        for _ in range(40):
+            deep = {"a": deep}
+        with self.assertRaises(ValueError):
+            infer_structure("r", [deep, deep], [], [])
+
+
 if __name__ == '__main__':
-    unittest.main()
\ No newline at end of file
+    unittest.main()
diff --git a/mooncake-wheel/tests/test_structured_object_store.py b/mooncake-wheel/tests/test_structured_object_store.py
new file mode 100644
index 00000000..d3f83376
--- /dev/null
+++ b/mooncake-wheel/tests/test_structured_object_store.py
@@ -0,0 +1,2734 @@
+from __future__ import annotations
+
+import ctypes
+import json
+import os
+import threading
+import time
+
+import numpy as np
+import pytest
+
+from mooncake.structured_object_store import (
+    BundleTransferPolicy,
+    MooncakeBundleTransfer,
+    RemoteBundleRef,
+    StructuredObjectPayload,
+    StructuredObjectReadSpec,
+    export_dataproto_ref,
+    import_dataproto_ref,
+    is_dataproto_ref_handle,
+    raw_destination,
+    tensor_object_buffer,
+)
+
+
+class SimpleDataProto:
+    def __init__(self, batch=None, non_tensor_batch=None, meta_info=None) -> None:
+        self.batch = {} if batch is None else batch
+        self.non_tensor_batch = {} if non_tensor_batch is None else non_tensor_batch
+        self.meta_info = {} if meta_info is None else meta_info
+
+    @classmethod
+    def from_dict(cls, batch, non_tensor_batch=None, meta_info=None):
+        return cls(batch=batch, non_tensor_batch=non_tensor_batch, meta_info=meta_info)
+
+
+class BadDataProto:
+    def __init__(self) -> None:
+        pass
+
+
+class InMemoryStore:
+    def __init__(self) -> None:
+        self.objects: dict[str, bytes] = {}
+        self.tensor_objects: dict[str, object] = {}
+        self.lock = threading.Lock()
+        self.registered: set[int] = set()
+        self.register_buffer_calls = 0
+        self.unregister_buffer_calls = 0
+        self.max_active_puts = 0
+        self.max_active_gets = 0
+        self.active_puts = 0
+        self.active_gets = 0
+        self.get_into_calls = 0
+        self.get_into_ranges_calls = 0
+        self.batch_get_into_calls = 0
+        self.batch_put_from_calls = 0
+        self.put_tensor_from_calls = 0
+        self.batch_remove_calls = 0
+        self.put_tensor_calls = 0
+        self.get_tensor_calls = 0
+
+    def _enter_put(self) -> None:
+        with self.lock:
+            self.active_puts += 1
+            self.max_active_puts = max(self.max_active_puts, self.active_puts)
+
+    def _exit_put(self) -> None:
+        with self.lock:
+            self.active_puts -= 1
+
+    def _enter_get(self, count: int = 1) -> None:
+        with self.lock:
+            self.active_gets += count
+            self.max_active_gets = max(self.max_active_gets, self.active_gets)
+
+    def _exit_get(self, count: int = 1) -> None:
+        with self.lock:
+            self.active_gets -= count
+
+    def put(self, key: str, value) -> int:
+        self._enter_put()
+        try:
+            time.sleep(0.01)
+            with self.lock:
+                self.objects[key] = bytes(value)
+            return 0
+        finally:
+            self._exit_put()
+
+    def get(self, key: str) -> bytes:
+        self._enter_get()
+        try:
+            time.sleep(0.01)
+            with self.lock:
+                return self.objects[key]
+        finally:
+            self._exit_get()
+
+    def remove(self, key: str, force: bool = False) -> int:
+        with self.lock:
+            self.objects.pop(key, None)
+            self.tensor_objects.pop(key, None)
+        return 0
+
+    def put_tensor(self, key: str, value) -> int:
+        self.put_tensor_calls += 1
+        with self.lock:
+            self.tensor_objects[key] = value.detach().clone()
+        return 0
+
+    def get_tensor(self, key: str):
+        self.get_tensor_calls += 1
+        with self.lock:
+            return self.tensor_objects[key].clone()
+
+    def batch_remove(self, keys: list[str], force: bool = False) -> list[int]:
+        self.batch_remove_calls += 1
+        for key in keys:
+            self.remove(key, force)
+        return [0 for _key in keys]
+
+    def batch_put_from(
+        self, keys: list[str], buffer_ptrs: list[int], sizes: list[int]
+    ) -> list[int]:
+        self.batch_put_from_calls += 1
+        results: list[int] = []
+        for key, ptr, size in zip(keys, buffer_ptrs, sizes):
+            data = ctypes.string_at(ptr, size)
+            results.append(self.put(key, data))
+        return results
+
+    def put_tensor_from(self, key: str, buffer_ptr: int, size: int) -> int:
+        self.put_tensor_from_calls += 1
+        if buffer_ptr not in self.registered:
+            return -1
+        return self.put(key, ctypes.string_at(buffer_ptr, size))
+
+    def register_buffer(self, buffer_ptr: int, size: int) -> int:
+        self.register_buffer_calls += 1
+        self.registered.add(buffer_ptr)
+        return 0
+
+    def unregister_buffer(self, buffer_ptr: int) -> int:
+        self.unregister_buffer_calls += 1
+        self.registered.remove(buffer_ptr)
+        return 0
+
+    def get_into(self, key: str, ptr: int, size: int) -> int:
+        self.get_into_calls += 1
+        self._enter_get()
+        try:
+            time.sleep(0.01)
+            with self.lock:
+                data = self.objects[key]
+            if len(data) > size:
+                return -1
+            ctypes.memmove(ptr, data, len(data))
+            return len(data)
+        finally:
+            self._exit_get()
+
+    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]]]:
+        self.get_into_ranges_calls += 1
+        total_keys = sum(len(keys) for keys in all_keys)
+        self._enter_get(total_keys)
+        try:
+            time.sleep(0.01)
+            results: list[list[list[int]]] = []
+            for base_ptr, keys, dst_groups, src_groups, size_groups in zip(
+                buffer_ptrs, all_keys, all_dst_offsets, all_src_offsets, all_sizes
+            ):
+                buffer_results: list[list[int]] = []
+                for key, dst_offsets, src_offsets, sizes in zip(
+                    keys, dst_groups, src_groups, size_groups
+                ):
+                    with self.lock:
+                        data = self.objects[key]
+                    key_results: list[int] = []
+                    for dst_offset, src_offset, size in zip(
+                        dst_offsets, src_offsets, sizes
+                    ):
+                        end = src_offset + size
+                        if end > len(data):
+                            key_results.append(-1)
+                            continue
+                        ctypes.memmove(
+                            base_ptr + dst_offset, data[src_offset:end], size
+                        )
+                        key_results.append(size)
+                    buffer_results.append(key_results)
+                results.append(buffer_results)
+            return results
+        finally:
+            self._exit_get(total_keys)
+
+    def batch_get_into(
+        self, keys: list[str], ptrs: list[int], sizes: list[int]
+    ) -> list[int]:
+        self.batch_get_into_calls += 1
+        self._enter_get(len(keys))
+        try:
+            time.sleep(0.01)
+            results = []
+            for key, ptr, size in zip(keys, ptrs, sizes):
+                with self.lock:
+                    data = self.objects[key]
+                if len(data) > size:
+                    results.append(-1)
+                    continue
+                ctypes.memmove(ptr, data, len(data))
+                results.append(len(data))
+            return results
+        finally:
+            self._exit_get(len(keys))
+
+
+class GetOnlyStore(InMemoryStore):
+    batch_get_into = None
+
+
+class MinimalStore(GetOnlyStore):
+    batch_put_from = None
+    register_buffer = None
+    unregister_buffer = None
+
+
+class ReadFastPathWithoutRegisterStore(MinimalStore):
+    batch_get_into = InMemoryStore.batch_get_into
+    get_into = InMemoryStore.get_into
+    get_into_ranges = InMemoryStore.get_into_ranges
+
+
+class PlainStore(MinimalStore):
+    get_into = None
+    get_into_ranges = None
+
+
+class PutTensorOnlyStore(InMemoryStore):
+    get_tensor = None
+
+
+class NoTensorFastPathStore(InMemoryStore):
+    put_tensor = None
+    get_tensor = None
+
+
+class FailingPutStore(InMemoryStore):
+    def __init__(self, fail_on_put: int) -> None:
+        super().__init__()
+        self.put_count = 0
+        self.fail_on_put = fail_on_put
+
+    def put(self, key: str, value) -> int:
+        self.put_count += 1
+        if self.put_count == self.fail_on_put:
+            raise RuntimeError("injected put failure")
+        return super().put(key, value)
+
+
+class FailingBatchGetStore(InMemoryStore):
+    def __init__(self) -> None:
+        super().__init__()
+        self.fail_key = ""
+
+    def batch_get_into(
+        self, keys: list[str], ptrs: list[int], sizes: list[int]
+    ) -> list[int]:
+        if self.fail_key in keys:
+            raise RuntimeError("injected get failure")
+        return super().batch_get_into(keys, ptrs, sizes)
+
+
+class FailingRemoveStore(FailingPutStore):
+    def remove(self, key: str, force: bool = False) -> int:
+        raise RuntimeError("injected remove failure")
+
+    batch_remove = None
+
+
+class ForceTrackingStore(InMemoryStore):
+    def __init__(self) -> None:
+        super().__init__()
+        self.batch_remove_forces: list[bool] = []
+
+    def batch_remove(self, keys: list[str], force: bool = False) -> list[int]:
+        self.batch_remove_forces.append(force)
+        return super().batch_remove(keys, force)
+
+
+class TransientBatchRemoveStore(InMemoryStore):
+    """batch_remove reports a transient failure for one key without deleting it.
+
+    The inherited per-key remove() retry then succeeds and deletes the key, so
+    cleanup should end with no outstanding error.
+    """
+
+    def batch_remove(self, keys: list[str], force: bool = False) -> list[int]:
+        self.batch_remove_calls += 1
+        results: list[int] = []
+        for index, key in enumerate(keys):
+            if index == 0:
+                results.append(-1)  # transient failure, key left in place
+            else:
+                self.remove(key, force)
+                results.append(0)
+        return results
+
+
+class StrictRegisterStore(InMemoryStore):
+    def register_buffer(self, buffer_ptr: int, size: int) -> int:
+        if buffer_ptr in self.registered:
+            return -600
+        return super().register_buffer(buffer_ptr, size)
+
+
+class FailingRegisterStore(InMemoryStore):
+    def __init__(self, fail_on_register: int) -> None:
+        super().__init__()
+        self.register_count = 0
+        self.fail_on_register = fail_on_register
+
+    def register_buffer(self, buffer_ptr: int, size: int) -> int:
+        self.register_count += 1
+        if self.register_count == self.fail_on_register:
+            return -1
+        return super().register_buffer(buffer_ptr, size)
+
+
+def make_transfer(
+    store: InMemoryStore | None = None,
+    *,
+    key_prefix: str = "test",
+    default_chunk_bytes: int | None = None,
+    buffer_pool=None,
+) -> tuple[InMemoryStore, MooncakeBundleTransfer]:
+    current_store = store or InMemoryStore()
+    kwargs = {"key_prefix": key_prefix, "buffer_pool": buffer_pool}
+    if default_chunk_bytes is not None:
+        kwargs["default_chunk_bytes"] = default_chunk_bytes
+    return current_store, MooncakeBundleTransfer(current_store, **kwargs)
+
+
+def real_transfer(key_prefix: str) -> tuple[object, MooncakeBundleTransfer]:
+    mooncake_store = pytest.importorskip("mooncake.store")
+    store = mooncake_store.MooncakeDistributedStore()
+    rc = store.setup(
+        os.getenv("LOCAL_HOSTNAME", "localhost"),
+        os.getenv("MC_METADATA_SERVER", "P2PHANDSHAKE"),
+        16 * 1024 * 1024,
+        4 * 1024 * 1024,
+        os.getenv("PROTOCOL", "tcp"),
+        os.getenv("DEVICE_NAME", ""),
+        os.getenv("MASTER_SERVER", "127.0.0.1:50051"),
+    )
+    if rc != 0:
+        pytest.skip(f"MooncakeDistributedStore setup failed: {rc}")
+    return store, MooncakeBundleTransfer(store, key_prefix=key_prefix)
+
+
+def structured_payload(
+    metadata: dict[str, object] | None = None, **buffers: object
+) -> StructuredObjectPayload:
+    return StructuredObjectPayload(metadata=metadata, buffers=buffers)
+
+
+def write_manifest(
+    store: InMemoryStore, manifest_key: str, manifest: dict[str, object]
+) -> None:
+    store.objects[manifest_key] = json.dumps(manifest, separators=(",", ":")).encode(
+        "utf-8"
+    )
+
+
+def test_put_object_roundtrips_numpy_and_torch_tensor_fields() -> None:
+    torch = pytest.importorskip("torch")
+    store, transfer = make_transfer()
+    array = np.arange(6, dtype=np.float32).reshape(2, 3)
+    tensor = torch.arange(6, dtype=torch.float32).reshape(2, 3)
+    scalar_tensor = torch.tensor(1.5, dtype=torch.float32)
+    bool_tensor = torch.tensor([True, False], dtype=torch.bool)
+
+    result = transfer.get_object(
+        transfer.put_object(
+            {
+                "array": array,
+                "tensor": tensor,
+                "scalar": scalar_tensor,
+                "bool_tensor": bool_tensor,
+            }
+        )
+    )
+
+    assert np.array_equal(result["array"], array)
+    assert torch.equal(result["tensor"], tensor)
+    assert torch.equal(result["scalar"], scalar_tensor)
+    assert torch.equal(result["bool_tensor"], bool_tensor)
+    assert store.put_tensor_calls == 3
+    assert store.get_tensor_calls == 3
+
+
+def test_put_object_roundtrips_wrapped_numpy_and_torch_values() -> None:
+    torch = pytest.importorskip("torch")
+    store, transfer = make_transfer()
+    array = np.arange(6, dtype=np.float32).reshape(2, 3)
+    tensor = torch.arange(6, dtype=torch.float32).reshape(2, 3)
+
+    assert np.array_equal(transfer.get_object(transfer.put_object(array)), array)
+    assert torch.equal(transfer.get_object(transfer.put_object(tensor)), tensor)
+    assert store.put_tensor_calls == 1
+    assert store.get_tensor_calls == 1
+
+
+def test_get_object_requires_get_tensor_for_tensor_payload() -> None:
+    torch = pytest.importorskip("torch")
+    store, transfer = make_transfer(PutTensorOnlyStore())
+    tensor = torch.arange(6, dtype=torch.float32).reshape(2, 3)
+
+    ref = transfer.put_object({"tensor": tensor})
+
+    with pytest.raises(RuntimeError, match="does not support get_tensor"):
+        transfer.get_object(ref)
+
+
+def test_put_object_torch_tensor_raw_fallback_roundtrip() -> None:
+    torch = pytest.importorskip("torch")
+    store, transfer = make_transfer(NoTensorFastPathStore())
+    tensor = torch.arange(6, dtype=torch.float32).reshape(2, 3)
+    scalar_tensor = torch.tensor(1.5, dtype=torch.float32)
+    bool_tensor = torch.tensor([True, False], dtype=torch.bool)
+
+    result = transfer.get_object(
+        transfer.put_object(
+            {
+                "tensor": tensor,
+                "scalar": scalar_tensor,
+                "bool_tensor": bool_tensor,
+            }
+        )
+    )
+
+    assert torch.equal(result["tensor"], tensor)
+    assert torch.equal(result["scalar"], scalar_tensor)
+    assert torch.equal(result["bool_tensor"], bool_tensor)
+    assert store.batch_get_into_calls > 0
+
+
+def test_bundle_read_spec_full_read_is_partial_special_case() -> None:
+    store, transfer = make_transfer()
+    array = np.arange(16, dtype=np.int32).reshape(4, 4)
+    payload = structured_payload({"type": "example"}, array=array, raw=b"abc")
+
+    ref = transfer.put_structured_object(payload)
+    result = transfer.materialize(transfer.read_spec(ref))
+
+    assert np.array_equal(result.objects["array"], array)
+    assert result.objects["raw"] == b"abc"
+    assert store.batch_get_into_calls > 0
+    assert store.registered == set()
+
+
+def test_structured_object_payload_metadata_defaults_empty() -> None:
+    store, transfer = make_transfer()
+    array = np.arange(8, dtype=np.int32).reshape(2, 4)
+
+    ref = transfer.put_structured_object(
+        StructuredObjectPayload(buffers={"array": array})
+    )
+    result = transfer.materialize(transfer.read_spec(ref))
+
+    assert np.array_equal(result.objects["array"], array)
+
+
+def test_bundle_chunked_full_read_via_read_spec() -> None:
+    store, transfer = make_transfer()
+    payload = bytes(range(128))
+
+    ref = transfer.put_structured_object(
+        structured_payload(payload=payload),
+        chunk_bytes=17,
+    )
+    result = transfer.materialize(transfer.read_spec(ref))
+
+    assert result.objects["payload"] == payload
+    assert len(ref.manifest["buffers"]["payload"]["chunks"]) > 1
+
+
+def test_bundle_put_falls_back_to_store_put_without_batch_put_support() -> None:
+    store, transfer = make_transfer(MinimalStore())
+    payload = bytes(range(128))
+
+    ref = transfer.put_structured_object(
+        structured_payload(payload=payload), chunk_bytes=17
+    )
+    result = transfer.materialize(transfer.read_spec(ref))
+
+    assert result.objects["payload"] == payload
+    assert len(ref.manifest["buffers"]["payload"]["chunks"]) > 1
+    assert store.registered == set()
+
+
+def test_bundle_uses_configurable_default_chunk_size() -> None:
+    store, transfer = make_transfer(default_chunk_bytes=17)
+    payload = bytes(range(128))
+
+    ref = transfer.put_bundle(b"meta", {"payload": payload})
+
+    assert len(ref.manifest["buffers"]["payload"]["chunks"]) > 1
+    assert ref.manifest["buffers"]["payload"]["chunks"][0]["bytes"] == 17
+
+
+def test_bundle_copy_mode_forces_store_put() -> None:
+    store, transfer = make_transfer()
+    payload = bytes(range(128))
+
+    ref = transfer.put_structured_object(
+        structured_payload(payload=payload),
+        chunk_bytes=17,
+        policy=BundleTransferPolicy(copy_mode="copy"),
+    )
+
+    assert store.batch_put_from_calls == 0
+    assert store.register_buffer_calls == 0
+    assert store.unregister_buffer_calls == 0
+
+    result = transfer.materialize(transfer.read_spec(ref))
+    assert result.objects["payload"] == payload
+
+
+def test_structured_object_copy_mode_skips_pre_registered_validation() -> None:
+    store, transfer = make_transfer()
+    payload = np.arange(64, dtype=np.uint8).reshape(8, 8).T
+
+    ref = transfer.put_structured_object(
+        structured_payload(payload=payload),
+        policy=BundleTransferPolicy(copy_mode="copy"),
+        pre_registered_buffers={"payload": True},
+    )
+
+    assert store.batch_put_from_calls == 0
+    assert store.register_buffer_calls == 0
+    assert store.unregister_buffer_calls == 0
+
+    result = transfer.materialize(transfer.read_spec(ref))
+    assert np.array_equal(result.objects["payload"], payload)
+
+
+def test_bundle_zero_copy_mode_requires_batch_put_support() -> None:
+    _store, transfer = make_transfer(MinimalStore())
+
+    with pytest.raises(RuntimeError, match="zero-copy put requested"):
+        transfer.put_structured_object(
+            structured_payload(payload=b"data"),
+            policy=BundleTransferPolicy(copy_mode="zero_copy"),
+        )
+
+
+def test_structured_object_pre_registered_buffers_passthrough() -> None:
+    store, transfer = make_transfer()
+    payload = np.arange(64, dtype=np.uint8).reshape(8, 8)
+    payload_ptr = ctypes.addressof(ctypes.c_char.from_buffer(payload))
+    assert store.register_buffer(payload_ptr, int(payload.nbytes)) == 0
+    store.register_buffer_calls = 0
+    store.unregister_buffer_calls = 0
+
+    ref = transfer.put_structured_object(
+        structured_payload(payload=payload),
+        pre_registered_buffers={"payload": True},
+    )
+
+    assert store.batch_put_from_calls > 0
+    assert store.register_buffer_calls == 0
+    assert store.unregister_buffer_calls == 0
+    assert payload_ptr in store.registered
+
+    result = transfer.materialize(transfer.read_spec(ref))
+    assert np.array_equal(result.objects["payload"], payload)
+    store.unregister_buffer(payload_ptr)
+
+
+def test_structured_object_pre_registered_rejects_copied_buffers() -> None:
+    _store, transfer = make_transfer()
+    non_contiguous = np.arange(64, dtype=np.uint8).reshape(8, 8).T
+
+    with pytest.raises(ValueError, match="pre-registered structured buffer"):
+        transfer.put_structured_object(
+            structured_payload(payload=non_contiguous),
+            pre_registered_buffers={"payload": True},
+        )
+
+    with pytest.raises(ValueError, match="pre-registered structured buffer"):
+        transfer.put_structured_object(
+            structured_payload(payload=b"readonly"),
+            pre_registered_buffers={"payload": True},
+        )
+
+    with pytest.raises(ValueError, match="unknown pre-registered"):
+        transfer.put_structured_object(
+            structured_payload(payload=bytearray(b"data")),
+            pre_registered_buffers={"missing": True},
+        )
+
+
+def test_structured_object_pre_registered_empty_buffer() -> None:
+    _store, transfer = make_transfer()
+    payload = bytearray()
+
+    ref = transfer.put_structured_object(
+        structured_payload(payload=payload),
+        pre_registered_buffers={"payload": True},
+    )
+
+    result = transfer.materialize(transfer.read_spec(ref))
+    assert result.objects["payload"] == b""
+
+
+def test_structured_object_roundtrip() -> None:
+    store, transfer = make_transfer()
+    array = np.arange(12, dtype=np.int16).reshape(3, 4)
+    payload = structured_payload(
+        {"type": "example", "fields": ["weights", "tokens"]},
+        weights=array,
+        tokens=b"abc",
+    )
+
+    ref = transfer.put_structured_object(payload)
+    result = transfer.materialize(transfer.read_spec(ref))
+
+    assert result.metadata == {
+        "layout": "structured",
+        "type": "example",
+        "fields": ["weights", "tokens"],
+        "__mooncake_structured_fields__": {
+            "weights": {
+                "encoding": "ndarray",
+                "dtype": array.dtype.str,
+                "shape": [3, 4],
+            },
+            "tokens": {"encoding": "bytes"},
+        },
+    }
+    assert np.array_equal(result.objects["weights"], array)
+    assert result.objects["weights"].flags["C_CONTIGUOUS"]
+    assert result.objects["tokens"] == b"abc"
+
+
+def test_structured_object_read_spec_select_members() -> None:
+    store, transfer = make_transfer()
+    array = np.arange(10, dtype=np.float32).reshape(2, 5)
+    payload = structured_payload(weights=array, raw=b"payload")
+
+    ref = transfer.put_structured_object(payload)
+    spec = transfer.read_spec(ref).select_members(["weights"])
+    assert isinstance(spec, StructuredObjectReadSpec)
+    batch_get_into_calls = store.batch_get_into_calls
+    result = transfer.materialize(spec)
+
+    assert list(result.objects) == ["weights"]
+    assert np.array_equal(result.objects["weights"], array)
+    assert store.get_into_calls > 0
+    assert store.get_into_ranges_calls == 0
+    assert store.batch_get_into_calls == batch_get_into_calls + 1
+
+
+def test_structured_object_multichunk_ndarray_uses_range_gather() -> None:
+    store, transfer = make_transfer()
+    array = np.arange(64, dtype=np.int16).reshape(8, 8)
+    payload = structured_payload(weights=array)
+
+    ref = transfer.put_structured_object(payload, chunk_bytes=32)
+    batch_get_into_calls = store.batch_get_into_calls
+    result = transfer.materialize(transfer.read_spec(ref))
+
+    assert np.array_equal(result.objects["weights"], array)
+    assert store.get_into_ranges_calls > 0
+    assert store.batch_get_into_calls == batch_get_into_calls + 1
+
+
+def test_structured_object_slice_member_uses_partial_range_reads() -> None:
+    store, transfer = make_transfer()
+    array = np.arange(96, dtype=np.int16).reshape(12, 8)
+    payload = structured_payload(weights=array, raw=b"payload")
+
+    ref = transfer.put_structured_object(payload, chunk_bytes=20)
+    spec = (
+        transfer.read_spec(ref)
+        .select_members(["weights"])
+        .slice_member("weights", axis=0, start=3, end=9)
+    )
+    before_range_reads = store.get_into_ranges_calls
+    result = transfer.materialize(spec)
+
+    assert np.array_equal(result.objects["weights"], array[3:9])
+    assert result.objects["weights"].shape == (6, 8)
+    assert store.get_into_ranges_calls == before_range_reads + 1
+
+
+def test_structured_object_slice_member_falls_back_to_plain_get_reads() -> None:
+    store, transfer = make_transfer(PlainStore())
+    array = np.arange(96, dtype=np.int16).reshape(12, 8)
+    payload = structured_payload(weights=array, raw=b"payload")
+
+    ref = transfer.put_structured_object(payload, chunk_bytes=20)
+    spec = (
+        transfer.read_spec(ref)
+        .select_members(["weights"])
+        .slice_member("weights", axis=0, start=3, end=9)
+    )
+    result = transfer.materialize(spec)
+
+    assert np.array_equal(result.objects["weights"], array[3:9])
+    assert result.objects["weights"].shape == (6, 8)
+    assert store.get_into_calls == 0
+    assert store.get_into_ranges_calls == 0
+    assert store.batch_get_into_calls == 0
+
+
+def test_structured_object_materialize_into_reuses_destination() -> None:
+    store, transfer = make_transfer()
+    array = np.arange(96, dtype=np.float32).reshape(12, 8)
+    payload = structured_payload(weights=array)
+
+    ref = transfer.put_structured_object(payload, chunk_bytes=40)
+    spec = transfer.read_spec(ref).slice_member("weights", axis=0, start=2, end=10)
+    destination = np.empty((8, 8), dtype=np.float32)
+    result = transfer.materialize_into(spec, destinations={"weights": destination})
+
+    assert result.objects["weights"] is destination
+    assert np.array_equal(destination, array[2:10])
+    assert store.get_into_ranges_calls > 0
+
+
+def test_structured_object_duplicate_destination_registration_is_tolerated() -> None:
+    store, transfer = make_transfer(StrictRegisterStore())
+    array = np.arange(96, dtype=np.float32).reshape(12, 8)
+    payload = structured_payload(weights=array)
+
+    ref = transfer.put_structured_object(payload, chunk_bytes=40)
+    spec = (
+        transfer.read_spec(ref)
+        .select_members(["weights"])
+        .slice_member("weights", axis=0, start=2, end=10)
+    )
+    destination = np.empty((8, 8), dtype=np.float32)
+    destination_ptr = ctypes.addressof(ctypes.c_char.from_buffer(destination))
+    assert store.register_buffer(destination_ptr, int(destination.nbytes)) == 0
+
+    result = transfer.materialize_into(spec, destinations={"weights": destination})
+
+    assert result.objects["weights"] is destination
+    assert np.array_equal(destination, array[2:10])
+    assert store.get_into_ranges_calls > 0
+    assert destination_ptr in store.registered
+    store.unregister_buffer(destination_ptr)
+
+
+def test_structured_object_slice_member_step_copy() -> None:
+    store, transfer = make_transfer()
+    array = np.arange(120, dtype=np.int32).reshape(15, 8)
+    payload = structured_payload(weights=array)
+
+    ref = transfer.put_structured_object(payload, chunk_bytes=24)
+    spec = transfer.read_spec(ref).slice_member(
+        "weights", axis=0, start=1, end=12, step=3
+    )
+    result = transfer.materialize(spec)
+
+    assert np.array_equal(result.objects["weights"], array[1:12:3])
+    assert result.objects["weights"].shape == array[1:12:3].shape
+    assert store.get_into_ranges_calls > 0
+
+
+def test_structured_object_invalid_slice_and_destination_raise() -> None:
+    store, transfer = make_transfer()
+    array = np.arange(24, dtype=np.int16).reshape(6, 4)
+    ref = transfer.put_structured_object(
+        structured_payload(weights=array, raw=b"abc"),
+        chunk_bytes=8,
+    )
+
+    with pytest.raises(ValueError, match="axis=0"):
+        transfer.materialize(
+            transfer.read_spec(ref).slice_member("weights", axis=1, start=0, end=2)
+        )
+    with pytest.raises(ValueError, match="step must be positive"):
+        transfer.materialize(
+            transfer.read_spec(ref).slice_member(
+                "weights", axis=0, start=0, end=2, step=0
+            )
+        )
+    with pytest.raises(ValueError, match="does not support slicing"):
+        transfer.materialize(
+            transfer.read_spec(ref).slice_member("raw", axis=0, start=0, end=1)
+        )
+    with pytest.raises(ValueError, match="shape mismatch"):
+        transfer.materialize_into(
+            transfer.read_spec(ref).slice_member("weights", axis=0, start=1, end=3),
+            destinations={"weights": np.empty((3, 4), dtype=np.int16)},
+        )
+    readonly_destination = np.empty((2, 4), dtype=np.int16)
+    readonly_destination.flags.writeable = False
+    with pytest.raises(ValueError, match="writeable"):
+        transfer.materialize_into(
+            transfer.read_spec(ref).slice_member("weights", axis=0, start=1, end=3),
+            destinations={"weights": readonly_destination},
+        )
+
+
+def test_bundle_remove_deletes_payload_and_manifest() -> None:
+    store, transfer = make_transfer()
+
+    ref = transfer.put_bundle(b"meta", {"payload": b"data"})
+    assert store.objects
+    transfer.remove_bundle(ref)
+
+    assert store.objects == {}
+    assert store.batch_remove_calls == 1
+
+
+def test_bundle_partial_put_failure_cleans_payloads() -> None:
+    store, transfer = make_transfer(FailingPutStore(fail_on_put=2))
+
+    with pytest.raises(RuntimeError, match="injected put failure"):
+        transfer.put_bundle(b"meta", {"payload": b"data"}, chunk_bytes=2)
+
+    assert store.objects == {}
+
+
+def test_bundle_cleanup_failure_preserves_put_error() -> None:
+    _store, transfer = make_transfer(FailingRemoveStore(fail_on_put=2))
+
+    with pytest.raises(RuntimeError, match="injected put failure"):
+        transfer.put_bundle(b"meta", {"payload": b"data"}, chunk_bytes=2)
+
+
+def test_bundle_remove_uses_force_batch_remove_when_available() -> None:
+    store, transfer = make_transfer(ForceTrackingStore())
+
+    ref = transfer.put_structured_object(
+        structured_payload(payload=bytes(range(64))),
+        chunk_bytes=8,
+    )
+    transfer.materialize(transfer.read_spec(ref))
+    transfer.remove_bundle(ref)
+
+    assert store.batch_remove_forces == [True]
+    assert store.objects == {}
+
+
+def test_bundle_remove_recovers_after_transient_batch_failure() -> None:
+    store, transfer = make_transfer(TransientBatchRemoveStore())
+
+    ref = transfer.put_bundle(b"meta", {"a": b"x", "b": b"y", "c": b"z"})
+    transfer.remove_bundle(ref)
+
+    assert store.objects == {}
+    assert store.batch_remove_calls == 1
+
+
+def test_bundle_concurrent_put_and_read_spec_full_read() -> None:
+    store, transfer = make_transfer(GetOnlyStore())
+    payload = bytes(range(128))
+
+    ref = transfer.put_structured_object(
+        structured_payload(payload=payload),
+        chunk_bytes=8,
+        policy=BundleTransferPolicy(max_inflight_put=4, put_mode="parallel"),
+    )
+    result = transfer.materialize(transfer.read_spec(ref))
+
+    assert result.objects["payload"] == payload
+    assert store.max_active_puts > 1
+    assert store.max_active_gets >= 1
+
+
+def test_bundle_duplicate_source_registration_is_tolerated() -> None:
+    store, transfer = make_transfer(StrictRegisterStore())
+    payload = np.arange(64, dtype=np.uint8).reshape(8, 8)
+    payload_ptr = ctypes.addressof(ctypes.c_char.from_buffer(payload))
+    assert store.register_buffer(payload_ptr, int(payload.nbytes)) == 0
+
+    ref = transfer.put_structured_object(structured_payload(payload=payload))
+
+    assert ref.manifest["buffers"]["payload"]["bytes"] == int(payload.nbytes)
+    assert payload_ptr in store.registered
+    store.unregister_buffer(payload_ptr)
+
+
+def test_bundle_partial_register_failure_unwinds_registered_buffers() -> None:
+    store, transfer = make_transfer(FailingRegisterStore(fail_on_register=2))
+    payload = bytes(range(64))
+
+    with pytest.raises(RuntimeError, match="register_buffer"):
+        transfer.put_structured_object(
+            structured_payload(payload=payload), chunk_bytes=32
+        )
+
+    assert store.registered == set()
+    assert store.objects == {}
+
+
+def test_bundle_batch_get_failure_unregisters_buffer() -> None:
+    store, transfer = make_transfer(FailingBatchGetStore())
+    ref = transfer.put_structured_object(
+        structured_payload(payload=bytes(range(64))),
+        chunk_bytes=8,
+    )
+    store.fail_key = ref.manifest["buffers"]["payload"]["chunks"][2]["key"]
+
+    with pytest.raises(RuntimeError, match="injected get failure"):
+        transfer.materialize(transfer.read_spec(ref))
+
+    assert store.registered == set()
+
+
+def test_bundle_falls_back_to_get_without_batch_get_into() -> None:
+    store, transfer = make_transfer(GetOnlyStore())
+
+    ref = transfer.put_structured_object(
+        structured_payload(payload=b"abcdef"),
+        chunk_bytes=2,
+    )
+    result = transfer.materialize(transfer.read_spec(ref))
+
+    assert result.objects["payload"] == b"abcdef"
+    assert store.max_active_gets >= 1
+
+
+def test_bundle_read_fast_paths_require_registration_support() -> None:
+    store, transfer = make_transfer(ReadFastPathWithoutRegisterStore())
+
+    ref = transfer.put_structured_object(
+        structured_payload(payload=bytes(range(64))),
+        chunk_bytes=8,
+    )
+    result = transfer.materialize(transfer.read_spec(ref))
+
+    assert result.objects["payload"] == bytes(range(64))
+    assert store.batch_get_into_calls == 0
+    assert store.get_into_calls == 0
+    assert store.get_into_ranges_calls == 0
+
+
+def test_bundle_invalid_policy_and_chunk_size_raise() -> None:
+    store, transfer = make_transfer()
+
+    with pytest.raises(ValueError, match="max_inflight_put"):
+        transfer.put_bundle(b"meta", {"payload": b"data"}, max_inflight_put=0)
+    with pytest.raises(ValueError, match="selected no members"):
+        transfer.materialize(
+            transfer.read_spec(
+                {"manifest_key": "test/default/example/manifest"}
+            ).select_members([])
+        )
+    with pytest.raises(ValueError, match="chunk_bytes"):
+        transfer.put_bundle(b"meta", {"payload": b"data"}, chunk_bytes=0)
+    with pytest.raises(ValueError, match="copy_mode"):
+        transfer.put_bundle(
+            b"meta",
+            {"payload": b"data"},
+            policy=BundleTransferPolicy(copy_mode="invalid"),
+        )
+
+
+def test_bundle_invalid_name_and_prefix_raise() -> None:
+    store, transfer = make_transfer()
+
+    with pytest.raises(ValueError, match="buffer name"):
+        transfer.put_bundle(b"meta", {"bad/name": b"data"})
+    with pytest.raises(ValueError, match="partition"):
+        transfer.put_bundle(b"meta", {"payload": b"data"}, partition="bad/name")
+    with pytest.raises(ValueError, match="key_prefix"):
+        MooncakeBundleTransfer(store, key_prefix="")
+    with pytest.raises(ValueError, match="chunk_bytes"):
+        MooncakeBundleTransfer(store, default_chunk_bytes=0)
+
+
+def test_structured_object_invalid_metadata_raises() -> None:
+    store, transfer = make_transfer()
+
+    with pytest.raises(TypeError, match="mapping"):
+        transfer.put_structured_object(
+            StructuredObjectPayload(metadata=["not", "a", "mapping"], buffers={})
+        )
+    with pytest.raises(TypeError, match="JSON-serializable"):
+        transfer.put_structured_object(
+            StructuredObjectPayload(metadata={"bad": object()}, buffers={})
+        )
+    with pytest.raises(ValueError, match="reserved"):
+        transfer.put_structured_object(
+            StructuredObjectPayload(
+                metadata={"__mooncake_structured_fields__": {}},
+                buffers={"payload": b"abc"},
+            )
+        )
+
+
+def test_structured_object_rejects_invalid_field_spec() -> None:
+    store, transfer = make_transfer()
+    manifest_key = "test/default/example/manifest"
+    meta_key = "test/default/example/meta"
+    payload_key = "test/default/example/buffer/payload"
+    metadata_blob = (
+        b'{"layout":"structured","__mooncake_structured_fields__":'
+        b'{"payload":{"encoding":"unknown"}}}'
+    )
+    store.objects[meta_key] = metadata_blob
+    store.objects[payload_key] = b"abc"
+    write_manifest(
+        store,
+        manifest_key,
+        {
+            "version": 1,
+            "layout": "bundle",
+            "object_id": "default/example",
+            "meta": {
+                "key": meta_key,
+                "bytes": len(metadata_blob),
+                "chunks": [{"key": meta_key, "bytes": len(metadata_blob)}],
+            },
+            "buffers": {
+                "payload": {
+                    "key": payload_key,
+                    "bytes": 3,
+                    "chunks": [{"key": payload_key, "bytes": 3}],
+                }
+            },
+        },
+    )
+
+    with pytest.raises(ValueError, match="unsupported structured field encoding"):
+        transfer.materialize(transfer.read_spec({"manifest_key": manifest_key}))
+
+
+def test_bundle_rejects_tampered_manifest() -> None:
+    store, transfer = make_transfer()
+    ref = transfer.put_bundle(b"meta", {"payload": b"abcdef"}, chunk_bytes=2)
+    tampered = dict(ref.manifest)
+    tampered["buffers"] = dict(ref.manifest["buffers"])
+    tampered["buffers"]["payload"] = dict(ref.manifest["buffers"]["payload"])
+    tampered["buffers"]["payload"]["chunks"] = [
+        {"key": "other/object", "bytes": 6},
+    ]
+    write_manifest(store, ref.manifest_key, tampered)
+
+    with pytest.raises(ValueError, match="namespace"):
+        transfer.materialize(transfer.read_spec({"manifest_key": ref.manifest_key}))
+
+    with pytest.raises(ValueError, match="manifest_key"):
+        transfer.remove_bundle(
+            RemoteBundleRef(manifest_key="test/other/manifest", manifest=ref.manifest)
+        )
+
+
+def test_structured_object_torch_tensor_falls_back_without_buffer_pool() -> None:
+    torch = pytest.importorskip("torch")
+    mooncake_store = pytest.importorskip("mooncake.store")
+    if not hasattr(mooncake_store, "_serialize_tensor") or not hasattr(
+        mooncake_store, "_deserialize_tensor"
+    ):
+        pytest.skip("built mooncake.store lacks tensor serialization helpers")
+    store, transfer = make_transfer()
+    tensor = torch.arange(12, dtype=torch.int64).reshape(3, 4)
+
+    ref = transfer.put_structured_object(
+        StructuredObjectPayload(buffers={"tensor": tensor})
+    )
+    result = transfer.materialize(transfer.read_spec(ref))
+
+    assert store.put_tensor_from_calls == 0
+    assert store.batch_put_from_calls == 0
+    assert torch.equal(result.objects["tensor"], tensor)
+
+
+def test_structured_object_direct_torch_tensor_materialize_into_uses_real_store() -> (
+    None
+):
+    torch = pytest.importorskip("torch")
+    mooncake_store = pytest.importorskip("mooncake.store")
+    if not hasattr(mooncake_store, "get_tensor_into") and not hasattr(
+        mooncake_store.MooncakeDistributedStore, "get_tensor_into"
+    ):
+        pytest.skip("built mooncake.store lacks get_tensor_into")
+    store = mooncake_store.MooncakeDistributedStore()
+    rc = store.setup(
+        os.getenv("LOCAL_HOSTNAME", "localhost"),
+        os.getenv("MC_METADATA_SERVER", "P2PHANDSHAKE"),
+        16 * 1024 * 1024,
+        4 * 1024 * 1024,
+        os.getenv("PROTOCOL", "tcp"),
+        os.getenv("DEVICE_NAME", ""),
+        os.getenv("MASTER_SERVER", "127.0.0.1:50051"),
+    )
+    if rc != 0:
+        pytest.skip(f"MooncakeDistributedStore setup failed: {rc}")
+    transfer = MooncakeBundleTransfer(store, key_prefix="structured-test-direct")
+    tensor = torch.arange(12, dtype=torch.int64).reshape(3, 4)
+    ref = transfer.put_structured_object(
+        StructuredObjectPayload(buffers={"tensor": tensor})
+    )
+    total_bytes = int(ref.manifest["buffers"]["tensor"]["bytes"])
+    destination = ctypes.create_string_buffer(total_bytes)
+
+    try:
+        result = transfer.materialize_into(
+            transfer.read_spec(ref),
+            {
+                "tensor": tensor_object_buffer(
+                    ctypes.addressof(destination),
+                    total_bytes,
+                    destination,
+                    batch_size=3,
+                )
+            },
+        )
+        assert torch.equal(result.objects["tensor"], tensor)
+    finally:
+        transfer.remove_bundle(ref)
+
+
+def test_structured_object_torch_tensor_zero_copy_rejects_plain_tensor() -> None:
+    torch = pytest.importorskip("torch")
+    mooncake_store = pytest.importorskip("mooncake.store")
+    if not hasattr(mooncake_store, "_serialize_tensor"):
+        pytest.skip("built mooncake.store lacks tensor serialization helpers")
+    _store, transfer = make_transfer()
+    tensor = torch.arange(12, dtype=torch.int64).reshape(3, 4)
+
+    with pytest.raises(ValueError, match="BufferPool tensor-object buffer"):
+        transfer.put_structured_object(
+            StructuredObjectPayload(buffers={"tensor": tensor}),
+            policy=BundleTransferPolicy(copy_mode="zero_copy"),
+        )
+
+
+def test_structured_object_torch_tensor_slice_uses_range_reads() -> None:
+    torch = pytest.importorskip("torch")
+    mooncake_store = pytest.importorskip("mooncake.store")
+    if not hasattr(mooncake_store, "_serialize_tensor") or not hasattr(
+        mooncake_store, "_deserialize_tensor"
+    ):
+        pytest.skip("built mooncake.store lacks tensor serialization helpers")
+    store, transfer = make_transfer(NoTensorFastPathStore())
+    tensor = torch.arange(48, dtype=torch.int64).reshape(12, 4)
+    ref = transfer.put_structured_object(
+        StructuredObjectPayload(buffers={"tensor": tensor}), chunk_bytes=40
+    )
+
+    result = transfer.materialize(
+        transfer.read_spec(ref)
+        .select_members(["tensor"])
+        .slice_member("tensor", axis=0, start=3, end=9)
+    )
+
+    assert torch.equal(result.objects["tensor"], tensor[3:9])
+    assert store.get_into_ranges_calls >= 2
+
+
+def test_structured_object_direct_torch_tensor_slice_uses_real_store_ranges() -> None:
+    torch = pytest.importorskip("torch")
+    mooncake_store = pytest.importorskip("mooncake.store")
+    if not hasattr(mooncake_store, "_serialize_tensor") or not hasattr(
+        mooncake_store, "_deserialize_tensor"
+    ):
+        pytest.skip("built mooncake.store lacks tensor serialization helpers")
+    _store, transfer = real_transfer("structured-test-slice")
+    tensor = torch.arange(48, dtype=torch.int64).reshape(12, 4)
+    ref = transfer.put_structured_object(StructuredObjectPayload(buffers={"tensor": tensor}))
+
+    try:
+        result = transfer.materialize(
+            transfer.read_spec(ref)
+            .select_members(["tensor"])
+            .slice_member("tensor", axis=0, start=3, end=9)
+        )
+        assert torch.equal(result.objects["tensor"], tensor[3:9])
+    finally:
+        transfer.remove_bundle(ref)
+
+
+def test_structured_object_tensor_object_buffer_uses_put_tensor_from() -> None:
+    store, transfer = make_transfer()
+    source = ctypes.create_string_buffer(b"tensor-payload")
+    store.register_buffer(ctypes.addressof(source), len(source.raw))
+
+    ref = transfer.put_structured_object(
+        StructuredObjectPayload(
+            buffers={
+                "tensor": tensor_object_buffer(
+                    ctypes.addressof(source), len(source.raw), source, batch_size=1
+                )
+            }
+        ),
+        policy=BundleTransferPolicy(copy_mode="zero_copy"),
+    )
+
+    payload = ref.manifest["buffers"]["tensor"]
+    assert store.put_tensor_from_calls == 1
+    assert store.batch_put_from_calls == 0
+    assert payload["bytes"] == len(source.raw)
+    assert payload["chunks"] == [{"key": payload["key"], "bytes": len(source.raw)}]
+
+
+def test_structured_object_tensor_object_buffer_materialize_into_uses_ranges() -> None:
+    store, transfer = make_transfer()
+    source = ctypes.create_string_buffer(b"tensor-payload")
+    store.register_buffer(ctypes.addressof(source), len(source.raw))
+    ref = transfer.put_structured_object(
+        StructuredObjectPayload(
+            buffers={
+                "tensor": tensor_object_buffer(
+                    ctypes.addressof(source), len(source.raw), source, batch_size=1
+                )
+            }
+        ),
+        policy=BundleTransferPolicy(copy_mode="zero_copy"),
+    )
+    destination = ctypes.create_string_buffer(len(source.raw))
+
+    result = transfer.materialize_into(
+        transfer.read_spec(ref),
+        {
+            "tensor": tensor_object_buffer(
+                ctypes.addressof(destination),
+                len(destination.raw),
+                destination,
+                batch_size=1,
+            )
+        },
+    )
+
+    assert result.objects["tensor"].ptr == ctypes.addressof(destination)
+    assert destination.raw == source.raw
+    assert store.get_into_ranges_calls == 1
+
+
+def test_structured_object_torch_tensor_zero_copy_uses_real_buffer_pool() -> None:
+    torch = pytest.importorskip("torch")
+    mooncake_store = pytest.importorskip("mooncake.store")
+    if (
+        not hasattr(mooncake_store, "BufferPool")
+        or not hasattr(mooncake_store, "_serialize_tensor")
+        or not hasattr(mooncake_store, "_deserialize_tensor")
+    ):
+        pytest.skip("built mooncake.store lacks BufferPool tensor helpers")
+    store = mooncake_store.MooncakeDistributedStore()
+    rc = store.setup(
+        os.getenv("LOCAL_HOSTNAME", "localhost"),
+        os.getenv("MC_METADATA_SERVER", "P2PHANDSHAKE"),
+        16 * 1024 * 1024,
+        4 * 1024 * 1024,
+        os.getenv("PROTOCOL", "tcp"),
+        os.getenv("DEVICE_NAME", ""),
+        os.getenv("MASTER_SERVER", "127.0.0.1:50051"),
+    )
+    if rc != 0:
+        pytest.skip(f"MooncakeDistributedStore setup failed: {rc}")
+    pool = mooncake_store.BufferPool(store, min_size_class=4096, alignment=4096)
+    transfer = MooncakeBundleTransfer(
+        store, key_prefix="structured-test", buffer_pool=pool
+    )
+    tensor = torch.arange(12, dtype=torch.int64).reshape(3, 4)
+    metadata, data_ptr, tensor_nbytes, owner = mooncake_store._serialize_tensor(tensor)
+    total_bytes = len(metadata) + int(tensor_nbytes)
+    lease = pool.acquire(total_bytes)
+    view = lease.buffer
+    view[: len(metadata)] = metadata
+    ctypes.memmove(lease.ptr + len(metadata), int(data_ptr), int(tensor_nbytes))
+    view.release()
+
+    try:
+        ref = transfer.put_structured_object(
+            StructuredObjectPayload(
+                buffers={
+                    "tensor": tensor_object_buffer(
+                        lease.ptr, total_bytes, lease, batch_size=3
+                    )
+                }
+            ),
+            policy=BundleTransferPolicy(copy_mode="zero_copy"),
+        )
+        result = transfer.materialize(transfer.read_spec(ref))
+        assert torch.equal(result.objects["tensor"], tensor)
+        _ = owner
+    finally:
+        lease.release()
+        pool.close()
+
+
+def test_dataproto_helper_requires_batch_size_for_tensor_object_buffer() -> None:
+    _store, transfer = make_transfer()
+    data = SimpleDataProto(batch={"tensor": tensor_object_buffer(1, 128, object())})
+
+    with pytest.raises(TypeError, match="batch_size"):
+        transfer.put_dataproto(data, policy=BundleTransferPolicy(copy_mode="zero_copy"))
+
+
+def test_dataproto_helper_roundtrip_uses_structured_object() -> None:
+    store, transfer = make_transfer()
+    data = SimpleDataProto(
+        batch={"input_ids": np.arange(12, dtype=np.int64).reshape(4, 3)},
+        non_tensor_batch={
+            "reward": np.asarray([1.0, 0.0, 0.5, -1.0], dtype=np.float32)
+        },
+        meta_info={"step": 7, "source": "unit"},
+    )
+
+    ref = transfer.put_dataproto(
+        data, namespace="roll", partition="train", stage="rollout"
+    )
+    result = transfer.get_dataproto(ref, data_cls=SimpleDataProto)
+
+    assert ref.batch_size == 4
+    assert set(ref.stage_refs) == {"rollout"}
+    assert set(ref.field_index) == {"input_ids", "reward"}
+    assert np.array_equal(result.batch["input_ids"], data.batch["input_ids"])
+    assert np.array_equal(
+        result.non_tensor_batch["reward"], data.non_tensor_batch["reward"]
+    )
+    assert result.meta_info == data.meta_info
+    stage_metadata = transfer.materialize(
+        transfer.read_spec(ref.stage_refs["rollout"]).select_members(
+            ["batch.input_ids"]
+        )
+    ).metadata
+    assert stage_metadata["dataproto"]["stage"] == "rollout"
+
+
+def test_dataproto_manifest_view_reuses_structured_manifest_specs() -> None:
+    _store, transfer = make_transfer()
+    data = SimpleDataProto(
+        batch={"input_ids": np.arange(12, dtype=np.int64).reshape(4, 3)},
+        non_tensor_batch={
+            "reward": np.asarray([1.0, 0.0, 0.5, -1.0], dtype=np.float32)
+        },
+        meta_info={"step": 7},
+    )
+
+    ref = transfer.put_dataproto(
+        data, namespace="roll", partition="train", stage="rollout"
+    )
+    view = transfer.dataproto_manifest_view(ref)
+
+    assert view["namespace"] == "roll"
+    assert view["partition"] == "train"
+    assert view["batch_size"] == 4
+    assert view["stages"]["rollout"]["dataproto"]["stage"] == "rollout"
+    assert view["batch_fields"]["input_ids"]["member"] == "batch.input_ids"
+    assert view["batch_fields"]["input_ids"]["spec"]["encoding"] == "ndarray"
+    assert view["batch_fields"]["input_ids"]["spec"]["shape"] == [4, 3]
+    assert view["non_tensor_fields"]["reward"]["spec"]["dtype"] == "<f4"
+    assert view["meta_info_keys"] == ["step"]
+
+
+def test_dataproto_helper_selects_fields_and_meta() -> None:
+    _store, transfer = make_transfer()
+    data = SimpleDataProto(
+        batch={
+            "input_ids": np.arange(8, dtype=np.int64).reshape(4, 2),
+            "attention_mask": np.ones((4, 2), dtype=np.int32),
+        },
+        non_tensor_batch={
+            "reward": np.asarray([1.0, 0.0, 0.5, -1.0], dtype=np.float32),
+            "uid": np.arange(4, dtype=np.int64),
+        },
+        meta_info={"step": 7, "source": "unit"},
+    )
+    ref = transfer.put_dataproto(data)
+
+    result = transfer.get_dataproto(
+        ref,
+        fields=["input_ids", "reward"],
+        meta_info_keys=["step"],
+    )
+    batch_only = transfer.get_dataproto(ref, batch_fields=["input_ids"])
+    non_tensor_only = transfer.get_dataproto(ref, non_tensor_fields=["reward"])
+
+    assert set(result["batch"]) == {"input_ids"}
+    assert set(result["non_tensor_batch"]) == {"reward"}
+    assert result["meta_info"] == {"step": 7}
+    assert np.array_equal(result["batch"]["input_ids"], data.batch["input_ids"])
+    assert np.array_equal(
+        result["non_tensor_batch"]["reward"], data.non_tensor_batch["reward"]
+    )
+    assert set(batch_only["batch"]) == {"input_ids"}
+    assert batch_only["non_tensor_batch"] == {}
+    assert np.array_equal(batch_only["batch"]["input_ids"], data.batch["input_ids"])
+    assert non_tensor_only["batch"] == {}
+    assert set(non_tensor_only["non_tensor_batch"]) == {"reward"}
+    assert np.array_equal(
+        non_tensor_only["non_tensor_batch"]["reward"], data.non_tensor_batch["reward"]
+    )
+
+
+def test_dataproto_helper_reads_rows_with_real_store_ranges() -> None:
+    torch = pytest.importorskip("torch")
+    mooncake_store = pytest.importorskip("mooncake.store")
+    if not hasattr(mooncake_store, "_serialize_tensor") or not hasattr(
+        mooncake_store, "_deserialize_tensor"
+    ):
+        pytest.skip("built mooncake.store lacks tensor serialization helpers")
+    _store, transfer = real_transfer("structured-test-dataproto-rows")
+    tensor = torch.arange(24, dtype=torch.int64).reshape(6, 4)
+    data = SimpleDataProto(
+        batch={
+            "tensor": tensor,
+            "array": np.arange(18, dtype=np.int64).reshape(6, 3),
+        },
+        non_tensor_batch={
+            "text": np.asarray(
+                ["a", "bb", "ccc", "dddd", "eeeee", "ffffff"], dtype=object
+            ),
+            "json": np.asarray(
+                [{"i": 0}, {"i": 1}, {"i": 2}, {"i": 3}, {"i": 4}, {"i": 5}],
+                dtype=object,
+            ),
+            "ragged": np.asarray(
+                [
+                    torch.arange(1, dtype=torch.float32),
+                    torch.arange(2, dtype=torch.float32),
+                    torch.arange(3, dtype=torch.float32),
+                    torch.arange(4, dtype=torch.float32),
+                    torch.arange(5, dtype=torch.float32),
+                    torch.arange(6, dtype=torch.float32),
+                ],
+                dtype=object,
+            ),
+        },
+    )
+    ref = transfer.put_dataproto(data)
+
+    try:
+        sliced = transfer.get_dataproto(ref, rows=slice(2, 5))
+        gathered = transfer.get_dataproto(ref, rows=[4, 1, 3])
+        array_dst = np.empty((3, 3), dtype=np.int64)
+        into = transfer.get_dataproto(
+            ref,
+            batch_fields=["array"],
+            rows=[4, 1, 3],
+            destinations={"array": array_dst},
+        )
+        pool = mooncake_store.BufferPool(_store, 1024 * 1024)
+        raw_array = pool.acquire(array_dst.nbytes)
+        raw_array_result = transfer.get_dataproto(
+            ref,
+            batch_fields=["array"],
+            rows=[4, 1, 3],
+            destinations={
+                "array": raw_destination(
+                    raw_array.ptr,
+                    raw_array.size,
+                    raw_array,
+                    pre_registered=True,
+                )
+            },
+        )
+        tensor_payload_bytes = int(
+            ref.stage_refs["default"].manifest["buffers"]["batch.tensor"]["metadata_bytes"]
+        ) + tensor[[4, 1, 3]].numel() * tensor.element_size()
+        raw_tensor = pool.acquire(tensor_payload_bytes)
+        raw_tensor_result = transfer.get_dataproto(
+            ref,
+            batch_fields=["tensor"],
+            rows=[4, 1, 3],
+            destinations={
+                "tensor": raw_destination(
+                    raw_tensor.ptr,
+                    raw_tensor.size,
+                    raw_tensor,
+                    pre_registered=True,
+                )
+            },
+        )
+
+        assert torch.equal(sliced["batch"]["tensor"], tensor[2:5])
+        assert np.array_equal(sliced["batch"]["array"], data.batch["array"][2:5])
+        assert sliced["non_tensor_batch"]["text"].tolist() == [
+            "ccc",
+            "dddd",
+            "eeeee",
+        ]
+        assert sliced["non_tensor_batch"]["json"].tolist() == [
+            {"i": 2},
+            {"i": 3},
+            {"i": 4},
+        ]
+        actual_ragged = sliced["non_tensor_batch"]["ragged"]
+        assert torch.equal(actual_ragged[0], data.non_tensor_batch["ragged"][2])
+        assert torch.equal(actual_ragged[1], data.non_tensor_batch["ragged"][3])
+        assert torch.equal(actual_ragged[2], data.non_tensor_batch["ragged"][4])
+        assert torch.equal(gathered["batch"]["tensor"], tensor[[4, 1, 3]])
+        assert np.array_equal(gathered["batch"]["array"], data.batch["array"][[4, 1, 3]])
+        assert gathered["non_tensor_batch"]["text"].tolist() == ["eeeee", "bb", "dddd"]
+        assert into["batch"]["array"] is array_dst
+        assert np.array_equal(array_dst, data.batch["array"][[4, 1, 3]])
+        assert np.array_equal(
+            raw_array_result["batch"]["array"], data.batch["array"][[4, 1, 3]]
+        )
+        assert raw_tensor_result["batch"]["tensor"].ptr == raw_tensor.ptr
+        decoded_tensor = mooncake_store._deserialize_tensor(
+            bytes(raw_tensor.buffer[:tensor_payload_bytes])
+        )
+        assert torch.equal(decoded_tensor, tensor[[4, 1, 3]])
+        raw_array.release()
+        raw_tensor.release()
+        pool.close()
+    finally:
+        transfer.cleanup_dataproto(ref)
+
+
+def test_dataproto_helper_supports_dict_cls_and_reports_bad_cls() -> None:
+    _store, transfer = make_transfer()
+    ref = transfer.put_dataproto(
+        SimpleDataProto(batch={"input_ids": np.arange(4)}, meta_info={"step": 1})
+    )
+
+    result = transfer.get_dataproto(ref, data_cls=dict)
+
+    assert result["meta_info"] == {"step": 1}
+    assert np.array_equal(result["batch"]["input_ids"], np.arange(4))
+    with pytest.raises(TypeError, match="cannot be constructed"):
+        transfer.get_dataproto(ref, data_cls=BadDataProto)
+
+
+def test_dataproto_helper_accepts_legacy_dict_inputs() -> None:
+    _store, transfer = make_transfer()
+    plain_ref = transfer.put_dataproto({"input_ids": np.arange(4)})
+    envelope_ref = transfer.put_dataproto(
+        {
+            "batch": {"tokens": np.arange(6).reshape(3, 2)},
+            "non_tensor_batch": {"uid": np.asarray(["a", "b", "c"], dtype=object)},
+            "meta_info": {"step": 3},
+        }
+    )
+
+    plain = transfer.get_dataproto(plain_ref)
+    envelope = transfer.get_dataproto(envelope_ref)
+
+    assert np.array_equal(plain["batch"]["input_ids"], np.arange(4))
+    assert np.array_equal(envelope["batch"]["tokens"], np.arange(6).reshape(3, 2))
+    assert envelope["non_tensor_batch"]["uid"].tolist() == ["a", "b", "c"]
+    assert envelope["meta_info"] == {"step": 3}
+
+
+def test_dataproto_helper_treats_reserved_plain_dict_keys_as_batch_fields() -> None:
+    _store, transfer = make_transfer()
+    ref = transfer.put_dataproto({"batch": np.arange(4), "meta_info": np.arange(4)})
+
+    result = transfer.get_dataproto(ref)
+
+    assert np.array_equal(result["batch"]["batch"], np.arange(4))
+    assert np.array_equal(result["batch"]["meta_info"], np.arange(4))
+    assert result["meta_info"] == {}
+
+
+def test_dataproto_helper_rejects_cross_section_field_name_collision() -> None:
+    _store, transfer = make_transfer()
+
+    with pytest.raises(ValueError, match="overlap"):
+        transfer.put_dataproto(
+            {
+                "batch": {"uid": np.arange(4)},
+                "non_tensor_batch": {"uid": np.asarray(["a", "b", "c", "d"])},
+            }
+        )
+
+
+def test_dataproto_ref_handle_roundtrip_materializes_and_cleans_up() -> None:
+    store, transfer = make_transfer()
+    input_ids = np.arange(6, dtype=np.int64).reshape(3, 2)
+    ref = transfer.put_dataproto(
+        SimpleDataProto(
+            batch={"input_ids": input_ids},
+            non_tensor_batch={"uid": np.asarray(["a", "b", "c"], dtype=object)},
+            meta_info={"step": 3, "tags": ["rollout"]},
+        ),
+        namespace="roll",
+        partition="train",
+        stage="rollout",
+    )
+    handle = export_dataproto_ref(ref)
+
+    assert is_dataproto_ref_handle(handle)
+    assert handle["stage_refs"] == {
+        "rollout": {"manifest_key": ref.stage_refs["rollout"].manifest_key}
+    }
+    imported = import_dataproto_ref(handle)
+    assert imported.stage_refs["rollout"].manifest == {}
+
+    result = transfer.get_dataproto(handle)
+    assert np.array_equal(result["batch"]["input_ids"], input_ids)
+    assert result["non_tensor_batch"]["uid"].tolist() == ["a", "b", "c"]
+    assert result["meta_info"] == {"step": 3, "tags": ["rollout"]}
+    assert transfer.dataproto_manifest_view(handle)["meta_info_keys"] == [
+        "step",
+        "tags",
+    ]
+
+    transfer.cleanup_dataproto(handle)
+    assert store.objects == {}
+
+
+def test_dataproto_ref_handle_exports_numpy_meta_and_validates_indexes() -> None:
+    _store, transfer = make_transfer()
+    ref = transfer.put_dataproto(
+        SimpleDataProto(
+            batch={"input_ids": np.arange(4)},
+            meta_info={"scores": np.asarray([1, 2], dtype=np.int64)},
+        )
+    )
+    ref.global_indexes = [0, 2]
+
+    handle = export_dataproto_ref(ref)
+    imported = import_dataproto_ref(handle)
+
+    assert handle["meta_info"] == {"scores": [1, 2]}
+    assert imported.global_indexes == [0, 2]
+    bad_indexes = dict(handle)
+    bad_indexes["global_indexes"] = "0,2"
+    with pytest.raises(ValueError, match="global_indexes"):
+        import_dataproto_ref(bad_indexes)
+
+
+def test_dataproto_ref_handle_rejects_unknown_field_stage() -> None:
+    _store, transfer = make_transfer()
+    ref = transfer.put_dataproto(SimpleDataProto(batch={"input_ids": np.arange(4)}))
+    handle = export_dataproto_ref(ref)
+    handle["field_index"]["input_ids"] = dict(handle["field_index"]["input_ids"])
+    handle["field_index"]["input_ids"]["stage"] = "missing"
+
+    with pytest.raises(ValueError, match="unknown stage"):
+        import_dataproto_ref(handle)
+
+
+def test_dataproto_helper_appends_stage_fields_without_rewriting_existing() -> None:
+    store, transfer = make_transfer()
+    rollout = SimpleDataProto(
+        batch={"input_ids": np.arange(8, dtype=np.int64).reshape(4, 2)},
+        meta_info={"step": 1},
+    )
+    ref = transfer.put_dataproto(rollout, stage="rollout")
+    initial_object_count = len(store.objects)
+
+    old_log_prob = SimpleDataProto(
+        batch={"old_log_probs": np.arange(4, dtype=np.float32)},
+        meta_info={"stage": "old_log_prob"},
+    )
+    ref = transfer.append_dataproto_fields(ref, old_log_prob, stage="old_log_prob")
+    result = transfer.get_dataproto(ref)
+
+    assert set(ref.stage_refs) == {"rollout", "old_log_prob"}
+    assert len(store.objects) > initial_object_count
+    assert np.array_equal(result["batch"]["input_ids"], rollout.batch["input_ids"])
+    assert np.array_equal(
+        result["batch"]["old_log_probs"], old_log_prob.batch["old_log_probs"]
+    )
+    assert result["meta_info"] == {"step": 1, "stage": "old_log_prob"}
+
+    with pytest.raises(ValueError, match="already exist"):
+        transfer.append_dataproto_fields(ref, rollout, stage="duplicate")
+
+
+def test_dataproto_helper_same_stage_append_reads_rows_with_real_store() -> None:
+    pytest.importorskip("mooncake.store")
+    _store, transfer = real_transfer("structured-test-dataproto-same-stage")
+    input_ids = np.arange(24, dtype=np.int64).reshape(6, 4)
+    old_log_probs = np.linspace(0.0, 1.0, 6, dtype=np.float32)
+    text = np.asarray(["", "bb", "ccc", "dddd", "eeeee", "ffffff"], dtype=object)
+    ref = transfer.put_dataproto(
+        SimpleDataProto(batch={"input_ids": input_ids}, meta_info={"step": 1}),
+        stage="rollout",
+    )
+    old_ref = ref
+    old_manifest_key = ref.stage_refs["rollout"].manifest_key
+    ref = transfer.append_dataproto_fields(
+        ref,
+        SimpleDataProto(
+            batch={"old_log_probs": old_log_probs},
+            non_tensor_batch={"text": text},
+            meta_info={"stage": "rollout-extra"},
+        ),
+        stage="rollout",
+    )
+
+    try:
+        old_result = transfer.get_dataproto(old_ref)
+        result = transfer.get_dataproto(ref)
+        sliced = transfer.get_dataproto(ref, rows=slice(2, 5))
+        gathered = transfer.get_dataproto(
+            ref, fields=["input_ids", "text"], rows=[4, 1, 3]
+        )
+        view = transfer.dataproto_manifest_view(ref)
+
+        assert set(ref.stage_refs) == {"rollout"}
+        assert ref.stage_refs["rollout"].manifest_key != old_manifest_key
+        assert np.array_equal(old_result["batch"]["input_ids"], input_ids)
+        assert np.array_equal(result["batch"]["input_ids"], input_ids)
+        assert np.array_equal(result["batch"]["old_log_probs"], old_log_probs)
+        assert result["non_tensor_batch"]["text"].tolist() == text.tolist()
+        assert result["meta_info"] == {"step": 1, "stage": "rollout-extra"}
+        assert np.array_equal(sliced["batch"]["input_ids"], input_ids[2:5])
+        assert np.array_equal(sliced["batch"]["old_log_probs"], old_log_probs[2:5])
+        assert sliced["non_tensor_batch"]["text"].tolist() == text[2:5].tolist()
+        assert np.array_equal(gathered["batch"]["input_ids"], input_ids[[4, 1, 3]])
+        assert gathered["non_tensor_batch"]["text"].tolist() == text[[4, 1, 3]].tolist()
+        assert view["batch_fields"]["input_ids"]["stage"] == "rollout"
+        assert view["batch_fields"]["old_log_probs"]["stage"] == "rollout"
+        assert view["non_tensor_fields"]["text"]["stage"] == "rollout"
+        assert set(ref.encoded_non_tensor) == {"text"}
+    finally:
+        transfer.cleanup_dataproto(ref)
+
+
+def test_dataproto_helper_reads_rollout_transfer_data_with_real_store() -> None:
+    pytest.importorskip("mooncake.store")
+    _store, transfer = real_transfer("structured-test-rollout-transfer")
+    row_ids = [f"rollout-row-{index}" for index in range(6)]
+    input_ids = np.arange(48, dtype=np.int64).reshape(6, 8)
+    attention_mask = (input_ids % 3 != 0).astype(np.int32)
+    position_ids = np.tile(np.arange(8, dtype=np.int64), (6, 1))
+    responses = (input_ids[:, -3:] + 100).astype(np.int64)
+    response_mask = np.ones((6, 3), dtype=np.int32)
+    prompts = np.asarray(
+        ["", "prompt-b", "prompt-c", "prompt-d", "prompt-e", "prompt-f"],
+        dtype=object,
+    )
+    sample_meta = np.asarray(
+        [{"row": index, "tag": row_id} for index, row_id in enumerate(row_ids)],
+        dtype=object,
+    )
+    old_log_probs = np.linspace(-0.5, 0.5, 18, dtype=np.float32).reshape(6, 3)
+    ref_log_probs = old_log_probs + np.float32(0.25)
+    rewards = np.linspace(0.0, 1.0, 6, dtype=np.float32)
+    advantages = rewards + np.float32(10.0)
+    returns = rewards + np.float32(20.0)
+    values = rewards + np.float32(30.0)
+    rollout_data = {
+        "batch": {
+            "input_ids": input_ids,
+            "attention_mask": attention_mask,
+            "position_ids": position_ids,
+            "responses": responses,
+            "response_mask": response_mask,
+        },
+        "non_tensor_batch": {
+            "prompts": prompts,
+            "sample_meta": sample_meta,
+        },
+        "meta_info": {"roll_row_ids": row_ids, "global_step": 7},
+    }
+    logprob_data = {
+        "batch": {
+            "old_log_probs": old_log_probs,
+            "ref_log_probs": ref_log_probs,
+        },
+        "non_tensor_batch": {},
+        "meta_info": {"logprob_stage": "complete"},
+    }
+    value_data = {
+        "batch": {
+            "rewards": rewards,
+            "advantages": advantages,
+            "returns": returns,
+            "values": values,
+        },
+        "non_tensor_batch": {},
+        "meta_info": {"critic_stage": "complete"},
+    }
+    ref = transfer.put_dataproto(
+        rollout_data,
+        namespace="roll",
+        partition="remote-batch",
+        stage="rollout",
+    )
+    ref = transfer.append_dataproto_fields(ref, logprob_data, stage="logprob")
+    ref = transfer.append_dataproto_fields(ref, value_data, stage="rollout")
+    handle = export_dataproto_ref(ref)
+
+    try:
+        full = transfer.get_dataproto(handle)
+        selected = transfer.get_dataproto(
+            handle,
+            fields=["input_ids", "old_log_probs", "advantages", "prompts"],
+            meta_info_keys=["roll_row_ids"],
+        )
+        sliced = transfer.get_dataproto(
+            handle,
+            fields=["responses", "ref_log_probs", "returns", "sample_meta"],
+            rows=slice(1, 5),
+        )
+        gathered = transfer.get_dataproto(
+            handle,
+            fields=["attention_mask", "values", "prompts", "sample_meta"],
+            rows=[5, 0, 3],
+        )
+        imported = import_dataproto_ref(handle)
+        imported_selected = transfer.get_dataproto(
+            imported,
+            batch_fields=["position_ids", "rewards"],
+            rows=[2, 4],
+        )
+        view = transfer.dataproto_manifest_view(handle)
+
+        assert full["meta_info"] == {
+            "roll_row_ids": row_ids,
+            "global_step": 7,
+            "logprob_stage": "complete",
+            "critic_stage": "complete",
+        }
+        assert np.array_equal(full["batch"]["input_ids"], input_ids)
+        assert np.array_equal(full["batch"]["attention_mask"], attention_mask)
+        assert np.array_equal(full["batch"]["position_ids"], position_ids)
+        assert np.array_equal(full["batch"]["responses"], responses)
+        assert np.array_equal(full["batch"]["response_mask"], response_mask)
+        assert np.array_equal(full["batch"]["old_log_probs"], old_log_probs)
+        assert np.array_equal(full["batch"]["ref_log_probs"], ref_log_probs)
+        assert np.array_equal(full["batch"]["rewards"], rewards)
+        assert np.array_equal(full["batch"]["advantages"], advantages)
+        assert np.array_equal(full["batch"]["returns"], returns)
+        assert np.array_equal(full["batch"]["values"], values)
+        assert full["non_tensor_batch"]["prompts"].tolist() == prompts.tolist()
+        assert full["non_tensor_batch"]["sample_meta"].tolist() == sample_meta.tolist()
+
+        assert set(selected["batch"]) == {"input_ids", "old_log_probs", "advantages"}
+        assert set(selected["non_tensor_batch"]) == {"prompts"}
+        assert selected["meta_info"] == {"roll_row_ids": row_ids}
+        assert np.array_equal(selected["batch"]["input_ids"], input_ids)
+        assert np.array_equal(selected["batch"]["old_log_probs"], old_log_probs)
+        assert np.array_equal(selected["batch"]["advantages"], advantages)
+        assert selected["non_tensor_batch"]["prompts"].tolist() == prompts.tolist()
+
+        assert np.array_equal(sliced["batch"]["responses"], responses[1:5])
+        assert np.array_equal(sliced["batch"]["ref_log_probs"], ref_log_probs[1:5])
+        assert np.array_equal(sliced["batch"]["returns"], returns[1:5])
+        assert sliced["non_tensor_batch"]["sample_meta"].tolist() == sample_meta[1:5].tolist()
+
+        assert np.array_equal(gathered["batch"]["attention_mask"], attention_mask[[5, 0, 3]])
+        assert np.array_equal(gathered["batch"]["values"], values[[5, 0, 3]])
+        assert gathered["non_tensor_batch"]["prompts"].tolist() == prompts[[5, 0, 3]].tolist()
+        assert gathered["non_tensor_batch"]["sample_meta"].tolist() == sample_meta[[5, 0, 3]].tolist()
+
+        assert np.array_equal(imported_selected["batch"]["position_ids"], position_ids[[2, 4]])
+        assert np.array_equal(imported_selected["batch"]["rewards"], rewards[[2, 4]])
+        assert view["batch_fields"]["input_ids"]["stage"] == "rollout"
+        assert view["batch_fields"]["old_log_probs"]["stage"] == "logprob"
+        assert view["batch_fields"]["advantages"]["stage"] == "rollout"
+        assert view["non_tensor_fields"]["prompts"]["stage"] == "rollout"
+    finally:
+        transfer.cleanup_dataproto(ref)
+
+
+def test_dataproto_helper_reads_large_rollout_transfer_data_with_real_store() -> None:
+    pytest.importorskip("mooncake.store")
+    _store, transfer = real_transfer("structured-test-large-rollout-transfer")
+    batch_size = 96
+    prompt_len = 128
+    response_len = 64
+    total_len = prompt_len + response_len
+    row_ids = [f"large-rollout-row-{index}" for index in range(batch_size)]
+    token_grid = np.arange(batch_size * total_len, dtype=np.int64).reshape(
+        batch_size, total_len
+    )
+    input_ids = token_grid + 1000
+    attention_mask = (token_grid % 7 != 0).astype(np.int32)
+    position_ids = np.tile(np.arange(total_len, dtype=np.int64), (batch_size, 1))
+    responses = input_ids[:, prompt_len:]
+    response_mask = np.ones((batch_size, response_len), dtype=np.int32)
+    action_log_probs = np.linspace(
+        -3.0, 3.0, batch_size * response_len, dtype=np.float32
+    ).reshape(batch_size, response_len)
+    ref_log_probs = action_log_probs + np.float32(0.125)
+    values = np.linspace(0.0, 1.0, batch_size, dtype=np.float32)
+    rewards = values + np.float32(1.0)
+    advantages = values + np.float32(2.0)
+    returns = values + np.float32(3.0)
+    prompts = np.asarray(
+        ["" if index % 17 == 0 else f"prompt-{index}-" + "x" * (index % 23) for index in range(batch_size)],
+        dtype=object,
+    )
+    sample_meta = np.asarray(
+        [
+            {"row": index, "row_id": row_id, "prompt_len": prompt_len, "response_len": response_len}
+            for index, row_id in enumerate(row_ids)
+        ],
+        dtype=object,
+    )
+    ref = transfer.put_dataproto(
+        {
+            "batch": {
+                "input_ids": input_ids,
+                "attention_mask": attention_mask,
+                "position_ids": position_ids,
+                "responses": responses,
+                "response_mask": response_mask,
+            },
+            "non_tensor_batch": {
+                "prompts": prompts,
+                "sample_meta": sample_meta,
+            },
+            "meta_info": {"roll_row_ids": row_ids, "global_step": 11},
+        },
+        namespace="roll",
+        partition="large-remote-batch",
+        stage="rollout",
+    )
+    ref = transfer.append_dataproto_fields(
+        ref,
+        {
+            "batch": {
+                "action_log_probs": action_log_probs,
+                "ref_log_probs": ref_log_probs,
+            },
+            "non_tensor_batch": {},
+            "meta_info": {"logprob_stage": "complete"},
+        },
+        stage="logprob",
+    )
+    ref = transfer.append_dataproto_fields(
+        ref,
+        {
+            "batch": {
+                "values": values,
+                "rewards": rewards,
+                "advantages": advantages,
+                "returns": returns,
+            },
+            "non_tensor_batch": {},
+            "meta_info": {"critic_stage": "complete"},
+        },
+        stage="rollout",
+    )
+    handle = export_dataproto_ref(ref)
+    gathered_rows = [95, 0, 63, 17, 42]
+
+    try:
+        selected = transfer.get_dataproto(
+            handle,
+            fields=["input_ids", "action_log_probs", "values", "prompts"],
+            rows=slice(12, 76),
+        )
+        gathered = transfer.get_dataproto(
+            handle,
+            fields=["responses", "ref_log_probs", "returns", "sample_meta"],
+            rows=gathered_rows,
+        )
+        destination = np.empty((len(gathered_rows), total_len), dtype=np.int64)
+        into = transfer.get_dataproto(
+            handle,
+            batch_fields=["position_ids"],
+            rows=gathered_rows,
+            destinations={"position_ids": destination},
+        )
+        meta_only = transfer.get_dataproto(handle, fields=[], meta_info_keys=["roll_row_ids"])
+        full_tail = transfer.get_dataproto(
+            handle,
+            fields=["attention_mask", "advantages"],
+            rows=slice(batch_size - 8, batch_size),
+        )
+
+        assert np.array_equal(selected["batch"]["input_ids"], input_ids[12:76])
+        assert np.array_equal(
+            selected["batch"]["action_log_probs"], action_log_probs[12:76]
+        )
+        assert np.array_equal(selected["batch"]["values"], values[12:76])
+        assert selected["non_tensor_batch"]["prompts"].tolist() == prompts[12:76].tolist()
+        assert np.array_equal(gathered["batch"]["responses"], responses[gathered_rows])
+        assert np.array_equal(
+            gathered["batch"]["ref_log_probs"], ref_log_probs[gathered_rows]
+        )
+        assert np.array_equal(gathered["batch"]["returns"], returns[gathered_rows])
+        assert gathered["non_tensor_batch"]["sample_meta"].tolist() == sample_meta[gathered_rows].tolist()
+        assert into["batch"]["position_ids"] is destination
+        assert np.array_equal(destination, position_ids[gathered_rows])
+        assert meta_only["batch"] == {}
+        assert meta_only["non_tensor_batch"] == {}
+        assert meta_only["meta_info"] == {"roll_row_ids": row_ids}
+        assert np.array_equal(
+            full_tail["batch"]["attention_mask"], attention_mask[batch_size - 8 :]
+        )
+        assert np.array_equal(full_tail["batch"]["advantages"], advantages[batch_size - 8 :])
+    finally:
+        transfer.cleanup_dataproto(ref)
+
+
+def test_dataproto_helper_selection_errors_and_destinations_with_real_store() -> None:
+    pytest.importorskip("mooncake.store")
+    _store, transfer = real_transfer("structured-test-selection-destinations")
+    batch_size = 12
+    input_ids = np.arange(batch_size * 12, dtype=np.int64).reshape(batch_size, 12)
+    scores = np.linspace(0.0, 1.0, batch_size, dtype=np.float32)
+    prompts = np.asarray([f"prompt-{index}" for index in range(batch_size)], dtype=object)
+    ref = transfer.put_dataproto(
+        {
+            "batch": {"input_ids": input_ids, "scores": scores},
+            "non_tensor_batch": {"prompts": prompts},
+            "meta_info": {"step": 1},
+        },
+        namespace="roll",
+        partition="selection-destinations",
+        stage="rollout",
+    )
+
+    try:
+        with pytest.raises(ValueError, match="fields cannot be combined"):
+            transfer.get_dataproto(ref, fields=["input_ids"], batch_fields=["scores"])
+        with pytest.raises(KeyError, match="unknown DataProto fields"):
+            transfer.get_dataproto(ref, fields=["missing"])
+        with pytest.raises(IndexError, match="out of range"):
+            transfer.get_dataproto(ref, fields=["input_ids"], rows=[batch_size])
+        with pytest.raises(TypeError, match="row indices"):
+            transfer.get_dataproto(ref, fields=["input_ids"], rows=["0"])
+        with pytest.raises(ValueError, match="step must be positive"):
+            transfer.get_dataproto(ref, fields=["input_ids"], rows=slice(None, None, -1))
+        with pytest.raises(ValueError, match="destination shape mismatch"):
+            transfer.get_dataproto(
+                ref,
+                batch_fields=["input_ids"],
+                rows=[0, 1, 2],
+                destinations={"input_ids": np.empty((2, 12), dtype=np.int64)},
+            )
+
+        scores_destination = np.empty((4,), dtype=np.float32)
+        result = transfer.get_dataproto(
+            ref,
+            batch_fields=["scores"],
+            rows=[11, 3, 3, 0],
+            destinations={"scores": scores_destination},
+        )
+        non_tensor_only = transfer.get_dataproto(
+            ref,
+            non_tensor_fields=["prompts"],
+            rows=[2, 4, 6],
+        )
+
+        assert result["batch"]["scores"] is scores_destination
+        assert np.array_equal(scores_destination, scores[[11, 3, 3, 0]])
+        assert result["non_tensor_batch"] == {}
+        assert non_tensor_only["batch"] == {}
+        assert non_tensor_only["non_tensor_batch"]["prompts"].tolist() == prompts[[2, 4, 6]].tolist()
+    finally:
+        transfer.cleanup_dataproto(ref)
+
+
+def test_dataproto_helper_imported_handle_same_stage_append_with_real_store() -> None:
+    pytest.importorskip("mooncake.store")
+    _store, transfer = real_transfer("structured-test-imported-handle-append")
+    batch_size = 8
+    input_ids = np.arange(batch_size * 6, dtype=np.int64).reshape(batch_size, 6)
+    masks = np.ones((batch_size, 6), dtype=np.int32)
+    values = np.linspace(1.0, 2.0, batch_size, dtype=np.float32)
+    rewards = values + np.float32(5.0)
+    ref = transfer.put_dataproto(
+        {
+            "batch": {"input_ids": input_ids, "masks": masks},
+            "non_tensor_batch": {},
+            "meta_info": {"step": 2},
+        },
+        namespace="roll",
+        partition="imported-handle-append",
+        stage="rollout",
+    )
+    imported = import_dataproto_ref(export_dataproto_ref(ref))
+    appended = transfer.append_dataproto_fields(
+        imported,
+        {
+            "batch": {"values": values, "rewards": rewards},
+            "non_tensor_batch": {},
+            "meta_info": {"critic": True},
+        },
+        stage="rollout",
+    )
+    handle = export_dataproto_ref(appended)
+
+    try:
+        result = transfer.get_dataproto(handle, fields=["input_ids", "values", "rewards"], rows=[7, 1, 4])
+        assert np.array_equal(result["batch"]["input_ids"], input_ids[[7, 1, 4]])
+        assert np.array_equal(result["batch"]["values"], values[[7, 1, 4]])
+        assert np.array_equal(result["batch"]["rewards"], rewards[[7, 1, 4]])
+        assert result["meta_info"] == {"step": 2, "critic": True}
+    finally:
+        transfer.cleanup_dataproto(appended)
+
+
+def test_bundle_manifest_rejects_tampered_cleanup_keys() -> None:
+    store, transfer = make_transfer()
+    ref = transfer.put_bundle(b"meta", {"payload": b"abcdef"})
+    tampered = dict(ref.manifest)
+    tampered["cleanup_keys"] = ["other/object"]
+    write_manifest(store, ref.manifest_key, tampered)
+
+    with pytest.raises(ValueError, match="cleanup_keys"):
+        transfer.remove_bundle({"manifest_key": ref.manifest_key})
+
+
+def test_dataproto_helper_rollout_edge_cases_with_real_store() -> None:
+    pytest.importorskip("mooncake.store")
+    _store, transfer = real_transfer("structured-test-rollout-edge-cases")
+    batch_size = 10
+    row_ids = [f"edge-row-{index}" for index in range(batch_size)]
+    input_ids = np.arange(batch_size * 16, dtype=np.int64).reshape(batch_size, 16)
+    attention_mask = (input_ids % 2 == 0).astype(np.int32)
+    prompts = np.asarray(
+        ["", "short", None, "三", "bytes", "long-" + "x" * 64, "7", "8", "9", "tail"],
+        dtype=object,
+    )
+    scores = np.linspace(-1.0, 1.0, batch_size, dtype=np.float32)
+    values = scores + np.float32(2.0)
+    replacement_scores = scores + np.float32(10.0)
+    ref = transfer.put_dataproto(
+        {
+            "batch": {"input_ids": input_ids, "attention_mask": attention_mask},
+            "non_tensor_batch": {"prompts": prompts},
+            "meta_info": {"roll_row_ids": row_ids},
+        },
+        namespace="roll",
+        partition="edge-remote-batch",
+        stage="rollout",
+    )
+    ref = transfer.append_dataproto_fields(
+        ref,
+        {"batch": {"scores": scores}, "non_tensor_batch": {}, "meta_info": {}},
+        stage="scores",
+    )
+
+    try:
+        with pytest.raises(ValueError, match="already exist"):
+            transfer.append_dataproto_fields(
+                ref,
+                {"batch": {"scores": scores}, "non_tensor_batch": {}, "meta_info": {}},
+                stage="other_scores",
+            )
+        with pytest.raises(ValueError, match="batch size"):
+            transfer.append_dataproto_fields(
+                ref,
+                {
+                    "batch": {"bad": np.arange(batch_size - 1, dtype=np.int64)},
+                    "non_tensor_batch": {},
+                    "meta_info": {},
+                },
+                stage="bad_size",
+            )
+
+        empty = transfer.get_dataproto(
+            ref,
+            fields=["input_ids", "scores", "prompts"],
+            rows=[],
+        )
+        tail_and_repeat = transfer.get_dataproto(
+            ref,
+            fields=["input_ids", "attention_mask", "scores", "prompts"],
+            rows=[-1, 0, -1, 3],
+        )
+        step_slice = transfer.get_dataproto(
+            ref,
+            batch_fields=["input_ids", "scores"],
+            rows=slice(1, 9, 2),
+        )
+        ref = transfer.append_dataproto_fields(
+            ref,
+            {
+                "batch": {"values": values},
+                "non_tensor_batch": {},
+                "meta_info": {"same_stage": True},
+            },
+            stage="rollout",
+        )
+        same_stage = transfer.get_dataproto(
+            ref,
+            fields=["input_ids", "values", "prompts"],
+            rows=[2, 5, 9],
+        )
+        ref = transfer.append_dataproto_fields(
+            ref,
+            {
+                "batch": {"scores": replacement_scores},
+                "non_tensor_batch": {},
+                "meta_info": {"scores_overwritten": True},
+            },
+            stage="scores",
+            overwrite=True,
+        )
+        overwritten = transfer.get_dataproto(
+            ref,
+            fields=["scores", "values"],
+            rows=[0, 4, 9],
+        )
+
+        assert empty["batch"]["input_ids"].shape == (0, 16)
+        assert empty["batch"]["scores"].shape == (0,)
+        assert empty["non_tensor_batch"]["prompts"].tolist() == []
+        assert np.array_equal(tail_and_repeat["batch"]["input_ids"], input_ids[[9, 0, 9, 3]])
+        assert np.array_equal(
+            tail_and_repeat["batch"]["attention_mask"], attention_mask[[9, 0, 9, 3]]
+        )
+        assert np.array_equal(tail_and_repeat["batch"]["scores"], scores[[9, 0, 9, 3]])
+        assert tail_and_repeat["non_tensor_batch"]["prompts"].tolist() == prompts[[9, 0, 9, 3]].tolist()
+        assert np.array_equal(step_slice["batch"]["input_ids"], input_ids[1:9:2])
+        assert np.array_equal(step_slice["batch"]["scores"], scores[1:9:2])
+        assert np.array_equal(same_stage["batch"]["input_ids"], input_ids[[2, 5, 9]])
+        assert np.array_equal(same_stage["batch"]["values"], values[[2, 5, 9]])
+        assert same_stage["non_tensor_batch"]["prompts"].tolist() == prompts[[2, 5, 9]].tolist()
+        assert np.array_equal(overwritten["batch"]["scores"], replacement_scores[[0, 4, 9]])
+        assert np.array_equal(overwritten["batch"]["values"], values[[0, 4, 9]])
+        assert overwritten["meta_info"]["same_stage"] is True
+        assert overwritten["meta_info"]["scores_overwritten"] is True
+    finally:
+        transfer.cleanup_dataproto(ref)
+
+    with pytest.raises(Exception):
+        transfer.get_dataproto(ref, fields=["input_ids"])
+
+
+def test_structured_object_zero_byte_payload_skips_store_put() -> None:
+    store, transfer = make_transfer()
+    ref = transfer.put_structured_object(
+        StructuredObjectPayload(buffers={"empty": np.empty((4, 0), dtype=np.float32)})
+    )
+    result = transfer.materialize(transfer.read_spec(ref))
+
+    assert result.objects["empty"].shape == (4, 0)
+    assert ref.manifest["buffers"]["empty"]["bytes"] == 0
+    assert ref.manifest["buffers"]["empty"]["chunks"] == []
+    assert store.objects == {ref.manifest_key: store.objects[ref.manifest_key], ref.manifest["meta"]["key"]: store.objects[ref.manifest["meta"]["key"]]}
+
+
+def test_dataproto_helper_multidim_boundary_reads_with_real_store() -> None:
+    pytest.importorskip("mooncake.store")
+    _store, transfer = real_transfer("structured-test-multidim-boundaries")
+    batch_size = 18
+    image_features = np.arange(batch_size * 3 * 16 * 16, dtype=np.float32).reshape(
+        batch_size, 3, 16, 16
+    )
+    logits = np.linspace(-2.0, 2.0, batch_size * 4 * 5 * 6, dtype=np.float32).reshape(
+        batch_size, 4, 5, 6
+    )
+    token_matrix = np.arange(batch_size * 7 * 9, dtype=np.int64).reshape(
+        batch_size, 7, 9
+    )
+    zero_width = np.empty((batch_size, 0), dtype=np.float32)
+    row_scores = np.linspace(10.0, 20.0, batch_size, dtype=np.float32)
+    byte_blobs = np.asarray(
+        [
+            b"",
+            b"alpha",
+            bytearray(b"beta"),
+            memoryview(b"gamma"),
+            *[f"blob-{index}".encode() for index in range(4, batch_size)],
+        ],
+        dtype=object,
+    )
+    json_meta = np.asarray(
+        [
+            None if index % 5 == 0 else {"row": index, "shape": [3, 16, 16]}
+            for index in range(batch_size)
+        ],
+        dtype=object,
+    )
+    ref = transfer.put_dataproto(
+        {
+            "batch": {
+                "image_features": image_features,
+                "logits": logits,
+                "token_matrix": token_matrix,
+                "zero_width": zero_width,
+            },
+            "non_tensor_batch": {
+                "byte_blobs": byte_blobs,
+                "json_meta": json_meta,
+            },
+            "meta_info": {"batch_size": batch_size, "layout": "multidim"},
+        },
+        namespace="roll",
+        partition="multidim-boundaries",
+        stage="rollout",
+    )
+    ref = transfer.append_dataproto_fields(
+        ref,
+        {
+            "batch": {"row_scores": row_scores},
+            "non_tensor_batch": {},
+            "meta_info": {"score_stage": "same-stage"},
+        },
+        stage="rollout",
+    )
+    handle = export_dataproto_ref(ref)
+    rows = [17, 0, 5, 17, 9]
+
+    try:
+        full = transfer.get_dataproto(
+            handle,
+            fields=["image_features", "zero_width", "byte_blobs", "json_meta"],
+        )
+        empty_slice = transfer.get_dataproto(
+            handle,
+            batch_fields=["image_features", "zero_width", "row_scores"],
+            rows=slice(4, 4),
+        )
+        duplicate_gather = transfer.get_dataproto(
+            handle,
+            fields=["logits", "token_matrix", "row_scores", "byte_blobs", "json_meta"],
+            rows=rows,
+        )
+        image_destination = np.empty((len(rows), 3, 16, 16), dtype=np.float32)
+        into = transfer.get_dataproto(
+            handle,
+            batch_fields=["image_features"],
+            rows=rows,
+            destinations={"image_features": image_destination},
+        )
+        step_batch = transfer.get_dataproto(
+            handle,
+            batch_fields=["logits", "token_matrix"],
+            rows=slice(2, 17, 3),
+        )
+        tail_non_tensor = transfer.get_dataproto(
+            handle,
+            non_tensor_fields=["byte_blobs", "json_meta"],
+            rows=slice(batch_size - 3, batch_size),
+        )
+        meta_selected = transfer.get_dataproto(
+            handle,
+            fields=[],
+            meta_info_keys=["layout", "score_stage"],
+        )
+
+        assert np.array_equal(full["batch"]["image_features"], image_features)
+        assert full["batch"]["zero_width"].shape == (batch_size, 0)
+        assert full["non_tensor_batch"]["byte_blobs"].tolist() == [
+            bytes(value) for value in byte_blobs.tolist()
+        ]
+        assert full["non_tensor_batch"]["json_meta"].tolist() == json_meta.tolist()
+        assert empty_slice["batch"]["image_features"].shape == (0, 3, 16, 16)
+        assert empty_slice["batch"]["zero_width"].shape == (0, 0)
+        assert empty_slice["batch"]["row_scores"].shape == (0,)
+        assert np.array_equal(duplicate_gather["batch"]["logits"], logits[rows])
+        assert np.array_equal(
+            duplicate_gather["batch"]["token_matrix"], token_matrix[rows]
+        )
+        assert np.array_equal(duplicate_gather["batch"]["row_scores"], row_scores[rows])
+        assert duplicate_gather["non_tensor_batch"]["byte_blobs"].tolist() == [
+            bytes(value) for value in byte_blobs[rows].tolist()
+        ]
+        assert duplicate_gather["non_tensor_batch"]["json_meta"].tolist() == json_meta[rows].tolist()
+        assert into["batch"]["image_features"] is image_destination
+        assert np.array_equal(image_destination, image_features[rows])
+        assert np.array_equal(step_batch["batch"]["logits"], logits[2:17:3])
+        assert np.array_equal(step_batch["batch"]["token_matrix"], token_matrix[2:17:3])
+        assert tail_non_tensor["batch"] == {}
+        assert tail_non_tensor["non_tensor_batch"]["byte_blobs"].tolist() == [
+            bytes(value) for value in byte_blobs[batch_size - 3 :].tolist()
+        ]
+        assert tail_non_tensor["non_tensor_batch"]["json_meta"].tolist() == json_meta[batch_size - 3 :].tolist()
+        assert meta_selected["batch"] == {}
+        assert meta_selected["non_tensor_batch"] == {}
+        assert meta_selected["meta_info"] == {
+            "layout": "multidim",
+            "score_stage": "same-stage",
+        }
+    finally:
+        transfer.cleanup_dataproto(ref)
+
+
+def test_dataproto_helper_rejects_inconsistent_batch_sizes() -> None:
+    _store, transfer = make_transfer()
+    data = SimpleDataProto(
+        batch={
+            "input_ids": np.arange(4, dtype=np.int64),
+            "attention_mask": np.arange(3, dtype=np.int64),
+        }
+    )
+
+    with pytest.raises(ValueError, match="inconsistent batch sizes"):
+        transfer.put_dataproto(data)
+
+
+def test_dataproto_helper_overwrite_rejects_partial_stage_replacement() -> None:
+    _store, transfer = make_transfer()
+    ref = transfer.put_dataproto(
+        SimpleDataProto(
+            batch={"input_ids": np.arange(4), "attention_mask": np.arange(4)}
+        ),
+        stage="rollout",
+    )
+
+    with pytest.raises(ValueError, match="must include existing fields"):
+        transfer.append_dataproto_fields(
+            ref,
+            SimpleDataProto(batch={"input_ids": np.arange(4)}),
+            stage="rollout",
+            overwrite=True,
+        )
+
+
+def test_dataproto_helper_overwrite_replaces_encoded_metadata() -> None:
+    store, transfer = make_transfer()
+    ref = transfer.put_dataproto(
+        SimpleDataProto(
+            non_tensor_batch={"text": np.asarray(["a", None], dtype=object)}
+        ),
+        stage="meta",
+    )
+    old_manifest_key = ref.stage_refs["meta"].manifest_key
+    ref = transfer.append_dataproto_fields(
+        ref,
+        SimpleDataProto(non_tensor_batch={"text": np.asarray([1, 2], dtype=np.int64)}),
+        stage="meta",
+        overwrite=True,
+    )
+    result = transfer.get_dataproto(ref)
+
+    assert "text" not in ref.encoded_non_tensor
+    assert old_manifest_key not in store.objects
+    assert ref.stage_refs["meta"].manifest_key in store.objects
+    assert np.array_equal(result["non_tensor_batch"]["text"], np.asarray([1, 2]))
+
+
+def test_dataproto_helper_cleanup_removes_all_stage_objects() -> None:
+    store, transfer = make_transfer()
+    ref = transfer.put_dataproto(
+        SimpleDataProto(batch={"input_ids": np.arange(4, dtype=np.int64)}),
+        stage="rollout",
+    )
+    ref = transfer.append_dataproto_fields(
+        ref,
+        SimpleDataProto(batch={"values": np.arange(4, dtype=np.float32)}),
+        stage="critic",
+    )
+
+    assert store.objects
+    transfer.cleanup_dataproto(ref)
+
+    assert store.objects == {}
+
+
+def test_dataproto_helper_object_non_tensor_codecs_roundtrip() -> None:
+    _store, transfer = make_transfer()
+    data = SimpleDataProto(
+        batch={"input_ids": np.arange(4, dtype=np.int64)},
+        non_tensor_batch={
+            "text": np.asarray(["hello", None, "world", "moon"], dtype=object),
+            "json": np.asarray(
+                [{"a": 1}, None, {"b": [2, 3]}, {"c": "x"}], dtype=object
+            ),
+            "blob": np.asarray(
+                [b"a", None, bytearray(b"bc"), memoryview(b"def")], dtype=object
+            ),
+            "nullable_int": np.asarray([1, None, 3, 4], dtype=object),
+        },
+    )
+
+    ref = transfer.put_dataproto(data)
+    result = transfer.get_dataproto(ref)
+
+    assert set(ref.encoded_non_tensor) == {"text", "json", "blob", "nullable_int"}
+    assert ref.encoded_non_tensor["json"]["codec"] == "msgpack_ragged"
+    assert result["non_tensor_batch"]["text"].tolist() == [
+        "hello",
+        None,
+        "world",
+        "moon",
+    ]
+    assert result["non_tensor_batch"]["json"].tolist() == [
+        {"a": 1},
+        None,
+        {"b": [2, 3]},
+        {"c": "x"},
+    ]
+    assert result["non_tensor_batch"]["blob"].tolist() == [b"a", None, b"bc", b"def"]
+    assert result["non_tensor_batch"]["nullable_int"].tolist() == [1, None, 3, 4]
+
+
+def test_dataproto_helper_rejects_unsupported_object_non_tensor() -> None:
+    _store, transfer = make_transfer()
+    data = SimpleDataProto(
+        non_tensor_batch={"fallback": np.asarray([object(), None], dtype=object)}
+    )
+
+    with pytest.raises(ValueError, match="unsupported structured non-tensor field"):
+        transfer.put_dataproto(data)
+
+
+def test_dataproto_helper_ragged_tensor_non_tensor_roundtrip() -> None:
+    torch = pytest.importorskip("torch")
+    _store, transfer = make_transfer()
+    ragged = np.asarray(
+        [
+            torch.arange(1, dtype=torch.float32),
+            None,
+            torch.arange(3, dtype=torch.float32),
+            torch.arange(2, dtype=torch.float32),
+        ],
+        dtype=object,
+    )
+    data = SimpleDataProto(non_tensor_batch={"ragged": ragged})
+
+    ref = transfer.put_dataproto(data)
+    result = transfer.get_dataproto(ref)
+
+    assert ref.encoded_non_tensor["ragged"]["codec"] == "ragged_tensor"
+    actual = result["non_tensor_batch"]["ragged"]
+    assert torch.equal(actual[0], ragged[0])
+    assert actual[1] is None
+    assert torch.equal(actual[2], ragged[2])
+    assert torch.equal(actual[3], ragged[3])
+
+
+def _assert_tensor_object_equal(actual, expected) -> None:
+    torch = pytest.importorskip("torch")
+    if expected is None:
+        assert actual is None
+        return
+    if isinstance(expected, dict):
+        assert isinstance(actual, dict)
+        assert set(actual) == set(expected)
+        for key, value in expected.items():
+            _assert_tensor_object_equal(actual[key], value)
+        return
+    if isinstance(expected, list):
+        assert isinstance(actual, list)
+        assert len(actual) == len(expected)
+        for actual_item, expected_item in zip(actual, expected):
+            _assert_tensor_object_equal(actual_item, expected_item)
+        return
+    if isinstance(expected, torch.Tensor):
+        assert torch.equal(actual, expected)
+        return
+    assert actual == expected
+
+
+def test_dataproto_helper_dict_of_tensors_object_array_roundtrip() -> None:
+    torch = pytest.importorskip("torch")
+    _store, transfer = make_transfer()
+    samples = np.asarray(
+        [
+            {"tokens": torch.arange(2, dtype=torch.int64), "reward": 1.0},
+            {"tokens": torch.arange(3, dtype=torch.int64), "reward": 2.0},
+            {"tokens": torch.arange(1, dtype=torch.int64), "reward": None},
+            {"tokens": None, "reward": 4.0},
+        ],
+        dtype=object,
+    )
+
+    ref = transfer.put_dataproto(
+        SimpleDataProto(
+            batch={"input_ids": np.arange(8, dtype=np.int64).reshape(4, 2)},
+            non_tensor_batch={"samples": samples},
+            meta_info={"source": "dict-of-tensors"},
+        )
+    )
+    result = transfer.get_dataproto(ref)
+    view = transfer.dataproto_manifest_view(ref)
+
+    assert ref.encoded_non_tensor["samples"]["codec"] == "structured_recursive"
+    assert view["non_tensor_fields"]["samples"]["spec"]["codec"] == "structured_recursive"
+    assert result["meta_info"] == {"source": "dict-of-tensors"}
+    actual = result["non_tensor_batch"]["samples"]
+    for row, expected in enumerate(samples):
+        _assert_tensor_object_equal(actual[row], expected)
+
+
+def test_dataproto_helper_dict_of_tensors_distinguishes_missing_keys_and_nulls() -> None:
+    torch = pytest.importorskip("torch")
+    _store, transfer = make_transfer()
+    samples = np.asarray(
+        [
+            {"tokens": torch.arange(2), "label": None},
+            {"label": 1},
+            None,
+            {"tokens": None, "label": 3},
+        ],
+        dtype=object,
+    )
+
+    ref = transfer.put_dataproto(SimpleDataProto(non_tensor_batch={"samples": samples}))
+    actual = transfer.get_dataproto(ref)["non_tensor_batch"]["samples"]
+
+    assert "tokens" in actual[0]
+    assert actual[0]["label"] is None
+    assert "tokens" not in actual[1]
+    assert actual[1]["label"] == 1
+    assert actual[2] is None
+    assert "tokens" in actual[3]
+    assert actual[3]["tokens"] is None
+    assert actual[3]["label"] == 3
+
+
+def test_dataproto_helper_nested_tensor_object_array_rows() -> None:
+    torch = pytest.importorskip("torch")
+    _store, transfer = make_transfer()
+    samples = np.asarray(
+        [
+            {"images": [{"pixels": torch.arange(2)}], "meta": {"rank": 0}},
+            {"images": [{"pixels": torch.arange(3)}, {"pixels": None}], "meta": {}},
+            {"images": [], "meta": {"rank": None}},
+        ],
+        dtype=object,
+    )
+
+    ref = transfer.put_dataproto(SimpleDataProto(non_tensor_batch={"samples": samples}))
+    actual = transfer.get_dataproto(ref)["non_tensor_batch"]["samples"]
+
+    for row, expected in enumerate(samples):
+        _assert_tensor_object_equal(actual[row], expected)
+
+
+def test_dataproto_helper_dict_of_native_object_leaves_uses_recursive_codec() -> None:
+    _store, transfer = make_transfer()
+    samples = np.asarray(
+        [
+            {
+                "media": [b"a", b"bc"],
+                "blob": b"payload-0",
+                "scores": np.asarray([1, 2], dtype=np.int64),
+                "label": "x",
+            },
+            {
+                "media": [],
+                "blob": b"payload-1",
+                "scores": np.asarray([3], dtype=np.int64),
+                "label": "y",
+            },
+            {"label": "missing-native"},
+            None,
+        ],
+        dtype=object,
+    )
+
+    ref = transfer.put_dataproto(SimpleDataProto(non_tensor_batch={"samples": samples}))
+    result = transfer.get_dataproto(ref)
+    actual = result["non_tensor_batch"]["samples"]
+
+    assert ref.encoded_non_tensor["samples"]["codec"] == "structured_recursive"
+    assert actual[0]["media"] == [b"a", b"bc"]
+    assert actual[0]["blob"] == b"payload-0"
+    assert actual[0]["scores"] == [1, 2]
+    assert actual[0]["label"] == "x"
+    assert actual[1]["media"] == []
+    assert actual[1]["blob"] == b"payload-1"
+    assert actual[1]["scores"] == [3]
+    assert actual[1]["label"] == "y"
+    assert actual[2] == {"label": "missing-native"}
+    assert actual[3] is None
+
+
+def test_dataproto_helper_reads_dict_of_tensors_rows_and_selected_field() -> None:
+    torch = pytest.importorskip("torch")
+    _store, transfer = make_transfer()
+    samples = np.asarray(
+        [
+            {"tokens": torch.arange(i + 1, dtype=torch.int64), "row": i}
+            for i in range(6)
+        ],
+        dtype=object,
+    )
+    samples[2] = {"tokens": None, "row": 2}
+    samples[4] = None
+    input_ids = np.arange(12, dtype=np.int64).reshape(6, 2)
+    ref = transfer.put_dataproto(
+        SimpleDataProto(
+            batch={"input_ids": input_ids},
+            non_tensor_batch={"samples": samples},
+            meta_info={"kind": "dict-tensor"},
+        )
+    )
+
+    sliced = transfer.get_dataproto(ref, fields=["samples"], rows=slice(1, 5))
+    gathered = transfer.get_dataproto(ref, fields=["input_ids", "samples"], rows=[5, 0, 2, 4])
+
+    assert sliced["batch"] == {}
+    assert sliced["meta_info"] == {"kind": "dict-tensor"}
+    for row, expected in enumerate(samples[1:5]):
+        _assert_tensor_object_equal(sliced["non_tensor_batch"]["samples"][row], expected)
+    assert np.array_equal(gathered["batch"]["input_ids"], input_ids[[5, 0, 2, 4]])
+    for row, expected_index in enumerate([5, 0, 2, 4]):
+        _assert_tensor_object_equal(
+            gathered["non_tensor_batch"]["samples"][row], samples[expected_index]
+        )
+
+
+def test_dataproto_helper_recursive_manifest_export_append_and_overwrite() -> None:
+    torch = pytest.importorskip("torch")
+    store, transfer = make_transfer()
+    input_ids = np.arange(4, dtype=np.int64)
+    ref = transfer.put_dataproto(SimpleDataProto(batch={"input_ids": input_ids}), stage="rollout")
+    samples = np.asarray(
+        [
+            {"tokens": torch.arange(1, dtype=torch.float32), "score": 0.0},
+            {"tokens": torch.arange(2, dtype=torch.float32), "score": 1.0},
+            {"tokens": torch.arange(3, dtype=torch.float32), "score": None},
+            {"tokens": None, "score": 3.0},
+        ],
+        dtype=object,
+    )
+
+    ref = transfer.append_dataproto_fields(
+        ref,
+        SimpleDataProto(non_tensor_batch={"samples": samples}, meta_info={"stage": "samples"}),
+        stage="rollout",
+    )
+    handle = export_dataproto_ref(ref)
+    json.dumps(handle)
+    imported = import_dataproto_ref(handle)
+    result = transfer.get_dataproto(imported)
+    view = transfer.dataproto_manifest_view(imported)
+
+    assert np.array_equal(result["batch"]["input_ids"], input_ids)
+    _assert_tensor_object_equal(result["non_tensor_batch"]["samples"][1], samples[1])
+    spec = view["non_tensor_fields"]["samples"]["spec"]
+    assert spec["codec"] == "structured_recursive"
+    assert spec["metadata"]["nodes"]
+    assert spec["metadata"]["leaves"]
+    assert any(name.endswith(".missing") for name in spec["payload_members"])
+
+    old_manifest_key = ref.stage_refs["rollout"].manifest_key
+    ref = transfer.append_dataproto_fields(
+        ref,
+        SimpleDataProto(
+            batch={"input_ids": input_ids},
+            non_tensor_batch={"samples": np.asarray(["a", "b", "c", "d"], dtype=object)},
+        ),
+        stage="rollout",
+        overwrite=True,
+    )
+    overwritten = transfer.get_dataproto(ref)
+    assert ref.encoded_non_tensor["samples"]["codec"] == "utf8_ragged"
+    assert overwritten["non_tensor_batch"]["samples"].tolist() == ["a", "b", "c", "d"]
+    assert old_manifest_key not in store.objects
diff --git a/mooncake-wheel/tests/transfer_engine_initiator_test.py b/mooncake-wheel/tests/transfer_engine_initiator_test.py
index f9551b75..2c2205dd 100644
--- a/mooncake-wheel/tests/transfer_engine_initiator_test.py
+++ b/mooncake-wheel/tests/transfer_engine_initiator_test.py
@@ -268,6 +268,51 @@ class TestVLLMAdaptorTransfer(unittest.TestCase):
             f"[✓] {circles} rounds of batch_write_async_read passed, batch size {batch_size}."
         )
 
+    def test_async_submit_to_unknown_peer_returns_zero(self):
+        """Async submit APIs return the documented failure sentinel 0.
+
+        openSegment failure used to surface as -1 cast through the unsigned
+        batch_id_t, i.e. 2**64 - 1, which passes the documented
+        `batch_id != 0` failure check and crashes get_batch_transfer_status.
+        """
+        # "unreachable_peer:9" is not registered with the metadata server, so
+        # openSegment fails before any buffer address is dereferenced.
+        batch_id = self.adaptor.batch_transfer_async_write(
+            "unreachable_peer:9", [0], [0], [16]
+        )
+        self.assertEqual(
+            batch_id,
+            0,
+            f"batch_transfer_async_write to unknown peer returned {batch_id}",
+        )
+
+        batch_id = self.adaptor.transfer_submit_write("unreachable_peer:9", 0, 0, 16)
+        self.assertEqual(
+            batch_id,
+            0,
+            f"transfer_submit_write to unknown peer returned {batch_id}",
+        )
+
+    def test_send_probe_reachable_peer(self):
+        """send_probe against a registered, reachable peer returns 0."""
+        rc = self.adaptor.send_probe(self.target_server_name)
+        self.assertEqual(
+            rc,
+            0,
+            f"send_probe to reachable peer {self.target_server_name} returned {rc}",
+        )
+
+    def test_send_probe_unknown_peer(self):
+        """send_probe against an unknown peer returns non-zero (does not crash)."""
+        # "unreachable_peer:9" is not registered with the metadata server, so the
+        # metadata lookup itself should fail and sendProbe returns ERR_METADATA.
+        rc = self.adaptor.send_probe("unreachable_peer:9")
+        self.assertNotEqual(
+            rc,
+            0,
+            "send_probe to unknown peer unexpectedly succeeded",
+        )
+
 
 if __name__ == "__main__":
     unittest.main()
diff --git a/requirements-dev.txt b/requirements-dev.txt
index 9e9dd019..a7e38e18 100644
--- a/requirements-dev.txt
+++ b/requirements-dev.txt
@@ -5,4 +5,4 @@ pre-commit==3.7.1
 ruff==0.6.9
 codespell==2.2.6
 cmake-format==0.6.13
-# clang-format is provided via system package (e.g., apt install clang-format) or toolchain
+# clang-format is provided via system package (e.g., sudo apt-get install -y clang-format-20) or toolchain
diff --git a/requirements_docs.txt b/requirements_docs.txt
index 9c267eda..2aacb268 100644
--- a/requirements_docs.txt
+++ b/requirements_docs.txt
@@ -4,6 +4,7 @@ sphinx-book-theme==1.1.4
 sphinx-copybutton==0.5.2
 sphinx-design==0.6.1
 sphinx-togglebutton==0.3.2
+sphinxcontrib-mermaid==2.0.1
 myst-parser==3.0.1  # `myst-parser==4.0.1` breaks inline code in titles
 msgspec
 snowballstemmer<3  # https://github.com/snowballstem/snowball/issues/229
diff --git a/scripts/ascend/dependencies_ascend.sh b/scripts/ascend/dependencies_ascend.sh
index d58d56f8..101aaaf4 100644
--- a/scripts/ascend/dependencies_ascend.sh
+++ b/scripts/ascend/dependencies_ascend.sh
@@ -60,7 +60,7 @@ if command -v apt-get &> /dev/null; then
             libunwind-dev \
             libnuma-dev \
             libpython3-dev \
-            libboost-all-dev \
+            libboost-dev \
             libssl-dev \
             libzstd-dev \
             libgrpc-dev \
diff --git a/scripts/ascend/dependencies_ascend_installation.sh b/scripts/ascend/dependencies_ascend_installation.sh
index 9fc316ec..eb9266dc 100644
--- a/scripts/ascend/dependencies_ascend_installation.sh
+++ b/scripts/ascend/dependencies_ascend_installation.sh
@@ -18,12 +18,29 @@
 
 #!/bin/bash
 
+# Try git clone with GitHub mirror fallback via https://ghfast.top/
+git_with_github_mirror_fallback() {
+    local repo_dir="$1"
+    local repo_url="$2"
+    shift 2
+
+    if git clone "$repo_url" "$repo_dir" "$@"; then
+        return 0
+    fi
+
+    echo "Direct clone failed, retrying with mirror https://ghfast.top/"
+    rm -rf "$repo_dir"
+    local mirror_url="https://ghfast.top/${repo_url}"
+    git clone "$mirror_url" "$repo_dir" "$@"
+}
+
 clone_repo_if_not_exists() {
-    local repo_dir=$1
-    local repo_url=$2
+    local repo_dir="$1"
+    local repo_url="$2"
+    shift 2
 
     if [ ! -d "$repo_dir" ]; then
-        git clone "$repo_url"
+        git_with_github_mirror_fallback "$repo_dir" "$repo_url" "$@"
     else
         echo "Directory $repo_dir already exists, skipping clone."
     fi
@@ -53,7 +70,7 @@ if command -v apt-get &> /dev/null; then
             libunwind-dev \
             libnuma-dev \
             libpython3-dev \
-            libboost-all-dev \
+            libboost-dev \
             libssl-dev \
             libgrpc-dev \
             libgrpc++-dev \
@@ -101,7 +118,7 @@ elif command -v yum &> /dev/null; then
     cd ../..
 
     # Install msgpack-c
-    clone_repo_if_not_exists "msgpack" "https://github.com/msgpack/msgpack-c.git"
+    clone_repo_if_not_exists "msgpack-c" "https://github.com/msgpack/msgpack-c.git"
     cd msgpack-c || exit
     git checkout cpp-7.0.0
     rm -rf build
diff --git a/scripts/ascend/dependencies_openeuler.sh b/scripts/ascend/dependencies_openeuler.sh
new file mode 100644
index 00000000..0d6b83f9
--- /dev/null
+++ b/scripts/ascend/dependencies_openeuler.sh
@@ -0,0 +1,310 @@
+#!/bin/bash
+# Copyright 2024-2026 KVCache.AI
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+set -euo pipefail
+
+GREEN="\033[0;32m"
+BLUE="\033[0;34m"
+YELLOW="\033[0;33m"
+RED="\033[0;31m"
+NC="\033[0m"
+
+SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
+REPO_ROOT="$(cd "${SCRIPT_DIR}/../.." && pwd)"
+DEPS_BUILD_DIR="${REPO_ROOT}/.mooncake-openeuler-deps"
+GOVER=1.25.9
+ASCEND_GITHUB_MIRROR_URLS="${ASCEND_GITHUB_MIRROR_URLS:-https://ghfast.top/}"
+
+print_section() {
+    echo -e "\n${BLUE}=== $1 ===${NC}"
+}
+
+print_success() {
+    echo -e "${GREEN}✓ $1${NC}"
+}
+
+print_error() {
+    echo -e "${RED}✗ ERROR: $1${NC}"
+    exit 1
+}
+
+print_warn() {
+    echo -e "${YELLOW}! $1${NC}"
+}
+
+check_success() {
+    if [ $? -ne 0 ]; then
+        print_error "$1"
+    fi
+}
+
+SKIP_CONFIRM=false
+for arg in "$@"; do
+    case $arg in
+        -y|--yes)
+            SKIP_CONFIRM=true
+            ;;
+        -h|--help)
+            echo -e "${YELLOW}Mooncake openEuler Dependencies Installer${NC}"
+            echo "Usage: bash scripts/ascend/dependencies_openeuler.sh [-y]"
+            exit 0
+            ;;
+    esac
+done
+
+if [ "$(id -u)" -ne 0 ]; then
+    print_error "Require root permission, try: sudo bash scripts/ascend/dependencies_openeuler.sh -y"
+fi
+
+echo -e "${YELLOW}Mooncake openEuler Dependencies Installer${NC}"
+echo "Repository root: ${REPO_ROOT}"
+echo "This script installs (same scope as root dependencies.sh):"
+echo "  - openEuler/RHEL system packages (dnf/yum)"
+echo "  - Git submodules (pybind11, yalantinglibs, ...)"
+echo "  - yalantinglibs (from extern/ submodule)"
+echo "  - Go ${GOVER} (for USE_ETCD / libetcd_wrapper.so)"
+echo "Run scripts/ascend/dependencies_ascend_installation.sh afterward for Ascend extras."
+echo
+
+if [ "$SKIP_CONFIRM" = false ]; then
+    read -p "Do you want to continue? [Y/n] " -n 1 -r
+    echo
+    if [[ ! $REPLY =~ ^[Yy]$ ]] && [[ ! $REPLY = "" ]]; then
+        echo -e "${YELLOW}Installation cancelled.${NC}"
+        exit 0
+    fi
+fi
+
+detect_pkg_manager() {
+    if command -v dnf >/dev/null 2>&1; then
+        echo dnf
+    elif command -v yum >/dev/null 2>&1; then
+        echo yum
+    else
+        print_error "Neither dnf nor yum found. This script is for openEuler/RHEL only."
+    fi
+}
+
+git_with_github_mirror_fallback() {
+    local repo_dir="$1"
+    local repo_url="$2"
+    shift 2
+
+    if git clone "$repo_url" "$repo_dir" "$@"; then
+        return 0
+    fi
+
+    print_warn "Direct clone failed, retrying with mirror https://ghfast.top/"
+    rm -rf "$repo_dir"
+    git clone "https://ghfast.top/${repo_url}" "$repo_dir" "$@"
+}
+
+clone_repo_if_not_exists() {
+    local repo_dir="$1"
+    local repo_url="$2"
+    shift 2
+
+    if [ -d "$repo_dir" ]; then
+        echo "Directory $repo_dir already exists, skipping clone."
+        return 0
+    fi
+    git_with_github_mirror_fallback "$repo_dir" "$repo_url" "$@"
+}
+
+install_rpm_packages() {
+    local pkg_mgr="$1"
+    shift
+    local packages=("$@")
+    local missing=()
+
+    print_section "Installing system packages via ${pkg_mgr}"
+    ${pkg_mgr} makecache -y || ${pkg_mgr} makecache
+    check_success "Failed to refresh ${pkg_mgr} metadata"
+
+    for pkg in "${packages[@]}"; do
+        if ! ${pkg_mgr} install -y "$pkg"; then
+            missing+=("$pkg")
+        fi
+    done
+
+    if [ ${#missing[@]} -gt 0 ]; then
+        print_warn "Some packages were not installed (may be optional on this openEuler release):"
+        printf '  - %s\n' "${missing[@]}"
+    fi
+}
+
+PKG_MGR="$(detect_pkg_manager)"
+
+# Core build dependencies aligned with root dependencies.sh (Debian names mapped to RHEL/openEuler).
+CORE_PACKAGES=(
+    gcc gcc-c++ make cmake ninja-build git wget unzip
+    gflags-devel glog-devel libibverbs-devel numactl-devel
+    gtest gtest-devel boost-devel openssl-devel hiredis-devel
+    libcurl-devel jsoncpp-devel libunwind-devel python3-devel
+    zstd-devel xxhash-devel pkgconf pkgconf-pkg-config patchelf
+    mpich mpich-devel glibc glibc-common
+)
+
+# Often provided by repos on openEuler; install best-effort.
+OPTIONAL_PACKAGES=(
+    grpc-devel grpc-plugins protobuf-devel protobuf-compiler
+    liburing-devel jemalloc-devel
+)
+
+install_rpm_packages "$PKG_MGR" "${CORE_PACKAGES[@]}"
+install_rpm_packages "$PKG_MGR" "${OPTIONAL_PACKAGES[@]}" || true
+
+# openEuler images may ship MPICH; remove OpenMPI if present to avoid conflicts.
+${PKG_MGR} remove -y openmpi openmpi-devel 2>/dev/null || true
+
+print_success "System package installation finished"
+
+export CPLUS_INCLUDE_PATH="$(
+    echo "${CPLUS_INCLUDE_PATH:-}" | tr ':' '\n' | grep -v "/usr/local/Ascend" | paste -sd: -
+)"
+
+print_section "Initializing Git submodules"
+cd "${REPO_ROOT}"
+if [ ! -f "${REPO_ROOT}/.gitmodules" ]; then
+    print_error "No .gitmodules found under ${REPO_ROOT}"
+fi
+
+submodule_updated=false
+if git submodule sync --recursive && git submodule update --init --recursive; then
+    submodule_updated=true
+elif [ -n "${ASCEND_GITHUB_MIRROR_URLS:-}" ]; then
+  while IFS= read -r raw; do
+    base="${raw#"${raw%%[![:space:]]*}"}"
+    base="${base%"${base##*[![:space:]]}"}"
+    [ -n "$base" ] || continue
+    [ "$base" = "https://github.com/" ] && continue
+    [ "$base" != "https://github.com/" ] && base="${base%/}/"
+    echo "Retrying submodule update with mirror ${base}"
+    if git -c url."${base}https://github.com/".insteadOf=https://github.com/ \
+        submodule sync --recursive && \
+       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
+    print_error "git submodule update failed (direct GitHub and mirrors exhausted)"
+fi
+print_success "Git submodules initialized"
+
+print_section "Installing yalantinglibs from extern/ submodule"
+if [ -d "${REPO_ROOT}/extern/yalantinglibs" ]; then
+    cd "${REPO_ROOT}/extern/yalantinglibs"
+    rm -rf build
+    mkdir -p build && cd build
+    cmake .. -DBUILD_EXAMPLES=OFF -DBUILD_BENCHMARK=OFF -DBUILD_UNIT_TESTS=OFF
+    check_success "Failed to configure yalantinglibs"
+    cmake --build . -j"$(nproc)"
+    check_success "Failed to build yalantinglibs"
+    cmake --install .
+    check_success "Failed to install yalantinglibs"
+    print_success "yalantinglibs installed from submodule"
+    cd "${REPO_ROOT}"
+else
+    print_warn "extern/yalantinglibs missing, building from upstream clone"
+    mkdir -p "${DEPS_BUILD_DIR}"
+    cd "${DEPS_BUILD_DIR}"
+    clone_repo_if_not_exists "yalantinglibs" "https://github.com/alibaba/yalantinglibs.git"
+    cd yalantinglibs
+    git checkout 0.5.5
+    rm -rf build && mkdir -p build && cd build
+    cmake .. -DBUILD_EXAMPLES=OFF -DBUILD_BENCHMARK=OFF -DBUILD_UNIT_TESTS=OFF
+    cmake --build . -j"$(nproc)"
+    cmake --install .
+    cd "${REPO_ROOT}"
+    print_success "yalantinglibs installed from git clone"
+fi
+
+print_section "Verifying essential build tools"
+for tool in getconf ldd patchelf cmake git; do
+    command -v "$tool" >/dev/null || print_error "${tool} not found after dependency installation"
+    print_success "${tool} found: $(command -v "$tool")"
+done
+
+print_section "Installing Go ${GOVER}"
+USED_CN_MIRROR=false
+
+install_go() {
+    local arch
+    arch="$(uname -m)"
+    case "$arch" in
+        aarch64) arch=arm64 ;;
+        x86_64) arch=amd64 ;;
+        *) print_error "Unsupported architecture: $arch" ;;
+    esac
+
+    local tarball="go${GOVER}.linux-${arch}.tar.gz"
+    local urls=(
+        "https://go.dev/dl/${tarball}"
+        "https://golang.google.cn/dl/${tarball}"
+        "https://mirrors.aliyun.com/golang/${tarball}"
+    )
+
+    local url downloaded=false
+    for url in "${urls[@]}"; do
+        echo "Downloading Go ${GOVER} from ${url}..."
+        if wget -q --show-progress --timeout=30 --tries=2 -O "${tarball}" "${url}"; then
+            downloaded=true
+            [[ "$url" != "https://go.dev/dl/${tarball}" ]] && USED_CN_MIRROR=true
+            break
+        fi
+        rm -f "${tarball}"
+    done
+
+    [ "$downloaded" = true ] || print_error "Failed to download Go ${GOVER}"
+
+    rm -rf /usr/local/go
+    tar -C /usr/local -xzf "${tarball}"
+    rm -f "${tarball}"
+    export PATH="/usr/local/go/bin:${PATH}"
+    print_success "Go ${GOVER} installed to /usr/local/go"
+}
+
+if command -v go >/dev/null 2>&1; then
+    current_go="$(go version | awk '{print $3}')"
+    if [[ "$current_go" == "go${GOVER}" ]]; then
+        print_success "Go ${GOVER} already installed (${current_go})"
+    else
+        print_warn "Found ${current_go}, installing Go ${GOVER}"
+        install_go
+    fi
+else
+    install_go
+fi
+
+export PATH="/usr/local/go/bin:${PATH}"
+if [ "$USED_CN_MIRROR" = true ] && [ -z "${GOPROXY:-}" ]; then
+    export GOPROXY=https://goproxy.cn,https://goproxy.io,direct
+    print_warn "GOPROXY set to ${GOPROXY} (restricted network detected)"
+fi
+
+go version
+print_success "Go is available: $(command -v go)"
+
+print_section "Installation complete"
+echo -e "${GREEN}openEuler dependencies are ready.${NC}"
+echo "Next steps (example NPU release build):"
+echo "  source /usr/local/Ascend/cann-9.0.0/set_env.sh"
+echo "  mkdir -p build && cd build"
+echo "  cmake .. -DUSE_ASCEND_DIRECT=ON -DUSE_CUDA=OFF -DWITH_EP=OFF -DUSE_HTTP=ON -DUSE_ETCD=ON -DSTORE_USE_ETCD=ON"
+echo "  cmake --build . -j\$(nproc) && cmake --install ."
\ No newline at end of file
diff --git a/scripts/bench_engram_store_27b.py b/scripts/bench_engram_store_27b.py
new file mode 100644
index 00000000..7a961ca3
--- /dev/null
+++ b/scripts/bench_engram_store_27b.py
@@ -0,0 +1,281 @@
+#!/usr/bin/env python3
+"""
+EngramStore backend benchmark.
+
+Scenario:
+  - 1 token per request (L=1)
+  - Batch sizes: 1, 4, 16, 64, 128, 256
+  - Benchmarks Mooncake's row-id lookup backend only
+
+Usage:
+  # Start master first:
+  # mooncake_master --default_kv_lease_ttl=500 --enable_http_metadata_server=true
+
+  export MOONCAKE_MASTER=127.0.0.1:50051
+  export MOONCAKE_TE_META_DATA_SERVER=http://127.0.0.1:8080/metadata
+  python scripts/bench_engram_store_27b.py
+
+Optional:
+  export ENGRAM_STORE_ALLOW_POPULATE_FALLBACK=1
+"""
+
+import ctypes
+import os
+import sys
+import time
+import uuid
+
+import numpy as np
+
+repo_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+build_dir = os.environ.get("MOONCAKE_BUILD_DIR", "build")
+build_store = os.path.join(repo_root, build_dir, "mooncake-integration")
+wheel_dir = os.path.join(repo_root, "mooncake-wheel")
+for d in [build_store, wheel_dir]:
+    if os.path.isdir(d) and d not in sys.path:
+        sys.path.insert(0, d)
+
+try:
+    import store
+    from mooncake.mooncake_config import MooncakeConfig
+except ImportError as e:
+    print(f"❌ Import failed: {e}")
+    sys.exit(1)
+
+BATCH_SIZES = [1, 4, 16, 64, 128, 256]
+NUM_WARMUP = 10
+NUM_ITER = 50
+TABLE_VOCAB_SIZES = [50000] * 8
+EMBEDDING_DIM = 16
+ALLOW_POPULATE_FALLBACK = (
+    os.environ.get("ENGRAM_STORE_ALLOW_POPULATE_FALLBACK", "0") == "1"
+)
+
+
+def create_engram_store_config():
+    cfg = store.EngramStoreConfig()
+    cfg.table_vocab_sizes = TABLE_VOCAB_SIZES
+    cfg.embedding_dim = EMBEDDING_DIM
+    return cfg
+
+
+def rollback_keys(store_obj, keys):
+    for key in keys:
+        rc = store_obj.remove(key, True)
+        if rc not in (0, -704):
+            print(f"Warning: failed to remove benchmark key {key}, rc={rc}")
+
+
+def populate_store_via_cxl_segment(store_obj, keys, embedding_buffers):
+    cxl_path = os.environ.get("MC_CXL_DEV_PATH")
+    if not cxl_path:
+        raise RuntimeError("MC_CXL_DEV_PATH is required for CXL fallback")
+
+    real_cxl_path = os.path.realpath(cxl_path)
+    base_addr = None
+    limit_addr = None
+    with open("/proc/self/maps", "r", encoding="utf-8") as fin:
+        for line in fin:
+            if real_cxl_path not in line:
+                continue
+            addr_range = line.split()[0]
+            start_hex, end_hex = addr_range.split("-")
+            base_addr = int(start_hex, 16)
+            limit_addr = int(end_hex, 16)
+            break
+
+    if base_addr is None or limit_addr is None:
+        raise RuntimeError(f"failed to find CXL mapping for {real_cxl_path}")
+
+    page_size = 4096
+    aligned_sizes = []
+    total_aligned_bytes = 0
+    for emb in embedding_buffers:
+        nbytes = emb.nbytes
+        aligned = (nbytes + page_size - 1) // page_size * page_size
+        aligned_sizes.append(aligned)
+        total_aligned_bytes += aligned
+
+    cursor = limit_addr - total_aligned_bytes
+    if cursor < base_addr:
+        raise RuntimeError("CXL mapping is too small for staging region")
+
+    published_keys = []
+    try:
+        for head_idx, emb in enumerate(embedding_buffers):
+            nbytes = emb.nbytes
+            aligned = aligned_sizes[head_idx]
+            if cursor + aligned > limit_addr:
+                raise RuntimeError("CXL mapping is too small for benchmark tables")
+            ctypes.memmove(cursor, emb.ctypes.data, nbytes)
+            key = keys[head_idx]
+            rc = store_obj.put_from(key, cursor, nbytes)
+            if rc != 0:
+                raise RuntimeError(
+                    f"CXL put_from fallback failed for {key}, rc={rc}"
+                )
+            published_keys.append(key)
+            cursor += aligned
+    except Exception:
+        rollback_keys(store_obj, published_keys)
+        raise
+
+
+def populate_store(engram_store, store_obj):
+    keys = engram_store.get_store_keys()
+    embedding_buffers = []
+    for vocab_size in engram_store.get_table_vocab_sizes():
+        emb = np.random.randn(vocab_size, engram_store.get_embedding_dim()).astype(np.float32)
+        embedding_buffers.append(emb)
+
+    t0 = time.perf_counter()
+    mode = "engram_store.populate"
+    try:
+        engram_store.populate(embedding_buffers)
+    except RuntimeError:
+        if not ALLOW_POPULATE_FALLBACK:
+            raise
+        protocol = os.environ.get("MOONCAKE_PROTOCOL", "")
+        if protocol == "cxl":
+            mode = "cxl put_from fallback"
+            populate_store_via_cxl_segment(store_obj, keys, embedding_buffers)
+        else:
+            mode = "store.put fallback"
+            published_keys = []
+            try:
+                for head_idx, emb in enumerate(embedding_buffers):
+                    key = keys[head_idx]
+                    rc = store_obj.put(key, emb)
+                    if rc != 0:
+                        raise RuntimeError(
+                            f"fallback populate failed for {key}, rc={rc}"
+                        )
+                    published_keys.append(key)
+            except Exception:
+                rollback_keys(store_obj, published_keys)
+                raise
+    populate_ms = (time.perf_counter() - t0) * 1000
+    return embedding_buffers, populate_ms, mode
+
+
+def make_row_ids(engram_store, batch_size, seq_len):
+    row_ids = []
+    for b in range(batch_size):
+        batch = []
+        for l in range(seq_len):
+            token_rows = []
+            for head, vocab_size in enumerate(engram_store.get_table_vocab_sizes()):
+                token_rows.append((b * seq_len + l + head) % vocab_size)
+            batch.append(token_rows)
+        row_ids.append(batch)
+    return row_ids
+
+
+def run_benchmark(engram_store, batch_size, num_warmup, num_iter):
+    seq_len = 1
+    row_ids = make_row_ids(engram_store, batch_size, seq_len)
+
+    for _ in range(num_warmup):
+        engram_store.lookup(row_ids)
+
+    total_ms_list = []
+    for _ in range(num_iter):
+        t0 = time.perf_counter()
+        engram_store.lookup(row_ids)
+        t1 = time.perf_counter()
+        total_ms_list.append((t1 - t0) * 1000)
+
+    total_ms = np.array(total_ms_list)
+    mean_total = np.mean(total_ms)
+    p50_total = np.percentile(total_ms, 50)
+    p99_total = np.percentile(total_ms, 99)
+
+    tokens_per_sec = batch_size / (mean_total / 1000)
+    bytes_per_request = (
+        batch_size
+        * seq_len
+        * engram_store.get_num_heads()
+        * engram_store.get_embedding_dim()
+        * 4
+    )
+    gbps = (bytes_per_request / 1e9) / (mean_total / 1000)
+
+    return {
+        "batch_size": batch_size,
+        "mean_total_ms": mean_total,
+        "p50_ms": p50_total,
+        "p99_ms": p99_total,
+        "tokens_per_sec": tokens_per_sec,
+        "gbps": gbps,
+    }
+
+
+def validate_lookup_correctness(engram_store, embedding_buffers):
+    row_ids = make_row_ids(engram_store, batch_size=1, seq_len=1)
+    output = np.asarray(engram_store.lookup(row_ids))
+
+    max_abs_err = 0.0
+    for head_idx in range(output.shape[2]):
+        row_idx = row_ids[0][0][head_idx]
+        expected = embedding_buffers[head_idx][row_idx]
+        actual = output[0, 0, head_idx]
+        max_abs_err = max(max_abs_err, float(np.max(np.abs(actual - expected))))
+    return max_abs_err
+
+
+def main():
+    print("=" * 60)
+    print("EngramStore Backend Benchmark")
+    print("  Config: 1 token, Mooncake row-id lookup only")
+    print("  Batch sizes: 1, 4, 16, 64, 128, 256")
+    print(f"  Build dir: {build_dir}")
+    print(f"  Protocol: {os.environ.get('MOONCAKE_PROTOCOL', '<unset>')}")
+    print(f"  Allow populate fallback: {ALLOW_POPULATE_FALLBACK}")
+    print("=" * 60)
+
+    config = MooncakeConfig.load_from_env()
+    store_obj = store.MooncakeDistributedStore()
+    rc = store_obj.setup(
+        config.local_hostname,
+        config.metadata_server,
+        config.global_segment_size,
+        config.local_buffer_size,
+        config.protocol,
+        config.device_name,
+        config.master_server_address,
+    )
+    if rc != 0:
+        raise RuntimeError(f"Failed to setup Mooncake store, rc={rc}")
+
+    engram_store = None
+    try:
+        cfg = create_engram_store_config()
+        layer_id = uuid.uuid4().int & 0x7FFFFFFF
+        engram_store = store.EngramStore(layer_id=layer_id, config=cfg, store=store_obj)
+        embedding_buffers, populate_ms, mode = populate_store(engram_store, store_obj)
+        print(f"Populate mode: {mode}, took {populate_ms:.2f} ms")
+
+        max_abs_err = validate_lookup_correctness(engram_store, embedding_buffers)
+        print(f"Max abs error: {max_abs_err:.6f}")
+
+        print("\nResults:")
+        print(
+            f"{'Batch':>8} {'Mean(ms)':>10} {'P50(ms)':>10} {'P99(ms)':>10} {'Tok/s':>12} {'GB/s':>10}"
+        )
+        for batch_size in BATCH_SIZES:
+            result = run_benchmark(engram_store, batch_size, NUM_WARMUP, NUM_ITER)
+            print(
+                f"{result['batch_size']:>8} {result['mean_total_ms']:>10.3f} {result['p50_ms']:>10.3f} "
+                f"{result['p99_ms']:>10.3f} {result['tokens_per_sec']:>12.1f} {result['gbps']:>10.3f}"
+            )
+    finally:
+        if engram_store is not None:
+            try:
+                engram_store.remove_from_store(force=True)
+            except Exception as exc:
+                print(f"Warning: failed to clean up benchmark layer: {exc}")
+        store_obj.close()
+
+
+if __name__ == "__main__":
+    main()
diff --git a/scripts/build_wheel.sh b/scripts/build_wheel.sh
index 24125b8a..e6bc964f 100755
--- a/scripts/build_wheel.sh
+++ b/scripts/build_wheel.sh
@@ -13,10 +13,11 @@ OUTPUT_DIR=${OUTPUT_DIR:-${2:-"dist"}}
 # CMake build directory (default: build).  EP/PG extensions are staged under
 # ${BUILD_DIR}/ep_pg_staging when the project was built with -DWITH_EP=ON.
 BUILD_DIR="${BUILD_DIR:-build}"
+BUILD_DIR_ABS="$(pwd)/${BUILD_DIR}"
 echo "Building wheel for Python ${PYTHON_VERSION} with output directory ${OUTPUT_DIR}"
 
 # Ensure LD_LIBRARY_PATH includes /usr/local/lib
-export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/$(pwd)/build/mooncake-common:/usr/local/lib
+export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:${BUILD_DIR_ABS}/mooncake-common:${BUILD_DIR_ABS}/mooncake-common/etcd:${BUILD_DIR_ABS}/mooncake-common/k8s-lease:/usr/local/lib
 
 echo "Cleaning wheel-build directory"
 rm -rf mooncake-wheel/mooncake_transfer_engine*
@@ -25,21 +26,24 @@ rm -f mooncake-wheel/mooncake/*.so
 
 echo "Creating directory structure..."
 
+# Copy shared allocator helper used by both CUDA and Ascend pluggable allocators.
+cp mooncake-integration/fabric_allocator_utils.py mooncake-wheel/mooncake/fabric_allocator_utils.py
+
 # Copy engine.so to mooncake directory (will be imported by transfer module)
-cp build/mooncake-integration/engine.*.so mooncake-wheel/mooncake/engine.so
+cp ${BUILD_DIR}/mooncake-integration/engine.*.so mooncake-wheel/mooncake/engine.so
 
 # Copy libasio.so to mooncake directory (runtime dependency of engine.so)
-cp build/mooncake-common/libasio.so mooncake-wheel/mooncake/libasio.so
+cp ${BUILD_DIR}/mooncake-common/libasio.so mooncake-wheel/mooncake/libasio.so
 
 # Copy store.so to mooncake directory
-if [ -f build/mooncake-integration/store.*.so ]; then
+if compgen -G "${BUILD_DIR}/mooncake-integration/store.*.so" >/dev/null; then
     echo "Copying store.so..."
-    cp build/mooncake-integration/store.*.so mooncake-wheel/mooncake/store.so
+    cp ${BUILD_DIR}/mooncake-integration/store.*.so mooncake-wheel/mooncake/store.so
     echo "Copying master binary..."
     # Copy master binary
-    cp build/mooncake-store/src/mooncake_master mooncake-wheel/mooncake/
+    cp ${BUILD_DIR}/mooncake-store/src/mooncake_master mooncake-wheel/mooncake/
     # Copy client binary
-    cp build/mooncake-store/src/mooncake_client mooncake-wheel/mooncake/
+    cp ${BUILD_DIR}/mooncake-store/src/mooncake_client mooncake-wheel/mooncake/
     # Copy async_store.py
     cp mooncake-integration/store/async_store.py mooncake-wheel/mooncake/async_store.py
 else
@@ -47,36 +51,48 @@ else
 fi
 
 # Copy libmooncake_store.so to mooncake directory (only when BUILD_SHARED_LIBS is set)
-if [ -f build/mooncake-store/src/libmooncake_store.so ]; then
+if [ -f ${BUILD_DIR}/mooncake-store/src/libmooncake_store.so ]; then
     echo "Copying libmooncake_store.so..."
-    cp build/mooncake-store/src/libmooncake_store.so mooncake-wheel/mooncake/libmooncake_store.so
+    cp ${BUILD_DIR}/mooncake-store/src/libmooncake_store.so mooncake-wheel/mooncake/libmooncake_store.so
+fi
+
+# Copy libmooncake_common.so to mooncake directory (only when BUILD_SHARED_LIBS is set)
+if [ -f ${BUILD_DIR}/mooncake-common/src/libmooncake_common.so ]; then
+    echo "Copying libmooncake_common.so..."
+    cp ${BUILD_DIR}/mooncake-common/src/libmooncake_common.so mooncake-wheel/mooncake/libmooncake_common.so
 fi
 
 # Copy libtransfer_engine.so to mooncake directory (only when USE_ETCD is set)
-if [ -f build/mooncake-common/etcd/libetcd_wrapper.so ]; then
+if [ -f ${BUILD_DIR}/mooncake-common/etcd/libetcd_wrapper.so ]; then
     echo "Copying libetcd_wrapper.so..."
-    cp build/mooncake-common/etcd/libetcd_wrapper.so mooncake-wheel/mooncake/libetcd_wrapper.so
+    cp ${BUILD_DIR}/mooncake-common/etcd/libetcd_wrapper.so mooncake-wheel/mooncake/libetcd_wrapper.so
+fi
+
+# Copy libk8s_lease_wrapper.so to mooncake directory (only when STORE_USE_K8S_LEASE is set)
+if [ -f ${BUILD_DIR}/mooncake-common/k8s-lease/libk8s_lease_wrapper.so ]; then
+    echo "Copying libk8s_lease_wrapper.so..."
+    cp ${BUILD_DIR}/mooncake-common/k8s-lease/libk8s_lease_wrapper.so mooncake-wheel/mooncake/libk8s_lease_wrapper.so
 fi
 
 # Copy libtransfer_engine.so to mooncake directory (only when BUILD_SHARED_LIBS is set)
-if [ -f build/mooncake-transfer-engine/src/libtransfer_engine.so ]; then
+if [ -f ${BUILD_DIR}/mooncake-transfer-engine/src/libtransfer_engine.so ]; then
     echo "Copying libtransfer_engine.so..."
-    cp build/mooncake-transfer-engine/src/libtransfer_engine.so mooncake-wheel/mooncake/libtransfer_engine.so
+    cp ${BUILD_DIR}/mooncake-transfer-engine/src/libtransfer_engine.so mooncake-wheel/mooncake/libtransfer_engine.so
 fi
 
 # Copy ascend_transport.so to mooncake directory (only when USE_ASCEND_DIRECT is set)
-if [ -f build/mooncake-transfer-engine/src/transport/ascend_transport/ascend_transport.so ]; then
+if [ -f ${BUILD_DIR}/mooncake-transfer-engine/src/transport/ascend_transport/ascend_transport.so ]; then
     echo "Copying ascend_transport.so..."
-    cp build/mooncake-transfer-engine/src/transport/ascend_transport/ascend_transport.so mooncake-wheel/mooncake/ascend_transport.so
+    cp ${BUILD_DIR}/mooncake-transfer-engine/src/transport/ascend_transport/ascend_transport.so mooncake-wheel/mooncake/ascend_transport.so
 fi
 
 # Copy nvlink-allocator.so to mooncake directory (only if it exists - CUDA builds only)
-if [ -f build/mooncake-transfer-engine/nvlink-allocator/nvlink_allocator.so ] \
+if [ -f ${BUILD_DIR}/mooncake-transfer-engine/nvlink-allocator/nvlink_allocator.so ] \
    || [ -f /usr/lib/libaccl_barex.so ] \
    || [ -f /usr/lib64/libaccl_barex.so ]; then
-    if [ -f build/mooncake-transfer-engine/nvlink-allocator/nvlink_allocator.so ]; then
+    if [ -f ${BUILD_DIR}/mooncake-transfer-engine/nvlink-allocator/nvlink_allocator.so ]; then
      echo "Copying CUDA nvlink_allocator.so..."
-     cp build/mooncake-transfer-engine/nvlink-allocator/nvlink_allocator.so mooncake-wheel/mooncake/nvlink_allocator.so
+     cp ${BUILD_DIR}/mooncake-transfer-engine/nvlink-allocator/nvlink_allocator.so mooncake-wheel/mooncake/nvlink_allocator.so
     fi
     echo "Copying allocator libraries..."
     # Copy allocator.py
@@ -86,9 +102,9 @@ else
 fi
 
 # Copy ubshmem_fabric_allocator.so to mooncake directory (only if it exists - NPU builds only)
-if [ -f build/mooncake-transfer-engine/ubshmem-allocator/ubshmem_fabric_allocator.so ]; then
+if [ -f ${BUILD_DIR}/mooncake-transfer-engine/ubshmem-allocator/ubshmem_fabric_allocator.so ]; then
     echo "Copying NPU ubshmem_fabric_allocator.so..."
-    cp build/mooncake-transfer-engine/ubshmem-allocator/ubshmem_fabric_allocator.so mooncake-wheel/mooncake/ubshmem_fabric_allocator.so
+    cp ${BUILD_DIR}/mooncake-transfer-engine/ubshmem-allocator/ubshmem_fabric_allocator.so mooncake-wheel/mooncake/ubshmem_fabric_allocator.so
     echo "Copying NPU allocator libraries..."
     # Copy allocator_ascend_npu.py
     cp mooncake-integration/allocator_ascend_npu.py mooncake-wheel/mooncake/allocator_ascend_npu.py
@@ -98,10 +114,10 @@ fi
 
 echo "Copying transfer_engine_bench..."
 # Copy transfer_engine_bench
-cp build/mooncake-transfer-engine/example/transfer_engine_bench mooncake-wheel/mooncake/
+cp ${BUILD_DIR}/mooncake-transfer-engine/example/transfer_engine_bench mooncake-wheel/mooncake/
 
-if [ -f "build/mooncake-transfer-engine/src/transport/ascend_transport/hccl_transport/ascend_transport_c/libascend_transport_mem.so" ]; then
-    cp build/mooncake-transfer-engine/src/transport/ascend_transport/hccl_transport/ascend_transport_c/libascend_transport_mem.so mooncake-wheel/mooncake/
+if [ -f "${BUILD_DIR}/mooncake-transfer-engine/src/transport/ascend_transport/hccl_transport/ascend_transport_c/libascend_transport_mem.so" ]; then
+    cp ${BUILD_DIR}/mooncake-transfer-engine/src/transport/ascend_transport/hccl_transport/ascend_transport_c/libascend_transport_mem.so mooncake-wheel/mooncake/
     echo "Copying ascend_transport_mem libraries..."
 else
     echo "Skipping libascend_transport_mem.so (not built - Ascend disabled)"
@@ -113,7 +129,7 @@ fi
 # so that patchelf never touches CUDA fatbins (see injection step below).
 # Use an absolute path: the script later `cd`s into mooncake-wheel/ and a
 # relative path would silently point to the wrong location.
-CUDA_EP_STAGING_DIR="$(pwd)/${BUILD_DIR}/ep_pg_staging"
+CUDA_EP_STAGING_DIR="${BUILD_DIR_ABS}/ep_pg_staging"
 
 # CI only: remove build/ to free disk before python -m build (set FREE_BUILD_DIR=1 to enable locally).
 # If EP/PG .so files were staged inside the build directory, preserve them in a
@@ -133,11 +149,39 @@ if [ "$CI" = "true" ] || [ "$FREE_BUILD_DIR" = "1" ]; then
     fi
 fi
 
+if [ "$NPU_BUILD" = "1" ]; then
+    echo "Stripping shared libraries to reduce wheel size..."
+    find mooncake-wheel/mooncake -name "*.so" -exec strip --strip-unneeded {} \;
+    echo "Pre-setting RPATH=\$ORIGIN on project ELF files so auditwheel resolves NEEDED from wheel-internal copies..."
+    find mooncake-wheel/mooncake -type f -exec sh -c 'file "$1" | grep -q ELF' _ {} \; -print0 | xargs -0 patchelf --force-rpath --set-rpath '$ORIGIN'
+fi
+
 echo "Building wheel package..."
 # Build the wheel package
 cd mooncake-wheel
 
-# Handle package name modification for non-CUDA builds
+BUILD_VARIANTS="NON_CUDA_BUILD CU13_BUILD NPU_BUILD EFA_BUILD EFA_NON_CUDA_BUILD MUSA_BUILD"
+BUILD_VARIANT_COUNT=0
+for build_variant in $BUILD_VARIANTS; do
+    if [ "${!build_variant}" = "1" ]; then
+        BUILD_VARIANT_COUNT=$((BUILD_VARIANT_COUNT + 1))
+    fi
+done
+if [ "$BUILD_VARIANT_COUNT" -gt 1 ]; then
+    echo "Error: only one of $BUILD_VARIANTS can be set"
+    exit 1
+fi
+
+# If a previous run was interrupted before the trailing restore (line ~481),
+# pyproject.toml is left in a renamed state and pyproject.toml.backup holds the
+# pristine original. Restore it first so the variant rename below always starts
+# from the clean file and the backup is never overwritten with modified content.
+if [ -f pyproject.toml.backup ]; then
+    echo "Restoring pyproject.toml from leftover backup of a previous run"
+    mv pyproject.toml.backup pyproject.toml
+fi
+
+# Handle package name modification for release build variants
 if [ "$NON_CUDA_BUILD" = "1" ]; then
     echo "Modifying package name for non-CUDA build"
     # Backup original pyproject.toml
@@ -147,12 +191,7 @@ if [ "$NON_CUDA_BUILD" = "1" ]; then
     sed -i 's/description = "Python binding of a Mooncake library using pybind11"/description = "Python binding of a Mooncake library using pybind11 (Non-CUDA version)"/' pyproject.toml
     sed -i 's/keywords = \["mooncake", "data transfer", "kv cache", "llm inference"\]/keywords = ["mooncake", "data transfer", "kv cache", "llm inference", "non-cuda"]/' pyproject.toml
     echo "Package name modified to: mooncake-transfer-engine-non-cuda"
-else
-    echo "Using standard package name: mooncake-transfer-engine"
-fi
-
-# Handle package name modification for CU13 builds
-if [ "$CU13_BUILD" = "1" ]; then
+elif [ "$CU13_BUILD" = "1" ]; then
     echo "Modifying package name for CU13 build"
     # Backup original pyproject.toml
     cp pyproject.toml pyproject.toml.backup
@@ -161,6 +200,42 @@ if [ "$CU13_BUILD" = "1" ]; then
     sed -i 's/description = "Python binding of a Mooncake library using pybind11"/description = "Python binding of a Mooncake library using pybind11 (CUDA 13 version)"/' pyproject.toml
     sed -i 's/keywords = \["mooncake", "data transfer", "kv cache", "llm inference"\]/keywords = ["mooncake", "data transfer", "kv cache", "llm inference", "cuda13"]/' pyproject.toml
     echo "Package name modified to: mooncake-transfer-engine-cuda13"
+elif [ "$NPU_BUILD" = "1" ]; then
+    echo "Modifying package name for Ascend NPU build"
+    # Backup original pyproject.toml
+    cp pyproject.toml pyproject.toml.backup
+    # Replace package name and description
+    sed -i 's/name = "mooncake-transfer-engine"/name = "mooncake-transfer-engine-npu"/' pyproject.toml
+    sed -i 's/description = "Python binding of a Mooncake library using pybind11"/description = "Python binding of a Mooncake library using pybind11 (Ascend NPU version)"/' pyproject.toml
+    sed -i 's/keywords = \["mooncake", "data transfer", "kv cache", "llm inference"\]/keywords = ["mooncake", "data transfer", "kv cache", "llm inference", "ascend", "npu"]/' pyproject.toml
+    echo "Package name modified to: mooncake-transfer-engine-npu"
+elif [ "$EFA_BUILD" = "1" ]; then
+    echo "Modifying package name for AWS EFA build (CUDA)"
+    # Backup original pyproject.toml
+    cp pyproject.toml pyproject.toml.backup
+    # Replace package name and description
+    sed -i 's/name = "mooncake-transfer-engine"/name = "mooncake-transfer-engine-efa"/' pyproject.toml
+    sed -i 's/description = "Python binding of a Mooncake library using pybind11"/description = "Python binding of a Mooncake library using pybind11 (AWS EFA, CUDA version)"/' pyproject.toml
+    sed -i 's/keywords = \["mooncake", "data transfer", "kv cache", "llm inference"\]/keywords = ["mooncake", "data transfer", "kv cache", "llm inference", "aws", "efa", "libfabric", "cuda"]/' pyproject.toml
+    echo "Package name modified to: mooncake-transfer-engine-efa"
+elif [ "$EFA_NON_CUDA_BUILD" = "1" ]; then
+    echo "Modifying package name for AWS EFA build (non-CUDA)"
+    # Backup original pyproject.toml
+    cp pyproject.toml pyproject.toml.backup
+    # Replace package name and description
+    sed -i 's/name = "mooncake-transfer-engine"/name = "mooncake-transfer-engine-efa-non-cuda"/' pyproject.toml
+    sed -i 's/description = "Python binding of a Mooncake library using pybind11"/description = "Python binding of a Mooncake library using pybind11 (AWS EFA, Non-CUDA version)"/' pyproject.toml
+    sed -i 's/keywords = \["mooncake", "data transfer", "kv cache", "llm inference"\]/keywords = ["mooncake", "data transfer", "kv cache", "llm inference", "aws", "efa", "libfabric", "non-cuda"]/' pyproject.toml
+    echo "Package name modified to: mooncake-transfer-engine-efa-non-cuda"
+elif [ "$MUSA_BUILD" = "1" ]; then
+    echo "Modifying package name for MUSA build"
+    # Backup original pyproject.toml
+    cp pyproject.toml pyproject.toml.backup
+    # Replace package name and description
+    sed -i 's/name = "mooncake-transfer-engine"/name = "mooncake-transfer-engine-musa"/' pyproject.toml
+    sed -i 's/description = "Python binding of a Mooncake library using pybind11"/description = "Python binding of a Mooncake library using pybind11 (MUSA version)"/' pyproject.toml
+    sed -i 's/keywords = \["mooncake", "data transfer", "kv cache", "llm inference"\]/keywords = ["mooncake", "data transfer", "kv cache", "llm inference", "musa", "moore-threads"]/' pyproject.toml
+    echo "Package name modified to: mooncake-transfer-engine-musa"
 else
     echo "Using standard package name: mooncake-transfer-engine"
 fi
@@ -170,7 +245,27 @@ rm -rf ${OUTPUT_DIR}/
 mkdir -p ${OUTPUT_DIR}
 
 echo "Installing required build packages"
-if command -v pip &>/dev/null; then
+if [ "$NPU_BUILD" = "1" ]; then
+    PYTHON_CMD="python${PYTHON_VERSION}"
+    if ! command -v "$PYTHON_CMD" &>/dev/null; then
+        echo "Error: $PYTHON_CMD not found for NPU wheel build"
+        exit 1
+    fi
+    max_attempts=3
+    attempt=1
+    while [ $attempt -le $max_attempts ]; do
+        if "$PYTHON_CMD" -m pip install --upgrade pip build setuptools wheel auditwheel; then
+            break
+        fi
+        echo "pip install attempt $attempt/$max_attempts failed, retrying in 5s..."
+        sleep 5
+        attempt=$((attempt + 1))
+    done
+    if [ $attempt -gt $max_attempts ]; then
+        echo "Error: pip install failed after $max_attempts attempts"
+        exit 1
+    fi
+elif command -v pip &>/dev/null; then
     python${PYTHON_VERSION} -m pip install --upgrade pip build setuptools wheel auditwheel
 elif command -v uv &>/dev/null; then
     uv pip install --upgrade pip
@@ -239,9 +334,18 @@ echo "Detected glibc version: $GLIBC_VERSION"
 echo "Using platform tag: $PLATFORM_TAG"
 
 echo "Repairing wheel with auditwheel for platform: $PLATFORM_TAG"
-python${PYTHON_VERSION} -m build --wheel --outdir ${OUTPUT_DIR}
-auditwheel repair ${OUTPUT_DIR}/*.whl \
+if [ "$NPU_BUILD" = "1" ]; then
+    python${PYTHON_VERSION} -m build --wheel --no-isolation --outdir ${OUTPUT_DIR}
+    AUDITWHEEL_CMD="python${PYTHON_VERSION} -m auditwheel"
+else
+    python${PYTHON_VERSION} -m build --wheel --outdir ${OUTPUT_DIR}
+    AUDITWHEEL_CMD="auditwheel"
+fi
+
+${AUDITWHEEL_CMD} repair ${OUTPUT_DIR}/*.whl \
     --exclude libcurl.so* \
+    --exclude libfabric.so* \
+    --exclude libefa.so* \
     --exclude libibverbs.so* \
     --exclude libmlx5.so* \
     --exclude libnuma.so* \
@@ -282,6 +386,8 @@ auditwheel repair ${OUTPUT_DIR}/*.whl \
     --exclude libffi.so* \
     --exclude libcuda.so* \
     --exclude libcudart.so* \
+    --exclude libmusa.so* \
+    --exclude libmusart.so* \
     --exclude libamdhip64.so* \
     --exclude libhsa-runtime64.so* \
     --exclude librocprofiler-register.so* \
@@ -325,6 +431,7 @@ auditwheel repair ${OUTPUT_DIR}/*.whl \
     --exclude libllm_datadist*.so \
     --exclude ascend_transport*.so \
     --exclude libaccl_barex.so* \
+    --exclude liburma.so* \
     -w ${REPAIRED_DIR}/ --plat ${PLATFORM_TAG}
 
 # Inject CUDA extensions into the repaired wheel.  patchelf (used by auditwheel)
@@ -356,10 +463,47 @@ if [ -n "$CUDA_EP_STAGING_TEMP" ]; then
     rm -rf "$CUDA_EP_STAGING_TEMP"
 fi
 
+# NPU only: move auditwheel-vendored .libs into mooncake/ and set RPATH=$ORIGIN
+# on all ELF files so everything resolves from a single directory.
+if [ "$NPU_BUILD" = "1" ]; then
+    REPAIRED_WHEEL=$(ls ${REPAIRED_DIR}/*.whl 2>/dev/null | head -1)
+    if [ -n "$REPAIRED_WHEEL" ]; then
+        WHEEL_UNPACK_DIR=$(mktemp -d)
+        python${PYTHON_VERSION} -m wheel unpack "$REPAIRED_WHEEL" -d "$WHEEL_UNPACK_DIR"
+        UNPACKED_PKG_DIR=$(find "$WHEEL_UNPACK_DIR" -mindepth 1 -maxdepth 1 -type d | head -1)
+        VENDORED_LIBS_DIR=$(find "$UNPACKED_PKG_DIR" -mindepth 1 -maxdepth 1 -type d -name "*.libs" | head -1)
+        if [ -n "$VENDORED_LIBS_DIR" ]; then
+            echo "Moving vendored libraries into mooncake/..."
+            cp "$VENDORED_LIBS_DIR"/*.so* "${UNPACKED_PKG_DIR}/mooncake/" 2>/dev/null || true
+            rm -rf "$VENDORED_LIBS_DIR"
+            rm -f "${UNPACKED_PKG_DIR}/$(basename "$VENDORED_LIBS_DIR").pth"
+        fi
+        echo "Setting RPATH=\$ORIGIN for all ELF files..."
+        find "${UNPACKED_PKG_DIR}/mooncake/" -type f -exec sh -c 'file "$1" | grep -q ELF' _ {} \; -print0 | xargs -0 patchelf --force-rpath --set-rpath '$ORIGIN'
+        echo "Verifying RPATH..."
+        verify_failed=0
+        while IFS= read -r -d '' f; do
+            rpath=$(patchelf --print-rpath "$f")
+            if [ "$rpath" != '$ORIGIN' ]; then
+                echo "Error: RPATH verification failed for $(basename "$f"): got '${rpath}'"
+                verify_failed=1
+            fi
+        done < <(find "${UNPACKED_PKG_DIR}/mooncake/" -type f -exec sh -c 'file "$1" | grep -q ELF' _ {} \; -print0)
+        if [ "$verify_failed" -ne 0 ]; then
+            exit 1
+        fi
+        rm "$REPAIRED_WHEEL"
+        python${PYTHON_VERSION} -m wheel pack "$UNPACKED_PKG_DIR" -d "${REPAIRED_DIR}/"
+        rm -rf "$WHEEL_UNPACK_DIR"
+    fi
+fi
+
 # Replace original wheel with repaired wheel
 rm -f ${OUTPUT_DIR}/*.whl
 mv ${REPAIRED_DIR}/*.whl ${OUTPUT_DIR}/
 
 cd ..
 
+[[ -f mooncake-wheel/pyproject.toml.backup ]] && mv mooncake-wheel/pyproject.toml.backup mooncake-wheel/pyproject.toml
+
 echo "Wheel package built and repaired successfully!"
diff --git a/scripts/check_hicache_hugepage_requirements.py b/scripts/check_hicache_hugepage_requirements.py
new file mode 100644
index 00000000..9080d1fe
--- /dev/null
+++ b/scripts/check_hicache_hugepage_requirements.py
@@ -0,0 +1,236 @@
+#!/usr/bin/env python3
+"""Estimate HugeTLB capacity for Mooncake HiCache launches."""
+
+from __future__ import annotations
+
+import argparse
+from dataclasses import dataclass
+from pathlib import Path
+import re
+import sys
+
+
+_UNIT_MAP = {
+    "b": 1,
+    "k": 1024,
+    "kb": 1024,
+    "kib": 1024,
+    "m": 1024**2,
+    "mb": 1024**2,
+    "mib": 1024**2,
+    "g": 1024**3,
+    "gb": 1024**3,
+    "gib": 1024**3,
+    "t": 1024**4,
+    "tb": 1024**4,
+    "tib": 1024**4,
+}
+
+
+@dataclass(frozen=True)
+class BudgetSummary:
+    baseline_per_rank_bytes: int
+    baseline_total_bytes: int
+    clean_arena_per_rank_bytes: int
+    clean_arena_total_bytes: int
+    available_bytes: int | None
+    status: str
+    exit_code: int
+
+
+def parse_size(value: str) -> int:
+    normalized = value.strip().lower().replace("_", "")
+    match = re.fullmatch(r"([0-9]+(?:\.[0-9]+)?)([a-z]*)", normalized)
+    if match is None:
+        raise ValueError(f"invalid size: {value!r}")
+
+    number = float(match.group(1))
+    unit = match.group(2) or "b"
+    if unit not in _UNIT_MAP:
+        raise ValueError(f"unsupported size unit in {value!r}")
+
+    return int(number * _UNIT_MAP[unit])
+
+
+def format_size(num_bytes: int) -> str:
+    for unit, scale in (("TiB", 1024**4), ("GiB", 1024**3), ("MiB", 1024**2)):
+        if num_bytes >= scale:
+            return f"{num_bytes / scale:.2f} {unit}"
+    if num_bytes >= 1024:
+        return f"{num_bytes / 1024:.2f} KiB"
+    return f"{num_bytes} B"
+
+
+def read_available_hugetlb(page_size_bytes: int) -> int | None:
+    if page_size_bytes % 1024 != 0:
+        return None
+
+    page_size_kib = page_size_bytes // 1024
+    path = Path(f"/sys/kernel/mm/hugepages/hugepages-{page_size_kib}kB/nr_hugepages")
+    try:
+        nr_hugepages = int(path.read_text(encoding="utf-8").strip())
+    except (FileNotFoundError, PermissionError, ValueError, OSError):
+        return None
+    return nr_hugepages * page_size_bytes
+
+
+def evaluate_budget(
+    *,
+    tp_size: int,
+    hicache_size_bytes: int,
+    global_segment_size_bytes: int,
+    arena_pool_size_bytes: int,
+    available_bytes: int | None,
+) -> BudgetSummary:
+    if tp_size <= 0:
+        raise ValueError("tp_size must be positive")
+
+    baseline_per_rank = hicache_size_bytes + global_segment_size_bytes
+    baseline_total = baseline_per_rank * tp_size
+    clean_arena_per_rank = baseline_per_rank + arena_pool_size_bytes
+    clean_arena_total = clean_arena_per_rank * tp_size
+
+    if available_bytes is None:
+        status = "planning_only"
+        exit_code = 0
+    elif available_bytes < baseline_total:
+        status = "insufficient_for_baseline"
+        exit_code = 2
+    elif available_bytes < clean_arena_total:
+        status = "baseline_fits_arena_may_fallback"
+        exit_code = 1
+    else:
+        status = "clean_arena_budget_available"
+        exit_code = 0
+
+    return BudgetSummary(
+        baseline_per_rank_bytes=baseline_per_rank,
+        baseline_total_bytes=baseline_total,
+        clean_arena_per_rank_bytes=clean_arena_per_rank,
+        clean_arena_total_bytes=clean_arena_total,
+        available_bytes=available_bytes,
+        status=status,
+        exit_code=exit_code,
+    )
+
+
+def build_parser() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(
+        description=(
+            "Estimate HugeTLB capacity needed for a Mooncake HiCache launch. "
+            "The tool reports both a baseline floor and a cleaner arena-backed "
+            "target derived from benchmark bring-up experience."
+        )
+    )
+    parser.add_argument("--tp-size", type=int, required=True)
+    parser.add_argument("--hicache-size", required=True)
+    parser.add_argument("--global-segment-size", required=True)
+    parser.add_argument("--arena-pool-size", default="0")
+    parser.add_argument("--hugepage-size", default="2mb")
+    parser.add_argument(
+        "--available-hugetlb",
+        help=(
+            "Override detected HugeTLB capacity with an explicit size such as "
+            "'96gb'. If omitted, the tool tries /sys/kernel/mm/hugepages."
+        ),
+    )
+    return parser
+
+
+def print_summary(args: argparse.Namespace, summary: BudgetSummary) -> None:
+    hicache_size = parse_size(args.hicache_size)
+    global_segment_size = parse_size(args.global_segment_size)
+    arena_pool_size = parse_size(args.arena_pool_size)
+    hugepage_size = parse_size(args.hugepage_size)
+
+    print("Mooncake HiCache HugeTLB sizing summary")
+    print("")
+    print(f"tp-size:                {args.tp_size}")
+    print(f"hicache-size:           {args.hicache_size} ({format_size(hicache_size)})")
+    print(
+        "global-segment-size:    "
+        f"{args.global_segment_size} ({format_size(global_segment_size)})"
+    )
+    print(
+        f"arena-pool-size:        {args.arena_pool_size} "
+        f"({format_size(arena_pool_size)})"
+    )
+    print(
+        f"hugepage-size:          {args.hugepage_size} "
+        f"({format_size(hugepage_size)})"
+    )
+    if summary.available_bytes is None:
+        print("available HugeTLB:      unavailable (planning-only mode)")
+    else:
+        print("available HugeTLB:      " f"{format_size(summary.available_bytes)}")
+    print("")
+    print("baseline floor / rank:  " f"{format_size(summary.baseline_per_rank_bytes)}")
+    print("baseline floor / total: " f"{format_size(summary.baseline_total_bytes)}")
+    print(
+        "clean arena / rank:     " f"{format_size(summary.clean_arena_per_rank_bytes)}"
+    )
+    print("clean arena / total:    " f"{format_size(summary.clean_arena_total_bytes)}")
+    print("")
+
+    if summary.status == "planning_only":
+        print("Status: planning only")
+        print(
+            "Meaning: baseline and clean-arena targets were computed, but "
+            "host HugeTLB capacity was not detected."
+        )
+    elif summary.status == "insufficient_for_baseline":
+        print("Status: insufficient for baseline")
+        print(
+            "Meaning: the current HugeTLB pool is below the baseline floor. "
+            "Expect startup or allocation failures."
+        )
+    elif summary.status == "baseline_fits_arena_may_fallback":
+        print("Status: baseline fits, arena may fall back")
+        print(
+            "Meaning: baseline should fit, but the cleaner fully arena-backed "
+            "target is above the current HugeTLB pool. Arena allocations may "
+            "partially or fully fall back to regular pages."
+        )
+    else:
+        print("Status: clean arena budget available")
+        print(
+            "Meaning: the current HugeTLB pool meets the cleaner target for "
+            "baseline plus arena-backed bring-up."
+        )
+
+    print("")
+    print(
+        "Note: the baseline floor is a conservative bring-up estimate "
+        "(hicache + global segment per rank). The clean-arena target adds the "
+        "arena pool per rank to show when arena-backed runs are less likely to "
+        "spill onto the regular-page fallback path."
+    )
+
+
+def main(argv: list[str] | None = None) -> int:
+    parser = build_parser()
+    args = parser.parse_args(argv)
+
+    try:
+        available_bytes = (
+            parse_size(args.available_hugetlb)
+            if args.available_hugetlb
+            else read_available_hugetlb(parse_size(args.hugepage_size))
+        )
+        summary = evaluate_budget(
+            tp_size=args.tp_size,
+            hicache_size_bytes=parse_size(args.hicache_size),
+            global_segment_size_bytes=parse_size(args.global_segment_size),
+            arena_pool_size_bytes=parse_size(args.arena_pool_size),
+            available_bytes=available_bytes,
+        )
+    except ValueError as exc:
+        print(f"error: {exc}", file=sys.stderr)
+        return 64
+
+    print_summary(args, summary)
+    return summary.exit_code
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())
diff --git a/scripts/code_format.sh b/scripts/code_format.sh
index d81a434f..a65f7974 100755
--- a/scripts/code_format.sh
+++ b/scripts/code_format.sh
@@ -67,24 +67,37 @@ print_error() {
 
 # Find clang-format binary (require version 20)
 find_clang_format() {
-    if command -v clang-format-20 &> /dev/null; then
-        echo "clang-format-20"
-        return 0
-    else
-        {
-            print_error "clang-format-20 not found. This project requires clang-format version 20."
-            echo ""
-            print_info "Installation instructions for Ubuntu:"
-            echo "  wget -qO- https://apt.llvm.org/llvm-snapshot.gpg.key | sudo tee /etc/apt/trusted.gpg.d/apt.llvm.org.asc > /dev/null"
-            echo "  echo \"deb http://apt.llvm.org/\$(lsb_release -cs)/ llvm-toolchain-\$(lsb_release -cs)-20 main\" | sudo tee /etc/apt/sources.list.d/llvm-20.list"
-            echo "  sudo apt-get update && sudo apt-get install -y clang-format-20"
-            echo ""
-            print_info "Or use the LLVM installation script:"
-            echo "  wget https://apt.llvm.org/llvm.sh && chmod +x llvm.sh && sudo ./llvm.sh 20"
-            echo "  sudo apt-get install -y clang-format-20"
-        } >&2
-        return 1
-    fi
+    # Prefer clang-format-20, but also accept generic clang-format if it is version 20
+    local candidates=("clang-format-20" "clang-format")
+    for candidate in "${candidates[@]}"; do
+        if command -v "${candidate}" &> /dev/null; then
+            local version_out
+            version_out=$("${candidate}" --version 2>/dev/null)
+            if [[ "${version_out}" =~ version[[:space:]]+([0-9]+) ]] && [[ "${BASH_REMATCH[1]}" -eq 20 ]]; then
+                echo "${candidate}"
+                return 0
+            fi
+        fi
+    done
+
+    {
+        print_error "clang-format version 20 not found."
+        if command -v clang-format &> /dev/null; then
+            local current
+            current=$(clang-format --version 2>/dev/null | head -1)
+            print_warn "Found: ${current}"
+        fi
+        echo ""
+        print_info "Installation instructions for Ubuntu:"
+        echo "  wget -qO- https://apt.llvm.org/llvm-snapshot.gpg.key | sudo tee /etc/apt/trusted.gpg.d/apt.llvm.org.asc > /dev/null"
+        echo "  echo \"deb http://apt.llvm.org/\$(lsb_release -cs)/ llvm-toolchain-\$(lsb_release -cs)-20 main\" | sudo tee /etc/apt/sources.list.d/llvm-20.list"
+        echo "  sudo apt-get update && sudo apt-get install -y clang-format-20"
+        echo ""
+        print_info "Or use the LLVM installation script:"
+        echo "  wget https://apt.llvm.org/llvm.sh && chmod +x llvm.sh && sudo ./llvm.sh 20"
+        echo "  sudo apt-get install -y clang-format-20"
+    } >&2
+    return 1
 }
 
 # Parse command line arguments
@@ -117,28 +130,28 @@ parse_args() {
 # Get list of all C/C++ files in the project
 get_all_files() {
     cd "${PROJECT_ROOT}"
-    
+
     # Find all files, respecting .gitignore, then filter by FILE_EXTENSIONS
     local all_files
     all_files=$(git ls-files --cached --others --exclude-standard 2>/dev/null || \
                 find . -type f | sed 's|^\./||')
-    
+
     # Filter by file extensions
     local result
     result=$(echo "${all_files}" | grep -E "${FILE_EXTENSIONS}" || true)
-    
+
     # Exclude specified directories
     for pattern in "${EXCLUDE_DIRS[@]}"; do
         result=$(echo "${result}" | grep -v "${pattern}" || true)
     done
-    
+
     echo "${result}"
 }
 
 # Get list of changed C/C++ files
 get_changed_files() {
     cd "${PROJECT_ROOT}"
-    
+
     # Get changed files compared to base branch
     local changed_files
     if git rev-parse --verify "${BASE_BRANCH}" &> /dev/null; then
@@ -150,15 +163,15 @@ get_changed_files() {
         changed_files=$(git diff --name-only "@{upstream}"...HEAD 2>/dev/null || \
                        git log --name-only --pretty=format: HEAD 2>/dev/null | sort -u || true)
     fi
-    
+
     # Filter C/C++ files and exclude specified directories
     local result
     result=$(echo "${changed_files}" | grep -E "${FILE_EXTENSIONS}" || true)
-    
+
     for pattern in "${EXCLUDE_DIRS[@]}"; do
         result=$(echo "${result}" | grep -v "${pattern}" || true)
     done
-    
+
     echo "${result}"
 }
 
@@ -168,7 +181,7 @@ format_files() {
     local clang_format="$2"
     local formatted_count=0
     local failed_count=0
-    
+
     if [[ -z "${files}" ]]; then
         if ${ALL_MODE}; then
             print_info "No C/C++ files found in the project."
@@ -177,7 +190,7 @@ format_files() {
         fi
         return 0
     fi
-    
+
     print_info "Using $(${clang_format} --version)"
     if ${ALL_MODE}; then
         print_info "Formatting all C/C++ files in the project"
@@ -185,7 +198,7 @@ format_files() {
         print_info "Comparing against: ${BASE_BRANCH}"
     fi
     echo ""
-    
+
     while IFS= read -r file; do
         if [[ -f "${PROJECT_ROOT}/${file}" ]]; then
             if ${CHECK_MODE}; then
@@ -215,7 +228,7 @@ format_files() {
             print_warn "File not found (deleted?): ${file}"
         fi
     done <<< "${files}"
-    
+
     echo ""
     if ${CHECK_MODE}; then
         if [[ ${failed_count} -gt 0 ]]; then
@@ -232,23 +245,23 @@ format_files() {
             return 1
         fi
     fi
-    
+
     return 0
 }
 
 # Main function
 main() {
     parse_args "$@"
-    
+
     print_info "Mooncake Code Format Script"
     echo ""
-    
+
     # Check for .clang-format file
     if [[ ! -f "${PROJECT_ROOT}/.clang-format" ]]; then
         print_error ".clang-format not found in project root"
         exit 1
     fi
-    
+
     # Find clang-format
     local clang_format
     if ! clang_format="$(find_clang_format)"; then
@@ -262,7 +275,7 @@ main() {
     else
         files=$(get_changed_files)
     fi
-    
+
     # Format files
     format_files "${files}" "${clang_format}"
 }
diff --git a/scripts/launch_cxi_transfer_test.sh b/scripts/launch_cxi_transfer_test.sh
new file mode 100644
index 00000000..b57f8537
--- /dev/null
+++ b/scripts/launch_cxi_transfer_test.sh
@@ -0,0 +1,25 @@
+#!/bin/bash
+
+ip=$(hostname -I | awk '{print $1}')
+echo $ip
+
+MC_MIN_PRC_PORT=25565 MC_MAX_PRC_PORT=25565 ./build/mooncake-transfer-engine/tests/cxi_transfer_test \
+    --mode target \
+    --use_device \
+    --server $ip \
+    --num_bufs 4 \
+    --buf_size_gb 1 &
+TARGET_PID=$!
+
+sleep 10
+
+./build/mooncake-transfer-engine/tests/cxi_transfer_test --mode initiator \
+    --server 127.0.0.1:12346 \
+    --target $ip:25565 \
+    --transfer_mb 128 \
+    --threads 4 \
+    --use_device \
+    --warmup 0
+
+kill -INT $TARGET_PID
+wait $TARGET_PID
\ No newline at end of file
diff --git a/scripts/run_tests.sh b/scripts/run_tests.sh
index 64adbe96..3aaf2c7b 100755
--- a/scripts/run_tests.sh
+++ b/scripts/run_tests.sh
@@ -40,10 +40,15 @@ DUMMY_TEST_PID_2=$!
 wait $DUMMY_TEST_PID_1 $DUMMY_TEST_PID_2
 kill $CLIENT_PID || true
 
-pip install torch numpy safetensors packaging
+pip install numpy safetensors packaging
+# Keep the test torch aligned with the EP/PG variants packaged into the CI wheel.
+pip install "${MOONCAKE_TEST_TORCH_SPEC:-torch==2.11.0+cu128}" \
+    --index-url "${MOONCAKE_TEST_TORCH_INDEX_URL:-https://download.pytorch.org/whl/cu128}" \
+    --extra-index-url https://pypi.org/simple
 MC_METADATA_SERVER=http://127.0.0.1:8080/metadata DEFAULT_KV_LEASE_TTL=500 python test_put_get_tensor.py
 MC_METADATA_SERVER=http://127.0.0.1:8080/metadata DEFAULT_KV_LEASE_TTL=500 python test_safetensor_functions.py
 kill $MASTER_PID || true
+wait $MASTER_PID 2>/dev/null || true
 
 
 # Check if MOONCAKE_STORAGE_ROOT_DIR is set and not empty
@@ -59,11 +64,46 @@ if [ -n "$TEST_SSD_OFFLOAD_IN_EVICT" ]; then
     sleep 1
     MC_METADATA_SERVER=http://127.0.0.1:8080/metadata DEFAULT_KV_LEASE_TTL=500 python test_ssd_offload_in_evict.py
     kill $MASTER_PID || true
+    wait $MASTER_PID 2>/dev/null || true
     rm -rf $TEST_ROOT_DIR
 else
     echo "Skipping test: MOONCAKE_STORAGE_ROOT_DIR environment variable is not set"
 fi
 
+if [ -n "$TEST_PROMOTION_ON_HIT" ]; then
+    TEST_ROOT_DIR="/tmp/mooncake_test_promotion"
+    mkdir -p $TEST_ROOT_DIR
+    echo "Running L2->L1 promotion-on-hit e2e test..."
+    # offload_on_evict drives the prerequisite SSD-only state; promotion_on_hit
+    # turns the read path into a promotion trigger; threshold=1 makes the test
+    # deterministic. --root_fs_dir is required so the master returns a non-
+    # empty fsdir from GetStorageConfig, which is the trigger that initializes
+    # the client's FileStorage (and therefore the offload heartbeat).
+    mooncake_master \
+        --default_kv_lease_ttl=500 \
+        --root_fs_dir=$TEST_ROOT_DIR \
+        --enable_offload=true \
+        --offload_on_evict=true \
+        --promotion_on_hit=true \
+        --promotion_admission_threshold=1 &
+    MASTER_PID=$!
+    sleep 1
+    # Lower bucket-flush thresholds so the test workload (~64 MB) actually
+    # writes to disk rather than sitting in the bucket backend's ungrouped
+    # pool until the default 500-key / 256-MB bucket fills.
+    MC_METADATA_SERVER=http://127.0.0.1:8080/metadata \
+        DEFAULT_KV_LEASE_TTL=500 \
+        MOONCAKE_OFFLOAD_FILE_STORAGE_PATH=$TEST_ROOT_DIR \
+        MOONCAKE_OFFLOAD_BUCKET_KEYS_LIMIT=10 \
+        MOONCAKE_OFFLOAD_BUCKET_SIZE_LIMIT_BYTES=10485760 \
+        python test_promotion_on_hit.py
+    kill $MASTER_PID || true
+    wait $MASTER_PID 2>/dev/null || true
+    rm -rf $TEST_ROOT_DIR
+else
+    echo "Skipping test: TEST_PROMOTION_ON_HIT environment variable is not set"
+fi
+
 echo "Running CXL protocol test (test_distributed_object_store_cxl.py)..."
 killall mooncake_master || true
 sleep 2
@@ -77,6 +117,7 @@ CXL_MASTER_PID=$!
 sleep 3
 MC_METADATA_SERVER=http://127.0.0.1:8080/metadata DEFAULT_KV_LEASE_TTL=500 python test_distributed_object_store_cxl.py
 kill $CXL_MASTER_PID || true
+wait $CXL_MASTER_PID 2>/dev/null || true
 sleep 2
 echo "CXL protocol test completed successfully!"
 
@@ -92,6 +133,7 @@ sleep 1
 MC_METADATA_SERVER=http://127.0.0.1:8080/metadata DEFAULT_KV_LEASE_TTL=500 python test_distributed_object_store.py
 sleep 1
 kill $MASTER_PID || true
+wait $MASTER_PID 2>/dev/null || true
 
 
 echo "All tests completed successfully!"
diff --git a/scripts/test_engram_store.py b/scripts/test_engram_store.py
new file mode 100755
index 00000000..759fcd7f
--- /dev/null
+++ b/scripts/test_engram_store.py
@@ -0,0 +1,419 @@
+#!/usr/bin/env python3
+"""
+Integration tests for the Mooncake EngramStore backend.
+
+This suite validates:
+1. metadata-only construction and store-backed construction
+2. populate / lookup / remove on the simplified Store interface
+3. Python-list and NumPy row-id lookup paths
+4. error handling and cleanup behavior
+"""
+
+import importlib
+import os
+import shutil
+import socket
+import subprocess
+import sys
+import tempfile
+import time
+import unittest
+import urllib.error
+import urllib.request
+import uuid
+from pathlib import Path
+
+import numpy as np
+
+REPO_ROOT = Path(__file__).resolve().parents[1]
+BUILD_DIR = os.environ.get("MOONCAKE_BUILD_DIR", "build")
+BUILD_STORE = REPO_ROOT / BUILD_DIR / "mooncake-integration"
+WHEEL_DIR = REPO_ROOT / "mooncake-wheel"
+MASTER_BINARY = REPO_ROOT / BUILD_DIR / "mooncake-store" / "src" / "mooncake_master"
+
+for path in (BUILD_STORE, WHEEL_DIR):
+    if path.is_dir() and str(path) not in sys.path:
+        sys.path.insert(0, str(path))
+
+from mooncake.mooncake_config import MooncakeConfig
+
+GLOBAL_STORE = None
+STORE_MODULE = None
+TEST_CONFIG = None
+MASTER_PROCESS = None
+MASTER_LOG_PATH = None
+MASTER_LOG_FILE = None
+
+
+def import_store_module():
+    if not BUILD_STORE.is_dir():
+        raise ImportError(
+            f"{BUILD_STORE} not found. Build Mooncake with store support: "
+            "cd build && cmake .. -DWITH_STORE=ON && make -j 128"
+        )
+
+    store_module = importlib.import_module("store")
+    print(f"✅ store.so imported successfully from {BUILD_STORE}")
+    return store_module
+
+
+def find_free_port():
+    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
+        sock.bind(("127.0.0.1", 0))
+        return sock.getsockname()[1]
+
+
+def resolve_master_binary():
+    if MASTER_BINARY.is_file():
+        return str(MASTER_BINARY)
+
+    binary = shutil.which("mooncake_master")
+    if binary:
+        return binary
+
+    raise FileNotFoundError(
+        "Cannot find mooncake_master. Build it first or install it into PATH."
+    )
+
+
+def wait_for_tcp_port(host, port, timeout=20.0):
+    deadline = time.time() + timeout
+    while time.time() < deadline:
+        try:
+            with socket.create_connection((host, port), timeout=0.5):
+                return
+        except OSError:
+            time.sleep(0.1)
+    raise RuntimeError(f"Timed out waiting for TCP port {host}:{port}")
+
+
+def wait_for_metadata_server(metadata_url, timeout=20.0):
+    deadline = time.time() + timeout
+    while time.time() < deadline:
+        try:
+            with urllib.request.urlopen(
+                metadata_url + "?key=engram_store_test_probe", timeout=1.0
+            ):
+                return
+        except urllib.error.HTTPError as exc:
+            if exc.code in (200, 400, 404):
+                return
+        except urllib.error.URLError:
+            time.sleep(0.1)
+        else:
+            return
+    raise RuntimeError(f"Timed out waiting for metadata server {metadata_url}")
+
+
+def read_master_log():
+    if MASTER_LOG_PATH is None or not os.path.exists(MASTER_LOG_PATH):
+        return "<no master log available>"
+    with open(MASTER_LOG_PATH, "r", encoding="utf-8", errors="replace") as fin:
+        return fin.read()
+
+
+def start_local_master():
+    global MASTER_PROCESS, MASTER_LOG_FILE, MASTER_LOG_PATH
+
+    rpc_port = find_free_port()
+    http_port = find_free_port()
+    metrics_port = find_free_port()
+    master_binary = resolve_master_binary()
+
+    fd, MASTER_LOG_PATH = tempfile.mkstemp(prefix="engram_store-master-", suffix=".log")
+    os.close(fd)
+
+    MASTER_LOG_FILE = open(MASTER_LOG_PATH, "w", encoding="utf-8")
+    cmd = [
+        master_binary,
+        "--default_kv_lease_ttl=500",
+        "--enable_http_metadata_server=true",
+        "--rpc_address=127.0.0.1",
+        f"--rpc_port={rpc_port}",
+        "--http_metadata_server_host=127.0.0.1",
+        f"--http_metadata_server_port={http_port}",
+        f"--metrics_port={metrics_port}",
+    ]
+    MASTER_PROCESS = subprocess.Popen(
+        cmd,
+        cwd=REPO_ROOT,
+        stdout=MASTER_LOG_FILE,
+        stderr=subprocess.STDOUT,
+        text=True,
+    )
+
+    metadata_url = f"http://127.0.0.1:{http_port}/metadata"
+    try:
+        wait_for_tcp_port("127.0.0.1", rpc_port)
+        wait_for_metadata_server(metadata_url)
+    except Exception:
+        log_output = read_master_log()
+        stop_local_master()
+        raise RuntimeError(
+            "Failed to start local mooncake_master for EngramStore tests.\n"
+            f"Log output:\n{log_output}"
+        )
+
+    print(
+        "✅ Started local mooncake_master for EngramStore tests at "
+        f"127.0.0.1:{rpc_port} with metadata {metadata_url}"
+    )
+    return MooncakeConfig(
+        local_hostname="127.0.0.1",
+        metadata_server=metadata_url,
+        global_segment_size=128 * 1024 * 1024,
+        local_buffer_size=64 * 1024 * 1024,
+        protocol="tcp",
+        device_name="",
+        master_server_address=f"127.0.0.1:{rpc_port}",
+    )
+
+
+def stop_local_master():
+    global MASTER_PROCESS, MASTER_LOG_FILE, MASTER_LOG_PATH
+
+    if MASTER_PROCESS is not None:
+        if MASTER_PROCESS.poll() is None:
+            MASTER_PROCESS.terminate()
+            try:
+                MASTER_PROCESS.wait(timeout=10)
+            except subprocess.TimeoutExpired:
+                MASTER_PROCESS.kill()
+                MASTER_PROCESS.wait(timeout=5)
+        MASTER_PROCESS = None
+
+    if MASTER_LOG_FILE is not None:
+        MASTER_LOG_FILE.close()
+        MASTER_LOG_FILE = None
+
+    if MASTER_LOG_PATH and os.path.exists(MASTER_LOG_PATH):
+        os.remove(MASTER_LOG_PATH)
+        MASTER_LOG_PATH = None
+
+
+def load_test_config():
+    if os.getenv("MOONCAKE_CONFIG_PATH") or os.getenv("MOONCAKE_MASTER"):
+        print("Using Mooncake configuration from environment")
+        return MooncakeConfig.load_from_env()
+
+    return start_local_master()
+
+
+def create_store_connection(store_module, config):
+    store = store_module.MooncakeDistributedStore()
+    print(
+        f"[{os.getpid()}] Connecting to Mooncake Master at "
+        f"{config.master_server_address}..."
+    )
+
+    rc = store.setup(
+        config.local_hostname,
+        config.metadata_server,
+        config.global_segment_size,
+        config.local_buffer_size,
+        config.protocol,
+        config.device_name,
+        config.master_server_address,
+    )
+    if rc != 0:
+        raise RuntimeError(f"Failed to setup mooncake store, error code: {rc}")
+
+    print("✅ Store connection established")
+    return store
+
+
+def setUpModule():
+    global GLOBAL_STORE, STORE_MODULE, TEST_CONFIG
+    STORE_MODULE = import_store_module()
+    try:
+        TEST_CONFIG = load_test_config()
+        GLOBAL_STORE = create_store_connection(STORE_MODULE, TEST_CONFIG)
+    except Exception:
+        stop_local_master()
+        raise
+
+
+def tearDownModule():
+    global GLOBAL_STORE
+    if GLOBAL_STORE:
+        print("\nClosing global store connection...")
+        GLOBAL_STORE.close()
+        GLOBAL_STORE = None
+    stop_local_master()
+
+
+class EngramStoreTestBase(unittest.TestCase):
+    def setUp(self):
+        if GLOBAL_STORE is None:
+            self.skipTest("Store not initialized")
+
+        self.store = GLOBAL_STORE
+        self.EngramStore = STORE_MODULE.EngramStore
+        self.EngramStoreConfig = STORE_MODULE.EngramStoreConfig
+        self._created_engram_stores = []
+        self._next_layer_id = uuid.uuid4().int & 0x7FFFFFFF
+
+    def tearDown(self):
+        for engram_store in reversed(self._created_engram_stores):
+            try:
+                engram_store.remove_from_store(force=True)
+            except Exception as exc:
+                print(f"Warning: failed to clean up EngramStore test layer: {exc}")
+
+    def create_config(self):
+        cfg = self.EngramStoreConfig()
+        cfg.table_vocab_sizes = [17, 19, 23, 29]
+        cfg.embedding_dim = 8
+        return cfg
+
+    def create_engram_store(self, layer_id=None, store_marker=Ellipsis):
+        if layer_id is None:
+            layer_id = self._next_layer_id
+            self._next_layer_id += 1
+
+        cfg = self.create_config()
+        if store_marker is Ellipsis:
+            engram_store = self.EngramStore(layer_id=layer_id, config=cfg, store=self.store)
+            self._created_engram_stores.append(engram_store)
+        elif store_marker is None:
+            engram_store = self.EngramStore(layer_id=layer_id, config=cfg)
+        else:
+            engram_store = self.EngramStore(layer_id=layer_id, config=cfg, store=store_marker)
+            self._created_engram_stores.append(engram_store)
+        return cfg, engram_store
+
+    def make_embedding_tables(self, engram_store):
+        embed_dim = engram_store.get_embedding_dim()
+        tables = []
+        for head_idx, vocab_size in enumerate(engram_store.get_table_vocab_sizes()):
+            base = np.arange(vocab_size * embed_dim, dtype=np.float32).reshape(
+                vocab_size, embed_dim
+            )
+            tables.append(base + head_idx * 1000)
+        return tables
+
+    def populate_store(self, engram_store):
+        tables = self.make_embedding_tables(engram_store)
+        engram_store.populate(tables)
+        return tables
+
+
+class TestEngramStoreMetadata(EngramStoreTestBase):
+    def test_creation_and_metadata(self):
+        cfg, engram_store = self.create_engram_store()
+        self.assertEqual(engram_store.get_num_heads(), len(cfg.table_vocab_sizes))
+        self.assertEqual(engram_store.get_embedding_dim(), cfg.embedding_dim)
+        self.assertEqual(engram_store.get_table_vocab_sizes(), cfg.table_vocab_sizes)
+        self.assertEqual(len(engram_store.get_store_keys()), len(cfg.table_vocab_sizes))
+
+    def test_creation_without_store_keeps_metadata_accessible(self):
+        layer_id = self._next_layer_id
+        cfg, engram_store = self.create_engram_store(layer_id=layer_id, store_marker=None)
+        self.assertEqual(engram_store.get_num_heads(), len(cfg.table_vocab_sizes))
+        self.assertEqual(engram_store.get_embedding_dim(), cfg.embedding_dim)
+        self.assertEqual(engram_store.get_store_keys()[0], f"engram:l{layer_id}:h0")
+
+
+class TestStorePopulateAndLookup(EngramStoreTestBase):
+    def test_populate_and_lookup_shape_with_python_lists(self):
+        _, engram_store = self.create_engram_store()
+        self.populate_store(engram_store)
+
+        row_ids = [
+            [[0, 1, 2, 3], [4, 5, 6, 7]],
+            [[1, 2, 3, 4], [8, 9, 10, 11]],
+        ]
+        output = engram_store.lookup(row_ids)
+
+        expected_shape = (
+            len(row_ids),
+            len(row_ids[0]),
+            engram_store.get_num_heads(),
+            engram_store.get_embedding_dim(),
+        )
+        self.assertEqual(output.shape, expected_shape)
+        self.assertFalse(np.any(np.isnan(output)))
+        self.assertFalse(np.any(np.isinf(output)))
+
+    def test_lookup_matches_stored_rows_from_numpy_ids(self):
+        _, engram_store = self.create_engram_store()
+        tables = self.populate_store(engram_store)
+
+        row_ids = np.array(
+            [[[0, 1, 2, 3], [4, 5, 6, 7], [8, 9, 10, 11]]], dtype=np.int64
+        )
+        output = np.asarray(engram_store.lookup(row_ids))
+
+        for pos in range(row_ids.shape[1]):
+            for head in range(engram_store.get_num_heads()):
+                idx = row_ids[0, pos, head]
+                np.testing.assert_allclose(output[0, pos, head], tables[head][idx])
+
+    def test_lookup_list_and_numpy_paths_match(self):
+        _, engram_store = self.create_engram_store()
+        self.populate_store(engram_store)
+
+        row_ids_list = [
+            [[0, 1, 2, 3], [4, 5, 6, 7]],
+            [[1, 2, 3, 4], [8, 9, 10, 11]],
+        ]
+        row_ids_numpy = np.asarray(row_ids_list, dtype=np.int64)
+
+        output_from_list = np.asarray(engram_store.lookup(row_ids_list))
+        output_from_numpy = np.asarray(engram_store.lookup(row_ids_numpy))
+        np.testing.assert_allclose(output_from_list, output_from_numpy)
+
+    def test_remove_from_store(self):
+        _, engram_store = self.create_engram_store()
+        self.populate_store(engram_store)
+
+        removed = engram_store.remove_from_store(force=True)
+        self.assertEqual(removed, engram_store.get_num_heads())
+
+        for key in engram_store.get_store_keys():
+            self.assertEqual(self.store.is_exist(key), 0)
+
+        self.assertEqual(engram_store.remove_from_store(force=True), 0)
+
+
+class TestErrorHandling(EngramStoreTestBase):
+    def test_populate_rejects_reusing_existing_layer_keys(self):
+        _, engram_store = self.create_engram_store()
+        self.populate_store(engram_store)
+        with self.assertRaises(Exception):
+            engram_store.populate(self.make_embedding_tables(engram_store))
+
+    def test_lookup_rejects_missing_tables(self):
+        _, engram_store = self.create_engram_store()
+        with self.assertRaises(Exception):
+            engram_store.lookup([[[0, 1, 2, 3]]])
+
+    def test_populate_rejects_wrong_table_shape(self):
+        _, engram_store = self.create_engram_store()
+        tables = self.make_embedding_tables(engram_store)
+        tables[0] = tables[0][:-1]
+        with self.assertRaises(Exception):
+            engram_store.populate(tables)
+
+    def test_lookup_rejects_empty_input(self):
+        _, engram_store = self.create_engram_store()
+        self.populate_store(engram_store)
+        with self.assertRaises(Exception):
+            engram_store.lookup([])
+
+    def test_lookup_rejects_wrong_head_dimension(self):
+        _, engram_store = self.create_engram_store()
+        self.populate_store(engram_store)
+        row_ids = np.zeros((1, 1, engram_store.get_num_heads() - 1), dtype=np.int64)
+        with self.assertRaises(Exception):
+            engram_store.lookup(row_ids)
+
+    def test_lookup_rejects_out_of_range_row_id(self):
+        _, engram_store = self.create_engram_store()
+        self.populate_store(engram_store)
+        with self.assertRaises(Exception):
+            engram_store.lookup([[[999, 1, 2, 3]]])
+
+
+if __name__ == "__main__":
+    unittest.main(verbosity=2)
diff --git a/scripts/test_hicache_hugepage_requirements.py b/scripts/test_hicache_hugepage_requirements.py
new file mode 100644
index 00000000..6de2c74e
--- /dev/null
+++ b/scripts/test_hicache_hugepage_requirements.py
@@ -0,0 +1,73 @@
+#!/usr/bin/env python3
+"""Unit tests for the HugeTLB sizing helper."""
+
+from __future__ import annotations
+
+from pathlib import Path
+import sys
+import unittest
+
+
+sys.path.insert(0, str(Path(__file__).resolve().parent))
+import check_hicache_hugepage_requirements as helper  # noqa: E402
+
+
+class ParseSizeTest(unittest.TestCase):
+    def test_parses_human_readable_units(self) -> None:
+        self.assertEqual(helper.parse_size("8gb"), 8 * 1024**3)
+        self.assertEqual(helper.parse_size("2MB"), 2 * 1024**2)
+        self.assertEqual(helper.parse_size("1.5tb"), int(1.5 * 1024**4))
+
+    def test_rejects_invalid_values(self) -> None:
+        with self.assertRaises(ValueError):
+            helper.parse_size("eight-gigabytes")
+
+
+class EvaluateBudgetTest(unittest.TestCase):
+    def test_detects_baseline_failure(self) -> None:
+        summary = helper.evaluate_budget(
+            tp_size=1,
+            hicache_size_bytes=80 * 1024**3,
+            global_segment_size_bytes=48 * 1024**3,
+            arena_pool_size_bytes=0,
+            available_bytes=96 * 1024**3,
+        )
+        self.assertEqual(summary.status, "insufficient_for_baseline")
+        self.assertEqual(summary.exit_code, 2)
+
+    def test_detects_partial_arena_budget(self) -> None:
+        summary = helper.evaluate_budget(
+            tp_size=4,
+            hicache_size_bytes=64 * 1024**3,
+            global_segment_size_bytes=8 * 1024**3,
+            arena_pool_size_bytes=56 * 1024**3,
+            available_bytes=320 * 1024**3,
+        )
+        self.assertEqual(summary.status, "baseline_fits_arena_may_fallback")
+        self.assertEqual(summary.exit_code, 1)
+
+    def test_detects_clean_arena_budget(self) -> None:
+        summary = helper.evaluate_budget(
+            tp_size=4,
+            hicache_size_bytes=64 * 1024**3,
+            global_segment_size_bytes=8 * 1024**3,
+            arena_pool_size_bytes=56 * 1024**3,
+            available_bytes=512 * 1024**3,
+        )
+        self.assertEqual(summary.status, "clean_arena_budget_available")
+        self.assertEqual(summary.exit_code, 0)
+
+    def test_supports_planning_only_mode(self) -> None:
+        summary = helper.evaluate_budget(
+            tp_size=2,
+            hicache_size_bytes=56 * 1024**3,
+            global_segment_size_bytes=8 * 1024**3,
+            arena_pool_size_bytes=56 * 1024**3,
+            available_bytes=None,
+        )
+        self.assertEqual(summary.status, "planning_only")
+        self.assertEqual(summary.exit_code, 0)
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/scripts/test_installation.sh b/scripts/test_installation.sh
index 8c65abdf..9e7f0776 100755
--- a/scripts/test_installation.sh
+++ b/scripts/test_installation.sh
@@ -37,7 +37,7 @@ echo "Running import structure test..."
 # Run the import structure test
 cp -r mooncake-wheel/tests test_env/
 cd test_env
-pip install torch numpy
+pip install torch==2.11.0 numpy
 python tests/test_import_structure.py
 
 echo "Running mooncake config test..."
diff --git a/scripts/test_tensor_api.py b/scripts/test_tensor_api.py
index a0972d2a..bdffc437 100644
--- a/scripts/test_tensor_api.py
+++ b/scripts/test_tensor_api.py
@@ -8,6 +8,7 @@ import unittest
 import torch
 import numpy as np
 from dataclasses import dataclass
+import mooncake.store as mooncake_store
 from mooncake.store import MooncakeDistributedStore
 from mooncake.store import ReplicateConfig
 
@@ -15,6 +16,14 @@ from mooncake.mooncake_config import MooncakeConfig
 
 import concurrent.futures
 
+
+@dataclass
+class WriterPartition:
+    rank: int
+    size: int
+    split_dim: int
+
+
 # ==========================================
 #  Global Variables & Configuration
 # ==========================================
@@ -29,12 +38,13 @@ DEFAULT_LOCAL_BUFFER_SIZE = 8 * 1024 * 1024 * 1024    # 8 GB
 DEFAULT_MASTER_METRICS_PORT = 9003
 DEFAULT_CHECK_SERVER = False
 
-# Must match C++ TensorMetadata: int32_t dtype + int32_t ndim + int64_t shape[4]
-TENSOR_METADATA_SIZE = 4 + 4 + 8 * 4  # 40 bytes
+# Must match current C++ TensorMetadata layout.
+# TensorObjectHeader (40) + TensorLayoutMetadata (264) = 304 bytes.
+TENSOR_METADATA_SIZE = 304
 
 
 def serialized_tensor_size(tensor):
-    """Size in bytes of [TensorMetadata][tensor data] as stored by get_tensor_into."""
+    """Size in bytes of [TensorObjectHeader+layout metadata][tensor data]."""
     return TENSOR_METADATA_SIZE + tensor.numel() * tensor.element_size()
 
 def verify_tensor_equality(original, received, rtol=0, atol=0, verbose=True):
@@ -143,6 +153,234 @@ def setUpModule():
         print(f"❌ Failed to establish global store connection: {e}")
         sys.exit(1)
 
+
+def require_unified_parallelism_api(test_case):
+    missing = [
+        name for name in ["ParallelAxis", "TensorParallelism", "ReadTarget"]
+        if not hasattr(mooncake_store, name)
+    ]
+    if missing:
+        test_case.skipTest(
+            "Unified parallelism Python bindings are not available in the installed extension: "
+            + ", ".join(missing)
+        )
+
+
+def make_parallel_axis(kind, rank, size, split_dim=None):
+    axis = mooncake_store.ParallelAxis()
+    axis.kind = kind
+    axis.rank = rank
+    axis.size = size
+    axis.split_dim = split_dim
+    return axis
+
+
+def make_tensor_parallelism(axes):
+    parallelism = mooncake_store.TensorParallelism()
+    parallelism.axes = axes
+    return parallelism
+
+
+def make_read_target(mode, parallelism=None):
+    target = mooncake_store.ReadTarget()
+    target.mode = mode
+    target.parallelism = parallelism
+    return target
+
+
+def make_writer_partition(rank, size, split_dim):
+    return WriterPartition(rank=rank, size=size, split_dim=split_dim)
+
+
+def build_tp_parallelism(tp_size, split_dim, rank=0):
+    return make_tensor_parallelism([
+        make_parallel_axis("tp", rank=rank, size=tp_size, split_dim=split_dim)
+    ])
+
+
+def build_dp_tp_parallelism(dp_rank, dp_size, tp_rank, tp_size, split_dim):
+    return make_tensor_parallelism([
+        make_parallel_axis("dp", rank=dp_rank, size=dp_size),
+        make_parallel_axis("tp", rank=tp_rank, size=tp_size, split_dim=split_dim),
+    ])
+
+
+def build_pp_tp_parallelism(pp_rank, pp_size, stage_id, tp_rank, tp_size, split_dim):
+    pp_axis = make_parallel_axis("pp", rank=pp_rank, size=pp_size)
+    pp_axis.stage_id = stage_id
+    return make_tensor_parallelism([
+        pp_axis,
+        make_parallel_axis("tp", rank=tp_rank, size=tp_size, split_dim=split_dim),
+    ])
+
+
+def build_ep_parallelism(ep_rank, ep_size, expert_id):
+    ep_axis = make_parallel_axis("ep", rank=ep_rank, size=ep_size)
+    ep_axis.expert_id = expert_id
+    return make_tensor_parallelism([ep_axis])
+
+
+def build_ep_tp_parallelism(ep_rank, ep_size, expert_id, tp_rank, tp_size, split_dim):
+    ep_axis = make_parallel_axis("ep", rank=ep_rank, size=ep_size)
+    ep_axis.expert_id = expert_id
+    return make_tensor_parallelism([
+        ep_axis,
+        make_parallel_axis("tp", rank=tp_rank, size=tp_size, split_dim=split_dim),
+    ])
+
+
+def build_parallelism_mode(mode, tp_rank, tp_size, split_dim):
+    if mode == "tp":
+        return build_tp_parallelism(tp_size, split_dim, rank=tp_rank)
+    if mode == "dp_tp":
+        return build_dp_tp_parallelism(0, 2, tp_rank, tp_size, split_dim)
+    if mode == "pp_tp":
+        return build_pp_tp_parallelism(0, 2, 7, tp_rank, tp_size, split_dim)
+    if mode == "ep_tp":
+        return build_ep_tp_parallelism(0, 2, 11, tp_rank, tp_size, split_dim)
+    raise ValueError(f"unsupported parallelism mode: {mode}")
+
+
+def put_full_tensor_with_parallelism_mode(store, key, tensor, split_dim, put_tp_size, mode):
+    if mode == "tp":
+        shards = tensor.chunk(put_tp_size, split_dim)
+        for rank, shard in enumerate(shards):
+            rc = store.put_tensor_with_parallelism(
+                key,
+                shard.contiguous(),
+                build_parallelism_mode(mode, rank, put_tp_size, split_dim),
+            )
+            if rc != 0:
+                return rc
+        return 0
+
+    for rank in range(put_tp_size):
+        rc = store.put_tensor_with_parallelism(
+            key,
+            tensor,
+            build_parallelism_mode(mode, rank, put_tp_size, split_dim),
+        )
+        if rc != 0:
+            return rc
+    return 0
+
+
+def clone_parallelism(parallelism):
+    cloned_axes = []
+    for axis in parallelism.axes:
+        cloned_axis = make_parallel_axis(
+            axis.kind,
+            rank=axis.rank,
+            size=axis.size,
+            split_dim=axis.split_dim,
+        )
+        if getattr(axis, "expert_id", None) is not None:
+            cloned_axis.expert_id = axis.expert_id
+        if getattr(axis, "stage_id", None) is not None:
+            cloned_axis.stage_id = axis.stage_id
+        cloned_axes.append(cloned_axis)
+    return make_tensor_parallelism(cloned_axes)
+
+
+def put_full_tensor_with_parallelism_template(store, key, tensor, parallelism):
+    tp_axis = next(axis for axis in parallelism.axes if axis.kind == "tp")
+    for rank in range(tp_axis.size):
+        shard_parallelism = clone_parallelism(parallelism)
+        for axis in shard_parallelism.axes:
+            if axis.kind == "tp":
+                axis.rank = rank
+                axis.size = tp_axis.size
+                axis.split_dim = tp_axis.split_dim
+                break
+        rc = store.put_tensor_with_parallelism(
+            key,
+            tensor,
+            shard_parallelism,
+        )
+        if rc != 0:
+            return rc
+    return 0
+
+
+def expected_target_shard(tensor, split_dim, target_tp_size, target_rank):
+    return chunk_tensor_for_rank(tensor, target_tp_size, split_dim, target_rank).contiguous()
+
+
+def reorder_parallelism_axes(parallelism, ordered_kinds):
+    axis_by_kind = {axis.kind: axis for axis in parallelism.axes}
+    return make_tensor_parallelism([axis_by_kind[kind] for kind in ordered_kinds])
+
+
+def make_deterministic_tensor(shape, dtype=torch.float32):
+    numel = 1
+    for dim in shape:
+        numel *= dim
+    return torch.arange(numel, dtype=dtype).view(*shape).contiguous()
+
+
+def chunk_tensor_for_rank(tensor, tp_size, split_dim, rank):
+    return tensor.chunk(tp_size, split_dim)[rank]
+
+
+def put_full_tensor_with_unified_tp(store, key, tensor, tp_size, split_dim):
+    for rank, shard in enumerate(tensor.chunk(tp_size, split_dim)):
+        rc = store.put_tensor_with_parallelism(
+            key,
+            shard.contiguous(),
+            build_tp_parallelism(tp_size, split_dim, rank=rank),
+        )
+        if rc != 0:
+            return rc
+    return 0
+
+
+def batch_put_full_tensors_with_unified_tp(store, keys, tensors, tp_size, split_dim):
+    results = [0 for _ in keys]
+    for rank in range(tp_size):
+        shard_tensors = [chunk_tensor_for_rank(tensor, tp_size, split_dim, rank).contiguous() for tensor in tensors]
+        rank_results = store.batch_put_tensor_with_parallelism(
+            keys,
+            shard_tensors,
+            [build_tp_parallelism(tp_size, split_dim, rank=rank) for _ in keys],
+        )
+        results = rank_results
+        if any(rc != 0 for rc in rank_results):
+            return rank_results
+    return results
+
+
+def reconstruction_case_name(shape, split_dim, tp_size):
+    shape_str = "x".join(str(dim) for dim in shape)
+    return f"shape={shape_str}, split_dim={split_dim}, tp_size={tp_size}"
+
+
+def is_uniform_tp_case(shape, split_dim, tp_size):
+    return shape[split_dim] % tp_size == 0
+
+
+def put_uniform_full_tensor_with_unified_tp(store, key, tensor, tp_size, split_dim):
+    if not is_uniform_tp_case(tuple(tensor.shape), split_dim, tp_size):
+        return -1
+    return put_full_tensor_with_unified_tp(store, key, tensor, tp_size, split_dim)
+
+
+def make_mb_scale_original_tensor(dtype=torch.float32):
+    return make_deterministic_tensor((16, 16, 16), dtype=dtype)
+
+
+def assert_tensor_matches_expected_shard(test_case, original_tensor, received_tensor,
+                                         split_dim, target_tp_size, target_rank,
+                                         context):
+    expected_tensor = expected_target_shard(
+        original_tensor, split_dim, target_tp_size, target_rank
+    )
+    test_case.assertIsNotNone(received_tensor, f"missing tensor for {context}")
+    test_case.assertTrue(
+        torch.equal(received_tensor, expected_tensor),
+        f"mismatch for {context}",
+    )
+
+
 def tearDownModule():
     """Executed once after all tests in this file: closes the global connection."""
     global GLOBAL_STORE
@@ -174,6 +412,110 @@ class MooncakeTestBase(unittest.TestCase):
 # ==========================================
 
 class TestMooncakeFunctional(MooncakeTestBase):
+    def assert_full_reconstruction_case(self, key, tensor, split_dim, tp_size):
+        rc = put_uniform_full_tensor_with_unified_tp(
+            self.store, key, tensor, tp_size, split_dim
+        )
+        self.assertEqual(rc, 0, f"put_tensor_with_parallelism failed for {key}")
+
+        shard_reads = []
+        expected_shards = tensor.chunk(tp_size, split_dim)
+        for rank, expected_shard in enumerate(expected_shards):
+            wrapper_shard = self.store.get_tensor_with_tp(
+                key, tp_rank=rank, tp_size=tp_size, split_dim=split_dim
+            )
+            self.assertIsNotNone(wrapper_shard, f"wrapper shard missing for rank {rank}")
+            self.assertTrue(
+                torch.equal(wrapper_shard, expected_shard),
+                f"wrapper shard mismatch for rank {rank}",
+            )
+
+            shard_parallelism = build_tp_parallelism(tp_size, split_dim, rank=rank)
+            shard_target = make_read_target("shard", shard_parallelism)
+            unified_shard = self.store.get_tensor_with_parallelism(key, shard_target)
+            self.assertIsNotNone(unified_shard, f"unified shard missing for rank {rank}")
+            self.assertTrue(
+                torch.equal(unified_shard, expected_shard),
+                f"unified shard mismatch for rank {rank}",
+            )
+            self.assertTrue(torch.equal(unified_shard, wrapper_shard))
+            shard_reads.append(unified_shard)
+
+        reconstructed_from_shards = torch.cat(shard_reads, dim=split_dim)
+        self.assertTrue(torch.equal(reconstructed_from_shards, tensor))
+
+        full_parallelism = build_tp_parallelism(tp_size, split_dim, rank=min(tp_size - 1, 1))
+        full_target = make_read_target("full", full_parallelism)
+        full_tensor = self.store.get_tensor_with_parallelism(key, full_target)
+        self.assertIsNotNone(full_tensor)
+        self.assertTrue(torch.equal(full_tensor, tensor))
+
+    def assert_full_into_reconstruction_case(self, key, tensor, split_dim, tp_size, buffer_spacing):
+        buffer = (ctypes.c_ubyte * buffer_spacing)()
+        buffer_ptr = ctypes.addressof(buffer)
+        self.assertEqual(self.store.register_buffer(buffer_ptr, buffer_spacing), 0)
+        try:
+            parallelism = build_tp_parallelism(tp_size, split_dim, rank=min(tp_size - 1, 1))
+            rc = put_uniform_full_tensor_with_unified_tp(
+                self.store, key, tensor, tp_size, split_dim
+            )
+            self.assertEqual(rc, 0)
+
+            target = make_read_target("full", parallelism)
+            full_tensor = self.store.get_tensor_with_parallelism(key, target)
+            reconstructed = self.store.get_tensor_with_parallelism_into(
+                key, buffer_ptr, buffer_spacing, target=target
+            )
+            self.assertIsNotNone(full_tensor)
+            self.assertIsNotNone(reconstructed)
+            self.assertTrue(torch.equal(full_tensor, tensor))
+            self.assertTrue(torch.equal(reconstructed, tensor))
+            self.assertTrue(torch.equal(reconstructed, full_tensor))
+        finally:
+            self.assertEqual(self.store.unregister_buffer(buffer_ptr), 0)
+
+    def assert_full_reconstruction_for_all_read_ranks(self, key, tensor, split_dim, tp_size):
+        rc = put_uniform_full_tensor_with_unified_tp(
+            self.store, key, tensor, tp_size, split_dim
+        )
+        self.assertEqual(rc, 0, f"put_tensor_with_parallelism failed for {key}")
+
+        expected_rank_count = len(tensor.chunk(tp_size, split_dim))
+        for rank in range(tp_size):
+            full_target = make_read_target("full", build_tp_parallelism(tp_size, split_dim, rank=rank))
+            full_tensor = self.store.get_tensor_with_parallelism(key, full_target)
+            self.assertIsNotNone(full_tensor, f"full tensor missing for read rank {rank}")
+            self.assertTrue(torch.equal(full_tensor, tensor), f"full tensor mismatch for read rank {rank}")
+
+            if rank < expected_rank_count:
+                shard_target = make_read_target("shard", build_tp_parallelism(tp_size, split_dim, rank=rank))
+                shard_tensor = self.store.get_tensor_with_parallelism(key, shard_target)
+                expected_shard = chunk_tensor_for_rank(tensor, tp_size, split_dim, rank)
+                self.assertIsNotNone(shard_tensor, f"shard missing for read rank {rank}")
+                self.assertTrue(torch.equal(shard_tensor, expected_shard), f"shard mismatch for read rank {rank}")
+
+    def assert_writer_shard_full_reconstruction_case(self, key, tensor, split_dim, shard_count):
+        for rank in range(shard_count):
+            rc = self.store.put_tensor_with_parallelism(
+                key,
+                tensor,
+                writer_partition=make_writer_partition(rank, shard_count, split_dim),
+            )
+            self.assertEqual(rc, 0, f"writer shard put failed for rank {rank}")
+
+        full_target = make_read_target("full")
+        full_tensor = self.store.get_tensor_with_parallelism(key, full_target)
+        self.assertIsNotNone(full_tensor)
+        self.assertTrue(torch.equal(full_tensor, tensor))
+
+        for rank in range(shard_count):
+            shard_target = make_read_target(
+                "shard", build_tp_parallelism(shard_count, split_dim, rank=rank)
+            )
+            shard_tensor = self.store.get_tensor_with_parallelism(key, shard_target)
+            self.assertIsNotNone(shard_tensor)
+            self.assertTrue(torch.equal(shard_tensor, tensor.chunk(shard_count, split_dim)[rank]))
+
     def test_01_basic_put_get(self):
         """Verify basic put and get functionality."""
         key = "func_test_single"
@@ -569,6 +911,1164 @@ class TestMooncakeFunctional(MooncakeTestBase):
             recon = torch.cat(reconstruction_parts, dim=split_dim)
             self.assertTrue(torch.equal(recon, original), f"Tensor {i} final reconstruction mismatch")
 
+    def test_12_unified_parallelism_as_stored_round_trip(self):
+        require_unified_parallelism_api(self)
+        key = "func_unified_as_stored"
+        tensor = torch.randn(16, 16, dtype=torch.float32)
+
+        rc = self.store.put_tensor_with_parallelism(key, tensor)
+        self.assertEqual(rc, 0, f"put_tensor_with_parallelism failed with rc={rc}")
+
+        retrieved_default = self.store.get_tensor_with_parallelism(key)
+        self.assertIsNotNone(retrieved_default)
+        self.assertTrue(torch.equal(retrieved_default, tensor))
+
+        target = make_read_target("as_stored")
+        retrieved_target = self.store.get_tensor_with_parallelism(key, target)
+        self.assertIsNotNone(retrieved_target)
+        self.assertTrue(torch.equal(retrieved_target, tensor))
+
+    def test_13_unified_parallelism_tp_matches_wrapper(self):
+        require_unified_parallelism_api(self)
+        key = "func_unified_tp_single"
+        tensor = torch.arange(24, dtype=torch.float32).view(4, 6).contiguous()
+        tp_size = 3
+        split_dim = 1
+
+        parallelism = make_tensor_parallelism([
+            make_parallel_axis("tp", rank=0, size=tp_size, split_dim=split_dim)
+        ])
+
+        shard = chunk_tensor_for_rank(tensor, tp_size, split_dim, 0).contiguous()
+        rc = self.store.put_tensor_with_parallelism(key, shard, parallelism)
+        self.assertEqual(rc, 0, f"put_tensor_with_parallelism failed with rc={rc}")
+
+        expected = shard
+        rank = 0
+        wrapper_shard = self.store.get_tensor_with_tp(
+            key, tp_rank=rank, tp_size=tp_size, split_dim=split_dim
+        )
+        self.assertIsNotNone(wrapper_shard)
+        self.assertTrue(torch.equal(wrapper_shard, expected))
+
+        shard_parallelism = make_tensor_parallelism([
+            make_parallel_axis("tp", rank=rank, size=tp_size, split_dim=split_dim)
+        ])
+
+        target = make_read_target("shard", shard_parallelism)
+        unified_shard = self.store.get_tensor_with_parallelism(key, target)
+        self.assertIsNotNone(unified_shard)
+        self.assertTrue(torch.equal(unified_shard, expected))
+        self.assertTrue(torch.equal(unified_shard, wrapper_shard))
+
+        for rank in range(1, tp_size):
+            wrapper_shard = self.store.get_tensor_with_tp(
+                key, tp_rank=rank, tp_size=tp_size, split_dim=split_dim
+            )
+            self.assertIsNone(wrapper_shard)
+
+            shard_parallelism = make_tensor_parallelism([
+                make_parallel_axis("tp", rank=rank, size=tp_size, split_dim=split_dim)
+            ])
+
+            target = make_read_target("shard", shard_parallelism)
+            unified_shard = self.store.get_tensor_with_parallelism(key, target)
+            self.assertIsNone(unified_shard)
+
+    def test_14_batch_unified_parallelism_tp_matches_wrapper(self):
+        require_unified_parallelism_api(self)
+        tp_size = 2
+        split_dim = 0
+        num_tensors = 2
+        keys, tensors = generate_tensors(num_tensors, 1)
+
+        results = batch_put_full_tensors_with_unified_tp(
+            self.store, keys, tensors, tp_size, split_dim
+        )
+        self.assertTrue(all(r == 0 for r in results), f"Batch unified put failed. Results: {results}")
+
+        parallelism = make_tensor_parallelism([
+            make_parallel_axis("tp", rank=0, size=tp_size, split_dim=split_dim)
+        ])
+        invalid_results = self.store.batch_put_tensor_with_parallelism(
+            keys,
+            tensors,
+            [parallelism for _ in keys],
+            writer_partitions=[make_writer_partition(0, 2, split_dim) for _ in keys],
+        )
+        self.assertTrue(all(r != 0 for r in invalid_results), f"expected invalid combined routing request: {invalid_results}")
+
+        for rank in range(tp_size):
+            wrapper_shards = self.store.batch_get_tensor_with_tp(
+                keys, tp_rank=rank, tp_size=tp_size
+            )
+
+            shard_parallelism = make_tensor_parallelism([
+                make_parallel_axis("tp", rank=rank, size=tp_size, split_dim=split_dim)
+            ])
+
+            target = make_read_target("shard", shard_parallelism)
+            unified_shards = self.store.batch_get_tensor_with_parallelism(
+                keys, [target for _ in keys]
+            )
+
+            self.assertEqual(len(wrapper_shards), num_tensors)
+            self.assertEqual(len(unified_shards), num_tensors)
+            for i in range(num_tensors):
+                expected = tensors[i].chunk(tp_size, split_dim)[rank]
+                self.assertTrue(torch.equal(wrapper_shards[i], expected))
+                self.assertTrue(torch.equal(unified_shards[i], expected))
+                self.assertTrue(torch.equal(unified_shards[i], wrapper_shards[i]))
+
+    def test_15_unified_parallelism_multi_axis_shard_round_trip(self):
+        require_unified_parallelism_api(self)
+        key = "func_unified_dp_tp_shard"
+        tensor = make_deterministic_tensor((4, 6))
+        parallelism = build_dp_tp_parallelism(
+            dp_rank=1, dp_size=2, tp_rank=0, tp_size=3, split_dim=1
+        )
+
+        rc = put_full_tensor_with_parallelism_template(
+            self.store, key, tensor, parallelism
+        )
+        self.assertEqual(rc, 0)
+
+        target = make_read_target("shard", parallelism)
+        result = self.store.get_tensor_with_parallelism(key, target)
+        self.assertIsNotNone(result)
+        self.assertTrue(torch.equal(result, tensor.chunk(3, 1)[0]))
+
+    def test_16_unified_parallelism_multi_axis_into_round_trip(self):
+        require_unified_parallelism_api(self)
+        key = "func_unified_pp_tp_into"
+        tensor = make_deterministic_tensor((6, 8))
+        parallelism = build_pp_tp_parallelism(
+            pp_rank=0, pp_size=2, stage_id=7, tp_rank=1, tp_size=2, split_dim=0
+        )
+        buffer_spacing = 1 * 1024 * 1024
+        buffer = (ctypes.c_ubyte * buffer_spacing)()
+        buffer_ptr = ctypes.addressof(buffer)
+        self.assertEqual(self.store.register_buffer(buffer_ptr, buffer_spacing), 0)
+        try:
+            rc = put_full_tensor_with_parallelism_template(
+                self.store, key, tensor, parallelism
+            )
+            self.assertEqual(rc, 0)
+            target = make_read_target("shard", parallelism)
+            result = self.store.get_tensor_with_parallelism_into(
+                key, buffer_ptr, buffer_spacing, target=target
+            )
+            self.assertIsNotNone(result)
+            self.assertTrue(torch.equal(result, tensor.chunk(2, 0)[1]))
+        finally:
+            self.assertEqual(self.store.unregister_buffer(buffer_ptr), 0)
+
+    def test_17_unified_parallelism_rejects_invalid_axis_fields(self):
+        require_unified_parallelism_api(self)
+        key = "func_unified_invalid_axis"
+        tensor = torch.randn(8, 8, dtype=torch.float32)
+
+        invalid_dp = make_tensor_parallelism([
+            make_parallel_axis("dp", rank=0, size=2, split_dim=0)
+        ])
+        self.assertNotEqual(
+            self.store.put_tensor_with_parallelism(key, tensor, invalid_dp), 0
+        )
+
+        invalid_ep = make_tensor_parallelism([
+            make_parallel_axis("ep", rank=0, size=2)
+        ])
+        self.assertNotEqual(
+            self.store.put_tensor_with_parallelism(key, tensor, invalid_ep), 0
+        )
+
+    def test_18_batch_unified_parallelism_multi_axis_round_trip(self):
+        require_unified_parallelism_api(self)
+        cases = [
+            (
+                "func_unified_batch_dp_tp_0",
+                make_deterministic_tensor((4, 6)),
+                build_dp_tp_parallelism(0, 2, 1, 3, 1),
+                lambda tensor: tensor.chunk(3, 1)[1],
+            ),
+            (
+                "func_unified_batch_pp_tp_1",
+                make_deterministic_tensor((6, 8)),
+                build_pp_tp_parallelism(1, 2, 9, 0, 2, 0),
+                lambda tensor: tensor.chunk(2, 0)[0],
+            ),
+        ]
+        keys = [key for key, _, _, _ in cases]
+        tensors = [tensor for _, tensor, _, _ in cases]
+        parallelisms = [parallelism for _, _, parallelism, _ in cases]
+        results = [
+            put_full_tensor_with_parallelism_template(
+                self.store, key, tensor, parallelism
+            )
+            for key, tensor, parallelism in zip(keys, tensors, parallelisms)
+        ]
+        self.assertTrue(all(r == 0 for r in results), f"batch put failed: {results}")
+
+        targets = [make_read_target("shard", parallelism) for parallelism in parallelisms]
+        shard_reads = self.store.batch_get_tensor_with_parallelism(keys, targets)
+        self.assertEqual(len(shard_reads), len(cases))
+        for i, (_, tensor, _, expected_fn) in enumerate(cases):
+            self.assertIsNotNone(shard_reads[i], f"batch shard missing for tensor {i}")
+            self.assertTrue(torch.equal(shard_reads[i], expected_fn(tensor)))
+
+    def test_19_batch_unified_parallelism_multi_axis_into_round_trip(self):
+        require_unified_parallelism_api(self)
+        cases = [
+            (
+                "func_unified_batch_into_dp_tp_0",
+                make_deterministic_tensor((4, 9)),
+                build_dp_tp_parallelism(1, 2, 2, 3, 1),
+                lambda tensor: tensor.chunk(3, 1)[2],
+            ),
+            (
+                "func_unified_batch_into_pp_tp_1",
+                make_deterministic_tensor((8, 6)),
+                build_pp_tp_parallelism(0, 2, 5, 1, 2, 0),
+                lambda tensor: tensor.chunk(2, 0)[1],
+            ),
+        ]
+        keys = [key for key, _, _, _ in cases]
+        tensors = [tensor for _, tensor, _, _ in cases]
+        parallelisms = [parallelism for _, _, parallelism, _ in cases]
+        results = [
+            put_full_tensor_with_parallelism_template(
+                self.store, key, tensor, parallelism
+            )
+            for key, tensor, parallelism in zip(keys, tensors, parallelisms)
+        ]
+        self.assertTrue(all(r == 0 for r in results), f"batch put failed: {results}")
+
+        buffer_sizes = [max(serialized_tensor_size(tensor) + 4096, 1 * 1024 * 1024) for tensor in tensors]
+        total_buffer_size = sum(buffer_sizes)
+        backing_buffer = (ctypes.c_ubyte * total_buffer_size)()
+        backing_buffer_ptr = ctypes.addressof(backing_buffer)
+        buffer_ptrs = []
+        offset = 0
+        for size in buffer_sizes:
+            buffer_ptrs.append(backing_buffer_ptr + offset)
+            offset += size
+
+        targets = [make_read_target("shard", parallelism) for parallelism in parallelisms]
+        self.assertEqual(self.store.register_buffer(backing_buffer_ptr, total_buffer_size), 0)
+        try:
+            shard_reads = self.store.batch_get_tensor_with_parallelism_into(
+                keys, buffer_ptrs, buffer_sizes, targets
+            )
+            self.assertEqual(len(shard_reads), len(cases))
+            for i, (_, tensor, _, expected_fn) in enumerate(cases):
+                self.assertIsNotNone(shard_reads[i], f"batch shard into missing for tensor {i}")
+                self.assertTrue(torch.equal(shard_reads[i], expected_fn(tensor)))
+        finally:
+            self.assertEqual(self.store.unregister_buffer(backing_buffer_ptr), 0)
+
+    def test_20_unified_parallelism_full_reconstructs_tensor(self):
+        require_unified_parallelism_api(self)
+        key = "func_unified_tp_full"
+        tensor = torch.arange(32, dtype=torch.float32).view(4, 8).contiguous()
+        tp_size = 4
+        split_dim = 1
+
+        rc = put_uniform_full_tensor_with_unified_tp(
+            self.store, key, tensor, tp_size, split_dim
+        )
+        self.assertEqual(rc, 0)
+
+        target = make_read_target("full", build_tp_parallelism(tp_size, split_dim, rank=0))
+        unified_full = self.store.get_tensor_with_parallelism(key, target)
+        self.assertIsNotNone(unified_full)
+        self.assertTrue(torch.equal(unified_full, tensor))
+
+    def test_21_unified_parallelism_multi_axis_full_reconstructs_tensor(self):
+        require_unified_parallelism_api(self)
+        key = "func_unified_multi_axis_full"
+        tensor = make_deterministic_tensor((6, 8))
+        parallelism = build_dp_tp_parallelism(
+            dp_rank=1, dp_size=2, tp_rank=0, tp_size=4, split_dim=1
+        )
+
+        rc = put_full_tensor_with_parallelism_template(
+            self.store, key, tensor, parallelism
+        )
+        self.assertEqual(rc, 0)
+
+        full_target = make_read_target(
+            "full",
+            build_dp_tp_parallelism(
+                dp_rank=1, dp_size=2, tp_rank=2, tp_size=4, split_dim=1
+            ),
+        )
+        result = self.store.get_tensor_with_parallelism(key, full_target)
+        self.assertIsNotNone(result)
+        self.assertTrue(torch.equal(result, tensor))
+
+    def test_22_unified_parallelism_multi_axis_full_into_reconstructs_tensor(self):
+        require_unified_parallelism_api(self)
+        key = "func_unified_multi_axis_full_into"
+        tensor = make_deterministic_tensor((8, 6))
+        parallelism = build_pp_tp_parallelism(
+            pp_rank=0, pp_size=2, stage_id=4, tp_rank=0, tp_size=2, split_dim=0
+        )
+        buffer_spacing = 1 * 1024 * 1024
+        buffer = (ctypes.c_ubyte * buffer_spacing)()
+        buffer_ptr = ctypes.addressof(buffer)
+        self.assertEqual(self.store.register_buffer(buffer_ptr, buffer_spacing), 0)
+        try:
+            rc = put_full_tensor_with_parallelism_template(
+                self.store, key, tensor, parallelism
+            )
+            self.assertEqual(rc, 0)
+            full_target = make_read_target(
+                "full",
+                build_pp_tp_parallelism(
+                    pp_rank=0, pp_size=2, stage_id=4, tp_rank=1, tp_size=2, split_dim=0
+                ),
+            )
+            result = self.store.get_tensor_with_parallelism_into(
+                key, buffer_ptr, buffer_spacing, target=full_target
+            )
+            self.assertIsNotNone(result)
+            self.assertTrue(torch.equal(result, tensor))
+        finally:
+            self.assertEqual(self.store.unregister_buffer(buffer_ptr), 0)
+
+    def test_23_unified_parallelism_into_matches_wrapper(self):
+        require_unified_parallelism_api(self)
+        key = "func_unified_tp_into"
+        tensor = torch.randn(16, 16, dtype=torch.float32)
+        tp_size = 2
+        split_dim = 0
+        buffer_spacing = 1 * 1024 * 1024
+        total_buffer_size = buffer_spacing * 2
+
+        large_buffer = (ctypes.c_ubyte * total_buffer_size)()
+        large_buffer_ptr = ctypes.addressof(large_buffer)
+        wrapper_ptr = large_buffer_ptr
+        unified_ptr = large_buffer_ptr + buffer_spacing
+
+        self.assertEqual(self.store.register_buffer(large_buffer_ptr, total_buffer_size), 0)
+        try:
+            parallelism = make_tensor_parallelism([
+                make_parallel_axis("tp", rank=0, size=tp_size, split_dim=split_dim)
+            ])
+            shard = chunk_tensor_for_rank(tensor, tp_size, split_dim, 0).contiguous()
+            rc = self.store.put_tensor_with_parallelism(key, shard, parallelism)
+            self.assertEqual(rc, 0)
+
+            target = make_read_target("shard", parallelism)
+            wrapper_tensor = self.store.get_tensor_with_tp_into(
+                key, wrapper_ptr, buffer_spacing, tp_rank=0, tp_size=tp_size, split_dim=split_dim
+            )
+            unified_tensor = self.store.get_tensor_with_parallelism_into(
+                key, unified_ptr, buffer_spacing, target=target
+            )
+            self.assertIsNotNone(wrapper_tensor)
+            self.assertIsNotNone(unified_tensor)
+            self.assertTrue(torch.equal(unified_tensor, wrapper_tensor))
+        finally:
+            self.assertEqual(self.store.unregister_buffer(large_buffer_ptr), 0)
+
+    def test_24_unified_parallelism_from_matches_wrapper(self):
+        require_unified_parallelism_api(self)
+        key = "func_unified_tp_from"
+        seed_key = "func_unified_tp_from_seed"
+        tensor = torch.randn(16, 16, dtype=torch.float32)
+        tp_size = 2
+        split_dim = 1
+        buffer_spacing = 1 * 1024 * 1024
+
+        buffer = (ctypes.c_ubyte * buffer_spacing)()
+        buffer_ptr = ctypes.addressof(buffer)
+        self.assertEqual(self.store.register_buffer(buffer_ptr, buffer_spacing), 0)
+        try:
+            rc = self.store.put_tensor(seed_key, tensor)
+            self.assertEqual(rc, 0)
+            full_tensor = self.store.get_tensor_into(seed_key, buffer_ptr, buffer_spacing)
+            self.assertIsNotNone(full_tensor)
+            full_size = serialized_tensor_size(full_tensor)
+
+            parallelism = make_tensor_parallelism([
+                make_parallel_axis("tp", rank=0, size=tp_size, split_dim=split_dim)
+            ])
+            rc = self.store.put_tensor_with_parallelism_from(
+                key, buffer_ptr, full_size, parallelism
+            )
+            self.assertEqual(rc, 0)
+
+            target = make_read_target("shard", parallelism)
+            unified_shard = self.store.get_tensor_with_parallelism(key, target)
+            wrapper_shard = self.store.get_tensor_with_tp(
+                key, tp_rank=0, tp_size=tp_size, split_dim=split_dim
+            )
+            self.assertIsNotNone(unified_shard)
+            self.assertTrue(torch.equal(unified_shard, wrapper_shard))
+        finally:
+            self.assertEqual(self.store.unregister_buffer(buffer_ptr), 0)
+
+    def test_25_unified_parallelism_full_into_reconstructs_tensor(self):
+        require_unified_parallelism_api(self)
+        key = "func_unified_tp_full_into"
+        tensor = torch.arange(4 * 9 * 6, dtype=torch.float32).view(4, 9, 6).contiguous()
+        self.assert_full_into_reconstruction_case(
+            key=key,
+            tensor=tensor,
+            split_dim=1,
+            tp_size=3,
+            buffer_spacing=1 * 1024 * 1024,
+        )
+
+    def test_26_unified_parallelism_full_reconstruction_matrix(self):
+        require_unified_parallelism_api(self)
+        cases = [
+            ((8, 12), 0, 4),
+            ((8, 12), 0, 8),
+            ((8, 12), 0, 16),
+            ((10, 12), 0, 4),
+            ((12, 9), 1, 3),
+            ((12, 9), 1, 9),
+            ((12, 9), 1, 16),
+            ((12, 10), 1, 4),
+            ((4, 10, 6), 1, 4),
+            ((4, 5, 7), 2, 8),
+        ]
+
+        for index, (shape, split_dim, tp_size) in enumerate(cases):
+            with self.subTest(case=reconstruction_case_name(shape, split_dim, tp_size)):
+                tensor = make_deterministic_tensor(shape)
+                key = f"func_unified_tp_matrix_{index}"
+                if is_uniform_tp_case(shape, split_dim, tp_size):
+                    self.assert_full_reconstruction_case(
+                        key=key,
+                        tensor=tensor,
+                        split_dim=split_dim,
+                        tp_size=tp_size,
+                    )
+                else:
+                    rc = put_uniform_full_tensor_with_unified_tp(
+                        self.store, key, tensor, tp_size, split_dim
+                    )
+                    self.assertNotEqual(rc, 0)
+
+    def test_27_unified_parallelism_full_into_matrix(self):
+        require_unified_parallelism_api(self)
+        cases = [
+            ((8, 12), 0, 16),
+            ((12, 10), 1, 4),
+            ((4, 10, 6), 1, 4),
+            ((4, 5, 7), 2, 8),
+        ]
+
+        for index, (shape, split_dim, tp_size) in enumerate(cases):
+            with self.subTest(case=reconstruction_case_name(shape, split_dim, tp_size)):
+                tensor = make_deterministic_tensor(shape)
+                key = f"func_unified_tp_full_into_matrix_{index}"
+                if is_uniform_tp_case(shape, split_dim, tp_size):
+                    buffer_spacing = max(serialized_tensor_size(tensor) + 4096, 1 * 1024 * 1024)
+                    self.assert_full_into_reconstruction_case(
+                        key=key,
+                        tensor=tensor,
+                        split_dim=split_dim,
+                        tp_size=tp_size,
+                        buffer_spacing=buffer_spacing,
+                    )
+                else:
+                    rc = put_uniform_full_tensor_with_unified_tp(
+                        self.store, key, tensor, tp_size, split_dim
+                    )
+                    self.assertNotEqual(rc, 0)
+
+    def test_27b_unified_parallelism_full_into_formula_all_split_dims_and_hetero_tp(self):
+        require_unified_parallelism_api(self)
+        cases = [
+            ((8, 12, 16), 0, 4, [2, 4, 8]),
+            ((8, 12, 16), 1, 4, [2, 3, 4, 6]),
+            ((8, 12, 16), 2, 4, [2, 4, 8, 16]),
+        ]
+        for case_index, (shape, split_dim, write_tp_size, read_tp_sizes) in enumerate(cases):
+            tensor = make_deterministic_tensor(shape)
+            for read_tp_size in read_tp_sizes:
+                with self.subTest(
+                    case=f"split_dim={split_dim},write_tp={write_tp_size},read_tp={read_tp_size}"
+                ):
+                    key = f"func_unified_formula_split_{case_index}_{read_tp_size}"
+                    self.assertEqual(shape[split_dim] % read_tp_size, 0)
+                    rc = put_uniform_full_tensor_with_unified_tp(
+                        self.store, key, tensor, write_tp_size, split_dim
+                    )
+                    self.assertEqual(rc, 0)
+                    target = make_read_target(
+                        "full",
+                        build_tp_parallelism(
+                            read_tp_size, split_dim, rank=read_tp_size - 1
+                        ),
+                    )
+                    full_result = self.store.get_tensor_with_parallelism(
+                        key, target
+                    )
+                    self.assertIsNotNone(full_result)
+                    self.assertTrue(torch.equal(full_result, tensor))
+
+                    buffer_spacing = max(
+                        serialized_tensor_size(tensor) + 4096, 1 * 1024 * 1024
+                    )
+                    buffer = (ctypes.c_ubyte * buffer_spacing)()
+                    buffer_ptr = ctypes.addressof(buffer)
+                    self.assertEqual(
+                        self.store.register_buffer(buffer_ptr, buffer_spacing), 0
+                    )
+                    try:
+                        into_result = self.store.get_tensor_with_parallelism_into(
+                            key, buffer_ptr, buffer_spacing, target=target
+                        )
+                        self.assertIsNotNone(into_result)
+                        self.assertTrue(torch.equal(into_result, tensor))
+                    finally:
+                        self.assertEqual(self.store.unregister_buffer(buffer_ptr), 0)
+
+    def test_28_writer_shard_full_reconstruction(self):
+        require_unified_parallelism_api(self)
+        key = "func_writer_shard_full"
+        tensor = make_deterministic_tensor((6, 8))
+        self.assert_writer_shard_full_reconstruction_case(
+            key=key, tensor=tensor, split_dim=1, shard_count=4
+        )
+
+    def test_29_writer_shard_full_into_reconstruction(self):
+        require_unified_parallelism_api(self)
+        key = "func_writer_shard_full_into"
+        tensor = make_deterministic_tensor((8, 6))
+        buffer_spacing = max(serialized_tensor_size(tensor) + 4096, 1 * 1024 * 1024)
+        buffer = (ctypes.c_ubyte * buffer_spacing)()
+        buffer_ptr = ctypes.addressof(buffer)
+        self.assertEqual(self.store.register_buffer(buffer_ptr, buffer_spacing), 0)
+        try:
+            for rank in range(3):
+                rc = self.store.put_tensor_with_parallelism(
+                    key,
+                    tensor,
+                    writer_partition=make_writer_partition(rank, 3, 1),
+                )
+                self.assertEqual(rc, 0)
+
+            result = self.store.get_tensor_with_parallelism_into(
+                key, buffer_ptr, buffer_spacing, target=make_read_target("full")
+            )
+            self.assertIsNotNone(result)
+            self.assertTrue(torch.equal(result, tensor))
+        finally:
+            self.assertEqual(self.store.unregister_buffer(buffer_ptr), 0)
+
+    def test_30_batch_writer_partition_full_reconstruction(self):
+        require_unified_parallelism_api(self)
+        tensors = [
+            make_deterministic_tensor((8, 12)),
+            make_deterministic_tensor((6, 9, 4)),
+            make_deterministic_tensor((12, 12)),
+        ]
+        keys = [f"func_writer_partition_batch_{i}" for i in range(len(tensors))]
+        writer_partitions = [
+            make_writer_partition(rank=0, size=4, split_dim=1),
+            make_writer_partition(rank=1, size=3, split_dim=0),
+            make_writer_partition(rank=2, size=4, split_dim=1),
+        ]
+
+        for rank in range(4):
+            batch_writer_partitions = [
+                make_writer_partition(rank=min(rank, part.size - 1), size=part.size, split_dim=part.split_dim)
+                for part in writer_partitions
+            ]
+            results = self.store.batch_put_tensor_with_parallelism(
+                keys,
+                tensors,
+                writer_partitions=batch_writer_partitions,
+            )
+            self.assertTrue(all(r == 0 for r in results), f"batch put failed: {results}")
+
+        full_targets = [make_read_target("full") for _ in tensors]
+        full_reads = self.store.batch_get_tensor_with_parallelism(keys, full_targets)
+        self.assertEqual(len(full_reads), len(tensors))
+        for i, tensor in enumerate(tensors):
+            self.assertTrue(torch.equal(full_reads[i], tensor), f"full reconstruction mismatch for tensor {i}")
+
+    def test_31_batch_writer_partition_into_round_trip(self):
+        require_unified_parallelism_api(self)
+        tensors = [
+            make_deterministic_tensor((7, 9)),
+            make_deterministic_tensor((8, 6)),
+        ]
+        keys = [f"func_writer_partition_batch_into_{i}" for i in range(len(tensors))]
+        writer_layouts = [
+            make_writer_partition(rank=0, size=3, split_dim=1),
+            make_writer_partition(rank=0, size=2, split_dim=0),
+        ]
+        for rank in range(3):
+            batch_writer_partitions = [
+                make_writer_partition(rank=min(rank, layout.size - 1), size=layout.size, split_dim=layout.split_dim)
+                for layout in writer_layouts
+            ]
+            results = self.store.batch_put_tensor_with_parallelism(
+                keys,
+                tensors,
+                writer_partitions=batch_writer_partitions,
+            )
+            self.assertTrue(all(r == 0 for r in results), f"batch put failed: {results}")
+
+        buffer_sizes = [max(serialized_tensor_size(tensor) + 4096, 1 * 1024 * 1024) for tensor in tensors]
+        total_buffer_size = sum(buffer_sizes)
+        backing_buffer = (ctypes.c_ubyte * total_buffer_size)()
+        backing_buffer_ptr = ctypes.addressof(backing_buffer)
+        buffer_ptrs = []
+        offset = 0
+        for size in buffer_sizes:
+            buffer_ptrs.append(backing_buffer_ptr + offset)
+            offset += size
+
+        self.assertEqual(self.store.register_buffer(backing_buffer_ptr, total_buffer_size), 0)
+        try:
+            full_reads = self.store.batch_get_tensor_with_parallelism_into(
+                keys, buffer_ptrs, buffer_sizes, [make_read_target("full") for _ in tensors]
+            )
+            self.assertEqual(len(full_reads), len(tensors))
+            for i, tensor in enumerate(tensors):
+                self.assertIsNotNone(full_reads[i], f"batch full into missing for tensor {i}")
+                self.assertTrue(torch.equal(full_reads[i], tensor), f"batch full into mismatch for tensor {i}")
+        finally:
+            self.assertEqual(self.store.unregister_buffer(backing_buffer_ptr), 0)
+
+    def test_32_batch_unified_parallelism_full_reconstruction(self):
+        require_unified_parallelism_api(self)
+        tp_size = 4
+        split_dim = 1
+        tensors = [
+            make_deterministic_tensor((8, 12)),
+            make_deterministic_tensor((6, 12, 4)),
+            make_deterministic_tensor((12, 12)),
+        ]
+        keys = [f"func_unified_tp_full_batch_{i}" for i in range(len(tensors))]
+        parallelisms = [build_tp_parallelism(tp_size, split_dim, rank=0) for _ in tensors]
+        results = batch_put_full_tensors_with_unified_tp(self.store, keys, tensors, tp_size, split_dim)
+        self.assertTrue(all(r == 0 for r in results), f"batch put failed: {results}")
+
+        full_targets = [make_read_target("full", build_tp_parallelism(tp_size, split_dim, rank=1)) for _ in tensors]
+        full_reads = self.store.batch_get_tensor_with_parallelism(keys, full_targets)
+        self.assertEqual(len(full_reads), len(tensors))
+        for i, tensor in enumerate(tensors):
+            self.assertTrue(torch.equal(full_reads[i], tensor), f"full reconstruction mismatch for tensor {i}")
+
+        shard_targets = [make_read_target("shard", build_tp_parallelism(tp_size, split_dim, rank=2)) for _ in tensors]
+        shard_reads = self.store.batch_get_tensor_with_parallelism(keys, shard_targets)
+        self.assertEqual(len(shard_reads), len(tensors))
+        for i, tensor in enumerate(tensors):
+            expected_shard = chunk_tensor_for_rank(tensor, tp_size, split_dim, 2)
+            self.assertTrue(torch.equal(shard_reads[i], expected_shard), f"shard mismatch for tensor {i}")
+
+    def test_31_batch_unified_parallelism_multi_axis_full_reconstruction(self):
+        require_unified_parallelism_api(self)
+        cases = [
+            (
+                "func_unified_batch_multi_full_0",
+                make_deterministic_tensor((6, 8)),
+                build_dp_tp_parallelism(1, 2, 0, 4, 1),
+                build_dp_tp_parallelism(1, 2, 2, 4, 1),
+            ),
+            (
+                "func_unified_batch_multi_full_1",
+                make_deterministic_tensor((8, 6)),
+                build_pp_tp_parallelism(0, 2, 6, 0, 2, 0),
+                build_pp_tp_parallelism(0, 2, 6, 1, 2, 0),
+            ),
+        ]
+        keys = [key for key, _, _, _ in cases]
+        tensors = [tensor for _, tensor, _, _ in cases]
+        write_parallelisms = [parallelism for _, _, parallelism, _ in cases]
+        results = [
+            put_full_tensor_with_parallelism_template(
+                self.store, key, tensor, parallelism
+            )
+            for key, tensor, parallelism in zip(keys, tensors, write_parallelisms)
+        ]
+        self.assertTrue(all(r == 0 for r in results), f"batch put failed: {results}")
+
+        full_targets = [make_read_target("full", target_parallelism) for _, _, _, target_parallelism in cases]
+        full_reads = self.store.batch_get_tensor_with_parallelism(keys, full_targets)
+        self.assertEqual(len(full_reads), len(cases))
+        for i, tensor in enumerate(tensors):
+            self.assertIsNotNone(full_reads[i], f"multi-axis full missing for tensor {i}")
+            self.assertTrue(torch.equal(full_reads[i], tensor), f"multi-axis full mismatch for tensor {i}")
+
+    def test_32_unified_parallelism_multi_axis_full_remaps_axis_order(self):
+        require_unified_parallelism_api(self)
+        key = "func_unified_multi_axis_full_remap_order"
+        tensor = make_deterministic_tensor((6, 8))
+        write_parallelism = build_dp_tp_parallelism(
+            dp_rank=1, dp_size=2, tp_rank=0, tp_size=4, split_dim=1
+        )
+        read_parallelism = reorder_parallelism_axes(write_parallelism, ["tp", "dp"])
+        read_parallelism.axes[0].rank = 3
+
+        rc = put_full_tensor_with_parallelism_template(
+            self.store, key, tensor, write_parallelism
+        )
+        self.assertEqual(rc, 0)
+
+        full_target = make_read_target("full", read_parallelism)
+        result = self.store.get_tensor_with_parallelism(key, full_target)
+        self.assertIsNotNone(result)
+        self.assertTrue(torch.equal(result, tensor))
+
+    def test_33_unified_parallelism_multi_axis_full_into_remaps_axis_order(self):
+        require_unified_parallelism_api(self)
+        key = "func_unified_multi_axis_full_into_remap_order"
+        tensor = make_deterministic_tensor((8, 6))
+        write_parallelism = build_pp_tp_parallelism(
+            pp_rank=0, pp_size=2, stage_id=4, tp_rank=0, tp_size=2, split_dim=0
+        )
+        read_parallelism = reorder_parallelism_axes(write_parallelism, ["tp", "pp"])
+        read_parallelism.axes[0].rank = 1
+        buffer_spacing = 1 * 1024 * 1024
+        buffer = (ctypes.c_ubyte * buffer_spacing)()
+        buffer_ptr = ctypes.addressof(buffer)
+        self.assertEqual(self.store.register_buffer(buffer_ptr, buffer_spacing), 0)
+        try:
+            rc = put_full_tensor_with_parallelism_template(
+                self.store, key, tensor, write_parallelism
+            )
+            self.assertEqual(rc, 0)
+            full_target = make_read_target("full", read_parallelism)
+            result = self.store.get_tensor_with_parallelism_into(
+                key, buffer_ptr, buffer_spacing, target=full_target
+            )
+            self.assertIsNotNone(result)
+            self.assertTrue(torch.equal(result, tensor))
+        finally:
+            self.assertEqual(self.store.unregister_buffer(buffer_ptr), 0)
+
+    def test_34_batch_unified_parallelism_multi_axis_full_remaps_axis_order(self):
+        require_unified_parallelism_api(self)
+        cases = [
+            (
+                "func_unified_batch_multi_full_remap_0",
+                make_deterministic_tensor((6, 8)),
+                build_dp_tp_parallelism(1, 2, 0, 4, 1),
+                reorder_parallelism_axes(build_dp_tp_parallelism(1, 2, 2, 4, 1), ["tp", "dp"]),
+            ),
+            (
+                "func_unified_batch_multi_full_remap_1",
+                make_deterministic_tensor((8, 6)),
+                build_pp_tp_parallelism(0, 2, 6, 0, 2, 0),
+                reorder_parallelism_axes(build_pp_tp_parallelism(0, 2, 6, 1, 2, 0), ["tp", "pp"]),
+            ),
+        ]
+        keys = [key for key, _, _, _ in cases]
+        tensors = [tensor for _, tensor, _, _ in cases]
+        write_parallelisms = [parallelism for _, _, parallelism, _ in cases]
+        read_parallelisms = [parallelism for _, _, _, parallelism in cases]
+        results = [
+            put_full_tensor_with_parallelism_template(
+                self.store, key, tensor, parallelism
+            )
+            for key, tensor, parallelism in zip(keys, tensors, write_parallelisms)
+        ]
+        self.assertTrue(all(r == 0 for r in results), f"batch put failed: {results}")
+
+        full_targets = [make_read_target("full", target_parallelism) for target_parallelism in read_parallelisms]
+        full_reads = self.store.batch_get_tensor_with_parallelism(keys, full_targets)
+        self.assertEqual(len(full_reads), len(cases))
+        for i, tensor in enumerate(tensors):
+            self.assertIsNotNone(full_reads[i], f"multi-axis remap full missing for tensor {i}")
+            self.assertTrue(torch.equal(full_reads[i], tensor), f"multi-axis remap full mismatch for tensor {i}")
+
+    def test_35_unified_parallelism_full_large_payload(self):
+        require_unified_parallelism_api(self)
+        key = "func_unified_tp_full_large"
+        tensor = make_deterministic_tensor((32, 32))
+        split_dim = 1
+        tp_size = 4
+        buffer_spacing = serialized_tensor_size(tensor) + 1 * 1024 * 1024
+
+        rc = put_uniform_full_tensor_with_unified_tp(
+            self.store, key, tensor, tp_size, split_dim
+        )
+        self.assertEqual(rc, 0)
+
+        full_target = make_read_target("full", build_tp_parallelism(tp_size, split_dim, rank=1))
+        full_tensor = self.store.get_tensor_with_parallelism(key, full_target)
+        self.assertIsNotNone(full_tensor)
+        self.assertTrue(torch.equal(full_tensor, tensor))
+
+        shard_rank = 3
+        shard_target = make_read_target("shard", build_tp_parallelism(tp_size, split_dim, rank=shard_rank))
+        shard_tensor = self.store.get_tensor_with_parallelism(key, shard_target)
+        expected_shard = chunk_tensor_for_rank(tensor, tp_size, split_dim, shard_rank)
+        self.assertIsNotNone(shard_tensor)
+        self.assertTrue(torch.equal(shard_tensor, expected_shard))
+
+        self.assert_full_into_reconstruction_case(
+            key=f"{key}_into",
+            tensor=tensor,
+            split_dim=split_dim,
+            tp_size=tp_size,
+            buffer_spacing=buffer_spacing,
+        )
+
+    def test_36_unified_parallelism_full_reconstruction_edge_matrix(self):
+        require_unified_parallelism_api(self)
+        cases = [
+            ((7,), 0, 8),
+            ((9,), 0, 4),
+            ((3, 5, 7), 2, 16),
+            ((5, 11, 3), 1, 8),
+        ]
+
+        for index, (shape, split_dim, tp_size) in enumerate(cases):
+            with self.subTest(case=reconstruction_case_name(shape, split_dim, tp_size)):
+                tensor = make_deterministic_tensor(shape)
+                key = f"func_unified_tp_edge_matrix_{index}"
+                rc = put_uniform_full_tensor_with_unified_tp(
+                    self.store, key, tensor, tp_size, split_dim
+                )
+                self.assertNotEqual(rc, 0)
+
+    def test_37_batch_unified_parallelism_full_into_matrix(self):
+        require_unified_parallelism_api(self)
+        cases = [
+            (make_deterministic_tensor((7,)), 0, 8),
+            (make_deterministic_tensor((8, 12)), 0, 16),
+            (make_deterministic_tensor((4, 5, 7)), 2, 8),
+        ]
+        for index, (tensor, split_dim, tp_size) in enumerate(cases):
+            with self.subTest(case=reconstruction_case_name(tuple(tensor.shape), split_dim, tp_size)):
+                rc = put_uniform_full_tensor_with_unified_tp(
+                    self.store,
+                    f"func_unified_tp_batch_full_into_{index}",
+                    tensor,
+                    tp_size,
+                    split_dim,
+                )
+                self.assertNotEqual(rc, 0)
+
+    def test_38_batch_unified_parallelism_full_ragged_cases(self):
+        require_unified_parallelism_api(self)
+        cases = [
+            (make_deterministic_tensor((7,)), 0, 8),
+            (make_deterministic_tensor((8, 12)), 0, 16),
+            (make_deterministic_tensor((11, 9)), 1, 4),
+            (make_deterministic_tensor((3, 5, 7)), 2, 16),
+        ]
+        for index, (tensor, split_dim, tp_size) in enumerate(cases):
+            with self.subTest(case=reconstruction_case_name(tuple(tensor.shape), split_dim, tp_size)):
+                rc = put_uniform_full_tensor_with_unified_tp(
+                    self.store,
+                    f"func_unified_tp_batch_ragged_{index}",
+                    tensor,
+                    tp_size,
+                    split_dim,
+                )
+                self.assertNotEqual(rc, 0)
+
+    def test_39_unified_parallelism_full_dtype_smoke(self):
+        require_unified_parallelism_api(self)
+        cases = [
+            (torch.float16, (8, 12), 1, 4),
+            (torch.bfloat16, (9,), 0, 16),
+        ]
+
+        for index, (dtype, shape, split_dim, tp_size) in enumerate(cases):
+            with self.subTest(dtype=str(dtype), case=reconstruction_case_name(shape, split_dim, tp_size)):
+                tensor = make_deterministic_tensor(shape, dtype=dtype)
+                key = f"func_unified_tp_dtype_{index}"
+                if is_uniform_tp_case(shape, split_dim, tp_size):
+                    self.assert_full_reconstruction_case(
+                        key=key,
+                        tensor=tensor,
+                        split_dim=split_dim,
+                        tp_size=tp_size,
+                    )
+                else:
+                    rc = put_uniform_full_tensor_with_unified_tp(
+                        self.store, key, tensor, tp_size, split_dim
+                    )
+                    self.assertNotEqual(rc, 0)
+
+    def test_40_unified_parallelism_target_shard_matrix(self):
+        require_unified_parallelism_api(self)
+        cases = [
+            ((8, 12, 6), 0, 4, 2, "tp"),
+            ((8, 12, 6), 1, 4, 2, "dp_tp"),
+            ((8, 12, 6), 2, 2, 4, "pp_tp"),
+            ((8, 12, 6), 1, 2, 4, "ep_tp"),
+            ((12, 8, 6), 0, 2, 2, "tp"),
+            ((12, 8, 8), 2, 4, 4, "dp_tp"),
+        ]
+
+        for index, (shape, split_dim, put_tp_size, get_tp_size, mode) in enumerate(cases):
+            with self.subTest(shape=shape, split_dim=split_dim, put_tp_size=put_tp_size, get_tp_size=get_tp_size, mode=mode):
+                tensor = make_deterministic_tensor(shape)
+                key = f"func_unified_target_shard_matrix_{index}"
+                rc = put_full_tensor_with_parallelism_mode(
+                    self.store, key, tensor, split_dim, put_tp_size, mode
+                )
+                self.assertEqual(rc, 0)
+
+                if not is_uniform_tp_case(shape, split_dim, get_tp_size):
+                    target_parallelism = build_parallelism_mode(
+                        mode, 0, get_tp_size, split_dim
+                    )
+                    target = make_read_target("shard", target_parallelism)
+                    shard = self.store.get_tensor_with_parallelism(key, target)
+                    self.assertIsNone(shard)
+                    continue
+
+                for target_rank in range(get_tp_size):
+                    target_parallelism = build_parallelism_mode(
+                        mode, target_rank, get_tp_size, split_dim
+                    )
+                    target = make_read_target("shard", target_parallelism)
+                    shard = self.store.get_tensor_with_parallelism(key, target)
+                    expected = expected_target_shard(
+                        tensor, split_dim, get_tp_size, target_rank
+                    )
+                    self.assertIsNotNone(shard)
+                    self.assertTrue(
+                        torch.equal(shard, expected),
+                        f"target shard mismatch for mode={mode}, rank={target_rank}",
+                    )
+
+    def test_41_unified_parallelism_target_shard_into_matrix(self):
+        require_unified_parallelism_api(self)
+        cases = [
+            ((8, 12, 6), 0, 4, 2, "tp"),
+            ((8, 12, 6), 1, 4, 2, "dp_tp"),
+            ((8, 12, 6), 2, 2, 4, "pp_tp"),
+            ((8, 12, 6), 1, 2, 4, "ep_tp"),
+        ]
+
+        for index, (shape, split_dim, put_tp_size, get_tp_size, mode) in enumerate(cases):
+            with self.subTest(shape=shape, split_dim=split_dim, put_tp_size=put_tp_size, get_tp_size=get_tp_size, mode=mode):
+                tensor = make_deterministic_tensor(shape)
+                key = f"func_unified_target_shard_into_matrix_{index}"
+                rc = put_full_tensor_with_parallelism_mode(
+                    self.store, key, tensor, split_dim, put_tp_size, mode
+                )
+                self.assertEqual(rc, 0)
+
+                if not is_uniform_tp_case(shape, split_dim, get_tp_size):
+                    buffer_spacing = 1 * 1024 * 1024
+                    buffer = (ctypes.c_ubyte * buffer_spacing)()
+                    buffer_ptr = ctypes.addressof(buffer)
+                    self.assertEqual(self.store.register_buffer(buffer_ptr, buffer_spacing), 0)
+                    try:
+                        target_parallelism = build_parallelism_mode(
+                            mode, 0, get_tp_size, split_dim
+                        )
+                        target = make_read_target("shard", target_parallelism)
+                        shard = self.store.get_tensor_with_parallelism_into(
+                            key, buffer_ptr, buffer_spacing, target=target
+                        )
+                        self.assertIsNone(shard)
+                    finally:
+                        self.assertEqual(self.store.unregister_buffer(buffer_ptr), 0)
+                    continue
+
+                for target_rank in range(get_tp_size):
+                    expected = expected_target_shard(
+                        tensor, split_dim, get_tp_size, target_rank
+                    )
+                    buffer_spacing = max(serialized_tensor_size(expected) + 4096, 1 * 1024 * 1024)
+                    buffer = (ctypes.c_ubyte * buffer_spacing)()
+                    buffer_ptr = ctypes.addressof(buffer)
+                    self.assertEqual(self.store.register_buffer(buffer_ptr, buffer_spacing), 0)
+                    try:
+                        target_parallelism = build_parallelism_mode(
+                            mode, target_rank, get_tp_size, split_dim
+                        )
+                        target = make_read_target("shard", target_parallelism)
+                        shard = self.store.get_tensor_with_parallelism_into(
+                            key, buffer_ptr, buffer_spacing, target=target
+                        )
+                        self.assertIsNotNone(shard)
+                        self.assertTrue(
+                            torch.equal(shard, expected),
+                            f"target shard into mismatch for mode={mode}, rank={target_rank}",
+                        )
+                    finally:
+                        self.assertEqual(self.store.unregister_buffer(buffer_ptr), 0)
+
+    def test_42_unified_parallelism_mb_scale_target_shard_matrix(self):
+        require_unified_parallelism_api(self)
+        original = make_mb_scale_original_tensor()
+        cases = [
+            (0, 4, 2, "tp"),
+            (1, 4, 2, "dp_tp"),
+            (2, 4, 2, "pp_tp"),
+            (2, 4, 2, "ep_tp"),
+        ]
+
+        for index, (split_dim, put_tp_size, get_tp_size, mode) in enumerate(cases):
+            key = f"func_unified_mb_target_shard_{index}"
+            with self.subTest(split_dim=split_dim, put_tp_size=put_tp_size,
+                              get_tp_size=get_tp_size, mode=mode):
+                rc = put_full_tensor_with_parallelism_mode(
+                    self.store, key, original, split_dim, put_tp_size, mode
+                )
+                self.assertEqual(rc, 0)
+                for target_rank in range(get_tp_size):
+                    target_parallelism = build_parallelism_mode(
+                        mode, target_rank, get_tp_size, split_dim
+                    )
+                    target = make_read_target("shard", target_parallelism)
+                    shard = self.store.get_tensor_with_parallelism(key, target)
+                    assert_tensor_matches_expected_shard(
+                        self, original, shard, split_dim, get_tp_size, target_rank,
+                        context=(
+                            f"mb shard read mode={mode} split_dim={split_dim} "
+                            f"target_rank={target_rank}"
+                        ),
+                    )
+
+    def test_43_unified_parallelism_mb_scale_target_shard_into_matrix(self):
+        require_unified_parallelism_api(self)
+        original = make_mb_scale_original_tensor()
+        cases = [
+            (0, 4, 2, "tp"),
+            (1, 4, 2, "dp_tp"),
+            (2, 4, 2, "pp_tp"),
+            (2, 4, 2, "ep_tp"),
+        ]
+
+        for index, (split_dim, put_tp_size, get_tp_size, mode) in enumerate(cases):
+            key = f"func_unified_mb_target_shard_into_{index}"
+            with self.subTest(split_dim=split_dim, put_tp_size=put_tp_size,
+                              get_tp_size=get_tp_size, mode=mode):
+                rc = put_full_tensor_with_parallelism_mode(
+                    self.store, key, original, split_dim, put_tp_size, mode
+                )
+                self.assertEqual(rc, 0)
+                for target_rank in range(get_tp_size):
+                    expected = expected_target_shard(
+                        original, split_dim, get_tp_size, target_rank
+                    )
+                    buffer_spacing = max(
+                        serialized_tensor_size(expected) + 4096,
+                        1 * 1024 * 1024,
+                    )
+                    buffer = (ctypes.c_ubyte * buffer_spacing)()
+                    buffer_ptr = ctypes.addressof(buffer)
+                    self.assertEqual(
+                        self.store.register_buffer(buffer_ptr, buffer_spacing), 0
+                    )
+                    try:
+                        target_parallelism = build_parallelism_mode(
+                            mode, target_rank, get_tp_size, split_dim
+                        )
+                        target = make_read_target("shard", target_parallelism)
+                        shard = self.store.get_tensor_with_parallelism_into(
+                            key, buffer_ptr, buffer_spacing, target=target
+                        )
+                        assert_tensor_matches_expected_shard(
+                            self, original, shard, split_dim, get_tp_size,
+                            target_rank,
+                            context=(
+                                f"mb shard into mode={mode} split_dim={split_dim} "
+                                f"target_rank={target_rank}"
+                            ),
+                        )
+                    finally:
+                        self.assertEqual(self.store.unregister_buffer(buffer_ptr), 0)
+
+    def test_44_unified_parallelism_concurrency_relation_matrix(self):
+        require_unified_parallelism_api(self)
+        original = make_mb_scale_original_tensor()
+        scenarios = [
+            (2, 4, [(0, 4, 2, "tp"), (1, 4, 2, "dp_tp"), (2, 4, 2, "pp_tp")]),
+            (3, 3, [(0, 4, 2, "tp"), (1, 4, 2, "dp_tp"), (2, 4, 2, "ep_tp")]),
+            (6, 3, [(0, 4, 2, "tp"), (1, 4, 2, "pp_tp"), (2, 4, 2, "ep_tp")]),
+        ]
+
+        def expand_jobs(cases, worker_count):
+            jobs = list(enumerate(cases))
+            if worker_count <= len(cases):
+                return jobs
+            extra_jobs = worker_count - len(cases)
+            for extra_index in range(extra_jobs):
+                case_index = extra_index % len(cases)
+                jobs.append((case_index, cases[case_index]))
+            return jobs
+
+        def put_worker(key, split_dim, put_tp_size, mode):
+            return put_full_tensor_with_parallelism_mode(
+                self.store, key, original, split_dim, put_tp_size, mode
+            )
+
+        def get_worker(key, split_dim, get_tp_size, mode):
+            for target_rank in range(get_tp_size):
+                target_parallelism = build_parallelism_mode(
+                    mode, target_rank, get_tp_size, split_dim
+                )
+                target = make_read_target("shard", target_parallelism)
+                shard = self.store.get_tensor_with_parallelism(key, target)
+                expected = expected_target_shard(
+                    original, split_dim, get_tp_size, target_rank
+                )
+                if shard is None or not torch.equal(shard, expected):
+                    return (
+                        f"concurrency shard mismatch key={key} split_dim={split_dim} "
+                        f"mode={mode} rank={target_rank}"
+                    )
+            return None
+
+        for scenario_index, (put_workers, get_workers, cases) in enumerate(scenarios):
+            with self.subTest(put_workers=put_workers, get_workers=get_workers):
+                put_jobs = [
+                    (
+                        case_index,
+                        f"func_unified_concurrency_{scenario_index}_put_{job_index}_case_{case_index}",
+                        split_dim,
+                        put_tp_size,
+                        get_tp_size,
+                        mode,
+                    )
+                    for job_index, (case_index, (split_dim, put_tp_size, get_tp_size, mode)) in enumerate(
+                        expand_jobs(cases, put_workers)
+                    )
+                ]
+                case_keys = {}
+                for case_index, key, split_dim, put_tp_size, get_tp_size, mode in put_jobs:
+                    case_keys.setdefault(case_index, []).append(
+                        (key, split_dim, put_tp_size, get_tp_size, mode)
+                    )
+
+                get_jobs = []
+                case_read_counts = {}
+                for case_index, (split_dim, put_tp_size, get_tp_size, mode) in expand_jobs(cases, get_workers):
+                    read_index = case_read_counts.get(case_index, 0)
+                    key, _split_dim, _put_tp_size, _get_tp_size, _mode = case_keys[case_index][
+                        read_index % len(case_keys[case_index])
+                    ]
+                    get_jobs.append((key, split_dim, put_tp_size, get_tp_size, mode))
+                    case_read_counts[case_index] = read_index + 1
+
+                with concurrent.futures.ThreadPoolExecutor(max_workers=max(put_workers, get_workers)) as executor:
+                    put_futures = [
+                        executor.submit(put_worker, key, split_dim, put_tp_size, mode)
+                        for _case_index, key, split_dim, put_tp_size, _get_tp_size, mode in put_jobs
+                    ]
+                    for future in concurrent.futures.as_completed(put_futures):
+                        self.assertEqual(future.result(), 0)
+
+                    get_futures = [
+                        executor.submit(get_worker, key, split_dim, get_tp_size, mode)
+                        for key, split_dim, _put_tp_size, get_tp_size, mode in get_jobs
+                    ]
+                    for future in concurrent.futures.as_completed(get_futures):
+                        error = future.result()
+                        if error:
+                            self.fail(error)
+
 # ==========================================
 #  Performance/Benchmark Tests
 # ==========================================
@@ -579,6 +2079,26 @@ class TestMooncakeBenchmark(MooncakeTestBase):
     TENSOR_SIZE_MB = 16
     TOTAL_SIZE_MB = 256
 
+    def _run_unified_full_manual_gather(self, key, tensor, split_dim, tp_size):
+        shards = []
+        for rank in range(tp_size):
+            shard = self.store.get_tensor_with_tp(
+                key, tp_rank=rank, tp_size=tp_size, split_dim=split_dim
+            )
+            self.assertIsNotNone(shard)
+            shards.append(shard)
+        reconstructed = torch.cat(shards, dim=split_dim)
+        self.assertTrue(torch.equal(reconstructed, tensor))
+        return reconstructed
+
+    def _run_unified_full_into(self, key, target, buffer_ptr, buffer_spacing, tensor):
+        reconstructed = self.store.get_tensor_with_parallelism_into(
+            key, buffer_ptr, buffer_spacing, target
+        )
+        self.assertIsNotNone(reconstructed)
+        self.assertTrue(torch.equal(reconstructed, tensor))
+        return reconstructed
+
     def setUp(self):
         """Benchmark-specific setUp."""
         # 1. Call parent setUp to clean the store (remove_all)
@@ -829,6 +2349,68 @@ class TestMooncakeBenchmark(MooncakeTestBase):
         self._print_perf(f"TP Batch Pub (TP={tp_size})", put_times)
         self._print_perf(f"TP Batch Get (TP={tp_size})", get_times)
 
+    def test_benchmark_07_unified_full_reconstruction(self):
+        require_unified_parallelism_api(self)
+        key = "bench_unified_full"
+        tensor = make_deterministic_tensor((4, 8))
+        tp_size = 4
+        split_dim = 1
+        full_target = make_read_target("full", build_tp_parallelism(tp_size, split_dim, rank=1))
+        manual_times = []
+        unified_times = []
+
+        print(f"--- Running Unified Full Reconstruction Benchmark (TP={tp_size}) ---")
+        for _ in range(self.BENCH_ITERATIONS):
+            self.store.remove_all()
+            rc = put_uniform_full_tensor_with_unified_tp(
+                self.store, key, tensor, tp_size, split_dim
+            )
+            self.assertEqual(rc, 0)
+
+            t0 = time.perf_counter()
+            unified_tensor = self.store.get_tensor_with_parallelism(key, full_target)
+            unified_times.append(time.perf_counter() - t0)
+            self.assertIsNotNone(unified_tensor)
+            self.assertTrue(torch.equal(unified_tensor, tensor))
+
+            t0 = time.perf_counter()
+            self._run_unified_full_manual_gather(key, tensor, split_dim, tp_size)
+            manual_times.append(time.perf_counter() - t0)
+
+        self.total_bits = tensor.numel() * tensor.element_size() * 8
+        self._print_perf("Unified Full Get", unified_times)
+        self._print_perf("Manual TP Gather", manual_times)
+
+    def test_benchmark_08_unified_full_into(self):
+        require_unified_parallelism_api(self)
+        key = "bench_unified_full_into"
+        tensor = self.tensors[0]
+        tp_size = 4
+        split_dim = 0
+        full_target = make_read_target("full", build_tp_parallelism(tp_size, split_dim, rank=2))
+        buffer_spacing = serialized_tensor_size(tensor) + 1 * 1024 * 1024
+        buffer = (ctypes.c_ubyte * buffer_spacing)()
+        buffer_ptr = ctypes.addressof(buffer)
+        self.assertEqual(self.store.register_buffer(buffer_ptr, buffer_spacing), 0)
+        into_times = []
+        try:
+            print(f"--- Running Unified Full Into Benchmark (TP={tp_size}) ---")
+            for _ in range(self.BENCH_ITERATIONS):
+                self.store.remove_all()
+                rc = put_uniform_full_tensor_with_unified_tp(
+                    self.store, key, tensor, tp_size, split_dim
+                )
+                self.assertEqual(rc, 0)
+
+                t0 = time.perf_counter()
+                self._run_unified_full_into(key, full_target, buffer_ptr, buffer_spacing, tensor)
+                into_times.append(time.perf_counter() - t0)
+        finally:
+            self.assertEqual(self.store.unregister_buffer(buffer_ptr), 0)
+
+        self.total_bits = tensor.numel() * tensor.element_size() * 8
+        self._print_perf("Unified Full Into", into_times)
+
 # ==========================================
 #  Stress/Concurrency Tests
 # ==========================================
@@ -841,6 +2423,51 @@ class TestMooncakeStress(MooncakeTestBase):
     TOTAL_ITEMS = 800    # Total number of items to process across all threads
     TENSOR_SIZE_MB = 4    # Size per tensor
 
+    def _run_unified_full_worker(self, worker_id, split_dim, tp_size, use_into=False, mixed_reads=False):
+        key = f"stress_unified_full_{worker_id}"
+        if split_dim == 0:
+            tensor = make_deterministic_tensor((8, 8))
+        elif split_dim == 1:
+            tensor = make_deterministic_tensor((4, 8))
+        else:
+            tensor = make_deterministic_tensor((2, 2, 8))
+
+        rc = put_uniform_full_tensor_with_unified_tp(
+            self.store, key, tensor, tp_size, split_dim
+        )
+        if rc != 0:
+            return f"put failed for {key}, rc={rc}"
+
+        full_target = make_read_target("full", build_tp_parallelism(tp_size, split_dim, rank=min(tp_size - 1, 1)))
+        full_tensor = self.store.get_tensor_with_parallelism(key, full_target)
+        if full_tensor is None or not torch.equal(full_tensor, tensor):
+            return f"full reconstruction mismatch for {key}, split_dim={split_dim}, tp_size={tp_size}"
+
+        if mixed_reads:
+            shard_rank = min(tp_size - 1, 2)
+            shard_target = make_read_target("shard", build_tp_parallelism(tp_size, split_dim, rank=shard_rank))
+            shard_tensor = self.store.get_tensor_with_parallelism(key, shard_target)
+            expected_shard = chunk_tensor_for_rank(tensor, tp_size, split_dim, shard_rank)
+            if shard_tensor is None or not torch.equal(shard_tensor, expected_shard):
+                return f"shard mismatch for {key}, split_dim={split_dim}, tp_size={tp_size}, rank={shard_rank}"
+
+        if use_into:
+            buffer_spacing = serialized_tensor_size(tensor) + 64 * 1024
+            buffer = (ctypes.c_ubyte * buffer_spacing)()
+            buffer_ptr = ctypes.addressof(buffer)
+            if self.store.register_buffer(buffer_ptr, buffer_spacing) != 0:
+                return f"register_buffer failed for {key}"
+            try:
+                into_tensor = self.store.get_tensor_with_parallelism_into(
+                    key, buffer_ptr, buffer_spacing, target=full_target
+                )
+                if into_tensor is None or not torch.equal(into_tensor, tensor):
+                    return f"full into mismatch for {key}, split_dim={split_dim}, tp_size={tp_size}"
+            finally:
+                self.store.unregister_buffer(buffer_ptr)
+
+        return None
+
     def _run_stress_worker(self, thread_id, items_per_thread):
         """
         Worker function:
@@ -889,6 +2516,10 @@ class TestMooncakeStress(MooncakeTestBase):
                 if not torch.equal(original_tensor, retrieved_tensor):
                     raise RuntimeError(f"Data Mismatch for {key}!")
 
+                remove_rc = self.store.remove(key, True)
+                if remove_rc != 0:
+                    raise RuntimeError(f"Remove failed for {key}, rc={remove_rc}")
+
                 ops_count += 1
 
         except Exception as e:
@@ -945,6 +2576,31 @@ class TestMooncakeStress(MooncakeTestBase):
             print(f"Avg QPS:        {total_ops / elapsed:.2f} ops/s")
             print(f"Avg Goodput:    {throughput_gbps:.2f} Gbps")
 
+    def test_stress_unified_parallelism_full_reads(self):
+        require_unified_parallelism_api(self)
+        worker_cases = [
+            (0, 0, 2, False, False),
+            (1, 1, 2, False, True),
+            (2, 0, 4, True, False),
+            (3, 1, 4, True, True),
+            (4, 2, 4, False, True),
+            (5, 2, 4, True, True),
+        ]
+        errors = []
+
+        with concurrent.futures.ThreadPoolExecutor(max_workers=len(worker_cases)) as executor:
+            futures = [
+                executor.submit(self._run_unified_full_worker, *case)
+                for case in worker_cases
+            ]
+            for future in concurrent.futures.as_completed(futures):
+                error = future.result()
+                if error:
+                    errors.append(error)
+
+        if errors:
+            self.fail(errors[0])
+
 
 # ==========================================
 #  Data Type & Precision Tests (Full Enum)
diff --git a/scripts/test_upsert_api.py b/scripts/test_upsert_api.py
index fbcb10ea..11aa2abc 100644
--- a/scripts/test_upsert_api.py
+++ b/scripts/test_upsert_api.py
@@ -31,9 +31,11 @@ import os
 import sys
 import time
 import unittest
+from dataclasses import dataclass
 
 import torch
 import numpy as np
+import mooncake.store as mooncake_store
 
 from mooncake.store import MooncakeDistributedStore, ReplicateConfig
 
@@ -41,12 +43,13 @@ from mooncake.store import MooncakeDistributedStore, ReplicateConfig
 #  Constants
 # ==========================================
 
-# Must match C++ TensorMetadata layout
-TENSOR_METADATA_SIZE = 4 + 4 + 8 * 4  # 40 bytes
+# Must match current C++ TensorMetadata layout.
+# TensorObjectHeader (40) + TensorLayoutMetadata (264) = 304 bytes.
+TENSOR_METADATA_SIZE = 304
 
 
 def serialized_tensor_size(tensor):
-    """Size of [TensorMetadata][tensor data] as stored by get_tensor_into."""
+    """Size of [TensorObjectHeader+layout metadata][tensor data]."""
     return TENSOR_METADATA_SIZE + tensor.numel() * tensor.element_size()
 
 
@@ -94,6 +97,51 @@ def tearDownModule():
         GLOBAL_STORE = None
 
 
+def require_unified_parallelism_api(test_case):
+    missing = [
+        name for name in ["ParallelAxis", "TensorParallelism", "ReadTarget"]
+        if not hasattr(mooncake_store, name)
+    ]
+    if missing:
+        test_case.skipTest(
+            "Unified parallelism Python bindings are not available in the installed extension: "
+            + ", ".join(missing)
+        )
+
+
+def make_parallel_axis(kind, rank, size, split_dim=None):
+    axis = mooncake_store.ParallelAxis()
+    axis.kind = kind
+    axis.rank = rank
+    axis.size = size
+    axis.split_dim = split_dim
+    return axis
+
+
+def make_tensor_parallelism(axes):
+    parallelism = mooncake_store.TensorParallelism()
+    parallelism.axes = axes
+    return parallelism
+
+
+def make_read_target(mode, parallelism=None):
+    target = mooncake_store.ReadTarget()
+    target.mode = mode
+    target.parallelism = parallelism
+    return target
+
+
+@dataclass
+class WriterPartition:
+    rank: int
+    size: int
+    split_dim: int
+
+
+def make_writer_partition(rank, size, split_dim):
+    return WriterPartition(rank=rank, size=size, split_dim=split_dim)
+
+
 class UpsertTestBase(unittest.TestCase):
     """Base class that provides self.store and cleans up before each test."""
 
@@ -466,6 +514,555 @@ class TestUpsertPubTensor(UpsertTestBase):
         self.assertTrue(torch.equal(t2, got))
 
 
+class TestUnifiedParallelismUpsert(UpsertTestBase):
+    def test_put_tensor_with_parallelism_rejects_nonuniform_multi_axis(self):
+        require_unified_parallelism_api(self)
+        key = "unified_nonuniform_multi_axis_reject"
+        tensor = torch.arange(6, dtype=torch.float32).view(2, 3).contiguous()
+        parallelism = make_tensor_parallelism([
+            make_parallel_axis("dp", rank=0, size=2),
+            make_parallel_axis("tp", rank=0, size=2, split_dim=1),
+        ])
+
+        rc = self.store.put_tensor_with_parallelism(key, tensor, parallelism)
+        self.assertNotEqual(rc, 0)
+        self.assertIsNone(
+            self.store.get_tensor_with_parallelism(key, make_read_target("shard", parallelism))
+        )
+
+    def test_put_tensor_with_parallelism_from_rejects_nonuniform_multi_axis(self):
+        require_unified_parallelism_api(self)
+        seed_key = "unified_nonuniform_multi_axis_from_seed"
+        key = "unified_nonuniform_multi_axis_from_reject"
+        tensor = torch.arange(6, dtype=torch.float32).view(2, 3).contiguous()
+        self.assertEqual(self.store.put_tensor(seed_key, tensor), 0)
+
+        buffer_size = serialized_tensor_size(tensor)
+        buf = ctypes.create_string_buffer(buffer_size)
+        ptr = ctypes.addressof(buf)
+        self.assertEqual(self.store.register_buffer(ptr, buffer_size), 0)
+        try:
+            got = self.store.get_tensor_into(seed_key, ptr, buffer_size)
+            self.assertIsNotNone(got)
+            parallelism = make_tensor_parallelism([
+                make_parallel_axis("dp", rank=1, size=2),
+                make_parallel_axis("tp", rank=1, size=2, split_dim=1),
+            ])
+            rc = self.store.put_tensor_with_parallelism_from(
+                key, ptr, buffer_size, parallelism
+            )
+            self.assertNotEqual(rc, 0)
+            self.assertIsNone(
+                self.store.get_tensor_with_parallelism(
+                    key, make_read_target("shard", parallelism)
+                )
+            )
+        finally:
+            self.store.unregister_buffer(ptr)
+
+    def test_put_tensor_with_writer_partition_rejects_nonuniform(self):
+        require_unified_parallelism_api(self)
+        key = "writer_partition_nonuniform_reject"
+        tensor = torch.arange(6, dtype=torch.float32).view(2, 3).contiguous()
+
+        rc = self.store.put_tensor_with_parallelism(
+            key,
+            tensor,
+            writer_partition=make_writer_partition(rank=0, size=2, split_dim=1),
+        )
+        self.assertNotEqual(rc, 0)
+        self.assertIsNone(self.store.get_tensor_with_parallelism(key, make_read_target("full")))
+
+    def test_batch_put_tensor_with_parallelism_rejects_nonuniform_and_preserves_uniform_data(self):
+        require_unified_parallelism_api(self)
+        keys = ["uniform_batch_ok", "nonuniform_batch_reject"]
+        tensors = [
+            torch.arange(8, dtype=torch.float32).view(2, 4).contiguous(),
+            torch.arange(6, dtype=torch.float32).view(2, 3).contiguous(),
+        ]
+        parallelisms = [
+            make_tensor_parallelism([
+                make_parallel_axis("dp", rank=0, size=2),
+                make_parallel_axis("tp", rank=1, size=2, split_dim=1),
+            ]),
+            make_tensor_parallelism([
+                make_parallel_axis("dp", rank=1, size=2),
+                make_parallel_axis("tp", rank=0, size=4, split_dim=1),
+            ]),
+        ]
+
+        results = list(
+            self.store.batch_put_tensor_with_parallelism(keys, tensors, parallelisms)
+        )
+        self.assertEqual(results[0], 0)
+        self.assertNotEqual(results[1], 0)
+
+        got_uniform = self.store.get_tensor_with_parallelism(
+            keys[0], make_read_target("shard", parallelisms[0])
+        )
+        self.assertIsNotNone(got_uniform)
+        self.assertTrue(torch.equal(got_uniform, tensors[0].chunk(2, 1)[1]))
+        self.assertIsNone(
+            self.store.get_tensor_with_parallelism(
+                keys[1], make_read_target("shard", parallelisms[1])
+            )
+        )
+
+    def test_upsert_tensor_with_parallelism_multi_axis(self):
+        require_unified_parallelism_api(self)
+        key = "unified_multi_axis_upsert"
+        original = torch.arange(8, dtype=torch.float32).view(2, 4).contiguous()
+        updated = (original + 50).contiguous()
+
+        dp_axis = make_parallel_axis("dp", rank=1, size=2)
+        tp_axis = make_parallel_axis("tp", rank=0, size=2, split_dim=1)
+        parallelism = make_tensor_parallelism([dp_axis, tp_axis])
+
+        for rank in range(2):
+            rank_parallelism = clone_parallelism(parallelism)
+            rank_parallelism.axes[1].rank = rank
+            self.assertEqual(
+                self.store.put_tensor_with_parallelism(key, original, rank_parallelism),
+                0,
+            )
+            self.assertEqual(
+                self.store.upsert_tensor_with_parallelism(key, updated, rank_parallelism),
+                0,
+            )
+
+        target = make_read_target("shard", parallelism)
+        got = self.store.get_tensor_with_parallelism(key, target)
+        self.assertIsNotNone(got)
+        self.assertTrue(torch.equal(got, updated.chunk(2, 1)[0]))
+
+    def test_batch_upsert_tensor_with_parallelism_multi_axis(self):
+        require_unified_parallelism_api(self)
+        keys = ["unified_multi_axis_batch_upsert_0", "unified_multi_axis_batch_upsert_1"]
+        originals = [
+            torch.arange(8, dtype=torch.float32).view(2, 4).contiguous(),
+            torch.arange(8, 16, dtype=torch.float32).view(4, 2).contiguous(),
+        ]
+        updates = [(tensor + 100).contiguous() for tensor in originals]
+
+        base_parallelisms = []
+        dp_tp = make_tensor_parallelism([
+            make_parallel_axis("dp", rank=1, size=2),
+            make_parallel_axis("tp", rank=1, size=2, split_dim=1),
+        ])
+        pp_tp = make_tensor_parallelism([
+            make_parallel_axis("pp", rank=0, size=2),
+            make_parallel_axis("tp", rank=1, size=2, split_dim=0),
+        ])
+        pp_tp.axes[0].stage_id = 3
+        base_parallelisms.extend([dp_tp, pp_tp])
+
+        rank_parallelisms = []
+        for base_parallelism in base_parallelisms:
+            tp_axis = next(axis for axis in base_parallelism.axes if axis.kind == "tp")
+            for rank in range(tp_axis.size):
+                parallelism = clone_parallelism(base_parallelism)
+                for axis in parallelism.axes:
+                    if axis.kind == "tp":
+                        axis.rank = rank
+                        break
+                rank_parallelisms.append(parallelism)
+        batch_keys = [keys[0]] * 2 + [keys[1]] * 2
+        batch_originals = [originals[0]] * 2 + [originals[1]] * 2
+        batch_updates = [updates[0]] * 2 + [updates[1]] * 2
+
+        self.assertEqual(
+            list(self.store.batch_put_tensor_with_parallelism(batch_keys, batch_originals, rank_parallelisms)),
+            [0, 0, 0, 0],
+        )
+        self.assertEqual(
+            list(self.store.batch_upsert_tensor_with_parallelism(batch_keys, batch_updates, rank_parallelisms)),
+            [0, 0, 0, 0],
+        )
+        invalid_results = list(
+            self.store.batch_put_tensor_with_parallelism(
+                keys,
+                originals,
+                parallelisms=base_parallelisms,
+                writer_partitions=[
+                    make_writer_partition(0, 2, 0),
+                    make_writer_partition(1, 2, 0),
+                ],
+            )
+        )
+        self.assertTrue(
+            all(result != 0 for result in invalid_results),
+            f"expected invalid combined routing request: {invalid_results}",
+        )
+
+        targets = []
+        for parallelism in base_parallelisms:
+            target = mooncake_store.ReadTarget()
+            target.mode = "shard"
+            target.parallelism = parallelism
+            targets.append(target)
+        results = self.store.batch_get_tensor_with_parallelism(keys, targets)
+        self.assertEqual(len(results), 2)
+        self.assertTrue(torch.equal(results[0], updates[0].chunk(2, 1)[1]))
+        self.assertTrue(torch.equal(results[1], updates[1].chunk(2, 0)[1]))
+
+    def test_upsert_tensor_with_parallelism_from_multi_axis(self):
+        require_unified_parallelism_api(self)
+        key = "unified_multi_axis_upsert_from"
+        seed_key = "unified_multi_axis_upsert_from_seed"
+        tensor = torch.arange(8, dtype=torch.float32).view(2, 4).contiguous()
+        updated = (tensor + 77).contiguous()
+
+        self.assertEqual(self.store.put_tensor(seed_key, updated), 0)
+        buffer_size = serialized_tensor_size(updated)
+        buf = ctypes.create_string_buffer(buffer_size)
+        ptr = ctypes.addressof(buf)
+        self.assertEqual(self.store.register_buffer(ptr, buffer_size), 0)
+        try:
+            got = self.store.get_tensor_into(seed_key, ptr, buffer_size)
+            self.assertIsNotNone(got)
+
+            dp_axis = make_parallel_axis("dp", rank=0, size=2)
+            tp_axis = make_parallel_axis("tp", rank=1, size=2, split_dim=1)
+            parallelism = make_tensor_parallelism([dp_axis, tp_axis])
+
+            for rank in range(2):
+                rank_parallelism = clone_parallelism(parallelism)
+                rank_parallelism.axes[1].rank = rank
+                self.assertEqual(
+                    self.store.upsert_tensor_with_parallelism_from(
+                        key, ptr, buffer_size, rank_parallelism
+                    ),
+                    0,
+                )
+
+            target = mooncake_store.ReadTarget()
+            target.mode = "shard"
+            target.parallelism = parallelism
+            result = self.store.get_tensor_with_parallelism(key, target)
+            self.assertIsNotNone(result)
+            self.assertTrue(torch.equal(result, updated.chunk(2, 1)[1]))
+        finally:
+            self.store.unregister_buffer(ptr)
+
+    def test_upsert_tensor_with_writer_partition(self):
+        require_unified_parallelism_api(self)
+        key = "writer_partition_upsert"
+        original = torch.arange(8, dtype=torch.float32).view(2, 4).contiguous()
+        updated = (original + 33).contiguous()
+
+        for rank in range(2):
+            self.assertEqual(
+                self.store.put_tensor_with_parallelism(
+                    key,
+                    original,
+                    writer_partition=make_writer_partition(rank, 2, 1),
+                ),
+                0,
+            )
+        for rank in range(2):
+            self.assertEqual(
+                self.store.upsert_tensor_with_parallelism(
+                    key,
+                    updated,
+                    writer_partition=make_writer_partition(rank, 2, 1),
+                ),
+                0,
+            )
+
+        target = make_read_target("full")
+        got = self.store.get_tensor_with_parallelism(key, target)
+        self.assertIsNotNone(got)
+        self.assertTrue(torch.equal(got, updated))
+
+    def test_upsert_tensor_with_writer_partition_from(self):
+        require_unified_parallelism_api(self)
+        key = "writer_partition_upsert_from"
+        seed_key = "writer_partition_upsert_from_seed"
+        tensor = torch.arange(16, dtype=torch.float32).view(4, 4).contiguous()
+
+        self.assertEqual(self.store.put_tensor(seed_key, tensor), 0)
+        for rank in range(4):
+            self.assertEqual(
+                self.store.put_tensor_with_parallelism(
+                    key,
+                    tensor,
+                    writer_partition=make_writer_partition(rank, 4, 0),
+                ),
+                0,
+            )
+        buffer_size = serialized_tensor_size(tensor)
+        buf = ctypes.create_string_buffer(buffer_size)
+        ptr = ctypes.addressof(buf)
+        self.assertEqual(self.store.register_buffer(ptr, buffer_size), 0)
+        try:
+            got = self.store.get_tensor_into(seed_key, ptr, buffer_size)
+            self.assertIsNotNone(got)
+
+            for rank in range(4):
+                self.assertEqual(
+                    self.store.upsert_tensor_with_parallelism_from(
+                        key,
+                        ptr,
+                        buffer_size,
+                        writer_partition=make_writer_partition(rank, 4, 0),
+                    ),
+                    0,
+                )
+
+            target = mooncake_store.ReadTarget()
+            target.mode = "full"
+            target.parallelism = None
+            result = self.store.get_tensor_with_parallelism(key, target)
+            self.assertIsNotNone(result)
+            self.assertTrue(torch.equal(result, tensor))
+        finally:
+            self.store.unregister_buffer(ptr)
+
+    def test_batch_upsert_tensor_with_writer_partitions(self):
+        require_unified_parallelism_api(self)
+        keys = ["writer_partition_batch_upsert_0", "writer_partition_batch_upsert_1"]
+        originals = [
+            torch.arange(8, dtype=torch.float32).view(2, 4).contiguous(),
+            torch.arange(8, 16, dtype=torch.float32).view(2, 4).contiguous(),
+        ]
+        updates = [(tensor + 100).contiguous() for tensor in originals]
+
+        for rank in range(2):
+            writer_partitions = [
+                make_writer_partition(rank=rank, size=2, split_dim=1),
+                make_writer_partition(rank=rank, size=2, split_dim=1),
+            ]
+            self.assertEqual(
+                list(
+                    self.store.batch_put_tensor_with_parallelism(
+                        keys,
+                        originals,
+                        writer_partitions=writer_partitions,
+                    )
+                ),
+                [0, 0],
+            )
+            self.assertEqual(
+                list(
+                    self.store.batch_upsert_tensor_with_parallelism(
+                        keys,
+                        updates,
+                        writer_partitions=writer_partitions,
+                    )
+                ),
+                [0, 0],
+            )
+
+        results = self.store.batch_get_tensor_with_parallelism(
+            keys, [make_read_target("full"), make_read_target("full")]
+        )
+        self.assertEqual(len(results), 2)
+        self.assertTrue(torch.equal(results[0], updates[0]))
+        self.assertTrue(torch.equal(results[1], updates[1]))
+
+    def test_batch_upsert_tensor_with_writer_partitions_from(self):
+        require_unified_parallelism_api(self)
+        keys = ["writer_partition_batch_upsert_from_0", "writer_partition_batch_upsert_from_1"]
+        seed_keys = [f"{key}_seed" for key in keys]
+        tensors = [
+            torch.arange(16, dtype=torch.float32).view(4, 4).contiguous(),
+            torch.arange(16, 32, dtype=torch.float32).view(4, 4).contiguous(),
+        ]
+
+        for seed_key, tensor in zip(seed_keys, tensors):
+            self.assertEqual(self.store.put_tensor(seed_key, tensor), 0)
+
+        for rank in range(4):
+            writer_partitions = [
+                make_writer_partition(rank=rank, size=4, split_dim=0),
+                make_writer_partition(rank=rank, size=4, split_dim=0),
+            ]
+            self.assertEqual(
+                list(
+                    self.store.batch_put_tensor_with_parallelism(
+                        keys,
+                        tensors,
+                        writer_partitions=writer_partitions,
+                    )
+                ),
+                [0, 0],
+            )
+
+        buffer_sizes = [serialized_tensor_size(tensor) for tensor in tensors]
+        backing_buffer = ctypes.create_string_buffer(sum(buffer_sizes))
+        base_ptr = ctypes.addressof(backing_buffer)
+        buffer_ptrs = []
+        offset = 0
+        for size in buffer_sizes:
+            buffer_ptrs.append(base_ptr + offset)
+            offset += size
+        self.assertEqual(self.store.register_buffer(base_ptr, len(backing_buffer)), 0)
+        try:
+            for seed_key, ptr, size in zip(seed_keys, buffer_ptrs, buffer_sizes):
+                got = self.store.get_tensor_into(seed_key, ptr, size)
+                self.assertIsNotNone(got)
+
+            for rank in range(4):
+                writer_partitions = [
+                    make_writer_partition(rank=rank, size=4, split_dim=0),
+                    make_writer_partition(rank=rank, size=4, split_dim=0),
+                ]
+                self.assertEqual(
+                    list(
+                        self.store.batch_upsert_tensor_with_parallelism_from(
+                            keys,
+                            buffer_ptrs,
+                            buffer_sizes,
+                            writer_partitions=writer_partitions,
+                        )
+                    ),
+                    [0, 0],
+                )
+
+            results = self.store.batch_get_tensor_with_parallelism(
+                keys, [make_read_target("full"), make_read_target("full")]
+            )
+            self.assertEqual(len(results), 2)
+            self.assertTrue(torch.equal(results[0], tensors[0]))
+            self.assertTrue(torch.equal(results[1], tensors[1]))
+        finally:
+            self.store.unregister_buffer(base_ptr)
+
+
+def clone_parallelism(parallelism):
+    cloned_axes = []
+    for axis in parallelism.axes:
+        cloned_axis = make_parallel_axis(
+            axis.kind,
+            rank=axis.rank,
+            size=axis.size,
+            split_dim=axis.split_dim,
+        )
+        if getattr(axis, "expert_id", None) is not None:
+            cloned_axis.expert_id = axis.expert_id
+        if getattr(axis, "stage_id", None) is not None:
+            cloned_axis.stage_id = axis.stage_id
+        cloned_axes.append(cloned_axis)
+    return make_tensor_parallelism(cloned_axes)
+
+
+class TestUpsertParallelTensor(UpsertTestBase):
+    """Parallelism-based tensor upsert coverage for TP cases."""
+
+    def test_upsert_tensor_with_parallelism_tp(self):
+        require_unified_parallelism_api(self)
+        tp_size = 2
+        split_dim = 1
+        key = "tp_upsert_single"
+        original = torch.arange(6, dtype=torch.float32).view(2, 3).contiguous()
+        updated = (original + 100).contiguous()
+        original_shard = original.chunk(tp_size, split_dim)[0].contiguous()
+        updated_shard = updated.chunk(tp_size, split_dim)[0].contiguous()
+        parallelism = make_tensor_parallelism(
+            [make_parallel_axis("tp", rank=0, size=tp_size, split_dim=split_dim)]
+        )
+        target = make_read_target("shard", parallelism)
+
+        self.assertEqual(
+            self.store.put_tensor_with_parallelism(key, original_shard, parallelism),
+            0,
+        )
+        self.assertEqual(
+            self.store.upsert_tensor_with_parallelism(key, updated_shard, parallelism),
+            0,
+        )
+
+        got = self.store.get_tensor_with_parallelism(key, target)
+        self.assertTrue(torch.equal(got, updated_shard))
+
+    def test_batch_upsert_tensor_with_parallelism_tp(self):
+        require_unified_parallelism_api(self)
+        tp_size = 2
+        split_dim = 0
+        keys = ["tp_batch_upsert_0", "tp_batch_upsert_1"]
+        originals = [torch.arange(16, dtype=torch.float32).view(4, 4).contiguous(),
+                     torch.arange(16, 32, dtype=torch.float32).view(4, 4).contiguous()]
+        updates = [(t + 1000).contiguous() for t in originals]
+        original_shards = [tensor.chunk(tp_size, split_dim)[0].contiguous() for tensor in originals]
+        update_shards = [tensor.chunk(tp_size, split_dim)[0].contiguous() for tensor in updates]
+        parallelism = make_tensor_parallelism(
+            [make_parallel_axis("tp", rank=0, size=tp_size, split_dim=split_dim)]
+        )
+        targets = [make_read_target("shard", parallelism) for _ in keys]
+
+        self.assertEqual(
+            list(self.store.batch_put_tensor_with_parallelism(keys, original_shards, [parallelism] * len(keys))),
+            [0, 0],
+        )
+        self.assertEqual(
+            list(self.store.batch_upsert_tensor_with_parallelism(keys, update_shards, [parallelism] * len(keys))),
+            [0, 0],
+        )
+
+        got_shards = self.store.batch_get_tensor_with_parallelism(keys, targets)
+        for got, expected in zip(got_shards, update_shards):
+            self.assertTrue(torch.equal(got, expected))
+
+    def test_upsert_tensor_with_parallelism_tp_with_config(self):
+        require_unified_parallelism_api(self)
+        tp_size = 2
+        split_dim = 1
+        key = "tp_upsert_pub_single"
+        config = ReplicateConfig()
+        config.replica_num = 1
+        base = torch.ones(4, 6, dtype=torch.float32).contiguous()
+        updated = torch.full((4, 6), 7.0, dtype=torch.float32).contiguous()
+        base_shard = base.chunk(tp_size, split_dim)[0].contiguous()
+        updated_shard = updated.chunk(tp_size, split_dim)[0].contiguous()
+        parallelism = make_tensor_parallelism(
+            [make_parallel_axis("tp", rank=0, size=tp_size, split_dim=split_dim)]
+        )
+        target = make_read_target("shard", parallelism)
+
+        self.assertEqual(
+            self.store.put_tensor_with_parallelism(key, base_shard, parallelism, config=config),
+            0,
+        )
+        self.assertEqual(
+            self.store.upsert_tensor_with_parallelism(key, updated_shard, parallelism, config=config),
+            0,
+        )
+
+        got = self.store.get_tensor_with_parallelism(key, target)
+        self.assertTrue(torch.equal(got, updated_shard))
+
+    def test_batch_upsert_tensor_with_parallelism_tp_with_config(self):
+        require_unified_parallelism_api(self)
+        tp_size = 2
+        split_dim = 0
+        config = ReplicateConfig()
+        config.replica_num = 1
+        keys = ["tp_upsert_pub_batch_0", "tp_upsert_pub_batch_1"]
+        base = [torch.ones(4, 4, dtype=torch.float32).contiguous(),
+                torch.full((4, 4), 2.0, dtype=torch.float32).contiguous()]
+        updates = [torch.full((4, 4), 9.0, dtype=torch.float32).contiguous(),
+                   torch.full((4, 4), 11.0, dtype=torch.float32).contiguous()]
+        base_shards = [tensor.chunk(tp_size, split_dim)[0].contiguous() for tensor in base]
+        update_shards = [tensor.chunk(tp_size, split_dim)[0].contiguous() for tensor in updates]
+        parallelism = make_tensor_parallelism(
+            [make_parallel_axis("tp", rank=0, size=tp_size, split_dim=split_dim)]
+        )
+        targets = [make_read_target("shard", parallelism) for _ in keys]
+
+        self.assertEqual(
+            list(self.store.batch_put_tensor_with_parallelism(keys, base_shards, [parallelism] * len(keys), config=config)),
+            [0, 0],
+        )
+        self.assertEqual(
+            list(self.store.batch_upsert_tensor_with_parallelism(keys, update_shards, [parallelism] * len(keys), config=config)),
+            [0, 0],
+        )
+
+        got_shards = self.store.batch_get_tensor_with_parallelism(keys, targets)
+        for got, expected in zip(got_shards, update_shards):
+            self.assertTrue(torch.equal(got, expected))
+
+
 # ==========================================
 #  6. Mixed Scenarios
 # ==========================================
@@ -522,7 +1119,9 @@ if __name__ == "__main__":
     print("   Raw bytes:  upsert, upsert_parts, upsert_batch,")
     print("               upsert_from, batch_upsert_from")
     print("   Tensor:     upsert_tensor, batch_upsert_tensor,")
-    print("               upsert_tensor_from, batch_upsert_tensor_from")
+    print("               upsert_tensor_from, batch_upsert_tensor_from,")
+    print("               upsert_tensor_with_parallelism,")
+    print("               batch_upsert_tensor_with_parallelism")
     print("   Pub tensor: upsert_pub_tensor, batch_upsert_pub_tensor")
     print("   Mixed:      cross-interface, sequential, size changes")
     unittest.main(verbosity=2)
diff --git a/scripts/tone_tests/python/test_hicache_storage_mooncake_backend.py b/scripts/tone_tests/python/test_hicache_storage_mooncake_backend.py
index 4cc35fe8..e3d751c3 100644
--- a/scripts/tone_tests/python/test_hicache_storage_mooncake_backend.py
+++ b/scripts/tone_tests/python/test_hicache_storage_mooncake_backend.py
@@ -258,6 +258,9 @@ class TestMooncakeBackendMLAModel(
         server_args, env_vars = super()._get_additional_server_args_and_env()
         server_args["--hicache-mem-layout"] = "page_first"
         server_args["--tp-size"] = 2
+        # Keep arena coverage explicit in the self-hosted integration suite.
+        env_vars.pop("MC_DISABLE_MMAP_ARENA", None)
+        env_vars["MC_MMAP_ARENA_POOL_SIZE"] = "8gb"
         return server_args, env_vars
 
 
@@ -274,6 +277,10 @@ class TestMooncakeBackendAccuracy(
         server_args["--tp-size"] = 2
         server_args["--hicache-mem-layout"] = "page_first_direct"
         server_args["--hicache-io-backend"] = "direct"
+        # Explicitly cover the allocator-disabled fallback path as part of the
+        # self-hosted HiCache integration suite.
+        env_vars.pop("MC_MMAP_ARENA_POOL_SIZE", None)
+        env_vars["MC_DISABLE_MMAP_ARENA"] = "1"
         return server_args, env_vars
 
     def test_eval_accuracy(self):